Merge pull request #348 from traceroot-ai/xinwei_agent_ui_improve_v3

[UI][Agent] Make the UI of agent chat more aligned with the trace and log section [1/n]

Author: XinweiHe

ui/src/app/explore/page.tsx

@@ -16,6 +16,40 @@ import AgentPanel from "@/components/agent-panel/AgentPanel";
 import { Span, Trace as TraceType } from "@/models/trace";
 import { initializeProviders } from "@/utils/provider";
 
+// Custom hook for persistent state with localStorage
+function usePersistentState<T>(
+  key: string,
+  defaultValue: T,
+): [T, (value: T | ((prev: T) => T)) => void] {
+  const [state, setState] = useState<T>(() => {
+    if (typeof window === "undefined") return defaultValue;
+    try {
+      const saved = localStorage.getItem(key);
+      return saved !== null ? JSON.parse(saved) : defaultValue;
+    } catch {
+      return defaultValue;
+    }
+  });
+
+  const setPersistentState = useCallback(
+    (value: T | ((prev: T) => T)) => {
+      setState((prev) => {
+        const newValue =
+          typeof value === "function" ? (value as (prev: T) => T)(prev) : value;
+        try {
+          localStorage.setItem(key, JSON.stringify(newValue));
+        } catch {
+          // Ignore localStorage errors
+        }
+        return newValue;
+      });
+    },
+    [key],
+  );
+
+  return [state, setPersistentState];
+}
+
 export default function Explore() {
   const [selectedTraceIds, setSelectedTraceIds] = useState<string[]>([]);
   const [selectedSpanIds, setSelectedSpanIds] = useState<string[]>([]);
@@ -39,6 +73,13 @@ export default function Explore() {
   const [loading, setLoading] = useState(true);
   const [hasTraceIdInUrl, setHasTraceIdInUrl] = useState(false);
 
+  // Agent panel state - persisted to localStorage
+  const [agentOpen, setAgentOpen] = usePersistentState("agentPanelOpen", false);
+
+  const handleAgentToggle = useCallback(() => {
+    setAgentOpen((prev) => !prev);
+  }, [setAgentOpen]);
+
   // Initialize providers
   useEffect(() => {
     initializeProviders();
@@ -181,32 +222,38 @@ export default function Explore() {
 
   // TODO (xinwei): Add ProtectedRoute
   return (
-    <AgentPanel
-      traceId={selectedTraceIds.length === 1 ? selectedTraceIds[0] : undefined}
-      traceIds={selectedTraceIds}
-      spanIds={selectedSpanIds}
-      queryStartTime={timeRange?.start}
-      queryEndTime={timeRange?.end}
-      onSpanSelect={(spanId) => handleSpanSelect([spanId])}
-    >
-      <div className="h-screen flex flex-col">
-        <ExploreHeader
-          onSearch={handleSearch}
-          onClearSearch={handleClearSearch}
-          onLogSearchValueChange={handleLogSearchValueChange}
-          onMetadataSearchTermsChange={handleMetadataSearchTermsChange}
-          searchDisabled={loading || hasTraceIdInUrl}
-          selectedTimeRange={selectedTimeRange}
-          onTimeRangeSelect={handleTimeRangeSelect}
-          onCustomTimeRangeSelect={handleCustomTimeRangeSelect}
-          currentTimezone={timezone}
-          timeDisabled={loading || hasTraceIdInUrl}
-          onRefresh={handleRefresh}
-          refreshDisabled={loading || hasTraceIdInUrl}
-          viewType={viewType}
-          onViewTypeChange={setViewType}
-        />
-        <div className="flex-1 overflow-hidden">
+    <div className="h-screen flex flex-col">
+      <ExploreHeader
+        onSearch={handleSearch}
+        onClearSearch={handleClearSearch}
+        onLogSearchValueChange={handleLogSearchValueChange}
+        onMetadataSearchTermsChange={handleMetadataSearchTermsChange}
+        searchDisabled={loading || hasTraceIdInUrl}
+        selectedTimeRange={selectedTimeRange}
+        onTimeRangeSelect={handleTimeRangeSelect}
+        onCustomTimeRangeSelect={handleCustomTimeRangeSelect}
+        currentTimezone={timezone}
+        timeDisabled={loading || hasTraceIdInUrl}
+        onRefresh={handleRefresh}
+        refreshDisabled={loading || hasTraceIdInUrl}
+        viewType={viewType}
+        onViewTypeChange={setViewType}
+        agentOpen={agentOpen}
+        onAgentToggle={handleAgentToggle}
+      />
+      <div className="flex-1 overflow-hidden">
+        <AgentPanel
+          traceId={
+            selectedTraceIds.length === 1 ? selectedTraceIds[0] : undefined
+          }
+          traceIds={selectedTraceIds}
+          spanIds={selectedSpanIds}
+          queryStartTime={timeRange?.start}
+          queryEndTime={timeRange?.end}
+          onSpanSelect={(spanId) => handleSpanSelect([spanId])}
+          isOpen={agentOpen}
+          onToggle={handleAgentToggle}
+        >
           <ResizablePanel
             leftPanel={
               <Trace
@@ -243,8 +290,8 @@ export default function Explore() {
             maxLeftWidth={60}
             defaultLeftWidth={46}
           />
-        </div>
+        </AgentPanel>
       </div>
-    </AgentPanel>
+    </div>
   );
 }

ui/src/components/agent-panel/AgentPanel.tsx

@@ -1,8 +1,6 @@
 "use client";
 
 import React, { useState, useRef, useEffect } from "react";
-import { RiRobot2Line } from "react-icons/ri";
-import { X } from "lucide-react";
 import Agent from "@/components/right-panel/agent/Agent";
 
 interface AgentPanelProps {
@@ -13,6 +11,8 @@ interface AgentPanelProps {
   queryEndTime?: Date;
   onSpanSelect?: (spanId: string) => void;
   onViewTypeChange?: (viewType: "log" | "trace") => void;
+  isOpen: boolean;
+  onToggle: () => void;
   children: React.ReactNode;
 }
 
@@ -24,33 +24,34 @@ export default function AgentPanel({
   queryEndTime,
   onSpanSelect,
   onViewTypeChange,
+  isOpen,
+  onToggle,
   children,
 }: AgentPanelProps) {
-  const [isOpen, setIsOpen] = useState(false);
-  const [isHovered, setIsHovered] = useState(false);
   const [panelWidth, setPanelWidth] = useState(30); // Track panel width percentage (default 30%)
   const [hasBeenOpened, setHasBeenOpened] = useState(false); // Track if panel has ever been opened
   const panelRef = useRef<HTMLDivElement>(null);
 
   const MIN_WIDTH = 20;
   const MAX_WIDTH = 35;
 
-  // Ensure panel stays closed unless explicitly opened
+  // Track when panel has been opened at least once
   useEffect(() => {
-    // Do NOT auto-open when traceId changes
-    // Panel should only open via user click on the robot icon
-  }, [traceId, traceIds]);
+    if (isOpen && !hasBeenOpened) {
+      setHasBeenOpened(true);
+    }
+  }, [isOpen, hasBeenOpened]);
 
   // Handle ESC key to close panel
   useEffect(() => {
     const handleEsc = (e: KeyboardEvent) => {
       if (e.key === "Escape" && isOpen) {
-        setIsOpen(false);
+        onToggle();
       }
     };
     window.addEventListener("keydown", handleEsc);
     return () => window.removeEventListener("keydown", handleEsc);
-  }, [isOpen]);
+  }, [isOpen, onToggle]);
 
   // Handle Cmd+Shift+I (Mac) or Ctrl+Shift+I (Windows) to toggle panel
   useEffect(() => {
@@ -63,39 +64,13 @@ export default function AgentPanel({
       ) {
         e.preventDefault();
         e.stopPropagation();
-        if (!isOpen) {
-          setHasBeenOpened(true);
-        }
-        setIsOpen((prev) => !prev);
+        onToggle();
       }
     };
     window.addEventListener("keydown", handleToggleShortcut, true); // Use capture phase
     return () =>
       window.removeEventListener("keydown", handleToggleShortcut, true);
-  }, [isOpen]);
-
-  // Calculate the icon position based on panel width when open
-  useEffect(() => {
-    if (isOpen && panelRef.current) {
-      const updateIconPosition = () => {
-        const windowWidth = window.innerWidth;
-        const panelPixelWidth = (windowWidth * panelWidth) / 100;
-        setPanelWidth(panelWidth);
-      };
-      updateIconPosition();
-      window.addEventListener("resize", updateIconPosition);
-      return () => window.removeEventListener("resize", updateIconPosition);
-    }
-  }, [isOpen, panelWidth]);
-
-  // Calculate icon right position: when closed = 0 (50% hidden), when open = at panel left edge
-  const getIconRightPosition = () => {
-    if (!isOpen) {
-      return isHovered ? "4px" : "-12px"; // Half-hidden when closed, centered on edge
-    }
-    // When open, position at the panel's left edge (border)
-    return `calc(${panelWidth}% - 16px)`; // Centered on the border line
-  };
+  }, [onToggle]);
 
   return (
     <div className="relative w-full h-full flex overflow-hidden">
@@ -119,7 +94,7 @@ export default function AgentPanel({
           ref={panelRef}
         >
           {/* Agent content area */}
-          <div className="h-full overflow-hidden pt-2">
+          <div className="h-full overflow-hidden pt-1">
             {/* Only render Agent when opened at least once */}
             {hasBeenOpened && (
               <Agent
@@ -171,20 +146,6 @@ export default function AgentPanel({
           }}
         />
       )}
-
-      {/* Robot icon trigger - only visible when closed */}
-      {!isOpen && (
-        <div
-          className="fixed bottom-4 right-4 z-50 cursor-pointer"
-          onClick={(e) => {
-            e.stopPropagation();
-            setHasBeenOpened(true);
-            setIsOpen(true);
-          }}
-        >
-          <RiRobot2Line className="h-7 w-7 text-black dark:text-white drop-shadow-lg" />
-        </div>
-      )}
     </div>
   );
 }

ui/src/components/explore/ExploreHeader.tsx

@@ -34,6 +34,10 @@ interface ExploreHeaderProps {
   // Mode toggle props
   viewType: ViewType;
   onViewTypeChange: (type: ViewType) => void;
+
+  // Agent panel props
+  agentOpen?: boolean;
+  onAgentToggle?: () => void;
 }
 
 export default function ExploreHeader({
@@ -51,6 +55,8 @@ export default function ExploreHeader({
   refreshDisabled = false,
   viewType,
   onViewTypeChange,
+  agentOpen,
+  onAgentToggle,
 }: ExploreHeaderProps) {
   return (
     <div className="sticky top-0 z-10 bg-white dark:bg-zinc-950 pt-1 pl-6 pr-2 pb-1 border-b border-zinc-200 dark:border-zinc-700">
@@ -73,7 +79,12 @@ export default function ExploreHeader({
             currentTimezone={currentTimezone}
             disabled={timeDisabled}
           />
-          <ModeToggle viewType={viewType} onViewTypeChange={onViewTypeChange} />
+          <ModeToggle
+            viewType={viewType}
+            onViewTypeChange={onViewTypeChange}
+            agentOpen={agentOpen}
+            onAgentToggle={onAgentToggle}
+          />
         </div>
       </div>
     </div>

ui/src/components/right-panel/ModeToggle.tsx

@@ -2,21 +2,27 @@
 
 import React from "react";
 import { Telescope, FileCode2 } from "lucide-react";
+import { RiRobot2Line } from "react-icons/ri";
 import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group";
+import { Toggle } from "@/components/ui/toggle";
 
 export type ViewType = "log" | "trace";
 
 interface ModeToggleProps {
   viewType: ViewType;
   onViewTypeChange: (type: ViewType) => void;
+  agentOpen?: boolean;
+  onAgentToggle?: () => void;
 }
 
 export default function ModeToggle({
   viewType,
   onViewTypeChange,
+  agentOpen = false,
+  onAgentToggle,
 }: ModeToggleProps) {
   return (
-    <div className="flex justify-end p-4">
+    <div className="flex justify-end items-center p-4 gap-3">
       <ToggleGroup
         type="single"
         value={viewType}
@@ -33,6 +39,17 @@ export default function ModeToggle({
           <Telescope size={22} />
         </ToggleGroupItem>
       </ToggleGroup>
+      {onAgentToggle && (
+        <Toggle
+          pressed={agentOpen}
+          onPressedChange={onAgentToggle}
+          aria-label="Toggle agent panel"
+          variant="outline"
+          size="lg"
+        >
+          <RiRobot2Line size={22} />
+        </Toggle>
+      )}
     </div>
   );
 }