Merge pull request #10 from abhigyanpatwari/embeddings_pipeline

Embeddings pipeline

.gitignore

@@ -38,3 +38,4 @@ coverage/
 
 
 
+

package.json

@@ -32,6 +32,7 @@
     "langchain": "^0.3.37",
     "lru-cache": "^11.2.4",
     "lucide-react": "^0.562.0",
+    "mermaid": "^11.12.2",
     "minisearch": "^7.2.0",
     "react": "^18.3.1",
     "react-dom": "^18.3.1",

src/components/CodeReferencesPanel.tsx

@@ -1,5 +1,5 @@
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import { Code, PanelLeftClose, PanelLeft, Trash2, X, Target } from 'lucide-react';
+import { Code, PanelLeftClose, PanelLeft, Trash2, X, Target, FileCode, Sparkles, MousePointerClick } from 'lucide-react';
 import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
 import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism';
 import { useAppState } from '../hooks/useAppState';
@@ -131,15 +131,22 @@ export const CodeReferencesPanel = ({ onFocusNode }: CodeReferencesPanelProps) =
       <aside className="h-full w-12 bg-surface border-r border-border-subtle flex flex-col items-center py-3 gap-2 flex-shrink-0">
         <button
           onClick={() => setIsCollapsed(false)}
-          className="p-2 text-text-secondary hover:text-text-primary hover:bg-hover rounded transition-colors"
+          className="p-2 text-text-secondary hover:text-cyan-400 hover:bg-cyan-500/10 rounded transition-colors"
           title="Expand Code Panel"
         >
           <PanelLeft className="w-5 h-5" />
         </button>
         <div className="w-6 h-px bg-border-subtle my-1" />
-        <div className="text-[10px] text-text-muted rotate-90 whitespace-nowrap font-mono">
-          {showCitations ? `${aiReferences.length} refs` : 'Selected'}
-        </div>
+        {showSelectedViewer && (
+          <div className="text-[9px] text-amber-400 rotate-90 whitespace-nowrap font-medium tracking-wide">
+            SELECTED
+          </div>
+        )}
+        {showCitations && (
+          <div className="text-[9px] text-cyan-400 rotate-90 whitespace-nowrap font-medium tracking-wide mt-4">
+            AI • {aiReferences.length}
+          </div>
+        )}
       </aside>
     );
   }
@@ -157,18 +164,17 @@ export const CodeReferencesPanel = ({ onFocusNode }: CodeReferencesPanelProps) =
         title="Drag to resize"
       />
       {/* Header */}
-      <div className="flex items-center justify-between px-3 py-2 border-b border-border-subtle bg-elevated/40">
+      <div className="flex items-center justify-between px-3 py-2.5 border-b border-border-subtle bg-gradient-to-r from-elevated/60 to-surface/60">
         <div className="flex items-center gap-2">
-          <Code className="w-4 h-4 text-cyan-300" />
-          <span className="text-sm font-medium">Code</span>
-          {showCitations && <span className="text-xs text-text-muted">• {aiReferences.length} refs</span>}
+          <Code className="w-4 h-4 text-cyan-400" />
+          <span className="text-sm font-semibold text-text-primary">Code Inspector</span>
         </div>
         <div className="flex items-center gap-1.5">
           {showCitations && (
             <button
               onClick={() => clearCodeReferences()}
-              className="p-1.5 text-text-muted hover:text-text-primary hover:bg-hover rounded transition-colors"
-              title="Clear all code references"
+              className="p-1.5 text-text-muted hover:text-red-400 hover:bg-red-500/10 rounded transition-colors"
+              title="Clear AI citations"
             >
               <Trash2 className="w-4 h-4" />
             </button>
@@ -187,14 +193,18 @@ export const CodeReferencesPanel = ({ onFocusNode }: CodeReferencesPanelProps) =
         {/* Top: Selected file viewer (when a node is selected) */}
         {showSelectedViewer && (
           <div className={`${showCitations ? 'h-[42%]' : 'flex-1'} min-h-0 flex flex-col`}>
-            <div className="px-3 py-2 bg-surface/40 border-b border-border-subtle flex items-center gap-2">
-              <span className="text-[11px] text-text-muted uppercase tracking-wide">Selected</span>
+            <div className="px-3 py-2 bg-gradient-to-r from-amber-500/8 to-orange-500/5 border-b border-amber-500/20 flex items-center gap-2">
+              <div className="flex items-center gap-1.5 px-2 py-0.5 bg-amber-500/15 rounded-md border border-amber-500/25">
+                <MousePointerClick className="w-3 h-3 text-amber-400" />
+                <span className="text-[10px] text-amber-300 font-semibold uppercase tracking-wide">Selected</span>
+              </div>
+              <FileCode className="w-3.5 h-3.5 text-amber-400/70 ml-1" />
               <span className="text-xs text-text-primary font-mono truncate flex-1">
-                {selectedNode?.properties?.filePath ?? selectedNode?.properties?.name}
+                {selectedNode?.properties?.filePath?.split('/').pop() ?? selectedNode?.properties?.name}
               </span>
               <button
                 onClick={() => setSelectedNode(null)}
-                className="p-1 text-text-muted hover:text-text-primary hover:bg-hover rounded transition-colors"
+                className="p-1 text-text-muted hover:text-amber-400 hover:bg-amber-500/10 rounded transition-colors"
                 title="Clear selection"
               >
                 <X className="w-4 h-4" />
@@ -255,12 +265,21 @@ export const CodeReferencesPanel = ({ onFocusNode }: CodeReferencesPanelProps) =
 
         {/* Divider between Selected viewer and AI refs (more visible) */}
         {showSelectedViewer && showCitations && (
-          <div className="h-2 bg-gradient-to-r from-cyan-500/0 via-cyan-400/35 to-cyan-500/0 border-y border-cyan-400/25" />
+          <div className="h-1.5 bg-gradient-to-r from-transparent via-border-subtle to-transparent" />
         )}
 
         {/* Bottom: AI citations list */}
         {showCitations && (
-          <div className="flex-1 min-h-0 overflow-y-auto scrollbar-thin p-3 space-y-3">
+          <div className="flex-1 min-h-0 flex flex-col">
+            {/* AI Citations Section Header */}
+            <div className="px-3 py-2 bg-gradient-to-r from-cyan-500/8 to-teal-500/5 border-b border-cyan-500/20 flex items-center gap-2">
+              <div className="flex items-center gap-1.5 px-2 py-0.5 bg-cyan-500/15 rounded-md border border-cyan-500/25">
+                <Sparkles className="w-3 h-3 text-cyan-400" />
+                <span className="text-[10px] text-cyan-300 font-semibold uppercase tracking-wide">AI Citations</span>
+              </div>
+              <span className="text-xs text-text-muted ml-1">{aiReferences.length} reference{aiReferences.length !== 1 ? 's' : ''}</span>
+            </div>
+            <div className="flex-1 min-h-0 overflow-y-auto scrollbar-thin p-3 space-y-3">
             {refsWithSnippets.map(({ ref, content, start, highlightStart, highlightEnd, totalLines }) => {
           const nodeColor = ref.label ? (NODE_COLORS as any)[ref.label] || '#6b7280' : '#6b7280';
           const hasRange = typeof ref.startLine === 'number';
@@ -368,11 +387,10 @@ export const CodeReferencesPanel = ({ onFocusNode }: CodeReferencesPanelProps) =
             </div>
           );
             })}
+            </div>
           </div>
         )}
       </div>
     </aside>
   );
 };
-
-

src/components/FileTreePanel.tsx

@@ -123,8 +123,8 @@ const TreeItem = ({
         onClick={handleClick}
         className={`
           w-full flex items-center gap-1.5 px-2 py-1 text-left text-sm
-          hover:bg-hover transition-colors rounded
-          ${isSelected ? 'bg-accent/20 text-accent' : 'text-text-secondary hover:text-text-primary'}
+          hover:bg-hover transition-colors rounded relative
+          ${isSelected ? 'bg-amber-500/15 text-amber-300 border-l-2 border-amber-400' : 'text-text-secondary hover:text-text-primary border-l-2 border-transparent'}
           ${matchesSearch ? 'bg-accent/10' : ''}
         `}
         style={{ paddingLeft: `${depth * 12 + 8}px` }}
@@ -216,6 +216,31 @@ export const FileTreePanel = ({ onFocusNode }: FileTreePanelProps) => {
     }
   }, [fileTree.length]); // Only run when tree first loads
 
+  // Auto-expand to selected file when selectedNode changes (e.g., from graph click)
+  useEffect(() => {
+    const path = selectedNode?.properties?.filePath;
+    if (!path) return;
+    
+    // Expand all parent folders leading to this file
+    const parts = path.split('/').filter(Boolean);
+    const pathsToExpand: string[] = [];
+    let currentPath = '';
+    
+    // Build all parent paths (exclude the last part if it's a file)
+    for (let i = 0; i < parts.length - 1; i++) {
+      currentPath = currentPath ? `${currentPath}/${parts[i]}` : parts[i];
+      pathsToExpand.push(currentPath);
+    }
+    
+    if (pathsToExpand.length > 0) {
+      setExpandedPaths(prev => {
+        const next = new Set(prev);
+        pathsToExpand.forEach(p => next.add(p));
+        return next;
+      });
+    }
+  }, [selectedNode?.id]); // Trigger when selected node changes
+
   const toggleExpanded = useCallback((path: string) => {
     setExpandedPaths(prev => {
       const next = new Set(prev);

src/components/GraphCanvas.tsx

@@ -1,5 +1,5 @@
-import { useEffect, useCallback, useState, forwardRef, useImperativeHandle } from 'react';
-import { ZoomIn, ZoomOut, Maximize2, Focus, RotateCcw, Play, Pause } from 'lucide-react';
+import { useEffect, useCallback, useMemo, useState, forwardRef, useImperativeHandle } from 'react';
+import { ZoomIn, ZoomOut, Maximize2, Focus, RotateCcw, Play, Pause, Lightbulb, LightbulbOff } from 'lucide-react';
 import { useSigma } from '../hooks/useSigma';
 import { useAppState } from '../hooks/useAppState';
 import { knowledgeGraphToGraphology, filterGraphByDepth, SigmaNodeAttributes, SigmaEdgeAttributes } from '../lib/graph-adapter';
@@ -11,8 +11,28 @@ export interface GraphCanvasHandle {
 }
 
 export const GraphCanvas = forwardRef<GraphCanvasHandle>((_, ref) => {
-  const { graph, setSelectedNode, selectedNode: appSelectedNode, visibleLabels, openCodePanel, depthFilter, highlightedNodeIds } = useAppState();
+  const {
+    graph,
+    setSelectedNode,
+    selectedNode: appSelectedNode,
+    visibleLabels,
+    openCodePanel,
+    depthFilter,
+    highlightedNodeIds,
+    aiCitationHighlightedNodeIds,
+    aiToolHighlightedNodeIds,
+    isAIHighlightsEnabled,
+    toggleAIHighlights,
+  } = useAppState();
   const [hoveredNodeName, setHoveredNodeName] = useState<string | null>(null);
+
+  const effectiveHighlightedNodeIds = useMemo(() => {
+    if (!isAIHighlightsEnabled) return highlightedNodeIds;
+    const next = new Set(highlightedNodeIds);
+    for (const id of aiCitationHighlightedNodeIds) next.add(id);
+    for (const id of aiToolHighlightedNodeIds) next.add(id);
+    return next;
+  }, [highlightedNodeIds, aiCitationHighlightedNodeIds, aiToolHighlightedNodeIds, isAIHighlightsEnabled]);
   
   const handleNodeClick = useCallback((nodeId: string) => {
     if (!graph) return;
@@ -55,7 +75,7 @@ export const GraphCanvas = forwardRef<GraphCanvasHandle>((_, ref) => {
     onNodeClick: handleNodeClick,
     onNodeHover: handleNodeHover,
     onStageClick: handleStageClick,
-    highlightedNodeIds,
+    highlightedNodeIds: effectiveHighlightedNodeIds,
   });
 
   // Expose focusNode to parent via ref
@@ -244,6 +264,21 @@ export const GraphCanvas = forwardRef<GraphCanvasHandle>((_, ref) => {
 
       {/* Query FAB */}
       <QueryFAB />
+
+      {/* AI Highlights toggle - Top Right */}
+      <div className="absolute top-4 right-4 z-20">
+        <button
+          onClick={toggleAIHighlights}
+          className={
+            isAIHighlightsEnabled
+              ? 'w-10 h-10 flex items-center justify-center bg-cyan-500/15 border border-cyan-400/40 rounded-lg text-cyan-200 hover:bg-cyan-500/20 hover:border-cyan-300/60 transition-colors'
+              : 'w-10 h-10 flex items-center justify-center bg-elevated border border-border-subtle rounded-lg text-text-muted hover:bg-hover hover:text-text-primary transition-colors'
+          }
+          title={isAIHighlightsEnabled ? 'Turn off AI highlights' : 'Turn on AI highlights'}
+        >
+          {isAIHighlightsEnabled ? <Lightbulb className="w-4 h-4" /> : <LightbulbOff className="w-4 h-4" />}
+        </button>
+      </div>
     </div>
   );
 });

src/components/MermaidDiagram.tsx

@@ -0,0 +1,139 @@
+import { useEffect, useRef, useState } from 'react';
+import mermaid from 'mermaid';
+import { AlertTriangle, Maximize2, Minimize2 } from 'lucide-react';
+
+// Initialize mermaid with dark theme
+mermaid.initialize({
+  startOnLoad: false,
+  theme: 'dark',
+  themeVariables: {
+    primaryColor: '#06b6d4',
+    primaryTextColor: '#e4e4ed',
+    primaryBorderColor: '#1e1e2a',
+    lineColor: '#3b3b54',
+    secondaryColor: '#1e1e2a',
+    tertiaryColor: '#0a0a10',
+    background: '#0a0a10',
+    mainBkg: '#0f0f18',
+    nodeBorder: '#3b3b54',
+    clusterBkg: '#1e1e2a',
+    titleColor: '#e4e4ed',
+    edgeLabelBackground: '#0f0f18',
+    nodeTextColor: '#e4e4ed',
+  },
+  flowchart: {
+    curve: 'basis',
+    padding: 15,
+    nodeSpacing: 50,
+    rankSpacing: 50,
+  },
+  sequence: {
+    actorMargin: 50,
+    boxMargin: 10,
+    boxTextMargin: 5,
+    noteMargin: 10,
+    messageMargin: 35,
+  },
+  fontFamily: '"JetBrains Mono", "Fira Code", monospace',
+  fontSize: 13,
+});
+
+interface MermaidDiagramProps {
+  code: string;
+}
+
+export const MermaidDiagram = ({ code }: MermaidDiagramProps) => {
+  const containerRef = useRef<HTMLDivElement>(null);
+  const [error, setError] = useState<string | null>(null);
+  const [isExpanded, setIsExpanded] = useState(false);
+  const [svg, setSvg] = useState<string>('');
+
+  useEffect(() => {
+    const renderDiagram = async () => {
+      if (!containerRef.current) return;
+
+      try {
+        // Generate unique ID for this diagram
+        const id = `mermaid-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+        
+        // Render the diagram
+        const { svg: renderedSvg } = await mermaid.render(id, code.trim());
+        setSvg(renderedSvg);
+        setError(null);
+      } catch (err) {
+        console.error('Mermaid render error:', err);
+        setError(err instanceof Error ? err.message : 'Failed to render diagram');
+        setSvg('');
+      }
+    };
+
+    renderDiagram();
+  }, [code]);
+
+  if (error) {
+    return (
+      <div className="my-3 p-4 bg-rose-500/10 border border-rose-500/30 rounded-lg">
+        <div className="flex items-center gap-2 text-rose-300 text-sm mb-2">
+          <AlertTriangle className="w-4 h-4" />
+          <span className="font-medium">Diagram Error</span>
+        </div>
+        <pre className="text-xs text-rose-200/70 font-mono whitespace-pre-wrap">{error}</pre>
+        <details className="mt-2">
+          <summary className="text-xs text-text-muted cursor-pointer hover:text-text-secondary">
+            Show source
+          </summary>
+          <pre className="mt-2 p-2 bg-surface rounded text-xs text-text-muted overflow-x-auto">
+            {code}
+          </pre>
+        </details>
+      </div>
+    );
+  }
+
+  return (
+    <div className={`my-3 relative group ${isExpanded ? 'fixed inset-4 z-50' : ''}`}>
+      {/* Backdrop for expanded view */}
+      {isExpanded && (
+        <div 
+          className="absolute inset-0 -m-4 bg-deep/95 backdrop-blur-sm"
+          onClick={() => setIsExpanded(false)}
+        />
+      )}
+      
+      <div className={`
+        relative bg-gradient-to-b from-surface to-elevated 
+        border border-border-subtle rounded-xl overflow-hidden
+        ${isExpanded ? 'h-full' : ''}
+      `}>
+        {/* Header */}
+        <div className="flex items-center justify-between px-3 py-2 bg-surface/60 border-b border-border-subtle">
+          <span className="text-[10px] text-text-muted uppercase tracking-wider font-medium">
+            Diagram
+          </span>
+          <button
+            onClick={() => setIsExpanded(!isExpanded)}
+            className="p-1 text-text-muted hover:text-text-primary hover:bg-hover rounded transition-colors"
+            title={isExpanded ? 'Minimize' : 'Expand'}
+          >
+            {isExpanded ? (
+              <Minimize2 className="w-3.5 h-3.5" />
+            ) : (
+              <Maximize2 className="w-3.5 h-3.5" />
+            )}
+          </button>
+        </div>
+        
+        {/* Diagram container */}
+        <div 
+          ref={containerRef}
+          className={`
+            flex items-center justify-center p-4 overflow-auto
+            ${isExpanded ? 'h-[calc(100%-40px)]' : 'max-h-[400px]'}
+          `}
+          dangerouslySetInnerHTML={{ __html: svg }}
+        />
+      </div>
+    </div>
+  );
+};
+

src/core/llm/agent.ts

@@ -33,7 +33,8 @@ import type {
 const SYSTEM_PROMPT = `You are Nexus, a code analysis agent. You explore codebases through a graph database and source files.
 
 **EXTREMELY IMPORTANT NOTE** : Even if there is a small chance of giving better context/understanding to the user using the highlight tool be extremely sure to use it. You can figure out yourself how do u use it in the specific context. 
-Always output in proper markdown formatting
+- Always output in proper markdown formatting
+- Always use grounding reference
 
 ## THINK ALOUD
 
@@ -68,6 +69,27 @@ You are diligent and tireless.
 \`execute_cypher\` - graph structure queries
 \`highlight_in_graph\` - highlight nodes for the user (they see a visual graph)
 
+## MERMAID DIAGRAMS
+
+Use mermaid diagrams when explaining:
+- **Architecture** - show component relationships with flowcharts or C4 diagrams
+- **Data flows** - illustrate how data moves through the system
+- **Call sequences** - show function call chains with sequence diagrams
+- **Class hierarchies** - display inheritance/composition with class diagrams
+- **State machines** - visualize state transitions
+
+Format: wrap in \`\`\`mermaid code blocks. Keep diagrams focused - 5-10 nodes max for clarity.
+
+Example:
+\`\`\`mermaid
+flowchart LR
+    A[API Handler] --> B[Service Layer]
+    B --> C[Database]
+    B --> D[Cache]
+\`\`\`
+
+Prefer diagrams over long textual explanations for structural concepts.
+
 ## GROUNDING REFERENCES
 
 When you cite code, include inline file references so the UI can surface the code automatically:

src/core/search/hybrid-search.ts

@@ -146,3 +146,4 @@ export const formatHybridResults = (results: HybridSearchResult[]): string => {
 
 
 
+

src/core/search/index.ts

@@ -22,3 +22,4 @@ export {
 
 
 
+

src/hooks/useAppState.tsx

@@ -67,6 +67,12 @@ interface AppState {
   // Query state
   highlightedNodeIds: Set<string>;
   setHighlightedNodeIds: (ids: Set<string>) => void;
+  // AI highlights (toggable)
+  aiCitationHighlightedNodeIds: Set<string>;
+  aiToolHighlightedNodeIds: Set<string>;
+  isAIHighlightsEnabled: boolean;
+  toggleAIHighlights: () => void;
+  clearAIToolHighlights: () => void;
   queryResult: QueryResult | null;
   setQueryResult: (result: QueryResult | null) => void;
   clearQueryHighlights: () => void;
@@ -165,6 +171,19 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
   // Query state
   const [highlightedNodeIds, setHighlightedNodeIds] = useState<Set<string>>(new Set());
   const [queryResult, setQueryResult] = useState<QueryResult | null>(null);
+
+  // AI highlights (separate from user/query highlights)
+  const [aiCitationHighlightedNodeIds, setAICitationHighlightedNodeIds] = useState<Set<string>>(new Set());
+  const [aiToolHighlightedNodeIds, setAIToolHighlightedNodeIds] = useState<Set<string>>(new Set());
+  const [isAIHighlightsEnabled, setAIHighlightsEnabled] = useState(true);
+
+  const toggleAIHighlights = useCallback(() => {
+    setAIHighlightsEnabled(prev => !prev);
+  }, []);
+
+  const clearAIToolHighlights = useCallback(() => {
+    setAIToolHighlightedNodeIds(new Set());
+  }, []);
   
   const clearQueryHighlights = useCallback(() => {
     setHighlightedNodeIds(new Set());
@@ -265,9 +284,9 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
     // Auto-open panel when references are added
     setCodePanelOpen(true);
     
-    // Also highlight in graph if nodeId provided
-    if (ref.nodeId) {
-      setHighlightedNodeIds(prev => new Set([...prev, ref.nodeId!]));
+    // Track AI highlights separately so they can be toggled off in the UI
+    if (ref.nodeId && ref.source === 'ai') {
+      setAICitationHighlightedNodeIds(prev => new Set([...prev, ref.nodeId!]));
     }
   }, []);
 
@@ -277,18 +296,12 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
       const removed = prev.filter(r => r.source === 'ai');
       const kept = prev.filter(r => r.source !== 'ai');
 
-      // Update graph highlights: remove nodeIds that were ONLY highlighted via AI refs
+      // Remove citation-based AI highlights for removed refs
       const removedNodeIds = new Set(removed.map(r => r.nodeId).filter(Boolean) as string[]);
       if (removedNodeIds.size > 0) {
-        setHighlightedNodeIds(prevHighlights => {
-          const next = new Set(prevHighlights);
-          for (const nodeId of removedNodeIds) {
-            const stillReferenced = kept.some(r => r.nodeId === nodeId);
-            const stillInQuery = queryResult?.nodeIds?.includes(nodeId) ?? false;
-            if (!stillReferenced && !stillInQuery) {
-              next.delete(nodeId);
-            }
-          }
+        setAICitationHighlightedNodeIds(prevIds => {
+          const next = new Set(prevIds);
+          for (const id of removedNodeIds) next.delete(id);
           return next;
         });
       }
@@ -490,6 +503,8 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
 
     // Refresh Code panel for the new question: keep user-pinned refs, clear old AI citations
     clearAICodeReferences();
+    // Also clear previous tool-driven AI highlights (highlight_in_graph)
+    clearAIToolHighlights();
 
     if (!isAgentReady) {
       // Try to initialize first
@@ -734,10 +749,10 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
                     }
                     
                     if (matchedIds.size > 0) {
-                      setHighlightedNodeIds(matchedIds);
+                      setAIToolHighlightedNodeIds(matchedIds);
                     }
                   } else if (rawIds.length > 0) {
-                    setHighlightedNodeIds(new Set(rawIds));
+                    setAIToolHighlightedNodeIds(new Set(rawIds));
                   }
                 }
               }
@@ -763,7 +778,7 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
       setIsChatLoading(false);
       setCurrentToolCalls([]);
     }
-  }, [chatMessages, isAgentReady, initializeAgent, resolveFilePath, findFileNodeId, addCodeReference, clearAICodeReferences]);
+  }, [chatMessages, isAgentReady, initializeAgent, resolveFilePath, findFileNodeId, addCodeReference, clearAICodeReferences, clearAIToolHighlights, graph]);
 
   const clearChat = useCallback(() => {
     setChatMessages([]);
@@ -776,11 +791,11 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
       const ref = prev.find(r => r.id === id);
       const newRefs = prev.filter(r => r.id !== id);
       
-      // Remove highlight if this was the only reference to that node
-      if (ref?.nodeId) {
-        const stillReferenced = newRefs.some(r => r.nodeId === ref.nodeId);
+      // Remove AI citation highlight if this was the only AI reference to that node
+      if (ref?.nodeId && ref.source === 'ai') {
+        const stillReferenced = newRefs.some(r => r.nodeId === ref.nodeId && r.source === 'ai');
         if (!stillReferenced) {
-          setHighlightedNodeIds(prev => {
+          setAICitationHighlightedNodeIds(prev => {
             const next = new Set(prev);
             next.delete(ref.nodeId!);
             return next;
@@ -833,6 +848,11 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
     setDepthFilter,
     highlightedNodeIds,
     setHighlightedNodeIds,
+    aiCitationHighlightedNodeIds,
+    aiToolHighlightedNodeIds,
+    isAIHighlightsEnabled,
+    toggleAIHighlights,
+    clearAIToolHighlights,
     queryResult,
     setQueryResult,
     clearQueryHighlights,