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:
parent
5e4d10b950
commit
d5c02fc61d
2 changed files with 33 additions and 67 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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 && (
|
||||
|
|
|
|||
Loading…
Reference in a new issue