experimental: tracing panel

CONTRIBUTING.md

@@ -95,6 +95,7 @@ NODE_OPTIONS=--max_old_space_size=8192 NODE_ENV=development make fe -B
 | `py`           | Setup     | Editable python install; only need to run once                 |
 | `install-all`  | Setup     | Install everything; takes a long time due to editable install  |
 | `fe`           | Build     | Package frontend into `marimo/`                                |
+| `fe-codegen`   | Build     | Build [api specification](./development_docs/openapi.md)       |
 | `wheel`        | Build     | Build wheel                                                    |
 | `check`        | Test      | Run all checks                                                 |
 | `check-test`   | Test      | Run all checks and tests                                       |

development_docs/openapi.md

@@ -24,3 +24,10 @@ marimo development openapi | openapi-spec-validator -
 ```bash
 make fe-codegen
 ```
+
+You will then need to reinstall the package in `/frontend`:
+
+```bash
+cd frontend
+pnpm update @marimo-team/marimo-api
+```

frontend/src/components/buttons/clear-button.tsx

@@ -0,0 +1,23 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { cn } from "@/utils/cn";
+
+interface ClearButtonProps {
+  className?: string;
+  dataTestId?: string;
+  onClick: () => void;
+}
+
+export const ClearButton: React.FC<ClearButtonProps> = (props) => (
+  <button
+    type="button"
+    data-testid={props.dataTestId}
+    className={cn(
+      "text-xs font-semibold text-accent-foreground",
+      props.className,
+    )}
+    onClick={props.onClick}
+  >
+    Clear
+  </button>
+);

frontend/src/components/editor/cell/CellStatus.tsx

@@ -321,7 +321,7 @@ export const CellStatusComponent: React.FC<CellStatusComponentProps> = ({
   return null;
 };
 
-const ElapsedTime = (props: { elapsedTime: string }) => {
+export const ElapsedTime = (props: { elapsedTime: string }) => {
   return (
     <span className="tracking-wide font-semibold">{props.elapsedTime}</span>
   );
@@ -346,7 +346,7 @@ const LastRanTime = (props: { lastRanTime: number }) => {
   );
 };
 
-function formatElapsedTime(elapsedTime: number | null) {
+export function formatElapsedTime(elapsedTime: number | null) {
   if (elapsedTime === null) {
     return "";
   }

frontend/src/components/editor/chrome/panels/logs-panel.tsx

@@ -6,6 +6,7 @@ import React from "react";
 import { FileTextIcon } from "lucide-react";
 import { CellLink } from "../../links/cell-link";
 import { PanelEmptyState } from "./empty-state";
+import { ClearButton } from "@/components/buttons/clear-button";
 
 interface Props {
   className?: string;
@@ -35,13 +36,7 @@ export const LogsPanel: React.FC = () => {
   return (
     <>
       <div className="flex flex-row justify-end px-2 py-1">
-        <button
-          data-testid="clear-logs-button"
-          className="text-xs font-semibold text-accent-foreground"
-          onClick={clearLogs}
-        >
-          Clear
-        </button>
+        <ClearButton dataTestId="clear-logs-button" onClick={clearLogs} />
       </div>
       <div className="overflow-auto flex-1">
         <LogViewer logs={logs} className="min-w-[300px]" />
@@ -84,7 +79,9 @@ function formatLog(log: CellLog) {
 
   return (
     <>
-      <span className="flex-shrink-0 text-[var(--gray-10)]">[{timestamp}]</span>
+      <span className="flex-shrink-0 text-[var(--gray-10)] dark:text-[var(--gray-11)]">
+        [{timestamp}]
+      </span>
       <span className={cn("flex-shrink-0", color)}>{level}</span>
       <span className="flex-shrink-0 text-[var(--gray-10)]">
         (<CellLink cellId={log.cellId} />)

frontend/src/components/editor/chrome/panels/tracing-panel.tsx

@@ -0,0 +1,7 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { Tracing } from "@/components/tracing/tracing";
+import React from "react";
+
+export const TracingPanel: React.FC = () => {
+  return <Tracing />;
+};

frontend/src/components/editor/chrome/types.ts

@@ -14,6 +14,7 @@ import {
   NotebookPenIcon,
   BoxIcon,
   BotMessageSquareIcon,
+  ActivityIcon,
 } from "lucide-react";
 
 export type PanelType =
@@ -22,6 +23,7 @@ export type PanelType =
   | "variables"
   | "outline"
   | "dependencies"
+  | "tracing"
   | "packages"
   | "documentation"
   | "snippets"
@@ -94,6 +96,13 @@ export const PANELS: PanelDescriptor[] = [
     tooltip: "Notebook logs",
     position: "sidebar",
   },
+  {
+    type: "tracing",
+    Icon: ActivityIcon,
+    tooltip: "Tracing",
+    position: "sidebar",
+    hidden: !getFeatureFlag("tracing"),
+  },
   {
     type: "snippets",
     Icon: SquareDashedBottomCodeIcon,

frontend/src/components/editor/chrome/wrapper/app-chrome.tsx

@@ -30,6 +30,7 @@ import { IfCapability } from "@/core/config/if-capability";
 import { PackagesPanel } from "../panels/packages-panel";
 import { ChatPanel } from "@/components/chat/chat-panel";
 import { TooltipProvider } from "@radix-ui/react-tooltip";
+import { TracingPanel } from "../panels/tracing-panel";
 
 const LazyTerminal = React.lazy(() => import("@/components/terminal/terminal"));
 
@@ -153,6 +154,7 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
             {selectedPanel === "scratchpad" && <ScratchpadPanel />}
             {selectedPanel === "chat" && <ChatPanel />}
             {selectedPanel === "logs" && <LogsPanel />}
+            {selectedPanel === "tracing" && <TracingPanel />}
           </TooltipProvider>
         </div>
       </Suspense>

frontend/src/components/tracing/tracing-spec.ts

@@ -0,0 +1,106 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import type { CellId } from "@/core/cells/ids";
+import type { CellRun } from "@/core/cells/runs";
+import type { ResolvedTheme } from "@/theme/useTheme";
+import type { TopLevelSpec } from "vega-lite";
+
+export const REACT_HOVERED_CELLID = "hoveredCellId";
+export const VEGA_HOVER_SIGNAL = "cellHover";
+
+export type ChartPosition = "sideBySide" | "above";
+export interface ChartValues {
+  cell: CellId;
+  cellNum: number;
+  startTimestamp: string;
+  endTimestamp: string;
+  elapsedTime: string;
+  status: CellRun["status"];
+}
+
+const cellNumField = "cellNum" satisfies keyof ChartValues;
+const cellField = "cell" satisfies keyof ChartValues;
+const startTimestampField = "startTimestamp" satisfies keyof ChartValues;
+const endTimestampField = "endTimestamp" satisfies keyof ChartValues;
+const statusField = "status" satisfies keyof ChartValues;
+
+export function createGanttBaseSpec(
+  chartValues: ChartValues[],
+  hiddenInputElementId: string,
+  chartPosition: ChartPosition,
+  theme: ResolvedTheme,
+): TopLevelSpec {
+  return {
+    $schema: "https://vega.github.io/schema/vega-lite/v5.json",
+    background: theme === "dark" ? "black" : undefined,
+    mark: {
+      type: "bar",
+      cornerRadius: 2,
+    },
+    params: [
+      {
+        name: REACT_HOVERED_CELLID,
+        bind: { element: `#${hiddenInputElementId}` },
+      },
+      {
+        name: VEGA_HOVER_SIGNAL,
+        select: {
+          type: "point",
+          on: "mouseover",
+          fields: [cellField],
+          clear: "mouseout",
+        },
+      },
+    ],
+    height: { step: 23 },
+    encoding: {
+      y: {
+        field: cellNumField,
+        scale: { paddingInner: 0.2 },
+        sort: { field: cellNumField },
+        title: "cell",
+        axis: chartPosition === "sideBySide" ? null : undefined,
+      },
+      x: {
+        field: startTimestampField,
+        type: "temporal",
+        axis: { orient: "top", title: null },
+      },
+      x2: {
+        field: endTimestampField,
+        type: "temporal",
+      },
+      tooltip: [
+        {
+          field: startTimestampField,
+          type: "temporal",
+          timeUnit: "dayhoursminutesseconds",
+          title: "Start",
+        },
+        {
+          field: endTimestampField,
+          type: "temporal",
+          timeUnit: "dayhoursminutesseconds",
+          title: "End",
+        },
+      ],
+      size: {
+        value: {
+          expr: `${REACT_HOVERED_CELLID} == toString(datum.cell) ? 19.5 : 18`,
+        },
+      },
+      color: {
+        field: statusField,
+        scale: { domain: ["success", "error"], range: ["#37BE5F", "red"] },
+        legend: null,
+      },
+    },
+    data: {
+      values: chartValues,
+    },
+    config: {
+      view: {
+        stroke: "transparent",
+      },
+    },
+  };
+}

frontend/src/components/tracing/tracing.test.tsx

@@ -0,0 +1,66 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { describe, it, expect, beforeAll, afterAll, vi } from "vitest";
+import { formatChartTime } from "./tracing";
+
+describe("formatChartTime", () => {
+  beforeAll(() => {
+    // Mock Date to always use UTC
+    vi.spyOn(global.Date.prototype, "getFullYear").mockImplementation(function (
+      this: Date,
+    ) {
+      return this.getUTCFullYear();
+    });
+    vi.spyOn(global.Date.prototype, "getMonth").mockImplementation(function (
+      this: Date,
+    ) {
+      return this.getUTCMonth();
+    });
+    vi.spyOn(global.Date.prototype, "getDate").mockImplementation(function (
+      this: Date,
+    ) {
+      return this.getUTCDate();
+    });
+    vi.spyOn(global.Date.prototype, "getHours").mockImplementation(function (
+      this: Date,
+    ) {
+      return this.getUTCHours();
+    });
+    vi.spyOn(global.Date.prototype, "getMinutes").mockImplementation(function (
+      this: Date,
+    ) {
+      return this.getUTCMinutes();
+    });
+    vi.spyOn(global.Date.prototype, "getSeconds").mockImplementation(function (
+      this: Date,
+    ) {
+      return this.getUTCSeconds();
+    });
+    vi.spyOn(global.Date.prototype, "getMilliseconds").mockImplementation(
+      function (this: Date) {
+        return this.getUTCMilliseconds();
+      },
+    );
+  });
+
+  afterAll(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("should handle a timestamp with milliseconds correctly", () => {
+    const timestamp = 1_704_067_200.123;
+    const formattedTime = formatChartTime(timestamp);
+    expect(formattedTime).toBe("2024-01-01 00:00:00.123");
+  });
+
+  it("should handle a timestamp at the start of the year correctly", () => {
+    const timestamp = 1_704_067_200;
+    const formattedTime = formatChartTime(timestamp);
+    expect(formattedTime).toBe("2024-01-01 00:00:00.000");
+  });
+
+  it("should handle a timestamp at the end of the year correctly", () => {
+    const timestamp = 1_704_067_199;
+    const formattedTime = formatChartTime(timestamp);
+    expect(formattedTime).toBe("2023-12-31 23:59:59.000");
+  });
+});