From 7b9f490ec9390c60901234695cb4b3d099dccb77 Mon Sep 17 00:00:00 2001
From: sxjeru <sxjeru@gmail.com>
Date: Mon, 17 Mar 2025 15:58:56 +0800
Subject: [PATCH] Update TokenTag.tsx

---
 .../ChatInput/ActionBar/Token/TokenTag.tsx       | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/features/ChatInput/ActionBar/Token/TokenTag.tsx b/src/features/ChatInput/ActionBar/Token/TokenTag.tsx
index eba65c6dacb0b..bb9b96559ab1f 100644
--- a/src/features/ChatInput/ActionBar/Token/TokenTag.tsx
+++ b/src/features/ChatInput/ActionBar/Token/TokenTag.tsx
@@ -2,7 +2,7 @@ import { TokenTag, Tooltip } from '@lobehub/ui';
 import { Popover } from 'antd';
 import { useTheme } from 'antd-style';
 import numeral from 'numeral';
-import { memo } from 'react';
+import { memo, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { Center, Flexbox } from 'react-layout-kit';
 
@@ -12,7 +12,7 @@ import { useTokenCount } from '@/hooks/useTokenCount';
 import { useAgentStore } from '@/store/agent';
 import { agentChatConfigSelectors, agentSelectors } from '@/store/agent/selectors';
 import { useChatStore } from '@/store/chat';
-import { topicSelectors } from '@/store/chat/selectors';
+import { chatSelectors, topicSelectors } from '@/store/chat/selectors';
 import { useToolStore } from '@/store/tool';
 import { toolSelectors } from '@/store/tool/selectors';
 
@@ -41,6 +41,11 @@ const Token = memo<TokenTagProps>(({ total: messageString }) => {
     ];
   });
 
+  const [historyCount, enableHistoryCount] = useAgentStore((s) => [
+    agentChatConfigSelectors.historyCount(s),
+    agentChatConfigSelectors.enableHistoryCount(s),
+  ]);
+
   const maxTokens = useModelContextWindowTokens(model, provider);
 
   // Tool usage token
@@ -60,7 +65,12 @@ const Token = memo<TokenTagProps>(({ total: messageString }) => {
   // Chat usage token
   const inputTokenCount = useTokenCount(input);
 
-  const chatsToken = useTokenCount(messageString) + inputTokenCount;
+  const chatsString = useMemo(() => {
+    const chats = chatSelectors.mainAIChatsWithHistoryConfig(useChatStore.getState());
+    return chats.map(chat => chat.content).join('');
+  }, [messageString, historyCount, enableHistoryCount]);
+
+  const chatsToken = useTokenCount(chatsString) + inputTokenCount;
 
   // SystemRole token
   const systemRoleToken = useTokenCount(systemRole);