From f75a7d46bbef4fd64ec3ca7920bfe88f1a99da0b Mon Sep 17 00:00:00 2001
From: weishu <twsxtd@gmail.com>
Date: Thu, 18 Dec 2025 08:50:07 +0800
Subject: [PATCH] feat: add model-aware context budget calculation for status
 bar warnings

---
 web/src/chat/modelConfig.ts                   | 28 +++++++++++++++++++
 .../AssistantChat/HappyComposer.tsx           |  1 +
 .../components/AssistantChat/StatusBar.tsx    | 23 ++++++++-------
 3 files changed, 42 insertions(+), 10 deletions(-)
 create mode 100644 web/src/chat/modelConfig.ts

diff --git a/web/src/chat/modelConfig.ts b/web/src/chat/modelConfig.ts
new file mode 100644
index 00000000..39ef20fa
--- /dev/null
+++ b/web/src/chat/modelConfig.ts
@@ -0,0 +1,28 @@
+import type { ModelMode } from '@/types/api'
+
+/**
+ * Context windows vary by model/provider and may change over time.
+ *
+ * The UI only needs this to compute a conservative "context remaining" warning.
+ * We intentionally keep a headroom budget to avoid false confidence near the limit
+ * (system prompts, tool overhead, and other hidden tokens can consume extra space).
+ *
+ * If/when the server provides an explicit per-session context limit, prefer that
+ * and use this only as a fallback.
+ */
+const CONTEXT_HEADROOM_TOKENS = 10_000
+
+const MODEL_CONTEXT_WINDOWS: Record<NonNullable<ModelMode>, number> = {
+    // Claude Code modes used in this app; currently treated as ~200k context.
+    default: 200_000,
+    sonnet: 200_000,
+    opus: 200_000
+}
+
+export function getContextBudgetTokens(modelMode: ModelMode): number | null {
+    const mode: NonNullable<ModelMode> = modelMode ?? 'default'
+    const windowTokens = MODEL_CONTEXT_WINDOWS[mode]
+    if (!windowTokens) return null
+    return Math.max(1, windowTokens - CONTEXT_HEADROOM_TOKENS)
+}
+
diff --git a/web/src/components/AssistantChat/HappyComposer.tsx b/web/src/components/AssistantChat/HappyComposer.tsx
index 5603c3f8..849c8d58 100644
--- a/web/src/components/AssistantChat/HappyComposer.tsx
+++ b/web/src/components/AssistantChat/HappyComposer.tsx
@@ -401,6 +401,7 @@ export function HappyComposer(props: {
                         thinking={thinking}
                         agentState={agentState}
                         contextSize={contextSize}
+                        modelMode={modelMode}
                         permissionMode={permissionMode}
                     />
 
diff --git a/web/src/components/AssistantChat/StatusBar.tsx b/web/src/components/AssistantChat/StatusBar.tsx
index fbb54547..d1ff5ec7 100644
--- a/web/src/components/AssistantChat/StatusBar.tsx
+++ b/web/src/components/AssistantChat/StatusBar.tsx
@@ -1,5 +1,6 @@
 import { useMemo } from 'react'
-import type { AgentState, PermissionMode } from '@/types/api'
+import type { AgentState, ModelMode, PermissionMode } from '@/types/api'
+import { getContextBudgetTokens } from '@/chat/modelConfig'
 
 const PERMISSION_MODE_LABELS: Record<string, string> = {
     default: 'Default',
@@ -8,9 +9,6 @@ const PERMISSION_MODE_LABELS: Record<string, string> = {
     bypassPermissions: 'Bypass All'
 }
 
-// Max context size for percentage calculation
-const MAX_CONTEXT_SIZE = 190000
-
 // Vibing messages for thinking state
 const VIBING_MESSAGES = [
     "Accomplishing", "Actioning", "Actualizing", "Baking", "Booping", "Brewing",
@@ -73,9 +71,9 @@ function getConnectionStatus(
     }
 }
 
-function getContextWarning(contextSize: number): { text: string; color: string } | null {
-    const percentageUsed = (contextSize / MAX_CONTEXT_SIZE) * 100
-    const percentageRemaining = 100 - percentageUsed
+function getContextWarning(contextSize: number, maxContextSize: number): { text: string; color: string } | null {
+    const percentageUsed = (contextSize / maxContextSize) * 100
+    const percentageRemaining = Math.max(0, 100 - percentageUsed)
 
     if (percentageRemaining <= 5) {
         return { text: `${Math.round(percentageRemaining)}% left`, color: 'text-red-500' }
@@ -91,6 +89,7 @@ export function StatusBar(props: {
     thinking: boolean
     agentState: AgentState | null | undefined
     contextSize?: number
+    modelMode?: ModelMode
     permissionMode?: PermissionMode
 }) {
     const connectionStatus = useMemo(
@@ -99,8 +98,13 @@ export function StatusBar(props: {
     )
 
     const contextWarning = useMemo(
-        () => props.contextSize !== undefined ? getContextWarning(props.contextSize) : null,
-        [props.contextSize]
+        () => {
+            if (props.contextSize === undefined) return null
+            const maxContextSize = getContextBudgetTokens(props.modelMode)
+            if (!maxContextSize) return null
+            return getContextWarning(props.contextSize, maxContextSize)
+        },
+        [props.contextSize, props.modelMode]
     )
 
     const permissionMode = props.permissionMode
@@ -136,4 +140,3 @@ export function StatusBar(props: {
         </div>
     )
 }
-