diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 1ce61f3caf..b4563e17d7 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -247,6 +247,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddReplicateDeploymentsJSONColumn(ctx, db); err != nil { return err } + if err := migrationAddRateLimitToTeamsAndCustomers(ctx, db); err != nil { + return err + } return nil } @@ -3357,3 +3360,52 @@ func migrationAddReplicateDeploymentsJSONColumn(ctx context.Context, db *gorm.DB } return nil } + +// migrationAddRateLimitToTeamsAndCustomers adds rate_limit_id column to governance_teams and governance_customers tables +func migrationAddRateLimitToTeamsAndCustomers(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_rate_limit_to_teams_and_customers", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migrator := tx.Migrator() + + // Add rate_limit_id to governance_teams table + if !migrator.HasColumn(&tables.TableTeam{}, "rate_limit_id") { + if err := migrator.AddColumn(&tables.TableTeam{}, "rate_limit_id"); err != nil { + return fmt.Errorf("failed to add rate_limit_id column to teams: %w", err) + } + } + + // Add rate_limit_id to governance_customers table + if !migrator.HasColumn(&tables.TableCustomer{}, "rate_limit_id") { + if err := migrator.AddColumn(&tables.TableCustomer{}, "rate_limit_id"); err != nil { + return fmt.Errorf("failed to add rate_limit_id column to customers: %w", err) + } + } + + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migrator := tx.Migrator() + + if migrator.HasColumn(&tables.TableTeam{}, "rate_limit_id") { + if err := migrator.DropColumn(&tables.TableTeam{}, "rate_limit_id"); err != nil { + return fmt.Errorf("failed to drop rate_limit_id column from teams: %w", err) + } + } + + if migrator.HasColumn(&tables.TableCustomer{}, "rate_limit_id") { + if err := migrator.DropColumn(&tables.TableCustomer{}, "rate_limit_id"); err != nil { + return fmt.Errorf("failed to drop rate_limit_id column from customers: %w", err) + } + } + + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running rate limit migration for teams and customers: %s", err.Error()) + } + return nil +} diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index cb3ea49aa9..7b1c6d811b 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -1812,7 +1812,7 @@ func (s *RDBConfigStore) DeleteVirtualKeyMCPConfig(ctx context.Context, id uint, // GetTeams retrieves all teams from the database. func (s *RDBConfigStore) GetTeams(ctx context.Context, customerID string) ([]tables.TableTeam, error) { // Preload relationships for complete information - query := s.db.WithContext(ctx).Preload("Customer").Preload("Budget") + query := s.db.WithContext(ctx).Preload("Customer").Preload("Budget").Preload("RateLimit") // Optional filtering by customer if customerID != "" { query = query.Where("customer_id = ?", customerID) @@ -1827,7 +1827,7 @@ func (s *RDBConfigStore) GetTeams(ctx context.Context, customerID string) ([]tab // GetTeam retrieves a specific team from the database. func (s *RDBConfigStore) GetTeam(ctx context.Context, id string) (*tables.TableTeam, error) { var team tables.TableTeam - if err := s.db.WithContext(ctx).Preload("Customer").Preload("Budget").First(&team, "id = ?", id).Error; err != nil { + if err := s.db.WithContext(ctx).Preload("Customer").Preload("Budget").Preload("RateLimit").First(&team, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, ErrNotFound } @@ -1868,7 +1868,7 @@ func (s *RDBConfigStore) UpdateTeam(ctx context.Context, team *tables.TableTeam, 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").First(&team, "id = ?", id).Error; err != nil { + if err := tx.WithContext(ctx).Preload("Budget").Preload("RateLimit").First(&team, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } @@ -1878,8 +1878,9 @@ 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 ID before deleting the team + // Store the budget and rate limit IDs before deleting the team budgetID := team.BudgetID + rateLimitID := team.RateLimitID // Delete the team first if err := tx.WithContext(ctx).Delete(&tables.TableTeam{}, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { @@ -1893,6 +1894,12 @@ func (s *RDBConfigStore) DeleteTeam(ctx context.Context, id string) error { 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 { + return err + } + } return nil }); err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { @@ -1906,7 +1913,7 @@ func (s *RDBConfigStore) DeleteTeam(ctx context.Context, id string) error { // GetCustomers retrieves all customers from the database. func (s *RDBConfigStore) GetCustomers(ctx context.Context) ([]tables.TableCustomer, error) { var customers []tables.TableCustomer - if err := s.db.WithContext(ctx).Preload("Teams").Preload("Budget").Order("created_at ASC").Find(&customers).Error; err != nil { + if err := s.db.WithContext(ctx).Preload("Teams").Preload("Budget").Preload("RateLimit").Order("created_at ASC").Find(&customers).Error; err != nil { return nil, err } return customers, nil @@ -1915,7 +1922,7 @@ func (s *RDBConfigStore) GetCustomers(ctx context.Context) ([]tables.TableCustom // GetCustomer retrieves a specific customer from the database. func (s *RDBConfigStore) GetCustomer(ctx context.Context, id string) (*tables.TableCustomer, error) { var customer tables.TableCustomer - if err := s.db.WithContext(ctx).Preload("Teams").Preload("Budget").First(&customer, "id = ?", id).Error; err != nil { + if err := s.db.WithContext(ctx).Preload("Teams").Preload("Budget").Preload("RateLimit").First(&customer, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, ErrNotFound } @@ -1956,7 +1963,7 @@ func (s *RDBConfigStore) UpdateCustomer(ctx context.Context, customer *tables.Ta func (s *RDBConfigStore) DeleteCustomer(ctx context.Context, id string) error { if err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var customer tables.TableCustomer - if err := tx.WithContext(ctx).Preload("Budget").First(&customer, "id = ?", id).Error; err != nil { + if err := tx.WithContext(ctx).Preload("Budget").Preload("RateLimit").First(&customer, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } @@ -1970,8 +1977,9 @@ func (s *RDBConfigStore) DeleteCustomer(ctx context.Context, id string) error { if err := tx.WithContext(ctx).Model(&tables.TableTeam{}).Where("customer_id = ?", id).Update("customer_id", nil).Error; err != nil { return err } - // Store the budget ID before deleting the customer + // Store the budget and rate limit IDs before deleting the customer budgetID := customer.BudgetID + rateLimitID := customer.RateLimitID // Delete the customer first if err := tx.WithContext(ctx).Delete(&tables.TableCustomer{}, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { @@ -1985,6 +1993,12 @@ func (s *RDBConfigStore) DeleteCustomer(ctx context.Context, id string) error { return err } } + // Delete the customer's rate limit if it exists + if rateLimitID != nil { + if err := tx.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil { + return err + } + } return nil }); err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { @@ -2005,9 +2019,15 @@ func (s *RDBConfigStore) GetRateLimits(ctx context.Context) ([]tables.TableRateL } // GetRateLimit retrieves a specific rate limit from the database. -func (s *RDBConfigStore) GetRateLimit(ctx context.Context, id string) (*tables.TableRateLimit, error) { +func (s *RDBConfigStore) GetRateLimit(ctx context.Context, id string, tx ...*gorm.DB) (*tables.TableRateLimit, error) { + var txDB *gorm.DB + if len(tx) > 0 { + txDB = tx[0] + } else { + txDB = s.db + } var rateLimit tables.TableRateLimit - if err := s.db.WithContext(ctx).First(&rateLimit, "id = ?", id).Error; err != nil { + if err := txDB.WithContext(ctx).First(&rateLimit, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, ErrNotFound } @@ -2060,6 +2080,20 @@ func (s *RDBConfigStore) UpdateRateLimits(ctx context.Context, rateLimits []*tab return nil } +// DeleteRateLimit deletes a rate limit from the database. +func (s *RDBConfigStore) DeleteRateLimit(ctx context.Context, id string, tx ...*gorm.DB) error { + var txDB *gorm.DB + if len(tx) > 0 { + txDB = tx[0] + } else { + txDB = s.db + } + if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", id).Error; err != nil { + return s.parseGormError(err) + } + return nil +} + // GetBudgets retrieves all budgets from the database. func (s *RDBConfigStore) GetBudgets(ctx context.Context) ([]tables.TableBudget, error) { var budgets []tables.TableBudget @@ -2131,6 +2165,20 @@ func (s *RDBConfigStore) UpdateBudget(ctx context.Context, budget *tables.TableB return nil } +// DeleteBudget deletes a budget from the database. +func (s *RDBConfigStore) DeleteBudget(ctx context.Context, id string, tx ...*gorm.DB) error { + var txDB *gorm.DB + if len(tx) > 0 { + txDB = tx[0] + } else { + txDB = s.db + } + if err := txDB.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", id).Error; err != nil { + return s.parseGormError(err) + } + return nil +} + // UpdateBudgetUsage updates only the current_usage field of a budget. // Uses SkipHooks to avoid triggering BeforeSave validation since we're only updating usage. func (s *RDBConfigStore) UpdateBudgetUsage(ctx context.Context, id string, currentUsage float64) error { diff --git a/framework/configstore/store.go b/framework/configstore/store.go index ae1aee6fff..6555ffcb23 100644 --- a/framework/configstore/store.go +++ b/framework/configstore/store.go @@ -102,10 +102,11 @@ type ConfigStore interface { // Rate limit CRUD GetRateLimits(ctx context.Context) ([]tables.TableRateLimit, error) - GetRateLimit(ctx context.Context, id string) (*tables.TableRateLimit, error) + GetRateLimit(ctx context.Context, id string, tx ...*gorm.DB) (*tables.TableRateLimit, error) CreateRateLimit(ctx context.Context, rateLimit *tables.TableRateLimit, tx ...*gorm.DB) error UpdateRateLimit(ctx context.Context, rateLimit *tables.TableRateLimit, tx ...*gorm.DB) error UpdateRateLimits(ctx context.Context, rateLimits []*tables.TableRateLimit, tx ...*gorm.DB) error + DeleteRateLimit(ctx context.Context, id string, tx ...*gorm.DB) error // Budget CRUD GetBudgets(ctx context.Context) ([]tables.TableBudget, error) @@ -113,6 +114,7 @@ type ConfigStore interface { CreateBudget(ctx context.Context, budget *tables.TableBudget, tx ...*gorm.DB) error UpdateBudget(ctx context.Context, budget *tables.TableBudget, tx ...*gorm.DB) error UpdateBudgets(ctx context.Context, budgets []*tables.TableBudget, tx ...*gorm.DB) error + DeleteBudget(ctx context.Context, id string, tx ...*gorm.DB) error UpdateBudgetUsage(ctx context.Context, id string, currentUsage float64) error UpdateRateLimitUsage(ctx context.Context, id string, tokenCurrentUsage int64, requestCurrentUsage int64) error diff --git a/framework/configstore/tables/customer.go b/framework/configstore/tables/customer.go index 0f5c920de6..c3b9a6f08c 100644 --- a/framework/configstore/tables/customer.go +++ b/framework/configstore/tables/customer.go @@ -2,14 +2,16 @@ package tables import "time" -// TableCustomer represents a customer entity with budget +// TableCustomer represents a customer entity with budget and rate limit type TableCustomer struct { - ID string `gorm:"primaryKey;type:varchar(255)" json:"id"` - Name string `gorm:"type:varchar(255);not null" json:"name"` - BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"` + ID string `gorm:"primaryKey;type:varchar(255)" json:"id"` + Name string `gorm:"type:varchar(255);not null" json:"name"` + BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"` + RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"` // Relationships Budget *TableBudget `gorm:"foreignKey:BudgetID" json:"budget,omitempty"` + RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID" json:"rate_limit,omitempty"` Teams []TableTeam `gorm:"foreignKey:CustomerID" json:"teams"` VirtualKeys []TableVirtualKey `gorm:"foreignKey:CustomerID" json:"virtual_keys"` diff --git a/framework/configstore/tables/team.go b/framework/configstore/tables/team.go index 72988b6ffe..3054263ce8 100644 --- a/framework/configstore/tables/team.go +++ b/framework/configstore/tables/team.go @@ -8,16 +8,18 @@ import ( "gorm.io/gorm" ) -// TableTeam represents a team entity with budget and customer association +// TableTeam represents a team entity with budget, rate limit and customer association 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"` + 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"` + RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID" json:"rate_limit,omitempty"` VirtualKeys []TableVirtualKey `gorm:"foreignKey:TeamID" json:"virtual_keys"` Profile *string `gorm:"type:text" json:"-"` diff --git a/plugins/governance/store.go b/plugins/governance/store.go index d97f52f8ae..5c8be2195a 100644 --- a/plugins/governance/store.go +++ b/plugins/governance/store.go @@ -224,7 +224,23 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData { if !ok || team == nil { return true // continue } - teams[key.(string)] = team + // Cross-reference live budget/rate limit from standalone maps + clone := *team + if clone.BudgetID != nil { + if liveBudget, exists := gs.budgets.Load(*clone.BudgetID); exists && liveBudget != nil { + if b, ok := liveBudget.(*configstoreTables.TableBudget); ok { + clone.Budget = b + } + } + } + if clone.RateLimitID != nil { + if liveRL, exists := gs.rateLimits.Load(*clone.RateLimitID); exists && liveRL != nil { + if rl, ok := liveRL.(*configstoreTables.TableRateLimit); ok { + clone.RateLimit = rl + } + } + } + teams[key.(string)] = &clone return true // continue iteration }) customers := make(map[string]*configstoreTables.TableCustomer) @@ -233,7 +249,23 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData { if !ok || customer == nil { return true // continue } - customers[key.(string)] = customer + // Cross-reference live budget/rate limit from standalone maps + clone := *customer + if clone.BudgetID != nil { + if liveBudget, exists := gs.budgets.Load(*clone.BudgetID); exists && liveBudget != nil { + if b, ok := liveBudget.(*configstoreTables.TableBudget); ok { + clone.Budget = b + } + } + } + if clone.RateLimitID != nil { + if liveRL, exists := gs.rateLimits.Load(*clone.RateLimitID); exists && liveRL != nil { + if rl, ok := liveRL.(*configstoreTables.TableRateLimit); ok { + clone.RateLimit = rl + } + } + } + customers[key.(string)] = &clone return true // continue iteration }) budgets := make(map[string]*configstoreTables.TableBudget) @@ -1387,6 +1419,30 @@ func (gs *LocalGovernanceStore) DumpRateLimits(ctx context.Context, tokenBaselin return true // continue }) + // Collect rate limit IDs from teams + gs.teams.Range(func(key, value interface{}) bool { + team, ok := value.(*configstoreTables.TableTeam) + if !ok || team == nil { + return true // continue + } + if team.RateLimitID != nil { + rateLimitIDs[*team.RateLimitID] = true + } + return true // continue + }) + + // Collect rate limit IDs from customers + gs.customers.Range(func(key, value interface{}) bool { + customer, ok := value.(*configstoreTables.TableCustomer) + if !ok || customer == nil { + return true // continue + } + if customer.RateLimitID != nil { + rateLimitIDs[*customer.RateLimitID] = true + } + return true // continue + }) + // Prepare rate limit usage updates with baselines type rateLimitUpdate struct { ID string @@ -1857,7 +1913,7 @@ func (gs *LocalGovernanceStore) rebuildInMemoryStructures(ctx context.Context, c // UTILITY FUNCTIONS -// collectRateLimitsFromHierarchy collects rate limits and their metadata from the hierarchy (Provider Configs → VK) +// collectRateLimitsFromHierarchy collects rate limits and their metadata from the hierarchy (Provider Configs → VK → Team → Customer) func (gs *LocalGovernanceStore) collectRateLimitsFromHierarchy(vk *configstoreTables.TableVirtualKey, requestedProvider schemas.ModelProvider) ([]*configstoreTables.TableRateLimit, []string) { if vk == nil { return nil, nil @@ -1886,6 +1942,56 @@ func (gs *LocalGovernanceStore) collectRateLimitsFromHierarchy(vk *configstoreTa } } + // Check Team rate limit if VK belongs to a team + var teamCustomerID string + 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.RateLimitID != nil { + if rateLimitValue, exists := gs.rateLimits.Load(*team.RateLimitID); exists && rateLimitValue != nil { + if rateLimit, ok := rateLimitValue.(*configstoreTables.TableRateLimit); ok && rateLimit != nil { + rateLimits = append(rateLimits, rateLimit) + rateLimitNames = append(rateLimitNames, "Team") + } + } + } + + // Check if team belongs to a customer + if team.CustomerID != nil { + teamCustomerID = *team.CustomerID + if customerValue, exists := gs.customers.Load(*team.CustomerID); exists && customerValue != nil { + if customer, ok := customerValue.(*configstoreTables.TableCustomer); ok && customer != nil { + if customer.RateLimitID != nil { + if rateLimitValue, exists := gs.rateLimits.Load(*customer.RateLimitID); exists && rateLimitValue != nil { + if rateLimit, ok := rateLimitValue.(*configstoreTables.TableRateLimit); ok && rateLimit != nil { + rateLimits = append(rateLimits, rateLimit) + rateLimitNames = append(rateLimitNames, "Customer") + } + } + } + } + } + } + } + } + } + + // Check Customer rate limit if VK directly belongs to a customer (skip if already collected via team) + if vk.CustomerID != nil && (teamCustomerID == "" || *vk.CustomerID != teamCustomerID) { + if customerValue, exists := gs.customers.Load(*vk.CustomerID); exists && customerValue != nil { + if customer, ok := customerValue.(*configstoreTables.TableCustomer); ok && customer != nil { + if customer.RateLimitID != nil { + if rateLimitValue, exists := gs.rateLimits.Load(*customer.RateLimitID); exists && rateLimitValue != nil { + if rateLimit, ok := rateLimitValue.(*configstoreTables.TableRateLimit); ok && rateLimit != nil { + rateLimits = append(rateLimits, rateLimit) + rateLimitNames = append(rateLimitNames, "Customer") + } + } + } + } + } + } + return rateLimits, rateLimitNames } @@ -1919,6 +2025,7 @@ func (gs *LocalGovernanceStore) collectBudgetsFromHierarchy(vk *configstoreTable } } + var teamCustomerID string if vk.TeamID != nil { if teamValue, exists := gs.teams.Load(*vk.TeamID); exists && teamValue != nil { if team, ok := teamValue.(*configstoreTables.TableTeam); ok && team != nil { @@ -1933,6 +2040,7 @@ func (gs *LocalGovernanceStore) collectBudgetsFromHierarchy(vk *configstoreTable // Check if team belongs to a customer if team.CustomerID != nil { + teamCustomerID = *team.CustomerID if customerValue, exists := gs.customers.Load(*team.CustomerID); exists && customerValue != nil { if customer, ok := customerValue.(*configstoreTables.TableCustomer); ok && customer != nil { if customer.BudgetID != nil { @@ -1950,7 +2058,8 @@ func (gs *LocalGovernanceStore) collectBudgetsFromHierarchy(vk *configstoreTable } } - if vk.CustomerID != nil { + // Check Customer budget if VK directly belongs to a customer (skip if already collected via team) + if vk.CustomerID != nil && (teamCustomerID == "" || *vk.CustomerID != teamCustomerID) { if customerValue, exists := gs.customers.Load(*vk.CustomerID); exists && customerValue != nil { if customer, ok := customerValue.(*configstoreTables.TableCustomer); ok && customer != nil { if customer.BudgetID != nil { @@ -2030,6 +2139,7 @@ func (gs *LocalGovernanceStore) UpdateVirtualKeyInMemory(vk *configstoreTables.T if vk == nil { return // Nothing to update } + // Do not update the current usage of the rate limit, as it will be updated by the usage tracker. // But update if max limit or reset duration changes. if existingVKValue, exists := gs.virtualKeys.Load(vk.Value); exists && existingVKValue != nil { @@ -2037,6 +2147,7 @@ func (gs *LocalGovernanceStore) UpdateVirtualKeyInMemory(vk *configstoreTables.T if !ok || existingVK == nil { return // Nothing to update } + // Create clone to avoid modifying the original clone := *vk // Update Budget for VK in memory store @@ -2186,6 +2297,11 @@ func (gs *LocalGovernanceStore) CreateTeamInMemory(team *configstoreTables.Table gs.budgets.Store(team.Budget.ID, team.Budget) } + // Create associated rate limit if exists + if team.RateLimit != nil { + gs.rateLimits.Store(team.RateLimit.ID, team.RateLimit) + } + gs.teams.Store(team.ID, team) } @@ -2194,12 +2310,14 @@ func (gs *LocalGovernanceStore) UpdateTeamInMemory(team *configstoreTables.Table if team == nil { return // Nothing to update } + // Check if there's an existing team to get current budget state if existingTeamValue, exists := gs.teams.Load(team.ID); exists && existingTeamValue != nil { existingTeam, ok := existingTeamValue.(*configstoreTables.TableTeam) if !ok || existingTeam == nil { return // Nothing to update } + // Create clone to avoid modifying the original clone := *team @@ -2219,6 +2337,24 @@ func (gs *LocalGovernanceStore) UpdateTeamInMemory(team *configstoreTables.Table gs.budgets.Delete(existingTeam.Budget.ID) } + // Handle rate limit updates with consistent logic + if clone.RateLimit != nil { + // Preserve existing usage from memory when updating team rate limit config + if existingRateLimitValue, exists := gs.rateLimits.Load(clone.RateLimit.ID); exists && existingRateLimitValue != nil { + if existingRateLimit, ok := existingRateLimitValue.(*configstoreTables.TableRateLimit); ok && existingRateLimit != nil { + // Preserve current usage and last reset time from existing in-memory rate limit + clone.RateLimit.TokenCurrentUsage = existingRateLimit.TokenCurrentUsage + clone.RateLimit.TokenLastReset = existingRateLimit.TokenLastReset + clone.RateLimit.RequestCurrentUsage = existingRateLimit.RequestCurrentUsage + clone.RateLimit.RequestLastReset = existingRateLimit.RequestLastReset + } + } + gs.rateLimits.Store(clone.RateLimit.ID, clone.RateLimit) + } else if existingTeam.RateLimit != nil { + // Rate limit was removed from the team, delete it from memory + gs.rateLimits.Delete(existingTeam.RateLimit.ID) + } + gs.teams.Store(team.ID, &clone) } else { gs.CreateTeamInMemory(team) @@ -2231,13 +2367,17 @@ func (gs *LocalGovernanceStore) DeleteTeamInMemory(teamID string) { return // Nothing to delete } - // Get team to check for associated budget + // Get team to check for associated budget 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 associated rate limit if exists + if team.RateLimitID != nil { + gs.rateLimits.Delete(*team.RateLimitID) + } } } @@ -2271,6 +2411,11 @@ func (gs *LocalGovernanceStore) CreateCustomerInMemory(customer *configstoreTabl gs.budgets.Store(customer.Budget.ID, customer.Budget) } + // Create associated rate limit if exists + if customer.RateLimit != nil { + gs.rateLimits.Store(customer.RateLimit.ID, customer.RateLimit) + } + gs.customers.Store(customer.ID, customer) } @@ -2304,6 +2449,24 @@ func (gs *LocalGovernanceStore) UpdateCustomerInMemory(customer *configstoreTabl gs.budgets.Delete(existingCustomer.Budget.ID) } + // Handle rate limit updates with consistent logic + if clone.RateLimit != nil { + // Preserve existing usage from memory when updating customer rate limit config + if existingRateLimitValue, exists := gs.rateLimits.Load(clone.RateLimit.ID); exists && existingRateLimitValue != nil { + if existingRateLimit, ok := existingRateLimitValue.(*configstoreTables.TableRateLimit); ok && existingRateLimit != nil { + // Preserve current usage and last reset time from existing in-memory rate limit + clone.RateLimit.TokenCurrentUsage = existingRateLimit.TokenCurrentUsage + clone.RateLimit.TokenLastReset = existingRateLimit.TokenLastReset + clone.RateLimit.RequestCurrentUsage = existingRateLimit.RequestCurrentUsage + clone.RateLimit.RequestLastReset = existingRateLimit.RequestLastReset + } + } + gs.rateLimits.Store(clone.RateLimit.ID, clone.RateLimit) + } else if existingCustomer.RateLimit != nil { + // Rate limit was removed from the customer, delete it from memory + gs.rateLimits.Delete(existingCustomer.RateLimit.ID) + } + gs.customers.Store(customer.ID, &clone) } else { gs.CreateCustomerInMemory(customer) @@ -2316,13 +2479,17 @@ func (gs *LocalGovernanceStore) DeleteCustomerInMemory(customerID string) { return // Nothing to delete } - // Get customer to check for associated budget + // Get customer to check for associated budget and rate limit if customerValue, exists := gs.customers.Load(customerID); exists && customerValue != nil { if customer, ok := customerValue.(*configstoreTables.TableCustomer); ok && customer != nil { // Delete associated budget if exists if customer.BudgetID != nil { gs.budgets.Delete(*customer.BudgetID) } + // Delete associated rate limit if exists + if customer.RateLimitID != nil { + gs.rateLimits.Delete(*customer.RateLimitID) + } } } @@ -2567,7 +2734,7 @@ func (gs *LocalGovernanceStore) updateBudgetReferences(resetBudget *configstoreT }) } -// updateRateLimitReferences updates all VKs and provider configs that reference a reset rate limit +// updateRateLimitReferences updates all VKs, teams, customers, users and provider configs that reference a reset rate limit func (gs *LocalGovernanceStore) updateRateLimitReferences(resetRateLimit *configstoreTables.TableRateLimit) { rateLimitID := resetRateLimit.ID // Update VKs that reference this rate limit @@ -2600,6 +2767,34 @@ func (gs *LocalGovernanceStore) updateRateLimitReferences(resetRateLimit *config } return true // continue }) + + // Update teams that reference this rate limit + gs.teams.Range(func(key, value interface{}) bool { + team, ok := value.(*configstoreTables.TableTeam) + if !ok || team == nil { + return true // continue + } + if team.RateLimitID != nil && *team.RateLimitID == rateLimitID { + clone := *team + clone.RateLimit = resetRateLimit + gs.teams.Store(key, &clone) + } + return true // continue + }) + + // Update customers that reference this rate limit + gs.customers.Range(func(key, value interface{}) bool { + customer, ok := value.(*configstoreTables.TableCustomer) + if !ok || customer == nil { + return true // continue + } + if customer.RateLimitID != nil && *customer.RateLimitID == rateLimitID { + clone := *customer + clone.RateLimit = resetRateLimit + gs.customers.Store(key, &clone) + } + return true // continue + }) } // HasRoutingRules checks if there are any routing rules configured diff --git a/plugins/governance/tracker.go b/plugins/governance/tracker.go index d53538233f..7e52f76ee9 100644 --- a/plugins/governance/tracker.go +++ b/plugins/governance/tracker.go @@ -109,8 +109,9 @@ func (t *UsageTracker) UpdateUsage(ctx context.Context, update *UsageUpdate) { return } - // Update rate limit usage (both provider-level and VK-level) if applicable - if vk.RateLimit != nil || len(vk.ProviderConfigs) > 0 { + // Update rate limit usage (VK-level, provider-config-level, team-level, customer-level) if applicable + // Include TeamID and CustomerID checks since rate limits can be configured at those levels + if vk.RateLimit != nil || len(vk.ProviderConfigs) > 0 || vk.TeamID != nil || vk.CustomerID != nil { if err := t.store.UpdateVirtualKeyRateLimitUsageInMemory(ctx, vk, update.Provider, update.TokensUsed, shouldUpdateTokens, shouldUpdateRequests); err != nil { t.logger.Error("failed to update rate limit usage for VK %s: %v", vk.ID, err) } diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index a2419cc7d2..d07fcedbab 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -168,28 +168,32 @@ type UpdateRateLimitRequest struct { // CreateTeamRequest represents the request body for creating a team 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 + 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 + RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` // Team can have its own rate limit } // UpdateTeamRequest represents the request body for updating a team type UpdateTeamRequest struct { - Name *string `json:"name,omitempty"` - CustomerID *string `json:"customer_id,omitempty"` - Budget *UpdateBudgetRequest `json:"budget,omitempty"` + Name *string `json:"name,omitempty"` + CustomerID *string `json:"customer_id,omitempty"` + Budget *UpdateBudgetRequest `json:"budget,omitempty"` + RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` } // CreateCustomerRequest represents the request body for creating a customer type CreateCustomerRequest struct { - Name string `json:"name" validate:"required"` - Budget *CreateBudgetRequest `json:"budget,omitempty"` + Name string `json:"name" validate:"required"` + Budget *CreateBudgetRequest `json:"budget,omitempty"` + RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` // Customer can have its own rate limit } // UpdateCustomerRequest represents the request body for updating a customer type UpdateCustomerRequest struct { - Name *string `json:"name,omitempty"` - Budget *UpdateBudgetRequest `json:"budget,omitempty"` + Name *string `json:"name,omitempty"` + Budget *UpdateBudgetRequest `json:"budget,omitempty"` + RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` } // CreateModelConfigRequest represents the request body for creating a model config @@ -1092,6 +1096,19 @@ func (h *GovernanceHandler) createTeam(ctx *fasthttp.RequestCtx) { return } } + // Validate rate limit if provided + if req.RateLimit != nil { + rateLimit := configstoreTables.TableRateLimit{ + TokenMaxLimit: req.RateLimit.TokenMaxLimit, + TokenResetDuration: req.RateLimit.TokenResetDuration, + RequestMaxLimit: req.RateLimit.RequestMaxLimit, + RequestResetDuration: req.RateLimit.RequestResetDuration, + } + if err := validateRateLimit(&rateLimit); err != nil { + SendError(ctx, 400, fmt.Sprintf("Invalid rate limit: %s", err.Error())) + return + } + } // Creating team in database var team configstoreTables.TableTeam if err := h.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error { @@ -1113,6 +1130,21 @@ func (h *GovernanceHandler) createTeam(ctx *fasthttp.RequestCtx) { } team.BudgetID = &budget.ID } + if req.RateLimit != nil { + rateLimit := configstoreTables.TableRateLimit{ + ID: uuid.NewString(), + TokenMaxLimit: req.RateLimit.TokenMaxLimit, + TokenResetDuration: req.RateLimit.TokenResetDuration, + RequestMaxLimit: req.RateLimit.RequestMaxLimit, + RequestResetDuration: req.RateLimit.RequestResetDuration, + TokenLastReset: time.Now(), + RequestLastReset: time.Now(), + } + if err := h.configStore.CreateRateLimit(ctx, &rateLimit, tx); err != nil { + return err + } + team.RateLimitID = &rateLimit.ID + } if err := h.configStore.CreateTeam(ctx, &team, tx); err != nil { return err } @@ -1190,6 +1222,9 @@ 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 + // Update fields if provided if req.Name != nil { team.Name = *req.Name @@ -1199,25 +1234,44 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { } // Handle budget updates if req.Budget != nil { - if team.BudgetID != nil { + // Check if budget limit is empty - means remove budget (reset duration doesn't matter) + budgetIsEmpty := req.Budget.MaxLimit == nil + 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 - budget, err := h.configStore.GetBudget(ctx, *team.BudgetID, tx) - if err != nil { - return err + if req.Budget.MaxLimit == nil || req.Budget.ResetDuration == nil { + return fmt.Errorf("both max_limit and reset_duration are required when updating a budget") } - if req.Budget.MaxLimit != nil { - budget.MaxLimit = *req.Budget.MaxLimit + budget := configstoreTables.TableBudget{} + if err := tx.First(&budget, "id = ?", *team.BudgetID).Error; err != nil { + return err } - if req.Budget.ResetDuration != nil { - budget.ResetDuration = *req.Budget.ResetDuration + budget.MaxLimit = *req.Budget.MaxLimit + budget.ResetDuration = *req.Budget.ResetDuration + if err := validateBudget(&budget); err != nil { + return err } - - if err := h.configStore.UpdateBudget(ctx, budget, tx); err != nil { + if err := h.configStore.UpdateBudget(ctx, &budget, tx); err != nil { return err } - team.Budget = budget + 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) + } + if _, err := configstoreTables.ParseDuration(*req.Budget.ResetDuration); err != nil { + return fmt.Errorf("invalid reset duration format: %s", *req.Budget.ResetDuration) + } budget := configstoreTables.TableBudget{ ID: uuid.NewString(), MaxLimit: *req.Budget.MaxLimit, @@ -1225,6 +1279,9 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { LastReset: time.Now(), CurrentUsage: 0, } + if err := validateBudget(&budget); err != nil { + return err + } if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil { return err } @@ -1232,9 +1289,71 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { team.Budget = &budget } } + // Handle rate limit updates + if req.RateLimit != nil { + // Check if rate limit values are empty - means remove rate limit (reset durations don't matter) + rateLimitIsEmpty := req.RateLimit.TokenMaxLimit == nil && req.RateLimit.RequestMaxLimit == nil + if rateLimitIsEmpty { + // Mark rate limit for deletion after FK is removed + if team.RateLimitID != nil { + rateLimitIDToDelete = *team.RateLimitID + team.RateLimitID = nil + team.RateLimit = nil + } + } else if team.RateLimitID != nil { + // Update existing rate limit + rateLimit := configstoreTables.TableRateLimit{} + if err := tx.First(&rateLimit, "id = ?", *team.RateLimitID).Error; err != nil { + return err + } + rateLimit.TokenMaxLimit = req.RateLimit.TokenMaxLimit + rateLimit.TokenResetDuration = req.RateLimit.TokenResetDuration + rateLimit.RequestMaxLimit = req.RateLimit.RequestMaxLimit + rateLimit.RequestResetDuration = req.RateLimit.RequestResetDuration + if err := validateRateLimit(&rateLimit); err != nil { + return err + } + if err := h.configStore.UpdateRateLimit(ctx, &rateLimit, tx); err != nil { + return err + } + team.RateLimit = &rateLimit + } else { + // Create new rate limit + rateLimit := configstoreTables.TableRateLimit{ + ID: uuid.NewString(), + TokenMaxLimit: req.RateLimit.TokenMaxLimit, + TokenResetDuration: req.RateLimit.TokenResetDuration, + RequestMaxLimit: req.RateLimit.RequestMaxLimit, + RequestResetDuration: req.RateLimit.RequestResetDuration, + TokenLastReset: time.Now(), + RequestLastReset: time.Now(), + } + if err := validateRateLimit(&rateLimit); err != nil { + return err + } + if err := h.configStore.CreateRateLimit(ctx, &rateLimit, tx); err != nil { + return err + } + team.RateLimitID = &rateLimit.ID + team.RateLimit = &rateLimit + } + } if err := h.configStore.UpdateTeam(ctx, team, tx); err != nil { 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 + } + } + if rateLimitIDToDelete != "" { + if err := tx.Delete(&configstoreTables.TableRateLimit{}, "id = ?", rateLimitIDToDelete).Error; err != nil { + return err + } + } + return nil }); err != nil { SendError(ctx, 500, "Failed to update team") @@ -1337,6 +1456,19 @@ func (h *GovernanceHandler) createCustomer(ctx *fasthttp.RequestCtx) { return } } + // Validate rate limit if provided + if req.RateLimit != nil { + rateLimit := configstoreTables.TableRateLimit{ + TokenMaxLimit: req.RateLimit.TokenMaxLimit, + TokenResetDuration: req.RateLimit.TokenResetDuration, + RequestMaxLimit: req.RateLimit.RequestMaxLimit, + RequestResetDuration: req.RateLimit.RequestResetDuration, + } + if err := validateRateLimit(&rateLimit); err != nil { + SendError(ctx, 400, fmt.Sprintf("Invalid rate limit: %s", err.Error())) + return + } + } var customer configstoreTables.TableCustomer if err := h.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error { customer = configstoreTables.TableCustomer{ @@ -1357,6 +1489,21 @@ func (h *GovernanceHandler) createCustomer(ctx *fasthttp.RequestCtx) { } customer.BudgetID = &budget.ID } + if req.RateLimit != nil { + rateLimit := configstoreTables.TableRateLimit{ + ID: uuid.NewString(), + TokenMaxLimit: req.RateLimit.TokenMaxLimit, + TokenResetDuration: req.RateLimit.TokenResetDuration, + RequestMaxLimit: req.RateLimit.RequestMaxLimit, + RequestResetDuration: req.RateLimit.RequestResetDuration, + TokenLastReset: time.Now(), + RequestLastReset: time.Now(), + } + if err := h.configStore.CreateRateLimit(ctx, &rateLimit, tx); err != nil { + return err + } + customer.RateLimitID = &rateLimit.ID + } if err := h.configStore.CreateCustomer(ctx, &customer, tx); err != nil { return err } @@ -1431,32 +1578,53 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { } // Updating customer in database if err := h.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error { + // Track IDs to delete after updating the customer (to avoid FK constraint) + var budgetIDToDelete, rateLimitIDToDelete string + // Update fields if provided if req.Name != nil { customer.Name = *req.Name } // Handle budget updates if req.Budget != nil { - if customer.BudgetID != nil { + // Check if budget limit is empty - means remove budget (reset duration doesn't matter) + budgetIsEmpty := req.Budget.MaxLimit == nil + if budgetIsEmpty { + // Mark budget for deletion after FK is removed + if customer.BudgetID != nil { + budgetIDToDelete = *customer.BudgetID + customer.BudgetID = nil + customer.Budget = nil + } + } else if customer.BudgetID != nil { // Update existing budget - budget, err := h.configStore.GetBudget(ctx, *customer.BudgetID, tx) - if err != nil { - return err + if req.Budget.MaxLimit == nil || req.Budget.ResetDuration == nil { + return fmt.Errorf("both max_limit and reset_duration are required when updating a budget") } - - if req.Budget.MaxLimit != nil { - budget.MaxLimit = *req.Budget.MaxLimit + budget := configstoreTables.TableBudget{} + if err := tx.First(&budget, "id = ?", *customer.BudgetID).Error; err != nil { + return err } - if req.Budget.ResetDuration != nil { - budget.ResetDuration = *req.Budget.ResetDuration + budget.MaxLimit = *req.Budget.MaxLimit + budget.ResetDuration = *req.Budget.ResetDuration + if err := validateBudget(&budget); err != nil { + return err } - - if err := h.configStore.UpdateBudget(ctx, budget, tx); err != nil { + if err := h.configStore.UpdateBudget(ctx, &budget, tx); err != nil { return err } - customer.Budget = budget + customer.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) + } + if _, err := configstoreTables.ParseDuration(*req.Budget.ResetDuration); err != nil { + return fmt.Errorf("invalid reset duration format: %s", *req.Budget.ResetDuration) + } budget := configstoreTables.TableBudget{ ID: uuid.NewString(), MaxLimit: *req.Budget.MaxLimit, @@ -1464,6 +1632,9 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { LastReset: time.Now(), CurrentUsage: 0, } + if err := validateBudget(&budget); err != nil { + return err + } if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil { return err } @@ -1471,9 +1642,71 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { customer.Budget = &budget } } + // Handle rate limit updates + if req.RateLimit != nil { + // Check if rate limit values are empty - means remove rate limit (reset durations don't matter) + rateLimitIsEmpty := req.RateLimit.TokenMaxLimit == nil && req.RateLimit.RequestMaxLimit == nil + if rateLimitIsEmpty { + // Mark rate limit for deletion after FK is removed + if customer.RateLimitID != nil { + rateLimitIDToDelete = *customer.RateLimitID + customer.RateLimitID = nil + customer.RateLimit = nil + } + } else if customer.RateLimitID != nil { + // Update existing rate limit + rateLimit := configstoreTables.TableRateLimit{} + if err := tx.First(&rateLimit, "id = ?", *customer.RateLimitID).Error; err != nil { + return err + } + rateLimit.TokenMaxLimit = req.RateLimit.TokenMaxLimit + rateLimit.TokenResetDuration = req.RateLimit.TokenResetDuration + rateLimit.RequestMaxLimit = req.RateLimit.RequestMaxLimit + rateLimit.RequestResetDuration = req.RateLimit.RequestResetDuration + if err := validateRateLimit(&rateLimit); err != nil { + return err + } + if err := h.configStore.UpdateRateLimit(ctx, &rateLimit, tx); err != nil { + return err + } + customer.RateLimit = &rateLimit + } else { + // Create new rate limit + rateLimit := configstoreTables.TableRateLimit{ + ID: uuid.NewString(), + TokenMaxLimit: req.RateLimit.TokenMaxLimit, + TokenResetDuration: req.RateLimit.TokenResetDuration, + RequestMaxLimit: req.RateLimit.RequestMaxLimit, + RequestResetDuration: req.RateLimit.RequestResetDuration, + TokenLastReset: time.Now(), + RequestLastReset: time.Now(), + } + if err := validateRateLimit(&rateLimit); err != nil { + return err + } + if err := h.configStore.CreateRateLimit(ctx, &rateLimit, tx); err != nil { + return err + } + customer.RateLimitID = &rateLimit.ID + customer.RateLimit = &rateLimit + } + } if err := h.configStore.UpdateCustomer(ctx, customer, tx); err != nil { 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 + } + } + if rateLimitIDToDelete != "" { + if err := tx.Delete(&configstoreTables.TableRateLimit{}, "id = ?", rateLimitIDToDelete).Error; err != nil { + return err + } + } + return nil }); err != nil { SendError(ctx, 500, "Failed to update customer") diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index de864cfe92..c45aa70c52 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -568,10 +568,38 @@ func (m *MockConfigStore) UpdateRateLimits(ctx context.Context, rateLimits []*ta return nil } -func (m *MockConfigStore) GetRateLimit(ctx context.Context, id string) (*tables.TableRateLimit, error) { +func (m *MockConfigStore) GetRateLimit(ctx context.Context, id string, tx ...*gorm.DB) (*tables.TableRateLimit, error) { return nil, nil } +func (m *MockConfigStore) DeleteRateLimit(ctx context.Context, id string, tx ...*gorm.DB) error { + if m.governanceConfig == nil || len(m.governanceConfig.RateLimits) == 0 { + return nil + } + filtered := make([]tables.TableRateLimit, 0, len(m.governanceConfig.RateLimits)) + for _, rl := range m.governanceConfig.RateLimits { + if rl.ID != id { + filtered = append(filtered, rl) + } + } + m.governanceConfig.RateLimits = filtered + return nil +} + +func (m *MockConfigStore) DeleteBudget(ctx context.Context, id string, tx ...*gorm.DB) error { + if m.governanceConfig == nil || len(m.governanceConfig.Budgets) == 0 { + return nil + } + filtered := make([]tables.TableBudget, 0, len(m.governanceConfig.Budgets)) + for _, b := range m.governanceConfig.Budgets { + if b.ID != id { + filtered = append(filtered, b) + } + } + m.governanceConfig.Budgets = filtered + return nil +} + func (m *MockConfigStore) GetRateLimits(ctx context.Context) ([]tables.TableRateLimit, error) { return []tables.TableRateLimit{}, nil } @@ -14972,6 +15000,7 @@ var excludedGoFields = map[string]map[string]bool{ "created_at": true, "updated_at": true, "budget": true, // GORM relation + "rate_limit": true, // GORM relation "teams": true, // GORM relation "virtual_keys": true, // GORM relation }, @@ -14980,6 +15009,7 @@ var excludedGoFields = map[string]map[string]bool{ "created_at": true, "updated_at": true, "budget": true, // GORM relation + "rate_limit": true, // GORM relation "customer": true, // GORM relation "virtual_keys": true, // GORM relation }, diff --git a/transports/config.schema.json b/transports/config.schema.json index 0b496f652b..2d4acbade0 100644 --- a/transports/config.schema.json +++ b/transports/config.schema.json @@ -296,6 +296,10 @@ "budget_id": { "type": "string", "description": "Associated budget ID" + }, + "rate_limit_id": { + "type": "string", + "description": "Associated rate limit ID" } }, "required": [ @@ -327,6 +331,10 @@ "type": "string", "description": "Associated budget ID" }, + "rate_limit_id": { + "type": "string", + "description": "Associated rate limit ID" + }, "profile": { "type": "object", "description": "Team profile data" diff --git a/ui/app/workspace/model-limits/views/modelLimitsTable.tsx b/ui/app/workspace/model-limits/views/modelLimitsTable.tsx index a50b2e8be0..3ef4009368 100644 --- a/ui/app/workspace/model-limits/views/modelLimitsTable.tsx +++ b/ui/app/workspace/model-limits/views/modelLimitsTable.tsx @@ -300,7 +300,7 @@ export default function ModelLimitsTable({ modelConfigs }: ModelLimitsTableProps )} e.stopPropagation()}> -
+
diff --git a/ui/app/workspace/user-groups/views/customerDialog.tsx b/ui/app/workspace/user-groups/views/customerDialog.tsx index 61007f0027..5ff9e0a1b4 100644 --- a/ui/app/workspace/user-groups/views/customerDialog.tsx +++ b/ui/app/workspace/user-groups/views/customerDialog.tsx @@ -25,9 +25,14 @@ interface CustomerDialogProps { interface CustomerFormData { name: string; - // Budget - budgetMaxLimit: number | undefined; + // Budget (stored as string to allow intermediate decimal states like "1.") + budgetMaxLimit: string; budgetResetDuration: string; + // Rate Limit (stored as string) + tokenMaxLimit: string; + tokenResetDuration: string; + requestMaxLimit: string; + requestResetDuration: string; isDirty: boolean; } @@ -35,9 +40,14 @@ interface CustomerFormData { const createInitialState = (customer?: Customer | null): Omit => { return { name: customer?.name || "", - // Budget - budgetMaxLimit: customer?.budget ? customer.budget.max_limit : undefined, // Already in dollars + // Budget (stored as string) + budgetMaxLimit: customer?.budget ? String(customer.budget.max_limit) : "", budgetResetDuration: customer?.budget?.reset_duration || "1M", + // Rate Limit (stored as string) + tokenMaxLimit: customer?.rate_limit?.token_max_limit ? String(customer.rate_limit.token_max_limit) : "", + tokenResetDuration: customer?.rate_limit?.token_reset_duration || "1h", + requestMaxLimit: customer?.rate_limit?.request_max_limit ? String(customer.rate_limit.request_max_limit) : "", + requestResetDuration: customer?.rate_limit?.request_reset_duration || "1h", }; }; @@ -64,12 +74,21 @@ export default function CustomerDialog({ customer, onSave, onCancel }: CustomerD name: formData.name, budgetMaxLimit: formData.budgetMaxLimit, budgetResetDuration: formData.budgetResetDuration, + tokenMaxLimit: formData.tokenMaxLimit, + tokenResetDuration: formData.tokenResetDuration, + requestMaxLimit: formData.requestMaxLimit, + requestResetDuration: formData.requestResetDuration, }; setFormData((prev) => ({ ...prev, isDirty: !isEqual(initialState, currentData), })); - }, [formData.name, formData.budgetMaxLimit, formData.budgetResetDuration, initialState]); + }, [formData.name, formData.budgetMaxLimit, formData.budgetResetDuration, formData.tokenMaxLimit, formData.tokenResetDuration, formData.requestMaxLimit, formData.requestResetDuration, initialState]); + + // Parse string values to numbers for validation and submission + const budgetMaxLimitNum = formData.budgetMaxLimit ? parseFloat(formData.budgetMaxLimit) : undefined; + const tokenMaxLimitNum = formData.tokenMaxLimit ? parseInt(formData.tokenMaxLimit) : undefined; + const requestMaxLimitNum = formData.requestMaxLimit ? parseInt(formData.requestMaxLimit) : undefined; // Validation const validator = useMemo( @@ -84,12 +103,28 @@ export default function CustomerDialog({ customer, onSave, onCancel }: CustomerD // Budget validation ...(formData.budgetMaxLimit ? [ - Validator.minValue(formData.budgetMaxLimit || 0, 0.01, "Budget max limit must be greater than $0.01"), + 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 + ? [ + 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 + ? [ + Validator.minValue(requestMaxLimitNum || 0, 1, "Request max limit must be at least 1"), + Validator.required(formData.requestResetDuration, "Request reset duration is required"), + ] + : []), ]), - [formData], + [formData, budgetMaxLimitNum, tokenMaxLimitNum, requestMaxLimitNum], ); const updateField = (field: K, value: CustomerFormData[K]) => { @@ -111,12 +146,30 @@ export default function CustomerDialog({ customer, onSave, onCancel }: CustomerD name: formData.name, }; - // Add budget if enabled - if (formData.budgetMaxLimit) { + // Detect budget changes using had/has pattern + const hadBudget = !!customer.budget; + const hasBudget = !!budgetMaxLimitNum; + if (hasBudget) { updateData.budget = { - max_limit: formData.budgetMaxLimit, // Already in dollars + max_limit: budgetMaxLimitNum, reset_duration: formData.budgetResetDuration, }; + } else if (hadBudget) { + updateData.budget = {} as UpdateCustomerRequest["budget"]; + } + + // Detect rate limit changes using had/has pattern + const hadRateLimit = !!customer.rate_limit; + const hasRateLimit = !!tokenMaxLimitNum || !!requestMaxLimitNum; + if (hasRateLimit) { + updateData.rate_limit = { + token_max_limit: tokenMaxLimitNum, + token_reset_duration: tokenMaxLimitNum ? formData.tokenResetDuration : undefined, + request_max_limit: requestMaxLimitNum, + request_reset_duration: requestMaxLimitNum ? formData.requestResetDuration : undefined, + }; + } else if (hadRateLimit) { + updateData.rate_limit = {} as UpdateCustomerRequest["rate_limit"]; } await updateCustomer({ customerId: customer.id, data: updateData }).unwrap(); @@ -128,13 +181,23 @@ export default function CustomerDialog({ customer, onSave, onCancel }: CustomerD }; // Add budget if enabled - if (formData.budgetMaxLimit) { + if (budgetMaxLimitNum) { createData.budget = { - max_limit: formData.budgetMaxLimit, // Already in dollars + max_limit: budgetMaxLimitNum, reset_duration: formData.budgetResetDuration, }; } + // Add rate limit if enabled (token or request limits) + if (tokenMaxLimitNum || requestMaxLimitNum) { + createData.rate_limit = { + token_max_limit: tokenMaxLimitNum, + token_reset_duration: tokenMaxLimitNum ? formData.tokenResetDuration : undefined, + request_max_limit: requestMaxLimitNum, + request_reset_duration: requestMaxLimitNum ? formData.requestResetDuration : undefined, + }; + } + await createCustomer(createData).unwrap(); toast.success("Customer created successfully"); } @@ -178,35 +241,97 @@ export default function CustomerDialog({ customer, onSave, onCancel }: CustomerD updateField("budgetMaxLimit", value === '' ? undefined : parseFloat(value))} + onChangeNumber={(value) => updateField("budgetMaxLimit", value)} onChangeSelect={(value) => updateField("budgetResetDuration", value)} options={resetDurationOptions} /> - {isEditing && customer?.budget && ( -
-
- Current Usage: -
- - {formatCurrency(customer.budget.current_usage)} / {formatCurrency(customer.budget.max_limit)} - - = customer.budget.max_limit ? "destructive" : "default"} - className="text-xs" - > - {Math.round((customer.budget.current_usage / customer.budget.max_limit) * 100)}% - -
-
-
- Last Reset: -
- - {formatDistanceToNow(new Date(customer.budget.last_reset), { addSuffix: true })} - -
+ + {/* Rate Limit Configuration - Token Limits */} + updateField("tokenMaxLimit", value)} + onChangeSelect={(value) => updateField("tokenResetDuration", value)} + options={resetDurationOptions} + /> + + {/* Rate Limit Configuration - Request Limits */} + updateField("requestMaxLimit", value)} + onChangeSelect={(value) => updateField("requestResetDuration", value)} + options={resetDurationOptions} + /> + + {/* Current Usage Section (only shown when editing with existing limits) */} + {isEditing && (customer?.budget || customer?.rate_limit) && ( +
+

Current Usage

+
+ {customer?.budget && ( +
+

Budget

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

+ Last Reset: {formatDistanceToNow(new Date(customer.budget.last_reset), { addSuffix: true })} +

+
+ )} + {customer?.rate_limit?.token_max_limit && ( +
+

Tokens

+
+ + {customer.rate_limit.token_current_usage.toLocaleString()} / {customer.rate_limit.token_max_limit.toLocaleString()} + + = customer.rate_limit.token_max_limit ? "destructive" : "default"} + className="text-xs" + > + {Math.round((customer.rate_limit.token_current_usage / customer.rate_limit.token_max_limit) * 100)}% + +
+

+ Last Reset: {formatDistanceToNow(new Date(customer.rate_limit.token_last_reset), { addSuffix: true })} +

+
+ )} + {customer?.rate_limit?.request_max_limit && ( +
+

Requests

+
+ + {customer.rate_limit.request_current_usage.toLocaleString()} / {customer.rate_limit.request_max_limit.toLocaleString()} + + = customer.rate_limit.request_max_limit ? "destructive" : "default"} + className="text-xs" + > + {Math.round((customer.rate_limit.request_current_usage / customer.rate_limit.request_max_limit) * 100)}% + +
+

+ Last Reset: {formatDistanceToNow(new Date(customer.rate_limit.request_last_reset), { addSuffix: true })} +

+
+ )}
)} diff --git a/ui/app/workspace/user-groups/views/customerTable.tsx b/ui/app/workspace/user-groups/views/customerTable.tsx index bb81011df6..198d6e4bd3 100644 --- a/ui/app/workspace/user-groups/views/customerTable.tsx +++ b/ui/app/workspace/user-groups/views/customerTable.tsx @@ -13,18 +13,25 @@ import { } from "@/components/ui/alertDialog"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { resetDurationLabels } from "@/lib/constants/governance"; import { getErrorMessage, useDeleteCustomerMutation } from "@/lib/store"; import { Customer, Team, VirtualKey } from "@/lib/types/governance"; import { cn } from "@/lib/utils"; -import { formatCurrency, parseResetPeriod } from "@/lib/utils/governance"; +import { formatCurrency } from "@/lib/utils/governance"; import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib"; import { Edit, Plus, Trash2 } from "lucide-react"; import { useState } from "react"; import { toast } from "sonner"; import CustomerDialog from "./customerDialog"; +// Helper to format reset duration for display +const formatResetDuration = (duration: string) => { + return resetDurationLabels[duration] || duration; +}; + interface CustomersTableProps { customers: Customer[]; teams: Team[]; @@ -75,93 +82,219 @@ export default function CustomersTable({ customers, teams, virtualKeys }: Custom return ( <> - {showCustomerDialog && ( - setShowCustomerDialog(false)} /> - )} - -
-
-
-

Manage customer accounts with their own teams, budgets, and access controls.

+ + {showCustomerDialog && ( + setShowCustomerDialog(false)} /> + )} + +
+
+
+

Manage customer accounts with their own teams, budgets, and access controls.

+
+
- -
-
- - - - Name - Teams - Budget - Reset Period - Virtual Keys - - - - - {customers?.length === 0 ? ( +
+
+ - - No customers found. Create your first customer to get started. - + Name + Teams + Budget + Rate Limit + Virtual Keys + - ) : ( - customers?.map((customer) => { - const teams = getTeamsForCustomer(customer.id); - const vks = getVirtualKeysForCustomer(customer.id); - - return ( - - -
{customer.name}
- - - {teams?.length > 0 ? ( -
- + + + {customers?.length === 0 ? ( + + + No customers found. Create your first customer to get started. + + + ) : ( + customers?.map((customer) => { + const customerTeams = getTeamsForCustomer(customer.id); + const vks = getVirtualKeysForCustomer(customer.id); + + // Budget calculations + const isBudgetExhausted = + customer.budget?.max_limit && customer.budget.max_limit > 0 && customer.budget.current_usage >= customer.budget.max_limit; + const budgetPercentage = + customer.budget?.max_limit && customer.budget.max_limit > 0 + ? Math.min((customer.budget.current_usage / customer.budget.max_limit) * 100, 100) + : 0; + + // Rate limit calculations + const isTokenLimitExhausted = + customer.rate_limit?.token_max_limit && + customer.rate_limit.token_max_limit > 0 && + customer.rate_limit.token_current_usage >= customer.rate_limit.token_max_limit; + const isRequestLimitExhausted = + customer.rate_limit?.request_max_limit && + customer.rate_limit.request_max_limit > 0 && + customer.rate_limit.request_current_usage >= customer.rate_limit.request_max_limit; + const isRateLimitExhausted = isTokenLimitExhausted || isRequestLimitExhausted; + const tokenPercentage = + customer.rate_limit?.token_max_limit && customer.rate_limit.token_max_limit > 0 + ? Math.min((customer.rate_limit.token_current_usage / customer.rate_limit.token_max_limit) * 100, 100) + : 0; + const requestPercentage = + customer.rate_limit?.request_max_limit && customer.rate_limit.request_max_limit > 0 + ? Math.min((customer.rate_limit.request_current_usage / customer.rate_limit.request_max_limit) * 100, 100) + : 0; + + const isExhausted = isBudgetExhausted || isRateLimitExhausted; + + return ( + + +
+ {customer.name} + {isExhausted && ( + + Limit Reached + + )} +
+ + + {customerTeams?.length > 0 ? ( +
- {teams.length} {teams.length === 1 ? "team" : "teams"} + {customerTeams.length} {customerTeams.length === 1 ? "team" : "teams"} - {teams.map((team) => team.name).join(", ")} + {customerTeams.map((team) => team.name).join(", ")} - -
- ) : ( - - - )} - - - {customer.budget ? ( - = customer.budget.max_limit && "text-destructive", - )} - > - {formatCurrency(customer.budget.current_usage)} / {formatCurrency(customer.budget.max_limit)} - - ) : ( - - - )} - - - {customer.budget ? ( - parseResetPeriod(customer.budget.reset_duration) - ) : ( - - - )} - - - {vks?.length > 0 ? ( -
- +
+ ) : ( + — + )} + + + {customer.budget ? ( + + +
+
+ {formatCurrency(customer.budget.max_limit)} + + {formatResetDuration(customer.budget.reset_duration)} + +
+ div]:bg-red-500/70" + : budgetPercentage > 80 + ? "[&>div]:bg-amber-500/70" + : "[&>div]:bg-emerald-500/70", + )} + /> +
+ + +

+ {formatCurrency(customer.budget.current_usage)} / {formatCurrency(customer.budget.max_limit)} +

+

+ Resets {formatResetDuration(customer.budget.reset_duration)} +

+ + + ) : ( + — + )} + + + {customer.rate_limit ? ( +
+ {customer.rate_limit.token_max_limit && ( + + +
+
+ {customer.rate_limit.token_max_limit.toLocaleString()} tokens + + {formatResetDuration(customer.rate_limit.token_reset_duration || "1h")} + +
+ div]:bg-red-500/70" + : tokenPercentage > 80 + ? "[&>div]:bg-amber-500/70" + : "[&>div]:bg-emerald-500/70", + )} + /> +
+ + +

+ {customer.rate_limit.token_current_usage.toLocaleString()} /{" "} + {customer.rate_limit.token_max_limit.toLocaleString()} tokens +

+

+ Resets {formatResetDuration(customer.rate_limit.token_reset_duration || "1h")} +

+ + + )} + {customer.rate_limit.request_max_limit && ( + + +
+
+ {customer.rate_limit.request_max_limit.toLocaleString()} req + + {formatResetDuration(customer.rate_limit.request_reset_duration || "1h")} + +
+ div]:bg-red-500/70" + : requestPercentage > 80 + ? "[&>div]:bg-amber-500/70" + : "[&>div]:bg-emerald-500/70", + )} + /> +
+ + +

+ {customer.rate_limit.request_current_usage.toLocaleString()} /{" "} + {customer.rate_limit.request_max_limit.toLocaleString()} requests +

+

+ Resets {formatResetDuration(customer.rate_limit.request_reset_duration || "1h")} +

+ + + )} +
+ ) : ( + — + )} + + + {vks?.length > 0 ? ( +
@@ -170,49 +303,74 @@ export default function CustomersTable({ customers, teams, virtualKeys }: Custom {vks.map((vk) => vk.name).join(", ")} - +
+ ) : ( + — + )} + + +
+ + + + + Edit + + + + + + + + + Delete + + + + Delete Customer + + Are you sure you want to delete "{customer.name}"? This will also delete all associated teams + and unassign any virtual keys. This action cannot be undone. + + + + Cancel + handleDelete(customer.id)} + disabled={isDeleting} + className="bg-red-600 hover:bg-red-700" + > + {isDeleting ? "Deleting..." : "Delete"} + + + +
- ) : ( - - - )} - - -
- - - - - - - - Delete Customer - - Are you sure you want to delete "{customer.name}"? This will also delete all associated teams - and unassign any virtual keys. This action cannot be undone. - - - - Cancel - handleDelete(customer.id)} disabled={isDeleting}> - {isDeleting ? "Deleting..." : "Delete"} - - - - -
- - - ); - }) - )} - -
+ + + ); + }) + )} + + +
-
+ ); } diff --git a/ui/app/workspace/user-groups/views/teamDialog.tsx b/ui/app/workspace/user-groups/views/teamDialog.tsx index e0230f6538..3dcd3f4363 100644 --- a/ui/app/workspace/user-groups/views/teamDialog.tsx +++ b/ui/app/workspace/user-groups/views/teamDialog.tsx @@ -29,9 +29,14 @@ interface TeamDialogProps { interface TeamFormData { name: string; customerId: string; - // Budget - budgetMaxLimit: number | undefined; + // Budget (stored as string to allow intermediate decimal states like "1.") + budgetMaxLimit: string; budgetResetDuration: string; + // Rate Limit + tokenMaxLimit: string; + tokenResetDuration: string; + requestMaxLimit: string; + requestResetDuration: string; isDirty: boolean; } @@ -40,9 +45,14 @@ const createInitialState = (team?: Team | null): Omit = return { name: team?.name || "", customerId: team?.customer_id || "", - // Budget - budgetMaxLimit: team?.budget ? team.budget.max_limit : undefined, // Already in dollars + // Budget (stored as string) + budgetMaxLimit: team?.budget ? String(team.budget.max_limit) : "", budgetResetDuration: team?.budget?.reset_duration || "1M", + // Rate Limit (stored as string) + tokenMaxLimit: team?.rate_limit?.token_max_limit ? String(team.rate_limit.token_max_limit) : "", + tokenResetDuration: team?.rate_limit?.token_reset_duration || "1h", + requestMaxLimit: team?.rate_limit?.request_max_limit ? String(team.rate_limit.request_max_limit) : "", + requestResetDuration: team?.rate_limit?.request_reset_duration || "1h", }; }; @@ -70,12 +80,21 @@ export default function TeamDialog({ team, customers, onSave, onCancel }: TeamDi customerId: formData.customerId, budgetMaxLimit: formData.budgetMaxLimit, budgetResetDuration: formData.budgetResetDuration, + tokenMaxLimit: formData.tokenMaxLimit, + tokenResetDuration: formData.tokenResetDuration, + requestMaxLimit: formData.requestMaxLimit, + requestResetDuration: formData.requestResetDuration, }; setFormData((prev) => ({ ...prev, isDirty: !isEqual(initialState, currentData), })); - }, [formData.name, formData.customerId, formData.budgetMaxLimit, formData.budgetResetDuration, initialState]); + }, [formData.name, formData.customerId, formData.budgetMaxLimit, formData.budgetResetDuration, formData.tokenMaxLimit, formData.tokenResetDuration, formData.requestMaxLimit, formData.requestResetDuration, initialState]); + + // Parse string values to numbers for validation and submission + const budgetMaxLimitNum = formData.budgetMaxLimit ? parseFloat(formData.budgetMaxLimit) : undefined; + const tokenMaxLimitNum = formData.tokenMaxLimit ? parseInt(formData.tokenMaxLimit) : undefined; + const requestMaxLimitNum = formData.requestMaxLimit ? parseInt(formData.requestMaxLimit) : undefined; // Validation const validator = useMemo( @@ -90,12 +109,28 @@ export default function TeamDialog({ team, customers, onSave, onCancel }: TeamDi // Budget validation ...(formData.budgetMaxLimit ? [ - Validator.minValue(formData.budgetMaxLimit || 0, 0.01, "Budget max limit must be greater than $0.01"), + 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 + ? [ + 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 + ? [ + Validator.minValue(requestMaxLimitNum || 0, 1, "Request max limit must be at least 1"), + Validator.required(formData.requestResetDuration, "Request reset duration is required"), + ] + : []), ]), - [formData], + [formData, budgetMaxLimitNum, tokenMaxLimitNum, requestMaxLimitNum], ); const updateField = (field: K, value: TeamFormData[K]) => { @@ -118,12 +153,30 @@ export default function TeamDialog({ team, customers, onSave, onCancel }: TeamDi customer_id: formData.customerId, }; - // Add budget if enabled - if (formData.budgetMaxLimit) { + // Detect budget changes using had/has pattern + const hadBudget = !!team.budget; + const hasBudget = !!budgetMaxLimitNum; + if (hasBudget) { updateData.budget = { - max_limit: formData.budgetMaxLimit, // Already in dollars + max_limit: budgetMaxLimitNum, reset_duration: formData.budgetResetDuration, }; + } else if (hadBudget) { + updateData.budget = {} as UpdateTeamRequest["budget"]; + } + + // Detect rate limit changes using had/has pattern + const hadRateLimit = !!team.rate_limit; + const hasRateLimit = !!tokenMaxLimitNum || !!requestMaxLimitNum; + if (hasRateLimit) { + updateData.rate_limit = { + token_max_limit: tokenMaxLimitNum, + token_reset_duration: tokenMaxLimitNum ? formData.tokenResetDuration : undefined, + request_max_limit: requestMaxLimitNum, + request_reset_duration: requestMaxLimitNum ? formData.requestResetDuration : undefined, + }; + } else if (hadRateLimit) { + updateData.rate_limit = {} as UpdateTeamRequest["rate_limit"]; } await updateTeam({ teamId: team.id, data: updateData }).unwrap(); @@ -136,13 +189,23 @@ export default function TeamDialog({ team, customers, onSave, onCancel }: TeamDi }; // Add budget if enabled - if (formData.budgetMaxLimit) { + if (budgetMaxLimitNum) { createData.budget = { - max_limit: formData.budgetMaxLimit, // Already in dollars + max_limit: budgetMaxLimitNum, reset_duration: formData.budgetResetDuration, }; } + // Add rate limit if enabled (token or request limits) + if (tokenMaxLimitNum || requestMaxLimitNum) { + createData.rate_limit = { + token_max_limit: tokenMaxLimitNum, + token_reset_duration: tokenMaxLimitNum ? formData.tokenResetDuration : undefined, + request_max_limit: requestMaxLimitNum, + request_reset_duration: requestMaxLimitNum ? formData.requestResetDuration : undefined, + }; + } + await createTeam(createData).unwrap(); toast.success("Team created successfully"); } @@ -206,35 +269,97 @@ export default function TeamDialog({ team, customers, onSave, onCancel }: TeamDi updateField("budgetMaxLimit", value === '' ? undefined : parseFloat(value))} + onChangeNumber={(value) => updateField("budgetMaxLimit", value)} onChangeSelect={(value) => updateField("budgetResetDuration", value)} options={resetDurationOptions} /> - {isEditing && team?.budget && ( -
-
- Current Usage: -
- - {formatCurrency(team.budget.current_usage)} /{" "} - {formatCurrency(team.budget.max_limit)} - - = team.budget.max_limit ? "destructive" : "default"} - className="font-mono text-xs" - > - {Math.round((team.budget.current_usage / team.budget.max_limit) * 100)}% - -
-
-
- Last Reset: -
- {formatDistanceToNow(new Date(team.budget.last_reset), { addSuffix: true })} -
+ {/* Rate Limit Configuration - Token Limits */} + updateField("tokenMaxLimit", value)} + onChangeSelect={(value) => updateField("tokenResetDuration", value)} + options={resetDurationOptions} + /> + + {/* Rate Limit Configuration - Request Limits */} + updateField("requestMaxLimit", value)} + onChangeSelect={(value) => updateField("requestResetDuration", value)} + options={resetDurationOptions} + /> + + {/* Current Usage Section (only shown when editing with existing limits) */} + {isEditing && (team?.budget || team?.rate_limit) && ( +
+

Current Usage

+
+ {team?.budget && ( +
+

Budget

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

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

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

Tokens

+
+ + {team.rate_limit.token_current_usage.toLocaleString()} / {team.rate_limit.token_max_limit.toLocaleString()} + + = team.rate_limit.token_max_limit ? "destructive" : "default"} + className="text-xs" + > + {Math.round((team.rate_limit.token_current_usage / team.rate_limit.token_max_limit) * 100)}% + +
+

+ Last Reset: {formatDistanceToNow(new Date(team.rate_limit.token_last_reset), { addSuffix: true })} +

+
+ )} + {team?.rate_limit?.request_max_limit && ( +
+

Requests

+
+ + {team.rate_limit.request_current_usage.toLocaleString()} / {team.rate_limit.request_max_limit.toLocaleString()} + + = team.rate_limit.request_max_limit ? "destructive" : "default"} + className="text-xs" + > + {Math.round((team.rate_limit.request_current_usage / team.rate_limit.request_max_limit) * 100)}% + +
+

+ Last Reset: {formatDistanceToNow(new Date(team.rate_limit.request_last_reset), { addSuffix: true })} +

+
+ )}
)} diff --git a/ui/app/workspace/user-groups/views/teamsTable.tsx b/ui/app/workspace/user-groups/views/teamsTable.tsx index d28f8bb6a4..f6cd20a702 100644 --- a/ui/app/workspace/user-groups/views/teamsTable.tsx +++ b/ui/app/workspace/user-groups/views/teamsTable.tsx @@ -13,18 +13,25 @@ import { } from "@/components/ui/alertDialog"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { resetDurationLabels } from "@/lib/constants/governance"; import { getErrorMessage, useDeleteTeamMutation } from "@/lib/store"; import { Customer, Team, VirtualKey } from "@/lib/types/governance"; import { cn } from "@/lib/utils"; -import { formatCurrency, parseResetPeriod } from "@/lib/utils/governance"; +import { formatCurrency } from "@/lib/utils/governance"; import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib"; import { Edit, Plus, Trash2 } from "lucide-react"; import { useState } from "react"; import { toast } from "sonner"; import TeamDialog from "./teamDialog"; +// Helper to format reset duration for display +const formatResetDuration = (duration: string) => { + return resetDurationLabels[duration] || duration; +}; + interface TeamsTableProps { teams: Team[]; customers: Customer[]; @@ -100,7 +107,7 @@ export default function TeamsTable({ teams, customers, virtualKeys }: TeamsTable Name Customer Budget - Reset Period + Rate Limit Virtual Keys @@ -117,32 +124,163 @@ export default function TeamsTable({ teams, customers, virtualKeys }: 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; + + // Rate limit calculations + const isTokenLimitExhausted = + team.rate_limit?.token_max_limit && + team.rate_limit.token_max_limit > 0 && + team.rate_limit.token_current_usage >= team.rate_limit.token_max_limit; + const isRequestLimitExhausted = + team.rate_limit?.request_max_limit && + team.rate_limit.request_max_limit > 0 && + team.rate_limit.request_current_usage >= team.rate_limit.request_max_limit; + const isRateLimitExhausted = isTokenLimitExhausted || isRequestLimitExhausted; + const tokenPercentage = + team.rate_limit?.token_max_limit && team.rate_limit.token_max_limit > 0 + ? Math.min((team.rate_limit.token_current_usage / team.rate_limit.token_max_limit) * 100, 100) + : 0; + const requestPercentage = + team.rate_limit?.request_max_limit && team.rate_limit.request_max_limit > 0 + ? Math.min((team.rate_limit.request_current_usage / team.rate_limit.request_max_limit) * 100, 100) + : 0; + + const isExhausted = isBudgetExhausted || isRateLimitExhausted; + return ( - - -
{team.name}
+ + +
+ {team.name} + {isExhausted && ( + + Limit Reached + + )} +
{customerName}
- + {team.budget ? ( - = team.budget.max_limit && "text-destructive")} - > - {formatCurrency(team.budget.current_usage)} / {formatCurrency(team.budget.max_limit)} - + + +
+
+ {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)} +

+ + ) : ( - - + — )} - - {team.budget ? ( - parseResetPeriod(team.budget.reset_duration) + + {team.rate_limit ? ( +
+ {team.rate_limit.token_max_limit && ( + + +
+
+ {team.rate_limit.token_max_limit.toLocaleString()} tokens + + {formatResetDuration(team.rate_limit.token_reset_duration || "1h")} + +
+ div]:bg-red-500/70" + : tokenPercentage > 80 + ? "[&>div]:bg-amber-500/70" + : "[&>div]:bg-emerald-500/70", + )} + /> +
+ + +

+ {team.rate_limit.token_current_usage.toLocaleString()} /{" "} + {team.rate_limit.token_max_limit.toLocaleString()} tokens +

+

+ Resets {formatResetDuration(team.rate_limit.token_reset_duration || "1h")} +

+ + + )} + {team.rate_limit.request_max_limit && ( + + +
+
+ {team.rate_limit.request_max_limit.toLocaleString()} req + + {formatResetDuration(team.rate_limit.request_reset_duration || "1h")} + +
+ div]:bg-red-500/70" + : requestPercentage > 80 + ? "[&>div]:bg-amber-500/70" + : "[&>div]:bg-emerald-500/70", + )} + /> +
+ + +

+ {team.rate_limit.request_current_usage.toLocaleString()} /{" "} + {team.rate_limit.request_max_limit.toLocaleString()} requests +

+

+ Resets {formatResetDuration(team.rate_limit.request_reset_duration || "1h")} +

+ + + )} +
) : ( - - + — )} @@ -158,20 +296,41 @@ export default function TeamsTable({ teams, customers, virtualKeys }: TeamsTable
) : ( - - + — )} -
- - - - - + + Edit + + + + + + + + + Delete + Delete Team @@ -182,7 +341,11 @@ export default function TeamsTable({ teams, customers, virtualKeys }: TeamsTable Cancel - handleDelete(team.id)} disabled={isDeleting}> + handleDelete(team.id)} + disabled={isDeleting} + className="bg-red-600 hover:bg-red-700" + > {isDeleting ? "Deleting..." : "Delete"} diff --git a/ui/components/formFooter.tsx b/ui/components/formFooter.tsx index ddb41004c2..0f4cf90c14 100644 --- a/ui/components/formFooter.tsx +++ b/ui/components/formFooter.tsx @@ -24,7 +24,7 @@ export default function FormFooter({ validator, label, onCancel, isLoading, isEd return ( - diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index 409a97a115..b0efad4851 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -29,18 +29,22 @@ export interface Team { name: string; customer_id?: string; budget_id?: string; + rate_limit_id?: string; // Populated relationships customer?: Customer; budget?: Budget; + rate_limit?: RateLimit; } export interface Customer { id: string; name: string; budget_id?: string; + rate_limit_id?: string; // Populated relationships teams?: Team[]; budget?: Budget; + rate_limit?: RateLimit; } export interface DBKey { @@ -172,22 +176,26 @@ export interface CreateTeamRequest { name: string; customer_id?: string; budget?: CreateBudgetRequest; + rate_limit?: CreateRateLimitRequest; } export interface UpdateTeamRequest { name?: string; customer_id?: string; budget?: UpdateBudgetRequest; + rate_limit?: UpdateRateLimitRequest; } export interface CreateCustomerRequest { name: string; budget?: CreateBudgetRequest; + rate_limit?: CreateRateLimitRequest; } export interface UpdateCustomerRequest { name?: string; budget?: UpdateBudgetRequest; + rate_limit?: UpdateRateLimitRequest; } export interface CreateBudgetRequest {