@marimo-team/frontend 0.24.1-dev83 → 0.24.1-dev85

package/src/components/editor/file-tree/requesting-tree.tsx

@@ -5,15 +5,24 @@ import { toast } from "@/components/ui/use-toast";
 import type {
   EditRequests,
   FileInfo,
+  FileRoot,
   FileUpdateResponse,
 } from "@/core/network/types";
 import { prettyError } from "@/utils/errors";
 import { Functions } from "@/utils/functions";
 import { type FilePath, PathBuilder } from "@/utils/paths";
-import { resolvePaths } from "@/utils/pathUtils";
 import { mapWithConcurrency } from "@/utils/semaphore";
 
 const FILE_OP_CONCURRENCY = 5;
+const WINDOWS_DRIVE_PATH = /^[A-Za-z]:[/\\]/;
+const WINDOWS_UNC_PATH = /^\\\\/;
+
+export type FileTreeNode = Omit<FileInfo, "children"> & {
+  children: FileTreeNode[];
+  isRoot: boolean;
+  rootPath: string;
+  isPrimaryRoot: boolean;
+};
 
 /**
  * Normalized result of a file mutation: the server response when successful,
@@ -35,8 +44,9 @@ export function handleFileResponse(
 }
 
 export class RequestingTree {
-  private delegate = new SimpleTree<FileInfo>([]);
-  private callbacks: {
+  private delegate = new SimpleTree<FileTreeNode>([]);
+  private readonly callbacks: {
+    getRoots: EditRequests["getFileRoots"];
     listFiles: EditRequests["sendListFiles"];
     createFileOrFolder: EditRequests["sendCreateFileOrFolder"];
     deleteFileOrFolder: EditRequests["sendDeleteFileOrFolder"];
@@ -45,6 +55,7 @@ export class RequestingTree {
   };
 
   constructor(callbacks: {
+    getRoots: EditRequests["getFileRoots"];
     listFiles: EditRequests["sendListFiles"];
     createFileOrFolder: EditRequests["sendCreateFileOrFolder"];
     deleteFileOrFolder: EditRequests["sendDeleteFileOrFolder"];
@@ -54,18 +65,35 @@ export class RequestingTree {
     this.callbacks = callbacks;
   }
 
-  private rootPath: FilePath = "" as FilePath;
-  private onChange: (data: FileInfo[]) => void = Functions.NOOP;
+  private roots: FileRoot[] = [];
+  private onChange: (data: FileTreeNode[]) => void = Functions.NOOP;
   private path = new PathBuilder("/");
 
-  initialize = async (onChange: (data: FileInfo[]) => void): Promise<void> => {
+  initialize = async (
+    onChange: (data: FileTreeNode[]) => void,
+  ): Promise<void> => {
     this.onChange = onChange;
     if (this.delegate.data.length === 0) {
       try {
-        const data = await this.callbacks.listFiles({ path: this.rootPath });
-        this.delegate = new SimpleTree(data.files);
-        this.rootPath = data.root as FilePath;
-        this.path = PathBuilder.guessDeliminator(data.root);
+        const { roots } = await this.callbacks.getRoots();
+        const primaryRoot = roots.find((root) => root.isPrimary);
+        if (!primaryRoot) {
+          throw new Error("File browser response is missing a primary root");
+        }
+        this.roots = roots;
+        this.path = PathBuilder.guessDeliminator(primaryRoot.path);
+        this.delegate = new SimpleTree(roots.map(toRootNode));
+        if (roots.length === 1) {
+          const data = await this.callbacks.listFiles({
+            path: primaryRoot.path,
+          });
+          this.delegate.update({
+            id: this.getPrimaryRootId(),
+            changes: {
+              children: annotateFiles(data.files, primaryRoot),
+            },
+          });
+        }
       } catch (error) {
         toast({
           title: "Failed",
@@ -74,119 +102,91 @@ export class RequestingTree {
       }
     }
 
-    this.onChange(this.delegate.data);
+    this.emitChange();
   };
 
   async expand(id: string): Promise<boolean> {
     const node = this.delegate.find(id);
-    if (!node) {
-      return false;
-    }
-    if (!node.data.isDirectory) {
+    if (!node?.data.isDirectory) {
       return false;
     }
 
-    // We may attempt to load empty directories multiple times
-    // but that is fine
     if (node.children && node.children.length > 0) {
-      // Already loaded
       return true;
     }
 
     const data = await this.callbacks.listFiles({ path: node.data.path });
-    this.delegate.update({ id, changes: { children: data.files } });
-    this.onChange(this.delegate.data);
+    this.delegate.update({
+      id,
+      changes: {
+        children: annotateFiles(data.files, this.getRootForNode(node.data)),
+      },
+    });
+    this.emitChange();
     return true;
   }
 
   async copy(id: string, newName: string): Promise<void> {
-    const node = this.delegate.find(id);
+    const node = this.getMutableNode(id);
     if (!node) {
-      toast({
-        title: "Failed",
-        description: `Node with id ${id} not found in the tree`,
-      });
       return;
     }
-    const { path, newPath } = resolvePaths({
-      path: node.data.path,
-      name: newName,
-      root: this.rootPath,
-    });
-    const parentPath = this.path.dirname(path);
-    const newFile = await this.callbacks
+    const path = node.data.path as FilePath;
+    const parentPath = this.getParentPath(node.data);
+    const newPath = joinPath(parentPath, newName);
+    const result = await this.callbacks
       .copyFileOrFolder({ path, newPath })
       .then(handleFileResponse);
-    if (!newFile?.info) {
-      return;
+    if (result) {
+      await this.refreshPath(parentPath);
     }
-    this.delegate.create({
-      parentId: node.parent?.id ?? null,
-      index: 0,
-      data: newFile.info,
-    });
-    this.onChange(this.delegate.data);
-    // Refresh the parent folder
-    await this.refreshAll([parentPath]);
   }
 
   async rename(id: string, name: string): Promise<void> {
-    const node = this.delegate.find(id);
+    const node = this.getMutableNode(id);
     if (!node) {
-      toast({
-        title: "Failed",
-        description: `Node with id ${id} not found in the tree`,
-      });
       return;
     }
-    const { path, newPath } = resolvePaths({
-      path: node.data.path,
-      name,
-      root: this.rootPath,
-    });
+    const path = node.data.path as FilePath;
+    const parentPath = this.getParentPath(node.data);
+    const newPath = joinPath(parentPath, name);
     const result = await this.callbacks
       .renameFileOrFolder({ path, newPath })
       .then(handleFileResponse);
-    if (!result) {
-      return;
+    if (result) {
+      await this.refreshPath(parentPath);
     }
-
-    this.delegate.update({ id, changes: { name, path: newPath } });
-    this.onChange(this.delegate.data);
-    // Rename all of its children
-    await this.refreshAll([newPath]);
   }
 
   async move(fromIds: string[], parentId: string | null): Promise<void> {
-    const parentPath = parentId
-      ? (this.delegate.find(parentId)?.data.path ?? parentId)
-      : this.rootPath;
+    const targetParentId = parentId ?? this.getPrimaryRootId();
+    const parent = this.delegate.find(targetParentId);
+    if (!parent?.data.isDirectory) {
+      return;
+    }
 
+    const refreshPaths = new Set<string>();
     await mapWithConcurrency(fromIds, FILE_OP_CONCURRENCY, async (id) => {
-      const node = this.delegate.find(id);
+      const node = this.getMutableNode(id, false);
       if (!node) {
         return;
       }
       const originalPath = node.data.path;
-      const newPath = this.path.join(
-        parentPath,
+      const sourceParentPath = this.getParentPath(node.data);
+      const newPath = joinPath(
+        parent.data.path,
         this.path.basename(originalPath as FilePath),
       );
       const result = await this.callbacks
         .renameFileOrFolder({ path: originalPath, newPath })
         .then(handleFileResponse);
-      if (!result) {
-        return;
+      if (result) {
+        refreshPaths.add(parent.data.path);
+        refreshPaths.add(sourceParentPath);
       }
-
-      this.delegate.move({ id, parentId, index: 0 });
-      this.delegate.update({ id, changes: { path: newPath } });
     });
 
-    this.onChange(this.delegate.data);
-
-    // Refresh the parent folder
-    await this.refreshAll([parentPath]);
+    await this.refreshPaths([...refreshPaths]);
   }
 
   async createFile({
@@ -198,140 +198,260 @@ export class RequestingTree {
     parentId: string | null;
     type?: "file" | "notebook";
   }): Promise<void> {
-    const parentPath = parentId
-      ? (this.delegate.find(parentId)?.data.path ?? parentId)
-      : this.rootPath;
-    const newFile = await this.callbacks
-      .createFileOrFolder({ path: parentPath, type: type, name: name })
-      .then(handleFileResponse);
-    if (!newFile?.info) {
+    const parent = this.getParentNode(parentId);
+    if (!parent) {
       return;
     }
-    this.delegate.create({
-      parentId,
-      index: 0,
-      data: newFile.info,
-    });
-    this.onChange(this.delegate.data);
-    // Refresh the parent folder
-    await this.refreshAll([parentPath]);
+    const result = await this.callbacks
+      .createFileOrFolder({
+        path: parent.data.path,
+        type,
+        name,
+      })
+      .then(handleFileResponse);
+    if (result) {
+      await this.refreshPath(parent.data.path as FilePath);
+    }
   }
 
   async createFolder(name: string, parentId: string | null): Promise<void> {
-    const parentPath = parentId
-      ? (this.delegate.find(parentId)?.data.path ?? parentId)
-      : this.rootPath;
-    const newFolder = await this.callbacks
-      .createFileOrFolder({ path: parentPath, type: "directory", name: name })
-      .then(handleFileResponse);
-    if (!newFolder?.info) {
+    const parent = this.getParentNode(parentId);
+    if (!parent) {
       return;
     }
-    this.delegate.create({
-      parentId,
-      index: 0,
-      data: newFolder.info,
-    });
-    this.onChange(this.delegate.data);
-    // Refresh the parent folder
-    await this.refreshAll([parentPath]);
+    const result = await this.callbacks
+      .createFileOrFolder({
+        path: parent.data.path,
+        type: "directory",
+        name,
+      })
+      .then(handleFileResponse);
+    if (result) {
+      await this.refreshPath(parent.data.path as FilePath);
+    }
   }
 
   async delete(id: string): Promise<void> {
-    const node = this.delegate.find(id);
+    const node = this.getMutableNode(id);
     if (!node) {
-      toast({
-        title: "Failed",
-        description: `Node with id ${id} not found in the tree`,
-      });
       return;
     }
-
+    const parentPath = this.getParentPath(node.data);
     const result = await this.callbacks
       .deleteFileOrFolder({ path: node.data.path })
       .then(handleFileResponse);
-    if (!result) {
-      return;
+    if (result) {
+      await this.refreshPath(parentPath);
     }
-    this.delegate.drop({ id });
-    this.onChange(this.delegate.data);
   }
 
   refreshAll = async (ids: string[]): Promise<void> => {
-    // For each open folder, refresh
-    const openFolders = [
-      this.rootPath,
-      ...ids.map((id) => this.delegate.find(id)?.data.path),
-    ].filter(Boolean);
-    // Request open folders with bounded concurrency; swallow per-folder errors.
-    const data = await mapWithConcurrency(
-      openFolders,
-      FILE_OP_CONCURRENCY,
-      (path) =>
-        this.callbacks.listFiles({ path: path }).catch(() => ({ files: [] })),
+    const paths = [
+      ...this.roots.map((root) => root.path),
+      ...ids
+        .map((id) => this.delegate.find(id)?.data.path)
+        .filter((path): path is string => Boolean(path)),
+    ];
+    await this.refreshPaths(paths);
+  };
+
+  refreshPath = async (path: FilePath): Promise<void> => {
+    await this.refreshPaths([path]);
+  };
+
+  public getPrimaryRelativePath = (path: FilePath): FilePath | null => {
+    const primaryRoot = this.getPrimaryRoot();
+    if (!primaryRoot) {
+      return null;
+    }
+    return relativePath(path, primaryRoot.path);
+  };
+
+  public getPrimaryRootPath = (): FilePath => {
+    return (this.getPrimaryRoot()?.path ?? "") as FilePath;
+  };
+
+  public getPrimaryRootId = (): string => {
+    const path = this.getPrimaryRoot()?.path ?? "";
+    return fileTreeNodeId(path, path);
+  };
+
+  public isPrimaryNode = (node: FileTreeNode): boolean => {
+    return node.isPrimaryRoot;
+  };
+
+  public getDisplayPath = (path: FilePath): string => {
+    const root = this.getRootForPath(path);
+    if (!root) {
+      return path;
+    }
+    const relative = relativePath(path, root.path);
+    if (!relative) {
+      return root.name;
+    }
+    return root.isPrimary ? relative : `${root.name}/${relative}`;
+  };
+
+  public isRootPath = (path: FilePath): boolean => {
+    return this.roots.some((root) => pathsEqual(path, root.path));
+  };
+
+  private getPrimaryRoot(): FileRoot | undefined {
+    return this.roots.find((root) => root.isPrimary);
+  }
+
+  private getRootForNode(node: FileTreeNode): FileRoot {
+    return (
+      this.roots.find((root) => pathsEqual(root.path, node.rootPath)) ??
+      this.getPrimaryRoot() ?? {
+        path: node.rootPath,
+        name: node.rootPath,
+        isPrimary: node.isPrimaryRoot,
+      }
     );
+  }
 
-    for (const [idx, openFolder] of openFolders.entries()) {
-      const datum = data[idx];
-      if (openFolder === this.rootPath) {
-        this.delegate = new SimpleTree(datum.files);
-      } else {
-        this.delegate.update({
-          id: openFolder,
-          changes: { children: datum.files },
+  private getRootForPath(path: FilePath): FileRoot | undefined {
+    return this.roots
+      .filter((root) => relativePath(path, root.path) !== null)
+      .toSorted((left, right) => right.path.length - left.path.length)[0];
+  }
+
+  private getParentNode(parentId: string | null) {
+    const resolvedId = parentId ?? this.getPrimaryRootId();
+    const node = this.delegate.find(resolvedId);
+    return node?.data.isDirectory ? node : null;
+  }
+
+  private getParentPath(node: FileTreeNode): FilePath {
+    const parent = this.delegate.find(node.id)?.parent;
+    return (parent?.data.path ?? node.rootPath) as FilePath;
+  }
+
+  private getMutableNode(id: string, showError = true) {
+    const node = this.delegate.find(id);
+    if (!node || node.data.isRoot) {
+      if (showError) {
+        toast({
+          title: "Failed",
+          description: node
+            ? "File browser roots cannot be modified"
+            : `Node with id ${id} not found in the tree`,
         });
       }
+      return null;
     }
+    return node;
+  }
 
-    this.onChange(this.delegate.data);
-  };
-
-  refreshPath = async (path: FilePath): Promise<void> => {
-    const data = await this.callbacks.listFiles({ path }).catch(() => null);
-    if (!data) {
+  private refreshPaths = async (paths: string[]): Promise<void> => {
+    const uniquePaths = [...new Set(paths)];
+    if (uniquePaths.length === 0) {
       return;
     }
+    const results = await mapWithConcurrency(
+      uniquePaths,
+      FILE_OP_CONCURRENCY,
+      (path) => this.callbacks.listFiles({ path }).catch(() => null),
+    );
 
-    if (path === this.rootPath) {
-      this.delegate = new SimpleTree(data.files);
-    } else {
-      const item = findFileByPath(this.delegate.data, path);
-      if (!item?.isDirectory) {
-        return;
+    for (const [index, path] of uniquePaths.entries()) {
+      const result = results[index];
+      if (!result) {
+        continue;
+      }
+      // The same absolute path may appear below multiple overlapping roots.
+      // Refresh every occurrence, while keeping each occurrence in its own ID
+      // namespace and preserving its root metadata.
+      for (const root of this.roots) {
+        const node = this.delegate.find(fileTreeNodeId(root.path, path));
+        if (!node?.data.isDirectory) {
+          continue;
+        }
+        this.delegate.update({
+          id: node.id,
+          changes: {
+            children: annotateFiles(result.files, root),
+          },
+        });
       }
-      this.delegate.update({ id: item.id, changes: { children: data.files } });
     }
-    this.onChange(this.delegate.data);
+    this.emitChange();
   };
 
-  public relativeFromRoot = (path: FilePath): FilePath => {
-    // Add a trailing delimiter to the root path if it doesn't have one
-    const root = this.rootPath.endsWith(this.path.deliminator)
-      ? this.rootPath
-      : `${this.rootPath}${this.path.deliminator}`;
+  private emitChange(): void {
+    const data = this.delegate.data;
+    this.onChange(this.roots.length === 1 ? (data[0]?.children ?? []) : data);
+  }
+}
 
-    if (path.startsWith(root)) {
-      return path.slice(root.length) as FilePath;
-    }
-    return path;
+function toRootNode(root: FileRoot): FileTreeNode {
+  return {
+    id: fileTreeNodeId(root.path, root.path),
+    path: root.path,
+    name: root.name,
+    isDirectory: true,
+    isMarimoFile: false,
+    children: [],
+    isRoot: true,
+    rootPath: root.path,
+    isPrimaryRoot: root.isPrimary,
   };
+}
 
-  public getRootPath = (): FilePath => {
-    return this.rootPath;
-  };
+function annotateFiles(files: FileInfo[], root: FileRoot): FileTreeNode[] {
+  return files.map((file) => ({
+    ...file,
+    id: fileTreeNodeId(root.path, file.path),
+    children: annotateFiles(file.children ?? [], root),
+    isRoot: false,
+    rootPath: root.path,
+    isPrimaryRoot: root.isPrimary,
+  }));
 }
 
-function findFileByPath(list: FileInfo[], path: string): FileInfo | undefined {
-  for (const item of list) {
-    if (item.path === path) {
-      return item;
-    }
-    const child = item.children
-      ? findFileByPath(item.children, path)
-      : undefined;
-    if (child) {
-      return child;
-    }
+/** A stable, globally unique tree ID for a path as viewed from one root. */
+export function fileTreeNodeId(rootPath: string, path: string): string {
+  return `${encodeURIComponent(rootPath)}:${encodeURIComponent(path)}`;
+}
+
+function relativePath(path: string, root: string): FilePath | null {
+  const windowsPath = isWindowsPath(root);
+  const normalizeCase = (value: string) =>
+    windowsPath ? value.toLowerCase() : value;
+  const comparedPath = normalizeCase(path);
+  const comparedRoot = normalizeCase(root);
+  if (comparedPath === comparedRoot) {
+    return "" as FilePath;
+  }
+
+  const delimiter = pathDelimiter(root);
+  const rootWithDelimiter = root.endsWith(delimiter)
+    ? root
+    : `${root}${delimiter}`;
+  const comparedPrefix = normalizeCase(rootWithDelimiter);
+  if (!comparedPath.startsWith(comparedPrefix)) {
+    return null;
+  }
+  return path.slice(rootWithDelimiter.length) as FilePath;
+}
+
+function pathsEqual(left: string, right: string): boolean {
+  return relativePath(left, right) === ("" as FilePath);
+}
+
+function joinPath(parent: string, name: string): FilePath {
+  const delimiter = pathDelimiter(parent);
+  return `${parent}${parent.endsWith(delimiter) ? "" : delimiter}${name}` as FilePath;
+}
+
+function isWindowsPath(path: string): boolean {
+  return WINDOWS_DRIVE_PATH.test(path) || WINDOWS_UNC_PATH.test(path);
+}
+
+function pathDelimiter(path: string): "/" | "\\" {
+  if (WINDOWS_UNC_PATH.test(path) || /^[A-Za-z]:\\/.test(path)) {
+    return "\\";
   }
-  return undefined;
+  return "/";
 }

package/src/components/editor/file-tree/state.tsx

@@ -11,6 +11,7 @@ export const treeAtom = atom<RequestingTree>((get) => {
   const client = get(requestClientAtom);
   invariant(client, "no requestClientAtom set");
   return new RequestingTree({
+    getRoots: client.getFileRoots,
     listFiles: client.sendListFiles,
     createFileOrFolder: client.sendCreateFileOrFolder,
     deleteFileOrFolder: client.sendDeleteFileOrFolder,

package/src/core/config/__tests__/config-schema.test.ts

@@ -77,6 +77,7 @@ test("default UserConfig - empty", () => {
         "theme": "light",
       },
       "experimental": {},
+      "file_browser": {},
       "formatting": {
         "line_length": 79,
       },
@@ -152,6 +153,7 @@ test("default UserConfig - one level", () => {
         "theme": "light",
       },
       "experimental": {},
+      "file_browser": {},
       "formatting": {
         "line_length": 79,
       },
@@ -212,6 +214,19 @@ test("default UserConfig with additional information", () => {
   );
 });
 
+test("UserConfig with file browser folders", () => {
+  const config = UserConfigSchema.parse({
+    file_browser: {
+      folders: [{ path: "/data", name: "Data" }, { path: "/shared" }],
+    },
+  });
+
+  expect(config.file_browser?.folders).toEqual([
+    { path: "/data", name: "Data" },
+    { path: "/shared" },
+  ]);
+});
+
 test("UserConfig with custom_providers", () => {
   const config = UserConfigSchema.parse({
     ai: {

package/src/core/config/config-schema.ts

@@ -214,6 +214,19 @@ export const UserConfigSchema = z
         sql_linter: z.boolean().optional(),
       })
       .prefault(() => ({})),
+    file_browser: z
+      .looseObject({
+        folders: z
+          .array(
+            z.object({
+              path: z.string(),
+              name: z.string().optional(),
+            }),
+          )
+          .optional(),
+      })
+      .optional()
+      .prefault({}),
     sharing: z
       .looseObject({
         html: z.boolean().optional(),
@@ -237,6 +250,7 @@ export const UserConfigSchema = z
     runtime: {},
     display: {},
     diagnostics: {},
+    file_browser: {},
     experimental: {},
     server: {},
     ai: {},
@@ -341,6 +355,7 @@ export function defaultUserConfig(): UserConfig {
     runtime: {},
     display: {},
     diagnostics: {},
+    file_browser: {},
     experimental: {},
     server: {},
     ai: {},