Refine session usage tracking

2025-11-25 12:03:33 +00:00
parent 48eb6b8982
commit bf32fcf136
9 changed files with 434 additions and 120 deletions
--- a/packages/ui/src/components/message-item.tsx
+++ b/packages/ui/src/components/message-item.tsx
@@ -1,6 +1,7 @@
-import { For, Show } from "solid-js"
+import { For, Show, createMemo } from "solid-js"
 import type { Message, SDKPart, MessageInfo, ClientPart } from "../types/message"
 import { partHasRenderableText } from "../types/message"
+import { formatTokenTotal } from "../lib/formatters"
 import MessagePart from "./message-part"

 interface MessageItemProps {
@@ -138,6 +139,41 @@ export default function MessageItem(props: MessageItemProps) {
    isUser()
      ? "message-item-base bg-[var(--message-user-bg)] border-l-4 border-[var(--message-user-border)]"
      : "message-item-base assistant-message bg-[var(--message-assistant-bg)] border-l-4 border-[var(--message-assistant-border)]"
+
+  const statChipClass =
+    "inline-flex items-center gap-1 rounded-full border border-[var(--border-base)] px-2 py-0.5 text-[10px]"
+  const statLabelClass = "uppercase text-[9px] tracking-wide text-[var(--text-muted)]"
+  const statValueClass = "font-semibold text-[var(--text-primary)]"
+
+  const usageStats = createMemo(() => {
+    const info = props.messageInfo
+    if (!info || info.role !== "assistant" || !info.tokens) {
+      return null
+    }
+
+    const tokens = info.tokens
+    const input = tokens.input ?? 0
+    const output = tokens.output ?? 0
+    const reasoning = tokens.reasoning ?? 0
+    if (input === 0 && output === 0 && reasoning === 0) {
+      return null
+    }
+
+    return {
+      input,
+      output,
+      reasoning,
+      cacheRead: tokens.cache?.read ?? 0,
+      cacheWrite: tokens.cache?.write ?? 0,
+      cost: info.cost ?? 0,
+    }
+  })
+
+  const formatCostValue = (value: number) => {
+    if (!value) return "$0.00"
+    if (value < 0.01) return `$${value.toPrecision(2)}`
+    return `$${value.toFixed(2)}`
+  }
 
  const agentIdentifier = () => {
    if (isUser()) return ""
@@ -193,6 +229,36 @@ export default function MessageItem(props: MessageItemProps) {
              Fork
            </button>
          </Show>
+          <Show when={usageStats()}>
+            {(usage) => (
+              <div class="flex flex-wrap items-center gap-1 text-[10px] text-[var(--text-muted)]">
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Input</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().input)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Output</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().output)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Reasoning</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().reasoning)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Cache Read</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().cacheRead)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Cache Write</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().cacheWrite)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Cost</span>
+                  <span class={statValueClass}>{formatCostValue(usage().cost)}</span>
+                </div>
+              </div>
+            )}
+          </Show>
          <span class="text-[11px] text-[var(--text-muted)]">{timestamp()}</span>
        </div>
      </div>
--- a/packages/ui/src/components/message-stream.tsx
+++ b/packages/ui/src/components/message-stream.tsx
@@ -32,6 +32,7 @@ import { sseManager } from "../lib/sse-manager"
 import Kbd from "./kbd"
 import { useConfig } from "../stores/preferences"
 import { getSessionInfo, computeDisplayParts, sessions, setActiveSession, setActiveParentSession } from "../stores/sessions"
+import { formatTokenTotal } from "../lib/formatters"
 import { setActiveInstanceId } from "../stores/instances"
 import { showCommandPalette } from "../stores/command-palette"

@@ -72,26 +73,9 @@ function navigateToTaskSession(location: TaskSessionLocation) {
  }
 }

-// Format tokens like TUI (e.g., "110K", "1.2M")
+// Format tokens like session sidebar (comma-separated totals)
 function formatTokens(tokens: number): string {
-  if (tokens >= 1000000) {
-    return `${(tokens / 1000000).toFixed(1)}M`
-  } else if (tokens >= 1000) {
-    return `${(tokens / 1000).toFixed(0)}K`
-  }
-  return tokens.toString()
-}
-
-// Format session info for the session view header
-function formatSessionInfo(usageTokens: number, contextWindow: number, usagePercent: number | null): string {
-  if (contextWindow > 0) {
-    const windowStr = formatTokens(contextWindow)
-    const usageStr = formatTokens(usageTokens)
-    const percent = usagePercent ?? Math.min(100, Math.max(0, Math.round((usageTokens / contextWindow) * 100)))
-    return `${usageStr} of ${windowStr} (${percent}%)`
-  }
-
-  return formatTokens(usageTokens)
+  return formatTokenTotal(tokens)
 }

 interface MessageStreamProps {
@@ -206,18 +190,27 @@ export default function MessageStream(props: MessageStreamProps) {

  const sessionInfo = createMemo(() =>
    getSessionInfo(props.instanceId, props.sessionId) ?? {
-      tokens: 0,
      cost: 0,
      contextWindow: 0,
      isSubscriptionModel: false,
-      contextUsageTokens: 0,
-      contextUsagePercent: null,
+      inputTokens: 0,
+      outputTokens: 0,
+      reasoningTokens: 0,
+      actualUsageTokens: 0,
+      modelOutputLimit: 0,
+      contextAvailableTokens: null,
    },
  )

-  const formattedSessionInfo = createMemo(() => {
+  const tokenStats = createMemo(() => {
    const info = sessionInfo()
-    return formatSessionInfo(info.contextUsageTokens, info.contextWindow, info.contextUsagePercent)
+    return {
+      input: info.inputTokens ?? 0,
+      output: info.outputTokens ?? 0,
+      cost: info.cost ?? 0,
+      used: info.actualUsageTokens ?? 0,
+      avail: info.contextAvailableTokens,
+    }
  })

  function isNearBottom(element: HTMLDivElement, offset = SCROLL_OFFSET) {
@@ -552,10 +545,20 @@ export default function MessageStream(props: MessageStreamProps) {

  return (
    <div class="message-stream-container">
-      <div class="connection-status">
-        <div class="connection-status-text connection-status-info flex items-center gap-2 text-sm font-medium">
-          <span>{formattedSessionInfo()}</span>
-        </div>
+        <div class="connection-status">
+          <div class="connection-status-text connection-status-info flex flex-wrap items-center gap-2 text-sm font-medium">
+            <div class="inline-flex items-center gap-1 rounded-full border border-base px-2 py-0.5 text-xs text-primary">
+              <span class="uppercase text-[10px] tracking-wide text-primary/70">Used</span>
+              <span class="font-semibold text-primary">{formatTokens(sessionInfo().actualUsageTokens ?? 0)}</span>
+            </div>
+            <div class="inline-flex items-center gap-1 rounded-full border border-base px-2 py-0.5 text-xs text-primary">
+              <span class="uppercase text-[10px] tracking-wide text-primary/70">Avail</span>
+              <span class="font-semibold text-primary">
+                {sessionInfo().contextAvailableTokens !== null ? formatTokens(sessionInfo().contextAvailableTokens ?? 0) : "--"}
+              </span>
+            </div>
+          </div>
+
        <div class="connection-status-text connection-status-shortcut">
          <div class="connection-status-shortcut-action">
            <button
--- a/packages/ui/src/components/session/context-usage-panel.tsx
+++ b/packages/ui/src/components/session/context-usage-panel.tsx
@@ -7,54 +7,71 @@ interface ContextUsagePanelProps {
  sessionId: string
 }

+const chipClass = "inline-flex items-center gap-1 rounded-full border border-base px-2 py-0.5 text-xs text-primary"
+const chipLabelClass = "uppercase text-[10px] tracking-wide text-primary/70"
+const headingClass = "text-xs font-semibold text-primary/70 uppercase tracking-wide"
+
 const ContextUsagePanel: Component<ContextUsagePanelProps> = (props) => {
  const info = createMemo(
    () =>
      getSessionInfo(props.instanceId, props.sessionId) ?? {
-        tokens: 0,
        cost: 0,
        contextWindow: 0,
        isSubscriptionModel: false,
-        contextUsageTokens: 0,
-        contextUsagePercent: null,
+        inputTokens: 0,
+        outputTokens: 0,
+        reasoningTokens: 0,
+        actualUsageTokens: 0,
+        modelOutputLimit: 0,
+        contextAvailableTokens: null,
      },
  )

-  const tokens = createMemo(() => info().tokens)
-  const contextUsageTokens = createMemo(() => info().contextUsageTokens ?? 0)
-  const contextWindow = createMemo(() => info().contextWindow)
-  const contextUsagePercent = createMemo(() => info().contextUsagePercent)
-
-  const costLabel = createMemo(() => {
-    if (info().isSubscriptionModel || info().cost <= 0) return "Included in plan"
-    return `$${info().cost.toFixed(2)} spent`
+  const inputTokens = createMemo(() => info().inputTokens ?? 0)
+  const outputTokens = createMemo(() => info().outputTokens ?? 0)
+  const actualUsageTokens = createMemo(() => info().actualUsageTokens ?? 0)
+  const availableTokens = createMemo(() => info().contextAvailableTokens)
+  const outputLimit = createMemo(() => info().modelOutputLimit ?? 0)
+  const costValue = createMemo(() => {
+    const value = info().isSubscriptionModel ? 0 : info().cost
+    return value > 0 ? value : 0
  })

+  const formatTokenValue = (value: number | null | undefined) => {
+    if (value === null || value === undefined) return "--"
+    return formatTokenTotal(value)
+  }
+
+  const costDisplay = createMemo(() => `$${costValue().toFixed(2)}`)
+
  return (
-    <div class="session-context-panel border-r border-base border-b px-3 py-3">
-      <div class="flex items-center justify-between gap-4">
-        <div>
-          <div class="text-xs font-semibold text-primary/70 uppercase tracking-wide">Tokens (last call)</div>
-          <div class="text-lg font-semibold text-primary">{formatTokenTotal(tokens())}</div>
+    <div class="session-context-panel border-r border-base border-b px-3 py-3 space-y-3">
+      <div class="flex flex-wrap items-center gap-2 text-xs text-primary/90">
+        <div class={headingClass}>Tokens</div>
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Input</span>
+          <span class="font-semibold text-primary">{formatTokenTotal(inputTokens())}</span>
        </div>
-        <div class="text-xs text-primary/70 text-right leading-tight">{costLabel()}</div>
-      </div>
-      <div class="mt-4">
-        <div class="flex items-center justify-between mb-1">
-          <div class="text-xs font-semibold text-primary/70 uppercase tracking-wide">Context window usage</div>
-          <div class="text-sm font-medium text-primary">{contextUsagePercent() !== null ? `${contextUsagePercent()}%` : "--"}</div>
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Output</span>
+          <span class="font-semibold text-primary">{formatTokenTotal(outputTokens())}</span>
        </div>
-        <div class="text-sm text-primary/90">
-          {contextWindow()
-            ? `${formatTokenTotal(contextUsageTokens())} of ${formatTokenTotal(contextWindow())}`
-            : "Window size unavailable"}
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Cost</span>
+          <span class="font-semibold text-primary">{costDisplay()}</span>
        </div>
      </div>
-      <div class="mt-3 h-1.5 rounded-full bg-base relative overflow-hidden">
-        <div
-          class="absolute inset-y-0 left-0 rounded-full bg-accent-primary transition-[width]"
-          style={{ width: contextUsagePercent() === null ? "0%" : `${contextUsagePercent()}%` }}
-        />
+
+      <div class="flex flex-wrap items-center gap-2 text-xs text-primary/90">
+        <div class={headingClass}>Context</div>
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Used</span>
+          <span class="font-semibold text-primary">{formatTokenTotal(actualUsageTokens())}</span>
+        </div>
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Avail</span>
+          <span class="font-semibold text-primary">{formatTokenValue(availableTokens())}</span>
+        </div>
      </div>
    </div>
  )
--- a/packages/ui/src/lib/formatters.ts
+++ b/packages/ui/src/lib/formatters.ts
@@ -1,4 +1,7 @@
 export function formatTokenTotal(value: number): string {
+  if (value >= 1_000_000_000) {
+    return `${(value / 1_000_000_000).toFixed(1)}B`
+  }
  if (value >= 1_000_000) {
    return `${(value / 1_000_000).toFixed(1)}M`
  }
--- a/packages/ui/src/stores/session-api.ts
+++ b/packages/ui/src/stores/session-api.ts
@@ -21,7 +21,7 @@ import {
  loading,
  setLoading,
 } from "./session-state"
-import { getDefaultModel, isModelValid } from "./session-models"
+import { DEFAULT_MODEL_OUTPUT_LIMIT, getDefaultModel, isModelValid } from "./session-models"
 import {
  computeDisplayParts,
  clearSessionIndex,
@@ -29,6 +29,7 @@ import {
  initializePartVersion,
  normalizeMessagePart,
  rebuildSessionIndex,
+  rebuildSessionUsage,
  updateSessionInfo,
 } from "./session-messages"

@@ -212,18 +213,25 @@ async function createSession(instanceId: string, agent?: string): Promise<Sessio
    const initialModel = initialProvider?.models.find((m) => m.id === session.model.modelId)
    const initialContextWindow = initialModel?.limit?.context ?? 0
    const initialSubscriptionModel = initialModel?.cost?.input === 0 && initialModel?.cost?.output === 0
-    const initialContextPercent = initialContextWindow > 0 ? 0 : null
+    const initialOutputLimit =
+      initialModel?.limit?.output && initialModel.limit.output > 0
+        ? initialModel.limit.output
+        : DEFAULT_MODEL_OUTPUT_LIMIT
+    const initialContextAvailable = initialContextWindow > 0 ? initialContextWindow : null

    setSessionInfoByInstance((prev) => {
      const next = new Map(prev)
      const instanceInfo = new Map(prev.get(instanceId))
      instanceInfo.set(session.id, {
-        tokens: 0,
        cost: 0,
        contextWindow: initialContextWindow,
        isSubscriptionModel: Boolean(initialSubscriptionModel),
-        contextUsageTokens: 0,
-        contextUsagePercent: initialContextPercent,
+        inputTokens: 0,
+        outputTokens: 0,
+        reasoningTokens: 0,
+        actualUsageTokens: 0,
+        modelOutputLimit: initialOutputLimit,
+        contextAvailableTokens: initialContextAvailable,
      })
      next.set(instanceId, instanceInfo)
      return next
@@ -310,18 +318,23 @@ async function forkSession(
  const forkModel = forkProvider?.models.find((m) => m.id === forkedSession.model.modelId)
  const forkContextWindow = forkModel?.limit?.context ?? 0
  const forkSubscriptionModel = forkModel?.cost?.input === 0 && forkModel?.cost?.output === 0
-  const forkContextPercent = forkContextWindow > 0 ? 0 : null
+  const forkOutputLimit =
+    forkModel?.limit?.output && forkModel.limit.output > 0 ? forkModel.limit.output : DEFAULT_MODEL_OUTPUT_LIMIT
+  const forkContextAvailable = forkContextWindow > 0 ? forkContextWindow : null

  setSessionInfoByInstance((prev) => {
    const next = new Map(prev)
    const instanceInfo = new Map(prev.get(instanceId))
    instanceInfo.set(forkedSession.id, {
-      tokens: 0,
      cost: 0,
      contextWindow: forkContextWindow,
      isSubscriptionModel: Boolean(forkSubscriptionModel),
-      contextUsageTokens: 0,
-      contextUsagePercent: forkContextPercent,
+      inputTokens: 0,
+      outputTokens: 0,
+      reasoningTokens: 0,
+      actualUsageTokens: 0,
+      modelOutputLimit: forkOutputLimit,
+      contextAvailableTokens: forkContextAvailable,
    })
    next.set(instanceId, instanceInfo)
    return next
@@ -587,6 +600,7 @@ async function loadMessages(instanceId: string, sessionId: string, force = false
    })

    rebuildSessionIndex(instanceId, sessionId, messages)
+    rebuildSessionUsage(instanceId, sessionId, messagesInfo)

    setMessagesLoaded((prev) => {
      const next = new Map(prev)
@@ -595,6 +609,7 @@ async function loadMessages(instanceId: string, sessionId: string, force = false
      next.set(instanceId, loadedSet)
      return next
    })
+
  } catch (error) {
    console.error("Failed to load messages:", error)
    throw error
@@ -608,17 +623,17 @@ async function loadMessages(instanceId: string, sessionId: string, force = false
      return next
    })
  }
- 
-   updateSessionInfo(instanceId, sessionId)
-   refreshPermissionsForSession(instanceId, sessionId)
- }

+  updateSessionInfo(instanceId, sessionId)
+  refreshPermissionsForSession(instanceId, sessionId)
+}

 export {
  createSession,
  deleteSession,
  fetchAgents,
  fetchProviders,
+
  fetchSessions,
  forkSession,
  loadMessages,
--- a/packages/ui/src/stores/session-events.ts
+++ b/packages/ui/src/stores/session-events.ts
@@ -30,6 +30,7 @@ import {
  normalizeMessagePart,
  rebuildSessionIndex,
  updateSessionInfo,
+  updateUsageFromMessageInfo,
 } from "./session-messages"
 import { loadMessages } from "./session-api"
 import { setSessionCompactionState } from "./session-compaction"
@@ -305,6 +306,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
    }

    session.messagesInfo.set(info.id, info)
+    updateUsageFromMessageInfo(instanceId, info.sessionID, info)
    withSession(instanceId, info.sessionID, () => {
      /* ensure reactivity */
    })
@@ -314,6 +316,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
  }
 }

+
 function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): void {
  const info = event.properties?.info
  if (!info) return
--- a/packages/ui/src/stores/session-messages.ts
+++ b/packages/ui/src/stores/session-messages.ts
@@ -1,9 +1,9 @@
 import type { Message, MessageDisplayParts } from "../types/message"
-import { partHasRenderableText } from "../types/message"
+import { partHasRenderableText, type MessageInfo } from "../types/message"
 import type { Provider } from "../types/session"

 import { decodeHtmlEntities } from "../lib/markdown"
-import { providers, sessions, setSessionInfoByInstance } from "./session-state"
+import { providers, sessions, sessionInfoByInstance, setSessionInfoByInstance } from "./session-state"
 import { DEFAULT_MODEL_OUTPUT_LIMIT } from "./session-models"

 interface SessionIndexCache {
@@ -11,7 +11,153 @@ interface SessionIndexCache {
  partIndex: Map<string, Map<string, number>>
 }

+interface AssistantUsageEntry {
+  info: MessageInfo
+  inputTokens: number
+  outputTokens: number
+  reasoningTokens: number
+  combinedTokens: number
+  cost: number
+  hasContextUsage: boolean
+  timestamp: number
+}
+
+interface SessionUsageState {
+  entries: Map<string, AssistantUsageEntry>
+  totalInputTokens: number
+  totalOutputTokens: number
+  totalReasoningTokens: number
+  totalCost: number
+  latestEntry: AssistantUsageEntry | null
+}
+
 const sessionIndexes = new Map<string, Map<string, SessionIndexCache>>()
+const sessionUsageStates = new Map<string, Map<string, SessionUsageState>>()
+
+function createEmptyUsageState(): SessionUsageState {
+  return {
+    entries: new Map(),
+    totalInputTokens: 0,
+    totalOutputTokens: 0,
+    totalReasoningTokens: 0,
+    totalCost: 0,
+    latestEntry: null,
+  }
+}
+
+function getUsageInstance(instanceId: string): Map<string, SessionUsageState> {
+  let usageMap = sessionUsageStates.get(instanceId)
+  if (!usageMap) {
+    usageMap = new Map()
+    sessionUsageStates.set(instanceId, usageMap)
+  }
+  return usageMap
+}
+
+function getSessionUsageState(instanceId: string, sessionId: string): SessionUsageState {
+  const usageMap = getUsageInstance(instanceId)
+  let state = usageMap.get(sessionId)
+  if (!state) {
+    state = createEmptyUsageState()
+    usageMap.set(sessionId, state)
+  }
+  return state
+}
+
+function recomputeLatestEntry(state: SessionUsageState) {
+  state.latestEntry = null
+  for (const entry of state.entries.values()) {
+    if (!state.latestEntry || entry.timestamp >= state.latestEntry.timestamp) {
+      state.latestEntry = entry
+    }
+  }
+}
+
+function extractAssistantUsage(info: MessageInfo): AssistantUsageEntry | null {
+  if (!info || info.role !== "assistant") return null
+  if (!info.tokens) return null
+  const tokens = info.tokens
+  const inputTokens = tokens.input ?? 0
+  const outputTokens = tokens.output ?? 0
+  const reasoningTokens = tokens.reasoning ?? 0
+  if (inputTokens === 0 && outputTokens === 0 && reasoningTokens === 0) {
+    return null
+  }
+  const cacheReadTokens = tokens.cache?.read ?? 0
+  const cacheWriteTokens = tokens.cache?.write ?? 0
+  const combinedTokens = info.summary
+    ? outputTokens
+    : inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens + reasoningTokens
+  const cost = info.cost ?? 0
+  const hasContextUsage = inputTokens + cacheReadTokens + cacheWriteTokens > 0
+  return {
+    info,
+    inputTokens,
+    outputTokens,
+    reasoningTokens,
+    combinedTokens,
+    cost,
+    hasContextUsage,
+    timestamp: info.time?.created ?? 0,
+  }
+}
+
+function removeUsageEntry(state: SessionUsageState, messageId: string | undefined) {
+  if (!messageId) return
+  const existing = state.entries.get(messageId)
+  if (!existing) return
+  state.entries.delete(messageId)
+  state.totalInputTokens -= existing.inputTokens
+  state.totalOutputTokens -= existing.outputTokens
+  state.totalReasoningTokens -= existing.reasoningTokens
+  state.totalCost -= existing.cost
+  if (state.latestEntry?.info.id === messageId) {
+    recomputeLatestEntry(state)
+  }
+}
+
+function addUsageEntry(state: SessionUsageState, entry: AssistantUsageEntry) {
+  state.entries.set(entry.info.id, entry)
+  state.totalInputTokens += entry.inputTokens
+  state.totalOutputTokens += entry.outputTokens
+  state.totalReasoningTokens += entry.reasoningTokens
+  state.totalCost += entry.cost
+  if (!state.latestEntry || entry.timestamp >= state.latestEntry.timestamp) {
+    state.latestEntry = entry
+  }
+}
+
+function updateUsageFromMessageInfo(instanceId: string, sessionId: string, info: MessageInfo) {
+  const messageId = typeof info.id === "string" ? info.id : undefined
+  if (!messageId) return
+  const state = getSessionUsageState(instanceId, sessionId)
+  removeUsageEntry(state, messageId)
+  const entry = extractAssistantUsage(info)
+  if (entry) {
+    addUsageEntry(state, entry)
+  }
+}
+
+function rebuildSessionUsage(instanceId: string, sessionId: string, messagesInfo: Map<string, MessageInfo>) {
+  const usageMap = getUsageInstance(instanceId)
+  const nextState = createEmptyUsageState()
+  for (const info of messagesInfo.values()) {
+    const entry = extractAssistantUsage(info)
+    if (entry) {
+      addUsageEntry(nextState, entry)
+    }
+  }
+  usageMap.set(sessionId, nextState)
+}
+
+function clearSessionUsage(instanceId: string, sessionId: string) {
+  const usageMap = sessionUsageStates.get(instanceId)
+  if (!usageMap) return
+  usageMap.delete(sessionId)
+  if (usageMap.size === 0) {
+    sessionUsageStates.delete(instanceId)
+  }
+}

 function decodeTextSegment(segment: any): any {
  if (typeof segment === "string") {
@@ -163,10 +309,12 @@ function clearSessionIndex(instanceId: string, sessionId: string) {
      sessionIndexes.delete(instanceId)
    }
  }
+  clearSessionUsage(instanceId, sessionId)
 }

 function removeSessionIndexes(instanceId: string) {
  sessionIndexes.delete(instanceId)
+  sessionUsageStates.delete(instanceId)
 }

 function updateSessionInfo(instanceId: string, sessionId: string) {
@@ -176,52 +324,67 @@ function updateSessionInfo(instanceId: string, sessionId: string) {
  const session = instanceSessions.get(sessionId)
  if (!session) return

-  let tokens = 0
-  let cost = 0
  let contextWindow = 0
  let isSubscriptionModel = false
  let modelID = ""
  let providerID = ""
  let actualUsageTokens = 0
-  let contextUsagePercent: number | null = null
-  let hasContextUsage = false

-  if (session.messagesInfo.size > 0) {
-    const messageArray = Array.from(session.messagesInfo.values()).reverse()
+  const usageState = getSessionUsageState(instanceId, sessionId)
+  const hasUsageEntries = usageState.entries.size > 0

-    for (const info of messageArray) {
-      if (info.role === "assistant" && info.tokens) {
-        const usage = info.tokens
+  let totalInputTokens = hasUsageEntries ? usageState.totalInputTokens : 0
+  let totalOutputTokens = hasUsageEntries ? usageState.totalOutputTokens : 0
+  let totalReasoningTokens = hasUsageEntries ? usageState.totalReasoningTokens : 0
+  let totalCost = hasUsageEntries ? usageState.totalCost : 0

-        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
+  let latestAssistantInfo: MessageInfo | null = usageState.latestEntry?.info ?? null
+  let latestHasContextUsage = usageState.latestEntry?.hasContextUsage ?? false
+  const previousInfo = sessionInfoByInstance().get(instanceId)?.get(sessionId)
+  let contextAvailableTokens: number | null = null
+  let contextAvailableFromPrevious = false

-          if (info.summary) {
-            tokens = outputTokens
-          } else {
-            tokens = inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens + reasoningTokens
-          }
+  if (latestAssistantInfo) {
+    const infoAny = latestAssistantInfo as any
+    actualUsageTokens = usageState.latestEntry?.combinedTokens ?? 0
+    modelID = infoAny.modelID || ""
+    providerID = infoAny.providerID || ""
+  } else if (previousInfo) {
+    totalInputTokens = previousInfo.inputTokens
+    totalOutputTokens = previousInfo.outputTokens
+    totalReasoningTokens = previousInfo.reasoningTokens
+    totalCost = previousInfo.cost
+    actualUsageTokens = previousInfo.actualUsageTokens

-          cost = info.cost || 0
-          actualUsageTokens = tokens
-          hasContextUsage = inputTokens + cacheReadTokens + cacheWriteTokens > 0
+    const previousContextWindow = previousInfo.contextWindow
+    const previousContextAvailable = previousInfo.contextAvailableTokens ?? null
+    const previousHasContextUsage =
+      previousContextAvailable !== null && previousContextWindow > 0
+        ? previousContextAvailable < previousContextWindow
+        : false

-          modelID = info.modelID || ""
-          providerID = info.providerID || ""
-          isSubscriptionModel = cost === 0
-
-          break
-        }
-      }
+    if (contextWindow === 0) {
+      contextWindow = previousContextWindow
    }
+
+    if (contextWindow !== previousContextWindow) {
+      contextAvailableTokens = null
+      contextAvailableFromPrevious = false
+      latestHasContextUsage = previousHasContextUsage
+    } else {
+      contextAvailableTokens = previousContextAvailable
+      contextAvailableFromPrevious = true
+      latestHasContextUsage = previousHasContextUsage
+    }
+
+    isSubscriptionModel = previousInfo.isSubscriptionModel
  }

  const instanceProviders = providers().get(instanceId) || []

+
+
+
  const sessionModel = session.model
  let selectedModel: Provider["models"][number] | undefined

@@ -252,30 +415,32 @@ function updateSessionInfo(instanceId: string, sessionId: string) {
  }

  const outputBudget = Math.min(modelOutputLimit, DEFAULT_MODEL_OUTPUT_LIMIT)
-  let contextUsageTokens = 0

-  if (hasContextUsage && actualUsageTokens > 0) {
-    contextUsageTokens = actualUsageTokens + outputBudget
+  if (!contextAvailableFromPrevious) {
    if (contextWindow > 0) {
-      const percent = Math.round((contextUsageTokens / contextWindow) * 100)
-      contextUsagePercent = Math.min(100, Math.max(0, percent))
+      if (latestHasContextUsage && actualUsageTokens > 0) {
+        contextAvailableTokens = Math.max(contextWindow - (actualUsageTokens + outputBudget), 0)
+      } else {
+        contextAvailableTokens = contextWindow
+      }
    } else {
-      contextUsagePercent = null
+      contextAvailableTokens = null
    }
-  } else {
-    contextUsagePercent = contextWindow > 0 ? 0 : null
  }

  setSessionInfoByInstance((prev) => {
    const next = new Map(prev)
    const instanceInfo = new Map(prev.get(instanceId))
    instanceInfo.set(sessionId, {
-      tokens,
-      cost,
+      cost: totalCost,
      contextWindow,
      isSubscriptionModel,
-      contextUsageTokens,
-      contextUsagePercent,
+      inputTokens: totalInputTokens,
+      outputTokens: totalOutputTokens,
+      reasoningTokens: totalReasoningTokens,
+      actualUsageTokens,
+      modelOutputLimit,
+      contextAvailableTokens,
    })
    next.set(instanceId, instanceInfo)
    return next
@@ -290,6 +455,8 @@ export {
  initializePartVersion,
  normalizeMessagePart,
  rebuildSessionIndex,
+  rebuildSessionUsage,
  removeSessionIndexes,
  updateSessionInfo,
+  updateUsageFromMessageInfo,
 }
--- a/packages/ui/src/stores/session-state.ts
+++ b/packages/ui/src/stores/session-state.ts
@@ -3,12 +3,15 @@ import { createSignal } from "solid-js"
 import type { Session, Agent, Provider } from "../types/session"

 export interface SessionInfo {
-  tokens: number
  cost: number
  contextWindow: number
  isSubscriptionModel: boolean
-  contextUsageTokens: number
-  contextUsagePercent: number | null
+  inputTokens: number
+  outputTokens: number
+  reasoningTokens: number
+  actualUsageTokens: number
+  modelOutputLimit: number
+  contextAvailableTokens: number | null
 }

 const [sessions, setSessions] = createSignal<Map<string, Map<string, Session>>>(new Map())
--- a/packages/ui/vite.config.js
+++ b/packages/ui/vite.config.js
@@ -0,0 +1,37 @@
+import { defineConfig } from "vite"
+import solid from "vite-plugin-solid"
+import { dirname, resolve } from "path"
+import { fileURLToPath } from "url"
+
+const __dirname = dirname(fileURLToPath(import.meta.url))
+
+export default defineConfig({
+  root: "./src/renderer",
+  plugins: [solid()],
+  css: {
+    postcss: "./postcss.config.js",
+  },
+  resolve: {
+    alias: {
+      "@": resolve(__dirname, "./src"),
+    },
+  },
+  optimizeDeps: {
+    exclude: ["lucide-solid"],
+  },
+  ssr: {
+    noExternal: ["lucide-solid"],
+  },
+  server: {
+    port: 3000,
+  },
+  build: {
+    outDir: "dist",
+    rollupOptions: {
+      input: {
+        main: resolve(__dirname, "./src/renderer/index.html"),
+        loading: resolve(__dirname, "./src/renderer/loading.html"),
+      },
+    },
+  },
+})