From 01ec0a77a8b0028679bed63720ad46e963aa9c1d Mon Sep 17 00:00:00 2001
From: abhigyanpatwari <abhigyan1.patwari@gmail.com>
Date: Tue, 6 Jan 2026 16:22:43 +0530
Subject: [PATCH] highlight tool fixes

---
 src/components/RightPanel.tsx   |  6 +++---
 src/components/ToolCallCard.tsx | 37 ++++++++++++++++++++++++++-------
 src/core/llm/tools.ts           | 12 +++++++++--
 src/hooks/useAppState.tsx       | 32 +++++++++++++++++++++++++---
 4 files changed, 71 insertions(+), 16 deletions(-)

diff --git a/src/components/RightPanel.tsx b/src/components/RightPanel.tsx
index 573b9ec201..8e4c07d3fb 100644
--- a/src/components/RightPanel.tsx
+++ b/src/components/RightPanel.tsx
@@ -383,9 +383,9 @@ export const RightPanel = () => {
                               
                               if (isInline) {
                                 return (
-                                  <code className="px-1 py-0.5 bg-surface rounded text-accent font-mono text-xs" {...props}>
-                                    {children}
-                                  </code>
+                                <code className="px-1 py-0.5 bg-surface rounded text-accent font-mono text-xs" {...props}>
+                                  {children}
+                                </code>
                                 );
                               }
                               
diff --git a/src/components/ToolCallCard.tsx b/src/components/ToolCallCard.tsx
index b5b7dedc2d..689d7492b1 100644
--- a/src/components/ToolCallCard.tsx
+++ b/src/components/ToolCallCard.tsx
@@ -5,7 +5,7 @@
  * Shows the tool name, status, and when expanded, the query/args and result.
  */
 
-import { useState, useCallback } from 'react';
+import { useState, useCallback, useMemo } from 'react';
 import { ChevronDown, ChevronRight, Sparkles, Check, Loader2, AlertCircle, Eye, EyeOff } from 'lucide-react';
 import type { ToolCallInfo } from '../core/llm/types';
 import { useAppState } from '../hooks/useAppState';
@@ -111,17 +111,38 @@ const extractHighlightNodeIds = (result: string | undefined): string[] => {
 
 export const ToolCallCard = ({ toolCall, defaultExpanded = false }: ToolCallCardProps) => {
   const [isExpanded, setIsExpanded] = useState(defaultExpanded);
-  const { highlightedNodeIds, setHighlightedNodeIds } = useAppState();
+  const { highlightedNodeIds, setHighlightedNodeIds, graph } = useAppState();
   const status = getStatusDisplay(toolCall.status);
   const formattedArgs = formatArgs(toolCall.args);
   
   // Check if this is a highlight tool and extract node IDs
   const isHighlightTool = toolCall.name === 'highlight_in_graph';
-  const highlightNodeIds = isHighlightTool ? extractHighlightNodeIds(toolCall.result) : [];
+  const rawHighlightNodeIds = isHighlightTool ? extractHighlightNodeIds(toolCall.result) : [];
+  
+  // Resolve raw IDs to actual graph node IDs (handles partial ID matching)
+  const resolvedNodeIds = useMemo(() => {
+    if (rawHighlightNodeIds.length === 0 || !graph) return rawHighlightNodeIds;
+    
+    const graphNodeIds = graph.nodes.map(n => n.id);
+    const resolved: string[] = [];
+    
+    for (const rawId of rawHighlightNodeIds) {
+      if (graphNodeIds.includes(rawId)) {
+        resolved.push(rawId);
+      } else {
+        // Try partial match - find node whose ID ends with the raw ID
+        const found = graphNodeIds.find(gid => 
+          gid.endsWith(rawId) || gid.endsWith(':' + rawId)
+        );
+        if (found) resolved.push(found);
+      }
+    }
+    return resolved;
+  }, [rawHighlightNodeIds, graph]);
   
   // Check if these specific nodes are currently highlighted
-  const isHighlightActive = highlightNodeIds.length > 0 && 
-    highlightNodeIds.some(id => highlightedNodeIds.has(id));
+  const isHighlightActive = resolvedNodeIds.length > 0 && 
+    resolvedNodeIds.some(id => highlightedNodeIds.has(id));
   
   // Toggle highlight on/off
   const toggleHighlight = useCallback((e: React.MouseEvent) => {
@@ -131,9 +152,9 @@ export const ToolCallCard = ({ toolCall, defaultExpanded = false }: ToolCallCard
       setHighlightedNodeIds(new Set());
     } else {
       // Turn on - set these nodes as highlighted
-      setHighlightedNodeIds(new Set(highlightNodeIds));
+      setHighlightedNodeIds(new Set(resolvedNodeIds));
     }
-  }, [isHighlightActive, highlightNodeIds, setHighlightedNodeIds]);
+  }, [isHighlightActive, resolvedNodeIds, setHighlightedNodeIds]);
   
   return (
     <div className={`rounded-lg border ${status.borderColor} ${status.bgColor} overflow-hidden transition-all`}>
@@ -153,7 +174,7 @@ export const ToolCallCard = ({ toolCall, defaultExpanded = false }: ToolCallCard
         </span>
         
         {/* Highlight toggle button - only for highlight_in_graph tool with results */}
-        {isHighlightTool && highlightNodeIds.length > 0 && (
+        {isHighlightTool && resolvedNodeIds.length > 0 && (
           <button
             onClick={toggleHighlight}
             className={`flex items-center gap-1 px-2 py-0.5 rounded text-xs transition-colors ${
diff --git a/src/core/llm/tools.ts b/src/core/llm/tools.ts
index 19ba0864a6..9375ab8e5c 100644
--- a/src/core/llm/tools.ts
+++ b/src/core/llm/tools.ts
@@ -600,9 +600,17 @@ export const createGraphRAGTools = (
     },
     {
       name: 'highlight_in_graph',
-      description: 'Highlight specific nodes in the visual knowledge graph that the user can see alongside this chat. Use this to visually show the user which code elements you are discussing. Pass the node IDs from previous search/query results.',
+      description: `Highlight specific nodes in the visual knowledge graph. Pass the EXACT node IDs from your query results.
+
+IMPORTANT: Node IDs include a label prefix! Format is: Label:filepath:name
+Examples:
+- Class:src/agents/base.py:BaseAgent
+- Function:src/utils.ts:calculateSum
+- File:src/main.py
+
+Copy the ID EXACTLY as it appears in query results (the "classId", "fnId", "fileId", etc. columns).`,
       schema: z.object({
-        nodeIds: z.array(z.string()).describe('Array of node IDs to highlight (from search results or queries)'),
+        nodeIds: z.array(z.string()).describe('Array of EXACT node IDs to highlight - must include the label prefix like "Class:" or "Function:"'),
         description: z.string().optional().nullable().describe('Brief description of what these nodes represent'),
       }),
     }
diff --git a/src/hooks/useAppState.tsx b/src/hooks/useAppState.tsx
index 49c2e3a3fd..55d9851437 100644
--- a/src/hooks/useAppState.tsx
+++ b/src/hooks/useAppState.tsx
@@ -527,9 +527,35 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
               if (tc.result) {
                 const highlightMatch = tc.result.match(/\[HIGHLIGHT_NODES:([^\]]+)\]/);
                 if (highlightMatch) {
-                  const nodeIds = highlightMatch[1].split(',').map((id: string) => id.trim()).filter(Boolean);
-                  if (nodeIds.length > 0) {
-                    setHighlightedNodeIds(new Set(nodeIds));
+                  const rawIds = highlightMatch[1].split(',').map((id: string) => id.trim()).filter(Boolean);
+                  if (rawIds.length > 0 && graph) {
+                    // Try to match IDs against actual graph nodes
+                    // This handles cases where the LLM passes partial IDs without the label prefix
+                    const matchedIds = new Set<string>();
+                    const graphNodeIds = graph.nodes.map(n => n.id);
+                    
+                    for (const rawId of rawIds) {
+                      // First try exact match
+                      if (graphNodeIds.includes(rawId)) {
+                        matchedIds.add(rawId);
+                      } else {
+                        // Try to find a node whose ID ends with the raw ID
+                        // e.g., "src/path:ClassName" should match "Class:src/path:ClassName"
+                        const found = graphNodeIds.find(gid => 
+                          gid.endsWith(rawId) || gid.endsWith(':' + rawId)
+                        );
+                        if (found) {
+                          matchedIds.add(found);
+                        }
+                      }
+                    }
+                    
+                    if (matchedIds.size > 0) {
+                      setHighlightedNodeIds(matchedIds);
+                    }
+                  } else if (rawIds.length > 0) {
+                    // Fallback: just use the IDs directly if no graph available
+                    setHighlightedNodeIds(new Set(rawIds));
                   }
                 }
               }