Extract formatSessionUnits util and reuse getRateLimitsByModel helper

Drops the inline `'rateLimitsByModel' in session` pattern from the
waiting-room screen and model selector now that `getRateLimitsByModel`
exists in common, and lifts `formatSessionUnits` out of waiting-room
so the session-ended banner doesn't carry a duplicate. Adds a regression
test for the ended view's full-snapshot quota behavior.

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

Author: jahooma

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

cli/src/components/session-ended-banner.tsx

@@ -10,13 +10,11 @@ import {
 } 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'
 
-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
@@ -43,10 +41,6 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
   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
@@ -112,7 +106,13 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
       <text style={{ fg: theme.foreground, wrapMode: 'word' }}>
         Your freebuff session has ended.
         {premiumQuota && (
-          <span fg={quotaColor}>
+          <span
+            fg={
+              premiumQuota.recentCount >= premiumQuota.limit
+                ? theme.secondary
+                : theme.muted
+            }
+          >
             {'  ·  '}
             {formatSessionUnits(premiumQuota.recentCount)} of{' '}
             {premiumQuota.limit} premium sessions used today

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

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 {
@@ -19,8 +20,6 @@ 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,

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)

common/src/types/freebuff-session.ts

@@ -33,7 +33,10 @@ export type FreebuffSessionRateLimitByModel = Record<
 
 /** 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. */
+ *  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 =>

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')!