Show premium-session quota in freebuff session-ended banner (#618)

jahooma · claude · web-flow · commit 84cc946c1b4b · 2026-05-07T17:43:15.000-07:00
Co-authored-by: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/cli/src/components/freebuff-model-selector.tsx b/cli/src/components/freebuff-model-selector.tsx
@@ -11,6 +11,7 @@ import {
   isFreebuffModelAvailable,
   isFreebuffPremiumModelId,
 } from '@codebuff/common/constants/freebuff-models'
+import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
 
 import { joinFreebuffQueue } from '../hooks/use-freebuff-session'
 import { useNow } from '../hooks/use-now'
@@ -127,10 +128,7 @@ export const FreebuffModelSelector: React.FC = () => {
   }, [now, selectedModel, session, setSelectedModel])
 
   const committedModelId = session?.status === 'queued' ? session.model : null
-  const rateLimitsByModel =
-    session && 'rateLimitsByModel' in session
-      ? session.rateLimitsByModel
-      : undefined
+  const rateLimitsByModel = getRateLimitsByModel(session)
 
   const BUTTON_CHROME = 4 // 2 border + 2 padding
   const NAME_GAP = 2 // spaces between name column and details column
diff --git a/cli/src/components/session-ended-banner.tsx b/cli/src/components/session-ended-banner.tsx
@@ -1,3 +1,4 @@
+import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
 import { TextAttributes } from '@opentui/core'
 import { useKeyboard } from '@opentui/react'
 import React, { useCallback, useState } from 'react'
@@ -8,6 +9,8 @@ import {
   returnToFreebuffLanding,
 } from '../hooks/use-freebuff-session'
 import { useTheme } from '../hooks/use-theme'
+import { useFreebuffSessionStore } from '../state/freebuff-session-store'
+import { formatSessionUnits } from '../utils/format-session-units'
 import { BORDER_CHARS } from '../utils/ui-constants'
 
 import type { KeyEvent } from '@opentui/core'
@@ -32,6 +35,19 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
     'waiting-room' | 'same-chat' | null
   >(null)
 
+  // All premium models share one daily pool; the server replicates the same
+  // snapshot under each premium model id, so the first entry has the right
+  // count.
+  const premiumQuota = useFreebuffSessionStore(
+    (s) => Object.values(getRateLimitsByModel(s.session) ?? {})[0] ?? null,
+  )
+  const isQuotaExhausted = premiumQuota
+    ? premiumQuota.recentCount >= premiumQuota.limit
+    : false
+  const bannerTitle = premiumQuota
+    ? `Session ended  ·  ${formatSessionUnits(premiumQuota.recentCount)} of ${premiumQuota.limit} premium sessions used today`
+    : 'Session ended'
+
   // While a request is still streaming, restart is disabled: it would
   // unmount <Chat> and abort the in-flight agent run. The promise is "we
   // let the agent finish" — honoring that means Enter does nothing until
@@ -78,12 +94,15 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
 
   return (
     <box
-      title="Session ended"
+      title={bannerTitle}
       titleAlignment="center"
       style={{
         width: '100%',
         borderStyle: 'single',
-        borderColor: theme.muted,
+        // Amber border doubles as the "you've hit the cap" signal now that
+        // the quota count lives in the title (which can't carry per-char
+        // color); muted otherwise.
+        borderColor: isQuotaExhausted ? theme.secondary : theme.muted,
         customBorderChars: BORDER_CHARS,
         paddingLeft: 1,
         paddingRight: 1,
@@ -93,9 +112,6 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
         gap: 0,
       }}
     >
-      <text style={{ fg: theme.foreground, wrapMode: 'word' }}>
-        Your freebuff session has ended.
-      </text>
       {isStreaming ? (
         <text style={{ fg: theme.muted, wrapMode: 'word' }}>
           Agent is wrapping up. Rejoin the wait room after it's finished.
@@ -115,7 +131,7 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
                 fg:
                   pendingAction === 'same-chat'
                     ? theme.muted
-                    : theme.primary,
+                    : theme.foreground,
               }}
               attributes={TextAttributes.BOLD}
             >
@@ -144,11 +160,14 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
                     ? theme.muted
                     : theme.foreground,
               }}
-              attributes={TextAttributes.BOLD}
             >
-              {pendingAction === 'waiting-room'
-                ? 'Opening model selection…'
-                : 'Change model (ESC)'}
+              {pendingAction === 'waiting-room' ? (
+                'Opening model selection…'
+              ) : (
+                <>
+                  Change model<span fg={theme.muted}>{'   Esc'}</span>
+                </>
+              )}
             </text>
           </Button>
         </box>
diff --git a/cli/src/components/waiting-room-screen.tsx b/cli/src/components/waiting-room-screen.tsx
@@ -15,8 +15,10 @@ import { useSheenAnimation } from '../hooks/use-sheen-animation'
 import { useTerminalDimensions } from '../hooks/use-terminal-dimensions'
 import { useTheme } from '../hooks/use-theme'
 import { exitFreebuffCleanly } from '../utils/freebuff-exit'
+import { formatSessionUnits } from '../utils/format-session-units'
 import { getLogoAccentColor, getLogoBlockColor } from '../utils/theme-system'
 import { FREEBUFF_PREMIUM_SESSION_LIMIT } from '@codebuff/common/constants/freebuff-models'
+import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
 
 import type { FreebuffSessionResponse } from '../types/freebuff-session'
 import type { FreebuffIpPrivacySignal } from '@codebuff/common/types/freebuff-session'
@@ -59,9 +61,6 @@ const formatRetryAfter = (ms: number): string => {
   return rem === 0 ? `${hours}h` : `${hours}h ${rem}m`
 }
 
-const formatSessionUnits = (units: number): string =>
-  Number.isInteger(units) ? String(units) : units.toFixed(1)
-
 const PRIVACY_SIGNAL_LABELS: Partial<Record<FreebuffIpPrivacySignal, string>> =
   {
     anonymous: 'anonymized network',
@@ -268,10 +267,7 @@ export const WaitingRoomScreen: React.FC<WaitingRoomScreenProps> = ({
   // pool; the server replicates the same snapshot under each premium model
   // id, so any entry has the right count. Renders amber when exhausted so
   // the limit reads as "you've hit it" rather than just another count.
-  const rateLimitsByModel =
-    session && 'rateLimitsByModel' in session
-      ? session.rateLimitsByModel
-      : undefined
+  const rateLimitsByModel = getRateLimitsByModel(session)
   const sharedPremiumUsed = rateLimitsByModel
     ? (Object.values(rateLimitsByModel)[0]?.recentCount ?? 0)
     : 0
diff --git a/cli/src/hooks/use-freebuff-session.ts b/cli/src/hooks/use-freebuff-session.ts
@@ -3,6 +3,7 @@ import {
   FALLBACK_FREEBUFF_MODEL_ID,
   resolveFreebuffModel,
 } from '@codebuff/common/constants/freebuff-models'
+import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
 import { useEffect } from 'react'
 
 import {
@@ -351,11 +352,16 @@ export function markFreebuffSessionCountryBlocked(params: {
 }
 
 /** Flip into the local `ended` state without an instanceId (server has lost
- *  our row). The chat surface stays mounted with the rejoin banner. */
+ *  our row). The chat surface stays mounted with the rejoin banner.
+ *  Preserves any `rateLimitsByModel` snapshot from the prior session so the
+ *  banner can show today's premium-session count without an extra fetch. */
 export function markFreebuffSessionEnded(): void {
   if (!IS_FREEBUFF) return
   controller?.abort()
-  controller?.apply({ status: 'ended' })
+  const rateLimitsByModel = getRateLimitsByModel(
+    useFreebuffSessionStore.getState().session,
+  )
+  controller?.apply({ status: 'ended', rateLimitsByModel })
 }
 
 interface UseFreebuffSessionResult {
@@ -508,12 +514,18 @@ export function useFreebuffSession(): UseFreebuffSessionResult {
         // active|ended → none means we've passed the server's hard cutoff.
         // Synthesize a no-instanceId ended state so the chat surface stays
         // mounted with the Enter-to-rejoin banner instead of looping back
-        // through the waiting room.
+        // through the waiting room. Carry forward whichever rate-limit
+        // snapshot we have — preferring the fresh `none` snapshot, falling
+        // back to whatever was on the prior active/ended row — so the
+        // banner's "N of M used today" line stays populated.
         if (
           (previousStatus === 'active' || previousStatus === 'ended') &&
           next.status === 'none'
         ) {
-          apply({ status: 'ended' })
+          const rateLimitsByModel =
+            next.rateLimitsByModel ??
+            getRateLimitsByModel(useFreebuffSessionStore.getState().session)
+          apply({ status: 'ended', rateLimitsByModel })
           return
         }
 
diff --git a/cli/src/utils/format-session-units.ts b/cli/src/utils/format-session-units.ts
@@ -0,0 +1,6 @@
+/** Premium-session counts come back from the server as `recentCount` units
+ *  that may be fractional (a long agent run can consume 1.3 sessions). Render
+ *  integers without a trailing `.0`, fractionals at one decimal — matches the
+ *  `limit` field which is always integer. */
+export const formatSessionUnits = (units: number): string =>
+  Number.isInteger(units) ? String(units) : units.toFixed(1)
diff --git a/common/src/types/freebuff-session.ts b/common/src/types/freebuff-session.ts
@@ -31,6 +31,20 @@ export type FreebuffSessionRateLimitByModel = Record<
   FreebuffSessionRateLimit
 >
 
+/** Pull the per-model premium quota snapshot off whichever session statuses
+ *  carry it (queued, active, ended, none). Returns undefined for terminal /
+ *  pre-join states that have no quota field. The parameter is intentionally
+ *  loose so the CLI can pass its `FreebuffSessionResponse` (which adds the
+ *  client-only `takeover_prompt` variant) without a discriminated-union
+ *  ceremony at every call site. */
+export const getRateLimitsByModel = (
+  session: { status: string } | null | undefined,
+): FreebuffSessionRateLimitByModel | undefined =>
+  session && 'rateLimitsByModel' in session
+    ? (session as { rateLimitsByModel?: FreebuffSessionRateLimitByModel })
+        .rateLimitsByModel
+    : undefined
+
 export type FreebuffCountryBlockReason =
   | 'country_not_allowed'
   | 'anonymized_or_unknown_country'
@@ -119,6 +133,10 @@ export type FreebuffSessionServerResponse =
       expiresAt?: string
       gracePeriodEndsAt?: string
       gracePeriodRemainingMs?: number
+      /** Snapshot of the user's premium-session quota at the moment the
+       *  session ended. Lets the post-session banner show "N of M premium
+       *  sessions used today" without an extra round-trip. */
+      rateLimitsByModel?: FreebuffSessionRateLimitByModel
     }
   | {
       /** Another CLI on the same account rotated our instance id. Polling
diff --git a/web/src/server/free-session/__tests__/public-api.test.ts b/web/src/server/free-session/__tests__/public-api.test.ts
@@ -960,6 +960,38 @@ describe('getSessionState', () => {
     expect(state.gracePeriodRemainingMs).toBe(GRACE_MS - 60_000)
   })
 
+  test('ended view carries the full premium-quota snapshot', async () => {
+    // The post-session banner reads any entry from rateLimitsByModel since
+    // all premium models share one daily pool. Unlike queued/active, the
+    // ended view ships the full unfiltered map so a single banner read is
+    // always safe.
+    await requestSession({ userId: 'u1', model: DEFAULT_MODEL, deps })
+    const row = deps.rows.get('u1')!
+    row.status = 'active'
+    row.admitted_at = new Date(deps._now().getTime() - SESSION_LEN - 60_000)
+    row.expires_at = new Date(deps._now().getTime() - 60_000)
+    deps.admits.push({
+      user_id: 'u1',
+      model: FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID,
+      admitted_at: new Date(deps._now().getTime() - 30 * 60_000),
+    })
+
+    const state = await getSessionState({
+      userId: 'u1',
+      claimedInstanceId: row.active_instance_id,
+      deps,
+    })
+    if (state.status !== 'ended') throw new Error('unreachable')
+    expect(
+      state.rateLimitsByModel?.[FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID],
+    ).toEqual(expectedRateLimit(FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID, 1))
+    // Every premium model is present (sharing the same recentCount) so the
+    // banner can read any entry without caring which model the user was on.
+    expect(state.rateLimitsByModel?.[FREEBUFF_KIMI_MODEL_ID]).toEqual(
+      expectedRateLimit(FREEBUFF_KIMI_MODEL_ID, 1),
+    )
+  })
+
   test('row past grace window returns none', async () => {
     await requestSession({ userId: 'u1', model: DEFAULT_MODEL, deps })
     const row = deps.rows.get('u1')!
diff --git a/web/src/server/free-session/public-api.ts b/web/src/server/free-session/public-api.ts
@@ -416,21 +416,31 @@ export async function requestSession(params: {
   return attachRateLimit(params.userId, view, deps)
 }
 
-/** Thread the current quota snapshot onto queued/active views so the CLI can
- *  render "N of M sessions used". Other statuses pass through unchanged.
- *  Called on both POST and GET so the line stays live across polls. */
+/** Thread the current quota snapshot onto queued/active/ended views so the
+ *  CLI can render "N of M sessions used" — both during the session and on
+ *  the post-session banner. Other statuses pass through unchanged. Called on
+ *  both POST and GET so the line stays live across polls. */
 async function attachRateLimit(
   userId: string,
   view: SessionStateResponse,
   deps: SessionDeps,
 ): Promise<SessionStateResponse> {
-  if (view.status !== 'queued' && view.status !== 'active') return view
-  if (view.status === 'active') {
-    const snapshot = await fetchRateLimitSnapshot(userId, view.model, deps)
-    return snapshot ? { ...view, rateLimit: snapshot.info } : view
+  if (
+    view.status !== 'queued' &&
+    view.status !== 'active' &&
+    view.status !== 'ended'
+  ) {
+    return view
   }
-
   const allRateLimitsByModel = await fetchRateLimitsByModel(userId, deps)
+  // The ended view doesn't carry a model id, so it gets the full snapshot
+  // unfiltered — the banner reads any entry's recentCount (they all share the
+  // same daily premium pool). Queued/active filter out unused models so the
+  // landing screen and waiting-room title don't list every premium model with
+  // a "0 used today" hint.
+  if (view.status === 'ended') {
+    return { ...view, rateLimitsByModel: allRateLimitsByModel }
+  }
   const rateLimit = allRateLimitsByModel[view.model]
   return {
     ...view,
