From c4a8a54bd76992fa29c24336ff6b23ba7010c601 Mon Sep 17 00:00:00 2001 From: Shantur Rathore Date: Sat, 25 Oct 2025 18:41:15 +0100 Subject: [PATCH] Add context window percentage and subscription model detection to session info --- src/components/message-stream.tsx | 109 ++++++++++++++++++++++++------ src/stores/sessions.ts | 2 + src/types/session.ts | 7 ++ 3 files changed, 97 insertions(+), 21 deletions(-) diff --git a/src/components/message-stream.tsx b/src/components/message-stream.tsx index 1910a567..e2c81e0a 100644 --- a/src/components/message-stream.tsx +++ b/src/components/message-stream.tsx @@ -5,28 +5,74 @@ import ToolCall from "./tool-call" import { sseManager } from "../lib/sse-manager" import Kbd from "./kbd" import { preferences } from "../stores/preferences" +import { providers } from "../stores/sessions" -// Calculate session tokens and cost from messagesInfo -function calculateSessionInfo(messagesInfo?: Map) { - if (!messagesInfo) return { tokens: 0, cost: 0 } +// 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 totalTokens = 0 - let totalCost = 0 + let tokens = 0 + let cost = 0 + let contextWindow = 0 + let isSubscriptionModel = false + let modelID = "" + let providerID = "" - for (const [, info] of messagesInfo) { + // 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.role === "assistant" && info.tokens) { - const tokens = info.tokens - totalTokens += - (tokens.input || 0) + - (tokens.cache?.read || 0) + - (tokens.cache?.write || 0) + - (tokens.output || 0) + - (tokens.reasoning || 0) - totalCost += info.cost || 0 + const usage = info.tokens + + if (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 + } } } - return { tokens: totalTokens, cost: totalCost } + // 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") @@ -39,11 +85,24 @@ function formatTokens(tokens: number): string { return tokens.toString() } -// Format session info like TUI (e.g., "110K • $0.42") -function formatSessionInfo(tokens: number, cost: number): string { +// Format session info like TUI (e.g., "110K/73% ($0.42)" or "110K/73%") +function formatSessionInfo(tokens: number, cost: number, contextWindow: number, isSubscriptionModel: boolean): string { const tokensStr = formatTokens(tokens) - const costStr = cost > 0 ? ` • $${cost.toFixed(2)}` : "" - return `${tokensStr}${costStr}` + + // Calculate percentage if we have context window + if (contextWindow > 0) { + const percentage = Math.round((tokens / contextWindow) * 100) + if (isSubscriptionModel) { + return `${tokensStr}/${percentage}%` + } + return `${tokensStr}/${percentage}% ($${cost.toFixed(2)})` + } + + // Fallback without context window + if (isSubscriptionModel) { + return tokensStr + } + return `${tokensStr} ($${cost.toFixed(2)})` } interface MessageStreamProps { @@ -155,8 +214,16 @@ export default function MessageStream(props: MessageStreamProps) {
{(() => { - const sessionInfo = calculateSessionInfo(props.messagesInfo) - return formatSessionInfo(sessionInfo.tokens, sessionInfo.cost) + const sessionInfo = calculateSessionInfo(props.messagesInfo, props.instanceId) + console.log("[MessageStream] sessionInfo:", sessionInfo) + const result = formatSessionInfo( + sessionInfo.tokens, + sessionInfo.cost, + sessionInfo.contextWindow, + sessionInfo.isSubscriptionModel, + ) + console.log("[MessageStream] formatted result:", result) + return result })()}
diff --git a/src/stores/sessions.ts b/src/stores/sessions.ts index dcf45607..12aef537 100644 --- a/src/stores/sessions.ts +++ b/src/stores/sessions.ts @@ -290,6 +290,8 @@ async function fetchProviders(instanceId: string): Promise { id, name: model.name, providerId: provider.id, + limit: model.limit, + cost: model.cost, })), })) diff --git a/src/types/session.ts b/src/types/session.ts index 553dc0cd..76fd83f5 100644 --- a/src/types/session.ts +++ b/src/types/session.ts @@ -45,4 +45,11 @@ export interface Model { id: string name: string providerId: string + limit?: { + context?: number + } + cost?: { + input?: number + output?: number + } }