vscode: portals in vscode

frontend/src/components/ui/context-menu.tsx

@@ -10,7 +10,11 @@ import type { VariantProps } from "class-variance-authority";
 import * as React from "react";
 import { StyleNamespace } from "@/theme/namespace";
 import { cn } from "@/utils/cn";
-import { withFullScreenAsRoot } from "./fullscreen";
+import {
+  MAX_HEIGHT_OFFSET,
+  withFullScreenAsRoot,
+  withSmartCollisionBoundary,
+} from "./fullscreen";
 import {
   MenuShortcut,
   menuContentCommon,
@@ -25,6 +29,12 @@ const ContextMenu = ContextMenuPrimitive.Root;
 const ContextMenuTrigger = ContextMenuPrimitive.Trigger;
 const ContextMenuGroup = ContextMenuPrimitive.Group;
 const ContextMenuPortal = withFullScreenAsRoot(ContextMenuPrimitive.Portal);
+const InternalContextMenuContent = withSmartCollisionBoundary(
+  ContextMenuPrimitive.Content,
+);
+const InternalContextMenuSubContent = withSmartCollisionBoundary(
+  ContextMenuPrimitive.SubContent,
+);
 const ContextMenuSub = ContextMenuPrimitive.Sub;
 const ContextMenuRadioGroup = ContextMenuPrimitive.RadioGroup;
 
@@ -52,7 +62,7 @@ const ContextMenuSubContent = React.forwardRef<
   React.ElementRef<typeof ContextMenuPrimitive.SubContent>,
   React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.SubContent>
 >(({ className, ...props }, ref) => (
-  <ContextMenuPrimitive.SubContent
+  <InternalContextMenuSubContent
     ref={ref}
     className={cn(
       menuContentCommon({ subcontent: true }),
@@ -72,7 +82,7 @@ const ContextMenuContent = React.forwardRef<
 >(({ className, scrollable = true, ...props }, ref) => (
   <ContextMenuPortal>
     <StyleNamespace>
-      <ContextMenuPrimitive.Content
+      <InternalContextMenuContent
         ref={ref}
         className={cn(
           menuContentCommon(),
@@ -83,7 +93,7 @@ const ContextMenuContent = React.forwardRef<
         style={{
           ...props.style,
           maxHeight: scrollable
-            ? "calc(var(--radix-context-menu-content-available-height) - 30px)"
+            ? `calc(var(--radix-context-menu-content-available-height) - ${MAX_HEIGHT_OFFSET}px)`
             : undefined,
         }}
         {...props}

frontend/src/components/ui/dropdown-menu.tsx

@@ -5,7 +5,11 @@ import { Check, ChevronRight, Circle } from "lucide-react";
 import React from "react";
 import { StyleNamespace } from "@/theme/namespace";
 import { cn } from "@/utils/cn";
-import { withFullScreenAsRoot } from "./fullscreen";
+import {
+  MAX_HEIGHT_OFFSET,
+  withFullScreenAsRoot,
+  withSmartCollisionBoundary,
+} from "./fullscreen";
 import {
   MenuShortcut,
   menuContentCommon,
@@ -21,6 +25,12 @@ const DropdownMenu = DropdownMenuPrimitive.Root;
 const DropdownMenuTrigger = DropdownMenuPrimitive.Trigger;
 const DropdownMenuGroup = DropdownMenuPrimitive.Group;
 const DropdownMenuPortal = withFullScreenAsRoot(DropdownMenuPrimitive.Portal);
+const InternalDropdownMenuContent = withSmartCollisionBoundary(
+  DropdownMenuPrimitive.Content,
+);
+const InternalDropdownMenuSubContent = withSmartCollisionBoundary(
+  DropdownMenuPrimitive.SubContent,
+);
 const DropdownMenuSub = DropdownMenuPrimitive.Sub;
 const DropdownMenuRadioGroup = DropdownMenuPrimitive.RadioGroup;
 
@@ -47,7 +57,7 @@ const DropdownMenuSubContent = React.forwardRef<
   React.ComponentRef<typeof DropdownMenuPrimitive.SubContent>,
   React.ComponentPropsWithoutRef<typeof DropdownMenuPrimitive.SubContent>
 >(({ className, ...props }, ref) => (
-  <DropdownMenuPrimitive.SubContent
+  <InternalDropdownMenuSubContent
     ref={ref}
     className={cn(
       menuContentCommon({ subcontent: true }),
@@ -68,7 +78,7 @@ const DropdownMenuContent = React.forwardRef<
 >(({ className, scrollable = true, sideOffset = 4, ...props }, ref) => (
   <DropdownMenuPortal>
     <StyleNamespace>
-      <DropdownMenuPrimitive.Content
+      <InternalDropdownMenuContent
         ref={ref}
         sideOffset={sideOffset}
         className={cn(
@@ -80,7 +90,7 @@ const DropdownMenuContent = React.forwardRef<
         style={{
           ...props.style,
           maxHeight: scrollable
-            ? "calc(var(--radix-dropdown-menu-content-available-height) - 30px)"
+            ? `calc(var(--radix-dropdown-menu-content-available-height) - ${MAX_HEIGHT_OFFSET}px)`
             : undefined,
         }}
         {...props}

frontend/src/components/ui/fullscreen.tsx

@@ -1,8 +1,16 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import { useState } from "react";
+import type { PopoverContentProps } from "@radix-ui/react-popover";
+import React, { useState } from "react";
+import { isInVscodeExtension } from "@/core/vscode/is-in-vscode";
 import { useEventListener } from "@/hooks/useEventListener";
 
+const VSCODE_OUTPUT_CONTAINER_SELECTOR = "[data-vscode-output-container]";
+
+// vscode has smaller viewport so we need all the max-height we can get.
+// Otherwise, we give a 30px buffer to the max-height.
+export const MAX_HEIGHT_OFFSET = isInVscodeExtension() ? 0 : 30;
+
 /**
  * Get the full screen element if we are in full screen mode
  */
@@ -25,14 +33,120 @@ export function withFullScreenAsRoot<
     container?: Element | DocumentFragment | null;
   },
 >(Component: React.ComponentType<T>) {
+  const FindClosestVscodeOutputContainer = (props: T) => {
+    const [closest, setClosest] = React.useState<Element | null>(null);
+    const el = React.useRef<HTMLDivElement>(null);
+
+    React.useLayoutEffect(() => {
+      if (!el.current) {
+        return;
+      }
+
+      const found = closestThroughShadowDOMs(
+        el.current,
+        VSCODE_OUTPUT_CONTAINER_SELECTOR,
+      );
+      setClosest(found);
+    }, []);
+
+    return (
+      <>
+        <div ref={el} className="contents invisible" />
+        <Component {...props} container={closest} />
+      </>
+    );
+  };
+
   const Comp = (props: T) => {
     const fullScreenElement = useFullScreenElement();
+
+    // If we are in the VSCode extension, we use the VSCode output container
+    const vscodeOutputContainer = isInVscodeExtension();
+    if (vscodeOutputContainer) {
+      return <FindClosestVscodeOutputContainer {...props} />;
+    }
+
     if (!fullScreenElement) {
       return <Component {...props} />;
     }
+
     return <Component {...props} container={fullScreenElement} />;
   };
 
   Comp.displayName = Component.displayName;
   return Comp;
 }
+
+/**
+ * HOC wrapping a PortalContent component to set a better collision boundary,
+ * when inside vscode.
+ */
+export function withSmartCollisionBoundary<
+  T extends {
+    collisionBoundary?: PopoverContentProps["collisionBoundary"];
+  },
+>(Component: React.ComponentType<T>) {
+  const FindClosestVscodeOutputContainer = (props: T) => {
+    const [closest, setClosest] = React.useState<Element | null>(null);
+    const el = React.useRef<HTMLDivElement>(null);
+
+    React.useLayoutEffect(() => {
+      if (!el.current) {
+        return;
+      }
+
+      const found = closestThroughShadowDOMs(
+        el.current,
+        VSCODE_OUTPUT_CONTAINER_SELECTOR,
+      );
+      setClosest(found);
+    }, []);
+
+    return (
+      <>
+        <div ref={el} className="contents invisible" />
+        <Component {...props} collisionBoundary={closest} />
+      </>
+    );
+  };
+
+  const Comp = (props: T) => {
+    // If we are in the VSCode extension, we use the VSCode output container
+    const vscodeOutputContainer = isInVscodeExtension();
+    if (vscodeOutputContainer) {
+      return <FindClosestVscodeOutputContainer {...props} />;
+    }
+
+    return <Component {...props} />;
+  };
+
+  Comp.displayName = Component.displayName;
+  return Comp;
+}
+
+/**
+ * Find the closest element (with .closest), but through shadow DOMs.
+ */
+function closestThroughShadowDOMs(
+  element: Element,
+  selector: string,
+): Element | null {
+  let currentElement: Element | null = element;
+
+  while (currentElement) {
+    const cellElement = currentElement.closest(selector);
+    if (cellElement) {
+      return cellElement;
+    }
+
+    const root = currentElement.getRootNode();
+    currentElement =
+      root instanceof ShadowRoot ? root.host : currentElement.parentElement;
+
+    if (currentElement === root) {
+      break;
+    }
+  }
+
+  return null;
+}

frontend/src/components/ui/popover.tsx

@@ -4,14 +4,22 @@ import * as PopoverPrimitive from "@radix-ui/react-popover";
 import * as React from "react";
 import { StyleNamespace } from "@/theme/namespace";
 import { cn } from "@/utils/cn";
-import { withFullScreenAsRoot } from "./fullscreen";
+import {
+  MAX_HEIGHT_OFFSET,
+  withFullScreenAsRoot,
+  withSmartCollisionBoundary,
+} from "./fullscreen";
 
 const Popover = PopoverPrimitive.Root;
 
 const PopoverTrigger = PopoverPrimitive.Trigger;
 const PopoverPortal = withFullScreenAsRoot(PopoverPrimitive.Portal);
 const PopoverClose = PopoverPrimitive.Close;
 
+const InternalPopoverContent = withSmartCollisionBoundary(
+  PopoverPrimitive.Content,
+);
+
 const PopoverContent = React.forwardRef<
   React.ElementRef<typeof PopoverPrimitive.Content>,
   React.ComponentPropsWithoutRef<typeof PopoverPrimitive.Content> & {
@@ -32,7 +40,7 @@ const PopoverContent = React.forwardRef<
   ) => {
     const content = (
       <StyleNamespace>
-        <PopoverPrimitive.Content
+        <InternalPopoverContent
           ref={ref}
           align={align}
           sideOffset={sideOffset}
@@ -44,7 +52,7 @@ const PopoverContent = React.forwardRef<
           style={{
             ...props.style,
             maxHeight: scrollable
-              ? "calc(var(--radix-popover-content-available-height) - 30px)"
+              ? `calc(var(--radix-popover-content-available-height) - ${MAX_HEIGHT_OFFSET}px)`
               : undefined,
           }}
           {...props}

frontend/src/components/ui/range-slider.tsx

@@ -1,14 +1,14 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import * as SliderPrimitive from "@radix-ui/react-slider";
-import { TooltipPortal } from "@radix-ui/react-tooltip";
 import * as React from "react";
 import { useLocale } from "react-aria";
 import { cn } from "@/utils/cn";
 import { prettyScientificNumber } from "@/utils/numbers";
 import { useBoolean } from "../../hooks/useBoolean";
 import {
   TooltipContent,
+  TooltipPortal,
   TooltipProvider,
   TooltipRoot,
   TooltipTrigger,

frontend/src/components/ui/select.tsx

@@ -12,7 +12,7 @@ import {
 import * as React from "react";
 import { StyleNamespace } from "@/theme/namespace";
 import { cn } from "@/utils/cn";
-import { withFullScreenAsRoot } from "./fullscreen";
+import { withFullScreenAsRoot, withSmartCollisionBoundary } from "./fullscreen";
 import { MENU_ITEM_DISABLED } from "./menu-items";
 import { selectStyles } from "./native-select";
 
@@ -60,13 +60,17 @@ const SelectTrigger = React.forwardRef<
 );
 SelectTrigger.displayName = SelectPrimitive.Trigger.displayName;
 
+const InternalSelectContent = withSmartCollisionBoundary(
+  SelectPrimitive.Content,
+);
+
 const SelectContent = React.forwardRef<
   React.ElementRef<typeof SelectPrimitive.Content>,
   React.ComponentPropsWithoutRef<typeof SelectPrimitive.Content>
 >(({ className, children, position = "popper", ...props }, ref) => (
   <SelectPortal>
     <StyleNamespace>
-      <SelectPrimitive.Content
+      <InternalSelectContent
         ref={ref}
         className={cn(
           "max-h-[300px] relative z-50 min-w-32 overflow-hidden rounded-md border bg-popover text-popover-foreground shadow-md data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2",
@@ -94,7 +98,7 @@ const SelectContent = React.forwardRef<
         <SelectPrimitive.ScrollDownButton className="flex items-center justify-center h-[20px] bg-background text-muted-foreground cursor-default">
           <ChevronDownIcon className="h-4 w-4 opacity-50" />
         </SelectPrimitive.ScrollDownButton>
-      </SelectPrimitive.Content>
+      </InternalSelectContent>
     </StyleNamespace>
   </SelectPortal>
 ));

frontend/src/components/ui/slider.tsx

@@ -1,14 +1,14 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import * as SliderPrimitive from "@radix-ui/react-slider";
-import { TooltipPortal } from "@radix-ui/react-tooltip";
 import * as React from "react";
 import { useLocale } from "react-aria";
 import { cn } from "@/utils/cn";
 import { prettyScientificNumber } from "@/utils/numbers";
 import { useBoolean } from "../../hooks/useBoolean";
 import {
   TooltipContent,
+  TooltipPortal,
   TooltipProvider,
   TooltipRoot,
   TooltipTrigger,