mirror of
https://github.com/coder/coder.git
synced 2026-11-08 14:42:12 +08:00
Yevhenii Shcherbina
1cada0649c
Closes https://linear.app/codercom/issue/AIGOV-287/add-effective-group-resolution Implements the effective AI budget resolution from the AI Governance cost-controls RFC: for a given user, a `user_ai_budget_overrides` row wins if present, otherwise the deployment budget policy (`highest`) picks the largest group budget across the user's groups, ties broken alphabetically. For now, I keep the logic under `coderd/aibridge/budget`, but that may change during the implementation of budget enforcement.
77 lines
2.5 KiB
Go
77 lines
2.5 KiB
Go
// Package budget resolves the effective AI spend budget for a user. A
|
|
// per-user override always wins; otherwise the deployment budget policy selects
|
|
// a budget from the groups the user belongs to.
|
|
package budget
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
|
|
"github.com/google/uuid"
|
|
"golang.org/x/xerrors"
|
|
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/codersdk"
|
|
)
|
|
|
|
// LimitSource identifies which tier produced an EffectiveBudget.
|
|
type LimitSource string
|
|
|
|
const (
|
|
// SourceUserOverride indicates the budget came from a per-user override.
|
|
SourceUserOverride LimitSource = "user_override"
|
|
// SourceGroup indicates the budget came from a group budget selected by the
|
|
// deployment policy.
|
|
SourceGroup LimitSource = "group"
|
|
)
|
|
|
|
// EffectiveBudget is the AI budget that applies to a user after override and
|
|
// policy resolution.
|
|
type EffectiveBudget struct {
|
|
// GroupID is the group the spend is attributed to.
|
|
GroupID uuid.UUID
|
|
// SpendLimitMicros is the effective spend limit in micro-units
|
|
// (1 unit = 1,000,000).
|
|
SpendLimitMicros int64
|
|
Source LimitSource
|
|
}
|
|
|
|
// ResolveUserAIBudget returns the effective AI budget for userID. The second
|
|
// return value is false when no budget is configured for the user. A per-user
|
|
// override wins unconditionally; otherwise the budget is selected from the
|
|
// user's groups according to policy.
|
|
func ResolveUserAIBudget(ctx context.Context, db database.Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveBudget, bool, error) {
|
|
// A per-user override always wins.
|
|
override, err := db.GetUserAIBudgetOverride(ctx, userID)
|
|
if err == nil {
|
|
return EffectiveBudget{
|
|
GroupID: override.GroupID,
|
|
SpendLimitMicros: override.SpendLimitMicros,
|
|
Source: SourceUserOverride,
|
|
}, true, nil
|
|
}
|
|
if !errors.Is(err, sql.ErrNoRows) {
|
|
return EffectiveBudget{}, false, xerrors.Errorf("get user AI budget override: %w", err)
|
|
}
|
|
|
|
// No override: select a group budget according to the deployment policy.
|
|
switch policy {
|
|
case codersdk.AIBudgetPolicyHighest:
|
|
row, err := db.GetHighestGroupAIBudgetByUser(ctx, userID)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return EffectiveBudget{}, false, nil
|
|
}
|
|
if err != nil {
|
|
return EffectiveBudget{}, false, xerrors.Errorf("get highest group AI budget: %w", err)
|
|
}
|
|
return EffectiveBudget{
|
|
GroupID: row.GroupID,
|
|
SpendLimitMicros: row.SpendLimitMicros,
|
|
Source: SourceGroup,
|
|
}, true, nil
|
|
default:
|
|
return EffectiveBudget{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy)
|
|
}
|
|
}
|