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

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";
@@ -40,6 +41,7 @@ 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";
@@ -107,6 +109,13 @@ startListening({
       const maxFreeAgentUsage = getMaxFreeAgentUsage(payload);
       dispatch(updateMaxAgentUsageAmount(maxFreeAgentUsage));
     }
+
+    if (isChatResponseChoice(payload)) {
+      const { usage } = payload;
+      if (usage) {
+        dispatch(setThreadUsage({ chatId: payload.id, usage }));
+      }
+    }
   },
 });
 

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

@@ -2,24 +2,21 @@ import React, { useCallback } from "react";
 import { Markdown } from "../Markdown";
 
 import { Container, Box } from "@radix-ui/themes";
-import { ToolCall, Usage } from "../../services/refact";
+import { ToolCall } from "../../services/refact";
 import { ToolContent } from "./ToolsContent";
 import { fallbackCopying } from "../../utils/fallbackCopying";
 import { telemetryApi } from "../../services/refact/telemetry";
 import { LikeButton } from "./LikeButton";
-import { UsageCounter } from "./UsageCounter";
 
 type ChatInputProps = {
   message: string | null;
   toolCalls?: ToolCall[] | null;
-  usage?: Usage | null;
   isLast?: boolean;
 };
 
 export const AssistantInput: React.FC<ChatInputProps> = ({
   message,
   toolCalls,
-  usage,
   isLast,
 }) => {
   const [sendTelemetryEvent] =
@@ -70,7 +67,6 @@ export const AssistantInput: React.FC<ChatInputProps> = ({
         </Box>
       )}
       {toolCalls && <ToolContent toolCalls={toolCalls} />}
-      {usage && <UsageCounter usage={usage} />}
       {isLast && <LikeButton />}
     </Container>
   );

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} />
@@ -189,7 +193,6 @@ function renderMessages(
         key={key}
         message={head.content}
         toolCalls={head.tool_calls}
-        usage={head.usage}
         isLast={isLast}
       />,
     ];

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

@@ -4,6 +4,7 @@ import { ArrowDownIcon, ArrowUpIcon } from "@radix-ui/react-icons";
 import { ScrollArea } from "../../ScrollArea";
 import { Usage } from "../../../services/refact";
 import styles from "./UsageCounter.module.css";
+import { calculateInputTokens } from "./UsageCounter.utils";
 
 type UsageCounterProps = {
   usage: Usage;
@@ -17,19 +18,25 @@ function formatNumber(num: number): string {
       : num.toString();
 }
 
-const calculateTokens = (usage: Usage, keys: (keyof Usage)[]): number =>
-  keys.reduce((acc, key) => {
-    const value = usage[key];
-    return acc + (typeof value === "number" ? value : 0);
-  }, 0);
+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 = calculateTokens(usage, [
+  const inputTokens = calculateInputTokens(usage, [
     "prompt_tokens",
     "cache_creation_input_tokens",
     "cache_read_input_tokens",
   ]);
-  const outputTokens = calculateTokens(usage, ["completion_tokens"]);
+  const outputTokens = calculateInputTokens(usage, ["completion_tokens"]);
 
   return (
     <HoverCard.Root>
@@ -59,43 +66,28 @@ export const UsageCounter: React.FC<UsageCounterProps> = ({ usage }) => {
             <Text size="2" mb="2">
               Tokens spent per message:
             </Text>
-            <Flex align="center" justify="between" width="100%">
-              <Text size="1" weight="bold">
-                Input tokens (in total):{" "}
-              </Text>
-              <Text size="1">{inputTokens}</Text>
-            </Flex>
+            <TokenDisplay
+              label="Input tokens (in total):"
+              value={inputTokens}
+            />
             {usage.cache_read_input_tokens && (
-              <Flex align="center" justify="between" width="100%">
-                <Text size="1" weight="bold">
-                  Cache read input tokens:{" "}
-                </Text>
-                <Text size="1">{usage.cache_read_input_tokens}</Text>
-              </Flex>
+              <TokenDisplay
+                label="Cache read input tokens:"
+                value={usage.cache_read_input_tokens}
+              />
             )}
             {usage.cache_creation_input_tokens && (
-              <Flex align="center" justify="between" width="100%">
-                <Text size="1" weight="bold">
-                  Cache creation input tokens:{" "}
-                </Text>
-                <Text size="1">{usage.cache_creation_input_tokens}</Text>
-              </Flex>
+              <TokenDisplay
+                label="Cache creation input tokens:"
+                value={usage.cache_creation_input_tokens}
+              />
             )}
-            <Flex align="center" justify="between" width="100%">
-              <Text size="1" weight="bold">
-                Completion tokens:{" "}
-              </Text>
-              <Text size="1">{outputTokens}</Text>
-            </Flex>
+            <TokenDisplay label="Completion tokens:" value={outputTokens} />
             {usage.completion_tokens_details && (
-              <Flex align="center" justify="between" width="100%">
-                <Text size="1" weight="bold">
-                  Reasoning tokens:{" "}
-                </Text>
-                <Text size="1">
-                  {usage.completion_tokens_details.reasoning_tokens}
-                </Text>
-              </Flex>
+              <TokenDisplay
+                label="Reasoning tokens:"
+                value={usage.completion_tokens_details.reasoning_tokens}
+              />
             )}
           </Flex>
         </HoverCard.Content>

refact-agent/gui/src/components/ChatContent/UsageCounter/UsageCounter.utils.ts

@@ -0,0 +1,10 @@
+import { Usage } from "../../../services/refact";
+
+export const calculateInputTokens = (
+  usage: Usage,
+  keys: (keyof Usage)[],
+): number =>
+  keys.reduce((acc, key) => {
+    const value = usage[key];
+    return acc + (typeof value === "number" ? value : 0);
+  }, 0);

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/features/Chat/Thread/actions.ts

@@ -8,6 +8,7 @@ import {
   LspChatMode,
   PayloadWithChatAndMessageId,
   PayloadWithChatAndBoolean,
+  PayloadWithChatAndUsage,
 } from "./types";
 import {
   isAssistantDelta,
@@ -62,6 +63,10 @@ export const setIsNewChatSuggested = createAction<PayloadWithChatAndBoolean>(
   "chatThread/setIsNewChatSuggested",
 );
 
+export const setThreadUsage = createAction<PayloadWithChatAndUsage>(
+  "chatThread/setThreadUsage",
+);
+
 export const setIsNewChatSuggestionRejected =
   createAction<PayloadWithChatAndBoolean>(
     "chatThread/setIsNewChatSuggestionRejected",

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

@@ -35,13 +35,17 @@ import {
   fixBrokenToolMessages,
   setIsNewChatSuggested,
   setIsNewChatSuggestionRejected,
+  setThreadUsage,
 } from "./actions";
 import { formatChatResponse } from "./utils";
 import {
   DEFAULT_MAX_NEW_TOKENS,
   isToolCallMessage,
   validateToolCall,
 } from "../../../services/refact";
+import { calculateInputTokens } from "../../../components/ChatContent/UsageCounter/UsageCounter.utils";
+
+const RECOMMENDED_MAXIMUM_PROMPT_TOKENS_AMOUNT = 30000;
 
 const createChatThread = (
   tool_use: ToolUse,
@@ -218,6 +222,31 @@ export const chatReducer = createReducer(initialState, (builder) => {
     };
   });
 
+  builder.addCase(setThreadUsage, (state, action) => {
+    if (state.thread.id !== action.payload.chatId) return state;
+
+    const { usage } = action.payload;
+    state.thread.usage = usage;
+
+    const inputTokensAmount = calculateInputTokens(usage, [
+      "prompt_tokens",
+      "cache_creation_input_tokens",
+      "cache_read_input_tokens",
+    ]);
+
+    if (inputTokensAmount >= RECOMMENDED_MAXIMUM_PROMPT_TOKENS_AMOUNT) {
+      const { wasSuggested, wasRejectedByUser } =
+        state.thread.new_chat_suggested;
+
+      state.thread.new_chat_suggested = {
+        wasSuggested: wasSuggested || !wasSuggested,
+        wasRejectedByUser: wasRejectedByUser
+          ? !wasRejectedByUser
+          : wasRejectedByUser,
+      };
+    }
+  });
+
   builder.addCase(setEnabledCheckpoints, (state, action) => {
     state.checkpoints_enabled = action.payload;
   });

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

@@ -18,6 +18,7 @@ export const selectCheckpointsEnabled = (state: RootState) =>
 
 export const selectThreadNewChatSuggested = (state: RootState) =>
   state.chat.thread.new_chat_suggested;
+export const selectThreadUsage = (state: RootState) => state.chat.thread.usage;
 export const selectIsWaiting = (state: RootState) =>
   state.chat.waiting_for_response;
 export const selectIsStreaming = (state: RootState) => state.chat.streaming;

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

@@ -1,3 +1,4 @@
+import { Usage } from "../../../services/refact";
 import { SystemPrompts } from "../../../services/refact/prompts";
 import { ChatMessages } from "../../../services/refact/types";
 import { parseOrElse } from "../../../utils/parseOrElse";
@@ -23,6 +24,7 @@ export type ChatThread = {
   project_name?: string;
   last_user_message_id?: string;
   new_chat_suggested: SuggestedChat;
+  usage?: Usage;
 };
 
 export type SuggestedChat = {
@@ -50,6 +52,7 @@ export type Chat = {
 export type PayloadWithId = { id: string };
 export type PayloadWithChatAndMessageId = { chatId: string; messageId: string };
 export type PayloadWithChatAndBoolean = { chatId: string; value: boolean };
+export type PayloadWithChatAndUsage = { chatId: string; usage: Usage };
 export type PayloadWithIdAndTitle = {
   title: string;
   isTitleGenerated: boolean;