Show premium-session quota in freebuff session-ended banner

Server now attaches `rateLimitsByModel` to active and ended responses (in
addition to queued), and the banner shows "N of 5 premium sessions used
today" when the session ends — matching the waiting-room landing line so
users can see how many free sessions they have left for the day.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: jahooma

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,10 +9,14 @@ import {
   returnToFreebuffLanding,
 } from '../hooks/use-freebuff-session'
 import { useTheme } from '../hooks/use-theme'
+import { useFreebuffSessionStore } from '../state/freebuff-session-store'
 import { BORDER_CHARS } from '../utils/ui-constants'
 
 import type { KeyEvent } from '@opentui/core'
 
+const formatSessionUnits = (units: number): string =>
+  Number.isInteger(units) ? String(units) : units.toFixed(1)
+
 interface SessionEndedBannerProps {
   /** True while an agent request is still streaming under the server-side
    *  grace window. Swaps the Enter-to-rejoin affordance for a "let it
@@ -32,6 +37,17 @@ 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 quotaColor =
+    premiumQuota && premiumQuota.recentCount >= premiumQuota.limit
+      ? theme.secondary
+      : theme.muted
+
   // 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
@@ -95,6 +111,13 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
     >
       <text style={{ fg: theme.foreground, wrapMode: 'word' }}>
         Your freebuff session has ended.
+        {premiumQuota && (
+          <span fg={quotaColor}>
+            {'  ·  '}
+            {formatSessionUnits(premiumQuota.recentCount)} of{' '}
+            {premiumQuota.limit} premium sessions used today
+          </span>
+        )}
       </text>
       {isStreaming ? (
         <text style={{ fg: theme.muted, wrapMode: 'word' }}>

cli/src/hooks/use-freebuff-session.ts

@@ -19,6 +19,8 @@ import {
 import { logger } from '../utils/logger'
 import { saveFreebuffModelPreference } from '../utils/settings'
 
+import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
+
 import type { FreebuffSessionResponse } from '../types/freebuff-session'
 import type {
   FreebuffCountryBlockReason,
@@ -351,11 +353,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 +515,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
         }
 

common/src/types/freebuff-session.ts

@@ -31,6 +31,17 @@ 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. */
+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 +130,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

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,