diff --git a/cli/src/components/freebuff-model-selector.tsx b/cli/src/components/freebuff-model-selector.tsx index edc889b1c..294a4b32f 100644 --- 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 index 7482cbdf5..278729f95 100644 --- 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 = ({ '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 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 = ({ return ( = ({ gap: 0, }} > - - Your freebuff session has ended. - {isStreaming ? ( Agent is wrapping up. Rejoin the wait room after it's finished. @@ -115,7 +131,7 @@ export const SessionEndedBanner: React.FC = ({ fg: pendingAction === 'same-chat' ? theme.muted - : theme.primary, + : theme.foreground, }} attributes={TextAttributes.BOLD} > @@ -144,11 +160,14 @@ export const SessionEndedBanner: React.FC = ({ ? theme.muted : theme.foreground, }} - attributes={TextAttributes.BOLD} > - {pendingAction === 'waiting-room' - ? 'Opening model selection…' - : 'Change model (ESC)'} + {pendingAction === 'waiting-room' ? ( + 'Opening model selection…' + ) : ( + <> + Change model{' Esc'} + + )} diff --git a/cli/src/components/waiting-room-screen.tsx b/cli/src/components/waiting-room-screen.tsx index e86b536ed..a07971cab 100644 --- 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> = { anonymous: 'anonymized network', @@ -268,10 +267,7 @@ export const WaitingRoomScreen: React.FC = ({ // 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 index baa8a2b13..3211acb7a 100644 --- 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 new file mode 100644 index 000000000..75532df80 --- /dev/null +++ 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 index 8d4eebd36..9dbf19149 100644 --- 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 index 2ac2ad75a..351e17ac0 100644 --- 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 index 59af4db81..68a0f59bc 100644 --- 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 { - 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,