diff --git a/ui/src/ui/views/chat.browser.test.ts b/ui/src/ui/views/chat.browser.test.ts
index c17525bb60b..0d3d6a58563 100644
--- a/ui/src/ui/views/chat.browser.test.ts
+++ b/ui/src/ui/views/chat.browser.test.ts
@@ -37,7 +37,7 @@ function createProps(overrides: Partial<ChatProps> = {}): ChatProps {
           key: "main",
           kind: "direct",
           updatedAt: null,
-          inputTokens: 3_800,
+          totalTokens: 3_800,
           contextTokens: 4_000,
         },
       ],
diff --git a/ui/src/ui/views/chat.ts b/ui/src/ui/views/chat.ts
index 88a712706f0..afa858fdd04 100644
--- a/ui/src/ui/views/chat.ts
+++ b/ui/src/ui/views/chat.ts
@@ -255,7 +255,8 @@ function renderContextNotice(
   session: GatewaySessionRow | undefined,
   defaultContextTokens: number | null,
 ) {
-  const used = session?.inputTokens ?? 0;
+  // Use the same notion as `session_status`: the fresh totalTokens
+  const used = session?.totalTokens ?? 0;
   const limit = session?.contextTokens ?? defaultContextTokens ?? 0;
   if (!used || !limit) {
     return nothing;