experimental: tracing panel (#3168)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->
Fixes #2898. A v1 tracing panel for observability efforts


https://github.com/user-attachments/assets/1c042e0b-2f70-4ec0-ad63-a0c59d5a89a4

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->
- view cell runs, status, code snippet and timestamps
- hover sync between react components and vega chart
- supports dark theme
- each run now has a run_id context from the backend
- note: the timestamps here should be 'more accurate' than the frontend
timers since they come from code execution in the backend. Hence, there
can also be a small discrepancy btwn them.

## 📋 Checklist

- [X] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [X] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [X] I have added tests for the changes made.
- [X] I have run the code and verified that it works as expected.

## 📜 Reviewers

<!--
Tag potential reviewers from the community or maintainers who might be
interested in reviewing this pull request.

Your PR will be reviewed more quickly if you can figure out the right
person to tag with @ -->

@akshayka OR @mscolnick

---------

Co-authored-by: Myles Scolnick <[email protected]>

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");
+  });
+});