Trace view fixes and improvements

apps/webapp/app/components/primitives/ShortcutKey.tsx

@@ -2,8 +2,14 @@ import { Fragment } from "react";
 import { Modifier, ShortcutDefinition } from "~/hooks/useShortcutKeys";
 import { cn } from "~/utils/cn";
 import { useOperatingSystem } from "./OperatingSystemProvider";
+import {
+  ChevronDownIcon,
+  ChevronLeftIcon,
+  ChevronRightIcon,
+  ChevronUpIcon,
+} from "@heroicons/react/20/solid";
 
-const variants = {
+export const variants = {
   small:
     "text-[0.6rem] font-medium min-w-[17px] rounded-[2px] px-1 ml-1 -mr-0.5 grid place-content-center border border-dimmed/40 text-text-dimmed group-hover:text-text-bright/80 group-hover:border-dimmed/60 transition uppercase",
   medium:
@@ -23,7 +29,7 @@ export function ShortcutKey({ shortcut, variant, className }: ShortcutKeyProps)
   const isMac = platform === "mac";
   let relevantShortcut = "mac" in shortcut ? (isMac ? shortcut.mac : shortcut.windows) : shortcut;
   const modifiers = relevantShortcut.modifiers ?? [];
-  const character = keyString(relevantShortcut.key, isMac);
+  const character = keyString(relevantShortcut.key, isMac, variant);
 
   return (
     <span className={cn(variants[variant], className)}>
@@ -35,10 +41,22 @@ export function ShortcutKey({ shortcut, variant, className }: ShortcutKeyProps)
   );
 }
 
-function keyString(key: String, isMac: boolean) {
+function keyString(key: String, isMac: boolean, size: "small" | "medium") {
+  key = key.toLowerCase();
+
+  const className = size === "small" ? "w-2.5 h-4" : "w-3 h-5";
+
   switch (key) {
     case "enter":
       return isMac ? "↵" : key;
+    case "arrowdown":
+      return <ChevronDownIcon className={className} />;
+    case "arrowup":
+      return <ChevronUpIcon className={className} />;
+    case "arrowleft":
+      return <ChevronLeftIcon className={className} />;
+    case "arrowright":
+      return <ChevronRightIcon className={className} />;
     default:
       return key;
   }

apps/webapp/app/components/primitives/TreeView/TreeView.tsx

@@ -1,10 +1,10 @@
 import { VirtualItem, Virtualizer, useVirtualizer } from "@tanstack/react-virtual";
+import { motion } from "framer-motion";
 import { MutableRefObject, RefObject, useCallback, useEffect, useReducer, useRef } from "react";
 import { UnmountClosed } from "react-collapse";
 import { cn } from "~/utils/cn";
 import { NodeState, NodesState, reducer } from "./reducer";
 import { applyFilterToState, concreteStateFromInput, selectedIdFromState } from "./utils";
-import { motion } from "framer-motion";
 
 export type TreeViewProps<TData> = {
   tree: FlatTree<TData>;
@@ -165,6 +165,11 @@ export type UseTreeStateOutput = {
   expandNode: (id: string, scrollToNode?: boolean) => void;
   collapseNode: (id: string) => void;
   toggleExpandNode: (id: string, scrollToNode?: boolean) => void;
+  expandAllBelowDepth: (depth: number) => void;
+  collapseAllBelowDepth: (depth: number) => void;
+  expandLevel: (level: number) => void;
+  collapseLevel: (level: number) => void;
+  toggleExpandLevel: (level: number) => void;
   selectFirstVisibleNode: (scrollToNode?: boolean) => void;
   selectLastVisibleNode: (scrollToNode?: boolean) => void;
   selectNextVisibleNode: (scrollToNode?: boolean) => void;
@@ -333,6 +338,41 @@ export function useTree<TData>({
     [state]
   );
 
+  const expandAllBelowDepth = useCallback(
+    (depth: number) => {
+      dispatch({ type: "EXPAND_ALL_BELOW_DEPTH", payload: { tree, depth } });
+    },
+    [state]
+  );
+
+  const collapseAllBelowDepth = useCallback(
+    (depth: number) => {
+      dispatch({ type: "COLLAPSE_ALL_BELOW_DEPTH", payload: { tree, depth } });
+    },
+    [state]
+  );
+
+  const expandLevel = useCallback(
+    (level: number) => {
+      dispatch({ type: "EXPAND_LEVEL", payload: { tree, level } });
+    },
+    [state]
+  );
+
+  const collapseLevel = useCallback(
+    (level: number) => {
+      dispatch({ type: "COLLAPSE_LEVEL", payload: { tree, level } });
+    },
+    [state]
+  );
+
+  const toggleExpandLevel = useCallback(
+    (level: number) => {
+      dispatch({ type: "TOGGLE_EXPAND_LEVEL", payload: { tree, level } });
+    },
+    [state]
+  );
+
   const getTreeProps = useCallback(() => {
     return {
       role: "tree",
@@ -368,25 +408,48 @@ export function useTree<TData>({
           }
           case "Left":
           case "ArrowLeft": {
+            e.preventDefault();
+
             const selected = selectedIdFromState(state.nodes);
             if (selected) {
               const treeNode = tree.find((node) => node.id === selected);
-              if (treeNode && treeNode.hasChildren && state.nodes[selected].expanded) {
+
+              if (e.altKey) {
+                if (treeNode && treeNode.hasChildren) {
+                  collapseLevel(treeNode.level);
+                }
+                break;
+              }
+
+              const shouldCollapse =
+                treeNode && treeNode.hasChildren && state.nodes[selected].expanded;
+              if (shouldCollapse) {
                 collapseNode(selected);
               } else {
                 selectParentNode(true);
               }
             }
-            e.preventDefault();
+
             break;
           }
           case "Right":
           case "ArrowRight": {
+            e.preventDefault();
+
             const selected = selectedIdFromState(state.nodes);
+
             if (selected) {
+              const treeNode = tree.find((node) => node.id === selected);
+
+              if (e.altKey) {
+                if (treeNode && treeNode.hasChildren) {
+                  expandLevel(treeNode.level);
+                }
+                break;
+              }
+
               expandNode(selected, true);
             }
-            e.preventDefault();
             break;
           }
           case "Escape": {
@@ -427,6 +490,11 @@ export function useTree<TData>({
     expandNode,
     collapseNode,
     toggleExpandNode,
+    expandAllBelowDepth,
+    collapseAllBelowDepth,
+    expandLevel,
+    collapseLevel,
+    toggleExpandLevel,
     selectFirstVisibleNode,
     selectLastVisibleNode,
     selectNextVisibleNode,

apps/webapp/app/components/primitives/TreeView/reducer.ts

@@ -91,6 +91,46 @@ type ToggleExpandNodeAction = {
   } & WithScrollToNode;
 };
 
+type ExpandAllBelowDepthAction = {
+  type: "EXPAND_ALL_BELOW_DEPTH";
+  payload: {
+    depth: number;
+    tree: FlatTree<any>;
+  };
+};
+
+type CollapseAllBelowDepthAction = {
+  type: "COLLAPSE_ALL_BELOW_DEPTH";
+  payload: {
+    depth: number;
+    tree: FlatTree<any>;
+  };
+};
+
+type ExpandLevelAction = {
+  type: "EXPAND_LEVEL";
+  payload: {
+    level: number;
+    tree: FlatTree<any>;
+  };
+};
+
+type CollapseLevelAction = {
+  type: "COLLAPSE_LEVEL";
+  payload: {
+    level: number;
+    tree: FlatTree<any>;
+  };
+};
+
+type ToggleExpandLevelAction = {
+  type: "TOGGLE_EXPAND_LEVEL";
+  payload: {
+    level: number;
+    tree: FlatTree<any>;
+  };
+};
+
 type SelectFirstVisibleNodeAction = {
   type: "SELECT_FIRST_VISIBLE_NODE";
   payload: {
@@ -135,6 +175,11 @@ export type Action =
   | ExpandNodeAction
   | CollapseNodeAction
   | ToggleExpandNodeAction
+  | ExpandAllBelowDepthAction
+  | CollapseAllBelowDepthAction
+  | ExpandLevelAction
+  | CollapseLevelAction
+  | ToggleExpandLevelAction
   | SelectFirstVisibleNodeAction
   | SelectLastVisibleNodeAction
   | SelectNextVisibleNodeAction
@@ -229,6 +274,109 @@ export function reducer(state: TreeState, action: Action): TreeState {
         });
       }
     }
+    case "EXPAND_ALL_BELOW_DEPTH": {
+      const nodesToExpand = action.payload.tree.filter(
+        (n) => n.level >= action.payload.depth && n.hasChildren
+      );
+
+      const newNodes = Object.fromEntries(
+        Object.entries(state.nodes).map(([key, value]) => [
+          key,
+          {
+            ...value,
+            expanded: nodesToExpand.find((n) => n.id === key) ? true : value.expanded,
+          },
+        ])
+      );
+
+      const visibleNodes = applyVisibility(action.payload.tree, newNodes);
+      return { nodes: visibleNodes, changes: generateChanges(state.nodes, visibleNodes) };
+    }
+    case "COLLAPSE_ALL_BELOW_DEPTH": {
+      const nodesToCollapse = action.payload.tree.filter(
+        (n) => n.level >= action.payload.depth && n.hasChildren
+      );
+
+      const newNodes = Object.fromEntries(
+        Object.entries(state.nodes).map(([key, value]) => [
+          key,
+          {
+            ...value,
+            expanded: nodesToCollapse.find((n) => n.id === key) ? false : value.expanded,
+          },
+        ])
+      );
+
+      const visibleNodes = applyVisibility(action.payload.tree, newNodes);
+      return { nodes: visibleNodes, changes: generateChanges(state.nodes, visibleNodes) };
+    }
+    case "EXPAND_LEVEL": {
+      const nodesToExpand = action.payload.tree.filter(
+        (n) => n.level <= action.payload.level && n.hasChildren
+      );
+
+      const newNodes = Object.fromEntries(
+        Object.entries(state.nodes).map(([key, value]) => [
+          key,
+          {
+            ...value,
+            expanded: nodesToExpand.find((n) => n.id === key) ? true : value.expanded,
+          },
+        ])
+      );
+
+      const visibleNodes = applyVisibility(action.payload.tree, newNodes);
+      return { nodes: visibleNodes, changes: generateChanges(state.nodes, visibleNodes) };
+    }
+    case "COLLAPSE_LEVEL": {
+      const nodesToCollapse = action.payload.tree.filter(
+        (n) => n.level === action.payload.level && n.hasChildren
+      );
+
+      const newNodes = Object.fromEntries(
+        Object.entries(state.nodes).map(([key, value]) => [
+          key,
+          {
+            ...value,
+            expanded: nodesToCollapse.find((n) => n.id === key) ? false : value.expanded,
+          },
+        ])
+      );
+
+      const visibleNodes = applyVisibility(action.payload.tree, newNodes);
+      return { nodes: visibleNodes, changes: generateChanges(state.nodes, visibleNodes) };
+    }
+    case "TOGGLE_EXPAND_LEVEL": {
+      //first get the first item at that level in the tree. If it is expanded, collapse all nodes at that level
+      //if it is collapsed, expand all nodes at that level
+      const nodesAtLevel = action.payload.tree.filter(
+        (n) => n.level === action.payload.level && n.hasChildren
+      );
+      const firstNode = nodesAtLevel[0];
+      if (!firstNode) {
+        return state;
+      }
+
+      const currentlyExpanded = state.nodes[firstNode.id]?.expanded ?? true;
+      const currentVisible = state.nodes[firstNode.id]?.visible ?? true;
+      if (currentlyExpanded && currentVisible) {
+        return reducer(state, {
+          type: "COLLAPSE_LEVEL",
+          payload: {
+            level: action.payload.level,
+            tree: action.payload.tree,
+          },
+        });
+      } else {
+        return reducer(state, {
+          type: "EXPAND_LEVEL",
+          payload: {
+            level: action.payload.level,
+            tree: action.payload.tree,
+          },
+        });
+      }
+    }
     case "SELECT_FIRST_VISIBLE_NODE": {
       const node = firstVisibleNode(action.payload.tree, state.nodes);
       if (node) {

apps/webapp/app/components/runs/v3/TaskRunsTable.tsx

@@ -1,10 +1,16 @@
+import { ArrowPathIcon, StopCircleIcon } from "@heroicons/react/20/solid";
 import { StopIcon } from "@heroicons/react/24/outline";
 import { BeakerIcon, BookOpenIcon, CheckIcon } from "@heroicons/react/24/solid";
+import { useLocation } from "@remix-run/react";
+import { formatDuration } from "@trigger.dev/core/v3";
 import { User } from "@trigger.dev/database";
+import { Button, LinkButton } from "~/components/primitives/Buttons";
+import { Dialog, DialogTrigger } from "~/components/primitives/Dialog";
+import { useEnvironments } from "~/hooks/useEnvironments";
 import { useOrganization } from "~/hooks/useOrganizations";
 import { useProject } from "~/hooks/useProject";
 import { RunListAppliedFilters, RunListItem } from "~/presenters/v3/RunListPresenter.server";
-import { docsPath, v3RunPath, v3TestPath } from "~/utils/pathBuilder";
+import { docsPath, v3RunSpanPath, v3TestPath } from "~/utils/pathBuilder";
 import { EnvironmentLabel } from "../../environments/EnvironmentLabel";
 import { DateTime } from "../../primitives/DateTime";
 import { Paragraph } from "../../primitives/Paragraph";
@@ -14,21 +20,14 @@ import {
   TableBlankRow,
   TableBody,
   TableCell,
-  TableCellChevron,
   TableCellMenu,
   TableHeader,
   TableHeaderCell,
   TableRow,
 } from "../../primitives/Table";
-import { formatDuration } from "@trigger.dev/core/v3";
-import { TaskRunStatusCombo } from "./TaskRunStatus";
-import { useEnvironments } from "~/hooks/useEnvironments";
-import { Button, LinkButton } from "~/components/primitives/Buttons";
-import { ArrowPathIcon, StopCircleIcon } from "@heroicons/react/20/solid";
-import { Dialog, DialogTrigger } from "~/components/primitives/Dialog";
 import { CancelRunDialog } from "./CancelRunDialog";
-import { useLocation } from "@remix-run/react";
 import { ReplayRunDialog } from "./ReplayRunDialog";
+import { TaskRunStatusCombo } from "./TaskRunStatus";
 
 type RunsTableProps = {
   total: number;
@@ -78,7 +77,7 @@ export function TaskRunsTable({
           <BlankState isLoading={isLoading} filters={filters} />
         ) : (
           runs.map((run) => {
-            const path = v3RunPath(organization, project, run);
+            const path = v3RunSpanPath(organization, project, run, { spanId: run.spanId });
             const usernameForEnv =
               currentUser.id !== run.environment.userId ? run.environment.userName : undefined;
             return (