team budget db structure refactor

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

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{{

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 {

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 {

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
 	}