Skip to content

Commit 9702f36

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

4 files changed

Lines changed: 73 additions & 12 deletions

File tree

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
12
import { TextAttributes } from '@opentui/core'
23
import { useKeyboard } from '@opentui/react'
34
import React, { useCallback, useState } from 'react'
@@ -8,10 +9,14 @@ import {
89
returnToFreebuffLanding,
910
} from '../hooks/use-freebuff-session'
1011
import { useTheme } from '../hooks/use-theme'
12+
import { useFreebuffSessionStore } from '../state/freebuff-session-store'
1113
import { BORDER_CHARS } from '../utils/ui-constants'
1214

1315
import type { KeyEvent } from '@opentui/core'
1416

17+
const formatSessionUnits = (units: number): string =>
18+
Number.isInteger(units) ? String(units) : units.toFixed(1)
19+
1520
interface SessionEndedBannerProps {
1621
/** True while an agent request is still streaming under the server-side
1722
* grace window. Swaps the Enter-to-rejoin affordance for a "let it
@@ -32,6 +37,17 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
3237
'waiting-room' | 'same-chat' | null
3338
>(null)
3439

40+
// All premium models share one daily pool; the server replicates the same
41+
// snapshot under each premium model id, so the first entry has the right
42+
// count.
43+
const premiumQuota = useFreebuffSessionStore(
44+
(s) => Object.values(getRateLimitsByModel(s.session) ?? {})[0] ?? null,
45+
)
46+
const quotaColor =
47+
premiumQuota && premiumQuota.recentCount >= premiumQuota.limit
48+
? theme.secondary
49+
: theme.muted
50+
3551
// While a request is still streaming, restart is disabled: it would
3652
// unmount <Chat> and abort the in-flight agent run. The promise is "we
3753
// let the agent finish" — honoring that means Enter does nothing until
@@ -95,6 +111,13 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
95111
>
96112
<text style={{ fg: theme.foreground, wrapMode: 'word' }}>
97113
Your freebuff session has ended.
114+
{premiumQuota && (
115+
<span fg={quotaColor}>
116+
{' · '}
117+
{formatSessionUnits(premiumQuota.recentCount)} of{' '}
118+
{premiumQuota.limit} premium sessions used today
119+
</span>
120+
)}
98121
</text>
99122
{isStreaming ? (
100123
<text style={{ fg: theme.muted, wrapMode: 'word' }}>

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

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ import {
1919
import { logger } from '../utils/logger'
2020
import { saveFreebuffModelPreference } from '../utils/settings'
2121

22+
import { getRateLimitsByModel } from '@codebuff/common/types/freebuff-session'
23+
2224
import type { FreebuffSessionResponse } from '../types/freebuff-session'
2325
import type {
2426
FreebuffCountryBlockReason,
@@ -351,11 +353,16 @@ export function markFreebuffSessionCountryBlocked(params: {
351353
}
352354

353355
/** Flip into the local `ended` state without an instanceId (server has lost
354-
* our row). The chat surface stays mounted with the rejoin banner. */
356+
* our row). The chat surface stays mounted with the rejoin banner.
357+
* Preserves any `rateLimitsByModel` snapshot from the prior session so the
358+
* banner can show today's premium-session count without an extra fetch. */
355359
export function markFreebuffSessionEnded(): void {
356360
if (!IS_FREEBUFF) return
357361
controller?.abort()
358-
controller?.apply({ status: 'ended' })
362+
const rateLimitsByModel = getRateLimitsByModel(
363+
useFreebuffSessionStore.getState().session,
364+
)
365+
controller?.apply({ status: 'ended', rateLimitsByModel })
359366
}
360367

361368
interface UseFreebuffSessionResult {
@@ -508,12 +515,18 @@ export function useFreebuffSession(): UseFreebuffSessionResult {
508515
// active|ended → none means we've passed the server's hard cutoff.
509516
// Synthesize a no-instanceId ended state so the chat surface stays
510517
// mounted with the Enter-to-rejoin banner instead of looping back
511-
// through the waiting room.
518+
// through the waiting room. Carry forward whichever rate-limit
519+
// snapshot we have — preferring the fresh `none` snapshot, falling
520+
// back to whatever was on the prior active/ended row — so the
521+
// banner's "N of M used today" line stays populated.
512522
if (
513523
(previousStatus === 'active' || previousStatus === 'ended') &&
514524
next.status === 'none'
515525
) {
516-
apply({ status: 'ended' })
526+
const rateLimitsByModel =
527+
next.rateLimitsByModel ??
528+
getRateLimitsByModel(useFreebuffSessionStore.getState().session)
529+
apply({ status: 'ended', rateLimitsByModel })
517530
return
518531
}
519532

common/src/types/freebuff-session.ts

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,17 @@ export type FreebuffSessionRateLimitByModel = Record<
3131
FreebuffSessionRateLimit
3232
>
3333

34+
/** Pull the per-model premium quota snapshot off whichever session statuses
35+
* carry it (queued, active, ended, none). Returns undefined for terminal /
36+
* pre-join states that have no quota field. */
37+
export const getRateLimitsByModel = (
38+
session: { status: string } | null | undefined,
39+
): FreebuffSessionRateLimitByModel | undefined =>
40+
session && 'rateLimitsByModel' in session
41+
? (session as { rateLimitsByModel?: FreebuffSessionRateLimitByModel })
42+
.rateLimitsByModel
43+
: undefined
44+
3445
export type FreebuffCountryBlockReason =
3546
| 'country_not_allowed'
3647
| 'anonymized_or_unknown_country'
@@ -119,6 +130,10 @@ export type FreebuffSessionServerResponse =
119130
expiresAt?: string
120131
gracePeriodEndsAt?: string
121132
gracePeriodRemainingMs?: number
133+
/** Snapshot of the user's premium-session quota at the moment the
134+
* session ended. Lets the post-session banner show "N of M premium
135+
* sessions used today" without an extra round-trip. */
136+
rateLimitsByModel?: FreebuffSessionRateLimitByModel
122137
}
123138
| {
124139
/** Another CLI on the same account rotated our instance id. Polling

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

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -416,21 +416,31 @@ export async function requestSession(params: {
416416
return attachRateLimit(params.userId, view, deps)
417417
}
418418

419-
/** Thread the current quota snapshot onto queued/active views so the CLI can
420-
* render "N of M sessions used". Other statuses pass through unchanged.
421-
* Called on both POST and GET so the line stays live across polls. */
419+
/** Thread the current quota snapshot onto queued/active/ended views so the
420+
* CLI can render "N of M sessions used" — both during the session and on
421+
* the post-session banner. Other statuses pass through unchanged. Called on
422+
* both POST and GET so the line stays live across polls. */
422423
async function attachRateLimit(
423424
userId: string,
424425
view: SessionStateResponse,
425426
deps: SessionDeps,
426427
): Promise<SessionStateResponse> {
427-
if (view.status !== 'queued' && view.status !== 'active') return view
428-
if (view.status === 'active') {
429-
const snapshot = await fetchRateLimitSnapshot(userId, view.model, deps)
430-
return snapshot ? { ...view, rateLimit: snapshot.info } : view
428+
if (
429+
view.status !== 'queued' &&
430+
view.status !== 'active' &&
431+
view.status !== 'ended'
432+
) {
433+
return view
431434
}
432-
433435
const allRateLimitsByModel = await fetchRateLimitsByModel(userId, deps)
436+
// The ended view doesn't carry a model id, so it gets the full snapshot
437+
// unfiltered — the banner reads any entry's recentCount (they all share the
438+
// same daily premium pool). Queued/active filter out unused models so the
439+
// landing screen and waiting-room title don't list every premium model with
440+
// a "0 used today" hint.
441+
if (view.status === 'ended') {
442+
return { ...view, rateLimitsByModel: allRateLimitsByModel }
443+
}
434444
const rateLimit = allRateLimitsByModel[view.model]
435445
return {
436446
...view,

0 commit comments

Comments
 (0)