refactor: add checkbox state management and tree feature configuration hooks

apps/app/src/client/components/Sidebar/PageTree/PageTreeSubstance.tsx

@@ -118,6 +118,7 @@ export const PageTreeContent = memo(({ isWipPageShown }: PageTreeContentProps) =
   return (
     <div className="pt-4">
       <SimplifiedItemsTree
+        enableRenaming
         isEnableActions={!isGuestUser}
         isReadOnlyUser={!!isReadOnlyUser}
         isWipPageShown={isWipPageShown}

apps/app/src/features/page-tree/components/SimplifiedItemsTree.tsx

@@ -1,12 +1,5 @@
 import type { FC } from 'react';
-import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import {
-  asyncDataLoaderFeature,
-  checkboxesFeature,
-  hotkeysCoreFeature,
-  renamingFeature,
-  selectionFeature,
-} from '@headless-tree/core';
+import { useMemo } from 'react';
 import { useTree } from '@headless-tree/react';
 import { useVirtualizer } from '@tanstack/react-virtual';
 
@@ -16,28 +9,17 @@ import { useSWRxRootPage } from '~/stores/page-listing';
 import { ROOT_PAGE_VIRTUAL_ID } from '../constants/_inner';
 import {
   useAutoExpandAncestors,
+  useCheckboxChangeNotification,
+  useCheckboxState,
   useDataLoader,
   useExpandParentOnCreate,
   useScrollToSelectedItem,
+  useTreeFeatures,
   useTreeItemHandlers,
+  useTreeRevalidation,
 } from '../hooks/_inner';
 import type { TreeItemProps } from '../interfaces';
 import { useTriggerTreeRebuild } from '../states/_inner';
-import {
-  usePageTreeInformationGeneration,
-  usePageTreeRevalidationEffect,
-} from '../states/page-tree-update';
-
-// Base features for all tree variants
-const BASE_FEATURES = [
-  asyncDataLoaderFeature,
-  selectionFeature,
-  hotkeysCoreFeature,
-  renamingFeature,
-];
-
-// Features with checkboxes support
-const FEATURES_WITH_CHECKBOXES = [...BASE_FEATURES, checkboxesFeature];
 
 // Stable createLoadingItemData function
 const createLoadingItemData = (): IPageForTreeItem => ({
@@ -59,7 +41,8 @@ type Props = {
   CustomTreeItem: React.FunctionComponent<TreeItemProps>;
   estimateTreeItemSize: () => number;
   scrollerElem?: HTMLElement | null;
-  // Checkbox feature options
+  // Feature options
+  enableRenaming?: boolean;
   enableCheckboxes?: boolean;
   initialCheckedItems?: string[];
   onCheckedItemsChange?: (checkedItems: IPageForTreeItem[]) => void;
@@ -75,6 +58,7 @@ export const SimplifiedItemsTree: FC<Props> = (props: Props) => {
     CustomTreeItem,
     estimateTreeItemSize,
     scrollerElem,
+    enableRenaming = false,
     enableCheckboxes = false,
     initialCheckedItems = [],
     onCheckedItemsChange,
@@ -90,10 +74,21 @@ export const SimplifiedItemsTree: FC<Props> = (props: Props) => {
   const dataLoader = useDataLoader(rootPageId, allPagesCount);
 
   // Tree item handlers (rename, create, etc.) with stable callbacks for headless-tree
-  // Note: triggerTreeRebuild is stable (from useSetAtom), so no need for useCallback wrapper
   const { getItemName, isItemFolder, handleRename, creatingParentId } =
     useTreeItemHandlers(triggerTreeRebuild);
 
+  // Configure tree features based on options
+  const features = useTreeFeatures({
+    enableRenaming,
+    enableCheckboxes,
+  });
+
+  // Manage checkbox state (must be called before useTree to get setCheckedItems)
+  const { checkedItemIds, setCheckedItems } = useCheckboxState({
+    enabled: enableCheckboxes,
+    initialCheckedItems,
+  });
+
   // Stable initial state
   // biome-ignore lint/correctness/useExhaustiveDependencies: initialCheckedItems is intentionally not in deps to avoid reinitializing on every change
   const initialState = useMemo(
@@ -104,15 +99,6 @@ export const SimplifiedItemsTree: FC<Props> = (props: Props) => {
     [enableCheckboxes],
   );
 
-  // State to track checked items for re-rendering
-  const [checkedItemIds, setCheckedItemIds] =
-    useState<string[]>(initialCheckedItems);
-
-  // Callback to update checked items state (triggers re-render)
-  const handleSetCheckedItems = useCallback((itemIds: string[]) => {
-    setCheckedItemIds(itemIds);
-  }, []);
-
   const tree = useTree<IPageForTreeItem>({
     rootItemId: ROOT_PAGE_VIRTUAL_ID,
     getItemName,
@@ -121,38 +107,24 @@ export const SimplifiedItemsTree: FC<Props> = (props: Props) => {
     createLoadingItemData,
     dataLoader,
     onRename: handleRename,
-    features: enableCheckboxes ? FEATURES_WITH_CHECKBOXES : BASE_FEATURES,
-    // Checkbox configuration: prevent folder auto-check to avoid selecting all descendants
+    features,
+    // Checkbox configuration
     canCheckFolders: enableCheckboxes,
     propagateCheckedState: false,
-    // Custom setter to track checked items changes
-    setCheckedItems: enableCheckboxes ? handleSetCheckedItems : undefined,
+    setCheckedItems,
   });
 
   // Notify parent when checked items change
-  useEffect(() => {
-    if (!enableCheckboxes || onCheckedItemsChange == null) {
-      return;
-    }
-
-    const checkedPages = checkedItemIds
-      .map((id) => tree.getItemInstance(id)?.getItemData())
-      .filter((page): page is IPageForTreeItem => page != null);
-    onCheckedItemsChange(checkedPages);
-  }, [enableCheckboxes, checkedItemIds, onCheckedItemsChange, tree]);
-
-  // Track local generation number
-  const localGenerationRef = useRef(1);
-  const globalGeneration = usePageTreeInformationGeneration();
-
-  // Refetch data when global generation is updated
-  usePageTreeRevalidationEffect(tree, localGenerationRef.current, {
-    // Update local generation number after revalidation
-    onRevalidated: () => {
-      localGenerationRef.current = globalGeneration;
-    },
+  useCheckboxChangeNotification({
+    enabled: enableCheckboxes,
+    checkedItemIds,
+    tree,
+    onCheckedItemsChange,
   });
 
+  // Handle tree revalidation and items count tracking
+  useTreeRevalidation({ tree, triggerTreeRebuild });
+
   // Expand parent item when page creation is initiated
   useExpandParentOnCreate({
     tree,
@@ -162,18 +134,7 @@ export const SimplifiedItemsTree: FC<Props> = (props: Props) => {
 
   const items = tree.getItems();
 
-  // Track items count to detect when async data loading completes
-  const prevItemsCountRef = useRef(items.length);
-  useEffect(() => {
-    if (items.length !== prevItemsCountRef.current) {
-      prevItemsCountRef.current = items.length;
-      // Trigger re-render when items count changes (e.g., after async load completes)
-      triggerTreeRebuild();
-    }
-  }, [items.length, triggerTreeRebuild]);
-
   // Auto-expand items that are ancestors of targetPath
-  // Note: triggerTreeRebuild is stable, no need for useCallback wrapper
   useAutoExpandAncestors({
     items,
     targetPath,
@@ -206,7 +167,7 @@ export const SimplifiedItemsTree: FC<Props> = (props: Props) => {
           return null;
         }
 
-        const props = item.getProps();
+        const treeItemProps = item.getProps();
 
         return (
           <div
@@ -214,8 +175,8 @@ export const SimplifiedItemsTree: FC<Props> = (props: Props) => {
             data-index={virtualItem.index}
             ref={(node) => {
               virtualizer.measureElement(node);
-              if (node && props.ref) {
-                (props.ref as (node: HTMLElement) => void)(node);
+              if (node && treeItemProps.ref) {
+                (treeItemProps.ref as (node: HTMLElement) => void)(node);
               }
             }}
           >

apps/app/src/features/page-tree/hooks/_inner/index.ts

@@ -1,5 +1,8 @@
 export * from './use-auto-expand-ancestors';
+export * from './use-checkbox-state';
 export * from './use-data-loader';
 export * from './use-expand-parent-on-create';
 export * from './use-scroll-to-selected-item';
+export * from './use-tree-features';
 export * from './use-tree-item-handlers';
+export * from './use-tree-revalidation';

apps/app/src/features/page-tree/hooks/_inner/use-checkbox-state.ts

@@ -0,0 +1,70 @@
+import { useCallback, useEffect, useState } from 'react';
+
+import type { IPageForTreeItem } from '~/interfaces/page';
+
+type TreeInstance = {
+  getItemInstance: (
+    id: string,
+  ) => { getItemData: () => IPageForTreeItem } | undefined;
+};
+
+export type UseCheckboxStateOptions = {
+  enabled: boolean;
+  initialCheckedItems: string[];
+};
+
+export type UseCheckboxStateReturn = {
+  checkedItemIds: string[];
+  setCheckedItems: ((itemIds: string[]) => void) | undefined;
+};
+
+/**
+ * Hook to manage checkbox state for headless-tree.
+ * Provides state tracking and setter callback for checked items.
+ */
+export const useCheckboxState = (
+  options: UseCheckboxStateOptions,
+): UseCheckboxStateReturn => {
+  const { enabled, initialCheckedItems } = options;
+
+  // State to track checked items for re-rendering
+  const [checkedItemIds, setCheckedItemIds] =
+    useState<string[]>(initialCheckedItems);
+
+  // Callback to update checked items state (triggers re-render)
+  const handleSetCheckedItems = useCallback((itemIds: string[]) => {
+    setCheckedItemIds(itemIds);
+  }, []);
+
+  return {
+    checkedItemIds,
+    setCheckedItems: enabled ? handleSetCheckedItems : undefined,
+  };
+};
+
+export type UseCheckboxChangeNotificationOptions = {
+  enabled: boolean;
+  checkedItemIds: string[];
+  tree: TreeInstance;
+  onCheckedItemsChange?: (checkedItems: IPageForTreeItem[]) => void;
+};
+
+/**
+ * Hook to notify parent when checked items change.
+ */
+export const useCheckboxChangeNotification = (
+  options: UseCheckboxChangeNotificationOptions,
+): void => {
+  const { enabled, checkedItemIds, tree, onCheckedItemsChange } = options;
+
+  useEffect(() => {
+    if (!enabled || onCheckedItemsChange == null) {
+      return;
+    }
+
+    const checkedPages = checkedItemIds
+      .map((id) => tree.getItemInstance(id)?.getItemData())
+      .filter((page): page is IPageForTreeItem => page != null);
+    onCheckedItemsChange(checkedPages);
+  }, [enabled, checkedItemIds, onCheckedItemsChange, tree]);
+};

apps/app/src/features/page-tree/hooks/_inner/use-tree-features.ts

@@ -0,0 +1,42 @@
+import { useMemo } from 'react';
+import type { FeatureImplementation } from '@headless-tree/core';
+import {
+  asyncDataLoaderFeature,
+  checkboxesFeature,
+  hotkeysCoreFeature,
+  renamingFeature,
+  selectionFeature,
+} from '@headless-tree/core';
+
+export type UseTreeFeaturesOptions = {
+  enableRenaming?: boolean;
+  enableCheckboxes?: boolean;
+};
+
+/**
+ * Hook to configure tree features based on options.
+ * Returns a stable array of features for use with headless-tree.
+ */
+export const useTreeFeatures = (
+  options: UseTreeFeaturesOptions = {},
+): FeatureImplementation<unknown>[] => {
+  const { enableRenaming = true, enableCheckboxes = false } = options;
+
+  return useMemo(() => {
+    const features: FeatureImplementation<unknown>[] = [
+      asyncDataLoaderFeature,
+      selectionFeature,
+      hotkeysCoreFeature,
+    ];
+
+    if (enableRenaming) {
+      features.push(renamingFeature);
+    }
+
+    if (enableCheckboxes) {
+      features.push(checkboxesFeature);
+    }
+
+    return features;
+  }, [enableRenaming, enableCheckboxes]);
+};

apps/app/src/features/page-tree/hooks/_inner/use-tree-revalidation.ts

@@ -0,0 +1,51 @@
+import { useEffect, useRef } from 'react';
+
+import {
+  usePageTreeInformationGeneration,
+  usePageTreeRevalidationEffect,
+} from '../../states/page-tree-update';
+
+type TreeInstance = {
+  getItems: () => unknown[];
+};
+
+type UseTreeRevalidationOptions = {
+  tree: TreeInstance;
+  triggerTreeRebuild: () => void;
+};
+
+/**
+ * Hook to handle tree revalidation when global generation changes
+ * and track items count changes for async data loading
+ */
+export const useTreeRevalidation = (options: UseTreeRevalidationOptions) => {
+  const { tree, triggerTreeRebuild } = options;
+
+  // Track local generation number
+  const localGenerationRef = useRef(1);
+  const globalGeneration = usePageTreeInformationGeneration();
+
+  // Refetch data when global generation is updated
+  usePageTreeRevalidationEffect(
+    tree as Parameters<typeof usePageTreeRevalidationEffect>[0],
+    localGenerationRef.current,
+    {
+      // Update local generation number after revalidation
+      onRevalidated: () => {
+        localGenerationRef.current = globalGeneration;
+      },
+    },
+  );
+
+  const items = tree.getItems();
+
+  // Track items count to detect when async data loading completes
+  const prevItemsCountRef = useRef(items.length);
+  useEffect(() => {
+    if (items.length !== prevItemsCountRef.current) {
+      prevItemsCountRef.current = items.length;
+      // Trigger re-render when items count changes (e.g., after async load completes)
+      triggerTreeRebuild();
+    }
+  }, [items.length, triggerTreeRebuild]);
+};