Skip to content

Commit 485b618

Browse files
jahoomaclaude
andcommitted
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>
1 parent 9702f36 commit 485b618

7 files changed

Lines changed: 56 additions & 22 deletions

File tree

cli/src/components/freebuff-model-selector.tsx

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
isFreebuffModelAvailable,
1212
isFreebuffPremiumModelId,
1313
} from '@codebuff/common/constants/freebuff-models'
14+
import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
1415

1516
import { joinFreebuffQueue } from '../hooks/use-freebuff-session'
1617
import { useNow } from '../hooks/use-now'
@@ -127,10 +128,7 @@ export const FreebuffModelSelector: React.FC = () => {
127128
}, [now, selectedModel, session, setSelectedModel])
128129

129130
const committedModelId = session?.status === 'queued' ? session.model : null
130-
const rateLimitsByModel =
131-
session && 'rateLimitsByModel' in session
132-
? session.rateLimitsByModel
133-
: undefined
131+
const rateLimitsByModel = getRateLimitsByModel(session)
134132

135133
const BUTTON_CHROME = 4 // 2 border + 2 padding
136134
const NAME_GAP = 2 // spaces between name column and details column

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,11 @@ import {
1010
} from '../hooks/use-freebuff-session'
1111
import { useTheme } from '../hooks/use-theme'
1212
import { useFreebuffSessionStore } from '../state/freebuff-session-store'
13+
import { formatSessionUnits } from '../utils/format-session-units'
1314
import { BORDER_CHARS } from '../utils/ui-constants'
1415

1516
import type { KeyEvent } from '@opentui/core'
1617

17-
const formatSessionUnits = (units: number): string =>
18-
Number.isInteger(units) ? String(units) : units.toFixed(1)
19-
2018
interface SessionEndedBannerProps {
2119
/** True while an agent request is still streaming under the server-side
2220
* grace window. Swaps the Enter-to-rejoin affordance for a "let it
@@ -43,10 +41,6 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
4341
const premiumQuota = useFreebuffSessionStore(
4442
(s) => Object.values(getRateLimitsByModel(s.session) ?? {})[0] ?? null,
4543
)
46-
const quotaColor =
47-
premiumQuota && premiumQuota.recentCount >= premiumQuota.limit
48-
? theme.secondary
49-
: theme.muted
5044

5145
// While a request is still streaming, restart is disabled: it would
5246
// unmount <Chat> and abort the in-flight agent run. The promise is "we
@@ -112,7 +106,13 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
112106
<text style={{ fg: theme.foreground, wrapMode: 'word' }}>
113107
Your freebuff session has ended.
114108
{premiumQuota && (
115-
<span fg={quotaColor}>
109+
<span
110+
fg={
111+
premiumQuota.recentCount >= premiumQuota.limit
112+
? theme.secondary
113+
: theme.muted
114+
}
115+
>
116116
{' · '}
117117
{formatSessionUnits(premiumQuota.recentCount)} of{' '}
118118
{premiumQuota.limit} premium sessions used today

cli/src/components/waiting-room-screen.tsx

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,10 @@ import { useSheenAnimation } from '../hooks/use-sheen-animation'
1515
import { useTerminalDimensions } from '../hooks/use-terminal-dimensions'
1616
import { useTheme } from '../hooks/use-theme'
1717
import { exitFreebuffCleanly } from '../utils/freebuff-exit'
18+
import { formatSessionUnits } from '../utils/format-session-units'
1819
import { getLogoAccentColor, getLogoBlockColor } from '../utils/theme-system'
1920
import { FREEBUFF_PREMIUM_SESSION_LIMIT } from '@codebuff/common/constants/freebuff-models'
21+
import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
2022

2123
import type { FreebuffSessionResponse } from '../types/freebuff-session'
2224
import type { FreebuffIpPrivacySignal } from '@codebuff/common/types/freebuff-session'
@@ -59,9 +61,6 @@ const formatRetryAfter = (ms: number): string => {
5961
return rem === 0 ? `${hours}h` : `${hours}h ${rem}m`
6062
}
6163

62-
const formatSessionUnits = (units: number): string =>
63-
Number.isInteger(units) ? String(units) : units.toFixed(1)
64-
6564
const PRIVACY_SIGNAL_LABELS: Partial<Record<FreebuffIpPrivacySignal, string>> =
6665
{
6766
anonymous: 'anonymized network',
@@ -268,10 +267,7 @@ export const WaitingRoomScreen: React.FC<WaitingRoomScreenProps> = ({
268267
// pool; the server replicates the same snapshot under each premium model
269268
// id, so any entry has the right count. Renders amber when exhausted so
270269
// the limit reads as "you've hit it" rather than just another count.
271-
const rateLimitsByModel =
272-
session && 'rateLimitsByModel' in session
273-
? session.rateLimitsByModel
274-
: undefined
270+
const rateLimitsByModel = getRateLimitsByModel(session)
275271
const sharedPremiumUsed = rateLimitsByModel
276272
? (Object.values(rateLimitsByModel)[0]?.recentCount ?? 0)
277273
: 0

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import {
33
FALLBACK_FREEBUFF_MODEL_ID,
44
resolveFreebuffModel,
55
} from '@codebuff/common/constants/freebuff-models'
6+
import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
67
import { useEffect } from 'react'
78

89
import {
@@ -19,8 +20,6 @@ import {
1920
import { logger } from '../utils/logger'
2021
import { saveFreebuffModelPreference } from '../utils/settings'
2122

22-
import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
23-
2423
import type { FreebuffSessionResponse } from '../types/freebuff-session'
2524
import type {
2625
FreebuffCountryBlockReason,
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
/** Premium-session counts come back from the server as `recentCount` units
2+
* that may be fractional (a long agent run can consume 1.3 sessions). Render
3+
* integers without a trailing `.0`, fractionals at one decimal — matches the
4+
* `limit` field which is always integer. */
5+
export const formatSessionUnits = (units: number): string =>
6+
Number.isInteger(units) ? String(units) : units.toFixed(1)

common/src/types/freebuff-session.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,10 @@ export type FreebuffSessionRateLimitByModel = Record<
3333

3434
/** Pull the per-model premium quota snapshot off whichever session statuses
3535
* carry it (queued, active, ended, none). Returns undefined for terminal /
36-
* pre-join states that have no quota field. */
36+
* pre-join states that have no quota field. The parameter is intentionally
37+
* loose so the CLI can pass its `FreebuffSessionResponse` (which adds the
38+
* client-only `takeover_prompt` variant) without a discriminated-union
39+
* ceremony at every call site. */
3740
export const getRateLimitsByModel = (
3841
session: { status: string } | null | undefined,
3942
): FreebuffSessionRateLimitByModel | undefined =>

web/src/server/free-session/__tests__/public-api.test.ts

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -960,6 +960,38 @@ describe('getSessionState', () => {
960960
expect(state.gracePeriodRemainingMs).toBe(GRACE_MS - 60_000)
961961
})
962962

963+
test('ended view carries the full premium-quota snapshot', async () => {
964+
// The post-session banner reads any entry from rateLimitsByModel since
965+
// all premium models share one daily pool. Unlike queued/active, the
966+
// ended view ships the full unfiltered map so a single banner read is
967+
// always safe.
968+
await requestSession({ userId: 'u1', model: DEFAULT_MODEL, deps })
969+
const row = deps.rows.get('u1')!
970+
row.status = 'active'
971+
row.admitted_at = new Date(deps._now().getTime() - SESSION_LEN - 60_000)
972+
row.expires_at = new Date(deps._now().getTime() - 60_000)
973+
deps.admits.push({
974+
user_id: 'u1',
975+
model: FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID,
976+
admitted_at: new Date(deps._now().getTime() - 30 * 60_000),
977+
})
978+
979+
const state = await getSessionState({
980+
userId: 'u1',
981+
claimedInstanceId: row.active_instance_id,
982+
deps,
983+
})
984+
if (state.status !== 'ended') throw new Error('unreachable')
985+
expect(
986+
state.rateLimitsByModel?.[FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID],
987+
).toEqual(expectedRateLimit(FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID, 1))
988+
// Every premium model is present (sharing the same recentCount) so the
989+
// banner can read any entry without caring which model the user was on.
990+
expect(state.rateLimitsByModel?.[FREEBUFF_KIMI_MODEL_ID]).toEqual(
991+
expectedRateLimit(FREEBUFF_KIMI_MODEL_ID, 1),
992+
)
993+
})
994+
963995
test('row past grace window returns none', async () => {
964996
await requestSession({ userId: 'u1', model: DEFAULT_MODEL, deps })
965997
const row = deps.rows.get('u1')!

0 commit comments

Comments
 (0)