fix: Move useStream uo to thread view to prevent remounting (#477)

* fix: Move useStream uo to thread view to prevent remounting

* ui fixes

apps/web/src/components/v2/actions-renderer.tsx

@@ -1,27 +1,14 @@
 import { isAIMessageSDK, isHumanMessageSDK } from "@/lib/langchain-messages";
 import { UseStream, useStream } from "@langchain/langgraph-sdk/react";
 import { AssistantMessage } from "../thread/messages/ai";
+import { Dispatch, SetStateAction, useEffect, useMemo, useState } from "react";
 import {
-  Dispatch,
-  SetStateAction,
-  useEffect,
-  useMemo,
-  useRef,
-  useState,
-} from "react";
-import { ManagerGraphState } from "@open-swe/shared/open-swe/manager/types";
-import { useCancelStream } from "@/hooks/useCancelStream";
-import {
-  isCustomNodeEvent,
   CustomNodeEvent,
   INITIALIZE_NODE_ID,
   ACCEPTED_PLAN_NODE_ID,
   mapCustomEventsToSteps,
 } from "@open-swe/shared/open-swe/custom-node-events";
-import {
-  DO_NOT_RENDER_ID_PREFIX,
-  PLANNER_GRAPH_ID,
-} from "@open-swe/shared/constants";
+import { DO_NOT_RENDER_ID_PREFIX } from "@open-swe/shared/constants";
 import { Message } from "@langchain/langgraph-sdk";
 import { InitializeStep } from "../gen-ui/initialize-step";
 import { AcceptedPlanStep } from "../gen-ui/accepted-plan-step";
@@ -33,7 +20,6 @@ import { Interrupt } from "../thread/messages/interrupt";
 import { AlertCircle } from "lucide-react";
 import { ErrorState } from "./types";
 import { CollapsibleAlert } from "./collapsible-alert";
-import { TokenUsage } from "./token-usage";
 
 interface AcceptedPlanEventData {
   planTitle: string;
@@ -72,16 +58,13 @@ function isAcceptedPlanEvents(
   return events.every(isAcceptedPlanEvent);
 }
 
-interface ActionsRendererProps {
-  graphId: string;
-  threadId: string;
+interface ActionsRendererProps<
+  StateType extends PlannerGraphState | GraphState,
+> {
   runId?: string;
-  setProgrammerSession?: (
-    session: ManagerGraphState["programmerSession"],
-  ) => void;
-  programmerSession?: ManagerGraphState["programmerSession"];
-  setSelectedTab?: Dispatch<SetStateAction<"planner" | "programmer">>;
-  onStreamReady: (cancelFn: (() => void) | undefined) => void;
+  customNodeEvents: CustomNodeEvent[];
+  setCustomNodeEvents: Dispatch<SetStateAction<CustomNodeEvent[]>>;
+  stream: ReturnType<typeof useStream<StateType>>;
 }
 
 const getCustomNodeEventsFromMessages = (
@@ -106,65 +89,17 @@ const getCustomNodeEventsFromMessages = (
     .flat();
 };
 
-export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
-  graphId,
-  threadId,
+export function ActionsRenderer<
+  StateType extends PlannerGraphState | GraphState,
+>({
   runId,
-  setProgrammerSession,
-  programmerSession,
-  setSelectedTab,
-  onStreamReady,
-}: ActionsRendererProps) {
-  const [customNodeEvents, setCustomNodeEvents] = useState<CustomNodeEvent[]>(
-    [],
-  );
-  const joinedRunId = useRef<string | undefined>(undefined);
-  const [streamLoading, setStreamLoading] = useState(false);
+  customNodeEvents,
+  setCustomNodeEvents,
+  stream,
+}: ActionsRendererProps<StateType>) {
+  const [streamLoading, setStreamLoading] = useState(stream.isLoading);
   const [errorState, setErrorState] = useState<ErrorState | null>(null);
 
-  const stream = useStream<State>({
-    apiUrl: process.env.NEXT_PUBLIC_API_URL,
-    assistantId: graphId,
-    reconnectOnMount: true,
-    threadId,
-    onCustomEvent: (event) => {
-      if (isCustomNodeEvent(event)) {
-        setCustomNodeEvents((prev) => [...prev, event]);
-      }
-    },
-    fetchStateHistory: false,
-  });
-
-  useEffect(() => {
-    if (stream.error) {
-      const rawErrorMessage =
-        typeof stream.error === "object" && "message" in stream.error
-          ? (stream.error.message as string)
-          : "An unknown error occurred in the manager";
-
-      if (rawErrorMessage.includes("overloaded_error")) {
-        setErrorState({
-          message:
-            "An Anthropic overloaded error occurred. This error occurs when Anthropic APIs experience high traffic across all users.",
-          details: rawErrorMessage,
-        });
-      } else {
-        setErrorState({
-          message: rawErrorMessage,
-        });
-      }
-    } else {
-      setErrorState(null);
-    }
-  }, [stream.error]);
-
-  const { cancelRun } = useCancelStream<State>({
-    stream,
-    threadId,
-    runId,
-    streamName: graphId === "planner" ? "Planner" : "Programmer",
-  });
-
   const initializeEvents = useMemo(
     () =>
       customNodeEvents.filter(
@@ -190,6 +125,20 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
     initStatus = "done";
   }
 
+  // Filter out human & do not render messages
+  const filteredMessages = stream.messages?.filter(
+    (m) =>
+      !isHumanMessageSDK(m) &&
+      !(m.id && m.id.startsWith(DO_NOT_RENDER_ID_PREFIX)),
+  );
+  const isLastMessageHidden = !!(
+    stream.messages?.length > 0 &&
+    stream.messages[stream.messages.length - 1].id &&
+    stream.messages[stream.messages.length - 1].id?.startsWith(
+      DO_NOT_RENDER_ID_PREFIX,
+    )
+  );
+
   useEffect(() => {
     const allCustomEvents = getCustomNodeEventsFromMessages(stream.messages);
     if (!allCustomEvents?.length) {
@@ -219,79 +168,42 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
   // Clear streamLoading as soon as we get any content (agent has started running)
   useEffect(() => {
     const hasContent =
-      (stream.messages && stream.messages.length > 0) ||
-      customNodeEvents.length > 0;
+      filteredMessages.length > 0 || customNodeEvents.length > 0;
 
     if (hasContent && streamLoading) {
       setStreamLoading(false);
     }
   }, [stream.messages, customNodeEvents, streamLoading]);
 
-  // TODO: If the SDK changes go in, use this instead:
-  // stream.joinStream(runId, undefined, { streamMode: ["values", "messages", "custom"]}).catch(console.error);
   useEffect(() => {
-    if (runId && runId !== joinedRunId.current) {
-      joinedRunId.current = runId;
-      setStreamLoading(true);
-      stream
-        .joinStream(runId)
-        .catch(console.error)
-        .finally(() => setStreamLoading(false));
-    } else if (!runId) {
-      joinedRunId.current = undefined;
-    }
-  }, [runId, stream]);
+    if (stream.error) {
+      const rawErrorMessage =
+        typeof stream.error === "object" && "message" in stream.error
+          ? (stream.error.message as string)
+          : "An unknown error occurred in the manager";
 
-  useEffect(() => {
-    if (stream.isLoading) {
-      onStreamReady(cancelRun);
-    } else {
-      onStreamReady(undefined);
-    }
-  }, [onStreamReady, runId]); // Depend on runId instead of cancelRun to avoid infinite loops
-
-  // Filter out human & do not render messages
-  const filteredMessages = stream.messages?.filter(
-    (m) =>
-      !isHumanMessageSDK(m) &&
-      !(m.id && m.id.startsWith(DO_NOT_RENDER_ID_PREFIX)),
-  );
-  const isLastMessageHidden = !!(
-    stream.messages?.length > 0 &&
-    stream.messages[stream.messages.length - 1].id &&
-    stream.messages[stream.messages.length - 1].id?.startsWith(
-      DO_NOT_RENDER_ID_PREFIX,
-    )
-  );
-
-  // TODO: Need a better way to handle this. Not great like this...
-  useEffect(() => {
-    if (
-      "programmerSession" in stream.values &&
-      stream.values.programmerSession &&
-      (stream.values.programmerSession.runId !== programmerSession?.runId ||
-        stream.values.programmerSession.threadId !==
-          programmerSession?.threadId)
-    ) {
-      setProgrammerSession?.(stream.values.programmerSession);
-
-      // Only switch tabs from the planner ActionsRenderer to ensure proper timing
-      // This allows the accepted plan step to be visible before switching
-      if (graphId === PLANNER_GRAPH_ID) {
-        // Add a small delay to allow the accepted plan step to render first
-        setTimeout(() => {
-          setSelectedTab?.("programmer");
-        }, 2000);
+      if (rawErrorMessage.includes("overloaded_error")) {
+        setErrorState({
+          message:
+            "An Anthropic overloaded error occurred. This error occurs when Anthropic APIs experience high traffic across all users.",
+          details: rawErrorMessage,
+        });
+      } else {
+        setErrorState({
+          message: rawErrorMessage,
+        });
       }
+    } else {
+      setErrorState(null);
     }
-  }, [stream.values, graphId]);
+  }, [stream.error]);
 
-  if (streamLoading) {
+  if (streamLoading && !errorState) {
     return <LoadingActionsCardContent />;
   }
 
   return (
-    <div className="flex w-full flex-col gap-2">
+    <div className="flex h-full w-full flex-col gap-2 overflow-y-auto py-4">
       {initializeEvents.length > 0 && steps.length > 0 && (
         <InitializeStep
           status={initStatus}
@@ -339,7 +251,6 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
           icon={<AlertCircle className="size-4" />}
         />
       ) : null}
-      <TokenUsage tokenData={stream.values.tokenData} />
     </div>
   );
 }

apps/web/src/components/v2/thread-view.tsx

@@ -11,7 +11,7 @@ import { ThreadMetadata } from "./types";
 import { useStream } from "@langchain/langgraph-sdk/react";
 import { ManagerGraphState } from "@open-swe/shared/open-swe/manager/types";
 import { PlannerGraphState } from "@open-swe/shared/open-swe/planner/types";
-import { GraphState } from "@open-swe/shared/open-swe/types";
+import { GraphState, CacheMetrics } from "@open-swe/shared/open-swe/types";
 import { ActionsRenderer } from "./actions-renderer";
 import { ThemeToggle } from "../theme-toggle";
 import { HumanMessage } from "@langchain/core/messages";
@@ -23,7 +23,6 @@ import {
 import { useThreadStatus } from "@/hooks/useThreadStatus";
 import { cn } from "@/lib/utils";
 
-import { StickToBottom } from "use-stick-to-bottom";
 import {
   StickyToBottomContent,
   ScrollToBottom,
@@ -31,6 +30,12 @@ import {
 import { ManagerChat } from "./manager-chat";
 import { CancelStreamButton } from "./cancel-stream-button";
 import { ErrorState } from "./types";
+import {
+  CustomNodeEvent,
+  isCustomNodeEvent,
+} from "@open-swe/shared/open-swe/custom-node-events";
+import { StickToBottom } from "use-stick-to-bottom";
+import { TokenUsage } from "./token-usage";
 
 interface ThreadViewProps {
   stream: ReturnType<typeof useStream<ManagerGraphState>>;
@@ -55,6 +60,65 @@ export function ThreadView({
     useState<ManagerGraphState["programmerSession"]>();
   const [errorState, setErrorState] = useState<ErrorState | null>(null);
 
+  const [customPlannerNodeEvents, setCustomPlannerNodeEvents] = useState<
+    CustomNodeEvent[]
+  >([]);
+  const [customProgrammerNodeEvents, setCustomProgrammerNodeEvents] = useState<
+    CustomNodeEvent[]
+  >([]);
+
+  const plannerStream = useStream<PlannerGraphState>({
+    apiUrl: process.env.NEXT_PUBLIC_API_URL,
+    assistantId: PLANNER_GRAPH_ID,
+    reconnectOnMount: true,
+    threadId: plannerSession?.threadId,
+    onCustomEvent: (event) => {
+      if (isCustomNodeEvent(event)) {
+        setCustomPlannerNodeEvents((prev) => [...prev, event]);
+      }
+    },
+    fetchStateHistory: false,
+  });
+
+  const joinedPlannerRunId = useRef<string | undefined>(undefined);
+  useEffect(() => {
+    if (
+      plannerSession?.runId &&
+      plannerSession.runId !== joinedPlannerRunId.current
+    ) {
+      joinedPlannerRunId.current = plannerSession.runId;
+      plannerStream.joinStream(plannerSession.runId).catch(console.error);
+    } else if (!plannerSession?.runId) {
+      joinedPlannerRunId.current = undefined;
+    }
+  }, [plannerSession]);
+
+  const programmerStream = useStream<GraphState>({
+    apiUrl: process.env.NEXT_PUBLIC_API_URL,
+    assistantId: PROGRAMMER_GRAPH_ID,
+    reconnectOnMount: true,
+    threadId: programmerSession?.threadId,
+    onCustomEvent: (event) => {
+      if (isCustomNodeEvent(event)) {
+        setCustomProgrammerNodeEvents((prev) => [...prev, event]);
+      }
+    },
+    fetchStateHistory: false,
+  });
+
+  const joinedProgrammerRunId = useRef<string | undefined>(undefined);
+  useEffect(() => {
+    if (
+      programmerSession?.runId &&
+      programmerSession.runId !== joinedProgrammerRunId.current
+    ) {
+      joinedProgrammerRunId.current = programmerSession.runId;
+      plannerStream.joinStream(programmerSession.runId).catch(console.error);
+    } else if (!programmerSession?.runId) {
+      joinedProgrammerRunId.current = undefined;
+    }
+  }, [programmerSession]);
+
   useEffect(() => {
     if (
       stream?.values?.plannerSession &&
@@ -94,6 +158,27 @@ export function ThreadView({
     }
   }, [stream.error]);
 
+  useEffect(() => {
+    if (
+      plannerStream.values.programmerSession &&
+      (plannerStream.values.programmerSession.runId !==
+        programmerSession?.runId ||
+        plannerStream.values.programmerSession.threadId !==
+          programmerSession?.threadId)
+    ) {
+      setProgrammerSession?.(plannerStream.values.programmerSession);
+
+      // Only switch tabs from the planner ActionsRenderer to ensure proper timing
+      // This allows the accepted plan step to be visible before switching
+      if (selectedTab === PLANNER_GRAPH_ID) {
+        // Add a small delay to allow the accepted plan step to render first
+        setTimeout(() => {
+          setSelectedTab?.("programmer");
+        }, 2000);
+      }
+    }
+  }, [plannerStream.values, selectedTab]);
+
   const { status: realTimeStatus } = useThreadStatus(displayThread.id);
 
   const getStatusDotColor = (status: string) => {
@@ -111,9 +196,6 @@ export function ThreadView({
     }
   };
 
-  const plannerCancelRef = useRef<(() => void) | null>(null);
-  const programmerCancelRef = useRef<(() => void) | null>(null);
-
   const cancelRun = () => {
     // TODO: ideally this calls stream.client.runs.cancel(threadId, runId)
     stream.stop();
@@ -187,7 +269,7 @@ export function ThreadView({
       </div>
 
       {/* Main Content - Split Layout */}
-      <div className="flex h-full w-full pt-12">
+      <div className="flex w-full pt-12">
         <ManagerChat
           messages={filteredMessages}
           chatInput={chatInput}
@@ -198,121 +280,146 @@ export function ThreadView({
           errorState={errorState}
         />
         {/* Right Side - Actions & Plan */}
-        <div className="flex h-full flex-1 flex-col">
-          <div className="relative flex-1">
-            <StickToBottom
-              className="absolute inset-0"
-              initial={true}
+        <div
+          className="flex flex-1 flex-col px-4 pt-4"
+          style={{ height: "calc(100vh - 3rem)" }}
+        >
+          <div className="min-h-0 flex-1">
+            <Tabs
+              defaultValue="planner"
+              className="flex h-full w-full flex-col"
+              value={selectedTab}
+              onValueChange={(value) =>
+                setSelectedTab(value as "planner" | "programmer")
+              }
             >
-              <StickyToBottomContent
-                className="h-full overflow-y-auto"
-                contentClassName="space-y-4 p-4"
-                content={
-                  <Tabs
-                    defaultValue="planner"
-                    className="w-full"
-                    value={selectedTab}
-                    onValueChange={(value) =>
-                      setSelectedTab(value as "planner" | "programmer")
+              <div className="flex flex-shrink-0 items-center justify-between">
+                <TabsList className="bg-muted/70 dark:bg-gray-800">
+                  <TabsTrigger value="planner">Planner</TabsTrigger>
+                  <TabsTrigger value="programmer">Programmer</TabsTrigger>
+                </TabsList>
+
+                <div className="flex items-center justify-center gap-2">
+                  {selectedTab === "planner" && plannerStream.isLoading && (
+                    <CancelStreamButton
+                      stream={plannerStream}
+                      threadId={plannerSession?.threadId}
+                      runId={plannerSession?.runId}
+                      streamName="Planner"
+                    />
+                  )}
+
+                  {selectedTab === "programmer" &&
+                    programmerStream.isLoading && (
+                      <CancelStreamButton
+                        stream={programmerStream}
+                        threadId={programmerSession?.threadId}
+                        runId={programmerSession?.runId}
+                        streamName="Programmer"
+                      />
+                    )}
+                  <TokenUsage
+                    tokenData={
+                      [
+                        plannerStream.values.tokenData,
+                        programmerStream.values.tokenData,
+                      ].filter(Boolean) as CacheMetrics[]
                     }
-                  >
-                    <div className="flex items-center justify-between">
-                      <TabsList className="bg-muted/70 dark:bg-gray-800">
-                        <TabsTrigger value="planner">Planner</TabsTrigger>
-                        <TabsTrigger value="programmer">Programmer</TabsTrigger>
-                      </TabsList>
+                  />
+                </div>
+              </div>
 
-                      <div className="flex gap-2">
-                        {selectedTab === "planner" &&
-                          plannerCancelRef.current && (
-                            <CancelStreamButton
-                              stream={stream}
-                              threadId={plannerSession?.threadId}
-                              runId={plannerSession?.runId}
-                              streamName="Planner"
-                            />
-                          )}
-
-                        {selectedTab === "programmer" &&
-                          programmerCancelRef.current && (
-                            <CancelStreamButton
-                              stream={stream}
-                              threadId={programmerSession?.threadId}
-                              runId={programmerSession?.runId}
-                              streamName="Programmer"
-                            />
-                          )}
-                      </div>
-                    </div>
-
-                    <TabsContent value="planner">
-                      <Card className="border-border bg-card px-0 py-4 dark:bg-gray-950">
-                        <CardContent className="space-y-2 p-3 pt-0">
-                          {plannerSession && (
-                            <ActionsRenderer<PlannerGraphState>
-                              graphId={PLANNER_GRAPH_ID}
-                              threadId={plannerSession.threadId}
-                              runId={plannerSession.runId}
-                              setProgrammerSession={setProgrammerSession}
-                              programmerSession={programmerSession}
-                              setSelectedTab={setSelectedTab}
-                              onStreamReady={(cancelFn) => {
-                                if (cancelFn) {
-                                  plannerCancelRef.current = cancelFn;
-                                } else {
-                                  plannerCancelRef.current = null;
-                                }
-                              }}
-                            />
-                          )}
-                          {!plannerSession && (
-                            <div className="flex items-center justify-center gap-2 py-8">
-                              <Clock className="text-muted-foreground size-4" />
-                              <span className="text-muted-foreground text-sm">
-                                No planner session
-                              </span>
-                            </div>
-                          )}
-                        </CardContent>
-                      </Card>
-                    </TabsContent>
-                    <TabsContent value="programmer">
-                      <Card className="border-border bg-card px-0 py-4 dark:bg-gray-950">
-                        <CardContent className="space-y-2 p-3 pt-0">
-                          {programmerSession && (
-                            <ActionsRenderer<GraphState>
-                              graphId={PROGRAMMER_GRAPH_ID}
-                              threadId={programmerSession.threadId}
-                              runId={programmerSession.runId}
-                              onStreamReady={(cancelFn) => {
-                                if (cancelFn) {
-                                  programmerCancelRef.current = cancelFn;
-                                } else {
-                                  programmerCancelRef.current = null;
-                                }
-                              }}
-                            />
-                          )}
-                          {!programmerSession && (
-                            <div className="flex items-center justify-center gap-2 py-8">
-                              <Terminal className="text-muted-foreground size-4" />
-                              <span className="text-muted-foreground text-sm">
-                                No programmer session
-                              </span>
-                            </div>
-                          )}
-                        </CardContent>
-                      </Card>
-                    </TabsContent>
-                  </Tabs>
-                }
-                footer={
-                  <div className="absolute bottom-4 left-1/2 -translate-x-1/2">
-                    <ScrollToBottom className="animate-in fade-in-0 zoom-in-95" />
-                  </div>
-                }
-              />
-            </StickToBottom>
+              <TabsContent
+                value="planner"
+                className="mb-10"
+              >
+                <Card className="border-border bg-card relative h-full p-0 dark:bg-gray-950">
+                  <CardContent className="h-full p-0">
+                    <StickToBottom
+                      className="absolute inset-0 h-full"
+                      initial={true}
+                    >
+                      <StickyToBottomContent
+                        className="scrollbar-pretty-auto h-full"
+                        content={
+                          <>
+                            {plannerSession ? (
+                              <div className="scrollbar-pretty-auto overflow-y-auto px-2">
+                                <ActionsRenderer<PlannerGraphState>
+                                  runId={plannerSession.runId}
+                                  customNodeEvents={customPlannerNodeEvents}
+                                  setCustomNodeEvents={
+                                    setCustomPlannerNodeEvents
+                                  }
+                                  stream={plannerStream}
+                                />
+                              </div>
+                            ) : (
+                              <div className="flex items-center justify-center gap-2 py-8">
+                                <Clock className="text-muted-foreground size-4" />
+                                <span className="text-muted-foreground text-sm">
+                                  No planner session
+                                </span>
+                              </div>
+                            )}
+                          </>
+                        }
+                        footer={
+                          <div className="absolute right-0 bottom-4 left-0 flex w-full justify-center">
+                            <ScrollToBottom className="animate-in fade-in-0 zoom-in-95" />
+                          </div>
+                        }
+                      />
+                    </StickToBottom>
+                  </CardContent>
+                </Card>
+              </TabsContent>
+              <TabsContent
+                value="programmer"
+                className="mb-10"
+              >
+                <Card className="border-border bg-card relative h-full p-0 dark:bg-gray-950">
+                  <CardContent className="h-full p-0">
+                    <StickToBottom
+                      className="absolute inset-0 h-full"
+                      initial={true}
+                    >
+                      <StickyToBottomContent
+                        className="scrollbar-pretty-auto h-full"
+                        content={
+                          <>
+                            {programmerSession ? (
+                              <div className="scrollbar-pretty-auto overflow-y-auto px-2">
+                                <ActionsRenderer<GraphState>
+                                  runId={programmerSession.runId}
+                                  customNodeEvents={customProgrammerNodeEvents}
+                                  setCustomNodeEvents={
+                                    setCustomProgrammerNodeEvents
+                                  }
+                                  stream={programmerStream}
+                                />
+                              </div>
+                            ) : (
+                              <div className="flex items-center justify-center gap-2 py-8">
+                                <Clock className="text-muted-foreground size-4" />
+                                <span className="text-muted-foreground text-sm">
+                                  No programmer session
+                                </span>
+                              </div>
+                            )}
+                          </>
+                        }
+                        footer={
+                          <div className="absolute right-0 bottom-4 left-0 flex w-full justify-center">
+                            <ScrollToBottom className="animate-in fade-in-0 zoom-in-95" />
+                          </div>
+                        }
+                      />
+                    </StickToBottom>
+                  </CardContent>
+                </Card>
+              </TabsContent>
+            </Tabs>
           </div>
         </div>
       </div>

apps/web/src/components/v2/token-usage.tsx

@@ -1,54 +1,135 @@
 import { CacheMetrics } from "@open-swe/shared/open-swe/types";
 import { calculateCostSavings } from "@open-swe/shared/caching";
+import { Badge } from "../ui/badge";
+import { Separator } from "../ui/separator";
 import {
-  Tooltip,
-  TooltipContent,
-  TooltipProvider,
-  TooltipTrigger,
-} from "../ui/tooltip";
-import { ChartNoAxesColumnIncreasing } from "lucide-react";
+  HoverCard,
+  HoverCardContent,
+  HoverCardTrigger,
+} from "../ui/hover-card";
+import {
+  ChartNoAxesColumnIncreasing,
+  Coins,
+  TrendingUp,
+  Zap,
+} from "lucide-react";
 
-export function TokenUsage({ tokenData }: { tokenData?: CacheMetrics }) {
-  if (!tokenData) return null;
+interface TokenUsageProps {
+  tokenData?: CacheMetrics[];
+}
 
-  const metrics = calculateCostSavings(tokenData);
-  return (
-    <div className="mt-4 ml-auto flex">
-      <TooltipProvider>
-        <Tooltip>
-          <TooltipTrigger>
-            <ChartNoAxesColumnIncreasing />
-          </TooltipTrigger>
-          <TooltipContent className="flex w-full flex-col gap-1 text-sm">
-            <p>Token usage data on actions where caching is enabled:</p>
-            <span className="flex w-full items-center justify-between">
-              <p>Input Tokens:</p>
-              <p>{metrics.totalInputTokens.toLocaleString()}</p>
-            </span>
-            <span className="flex w-full items-center justify-between">
-              <p>Output Tokens:</p>
-              <p>{metrics.totalOutputTokens.toLocaleString()}</p>
-            </span>
-            <span className="flex w-full items-center justify-between">
-              <p>Total Tokens:</p>
-              <p>{metrics.totalTokens.toLocaleString()}</p>
-            </span>
-
-            <span className="flex w-full items-center justify-between">
-              <p>Output Tokens Cost:</p>
-              <p>${metrics.totalOutputTokensCost.toFixed(2)}</p>
-            </span>
-            <span className="flex w-full items-center justify-between">
-              <p>Cache Savings:</p>
-              <p>${metrics.totalSavings.toFixed(2)}</p>
-            </span>
-            <span className="flex w-full items-center justify-between">
-              <p>Total Cost:</p>
-              <p>${metrics.totalCost.toFixed(2)}</p>
-            </span>
-          </TooltipContent>
-        </Tooltip>
-      </TooltipProvider>
-    </div>
+function mergeTokenData(tokenDataArray: CacheMetrics[]): CacheMetrics {
+  return tokenDataArray.reduce(
+    (merged, current) => ({
+      cacheCreationInputTokens:
+        merged.cacheCreationInputTokens + current.cacheCreationInputTokens,
+      cacheReadInputTokens:
+        merged.cacheReadInputTokens + current.cacheReadInputTokens,
+      inputTokens: merged.inputTokens + current.inputTokens,
+      outputTokens: merged.outputTokens + current.outputTokens,
+    }),
+    {
+      cacheCreationInputTokens: 0,
+      cacheReadInputTokens: 0,
+      inputTokens: 0,
+      outputTokens: 0,
+    },
+  );
+}
+
+export function TokenUsage({ tokenData }: TokenUsageProps) {
+  if (!tokenData || tokenData.length === 0) return null;
+
+  const mergedTokenData = mergeTokenData(tokenData);
+  const metrics = calculateCostSavings(mergedTokenData);
+
+  return (
+    <HoverCard>
+      <HoverCardTrigger asChild>
+        <button className="hover:bg-muted/50 ml-auto flex items-center gap-2 rounded-md p-2 transition-colors">
+          <ChartNoAxesColumnIncreasing className="h-4 w-4" />
+          <Badge
+            variant="secondary"
+            className="text-xs"
+          >
+            {tokenData.length} agent{tokenData.length !== 1 ? "s" : ""}
+          </Badge>
+        </button>
+      </HoverCardTrigger>
+      <HoverCardContent className="w-80">
+        <div className="space-y-4">
+          <div className="flex items-center gap-2">
+            <ChartNoAxesColumnIncreasing className="h-4 w-4" />
+            <h4 className="text-sm font-semibold">Token Usage</h4>
+          </div>
+
+          <div className="grid grid-cols-2 gap-4">
+            <div className="space-y-1">
+              <div className="flex items-center gap-1.5">
+                <Zap className="h-3 w-3 text-blue-500" />
+                <span className="text-muted-foreground text-xs font-medium">
+                  Input
+                </span>
+              </div>
+              <p className="text-sm font-semibold">
+                {metrics.totalInputTokens.toLocaleString()}
+              </p>
+            </div>
+            <div className="space-y-1">
+              <div className="flex items-center gap-1.5">
+                <TrendingUp className="h-3 w-3 text-green-500" />
+                <span className="text-muted-foreground text-xs font-medium">
+                  Output
+                </span>
+              </div>
+              <p className="text-sm font-semibold">
+                {metrics.totalOutputTokens.toLocaleString()}
+              </p>
+            </div>
+          </div>
+
+          <Separator />
+
+          <div className="space-y-3">
+            <div className="flex items-center justify-between">
+              <span className="text-muted-foreground text-xs font-medium">
+                Total Tokens
+              </span>
+              <span className="text-sm font-semibold">
+                {metrics.totalTokens.toLocaleString()}
+              </span>
+            </div>
+
+            <div className="space-y-2">
+              <div className="flex items-center justify-between">
+                <div className="flex items-center gap-1.5">
+                  <Coins className="h-3 w-3 text-amber-500" />
+                  <span className="text-muted-foreground text-xs font-medium">
+                    Cost
+                  </span>
+                </div>
+                <span className="text-sm font-semibold">
+                  ${metrics.totalCost.toFixed(2)}
+                </span>
+              </div>
+
+              {metrics.totalSavings > 0 && (
+                <div className="flex items-center justify-between">
+                  <span className="text-xs font-medium text-green-600">
+                    Cache Savings
+                  </span>
+                  <Badge
+                    variant="outline"
+                    className="border-green-200 text-green-600"
+                  >
+                    -${metrics.totalSavings.toFixed(2)}
+                  </Badge>
+                </div>
+              )}
+            </div>
+          </div>
+        </div>
+      </HoverCardContent>
+    </HoverCard>
   );
 }