feat: usage counter stories and basic UI (#516)

* feat: usage counter stories and basic UI

* chore: fixed prettier check

* feat: saving usage to assistant message on stream & reasoning mode selection fix

* fix: stubs for usage counters, fixed stories and refactored UsageCounters

* chore: removed console.log

* feat: storing usage object for chat thread instead of for each assistant message & showing suggestion to open new chat, if prompt tokens are exceeding 30k limit boundary

* chore: better place for calculateInputTokens utility function

* refactoring(middleware): simplified chat response middleware

* chore: remove unused imports

* refactoring: remove duplicate code

* feat: calculating recommended maximum input limit based on model n_ctx

* fix: fix tests

* fix: moved UsageCounter to separate component modlet out of ChatContent scope, removed comments

Author: alashchev17

refact-agent/gui/src/app/middleware.ts

@@ -11,6 +11,7 @@ import {
   restoreChat,
   newIntegrationChat,
   chatResponse,
+  setThreadUsage,
 } from "../features/Chat/Thread";
 import { statisticsApi } from "../services/refact/statistics";
 import { integrationsApi } from "../services/refact/integrations";
@@ -32,14 +33,11 @@ import { nextTip } from "../features/TipOfTheDay";
 import { telemetryApi } from "../services/refact/telemetry";
 import { CONFIG_PATH_URL, FULL_PATH_URL } from "../services/refact/consts";
 import { resetConfirmationInteractedState } from "../features/ToolConfirmation/confirmationSlice";
-import {
-  getAgentUsageCounter,
-  getMaxFreeAgentUsage,
-} from "../features/Chat/Thread/utils";
 import {
   updateAgentUsage,
   updateMaxAgentUsageAmount,
 } from "../features/AgentUsage/agentUsageSlice";
+import { isChatResponseChoice } from "../services/refact";
 
 const AUTH_ERROR_MESSAGE =
   "There is an issue with your API key. Check out your API Key or re-login";
@@ -97,15 +95,22 @@ startListening({
     // saving to store agent_usage counter from the backend, only one chunk has this field.
     const { payload } = action;
 
-    if ("refact_agent_request_available" in payload) {
-      const agentUsageCounter = getAgentUsageCounter(payload);
-
-      dispatch(updateAgentUsage(agentUsageCounter ?? null));
-    }
-
-    if ("refact_agent_max_request_num" in payload) {
-      const maxFreeAgentUsage = getMaxFreeAgentUsage(payload);
-      dispatch(updateMaxAgentUsageAmount(maxFreeAgentUsage));
+    if (isChatResponseChoice(payload)) {
+      const {
+        usage,
+        refact_agent_max_request_num,
+        refact_agent_request_available,
+      } = payload;
+      const actions = [
+        updateAgentUsage(refact_agent_request_available),
+        updateMaxAgentUsageAmount(refact_agent_max_request_num),
+      ];
+
+      actions.forEach((action) => dispatch(action));
+
+      if (usage) {
+        dispatch(setThreadUsage({ chatId: payload.id, usage }));
+      }
     }
   },
 });

refact-agent/gui/src/components/ChatContent/AssistantInput.tsx

@@ -21,6 +21,7 @@ export const AssistantInput: React.FC<ChatInputProps> = ({
 }) => {
   const [sendTelemetryEvent] =
     telemetryApi.useLazySendTelemetryChatEventQuery();
+
   const handleCopy = useCallback(
     (text: string) => {
       // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition

refact-agent/gui/src/components/ChatContent/ChatContent.tsx

@@ -24,6 +24,7 @@ import {
   selectIsWaiting,
   selectMessages,
   selectThread,
+  selectThreadUsage,
 } from "../../features/Chat/Thread/selectors";
 import { takeWhile } from "../../utils";
 import { GroupedDiffs } from "./DiffContent";
@@ -32,6 +33,7 @@ import { popBackTo } from "../../features/Pages/pagesSlice";
 import { ChatLinks, UncommittedChangesWarning } from "../ChatLinks";
 import { telemetryApi } from "../../services/refact/telemetry";
 import { PlaceHolderText } from "./PlaceHolderText";
+import { UsageCounter } from "../UsageCounter";
 
 export type ChatContentProps = {
   onRetry: (index: number, question: UserMessage["content"]) => void;
@@ -47,6 +49,7 @@ export const ChatContent: React.FC<ChatContentProps> = ({
   const messages = useAppSelector(selectMessages);
   const isStreaming = useAppSelector(selectIsStreaming);
   const thread = useAppSelector(selectThread);
+  const threadUsage = useAppSelector(selectThreadUsage);
   const isConfig = thread.mode === "CONFIGURE";
   const isWaiting = useAppSelector(selectIsWaiting);
   const [sendTelemetryEvent] =
@@ -113,6 +116,7 @@ export const ChatContent: React.FC<ChatContentProps> = ({
         {messages.length === 0 && <PlaceHolderText />}
         {renderMessages(messages, onRetryWrapper)}
         <UncommittedChangesWarning />
+        {threadUsage && <UsageCounter usage={threadUsage} />}
 
         <Container py="4">
           <Spinner spinning={isWaiting} />

refact-agent/gui/src/components/ChatForm/SuggestNewChat/SuggestNewChat.tsx

@@ -96,7 +96,7 @@ export const SuggestNewChat = ({
           usage limits faster.
         </Text>
         <Flex align="center" gap="3" flexShrink="0">
-          <Link size="1" onClick={onCreateNewChat}>
+          <Link size="1" onClick={onCreateNewChat} color="indigo">
             Start a new chat
           </Link>
           <IconButton

refact-agent/gui/src/components/UsageCounter/UsageCounter.fixtures.ts

@@ -0,0 +1,27 @@
+import { Usage } from "../../services/refact";
+
+export const USAGE_COUNTER_STUB_GPT: Usage = {
+  completion_tokens: 30,
+  prompt_tokens: 3391,
+  total_tokens: 3421,
+  completion_tokens_details: {
+    accepted_prediction_tokens: 0,
+    audio_tokens: 0,
+    reasoning_tokens: 0,
+    rejected_prediction_tokens: 0,
+  },
+  prompt_tokens_details: {
+    audio_tokens: 0,
+    cached_tokens: 3328,
+  },
+};
+
+export const USAGE_COUNTER_STUB_ANTHROPIC: Usage = {
+  completion_tokens: 142,
+  prompt_tokens: 5,
+  total_tokens: 147,
+  completion_tokens_details: null,
+  prompt_tokens_details: null,
+  cache_creation_input_tokens: 3291,
+  cache_read_input_tokens: 3608,
+};

refact-agent/gui/src/components/UsageCounter/UsageCounter.module.css

@@ -0,0 +1,12 @@
+.usageCounterContainer {
+  /* position: absolute; */
+  /* top: 0;
+  right: 0; */
+  margin-left: auto;
+  display: flex;
+  align-items: center;
+  padding: var(--space-2) var(--space-3);
+  gap: 8px;
+  max-width: max-content;
+  opacity: 0.7;
+}

refact-agent/gui/src/components/UsageCounter/UsageCounter.stories.tsx

@@ -0,0 +1,57 @@
+import React from "react";
+import type { Meta, StoryObj } from "@storybook/react";
+import { Provider } from "react-redux";
+
+import { setUpStore } from "../../app/store";
+import { Theme } from "../Theme";
+import { AbortControllerProvider } from "../../contexts/AbortControllers";
+
+import { UsageCounter } from ".";
+import { Usage } from "../../services/refact";
+import {
+  USAGE_COUNTER_STUB_ANTHROPIC,
+  USAGE_COUNTER_STUB_GPT,
+} from "./UsageCounter.fixtures";
+
+const MockedStore: React.FC<{ usage: Usage }> = ({ usage }) => {
+  const store = setUpStore({
+    config: {
+      themeProps: {
+        appearance: "dark",
+      },
+      host: "web",
+      lspPort: 8001,
+    },
+  });
+
+  return (
+    <Provider store={store}>
+      <AbortControllerProvider>
+        <Theme accentColor="gray">
+          <UsageCounter usage={usage} />
+        </Theme>
+      </AbortControllerProvider>
+    </Provider>
+  );
+};
+
+const meta: Meta<typeof MockedStore> = {
+  title: "UsageCounter",
+  component: MockedStore,
+  args: {
+    usage: USAGE_COUNTER_STUB_GPT,
+  },
+};
+
+export default meta;
+
+export const GPTUsageCounter: StoryObj<typeof UsageCounter> = {
+  args: {
+    usage: USAGE_COUNTER_STUB_GPT,
+  },
+};
+export const AnthropicUsageCounter: StoryObj<typeof UsageCounter> = {
+  args: {
+    usage: USAGE_COUNTER_STUB_ANTHROPIC,
+  },
+};

refact-agent/gui/src/components/UsageCounter/UsageCounter.tsx

@@ -0,0 +1,99 @@
+import React from "react";
+import { Card, Flex, HoverCard, Text } from "@radix-ui/themes";
+import { ArrowDownIcon, ArrowUpIcon } from "@radix-ui/react-icons";
+
+import { ScrollArea } from "../ScrollArea";
+import { calculateUsageInputTokens } from "../../utils/calculateUsageInputTokens";
+import type { Usage } from "../../services/refact";
+
+import styles from "./UsageCounter.module.css";
+
+type UsageCounterProps = {
+  usage: Usage;
+};
+
+function formatNumber(num: number): string {
+  return num >= 1_000_000
+    ? (num / 1_000_000).toFixed(1) + "M"
+    : num >= 1_000
+      ? (num / 1_000).toFixed(2) + "k"
+      : num.toString();
+}
+
+const TokenDisplay: React.FC<{ label: string; value: number }> = ({
+  label,
+  value,
+}) => (
+  <Flex align="center" justify="between" width="100%">
+    <Text size="1" weight="bold">
+      {label}
+    </Text>
+    <Text size="1">{value}</Text>
+  </Flex>
+);
+
+export const UsageCounter: React.FC<UsageCounterProps> = ({ usage }) => {
+  const inputTokens = calculateUsageInputTokens(usage, [
+    "prompt_tokens",
+    "cache_creation_input_tokens",
+    "cache_read_input_tokens",
+  ]);
+  const outputTokens = calculateUsageInputTokens(usage, ["completion_tokens"]);
+
+  return (
+    <HoverCard.Root>
+      <HoverCard.Trigger>
+        <Card className={styles.usageCounterContainer}>
+          <Flex align="center">
+            <ArrowUpIcon width="12" height="12" />
+            <Text size="1">{formatNumber(inputTokens)}</Text>
+          </Flex>
+          <Flex align="center">
+            <ArrowDownIcon width="12" height="12" />
+            <Text size="1">{outputTokens}</Text>
+          </Flex>
+        </Card>
+      </HoverCard.Trigger>
+      <ScrollArea scrollbars="both" asChild>
+        <HoverCard.Content
+          size="1"
+          maxHeight="50vh"
+          maxWidth="90vw"
+          minWidth="300px"
+          avoidCollisions
+          align="end"
+          side="top"
+        >
+          <Flex direction="column" align="start" gap="2">
+            <Text size="2" mb="2">
+              Tokens spent per message:
+            </Text>
+            <TokenDisplay
+              label="Input tokens (in total):"
+              value={inputTokens}
+            />
+            {usage.cache_read_input_tokens !== undefined && (
+              <TokenDisplay
+                label="Cache read input tokens:"
+                value={usage.cache_read_input_tokens}
+              />
+            )}
+            {usage.cache_creation_input_tokens !== undefined && (
+              <TokenDisplay
+                label="Cache creation input tokens:"
+                value={usage.cache_creation_input_tokens}
+              />
+            )}
+            <TokenDisplay label="Completion tokens:" value={outputTokens} />
+            {usage.completion_tokens_details && (
+              <TokenDisplay
+                label="Reasoning tokens:"
+                value={usage.completion_tokens_details.reasoning_tokens}
+              />
+            )}
+          </Flex>
+        </HoverCard.Content>
+      </ScrollArea>
+    </HoverCard.Root>
+  );
+};

refact-agent/gui/src/components/UsageCounter/index.ts

@@ -0,0 +1 @@
+export { UsageCounter } from "./UsageCounter";

refact-agent/gui/src/features/Chat/Thread/actions.ts

@@ -8,6 +8,8 @@ import {
   LspChatMode,
   PayloadWithChatAndMessageId,
   PayloadWithChatAndBoolean,
+  PayloadWithChatAndUsage,
+  PayloadWithChatAndNumber,
 } from "./types";
 import {
   isAssistantDelta,
@@ -58,10 +60,19 @@ export const setLastUserMessageId = createAction<PayloadWithChatAndMessageId>(
   "chatThread/setLastUserMessageId",
 );
 
+export const updateMaximumContextTokens =
+  createAction<PayloadWithChatAndNumber>(
+    "chatThread/updateMaximumContextTokens",
+  );
+
 export const setIsNewChatSuggested = createAction<PayloadWithChatAndBoolean>(
   "chatThread/setIsNewChatSuggested",
 );
 
+export const setThreadUsage = createAction<PayloadWithChatAndUsage>(
+  "chatThread/setThreadUsage",
+);
+
 export const setIsNewChatSuggestionRejected =
   createAction<PayloadWithChatAndBoolean>(
     "chatThread/setIsNewChatSuggestionRejected",