From f42d0a285dfbac04f0a6f4f16a3c1a1e0df07e4d Mon Sep 17 00:00:00 2001 From: Shantur Rathore Date: Thu, 13 Nov 2025 11:17:26 +0000 Subject: [PATCH] Propagate session context usage percent --- src/App.tsx | 13 +--- src/components/message-stream.tsx | 119 ++++-------------------------- src/stores/sessions.ts | 58 +++++++++------ 3 files changed, 56 insertions(+), 134 deletions(-) diff --git a/src/App.tsx b/src/App.tsx index add8766c..af0f7f72 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -302,19 +302,14 @@ const ContextUsagePanel: Component<{ instanceId: string; sessionId: string }> = contextWindow: 0, isSubscriptionModel: false, contextUsageTokens: 0, + contextUsagePercent: null, }, ) const tokens = createMemo(() => info().tokens) const contextUsageTokens = createMemo(() => info().contextUsageTokens ?? 0) const contextWindow = createMemo(() => info().contextWindow) - const percentage = createMemo(() => { - const windowSize = contextWindow() - if (!windowSize || windowSize <= 0) return null - const usage = contextUsageTokens() - const percent = Math.round((usage / windowSize) * 100) - return Math.min(100, Math.max(0, percent)) - }) + const contextUsagePercent = createMemo(() => info().contextUsagePercent) const costLabel = createMemo(() => { if (info().isSubscriptionModel || info().cost <= 0) return "Included in plan" @@ -333,7 +328,7 @@ const ContextUsagePanel: Component<{ instanceId: string; sessionId: string }> =
Context window usage
-
{percentage() !== null ? `${percentage()}%` : "--"}
+
{contextUsagePercent() !== null ? `${contextUsagePercent()}%` : "--"}
{contextWindow() @@ -344,7 +339,7 @@ const ContextUsagePanel: Component<{ instanceId: string; sessionId: string }> =
diff --git a/src/components/message-stream.tsx b/src/components/message-stream.tsx index a62c9f5c..265b0239 100644 --- a/src/components/message-stream.tsx +++ b/src/components/message-stream.tsx @@ -31,14 +31,7 @@ import ToolCall from "./tool-call" import { sseManager } from "../lib/sse-manager" import Kbd from "./kbd" import { preferences } from "../stores/preferences" -import { - providers, - getSessionInfo, - computeDisplayParts, - sessions, - setActiveSession, - setActiveParentSession, -} from "../stores/sessions" +import { getSessionInfo, computeDisplayParts, sessions, setActiveSession, setActiveParentSession } from "../stores/sessions" import { setActiveInstanceId } from "../stores/instances" const SCROLL_OFFSET = 64 @@ -76,75 +69,6 @@ function navigateToTaskSession(location: TaskSessionLocation) { } } -// Calculate session tokens and cost from messagesInfo (matches TUI logic) -function calculateSessionInfo(messagesInfo?: Map, instanceId?: string) { - if (!messagesInfo || messagesInfo.size === 0) - return { tokens: 0, cost: 0, contextWindow: 0, isSubscriptionModel: false } - - let tokens = 0 - let cost = 0 - let contextWindow = 0 - let modelID = "" - let providerID = "" - let isSubscriptionModel = false - - // Go backwards through messages to find the last relevant assistant message (like TUI) - const messageArray = Array.from(messagesInfo.values()).reverse() - - for (const info of messageArray) { - if (!info) continue - if (info.role === "assistant" && info.tokens) { - const usage = info.tokens - - if (usage.output && usage.output > 0) { - if (info.summary) { - // If summary message, only count output tokens and stop (like TUI) - tokens = usage.output || 0 - cost = info.cost || 0 - } else { - // Regular message - count all token types (like TUI) - tokens = - (usage.input || 0) + - (usage.cache?.read || 0) + - (usage.cache?.write || 0) + - (usage.output || 0) + - (usage.reasoning || 0) - cost = info.cost || 0 - } - - // Get model info for context window and subscription check - modelID = info.modelID || "" - providerID = info.providerID || "" - isSubscriptionModel = cost === 0 - - break - } - } - } - - // Try to get context window from providers - if (instanceId && modelID && providerID) { - const instanceProviders = providers().get(instanceId) || [] - console.log("[calculateSessionInfo] instanceProviders:", instanceProviders) - console.log("[calculateSessionInfo] looking for providerID:", providerID, "modelID:", modelID) - const provider = instanceProviders.find((p) => p.id === providerID) - console.log("[calculateSessionInfo] found provider:", provider) - if (provider) { - const model = provider.models.find((m) => m.id === modelID) - console.log("[calculateSessionInfo] found model:", model) - if (model?.limit?.context) { - contextWindow = model.limit.context - } - // Check if it's a subscription model (cost is 0 for both input and output) - if (model?.cost?.input === 0 && model?.cost?.output === 0) { - isSubscriptionModel = true - } - } - } - - return { tokens, cost, contextWindow, isSubscriptionModel } -} - // Format tokens like TUI (e.g., "110K", "1.2M") function formatTokens(tokens: number): string { if (tokens >= 1000000) { @@ -156,16 +80,15 @@ function formatTokens(tokens: number): string { } // Format session info for the session view header -function formatSessionInfo(tokens: number, _cost: number, contextWindow: number, _isSubscriptionModel: boolean): string { - const tokensStr = formatTokens(tokens) - +function formatSessionInfo(usageTokens: number, contextWindow: number, usagePercent: number | null): string { if (contextWindow > 0) { const windowStr = formatTokens(contextWindow) - const percentage = Math.min(100, Math.max(0, Math.round((tokens / contextWindow) * 100))) - return `${tokensStr} of ${windowStr} (${percentage}%)` + const usageStr = formatTokens(usageTokens) + const percent = usagePercent ?? Math.min(100, Math.max(0, Math.round((usageTokens / contextWindow) * 100))) + return `${usageStr} of ${windowStr} (${percent}%)` } - return tokensStr + return formatTokens(usageTokens) } interface MessageStreamProps { @@ -274,30 +197,20 @@ export default function MessageStream(props: MessageStreamProps) { return `${toolPart.id}:${version}:${status}` } - const sessionInfo = createMemo(() => { - return ( - getSessionInfo(props.instanceId, props.sessionId) || { - tokens: 0, - cost: 0, - contextWindow: 0, - isSubscriptionModel: false, - } - ) - }) - - const formattedSessionInfo = createMemo(() => { - const sessionInfo = getSessionInfo(props.instanceId, props.sessionId) || { + const sessionInfo = createMemo(() => + getSessionInfo(props.instanceId, props.sessionId) ?? { tokens: 0, cost: 0, contextWindow: 0, isSubscriptionModel: false, - } - return formatSessionInfo( - sessionInfo.tokens, - sessionInfo.cost, - sessionInfo.contextWindow, - sessionInfo.isSubscriptionModel, - ) + contextUsageTokens: 0, + contextUsagePercent: null, + }, + ) + + const formattedSessionInfo = createMemo(() => { + const info = sessionInfo() + return formatSessionInfo(info.contextUsageTokens, info.contextWindow, info.contextUsagePercent) }) function isNearBottom(element: HTMLDivElement, offset = SCROLL_OFFSET) { diff --git a/src/stores/sessions.ts b/src/stores/sessions.ts index b7417b5c..abad0b65 100644 --- a/src/stores/sessions.ts +++ b/src/stores/sessions.ts @@ -34,6 +34,7 @@ interface SessionInfo { contextWindow: number isSubscriptionModel: boolean contextUsageTokens: number + contextUsagePercent: number | null } interface SessionForkResponse { @@ -552,8 +553,9 @@ function updateSessionInfo(instanceId: string, sessionId: string) { let isSubscriptionModel = false let modelID = "" let providerID = "" - let inputTokensForUsage = 0 - let cacheReadTokensForUsage = 0 + let actualUsageTokens = 0 + let contextUsagePercent: number | null = null + let hasContextUsage = false // Calculate from last assistant message in this session only if (session.messagesInfo.size > 0) { @@ -565,23 +567,23 @@ function updateSessionInfo(instanceId: string, sessionId: string) { const usage = info.tokens if (usage.output > 0) { + const inputTokens = usage.input || 0 + const reasoningTokens = usage.reasoning || 0 + const cacheReadTokens = usage.cache?.read || 0 + const cacheWriteTokens = usage.cache?.write || 0 + const outputTokens = usage.output || 0 + if (info.summary) { // If summary message, only count output tokens and stop (like TUI) - tokens = usage.output || 0 - cost = info.cost || 0 + tokens = outputTokens } else { // Regular message - count all token types (like TUI) - tokens = - (usage.input || 0) + - (usage.cache?.read || 0) + - (usage.cache?.write || 0) + - (usage.output || 0) + - (usage.reasoning || 0) - cost = info.cost || 0 + tokens = inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens + reasoningTokens } - inputTokensForUsage = usage.input || 0 - cacheReadTokensForUsage = usage.cache?.read || 0 + cost = info.cost || 0 + actualUsageTokens = tokens + hasContextUsage = inputTokens + cacheReadTokens + cacheWriteTokens > 0 // Get model info identifiers for context lookups modelID = info.modelID || "" @@ -625,7 +627,20 @@ function updateSessionInfo(instanceId: string, sessionId: string) { } } - const contextUsageTokens = inputTokensForUsage + cacheReadTokensForUsage + modelOutputLimit + const outputBudget = Math.min(modelOutputLimit, DEFAULT_MODEL_OUTPUT_LIMIT) + let contextUsageTokens = 0 + + if (hasContextUsage && actualUsageTokens > 0) { + contextUsageTokens = actualUsageTokens + outputBudget + if (contextWindow > 0) { + const percent = Math.round((contextUsageTokens / contextWindow) * 100) + contextUsagePercent = Math.min(100, Math.max(0, percent)) + } else { + contextUsagePercent = null + } + } else { + contextUsagePercent = contextWindow > 0 ? 0 : null + } setSessionInfoByInstance((prev) => { const next = new Map(prev) @@ -636,6 +651,7 @@ function updateSessionInfo(instanceId: string, sessionId: string) { contextWindow, isSubscriptionModel, contextUsageTokens, + contextUsagePercent, }) next.set(instanceId, instanceInfo) return next @@ -708,11 +724,8 @@ async function createSession(instanceId: string, agent?: string): Promise p.id === session.model.providerId) const initialModel = initialProvider?.models.find((m) => m.id === session.model.modelId) const initialContextWindow = initialModel?.limit?.context ?? 0 - const initialOutputLimit = - initialModel?.limit?.output && initialModel.limit.output > 0 - ? initialModel.limit.output - : DEFAULT_MODEL_OUTPUT_LIMIT const initialSubscriptionModel = initialModel?.cost?.input === 0 && initialModel?.cost?.output === 0 + const initialContextPercent = initialContextWindow > 0 ? 0 : null setSessionInfoByInstance((prev) => { const next = new Map(prev) @@ -722,7 +735,8 @@ async function createSession(instanceId: string, agent?: string): Promise p.id === forkedSession.model.providerId) const forkModel = forkProvider?.models.find((m) => m.id === forkedSession.model.modelId) const forkContextWindow = forkModel?.limit?.context ?? 0 - const forkOutputLimit = - forkModel?.limit?.output && forkModel.limit.output > 0 ? forkModel.limit.output : DEFAULT_MODEL_OUTPUT_LIMIT const forkSubscriptionModel = forkModel?.cost?.input === 0 && forkModel?.cost?.output === 0 + const forkContextPercent = forkContextWindow > 0 ? 0 : null setSessionInfoByInstance((prev) => { const next = new Map(prev) @@ -822,7 +835,8 @@ async function forkSession( cost: 0, contextWindow: forkContextWindow, isSubscriptionModel: Boolean(forkSubscriptionModel), - contextUsageTokens: forkOutputLimit, + contextUsageTokens: 0, + contextUsagePercent: forkContextPercent, }) next.set(instanceId, instanceInfo) return next