diff --git a/framework/configstore/clientconfig.go b/framework/configstore/clientconfig.go index 90a27db25e..2b6a039552 100644 --- a/framework/configstore/clientconfig.go +++ b/framework/configstore/clientconfig.go @@ -910,9 +910,21 @@ func GenerateTeamHash(t tables.TableTeam) (string, error) { hash.Write([]byte("customerID:" + *t.CustomerID)) } - // Hash BudgetID - if t.BudgetID != nil { - hash.Write([]byte("budgetID:" + *t.BudgetID)) + // Hash sorted budget IDs — team now owns multiple budgets; slice order must not + // affect the hash, otherwise config-sync would flip the hash on every reload. + if len(t.Budgets) > 0 { + ids := make([]string, len(t.Budgets)) + for i, b := range t.Budgets { + ids[i] = b.ID + } + sort.Strings(ids) + hash.Write([]byte("budgetIDs:")) + for i, id := range ids { + if i > 0 { + hash.Write([]byte{','}) + } + hash.Write([]byte(id)) + } } // Hash Profile - use Profile if set, else marshal ParsedProfile diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 05f33a3bec..303a3cc8c3 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -417,6 +417,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrateCalendarAlignedToBudgetsAndRateLimitsTable(ctx, db); err != nil { return err } + if err := migrationAddTeamBudgetsToBudgetsTable(ctx, db); err != nil { + return err + } return nil } @@ -6081,6 +6084,133 @@ func migrationAddMultiBudgetTables(ctx context.Context, db *gorm.DB) error { return nil } +// migrationAddTeamBudgetsToBudgetsTable pivots team budgets from a single-FK on +// governance_teams.budget_id to multi-budget ownership via governance_budgets.team_id, +// mirroring how VK/ProviderConfig budgets were restructured in migrationAddMultiBudgetTables. +func migrationAddTeamBudgetsToBudgetsTable(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_team_budgets_to_budgets_table", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + + // Add team_id FK column on governance_budgets + if !mg.HasColumn(&tables.TableBudget{}, "team_id") { + if err := mg.AddColumn(&tables.TableBudget{}, "TeamID"); err != nil { + return fmt.Errorf("failed to add team_id column to governance_budgets: %w", err) + } + } + + // Create index on the new FK column (AddColumn doesn't create indexes from struct tags) + if !mg.HasIndex(&tables.TableBudget{}, "idx_governance_budgets_team_id") { + if err := mg.CreateIndex(&tables.TableBudget{}, "TeamID"); err != nil { + return fmt.Errorf("failed to create index on governance_budgets.team_id: %w", err) + } + } + + // Create FK constraint with CASCADE delete (defined on TableTeam.Budgets) + if !mg.HasConstraint(&tables.TableTeam{}, "Budgets") { + if err := mg.CreateConstraint(&tables.TableTeam{}, "Budgets"); err != nil { + return fmt.Errorf("failed to create FK constraint for Team -> Budgets: %w", err) + } + } + + // Backfill: set team_id from legacy governance_teams.budget_id (if column still exists) + if mg.HasColumn(&tables.TableTeam{}, "budget_id") { + // Preflight: raw SQL below bypasses TableBudget.BeforeSave (which now + // enforces exactly-one-of {TeamID, VirtualKeyID, ProviderConfigID}). + // Fail fast if any team-referenced budget is already owned by a VK or + // ProviderConfig, rather than silently producing a multi-owner row + // that would later be rejected by the hook on its next update. + var conflictCount int64 + if err := tx.Raw(` + SELECT COUNT(*) FROM governance_budgets b + WHERE (b.virtual_key_id IS NOT NULL OR b.provider_config_id IS NOT NULL) + AND EXISTS (SELECT 1 FROM governance_teams t WHERE t.budget_id = b.id) + `).Scan(&conflictCount).Error; err != nil { + return fmt.Errorf("failed to check for multi-owner team budget conflicts: %w", err) + } + if conflictCount > 0 { + return fmt.Errorf( + "cannot migrate team budgets: %d budget row(s) referenced by a team are already owned by a virtual key or provider config; resolve manually before re-running", + conflictCount, + ) + } + + if err := tx.Exec(` + UPDATE governance_budgets SET team_id = ( + SELECT id FROM governance_teams + WHERE governance_teams.budget_id = governance_budgets.id + ) WHERE team_id IS NULL AND EXISTS ( + SELECT 1 FROM governance_teams + WHERE governance_teams.budget_id = governance_budgets.id + ) + `).Error; err != nil { + return fmt.Errorf("failed to backfill team budget team_id: %w", err) + } + + // Drop legacy budget_id column from governance_teams (raw SQL to avoid GORM FK lookup issues) + _ = tx.Exec("ALTER TABLE governance_teams DROP COLUMN IF EXISTS budget_id") + } + + // Refresh config_hash for teams whose budgets just got linked. GenerateTeamHash + // now includes sorted budget IDs, so hashes written by the earlier + // migrationAddConfigHashColumn (which ran before budgets were associated) + // are stale and would cause phantom drift on the next config.json sync. + var teamsToRehash []tables.TableTeam + if err := tx.Preload("Budgets").Find(&teamsToRehash).Error; err != nil { + return fmt.Errorf("failed to fetch teams for hash refresh: %w", err) + } + for _, team := range teamsToRehash { + if len(team.Budgets) == 0 { + continue // hash did not change; skip + } + hash, err := GenerateTeamHash(team) + if err != nil { + return fmt.Errorf("failed to generate hash for team %s: %w", team.ID, err) + } + if err := tx.Model(&tables.TableTeam{}).Where("id = ?", team.ID).Update("config_hash", hash).Error; err != nil { + return fmt.Errorf("failed to update hash for team %s: %w", team.ID, err) + } + } + + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + if mg.HasColumn(&tables.TableBudget{}, "team_id") { + if err := mg.DropColumn(&tables.TableBudget{}, "team_id"); err != nil { + return err + } + } + return nil + }, + }}) + // SQLite workaround — same reasoning as migrationAddMultiBudgetTables. + if db.Dialector.Name() == "sqlite" { + sqlDB, err := db.DB() + if err != nil { + return fmt.Errorf("failed to get underlying sql.DB: %w", err) + } + sqlDB.SetMaxOpenConns(1) + defer sqlDB.SetMaxOpenConns(0) + + if err := db.Exec("PRAGMA foreign_keys = OFF").Error; err != nil { + return fmt.Errorf("failed to disable SQLite foreign keys: %w", err) + } + defer func() { + if err := db.Exec("PRAGMA foreign_keys = ON").Error; err != nil { + log.Fatalf("[Migration] FATAL: failed to re-enable SQLite foreign keys: %v", err) + } + }() + } + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running add_team_budgets_to_budgets_table migration: %s", err.Error()) + } + return nil +} + // migrationAddPerUserOAuthTables adds the oauth_user_sessions and oauth_user_tokens tables func migrationAddPerUserOAuthTables(ctx context.Context, db *gorm.DB) error { m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index b19a6b3d86..0eff1506a9 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -1915,6 +1915,7 @@ func preloadCustomerRelations(db *gorm.DB, prefix string) *gorm.DB { } return db. Preload(relation("Teams")). + Preload(relation("Teams.Budgets")). Preload(relation("Budget")). Preload(relation("RateLimit")). Preload(relation("VirtualKeys")) @@ -2589,7 +2590,7 @@ func (s *RDBConfigStore) GetTeams(ctx context.Context, customerID string) ([]tab // Preload relationships for complete information query := s.DB().WithContext(ctx). Select(teamSelectWithVKCount). - Preload("Customer").Preload("Budget").Preload("RateLimit") + Preload("Customer").Preload("Budgets").Preload("RateLimit") // Optional filtering by customer if customerID != "" { query = query.Where("customer_id = ?", customerID) @@ -2632,7 +2633,7 @@ func (s *RDBConfigStore) GetTeamsPaginated(ctx context.Context, params TeamsQuer var teams []tables.TableTeam if err := baseQuery. Select(teamSelectWithVKCount). - Preload("Customer").Preload("Budget").Preload("RateLimit"). + Preload("Customer").Preload("Budgets").Preload("RateLimit"). Order("created_at ASC, id ASC"). Offset(offset).Limit(limit). Find(&teams).Error; err != nil { @@ -2647,7 +2648,7 @@ func (s *RDBConfigStore) GetTeam(ctx context.Context, id string) (*tables.TableT var team tables.TableTeam if err := s.DB().WithContext(ctx). Select(teamSelectWithVKCount). - Preload("Customer").Preload("Budget").Preload("RateLimit"). + Preload("Customer").Preload("Budgets").Preload("RateLimit"). First(&team, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, ErrNotFound @@ -2686,10 +2687,12 @@ func (s *RDBConfigStore) UpdateTeam(ctx context.Context, team *tables.TableTeam, } // DeleteTeam deletes a team from the database. +// Owned budgets cascade via the governance_budgets.team_id FK. +// Rate limit is a sibling row (team holds a FK to it) — deleted explicitly. func (s *RDBConfigStore) DeleteTeam(ctx context.Context, id string) error { if err := s.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error { var team tables.TableTeam - if err := tx.WithContext(ctx).Preload("Budget").Preload("RateLimit").First(&team, "id = ?", id).Error; err != nil { + if err := tx.WithContext(ctx).Preload("RateLimit").First(&team, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } @@ -2699,22 +2702,14 @@ func (s *RDBConfigStore) DeleteTeam(ctx context.Context, id string) error { if err := tx.WithContext(ctx).Model(&tables.TableVirtualKey{}).Where("team_id = ?", id).Update("team_id", nil).Error; err != nil { return err } - // Store the budget and rate limit IDs before deleting the team - budgetID := team.BudgetID rateLimitID := team.RateLimitID - // Delete the team first + // Delete the team — owned budgets cascade via FK on governance_budgets.team_id if err := tx.WithContext(ctx).Delete(&tables.TableTeam{}, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } return err } - // Delete the team's budget if it exists - if budgetID != nil { - if err := tx.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", *budgetID).Error; err != nil { - return err - } - } // Delete the team's rate limit if it exists if rateLimitID != nil { if err := tx.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil { diff --git a/framework/configstore/tables/budget.go b/framework/configstore/tables/budget.go index 897cfec4d9..d2bbe8a68e 100644 --- a/framework/configstore/tables/budget.go +++ b/framework/configstore/tables/budget.go @@ -15,7 +15,8 @@ type TableBudget struct { LastReset time.Time `gorm:"index" json:"last_reset"` // Last time budget was reset CurrentUsage float64 `gorm:"default:0" json:"current_usage"` // Current usage in dollars - // Owner FKs: a budget belongs to at most one VK or one ProviderConfig + // Owner FKs: a budget belongs to at most one Team, one VK, or one ProviderConfig + TeamID *string `gorm:"type:varchar(255);index" json:"team_id,omitempty"` VirtualKeyID *string `gorm:"type:varchar(255);index" json:"virtual_key_id,omitempty"` ProviderConfigID *uint `gorm:"index" json:"provider_config_id,omitempty"` @@ -35,8 +36,18 @@ func (TableBudget) TableName() string { return "governance_budgets" } // BeforeSave hook for Budget to validate reset duration format and max limit func (b *TableBudget) BeforeSave(tx *gorm.DB) error { // A budget belongs to at most one owner type - if b.VirtualKeyID != nil && b.ProviderConfigID != nil { - return fmt.Errorf("budget cannot belong to both a virtual key and a provider config") + owners := 0 + if b.TeamID != nil { + owners++ + } + if b.VirtualKeyID != nil { + owners++ + } + if b.ProviderConfigID != nil { + owners++ + } + if owners > 1 { + return fmt.Errorf("budget cannot have more than one owner (team/virtual key/provider config)") } // Validate that ResetDuration is in correct format (e.g., "30s", "5m", "1h", "1d", "1w", "1M", "1Y") if d, err := ParseDuration(b.ResetDuration); err != nil { diff --git a/framework/configstore/tables/team.go b/framework/configstore/tables/team.go index 4beee97ab9..271f3bc8e3 100644 --- a/framework/configstore/tables/team.go +++ b/framework/configstore/tables/team.go @@ -4,7 +4,6 @@ import ( "encoding/json" "time" - bifrost "github.com/maximhq/bifrost/core" "gorm.io/gorm" ) @@ -13,12 +12,11 @@ type TableTeam struct { ID string `gorm:"primaryKey;type:varchar(255)" json:"id"` Name string `gorm:"type:varchar(255);not null" json:"name"` CustomerID *string `gorm:"type:varchar(255);index" json:"customer_id,omitempty"` // A team can belong to a customer - BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"` RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"` // Relationships Customer *TableCustomer `gorm:"foreignKey:CustomerID" json:"customer,omitempty"` - Budget *TableBudget `gorm:"foreignKey:BudgetID" json:"budget,omitempty"` + Budgets []TableBudget `gorm:"foreignKey:TeamID;constraint:OnDelete:CASCADE" json:"budgets,omitempty"` // Multiple budgets with different reset intervals RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID" json:"rate_limit,omitempty"` VirtualKeys []TableVirtualKey `gorm:"foreignKey:TeamID" json:"virtual_keys,omitempty"` @@ -52,7 +50,7 @@ func (t *TableTeam) BeforeSave(tx *gorm.DB) error { if err != nil { return err } - t.Profile = bifrost.Ptr(string(data)) + t.Profile = new(string(data)) } else { t.Profile = nil } @@ -61,7 +59,7 @@ func (t *TableTeam) BeforeSave(tx *gorm.DB) error { if err != nil { return err } - t.Config = bifrost.Ptr(string(data)) + t.Config = new(string(data)) } else { t.Config = nil } @@ -70,7 +68,7 @@ func (t *TableTeam) BeforeSave(tx *gorm.DB) error { if err != nil { return err } - t.Claims = bifrost.Ptr(string(data)) + t.Claims = new(string(data)) } else { t.Claims = nil } diff --git a/plugins/governance/store.go b/plugins/governance/store.go index 1cd0031d1c..26f317d96b 100644 --- a/plugins/governance/store.go +++ b/plugins/governance/store.go @@ -310,12 +310,20 @@ func (gs *LocalGovernanceStore) GetGovernanceData(ctx context.Context) *Governan if team == nil { return } - if team.BudgetID != nil { - if liveBudget, exists := gs.budgets.Load(*team.BudgetID); exists && liveBudget != nil { - if b, ok := liveBudget.(*configstoreTables.TableBudget); ok { - team.Budget = b + // Allocate a fresh slice — shallow-copying `team` (via `clone := *team` at + // the caller) reuses the backing array, so in-place writes would mutate + // the live gs.teams entry under concurrent reads. Mirrors the VK pattern + // above. Budgets missing from gs.budgets are dropped rather than kept stale. + if len(team.Budgets) > 0 { + liveBudgets := make([]configstoreTables.TableBudget, 0, len(team.Budgets)) + for _, b := range team.Budgets { + if lb, exists := gs.budgets.Load(b.ID); exists && lb != nil { + if budget, ok := lb.(*configstoreTables.TableBudget); ok { + liveBudgets = append(liveBudgets, *budget) + } } } + team.Budgets = liveBudgets } if team.RateLimitID != nil { if liveRL, exists := gs.rateLimits.Load(*team.RateLimitID); exists && liveRL != nil { @@ -811,16 +819,20 @@ func (gs *LocalGovernanceStore) CheckTeamBudget(ctx context.Context, teamID stri return DecisionAllow, nil } team, ok := teamValue.(*configstoreTables.TableTeam) - if !ok || team.BudgetID == nil { + if !ok || len(team.Budgets) == 0 { return DecisionAllow, nil } - teamBudget := gs.LoadBudget(ctx, *team.BudgetID) - if teamBudget == nil { + list := make([]*configstoreTables.TableBudget, 0, len(team.Budgets)) + for _, b := range team.Budgets { + if hot := gs.LoadBudget(ctx, b.ID); hot != nil { + list = append(list, hot) + } + } + if len(list) == 0 { return DecisionAllow, nil } key := fmt.Sprintf("Team:%s", teamID) - entityWiseBudgets := EntityWiseBudgets{key: {teamBudget}} - return gs.CheckBudget(ctx, entityWiseBudgets, baselines) + return gs.CheckBudget(ctx, EntityWiseBudgets{key: list}, baselines) } // CheckTeamRateLimit checks team-level rate limit and returns evaluation result if violated @@ -2087,8 +2099,11 @@ func (gs *LocalGovernanceStore) collectBudgetsFromHierarchy(_ context.Context, v if vk.TeamID != nil { if teamValue, exists := gs.teams.Load(*vk.TeamID); exists && teamValue != nil { if team, ok := teamValue.(*configstoreTables.TableTeam); ok && team != nil { - if team.BudgetID != nil { - if budgetValue, exists := gs.budgets.Load(*team.BudgetID); exists && budgetValue != nil { + for _, tb := range team.Budgets { + if seen[tb.ID] { + continue + } + if budgetValue, exists := gs.budgets.Load(tb.ID); exists && budgetValue != nil { if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil { if categoryBudgets := entityWiseBudgets["Team"]; categoryBudgets == nil { entityWiseBudgets["Team"] = []*configstoreTables.TableBudget{} @@ -2372,9 +2387,10 @@ func (gs *LocalGovernanceStore) CreateTeamInMemory(ctx context.Context, team *co return // Nothing to create } - // Create associated budget if exists - if team.Budget != nil { - gs.budgets.Store(team.Budget.ID, team.Budget) + // Create associated budgets if they exist + for i := range team.Budgets { + b := team.Budgets[i] + gs.budgets.Store(b.ID, &b) } // Create associated rate limit if exists @@ -2401,20 +2417,30 @@ func (gs *LocalGovernanceStore) UpdateTeamInMemory(ctx context.Context, team *co // Create clone to avoid modifying the original clone := *team - // Handle budget updates with consistent logic - if clone.Budget != nil { - // Preserve existing usage from memory when updating team budget config - if existingBudgetValue, exists := gs.budgets.Load(clone.Budget.ID); exists && existingBudgetValue != nil { - if existingBudget, ok := existingBudgetValue.(*configstoreTables.TableBudget); ok && existingBudget != nil { - // Preserve current usage and last reset time from existing in-memory budget - clone.Budget.CurrentUsage = existingBudget.CurrentUsage - clone.Budget.LastReset = existingBudget.LastReset + // Reconcile multi-budget slice by ID: preserve live usage on matches, + // evict budgets that disappeared from the team (owned-FK semantics — + // a team's budgets are team-scoped, so dropping the association means + // the budget no longer exists for anyone). + existingBudgetIDs := map[string]struct{}{} + for _, b := range existingTeam.Budgets { + existingBudgetIDs[b.ID] = struct{}{} + } + nextBudgetIDs := map[string]struct{}{} + for i := range clone.Budgets { + b := &clone.Budgets[i] + nextBudgetIDs[b.ID] = struct{}{} + if live, exists := gs.budgets.Load(b.ID); exists && live != nil { + if lb, ok := live.(*configstoreTables.TableBudget); ok && lb != nil { + b.CurrentUsage = lb.CurrentUsage + b.LastReset = lb.LastReset } } - gs.budgets.Store(clone.Budget.ID, clone.Budget) - } else if existingTeam.Budget != nil { - // Budget was removed from the team, delete it from memory - gs.budgets.Delete(existingTeam.Budget.ID) + gs.budgets.Store(b.ID, b) + } + for id := range existingBudgetIDs { + if _, stillThere := nextBudgetIDs[id]; !stillThere { + gs.budgets.Delete(id) + } } // Handle rate limit updates with consistent logic @@ -2447,12 +2473,12 @@ func (gs *LocalGovernanceStore) DeleteTeamInMemory(ctx context.Context, teamID s return // Nothing to delete } - // Get team to check for associated budget and rate limit + // Get team to check for associated budgets and rate limit if teamValue, exists := gs.teams.Load(teamID); exists && teamValue != nil { if team, ok := teamValue.(*configstoreTables.TableTeam); ok && team != nil { - // Delete associated budget if exists - if team.BudgetID != nil { - gs.budgets.Delete(*team.BudgetID) + // Delete all associated budgets + for _, b := range team.Budgets { + gs.budgets.Delete(b.ID) } // Delete associated rate limit if exists if team.RateLimitID != nil { @@ -2807,10 +2833,14 @@ func (gs *LocalGovernanceStore) updateBudgetReferences(ctx context.Context, rese if !ok || team == nil { return true // continue } - if team.BudgetID != nil && *team.BudgetID == budgetID { - clone := *team - clone.Budget = resetBudget - gs.teams.Store(key, &clone) + for i := range team.Budgets { + if team.Budgets[i].ID == budgetID { + clone := *team + clone.Budgets = append([]configstoreTables.TableBudget(nil), team.Budgets...) + clone.Budgets[i] = *resetBudget + gs.teams.Store(key, &clone) + break + } } return true // continue }) diff --git a/plugins/governance/test_utils.go b/plugins/governance/test_utils.go index b7ad6dbad3..51533ba075 100644 --- a/plugins/governance/test_utils.go +++ b/plugins/governance/test_utils.go @@ -172,8 +172,8 @@ func buildTeam(id, name string, budget *configstoreTables.TableBudget) *configst Name: name, } if budget != nil { - team.Budget = budget - team.BudgetID = &budget.ID + budget.TeamID = &team.ID + team.Budgets = []configstoreTables.TableBudget{*budget} } return team } diff --git a/tests/governance/advancedscenarios_test.go b/tests/governance/advancedscenarios_test.go index aa928c7dc9..bfc108ec8a 100644 --- a/tests/governance/advancedscenarios_test.go +++ b/tests/governance/advancedscenarios_test.go @@ -24,10 +24,10 @@ func TestVKSwitchTeamAfterBudgetExhaustion(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: team1Name, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: team1Budget, ResetDuration: "1h", - }, + }}, }, }) @@ -46,10 +46,10 @@ func TestVKSwitchTeamAfterBudgetExhaustion(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: team2Name, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: team2Budget, ResetDuration: "1h", - }, + }}, }, }) @@ -372,10 +372,10 @@ func TestHierarchicalChainBudgetSwitch(t *testing.T) { Body: CreateTeamRequest{ Name: team1Name, CustomerID: &customer1ID, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: 100.0, // High budget - customer is limiting ResetDuration: "1h", - }, + }}, }, }) @@ -415,10 +415,10 @@ func TestHierarchicalChainBudgetSwitch(t *testing.T) { Body: CreateTeamRequest{ Name: team2Name, CustomerID: &customer2ID, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: 100.0, // High budget ResetDuration: "1h", - }, + }}, }, }) @@ -674,10 +674,10 @@ func TestTeamBudgetUpdateAfterExhaustion(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: teamName, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: initialBudget, ResetDuration: "1h", - }, + }}, }, }) @@ -762,10 +762,10 @@ func TestTeamBudgetUpdateAfterExhaustion(t *testing.T) { Method: "PUT", Path: "/api/governance/teams/" + teamID, Body: UpdateTeamRequest{ - Budget: &UpdateBudgetRequest{ - MaxLimit: &newBudget, - ResetDuration: &resetDuration, - }, + Budgets: &[]BudgetRequest{{ + MaxLimit: newBudget, + ResetDuration: resetDuration, + }}, }, }) @@ -1291,10 +1291,10 @@ func TestTeamDeletionDeletesBudget(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: teamName, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: 100.0, ResetDuration: "1h", - }, + }}, }, }) @@ -1323,7 +1323,12 @@ func TestTeamDeletionDeletesBudget(t *testing.T) { budgetsMap1 := getBudgetsResp1.Body["budgets"].(map[string]interface{}) teamData1 := teamsMap1[teamID].(map[string]interface{}) - budgetID := teamData1["budget_id"].(string) + // Teams now expose a `budgets` array instead of a single `budget_id`. + budgetsList, ok := teamData1["budgets"].([]interface{}) + if !ok || len(budgetsList) == 0 { + t.Fatalf("Team %s has no budgets in memory before deletion", teamID) + } + budgetID := budgetsList[0].(map[string]interface{})["id"].(string) _, budgetExists := budgetsMap1[budgetID] if !budgetExists { diff --git a/tests/governance/configupdatesync_test.go b/tests/governance/configupdatesync_test.go index 994d0448f6..a7008e80ed 100644 --- a/tests/governance/configupdatesync_test.go +++ b/tests/governance/configupdatesync_test.go @@ -582,10 +582,10 @@ func TestTeamBudgetUpdateSyncToMemory(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: teamName, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: initialBudget, ResetDuration: resetDuration, - }, + }}, }, }) @@ -627,7 +627,12 @@ func TestTeamBudgetUpdateSyncToMemory(t *testing.T) { teamsMap1 := getTeamsResp1.Body["teams"].(map[string]interface{}) teamData1 := teamsMap1[teamID].(map[string]interface{}) - budgetID, _ := teamData1["budget_id"].(string) + // Teams now expose a `budgets` array instead of a single `budget_id`. + budgetsList, ok := teamData1["budgets"].([]interface{}) + if !ok || len(budgetsList) == 0 { + t.Fatalf("Team %s has no budgets in memory", teamID) + } + budgetID, _ := budgetsList[0].(map[string]interface{})["id"].(string) getBudgetsResp1 := MakeRequest(t, APIRequest{ Method: "GET", @@ -697,10 +702,10 @@ func TestTeamBudgetUpdateSyncToMemory(t *testing.T) { Method: "PUT", Path: "/api/governance/teams/" + teamID, Body: UpdateTeamRequest{ - Budget: &UpdateBudgetRequest{ - MaxLimit: &newLowerBudget, - ResetDuration: &resetDurationPtr, - }, + Budgets: &[]BudgetRequest{{ + MaxLimit: newLowerBudget, + ResetDuration: resetDurationPtr, + }}, }, }) diff --git a/tests/governance/customerbudget_test.go b/tests/governance/customerbudget_test.go index 3917136502..c4c9bfe1e7 100644 --- a/tests/governance/customerbudget_test.go +++ b/tests/governance/customerbudget_test.go @@ -197,10 +197,10 @@ func TestCustomerBudgetExceededWithMultipleTeams(t *testing.T) { Body: CreateTeamRequest{ Name: "test-team-" + generateRandomID(), CustomerID: &customerID, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: 1.0, // High team budget so customer is the limiting factor ResetDuration: "1h", - }, + }}, }, }) diff --git a/tests/governance/e2e_test.go b/tests/governance/e2e_test.go index bab26fff30..26b59d57ea 100644 --- a/tests/governance/e2e_test.go +++ b/tests/governance/e2e_test.go @@ -29,10 +29,10 @@ func TestMultipleVKsSharingTeamBudgetFairness(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: teamName, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: teamBudget, ResetDuration: teamResetDuration, - }, + }}, }, }) @@ -221,10 +221,10 @@ func TestFullBudgetHierarchyEnforcement(t *testing.T) { Body: CreateTeamRequest{ Name: teamName, CustomerID: &customerID, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: 100.0, // Medium ResetDuration: "1h", - }, + }}, }, }) @@ -1086,10 +1086,10 @@ func TestTeamDeletionCascade(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: teamName, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: 100.0, ResetDuration: "1h", - }, + }}, }, }) diff --git a/tests/governance/edgecases_test.go b/tests/governance/edgecases_test.go index 32dad1bbe6..b70d9738c0 100644 --- a/tests/governance/edgecases_test.go +++ b/tests/governance/edgecases_test.go @@ -43,10 +43,10 @@ func TestCrissCrossComplexBudgetHierarchy(t *testing.T) { Body: CreateTeamRequest{ Name: "test-team-criss-cross-" + generateRandomID(), CustomerID: &customerID, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: teamBudget, ResetDuration: "1h", - }, + }}, }, }) diff --git a/tests/governance/inmemorysync_test.go b/tests/governance/inmemorysync_test.go index 4ca4594c22..5e18eb6c1d 100644 --- a/tests/governance/inmemorysync_test.go +++ b/tests/governance/inmemorysync_test.go @@ -143,10 +143,10 @@ func TestInMemorySyncTeamUpdate(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: teamName, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: initialBudget, ResetDuration: "1h", - }, + }}, }, }) @@ -184,9 +184,10 @@ func TestInMemorySyncTeamUpdate(t *testing.T) { Method: "PUT", Path: "/api/governance/teams/" + teamID, Body: UpdateTeamRequest{ - Budget: &UpdateBudgetRequest{ - MaxLimit: &newTeamBudget, - }, + Budgets: &[]BudgetRequest{{ + MaxLimit: newTeamBudget, + ResetDuration: "1h", + }}, }, }) @@ -223,7 +224,13 @@ func TestInMemorySyncTeamUpdate(t *testing.T) { } teamDataMap := teamData2.(map[string]interface{}) - budgetID, _ := teamDataMap["budget_id"].(string) + // Teams now expose a `budgets` array instead of a single `budget_id` — read the first. + var budgetID string + if budgetsList, ok := teamDataMap["budgets"].([]interface{}); ok && len(budgetsList) > 0 { + if b, ok := budgetsList[0].(map[string]interface{}); ok { + budgetID, _ = b["id"].(string) + } + } if budgetID != "" { budgetData, budgetExists := budgetsMap2[budgetID] @@ -477,10 +484,10 @@ func TestDataEndpointConsistency(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: teamName, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: 30.0, ResetDuration: "1h", - }, + }}, }, }) diff --git a/tests/governance/teambudget_test.go b/tests/governance/teambudget_test.go index 191a39fb1c..301ca63002 100644 --- a/tests/governance/teambudget_test.go +++ b/tests/governance/teambudget_test.go @@ -20,10 +20,10 @@ func TestTeamBudgetExceededWithMultipleVKs(t *testing.T) { Path: "/api/governance/teams", Body: CreateTeamRequest{ Name: teamName, - Budget: &BudgetRequest{ + Budgets: []BudgetRequest{{ MaxLimit: teamBudget, ResetDuration: "1h", - }, + }}, }, }) diff --git a/tests/governance/test_utils.go b/tests/governance/test_utils.go index 487598ef46..b4ada19490 100644 --- a/tests/governance/test_utils.go +++ b/tests/governance/test_utils.go @@ -241,9 +241,9 @@ type BudgetRequest struct { // CreateTeamRequest represents a request to create a team type CreateTeamRequest struct { - Name string `json:"name"` - CustomerID *string `json:"customer_id,omitempty"` - Budget *BudgetRequest `json:"budget,omitempty"` + Name string `json:"name"` + CustomerID *string `json:"customer_id,omitempty"` + Budgets []BudgetRequest `json:"budgets,omitempty"` } // CreateCustomerRequest represents a request to create a customer @@ -279,8 +279,12 @@ type UpdateVirtualKeyRequest struct { // UpdateTeamRequest represents a request to update a team type UpdateTeamRequest struct { - Name *string `json:"name,omitempty"` - Budget *UpdateBudgetRequest `json:"budget,omitempty"` + Name *string `json:"name,omitempty"` + // Pointer-to-slice so tests can distinguish: + // nil → field omitted (budgets untouched by server) + // &[]BudgetRequest{} → explicit empty array (server clears all budgets) + // &[]BudgetRequest{…} → replace with the provided budgets + Budgets *[]BudgetRequest `json:"budgets,omitempty"` } // UpdateCustomerRequest represents a request to update a customer diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index c5ad8de9ad..b30ae8dd99 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -222,7 +222,7 @@ func collectProviderConfigDeleteIDs( type CreateTeamRequest struct { Name string `json:"name" validate:"required"` CustomerID *string `json:"customer_id,omitempty"` // Team can belong to a customer - Budget *CreateBudgetRequest `json:"budget,omitempty"` // Team can have its own budget + Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: each must have a unique reset_duration RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` // Team can have its own rate limit } @@ -230,7 +230,7 @@ type CreateTeamRequest struct { type UpdateTeamRequest struct { Name *string `json:"name,omitempty"` CustomerID *string `json:"customer_id,omitempty"` - Budget *UpdateBudgetRequest `json:"budget,omitempty"` + Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: replaces all team budgets RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` } @@ -1407,18 +1407,6 @@ func (h *GovernanceHandler) createTeam(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Team name is required") return } - // Validate budget if provided - if req.Budget != nil { - if req.Budget.MaxLimit < 0 { - SendError(ctx, 400, fmt.Sprintf("Budget max_limit cannot be negative: %.2f", req.Budget.MaxLimit)) - return - } - // Validate reset duration format - if _, err := configstoreTables.ParseDuration(req.Budget.ResetDuration); err != nil { - SendError(ctx, 400, fmt.Sprintf("Invalid reset duration format: %s", req.Budget.ResetDuration)) - return - } - } // Validate rate limit if provided if req.RateLimit != nil { rateLimit := configstoreTables.TableRateLimit{ @@ -1440,22 +1428,6 @@ func (h *GovernanceHandler) createTeam(ctx *fasthttp.RequestCtx) { Name: req.Name, CustomerID: req.CustomerID, } - if req.Budget != nil { - budget := configstoreTables.TableBudget{ - ID: uuid.NewString(), - MaxLimit: req.Budget.MaxLimit, - ResetDuration: req.Budget.ResetDuration, - LastReset: budgetLastReset(false, req.Budget.ResetDuration), - CurrentUsage: 0, - } - if err := validateBudget(&budget); err != nil { - return err - } - if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil { - return err - } - team.BudgetID = &budget.ID - } if req.RateLimit != nil { rateLimit := configstoreTables.TableRateLimit{ ID: uuid.NewString(), @@ -1471,11 +1443,47 @@ func (h *GovernanceHandler) createTeam(ctx *fasthttp.RequestCtx) { } team.RateLimitID = &rateLimit.ID } + // Team row must exist before child budgets (FK on governance_budgets.team_id) if err := h.configStore.CreateTeam(ctx, &team, tx); err != nil { return err } + // Create owned multi-budgets; enforce unique reset_duration per team + seenDurations := make(map[string]bool) + for _, b := range req.Budgets { + if b.MaxLimit < 0 { + return &badRequestError{err: fmt.Errorf("budget max_limit cannot be negative: %.2f", b.MaxLimit)} + } + if _, err := configstoreTables.ParseDuration(b.ResetDuration); err != nil { + return &badRequestError{err: fmt.Errorf("invalid reset duration format: %s", b.ResetDuration)} + } + if seenDurations[b.ResetDuration] { + return &badRequestError{err: fmt.Errorf("duplicate reset_duration in budgets: %s", b.ResetDuration)} + } + seenDurations[b.ResetDuration] = true + budget := configstoreTables.TableBudget{ + ID: uuid.NewString(), + MaxLimit: b.MaxLimit, + ResetDuration: b.ResetDuration, + LastReset: budgetLastReset(b.CalendarAligned, b.ResetDuration), + CurrentUsage: 0, + CalendarAligned: b.CalendarAligned, + TeamID: &team.ID, + } + if err := validateBudget(&budget); err != nil { + return err + } + if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil { + return err + } + team.Budgets = append(team.Budgets, budget) + } return nil }); err != nil { + var badReqErr *badRequestError + if errors.As(err, &badReqErr) { + SendError(ctx, 400, err.Error()) + return + } logger.Error("failed to create team: %v", err) SendError(ctx, 500, "failed to create team") return @@ -1548,8 +1556,8 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { } // Updating team in database if err := h.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error { - // Track IDs to delete after updating the team (to avoid FK constraint) - var budgetIDToDelete, rateLimitIDToDelete string + // Track rate-limit ID to delete after updating the team (to avoid FK constraint) + var rateLimitIDToDelete string // Update fields if provided if req.Name != nil { @@ -1562,63 +1570,81 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { team.CustomerID = req.CustomerID } } - // Handle budget updates - if req.Budget != nil { - // Check if budget removal is requested (all fields nil) - budgetIsEmpty := isBudgetRemovalRequest(req.Budget) - if budgetIsEmpty { - // Mark budget for deletion after FK is removed - if team.BudgetID != nil { - budgetIDToDelete = *team.BudgetID - team.BudgetID = nil - team.Budget = nil - } - } else if team.BudgetID != nil { - // Update existing budget — all fields are optional (partial update) - budget := configstoreTables.TableBudget{} - if err := tx.First(&budget, "id = ?", *team.BudgetID).Error; err != nil { - return err - } - if req.Budget.MaxLimit != nil { - budget.MaxLimit = *req.Budget.MaxLimit - } - if req.Budget.ResetDuration != nil { - budget.ResetDuration = *req.Budget.ResetDuration - } - if err := validateBudget(&budget); err != nil { - return err - } - if err := h.configStore.UpdateBudget(ctx, &budget, tx); err != nil { - return err - } - team.Budget = &budget - } else { - // Create new budget - if req.Budget.MaxLimit == nil || req.Budget.ResetDuration == nil { - return fmt.Errorf("both max_limit and reset_duration are required when creating a new budget") - } - if *req.Budget.MaxLimit < 0 { - return fmt.Errorf("budget max_limit cannot be negative: %.2f", *req.Budget.MaxLimit) + // Multi-budget reconciliation: match by reset_duration, preserve usage on update, + // create new budgets for new durations, delete unmatched existing budgets. + // Mirrors VK multi-budget handling above. + if req.Budgets != nil { + // Validate incoming budgets + seenDurations := make(map[string]bool) + for _, b := range req.Budgets { + if b.MaxLimit < 0 { + return &badRequestError{err: fmt.Errorf("budget max_limit cannot be negative: %.2f", b.MaxLimit)} } - if _, err := configstoreTables.ParseDuration(*req.Budget.ResetDuration); err != nil { - return fmt.Errorf("invalid reset duration format: %s", *req.Budget.ResetDuration) + if _, err := configstoreTables.ParseDuration(b.ResetDuration); err != nil { + return &badRequestError{err: fmt.Errorf("invalid reset duration format: %s", b.ResetDuration)} } - budget := configstoreTables.TableBudget{ - ID: uuid.NewString(), - MaxLimit: *req.Budget.MaxLimit, - ResetDuration: *req.Budget.ResetDuration, - LastReset: budgetLastReset(false, *req.Budget.ResetDuration), - CurrentUsage: 0, + if seenDurations[b.ResetDuration] { + return &badRequestError{err: fmt.Errorf("duplicate reset_duration in budgets: %s", b.ResetDuration)} } - if err := validateBudget(&budget); err != nil { - return err + seenDurations[b.ResetDuration] = true + } + + existingByDuration := make(map[string]configstoreTables.TableBudget) + for _, existing := range team.Budgets { + existingByDuration[existing.ResetDuration] = existing + } + + var reconciledBudgets []configstoreTables.TableBudget + matchedIDs := make(map[string]bool) + for _, b := range req.Budgets { + if existing, found := existingByDuration[b.ResetDuration]; found { + wasCalendarAligned := existing.CalendarAligned + existing.MaxLimit = b.MaxLimit + existing.CalendarAligned = b.CalendarAligned + // Match the UI's calendar-alignment confirmation promise: on the + // false → true transition, snap LastReset to the current period + // start and zero out CurrentUsage now, instead of lazily waiting + // for the next period boundary in ResetExpiredBudgetsInMemory. + if b.CalendarAligned && !wasCalendarAligned { + existing.LastReset = configstoreTables.GetCalendarPeriodStart(b.ResetDuration, time.Now()) + existing.CurrentUsage = 0 + } + if err := validateBudget(&existing); err != nil { + return err + } + if err := h.configStore.UpdateBudget(ctx, &existing, tx); err != nil { + return err + } + reconciledBudgets = append(reconciledBudgets, existing) + matchedIDs[existing.ID] = true + } else { + budget := configstoreTables.TableBudget{ + ID: uuid.NewString(), + MaxLimit: b.MaxLimit, + ResetDuration: b.ResetDuration, + LastReset: budgetLastReset(b.CalendarAligned, b.ResetDuration), + CurrentUsage: 0, + CalendarAligned: b.CalendarAligned, + TeamID: &team.ID, + } + if err := validateBudget(&budget); err != nil { + return err + } + if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil { + return err + } + reconciledBudgets = append(reconciledBudgets, budget) } - if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil { - return err + } + // Delete budgets that are no longer present + for _, existing := range team.Budgets { + if !matchedIDs[existing.ID] { + if err := h.configStore.DeleteBudget(ctx, existing.ID, tx); err != nil { + return fmt.Errorf("failed to delete removed team budget: %w", err) + } } - team.BudgetID = &budget.ID - team.Budget = &budget } + team.Budgets = reconciledBudgets } // Handle rate limit updates if req.RateLimit != nil { @@ -1673,12 +1699,9 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { return err } - // Now that FK references are removed, delete the orphaned budget/rate limit - if budgetIDToDelete != "" { - if err := tx.Delete(&configstoreTables.TableBudget{}, "id = ?", budgetIDToDelete).Error; err != nil { - return err - } - } + // Now that FK references are removed, delete the orphaned rate limit. + // Budgets are reconciled above (deletion of unmatched rows happens inside + // the reconciliation loop), so nothing to clean up here. if rateLimitIDToDelete != "" { if err := tx.Delete(&configstoreTables.TableRateLimit{}, "id = ?", rateLimitIDToDelete).Error; err != nil { return err @@ -1687,6 +1710,12 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { return nil }); err != nil { + var badReqErr *badRequestError + if errors.As(err, &badReqErr) { + SendError(ctx, 400, err.Error()) + return + } + logger.Error("failed to update team: %v", err) SendError(ctx, 500, "Failed to update team") return } @@ -1809,18 +1838,6 @@ func (h *GovernanceHandler) createCustomer(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Customer name is required") return } - // Validate budget if provided - if req.Budget != nil { - if req.Budget.MaxLimit < 0 { - SendError(ctx, 400, fmt.Sprintf("Budget max_limit cannot be negative: %.2f", req.Budget.MaxLimit)) - return - } - // Validate reset duration format - if _, err := configstoreTables.ParseDuration(req.Budget.ResetDuration); err != nil { - SendError(ctx, 400, fmt.Sprintf("Invalid reset duration format: %s", req.Budget.ResetDuration)) - return - } - } // Validate rate limit if provided if req.RateLimit != nil { rateLimit := configstoreTables.TableRateLimit{ diff --git a/ui/app/workspace/governance/views/teamDialog.tsx b/ui/app/workspace/governance/views/teamDialog.tsx index 6d282d1229..78dae8e395 100644 --- a/ui/app/workspace/governance/views/teamDialog.tsx +++ b/ui/app/workspace/governance/views/teamDialog.tsx @@ -50,6 +50,7 @@ import { formatDistanceToNow } from "date-fns"; import isEqual from "lodash.isequal"; import { useEffect, useMemo, useState } from "react"; import { toast } from "sonner"; +import { v4 as uuid } from "uuid"; interface TeamDialogProps { team?: Team | null; @@ -58,13 +59,23 @@ interface TeamDialogProps { onCancel: () => void; } +// One editable budget row; teams own multiple, each keyed by reset_duration +// on the wire. The client-side `id` is stable across re-renders and equals +// the persisted budget's id for existing rows, or a fresh UUID for new ones — +// used as the React key and for matching against `team.budgets` when we need +// to distinguish "already persisted" from "just added in the form". +interface TeamBudgetRow { + id: string; + maxLimit: number | undefined; + resetDuration: string; + calendarAligned: boolean; +} + interface TeamFormData { name: string; customerId: string; - // Budget - budgetMaxLimit: number | undefined; - budgetResetDuration: string; - budgetCalendarAligned: boolean; + // Multi-budget: each row has a unique reset_duration on submit + budgets: TeamBudgetRow[]; // Rate Limit tokenMaxLimit: number | undefined; tokenResetDuration: string; @@ -80,10 +91,13 @@ const createInitialState = ( return { name: team?.name || "", customerId: team?.customer_id || "", - // Budget - budgetMaxLimit: team?.budget?.max_limit ?? undefined, - budgetResetDuration: team?.budget?.reset_duration || "1M", - budgetCalendarAligned: team?.budget?.calendar_aligned ?? false, + budgets: + team?.budgets?.map((b) => ({ + id: b.id, + maxLimit: b.max_limit, + resetDuration: b.reset_duration, + calendarAligned: b.calendar_aligned ?? false, + })) ?? [], // Rate Limit tokenMaxLimit: team?.rate_limit?.token_max_limit ?? undefined, tokenResetDuration: team?.rate_limit?.token_reset_duration || "1h", @@ -111,7 +125,7 @@ export default function TeamDialog({ const nextInitial = createInitialState(team); setInitialState(nextInitial); setFormData({ ...nextInitial, isDirty: false }); - setShowCalendarAlignWarning(false); + setPendingCalendarAlignIdx(null); }, [team]); const hasCreateAccess = useRbac(RbacResource.Teams, RbacOperation.Create); @@ -123,25 +137,63 @@ export default function TeamDialog({ const [updateTeam, { isLoading: isUpdating }] = useUpdateTeamMutation(); const loading = isCreating || isUpdating; - const [showCalendarAlignWarning, setShowCalendarAlignWarning] = - useState(false); + // Tracks which row (by index) is awaiting calendar-align confirmation. + const [pendingCalendarAlignIdx, setPendingCalendarAlignIdx] = useState< + number | null + >(null); + const showCalendarAlignWarning = pendingCalendarAlignIdx !== null; + + const updateBudgetRow = (idx: number, patch: Partial) => { + setFormData((prev) => { + const next = prev.budgets.map((row, i) => + i === idx ? { ...row, ...patch } : row, + ); + return { ...prev, budgets: next }; + }); + }; + + const addBudgetRow = () => { + setFormData((prev) => ({ + ...prev, + budgets: [ + ...prev.budgets, + { + id: uuid(), + maxLimit: undefined, + resetDuration: "1M", + calendarAligned: false, + }, + ], + })); + }; - const handleCalendarAlignedChange = (checked: boolean) => { - if (checked && isEditing && team?.budget && !team.budget.calendar_aligned) { - setShowCalendarAlignWarning(true); + const removeBudgetRow = (idx: number) => { + setFormData((prev) => ({ + ...prev, + budgets: prev.budgets.filter((_, i) => i !== idx), + })); + }; + + const handleCalendarAlignedChange = (idx: number, checked: boolean) => { + // Match the persisted budget by stable row id — for seeded rows this equals + // the server-side budget id; for newly-added rows it's a client-only UUID + // that won't match anything in team.budgets (correctly: no warning for new rows). + // Avoids the reset_duration-duplicate ambiguity before validation resolves. + const rowId = formData.budgets[idx]?.id; + const existingBudget = team?.budgets?.find((b) => b.id === rowId); + if (checked && isEditing && existingBudget && !existingBudget.calendar_aligned) { + setPendingCalendarAlignIdx(idx); } else { - updateField("budgetCalendarAligned", checked); + updateBudgetRow(idx, { calendarAligned: checked }); } }; // Track isDirty state useEffect(() => { - const currentData = { + const currentData: Omit = { name: formData.name, customerId: formData.customerId, - budgetMaxLimit: formData.budgetMaxLimit, - budgetResetDuration: formData.budgetResetDuration, - budgetCalendarAligned: formData.budgetCalendarAligned, + budgets: formData.budgets, tokenMaxLimit: formData.tokenMaxLimit, tokenResetDuration: formData.tokenResetDuration, requestMaxLimit: formData.requestMaxLimit, @@ -154,9 +206,7 @@ export default function TeamDialog({ }, [ formData.name, formData.customerId, - formData.budgetMaxLimit, - formData.budgetResetDuration, - formData.budgetCalendarAligned, + formData.budgets, formData.tokenMaxLimit, formData.tokenResetDuration, formData.requestMaxLimit, @@ -164,71 +214,73 @@ export default function TeamDialog({ initialState, ]); - // Values for validation and submission (already numbers) - const budgetMaxLimitNum = formData.budgetMaxLimit; const tokenMaxLimitNum = formData.tokenMaxLimit; const requestMaxLimitNum = formData.requestMaxLimit; // Validation - const validator = useMemo( - () => - new Validator([ - // Basic validation - Validator.required(formData.name.trim(), "Team name is required"), - - // Check if anything is dirty - Validator.custom(formData.isDirty, "No changes to save"), - - // Budget validation - ...(formData.budgetMaxLimit !== undefined && - formData.budgetMaxLimit !== null - ? [ - Validator.minValue( - budgetMaxLimitNum || 0, - 0.01, - "Budget max limit must be greater than $0.01", - ), - Validator.required( - formData.budgetResetDuration, - "Budget reset duration is required", - ), - ] - : []), - - // Rate limit validation - token limits - ...(formData.tokenMaxLimit !== undefined && - formData.tokenMaxLimit !== null - ? [ - Validator.minValue( - tokenMaxLimitNum || 0, - 1, - "Token max limit must be at least 1", - ), - Validator.required( - formData.tokenResetDuration, - "Token reset duration is required", - ), - ] - : []), - - // Rate limit validation - request limits - ...(formData.requestMaxLimit !== undefined && - formData.requestMaxLimit !== null - ? [ - Validator.minValue( - requestMaxLimitNum || 0, - 1, - "Request max limit must be at least 1", - ), - Validator.required( - formData.requestResetDuration, - "Request reset duration is required", - ), - ] - : []), - ]), - [formData, budgetMaxLimitNum, tokenMaxLimitNum, requestMaxLimitNum], - ); + const validator = useMemo(() => { + // Per-row budget validation plus cross-row uniqueness on reset_duration. + const budgetValidators = formData.budgets.flatMap((row, idx) => { + if (row.maxLimit === undefined || row.maxLimit === null) return []; + return [ + Validator.minValue( + row.maxLimit, + 0.01, + `Budget #${idx + 1} max limit must be greater than $0.01`, + ), + Validator.required( + row.resetDuration, + `Budget #${idx + 1} reset duration is required`, + ), + ]; + }); + const populatedDurations = formData.budgets + .filter((r) => r.maxLimit !== undefined && r.maxLimit !== null) + .map((r) => r.resetDuration); + const uniqueDurations = new Set(populatedDurations).size; + + return new Validator([ + Validator.required(formData.name.trim(), "Team name is required"), + Validator.custom(formData.isDirty, "No changes to save"), + ...budgetValidators, + Validator.custom( + uniqueDurations === populatedDurations.length, + "Each budget must have a distinct reset duration", + ), + + // Rate limit validation - token limits + ...(formData.tokenMaxLimit !== undefined && + formData.tokenMaxLimit !== null + ? [ + Validator.minValue( + tokenMaxLimitNum || 0, + 1, + "Token max limit must be at least 1", + ), + Validator.required( + formData.tokenResetDuration, + "Token reset duration is required", + ), + ] + : []), + + // Rate limit validation - request limits + ...(formData.requestMaxLimit !== undefined && + formData.requestMaxLimit !== null + ? [ + Validator.minValue( + requestMaxLimitNum || 0, + 1, + "Request max limit must be at least 1", + ), + Validator.required( + formData.requestResetDuration, + "Request reset duration is required", + ), + ] + : []), + ]); + }, [formData, tokenMaxLimitNum, requestMaxLimitNum]); const updateField = ( field: K, @@ -245,28 +297,26 @@ export default function TeamDialog({ return; } + // Serialize budget rows whose max_limit was filled in — rows left blank + // are silently dropped (the backend treats the slice as authoritative). + const submittableBudgets = formData.budgets + .filter((r) => r.maxLimit !== undefined && r.maxLimit !== null) + .map((r) => ({ + max_limit: r.maxLimit as number, + reset_duration: r.resetDuration, + calendar_aligned: r.calendarAligned, + })); + try { if (isEditing && team) { // Update existing team const updateData: UpdateTeamRequest = { name: formData.name, customer_id: formData.customerId || undefined, + // Always send: backend treats `budgets` as a full replacement. + budgets: submittableBudgets, }; - // Detect budget changes using had/has pattern - const hadBudget = !!team.budget; - const hasBudget = - budgetMaxLimitNum !== undefined && budgetMaxLimitNum !== null; - if (hasBudget) { - updateData.budget = { - max_limit: budgetMaxLimitNum, - reset_duration: formData.budgetResetDuration, - calendar_aligned: formData.budgetCalendarAligned, - }; - } else if (hadBudget) { - updateData.budget = {} as UpdateTeamRequest["budget"]; - } - // Detect rate limit changes using had/has pattern const hadRateLimit = !!team.rate_limit; const hasRateLimit = @@ -296,17 +346,10 @@ export default function TeamDialog({ const createData: CreateTeamRequest = { name: formData.name, customer_id: formData.customerId || undefined, + budgets: + submittableBudgets.length > 0 ? submittableBudgets : undefined, }; - // Add budget if enabled - if (budgetMaxLimitNum !== undefined && budgetMaxLimitNum !== null) { - createData.budget = { - max_limit: budgetMaxLimitNum, - reset_duration: formData.budgetResetDuration, - calendar_aligned: formData.budgetCalendarAligned, - }; - } - // Add rate limit if enabled (token or request limits) if ( (tokenMaxLimitNum !== undefined && tokenMaxLimitNum !== null) || @@ -411,56 +454,98 @@ export default function TeamDialog({ )} - {/* Budget Configuration */} - updateField("budgetMaxLimit", value)} - onChangeSelect={(value) => { - updateField("budgetResetDuration", value); - if (!supportsCalendarAlignment(value)) { - updateField("budgetCalendarAligned", false); - } - }} - options={resetDurationOptions} - dataTestId="budget-max-limit-input" - /> - - {/* Calendar alignment toggle — only shown when a budget is set and the period supports alignment */} - {formData.budgetMaxLimit && - supportsCalendarAlignment(formData.budgetResetDuration) && ( -
-
- -

+

+ + +
+ {formData.budgets.length === 0 && ( +

+ No budgets. Click "Add budget" to enforce a spend limit. +

+ )} + {formData.budgets.map((row, idx) => ( +
+
+
+ + updateBudgetRow(idx, { maxLimit: value }) + } + onChangeSelect={(value) => { + const patch: Partial = { + resetDuration: value, + }; + if (!supportsCalendarAlignment(value)) { + patch.calendarAligned = false; + } + updateBudgetRow(idx, patch); + }} + options={resetDurationOptions} + dataTestId={`budget-max-limit-input-${idx}`} + /> +
+
- + + {row.maxLimit !== undefined && + supportsCalendarAlignment(row.resetDuration) && ( +
+
+ +

+ Reset at the start of each period (e.g. 1st of + month) instead of rolling from creation date +

+
+ + handleCalendarAlignedChange(idx, checked) + } + data-testid={`team-budget-calendar-aligned-toggle-${idx}`} + /> +
+ )}
- )} + ))} +
{/* Warning dialog shown when enabling calendar alignment on an existing budget */} { + if (!open) setPendingCalendarAlignIdx(null); + }} > @@ -470,13 +555,21 @@ export default function TeamDialog({ current usage to{" "} $0.00 and snap the reset date to the start of the current{" "} - {formData.budgetResetDuration === "1d" + {pendingCalendarAlignIdx !== null && + formData.budgets[pendingCalendarAlignIdx]?.resetDuration === + "1d" ? "day" - : formData.budgetResetDuration === "1w" + : pendingCalendarAlignIdx !== null && + formData.budgets[pendingCalendarAlignIdx] + ?.resetDuration === "1w" ? "week" - : formData.budgetResetDuration === "1M" + : pendingCalendarAlignIdx !== null && + formData.budgets[pendingCalendarAlignIdx] + ?.resetDuration === "1M" ? "month" - : formData.budgetResetDuration === "1Y" + : pendingCalendarAlignIdx !== null && + formData.budgets[pendingCalendarAlignIdx] + ?.resetDuration === "1Y" ? "year" : "period"} . The usage reset to $0.00 cannot be undone, but calendar @@ -485,14 +578,21 @@ export default function TeamDialog({ - + setPendingCalendarAlignIdx(null)} + > Cancel { - updateField("budgetCalendarAligned", true); - setShowCalendarAlignWarning(false); + if (pendingCalendarAlignIdx !== null) { + updateBudgetRow(pendingCalendarAlignIdx, { + calendarAligned: true, + }); + } + setPendingCalendarAlignIdx(null); }} > Enable Calendar Alignment @@ -528,45 +628,44 @@ export default function TeamDialog({ /> {/* Current Usage Section (only shown when editing with existing limits) */} - {isEditing && (team?.budget || team?.rate_limit) && ( + {isEditing && + ((team?.budgets && team.budgets.length > 0) || + team?.rate_limit) && (

Current Usage

- {team?.budget && ( -
-

Budget

+ {team?.budgets?.map((b) => ( +
+

+ Budget ({b.reset_duration}) +

- {formatCurrency(team.budget.current_usage)} /{" "} - {formatCurrency(team.budget.max_limit)} + {formatCurrency(b.current_usage)} /{" "} + {formatCurrency(b.max_limit)} 0 && - team.budget.current_usage >= team.budget.max_limit + b.max_limit > 0 && b.current_usage >= b.max_limit ? "destructive" : "default" } className="text-xs" > - {team.budget.max_limit > 0 - ? Math.round( - (team.budget.current_usage / - team.budget.max_limit) * - 100, - ) + {b.max_limit > 0 + ? Math.round((b.current_usage / b.max_limit) * 100) : 0} %

Last Reset:{" "} - {formatDistanceToNow(new Date(team.budget.last_reset), { + {formatDistanceToNow(new Date(b.last_reset), { addSuffix: true, })}

- )} + ))} {team?.rate_limit?.token_max_limit && (

Tokens

diff --git a/ui/app/workspace/governance/views/teamsTable.tsx b/ui/app/workspace/governance/views/teamsTable.tsx index 14184be43c..dbd646488e 100644 --- a/ui/app/workspace/governance/views/teamsTable.tsx +++ b/ui/app/workspace/governance/views/teamsTable.tsx @@ -173,13 +173,11 @@ export default function TeamsTable({ const vks = getVirtualKeysForTeam(team.id); const customerName = getCustomerName(team.customer_id); - // Budget calculations - const isBudgetExhausted = - team.budget?.max_limit && team.budget.max_limit > 0 && team.budget.current_usage >= team.budget.max_limit; - const budgetPercentage = - team.budget?.max_limit && team.budget.max_limit > 0 - ? Math.min((team.budget.current_usage / team.budget.max_limit) * 100, 100) - : 0; + // Budget calculations — any of the team's budgets exhausted + const teamBudgets = team.budgets ?? []; + const isBudgetExhausted = teamBudgets.some( + (b) => b.max_limit > 0 && b.current_usage >= b.max_limit, + ); // Rate limit calculations const isTokenLimitExhausted = @@ -224,36 +222,47 @@ export default function TeamsTable({
- {team.budget ? ( - - -
-
- {formatCurrency(team.budget.max_limit)} - {formatResetDuration(team.budget.reset_duration)} -
- div]:bg-red-500/70" - : budgetPercentage > 80 - ? "[&>div]:bg-amber-500/70" - : "[&>div]:bg-emerald-500/70", - )} - /> -
- - -

- {formatCurrency(team.budget.current_usage)} / {formatCurrency(team.budget.max_limit)} -

-

- Resets {formatResetDuration(team.budget.reset_duration)} -

- - + {teamBudgets.length > 0 ? ( +
+ {teamBudgets.map((b) => { + const budgetPercentage = + b.max_limit > 0 ? Math.min((b.current_usage / b.max_limit) * 100, 100) : 0; + const isExhausted = b.max_limit > 0 && b.current_usage >= b.max_limit; + return ( + + +
+
+ {formatCurrency(b.max_limit)} + + {formatResetDuration(b.reset_duration)} + +
+ div]:bg-red-500/70" + : budgetPercentage > 80 + ? "[&>div]:bg-amber-500/70" + : "[&>div]:bg-emerald-500/70", + )} + /> +
+ + +

+ {formatCurrency(b.current_usage)} / {formatCurrency(b.max_limit)} +

+

+ Resets {formatResetDuration(b.reset_duration)} +

+ + + ); + })} +
) : ( - )} diff --git a/ui/components/sidebar.tsx b/ui/components/sidebar.tsx index b9c4cef00b..cb26b23bda 100644 --- a/ui/components/sidebar.tsx +++ b/ui/components/sidebar.tsx @@ -1,70 +1,83 @@ import { - ArrowUpRight, - BookUser, - Boxes, - BoxIcon, - BugIcon, - Building, - Building2, - ChartColumnBig, - ChevronsLeftRightEllipsis, - Construction, - DatabaseZap, - FlaskConical, - FolderGit, - Globe, - KeyRound, - Landmark, - LayoutGrid, - LogOut, - Logs, - Network, - PanelLeftClose, - Plug, - Puzzle, - ScrollText, - Search, - SearchCheck, - Settings, - Settings2Icon, - ShieldCheck, - ShieldUser, - Shuffle, - SlidersHorizontal, - Telescope, - ToolCase, - TrendingUp, - User, - UserRoundCheck, - Users, - Wallet, - WalletCards + ArrowUpRight, + BookUser, + Boxes, + BoxIcon, + BugIcon, + Building, + Building2, + ChartColumnBig, + ChevronsLeftRightEllipsis, + Construction, + DatabaseZap, + FlaskConical, + FolderGit, + Globe, + KeyRound, + Landmark, + LayoutGrid, + LogOut, + Logs, + Network, + PanelLeftClose, + Plug, + Puzzle, + ScrollText, + Search, + SearchCheck, + Settings, + Settings2Icon, + ShieldCheck, + ShieldUser, + Shuffle, + SlidersHorizontal, + Telescope, + ToolCase, + TrendingUp, + User, + UserRoundCheck, + Users, + Wallet, + WalletCards, } from "lucide-react"; -import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; +import { + Popover, + PopoverContent, + PopoverTrigger, +} from "@/components/ui/popover"; import { Separator } from "@/components/ui/separator"; import { - Sidebar, - SidebarContent, - SidebarGroup, - SidebarGroupContent, - SidebarHeader, - SidebarMenu, - SidebarMenuButton, - SidebarMenuItem, - SidebarMenuSub, - SidebarMenuSubButton, - SidebarMenuSubItem, - useSidebar, + Sidebar, + SidebarContent, + SidebarGroup, + SidebarGroupContent, + SidebarHeader, + SidebarMenu, + SidebarMenuButton, + SidebarMenuItem, + SidebarMenuSub, + SidebarMenuSubButton, + SidebarMenuSubItem, + useSidebar, } from "@/components/ui/sidebar"; import { useWebSocket } from "@/hooks/useWebSocket"; import { IS_ENTERPRISE, TRIAL_EXPIRY } from "@/lib/constants/config"; -import { useGetCoreConfigQuery, useGetLatestReleaseQuery, useGetVersionQuery, useLogoutMutation } from "@/lib/store"; +import { + useGetCoreConfigQuery, + useGetLatestReleaseQuery, + useGetVersionQuery, + useLogoutMutation, +} from "@/lib/store"; import { cn } from "@/lib/utils"; import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib"; import type { UserInfo } from "@enterprise/lib/store/utils/tokenManager"; import { getUserInfo } from "@enterprise/lib/store/utils/tokenManager"; -import { BooksIcon, DiscordLogoIcon, GithubLogoIcon } from "@phosphor-icons/react"; +import { + BooksIcon, + DiscordLogoIcon, + GithubLogoIcon, +} from "@phosphor-icons/react"; import { Link, useLocation, useNavigate } from "@tanstack/react-router"; import { differenceInDays } from "date-fns"; import { ChevronRight } from "lucide-react"; @@ -80,1162 +93,1371 @@ const PRODUCTION_SETUP_DISMISSED_COOKIE = "bifrost_production_setup_dismissed"; // Custom MCP Icon Component const MCPIcon = ({ className }: { className?: string }) => ( - - MCP clients icon - - - + + MCP clients icon + + + ); // Main navigation items // External links const externalLinks = [ - { - title: "Discord Server", - url: "https://discord.gg/exN5KAydbU", - icon: DiscordLogoIcon, - }, - { - title: "GitHub Repository", - url: "https://github.com/maximhq/bifrost", - icon: GithubLogoIcon, - }, - { - title: "Report a bug", - url: "https://github.com/maximhq/bifrost/issues/new?title=[Bug Report]&labels=bug&type=bug&projects=maximhq/1", - icon: BugIcon, - strokeWidth: 1.5, - }, - { - title: "Full Documentation", - url: "https://docs.getbifrost.ai", - icon: BooksIcon, - strokeWidth: 1, - }, + { + title: "Discord Server", + url: "https://discord.gg/exN5KAydbU", + icon: DiscordLogoIcon, + }, + { + title: "GitHub Repository", + url: "https://github.com/maximhq/bifrost", + icon: GithubLogoIcon, + }, + { + title: "Report a bug", + url: "https://github.com/maximhq/bifrost/issues/new?title=[Bug Report]&labels=bug&type=bug&projects=maximhq/1", + icon: BugIcon, + strokeWidth: 1.5, + }, + { + title: "Full Documentation", + url: "https://docs.getbifrost.ai", + icon: BooksIcon, + strokeWidth: 1, + }, ]; // Base promotional card (memoized outside component to prevent recreation) const productionSetupHelpCard = { - id: "production-setup", - title: "Need help with production setup?", - description: ( - <> - We offer help with production setup including custom integrations and dedicated support. -
-
- Book a demo with our team{" "} - - here - - . - - ), - dismissible: true, + id: "production-setup", + title: "Need help with production setup?", + description: ( + <> + We offer help with production setup including custom integrations and + dedicated support. +
+
+ Book a demo with our team{" "} + + here + + . + + ), + dismissible: true, }; // Sidebar item interface interface SidebarItem { - title: string; - url: string; - icon: React.ComponentType<{ className?: string }>; - description: string; - isAllowed?: boolean; - hasAccess: boolean; - subItems?: SidebarItem[]; - tag?: string; - isExternal?: boolean; - queryParam?: string; // Optional: for tab-based subitems (e.g., "client-settings") + title: string; + url: string; + icon: React.ComponentType<{ className?: string }>; + description: string; + isAllowed?: boolean; + hasAccess: boolean; + subItems?: SidebarItem[]; + tag?: string; + isExternal?: boolean; + queryParam?: string; // Optional: for tab-based subitems (e.g., "client-settings") } const getSidebarItemHref = (item: Pick) => { - return item.queryParam ? `${item.url}?tab=${item.queryParam}` : item.url; + return item.queryParam ? `${item.url}?tab=${item.queryParam}` : item.url; }; const SidebarItemView = ({ - item, - isActive, - isExternal, - isWebSocketConnected, - isExpanded, - onToggle, - pathname, - isSidebarCollapsed, - expandSidebar, - highlightedUrl, + item, + isActive, + isExternal, + isWebSocketConnected, + isExpanded, + onToggle, + pathname, + isSidebarCollapsed, + expandSidebar, + highlightedUrl, }: { - item: SidebarItem; - isActive: boolean; - isExternal?: boolean; - isWebSocketConnected: boolean; - isExpanded?: boolean; - onToggle?: () => void; - pathname: string; - isSidebarCollapsed: boolean; - expandSidebar: () => void; - highlightedUrl?: string; + item: SidebarItem; + isActive: boolean; + isExternal?: boolean; + isWebSocketConnected: boolean; + isExpanded?: boolean; + onToggle?: () => void; + pathname: string; + isSidebarCollapsed: boolean; + expandSidebar: () => void; + highlightedUrl?: string; }) => { - const hasSubItems = "subItems" in item && item.subItems && item.subItems.length > 0; - const isRouteMatch = (url: string) => { - if (url === "/workspace/custom-pricing") return pathname === url; - return pathname.startsWith(url); - }; - const isAnySubItemActive = - hasSubItems && - item.subItems?.some((subItem) => { - return isRouteMatch(subItem.url); - }); + const hasSubItems = + "subItems" in item && item.subItems && item.subItems.length > 0; + const isRouteMatch = (url: string) => { + if (url === "/workspace/custom-pricing") return pathname === url; + return pathname.startsWith(url); + }; + const isAnySubItemActive = + hasSubItems && + item.subItems?.some((subItem) => { + return isRouteMatch(subItem.url); + }); - const handleClick = (e: React.MouseEvent) => { - if (hasSubItems && item.hasAccess) { - e.preventDefault(); - // If sidebar is collapsed, expand it first then toggle the submenu - if (isSidebarCollapsed) { - expandSidebar(); - // Small delay to allow sidebar to expand before toggling submenu - setTimeout(() => { - if (onToggle) onToggle(); - }, 100); - } else if (onToggle) { - onToggle(); - } - } - }; + const handleClick = (e: React.MouseEvent) => { + if (hasSubItems && item.hasAccess) { + e.preventDefault(); + // If sidebar is collapsed, expand it first then toggle the submenu + if (isSidebarCollapsed) { + expandSidebar(); + // Small delay to allow sidebar to expand before toggling submenu + setTimeout(() => { + if (onToggle) onToggle(); + }, 100); + } else if (onToggle) { + onToggle(); + } + } + }; - const isHighlighted = !hasSubItems && highlightedUrl === item.url; + const isHighlighted = !hasSubItems && highlightedUrl === item.url; - const buttonClassName = `relative h-7.5 cursor-pointer rounded-sm border px-3 transition-all duration-200 ${ - isHighlighted - ? "bg-sidebar-accent text-accent-foreground border-primary/20" - : isActive || isAnySubItemActive - ? "bg-sidebar-accent text-primary border-primary/20" - : item.hasAccess - ? "hover:bg-sidebar-accent hover:text-accent-foreground border-transparent text-slate-500 dark:text-zinc-400" - : "hover:bg-destructive/5 hover:text-muted-foreground text-muted-foreground cursor-not-allowed border-transparent" - } `; + const buttonClassName = `relative h-7.5 cursor-pointer rounded-sm border px-3 transition-all duration-200 ${ + isHighlighted + ? "bg-sidebar-accent text-accent-foreground border-primary/20" + : isActive || isAnySubItemActive + ? "bg-sidebar-accent text-primary border-primary/20" + : item.hasAccess + ? "hover:bg-sidebar-accent hover:text-accent-foreground border-transparent text-slate-500 dark:text-zinc-400" + : "hover:bg-destructive/5 hover:text-muted-foreground text-muted-foreground cursor-not-allowed border-transparent" + } `; - const innerContent = ( -
-
- - - {item.title} - - {item.tag && ( - - {item.tag} - - )} -
- {hasSubItems && ( - - )} - {!hasSubItems && item.url === "/logs" && isWebSocketConnected && ( -
- )} - {isExternal && } -
- ); + const innerContent = ( +
+
+ + + {item.title} + + {item.tag && ( + + {item.tag} + + )} +
+ {hasSubItems && ( + + )} + {!hasSubItems && item.url === "/logs" && isWebSocketConnected && ( +
+ )} + {isExternal && ( + + )} +
+ ); - // Render strategy: - // - Items with sub-items: -
- {/* Collapsed state: vertical layout */} -
- Bifrost -
- -
-
- - { - setSearchQuery(e.target.value); - setFocusedIndex(-1); - }} - onKeyDown={handleSearchKeyDown} - className="border-input text-foreground placeholder:text-shadow-muted-foreground focus:ring-ring h-8 w-full rounded-sm border bg-transparent pr-14 pl-8 text-sm outline-none focus:bg-transparent" - /> - - ⌘ - K - -
-
- - - - - {filteredItems.map((item) => { - const isActive = isActiveRoute(item.url); + return ( + + + {/* Expanded state: horizontal layout */} +
+ + Bifrost + + +
+ {/* Collapsed state: vertical layout */} +
+ Bifrost +
+ +
+
+ + { + setSearchQuery(e.target.value); + setFocusedIndex(-1); + }} + onKeyDown={handleSearchKeyDown} + className="border-input text-foreground placeholder:text-shadow-muted-foreground focus:ring-ring h-8 w-full rounded-sm border bg-transparent pr-14 pl-8 text-sm outline-none focus:bg-transparent" + /> + + + ⌘ + + + K + + +
+
+ + + + + {filteredItems.map((item) => { + const isActive = isActiveRoute(item.url); - const highlightedUrl = focusedIndex >= 0 ? navigableItems[focusedIndex]?.url : undefined; - return ( - toggleItem(item.title)} - pathname={pathname} - isSidebarCollapsed={sidebarState === "collapsed"} - expandSidebar={() => toggleSidebar()} - highlightedUrl={highlightedUrl} - /> - ); - })} - - - -
-
- -
-
-
- {externalLinks.map((item, index) => ( - -
- -
- - ))} - - {IS_ENTERPRISE && userInfo && (userInfo.name || userInfo.email) ? ( - - - - - -
-
-

{userInfo.name || userInfo.email || "User"}

-
- - -
- - - ) : isAuthEnabled && !IS_ENTERPRISE ? ( -
- -
- ) : null} -
-
-
-
{version ?? ""}
- {trialDaysRemaining !== null && ( -
- {trialDaysRemaining} {trialDaysRemaining === 1 ? "day" : "days"} remaining -
- )} -
-
- - - ); -} \ No newline at end of file + const highlightedUrl = + focusedIndex >= 0 + ? navigableItems[focusedIndex]?.url + : undefined; + return ( + toggleItem(item.title)} + pathname={pathname} + isSidebarCollapsed={sidebarState === "collapsed"} + expandSidebar={() => toggleSidebar()} + highlightedUrl={highlightedUrl} + /> + ); + })} + + + +
+
+ +
+
+
+ {externalLinks.map((item, index) => ( + +
+ +
+ + ))} + + {IS_ENTERPRISE && + userInfo && + (userInfo.name || userInfo.email) ? ( + + + + + +
+
+

+ {userInfo.name || userInfo.email || "User"} +

+
+ + +
+ + + ) : isAuthEnabled && !IS_ENTERPRISE ? ( +
+ +
+ ) : null} +
+
+
+
{version ?? ""}
+ {trialDaysRemaining !== null && ( +
+ {trialDaysRemaining} {trialDaysRemaining === 1 ? "day" : "days"}{" "} + remaining +
+ )} +
+
+ + + ); +} diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index 6662297330..8208cb2de7 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -29,11 +29,10 @@ export interface Team { id: string; name: string; customer_id?: string; - budget_id?: string; rate_limit_id?: string; // Populated relationships customer?: Customer; - budget?: Budget; + budgets?: Budget[]; // Multi-budget: each with a distinct reset_duration rate_limit?: RateLimit; } @@ -185,14 +184,14 @@ export interface UpdateVirtualKeyRequest { export interface CreateTeamRequest { name: string; customer_id?: string; - budget?: CreateBudgetRequest; + budgets?: CreateBudgetRequest[]; // Multi-budget: each must have a unique reset_duration rate_limit?: CreateRateLimitRequest; } export interface UpdateTeamRequest { name?: string; customer_id?: string; - budget?: UpdateBudgetRequest; + budgets?: CreateBudgetRequest[]; // Replaces all team budgets; empty array clears rate_limit?: UpdateRateLimitRequest; }