Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 63 additions & 36 deletions cli/src/components/freebuff-model-selector.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react'
import { Button } from './button'
import {
FALLBACK_FREEBUFF_MODEL_ID,
FREEBUFF_MODELS,
getFreebuffDeploymentAvailabilityLabel,
getFreebuffModelsForAccessTier,
isFreebuffModelAvailable,
isFreebuffPremiumModelId,
} from '@codebuff/common/constants/freebuff-models'
Expand All @@ -26,39 +26,18 @@ import {
import type { FreebuffModelOption } from '@codebuff/common/constants/freebuff-models'
import type { KeyEvent } from '@opentui/core'

const FREEBUFF_MODEL_IDS = FREEBUFF_MODELS.map((m) => m.id)

// Section grouping: premium models share one quota pool, unlimited has none.
// Putting the tier on a section header lets each row drop its redundant
// "Premium"/"Unlimited" chip. The shared 0/5 counter lives in the page title
// (rendered by the parent), not the section header — this picker is purely a
// list of choices grouped by tier. Empty sections are filtered so a model set
// with no premium (or no unlimited) entries doesn't render an orphan header.
type Section = {
key: 'premium' | 'unlimited'
key: 'premium' | 'unlimited' | 'limited'
label: string
models: readonly FreebuffModelOption[]
}

const SECTIONS: readonly Section[] = (
[
{
key: 'premium',
label: 'PREMIUM',
models: FREEBUFF_MODELS.filter((m) =>
isFreebuffPremiumModelId(m.id),
),
},
{
key: 'unlimited',
label: 'UNLIMITED',
models: FREEBUFF_MODELS.filter(
(m) => !isFreebuffPremiumModelId(m.id),
),
},
] satisfies readonly Section[]
).filter((section) => section.models.length > 0)

/**
* Dual-purpose model picker:
* - Pre-chat landing (session 'none'): user hasn't joined any queue. Picking
Expand Down Expand Up @@ -86,6 +65,8 @@ export const FreebuffModelSelector: React.FC = () => {
const selectedModel = useFreebuffModelStore((s) => s.selectedModel)
const setSelectedModel = useFreebuffModelStore((s) => s.setSelectedModel)
const session = useFreebuffSessionStore((s) => s.session)
const accessTier =
session && 'accessTier' in session ? session.accessTier : 'full'
const now = useNow(60_000)
const deploymentAvailabilityLabel = useMemo(
() => getFreebuffDeploymentAvailabilityLabel(new Date(now)),
Expand All @@ -98,9 +79,48 @@ export const FreebuffModelSelector: React.FC = () => {
// selected model whenever the selection changes (after a successful switch
// or an external selectedModel update).
const [focusedId, setFocusedId] = useState<string>(selectedModel)
const availableModels = useMemo(
() => getFreebuffModelsForAccessTier(accessTier),
[accessTier],
)
const availableModelIds = useMemo(
() => availableModels.map((m) => m.id),
[availableModels],
)
const sections = useMemo(() => {
if (accessTier === 'limited') {
return [
{
key: 'limited',
label: 'LIMITED',
models: availableModels,
},
] satisfies readonly Section[]
}
return (
[
{
key: 'premium',
label: 'PREMIUM',
models: availableModels.filter((m) => isFreebuffPremiumModelId(m.id)),
},
{
key: 'unlimited',
label: 'UNLIMITED',
models: availableModels.filter(
(m) => !isFreebuffPremiumModelId(m.id),
),
},
] satisfies readonly Section[]
).filter((section) => section.models.length > 0)
}, [accessTier, availableModels])
useEffect(() => {
setFocusedId(selectedModel)
}, [selectedModel])
setFocusedId(
availableModelIds.includes(selectedModel)
? selectedModel
: availableModelIds[0]!,
)
}, [availableModelIds, selectedModel])

useEffect(() => {
// Landing-screen safety net: if the in-memory selection becomes
Expand All @@ -110,11 +130,12 @@ export const FreebuffModelSelector: React.FC = () => {
// preference (e.g. Kimi or DeepSeek) is preserved for the next launch.
if (
(session?.status === 'none' || !session) &&
!isFreebuffModelAvailable(selectedModel, new Date(now))
(!availableModelIds.includes(selectedModel) ||
!isFreebuffModelAvailable(selectedModel, new Date(now)))
) {
setSelectedModel(FALLBACK_FREEBUFF_MODEL_ID)
setSelectedModel(availableModelIds[0] ?? FALLBACK_FREEBUFF_MODEL_ID)
}
}, [now, selectedModel, session, setSelectedModel])
}, [availableModelIds, now, selectedModel, session, setSelectedModel])

const committedModelId = session?.status === 'queued' ? session.model : null
const rateLimitsByModel = getRateLimitsByModel(session)
Expand All @@ -128,7 +149,7 @@ export const FreebuffModelSelector: React.FC = () => {
// terminals where the secondary details spill to an indented second line.
const { wrapDetails, buttonOuterWidth, nameColumnWidth } = useMemo(() => {
const nameLen = (m: FreebuffModelOption) => m.displayName.length
const maxNameLen = Math.max(...FREEBUFF_MODELS.map(nameLen))
const maxNameLen = Math.max(...availableModels.map(nameLen))

const detailsParts = (model: FreebuffModelOption): number[] => {
const parts = [model.tagline.length]
Expand All @@ -149,8 +170,7 @@ export const FreebuffModelSelector: React.FC = () => {
joinedLen(detailsParts(model))

const maxOneLineOuter =
Math.max(...FREEBUFF_MODELS.map(oneLineLen)) +
BUTTON_CHROME
Math.max(...availableModels.map(oneLineLen)) + BUTTON_CHROME
if (maxOneLineOuter <= contentMaxWidth) {
return {
wrapDetails: false,
Expand All @@ -173,7 +193,7 @@ export const FreebuffModelSelector: React.FC = () => {
return parts.length === 0 ? 0 : 2 /* indent */ + joinedLen(parts)
}
const maxTwoLineInner = Math.max(
...FREEBUFF_MODELS.map((m) =>
...availableModels.map((m) =>
Math.max(labelLineLen(m), detailsLineLen(m)),
),
)
Expand All @@ -185,7 +205,7 @@ export const FreebuffModelSelector: React.FC = () => {
),
nameColumnWidth: maxNameLen,
}
}, [contentMaxWidth, deploymentAvailabilityLabel])
}, [availableModels, contentMaxWidth, deploymentAvailabilityLabel])

const isJoinable = useCallback(
(modelId: string) => {
Expand Down Expand Up @@ -228,7 +248,7 @@ export const FreebuffModelSelector: React.FC = () => {
}
if (!direction) return
const targetId = nextFreebuffModelId({
modelIds: FREEBUFF_MODEL_IDS,
modelIds: availableModelIds,
focusedId,
direction,
})
Expand All @@ -238,7 +258,14 @@ export const FreebuffModelSelector: React.FC = () => {
setFocusedId(targetId)
}
},
[pending, pick, focusedId, committedModelId, isJoinable],
[
pending,
pick,
focusedId,
committedModelId,
isJoinable,
availableModelIds,
],
),
)

Expand Down Expand Up @@ -345,7 +372,7 @@ export const FreebuffModelSelector: React.FC = () => {
gap: 0,
}}
>
{SECTIONS.map((section, sectionIdx) => (
{sections.map((section, sectionIdx) => (
<box
key={section.key}
style={{
Expand Down
7 changes: 6 additions & 1 deletion cli/src/components/session-ended-banner.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,13 @@ export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
const isQuotaExhausted = premiumQuota
? premiumQuota.recentCount >= premiumQuota.limit
: false
const accessTier = useFreebuffSessionStore((s) =>
s.session && 'accessTier' in s.session ? s.session.accessTier : 'full',
)
const quotaLabel =
accessTier === 'limited' ? 'limited sessions' : 'premium sessions'
const bannerTitle = premiumQuota
? `Session ended · ${formatSessionUnits(premiumQuota.recentCount)} of ${premiumQuota.limit} premium sessions used today`
? `Session ended · ${formatSessionUnits(premiumQuota.recentCount)} of ${premiumQuota.limit} ${quotaLabel} used today`
: 'Session ended'

// While a request is still streaming, restart is disabled: it would
Expand Down
34 changes: 25 additions & 9 deletions cli/src/components/waiting-room-screen.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,10 @@ import {
} from '../utils/freebuff-premium-reset'
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 {
FREEBUFF_LIMITED_SESSION_LIMIT,
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'
Expand Down Expand Up @@ -255,6 +258,8 @@ export const WaitingRoomScreen: React.FC<WaitingRoomScreenProps> = ({
const [exitHover, setExitHover] = useState(false)

const isQueued = session?.status === 'queued'
const accessTier =
session && 'accessTier' in session ? session.accessTier : 'full'
// 'none' = user hasn't joined any queue yet. We're in the pre-chat landing
// state: show the picker with live N-in-line hints and a prompt. Picking a
// model triggers joinFreebuffQueue, which POSTs and transitions us to
Expand All @@ -280,14 +285,22 @@ export const WaitingRoomScreen: React.FC<WaitingRoomScreenProps> = ({
: undefined
const sharedPremiumUsed = premiumRateLimit?.recentCount ?? 0
const isPremiumExhausted =
sharedPremiumUsed >= FREEBUFF_PREMIUM_SESSION_LIMIT
sharedPremiumUsed >=
(accessTier === 'limited'
? FREEBUFF_LIMITED_SESSION_LIMIT
: FREEBUFF_PREMIUM_SESSION_LIMIT)
const premiumUsedColor = isPremiumExhausted ? theme.secondary : theme.muted
// Pad the used count so the title's centered container doesn't shift width
// as the count ticks from "0" → "1.3" → "2" while loading.
const sessionUnitWidth = String(FREEBUFF_PREMIUM_SESSION_LIMIT).length + 2
const formattedSharedPremiumUsed = formatSessionUnits(
sharedPremiumUsed,
).padStart(sessionUnitWidth)
const sessionLimit =
accessTier === 'limited'
? FREEBUFF_LIMITED_SESSION_LIMIT
: FREEBUFF_PREMIUM_SESSION_LIMIT
const sessionLabel =
accessTier === 'limited' ? 'limited sessions' : 'premium sessions'
const sessionUnitWidth = String(sessionLimit).length + 2
const formattedSharedPremiumUsed =
formatSessionUnits(sharedPremiumUsed).padStart(sessionUnitWidth)
const premiumResetAt = getFreebuffPremiumResetAt({
rateLimitsByModel,
nowMs: now,
Expand Down Expand Up @@ -399,8 +412,8 @@ export const WaitingRoomScreen: React.FC<WaitingRoomScreenProps> = ({
style={{ fg: theme.muted, marginBottom: 1, wrapMode: 'word' }}
>
<span fg={premiumUsedColor}>
{formattedSharedPremiumUsed} of{' '}
{FREEBUFF_PREMIUM_SESSION_LIMIT} premium sessions used
{formattedSharedPremiumUsed} of {sessionLimit} {sessionLabel}{' '}
used
</span>
<span fg={theme.muted}>
{' · '}
Expand Down Expand Up @@ -540,7 +553,10 @@ export const WaitingRoomScreen: React.FC<WaitingRoomScreenProps> = ({
<span fg={theme.foreground}>
{formatSessionUnits(session.recentCount)} of {session.limit}
</span>{' '}
premium sessions today. Try again in{' '}
{session.accessTier === 'limited'
? 'limited sessions'
: 'premium sessions'}{' '}
today. Try again in{' '}
<span fg={theme.foreground}>
{formatRetryAfter(session.retryAfterMs)}
</span>
Expand Down
24 changes: 16 additions & 8 deletions cli/src/hooks/helpers/send-message.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,9 @@ export type ResetEarlyReturnStateParams = {
isQueuePausedRef?: MutableRefObject<boolean>
}

export const resetEarlyReturnState = (params: ResetEarlyReturnStateParams): void => {
export const resetEarlyReturnState = (
params: ResetEarlyReturnStateParams,
): void => {
const {
setCanProcessQueue,
updateChainInProgress,
Expand Down Expand Up @@ -186,11 +188,12 @@ export const prepareUserMessage = async (params: {
}
}

const { attachments: imageAttachments, messageContent } = await processImagesForMessage({
content: finalContent,
pendingImages,
projectRoot: getProjectRoot(),
})
const { attachments: imageAttachments, messageContent } =
await processImagesForMessage({
content: finalContent,
pendingImages,
projectRoot: getProjectRoot(),
})

const shouldInsertDivider =
lastMessageMode === null || lastMessageMode !== agentMode
Expand All @@ -214,7 +217,12 @@ export const prepareUserMessage = async (params: {
}))

// Pass original content (not finalContent) for display, but finalContent goes to agent
const userMessage = getUserMessage(content, imageAttachments, textAttachmentsForMessage, fileAttachmentsForMessage)
const userMessage = getUserMessage(
content,
imageAttachments,
textAttachmentsForMessage,
fileAttachmentsForMessage,
)
const userMessageId = userMessage.id
if (imageAttachments.length > 0) {
userMessage.attachments = imageAttachments
Expand Down Expand Up @@ -381,7 +389,6 @@ export const handleRunCompletion = (params: {
}

if (output.type === 'error') {

if (isOutOfCreditsError(output)) {
updater.setError(OUT_OF_CREDITS_MESSAGE)
useChatStore.getState().setInputMode('outOfCredits')
Expand Down Expand Up @@ -527,6 +534,7 @@ function handleFreebuffGateError(
switch (kind) {
case 'session_expired':
case 'waiting_room_required':
case 'session_model_mismatch':
// Our seat is gone mid-chat. Finalize the AI message so its streaming
// indicator stops — otherwise `isComplete` stays false and the message
// keeps rendering a blinking cursor forever, making the user think the
Expand Down
Loading
Loading