From c97dc131f17edd04d19fe0e5f624a0660728aeef Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Fri, 8 Aug 2025 13:31:55 +0000
Subject: [PATCH] Update MetricToolbar to support preselected metric and
 improve modal handling

Co-authored-by: nate <nate@buster.so>
---
 .../elements/MetricElement/MetricElement.tsx  |  2 +-
 .../elements/MetricElement/MetricToolbar.tsx  | 30 +++++++++++--------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/apps/web/src/components/ui/report/elements/MetricElement/MetricElement.tsx b/apps/web/src/components/ui/report/elements/MetricElement/MetricElement.tsx
index 899d9b804..c7d9d8ba4 100644
--- a/apps/web/src/components/ui/report/elements/MetricElement/MetricElement.tsx
+++ b/apps/web/src/components/ui/report/elements/MetricElement/MetricElement.tsx
@@ -32,7 +32,7 @@ export const MetricElement = withHOC(
     const readOnly = useReadOnly();
 
     const content = metricId ? (
-      <MetricToolbar>
+      <MetricToolbar selectedMetricId={metricId}>
         <MetricResizeContainer>
           <MetricContent
             metricId={metricId}
diff --git a/apps/web/src/components/ui/report/elements/MetricElement/MetricToolbar.tsx b/apps/web/src/components/ui/report/elements/MetricElement/MetricToolbar.tsx
index fa37eb10e..ef34db7e7 100644
--- a/apps/web/src/components/ui/report/elements/MetricElement/MetricToolbar.tsx
+++ b/apps/web/src/components/ui/report/elements/MetricElement/MetricToolbar.tsx
@@ -9,7 +9,7 @@ import { NodeTypeLabels } from '../../config/labels';
 import { AddMetricModal } from '@/components/features/modal/AddMetricModal';
 import { MetricPlugin, type TMetricElement } from '../../plugins/metric-plugin';
 
-export function MetricToolbar({ children }: { children: React.ReactNode }) {
+export function MetricToolbar({ children, selectedMetricId }: { children: React.ReactNode; selectedMetricId?: string }) {
   const editor = useEditorRef();
   const readOnly = useReadOnly();
   const selected = useSelected();
@@ -24,6 +24,10 @@ export function MetricToolbar({ children }: { children: React.ReactNode }) {
 
   const [openEditModal, setOpenEditModal] = React.useState(false);
 
+  const preselectedMetrics = React.useMemo(() => {
+    return selectedMetricId ? [{ id: selectedMetricId, name: '' }] : [];
+  }, [selectedMetricId]);
+
   const onOpenEdit = React.useCallback(() => {
     editor.tf.select();
     setOpenEditModal(true);
@@ -33,6 +37,17 @@ export function MetricToolbar({ children }: { children: React.ReactNode }) {
     setOpenEditModal(false);
   }, []);
 
+  const handleAddMetrics = React.useCallback(
+    async (metrics: { id: string; name: string }[]) => {
+      const id = metrics?.[0]?.id;
+      const at = editor.api.findPath(element);
+      if (!id || !at) return onCloseEdit();
+      plugin.api.metric.updateMetric(id, { at });
+      onCloseEdit();
+    },
+    [editor, element, onCloseEdit, plugin.api.metric]
+  );
+
   return (
     <PopoverBase open={isOpen} modal={false}>
       <PopoverAnchor>{children}</PopoverAnchor>
@@ -50,18 +65,9 @@ export function MetricToolbar({ children }: { children: React.ReactNode }) {
       <AddMetricModal
         open={openEditModal}
         loading={false}
-        selectedMetrics={[]}
+        selectedMetrics={preselectedMetrics}
         onClose={onCloseEdit}
-        onAddMetrics={async (metrics) => {
-          const selectedMetricId = metrics?.[0]?.id;
-          if (!selectedMetricId) return onCloseEdit();
-
-          const at = editor.api.findPath(element);
-          if (!at) return onCloseEdit();
-
-          plugin.api.metric.updateMetric(selectedMetricId, { at });
-          onCloseEdit();
-        }}
+        onAddMetrics={handleAddMetrics}
         selectionMode="single"
         saveButtonText="Update metric"
       />