From da7b70336c88b2cc583d2d88c3f7403d0c8f5395 Mon Sep 17 00:00:00 2001
From: Aleksandar Petrov <8142643+aleks-p@users.noreply.github.com>
Date: Mon, 12 Jan 2026 16:48:37 -0400
Subject: [PATCH] Improve click handling in split view

---
 .../FlameGraphCallTreeContainer.story.tsx     |  3 ++
 .../CallTree/FlameGraphCallTreeContainer.tsx  | 38 ++++++++++++++++---
 .../src/FlameGraphContainer.tsx               | 27 ++++++++++++-
 3 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.story.tsx b/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.story.tsx
index ceac052f409..75c96634cb1 100644
--- a/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.story.tsx
+++ b/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.story.tsx
@@ -40,6 +40,9 @@ export const Basic: StoryObj<typeof meta> = {
         onSandwich={(item) => {
           console.log('Sandwich:', item);
         }}
+        onSearch={(symbol) => {
+          console.log('Search:', symbol);
+        }}
       />
     );
   },
diff --git a/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.tsx b/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.tsx
index 7ca015152e2..eab15356f13 100644
--- a/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.tsx
+++ b/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.tsx
@@ -25,10 +25,11 @@ type Props = {
   colorScheme: ColorScheme | ColorSchemeDiff;
   search: string;
   compact?: boolean;
+  onSearch?: (symbol: string) => void;
 };
 
 const FlameGraphCallTreeContainer = memo(
-  ({ data, onSymbolClick, sandwichItem, onSandwich, onTableSort, colorScheme: initialColorScheme, search, compact = false }: Props) => {
+  ({ data, onSymbolClick, sandwichItem, onSandwich, onTableSort, colorScheme: initialColorScheme, search, compact = false, onSearch }: Props) => {
     const styles = useStyles2(getStyles);
     const theme = useTheme2();
 
@@ -354,13 +355,15 @@ const FlameGraphCallTreeContainer = memo(
                 row={row}
                 onFocus={handleSetFocusMode}
                 onShowCallers={handleSetCallersMode}
+                onSearch={onSearch}
                 focusedNodeId={focusedNodeId}
                 callersNodeLabel={callersNodeLabel}
                 styles={styles}
+                searchNodes={searchNodes}
               />
             ),
-            width: 50,
-            minWidth: 50,
+            width: onSearch ? 75 : 50,
+            minWidth: onSearch ? 75 : 50,
             disableSortBy: true,
           },
           {
@@ -435,13 +438,15 @@ const FlameGraphCallTreeContainer = memo(
                 row={row}
                 onFocus={handleSetFocusMode}
                 onShowCallers={handleSetCallersMode}
+                onSearch={onSearch}
                 focusedNodeId={focusedNodeId}
                 callersNodeLabel={callersNodeLabel}
                 styles={styles}
+                searchNodes={searchNodes}
               />
             ),
-            width: 50,
-            minWidth: 50,
+            width: onSearch ? 75 : 50,
+            minWidth: onSearch ? 75 : 50,
             disableSortBy: true,
           },
           {
@@ -755,16 +760,20 @@ function ActionsCell({
   row,
   onFocus,
   onShowCallers,
+  onSearch,
   focusedNodeId,
   callersNodeLabel,
   styles,
+  searchNodes,
 }: {
   row: Row<CallTreeNode> & UseExpandedRowProps<CallTreeNode>;
   onFocus: (nodeIdOrLabel: string, isLabel?: boolean) => void;
   onShowCallers: (label: string) => void;
+  onSearch?: (symbol: string) => void;
   focusedNodeId: string | undefined;
   callersNodeLabel: string | undefined;
   styles: any;
+  searchNodes?: string[];
 }) {
   const hasChildren = row.original.hasChildren;
   const isTheFocusedNode = row.original.id === focusedNodeId ||
@@ -773,6 +782,7 @@ function ActionsCell({
   const inCallersMode = callersNodeLabel !== undefined;
   const inFocusMode = focusedNodeId !== undefined;
   const isRootNode = row.original.depth === 0 && !row.original.parentId;
+  const isSearchMatch = searchNodes?.includes(row.original.id) ?? false;
 
   // Show focus button if:
   // - Node has children AND
@@ -830,6 +840,22 @@ function ActionsCell({
           <div className={styles.actionButtonPlaceholder} />
         )}
       </div>
+      {onSearch && !isSearchMatch && (
+        <div className={styles.actionButtonSlot}>
+          <Button
+            icon="search"
+            fill="text"
+            size="sm"
+            onClick={(e) => {
+              e.stopPropagation();
+              onSearch(row.original.label);
+            }}
+            tooltip="Search for this function"
+            aria-label="Search"
+            className={styles.actionButton}
+          />
+        </div>
+      )}
     </div>
   );
 }
@@ -1239,7 +1265,7 @@ function getStyles(theme: GrafanaTheme2) {
       justifyContent: 'flex-end',
       gap: theme.spacing(0.5),
       height: '20px',
-      minWidth: '60px', // Fixed width to ensure consistent alignment
+      minWidth: '60px', // Fixed width to ensure consistent alignment (72px when search button is present)
     }),
     actionButtonSlot: css({
       display: 'flex',
diff --git a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx b/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
index e7fd36c64ff..20cf606aee8 100644
--- a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
@@ -207,6 +207,30 @@ const FlameGraphContainer = ({
     [setSearch, resetFocus, onTableSymbolClick, search]
   );
 
+  // Separate callback for CallTree that doesn't trigger search
+  const onCallTreeSymbolClick = useCallback(
+    (symbol: string) => {
+      onTableSymbolClick?.(symbol);
+    },
+    [onTableSymbolClick]
+  );
+
+  // Search callback for CallTree search button
+  const onCallTreeSearch = useCallback(
+    (symbol: string) => {
+      const anchored = `^${escapeStringForRegex(symbol)}$`;
+
+      if (search === anchored) {
+        setSearch('');
+      } else {
+        onTableSymbolClick?.(symbol);
+        setSearch(anchored);
+        resetFocus();
+      }
+    },
+    [setSearch, resetFocus, onTableSymbolClick, search]
+  );
+
   if (!dataContainer) {
     return null;
   }
@@ -275,13 +299,14 @@ const FlameGraphContainer = ({
   const callTree = (
     <FlameGraphCallTreeContainer
       data={dataContainer}
-      onSymbolClick={onSymbolClick}
+      onSymbolClick={onCallTreeSymbolClick}
       sandwichItem={sandwichItem}
       onSandwich={setSandwichItem}
       onTableSort={onTableSort}
       colorScheme={colorScheme}
       search={search}
       compact={isCallTreeInSplitView}
+      onSearch={onCallTreeSearch}
     />
   );