Skip to content

fix: show model budget count in provider card header summary #5704

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
akshaydeo merged 2 commits into from
Aug 11, 2026
Merged

fix: show model budget count in provider card header summary #5704

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ad9a9ea
feat: add virtual key per model budget creation on the virtual key sheet
BearTS Jul 30, 2026
1e1f23a
feat: provider config card to show model budgets
BearTS Jul 30, 2026
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
13 changes: 13 additions & 0 deletions framework/configstore/tables/virtualkey.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,19 @@ type TableVirtualKeyProviderConfig struct {
RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID;onDelete:CASCADE" json:"rate_limit,omitempty"`
Budgets []TableBudget `gorm:"foreignKey:ProviderConfigID;constraint:OnDelete:CASCADE" json:"budgets,omitempty"` // Multiple budgets with different reset intervals
Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Empty means all keys allowed for this provider

// ModelBudgets carries per-model budgets/rate-limits under this provider for serialization
// only. They live in VK-scoped model configs (the source of truth), not this table; the
// handler hydrates this field when returning a VK so the sheet can render/edit them.
ModelBudgets []VKProviderModelBudget `gorm:"-" json:"model_budgets,omitempty"`
}

// VKProviderModelBudget is one per-model budget/rate-limit group under a VK provider config,
// used purely for serialization (reverse-mapped from a VK-scoped model config).
type VKProviderModelBudget struct {
ModelName string `json:"model_name"`
Budgets []TableBudget `json:"budgets,omitempty"`
RateLimit *TableRateLimit `json:"rate_limit,omitempty"`
}

// TableName sets the table name for each model
Expand Down
307 changes: 243 additions & 64 deletions transports/bifrost-http/handlers/governance.go
View file
Open in desktop

Large diffs are not rendered by default.

21 changes: 19 additions & 2 deletions transports/bifrost-http/handlers/governance_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,23 @@ func (m *mockRotateConfigStore) GetModelConfig(_ context.Context, scope string,
return lookupVKModelConfig(m.modelConfigs, scope, scopeID, modelName, provider)
}

// GetModelConfigsByScopeAndScopeIDs returns the stored configs matching the scope and scope IDs,
// mirroring the bulk load hydrateVKGovernance performs.
func (m *mockRotateConfigStore) GetModelConfigsByScopeAndScopeIDs(_ context.Context, scope string, scopeIDs []string) ([]configstoreTables.TableModelConfig, error) {
idset := make(map[string]bool, len(scopeIDs))
for _, id := range scopeIDs {
idset[id] = true
}
var out []configstoreTables.TableModelConfig
for _, mc := range m.modelConfigs {
if mc == nil || mc.Scope != scope || mc.ScopeID == nil || !idset[*mc.ScopeID] {
continue
}
out = append(out, *mc)
}
return out, nil
}

type mockRotateGovernanceManager struct {
GovernanceManager
store *mockRotateConfigStore
Expand Down Expand Up @@ -3341,7 +3358,7 @@ func TestApplyVKGovernanceFromModelConfigs_PreservesDirectlyAttachedBudget(t *te
}

// No VK-scoped model config exists for this VK.
applyVKGovernanceFromModelConfigs(vk, map[string]*configstoreTables.TableModelConfig{})
applyVKGovernanceFromModelConfigs(vk, map[string]*configstoreTables.TableModelConfig{}, nil)

if len(vk.Budgets) != 1 || vk.Budgets[0].ID != "bud-direct" {
t.Fatalf("directly attached budget was wiped: got %+v", vk.Budgets)
Expand Down Expand Up @@ -3376,7 +3393,7 @@ func TestApplyVKGovernanceFromModelConfigs_OverlaysModelConfigGovernance(t *test
},
}

applyVKGovernanceFromModelConfigs(vk, byKey)
applyVKGovernanceFromModelConfigs(vk, byKey, nil)

if len(vk.Budgets) != 1 || vk.Budgets[0].ID != "bud-mc" {
t.Fatalf("expected model-config budget overlaid, got %+v", vk.Budgets)
Expand Down
85 changes: 85 additions & 0 deletions ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -414,6 +414,91 @@ export default function VirtualKeyDetailSheet({
</div>
</>
)}

{/* Model Budgets — per-model caps/rate-limits under this provider */}
{config.model_budgets && config.model_budgets.length > 0 && (
<>
<DottedSeparator />
<div className="space-y-3">
<h4 className="text-sm font-medium">Model Budgets</h4>
{config.model_budgets.map((mb, mbIdx) => (
<div key={`${mb.model_name}-${mbIdx}`} className="space-y-3 rounded-md border p-3">
<span className="text-sm font-medium">{mb.model_name}</span>

{/* Budgets */}
{mb.budgets && mb.budgets.length > 0
? mb.budgets.map((b, bIdx) => (
<div key={bIdx} className="space-y-2">
{!isManagedByProfile && b.id ? (
<div className="flex justify-end">
<BudgetOverrideDialog
budget={b}
onSave={(data) => saveBudgetOverride(b.id, data)}
onRemove={() => clearBudgetOverride(b.id)}
disabled={!canUpdateVirtualKeys}
calendarAligned={virtualKey.calendar_aligned}
/>
</div>
) : null}
<UsageLine current={b.current_usage} max={getEffectiveBudgetLimit(b)} format={formatCurrency} />
{hasActiveBudgetOverride(b) ? (
<p className="text-muted-foreground text-xs">
Base {formatCurrency(b.max_limit)} + {formatCurrency(b.override_amount ?? 0)} override
</p>
) : null}
<div className="text-muted-foreground flex items-center justify-between text-xs">
<span>
Resets {parseResetPeriod(b.reset_duration)}
{virtualKey.calendar_aligned && supportsCalendarAlignment(b.reset_duration) && " (calendar)"}
</span>
{b.last_reset ? (
<span>Last reset {formatDistanceToNow(new Date(b.last_reset), { addSuffix: true })}</span>
) : null}
</div>
</div>
))
: null}

{/* Token Limits */}
{mb.rate_limit?.token_max_limit != null ? (
<div className="space-y-2">
<span className="text-muted-foreground text-xs font-medium">TOKEN LIMITS</span>
<UsageLine
current={mb.rate_limit.token_current_usage}
max={mb.rate_limit.token_max_limit}
format={(n) => n.toLocaleString()}
/>
<div className="text-muted-foreground text-xs">
Resets {parseResetPeriod(mb.rate_limit.token_reset_duration || "")}
{virtualKey.calendar_aligned &&
supportsCalendarAlignment(mb.rate_limit.token_reset_duration || "") &&
" (calendar)"}
</div>
</div>
) : null}

{/* Request Limits */}
{mb.rate_limit?.request_max_limit != null ? (
<div className="space-y-2">
<span className="text-muted-foreground text-xs font-medium">REQUEST LIMITS</span>
<UsageLine
current={mb.rate_limit.request_current_usage}
max={mb.rate_limit.request_max_limit}
format={(n) => n.toLocaleString()}
/>
<div className="text-muted-foreground text-xs">
Resets {parseResetPeriod(mb.rate_limit.request_reset_duration || "")}
{virtualKey.calendar_aligned &&
supportsCalendarAlignment(mb.rate_limit.request_reset_duration || "") &&
" (calendar)"}
</div>
</div>
) : null}
</div>
))}
</div>
</>
)}
</div>
</div>
))}
Expand Down
115 changes: 91 additions & 24 deletions ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,31 @@ const providerConfigSchema = z.object({
request_reset_duration: z.string().optional(),
})
.optional(),
// Per-model budgets/rate-limits under this provider
model_budgets: z
.array(
z.object({
model_name: z.string().trim().min(1, "Model name is required"),
budgets: z
.array(
z.object({
id: z.string().optional(),
max_limit: z.number().nonnegative().optional(),
reset_duration: z.string().optional(),
}),
)
.optional(),
rate_limit: z
.object({
token_max_limit: z.number().int().nonnegative().optional(),
token_reset_duration: z.string().optional(),
request_max_limit: z.number().int().nonnegative().optional(),
request_reset_duration: z.string().optional(),
})
.optional(),
}),
)
.optional(),
});

const mcpConfigSchema = z.object({
Expand Down Expand Up @@ -346,6 +371,22 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC
request_reset_duration: config.rate_limit.request_reset_duration,
}
: undefined,
model_budgets: config.model_budgets?.map((mb) => ({
model_name: mb.model_name,
budgets: mb.budgets?.map((b) => ({
id: b.id,
max_limit: b.max_limit,
reset_duration: b.reset_duration,
})),
rate_limit: mb.rate_limit
? {
token_max_limit: mb.rate_limit.token_max_limit ?? undefined,
token_reset_duration: mb.rate_limit.token_reset_duration,
request_max_limit: mb.rate_limit.request_max_limit ?? undefined,
request_reset_duration: mb.rate_limit.request_reset_duration,
}
: undefined,
})),
})) || [],
mcpConfigs:
virtualKey?.mcp_configs?.map((config) => ({
Expand Down Expand Up @@ -534,31 +575,50 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC
form.setValue("requestResetDuration", "1h", { shouldDirty: true });
};

const normalizeProviderConfigs = (configs: typeof providerConfigs, existingConfigs?: VirtualKey["provider_configs"]): any[] => {
return configs.map((config) => ({
...config,
budgets: config.budgets?.filter((b): b is { id?: string; max_limit: number; reset_duration: string } => b.max_limit !== undefined),
weight: config.weight ?? null,
rate_limit: (() => {
const hasTokenMaxLimit = config.rate_limit?.token_max_limit !== undefined;
const hasRequestMaxLimit = config.rate_limit?.request_max_limit !== undefined;
if (hasTokenMaxLimit || hasRequestMaxLimit) {
return {
token_max_limit: config.rate_limit?.token_max_limit ?? null,
token_reset_duration: hasTokenMaxLimit ? config.rate_limit?.token_reset_duration || "1h" : null,
request_max_limit: config.rate_limit?.request_max_limit ?? null,
request_reset_duration: hasRequestMaxLimit ? config.rate_limit?.request_reset_duration || "1h" : null,
};
}

const existingConfig = existingConfigs?.find((item) => (config.id ? item.id === config.id : item.provider === config.provider));
if (existingConfig?.rate_limit) {
return {};
}
// Build a request rate-limit payload from the form's rate-limit fields. Returns the field
// values when a limit is set, {} to clear an existing rate limit (removal), or undefined.
const normalizeRateLimit = (
rl: { token_max_limit?: number; token_reset_duration?: string; request_max_limit?: number; request_reset_duration?: string } | undefined,
hadExisting: boolean,
) => {
const hasToken = rl?.token_max_limit !== undefined;
const hasRequest = rl?.request_max_limit !== undefined;
if (hasToken || hasRequest) {
return {
token_max_limit: rl?.token_max_limit ?? null,
token_reset_duration: hasToken ? rl?.token_reset_duration || "1h" : null,
request_max_limit: rl?.request_max_limit ?? null,
request_reset_duration: hasRequest ? rl?.request_reset_duration || "1h" : null,
};
}
return hadExisting ? {} : undefined;
};

return undefined;
})(),
}));
const normalizeProviderConfigs = (configs: typeof providerConfigs, existingConfigs?: VirtualKey["provider_configs"]): any[] => {
return configs.map((config) => {
const existingConfig = existingConfigs?.find((item) => (config.id ? item.id === config.id : item.provider === config.provider));
return {
...config,
budgets: config.budgets?.filter((b): b is { id?: string; max_limit: number; reset_duration: string } => b.max_limit !== undefined),
weight: config.weight ?? null,
rate_limit: normalizeRateLimit(config.rate_limit, !!existingConfig?.rate_limit),
// Full desired per-model set: drop unfilled models, keep an empty array so the
// backend prunes any per-model budgets removed here.
model_budgets: (config.model_budgets || [])
.filter((mb) => mb.model_name && mb.model_name.trim() !== "")
.map((mb) => {
const existingMB = existingConfig?.model_budgets?.find((m) => m.model_name === mb.model_name.trim());
return {
model_name: mb.model_name.trim(),
budgets: (mb.budgets || []).filter(
(b): b is { id?: string; max_limit: number; reset_duration: string } => b.max_limit !== undefined,
),
rate_limit: normalizeRateLimit(mb.rate_limit, !!existingMB?.rate_limit),
};
})
.filter((mb) => mb.budgets.length > 0 || mb.rate_limit !== undefined),
};
});
};

const parseResetDurationMs = (duration?: string) => {
Expand Down Expand Up @@ -1137,6 +1197,7 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC
providerLabel={providerLabel}
iconProvider={iconProvider}
providerKeys={providerKeys}
showModelBudgets
onRemove={() => handleRemoveProvider(index)}
value={{
providerName: config.provider,
Expand All @@ -1146,6 +1207,11 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC
keyIds: config.key_ids || [],
budgets: config.budgets || [],
rateLimit: config.rate_limit ?? null,
modelBudgets: (config.model_budgets || []).map((mb) => ({
model_name: mb.model_name,
budgets: mb.budgets || [],
rate_limit: mb.rate_limit,
})),
}}
onChange={(next) => {
const updated = [...providerConfigs];
Expand All @@ -1162,6 +1228,7 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC
reset_config: l.reset_config,
})),
rate_limit: next.rateLimit ?? undefined,
model_budgets: next.modelBudgets,
};
form.setValue("providerConfigs", updated, {
shouldDirty: true,
Expand Down
16 changes: 14 additions & 2 deletions ui/components/ui/providerConfigCard.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,17 @@ export function ProviderConfigCard({

const modelBudgets = value.modelBudgets || [];
const capLabel = budgetLinesLabel(value.budgets);
// Header summary: the provider cap and/or a model-budget count, falling back to
// "No budget" only when neither is set — so a provider with only model budgets
// doesn't read as "No budget".
const modelBudgetCount = showModelBudgets ? modelBudgets.length : 0;
const headerSummary =
[
budgetLinesLabel(value.budgets, ""),
modelBudgetCount > 0 ? `${modelBudgetCount} model budget${modelBudgetCount === 1 ? "" : "s"}` : "",
]
.filter(Boolean)
.join(" · ") || "No budget";
const ws = globalProviderCap;

// Key scope handed to ModelMultiselect so model suggestions match the keys
Expand Down Expand Up @@ -216,8 +227,9 @@ export function ProviderConfigCard({
>
<RenderProviderIcon provider={iconProvider} size="sm" className="h-4 w-4 shrink-0" />
<span className="shrink-0 text-sm font-medium whitespace-nowrap">{providerLabel}</span>
<span className="min-w-0 flex-1" />
<span className="text-muted-foreground shrink-0 text-sm whitespace-nowrap">{capLabel}</span>
<span className="text-muted-foreground min-w-0 flex-1 truncate text-right text-sm" title={headerSummary}>
{headerSummary}
</span>
<button
type="button"
onClick={(e) => {
Expand Down
18 changes: 18 additions & 0 deletions ui/lib/types/governance.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,13 @@ export interface VirtualKey {
config_hash?: string; // Present when config is synced from config.json
}

// Per-model budgets/rate-limits under a provider config, surfaced on the VK for display/edit.
export interface VirtualKeyModelBudget {
model_name: string;
budgets?: Budget[];
rate_limit?: RateLimit;
}

export interface VirtualKeyProviderConfig {
id?: number;
provider: string;
Expand All @@ -129,6 +136,7 @@ export interface VirtualKeyProviderConfig {
allow_all_keys: boolean; // True means all keys allowed; false with empty keys means no keys allowed
budgets?: Budget[];
rate_limit?: RateLimit;
model_budgets?: VirtualKeyModelBudget[]; // Per-model budgets/rate-limits under this provider
keys?: DBKey[]; // Associated database keys for this provider (only used when allow_all_keys is false)
}

Expand Down Expand Up @@ -165,6 +173,14 @@ export interface UsageStats {
requests_last_reset: string;
}

// One per-model budget/rate-limit group in a provider-config request. model_name must be a
// concrete model (not the "*" wildcard, which is the provider-level tier).
export interface VirtualKeyModelBudgetRequest {
model_name: string;
budgets?: CreateBudgetRequest[];
rate_limit?: CreateRateLimitRequest;
}

// Request interfaces for provider config operations
export interface VirtualKeyProviderConfigRequest {
provider: string;
Expand All @@ -173,6 +189,7 @@ export interface VirtualKeyProviderConfigRequest {
blacklisted_models?: string[];
budgets?: CreateBudgetRequest[];
rate_limit?: CreateRateLimitRequest;
model_budgets?: VirtualKeyModelBudgetRequest[];
key_ids?: string[]; // List of DBKey UUIDs to associate with this provider config
}

Expand All @@ -184,6 +201,7 @@ export interface VirtualKeyProviderConfigUpdateRequest {
blacklisted_models?: string[];
budgets?: CreateBudgetRequest[];
rate_limit?: UpdateRateLimitRequest;
model_budgets?: VirtualKeyModelBudgetRequest[]; // Full desired per-model set when provider_configs is supplied
key_ids?: string[]; // List of DBKey UUIDs to associate with this provider config
}

Expand Down
Loading
Loading