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 { enhanceAIGroup, type PrecedingSlashInfo } from '@renderer/utils/aiGroupEnhancer';
import { extractSlashInfo, isCommandContent } from '@shared/utils/contentSanitizer'; import { extractSlashInfo, isCommandContent } from '@shared/utils/contentSanitizer';
import { getModelColorClass } from '@shared/utils/modelParser'; import { getModelColorClass } from '@shared/utils/modelParser';
import { estimateTokens } from '@shared/utils/tokenFormatting';
import { format } from 'date-fns'; import { format } from 'date-fns';
import { Bot, ChevronDown, Clock } from 'lucide-react'; import { Bot, ChevronDown, Clock } from 'lucide-react';
import { useShallow } from 'zustand/react/shallow'; import { useShallow } from 'zustand/react/shallow';
@ -248,28 +247,6 @@ const AIChatGroupInner = ({
// Get the total cost // Get the total cost
const costUSD = aiGroup.metrics.costUsd; 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 // Auto-expand if contains error or search result, or if manually expanded
const isExpanded = const isExpanded =
isAIGroupExpandedForTab(aiGroup.id) || containsHighlightedError || shouldExpandForSearch; isAIGroupExpandedForTab(aiGroup.id) || containsHighlightedError || shouldExpandForSearch;
@ -470,8 +447,6 @@ const AIChatGroupInner = ({
outputTokens={lastUsage.output_tokens} outputTokens={lastUsage.output_tokens}
cacheReadTokens={lastUsage.cache_read_input_tokens ?? 0} cacheReadTokens={lastUsage.cache_read_input_tokens ?? 0}
cacheCreationTokens={lastUsage.cache_creation_input_tokens ?? 0} cacheCreationTokens={lastUsage.cache_creation_input_tokens ?? 0}
thinkingTokens={thinkingTokens}
textOutputTokens={textOutputTokens}
modelName={enhanced.mainModel?.name} modelName={enhanced.mainModel?.name}
modelFamily={enhanced.mainModel?.family} modelFamily={enhanced.mainModel?.family}
size="sm" size="sm"

View file

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