refactor: remove token estimation logic from AIChatGroup and update TokenUsageDisplay for clarity

- Eliminated the calculation of thinking and text output tokens from AIChatGroup, simplifying the component.
- Updated TokenUsageDisplay to reflect changes, removing unnecessary props and adjusting context calculations.
- Ensured that context percentages are now based solely on total input tokens for improved accuracy.
This commit is contained in:
iliya 2026-03-10 12:47:32 +02:00
parent 5e4d10b950
commit d5c02fc61d
2 changed files with 33 additions and 67 deletions

View file

@ -6,7 +6,6 @@ import { useStore } from '@renderer/store';
import { enhanceAIGroup, type PrecedingSlashInfo } from '@renderer/utils/aiGroupEnhancer';
import { extractSlashInfo, isCommandContent } from '@shared/utils/contentSanitizer';
import { getModelColorClass } from '@shared/utils/modelParser';
import { estimateTokens } from '@shared/utils/tokenFormatting';
import { format } from 'date-fns';
import { Bot, ChevronDown, Clock } from 'lucide-react';
import { useShallow } from 'zustand/react/shallow';
@ -248,28 +247,6 @@ const AIChatGroupInner = ({
// Get the total cost
const costUSD = aiGroup.metrics.costUsd;
// Calculate thinking and text output tokens from assistant message content blocks
// These are estimated from the actual content, providing breakdown of output token usage
const { thinkingTokens, textOutputTokens } = useMemo(() => {
let thinking = 0;
let textOutput = 0;
const responses = aiGroup.responses || [];
for (const msg of responses) {
if (msg.type === 'assistant' && Array.isArray(msg.content)) {
for (const block of msg.content) {
if (block.type === 'thinking' && block.thinking) {
thinking += estimateTokens(block.thinking);
} else if (block.type === 'text' && block.text) {
textOutput += estimateTokens(block.text);
}
}
}
}
return { thinkingTokens: thinking, textOutputTokens: textOutput };
}, [aiGroup.responses]);
// Auto-expand if contains error or search result, or if manually expanded
const isExpanded =
isAIGroupExpandedForTab(aiGroup.id) || containsHighlightedError || shouldExpandForSearch;
@ -470,8 +447,6 @@ const AIChatGroupInner = ({
outputTokens={lastUsage.output_tokens}
cacheReadTokens={lastUsage.cache_read_input_tokens ?? 0}
cacheCreationTokens={lastUsage.cache_creation_input_tokens ?? 0}
thinkingTokens={thinkingTokens}
textOutputTokens={textOutputTokens}
modelName={enhanced.mainModel?.name}
modelFamily={enhanced.mainModel?.family}
size="sm"

View file

@ -32,10 +32,6 @@ interface TokenUsageDisplayProps {
cacheReadTokens: number;
/** Cache creation/write tokens count */
cacheCreationTokens: number;
/** Thinking tokens (extended thinking content) - estimated from content */
thinkingTokens?: number;
/** Text output tokens (Claude's text responses) - estimated from content */
textOutputTokens?: number;
/** Optional model name for display */
modelName?: string;
/** Optional model family for color styling */
@ -60,24 +56,22 @@ interface TokenUsageDisplayProps {
*/
const SessionContextSection = ({
contextStats,
totalTokens,
thinkingTokens = 0,
textOutputTokens = 0,
totalInputTokens,
}: Readonly<{
contextStats: ContextStats;
totalTokens: number;
thinkingTokens?: number;
textOutputTokens?: number;
totalInputTokens: number;
}>): React.JSX.Element => {
const [expanded, setExpanded] = useState(false);
const { tokensByCategory } = contextStats;
// Calculate combined thinking+text tokens and include in context total
const thinkingTextTokens = thinkingTokens + textOutputTokens;
const adjustedContextTotal = contextStats.totalEstimatedTokens + thinkingTextTokens;
// contextStats.totalEstimatedTokens already includes all categories (CLAUDE.md, @files,
// tool outputs, thinking+text, task coordination, user messages) — no manual adjustment needed.
// Denominator is total input tokens only (not output), since visible context is part of input.
const contextPercent =
totalTokens > 0 ? Math.min((adjustedContextTotal / totalTokens) * 100, 100).toFixed(1) : '0.0';
totalInputTokens > 0
? Math.min((contextStats.totalEstimatedTokens / totalInputTokens) * 100, 100).toFixed(1)
: '0.0';
// Count accumulated injections by category
const claudeMdCount = contextStats.accumulatedInjections.filter(
@ -96,28 +90,30 @@ const SessionContextSection = ({
(inj) => inj.category === 'user-message'
).length;
// Calculate percentages for each category
// Calculate percentages for each category (relative to total input tokens)
const claudeMdPercent =
totalTokens > 0
? Math.min((tokensByCategory.claudeMd / totalTokens) * 100, 100).toFixed(1)
totalInputTokens > 0
? Math.min((tokensByCategory.claudeMd / totalInputTokens) * 100, 100).toFixed(1)
: '0.0';
const mentionedFilesPercent =
totalTokens > 0
? Math.min((tokensByCategory.mentionedFiles / totalTokens) * 100, 100).toFixed(1)
totalInputTokens > 0
? Math.min((tokensByCategory.mentionedFiles / totalInputTokens) * 100, 100).toFixed(1)
: '0.0';
const toolOutputsPercent =
totalTokens > 0
? Math.min((tokensByCategory.toolOutputs / totalTokens) * 100, 100).toFixed(1)
totalInputTokens > 0
? Math.min((tokensByCategory.toolOutputs / totalInputTokens) * 100, 100).toFixed(1)
: '0.0';
const thinkingTextPercent =
totalTokens > 0 ? Math.min((thinkingTextTokens / totalTokens) * 100, 100).toFixed(1) : '0.0';
totalInputTokens > 0
? Math.min((tokensByCategory.thinkingText / totalInputTokens) * 100, 100).toFixed(1)
: '0.0';
const taskCoordinationPercent =
totalTokens > 0
? Math.min((tokensByCategory.taskCoordination / totalTokens) * 100, 100).toFixed(1)
totalInputTokens > 0
? Math.min((tokensByCategory.taskCoordination / totalInputTokens) * 100, 100).toFixed(1)
: '0.0';
const userMessagesPercent =
totalTokens > 0
? Math.min((tokensByCategory.userMessages / totalTokens) * 100, 100).toFixed(1)
totalInputTokens > 0
? Math.min((tokensByCategory.userMessages / totalInputTokens) * 100, 100).toFixed(1)
: '0.0';
return (
@ -148,7 +144,7 @@ const SessionContextSection = ({
className="whitespace-nowrap text-[10px] tabular-nums"
style={{ color: COLOR_TEXT_MUTED }}
>
{formatTokens(adjustedContextTotal)} ({contextPercent}%)
{formatTokens(contextStats.totalEstimatedTokens)} ({contextPercent}%)
</span>
</div>
@ -221,11 +217,11 @@ const SessionContextSection = ({
)}
{/* Thinking + Text */}
{thinkingTextTokens > 0 && (
{tokensByCategory.thinkingText > 0 && (
<div className="flex items-center justify-between text-[10px]">
<span style={{ color: COLOR_TEXT_MUTED }}>Thinking + Text</span>
<span className="tabular-nums" style={{ color: COLOR_TEXT_SECONDARY }}>
{formatTokens(thinkingTextTokens)}{' '}
{formatTokens(tokensByCategory.thinkingText)}{' '}
<span className="opacity-60">({thinkingTextPercent}%)</span>
</span>
</div>
@ -249,8 +245,6 @@ export const TokenUsageDisplay = ({
outputTokens,
cacheReadTokens,
cacheCreationTokens,
thinkingTokens = 0,
textOutputTokens = 0,
modelName,
modelFamily,
size = 'sm',
@ -261,6 +255,8 @@ export const TokenUsageDisplay = ({
costUsd,
}: Readonly<TokenUsageDisplayProps>): React.JSX.Element => {
const totalTokens = inputTokens + cacheReadTokens + cacheCreationTokens + outputTokens;
// Total input tokens only (without output) — used as denominator for visible context %
const totalInputTokens = inputTokens + cacheReadTokens + cacheCreationTokens;
const formattedTotal = formatTokens(totalTokens);
// Size-based classes
@ -531,17 +527,12 @@ export const TokenUsageDisplay = ({
)}
{/* Visible Context Breakdown - expandable section */}
{contextStats &&
(contextStats.totalEstimatedTokens > 0 ||
thinkingTokens > 0 ||
textOutputTokens > 0) && (
<SessionContextSection
contextStats={contextStats}
totalTokens={totalTokens}
thinkingTokens={thinkingTokens}
textOutputTokens={textOutputTokens}
/>
)}
{contextStats && contextStats.totalEstimatedTokens > 0 && (
<SessionContextSection
contextStats={contextStats}
totalInputTokens={totalInputTokens}
/>
)}
{/* CLAUDE.md Breakdown - fallback when contextStats not provided (deprecated) */}
{!contextStats && claudeMdStats && (