mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add user AI budget override endpoints (#25439)
Implements https://linear.app/codercom/issue/AIGOV-285 Follow the structure established in https://github.com/coder/coder/pull/25203 ## Summary Adds the `user_ai_budget_overrides` table and CRUD API at `/api/v2/users/{user}/ai/budget`. An override sets a custom per-user spend cap that supersedes group-budget resolution, attributing spend to a specific group. ## Schema ```sql CREATE TABLE user_ai_budget_overrides ( user_id UUID PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE, group_id UUID NOT NULL REFERENCES groups(id) ON DELETE CASCADE, spend_limit_micros BIGINT NOT NULL CHECK (spend_limit_micros >= 0), created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); ``` ## Membership lifecycle The membership invariant — a user must be a member of the attributed group, including when that group is "Everyone" — would naturally be expressed as a composite FK on `(user_id, group_id) → group_members_expanded(user_id, group_id)`. PostgreSQL doesn't allow foreign keys to reference views, so enforcement is split across two mechanisms: - **Write-time check.** A CHECK constraint on the table (`user_ai_budget_overrides_must_be_group_member`) calls a `STABLE` function `is_group_member(user_id, group_id)` that queries `group_members_expanded`. The view surfaces both regular group memberships and the implicit "Everyone" group memberships from `organization_members`. Any INSERT or UPDATE that violates the predicate is rejected with a Postgres `check_violation`, which the handler maps to a 400. `is_group_member` is defined as a general predicate, reusable by any future table that needs the same check. - **Cascade on removal.** Two `BEFORE DELETE` triggers handle membership loss: - `trigger_delete_user_ai_budget_overrides_on_group_member_delete` on `group_members` — covers regular group removals (admin action, OIDC sync). - `trigger_delete_user_ai_budget_overrides_on_org_member_delete` on `organization_members` — covers the "Everyone" group, whose membership lives in `organization_members`. The single-column FKs on `users(id)` and `groups(id)` remain to cascade on user or group deletion (those paths don't pass through `group_members`). ## Authorization The dbauthz layer gates each operation against the `User` and (for writes) `Group` resources: | Operation | User resource | Group resource | |-----------|----------------|----------------| | `GET` | `ActionRead` | — | | `PUT` | `ActionUpdate` | `ActionUpdate` | | `DELETE` | `ActionUpdate` | `ActionUpdate` | For `DELETE`, the dbauthz layer fetches the existing override first to learn the attributed `group_id`, then runs both checks. ### Role matrix | Role | GET | PUT | DELETE | |--------------|-----|-----|--------| | Owner | ✅ | ✅ | ✅ | | UserAdmin | ✅ | ✅ | ✅ | | OrgAdmin | ✅ | ❌ | ❌ | | OrgUserAdmin | ✅ | ❌ | ❌ | Internal discussion: https://codercom.slack.com/archives/C096PFVBZKN/p1779392747885359 ## Audit logs Audit logs will be addressed in a follow-up PR.
This commit is contained in:
Generated
+65
@@ -2449,6 +2449,23 @@ func (q *sqlQuerier) DeleteGroupAIBudget(ctx context.Context, groupID uuid.UUID)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const deleteUserAIBudgetOverride = `-- name: DeleteUserAIBudgetOverride :one
|
||||
DELETE FROM user_ai_budget_overrides WHERE user_id = $1 RETURNING user_id, group_id, spend_limit_micros, created_at, updated_at
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error) {
|
||||
row := q.db.QueryRowContext(ctx, deleteUserAIBudgetOverride, userID)
|
||||
var i UserAiBudgetOverride
|
||||
err := row.Scan(
|
||||
&i.UserID,
|
||||
&i.GroupID,
|
||||
&i.SpendLimitMicros,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel :one
|
||||
SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at
|
||||
FROM ai_model_prices
|
||||
@@ -2494,6 +2511,25 @@ func (q *sqlQuerier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (G
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getUserAIBudgetOverride = `-- name: GetUserAIBudgetOverride :one
|
||||
SELECT user_id, group_id, spend_limit_micros, created_at, updated_at
|
||||
FROM user_ai_budget_overrides
|
||||
WHERE user_id = $1
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error) {
|
||||
row := q.db.QueryRowContext(ctx, getUserAIBudgetOverride, userID)
|
||||
var i UserAiBudgetOverride
|
||||
err := row.Scan(
|
||||
&i.UserID,
|
||||
&i.GroupID,
|
||||
&i.SpendLimitMicros,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const upsertAIModelPrices = `-- name: UpsertAIModelPrices :exec
|
||||
INSERT INTO ai_model_prices (
|
||||
provider, model, input_price, output_price, cache_read_price, cache_write_price
|
||||
@@ -2548,6 +2584,35 @@ func (q *sqlQuerier) UpsertGroupAIBudget(ctx context.Context, arg UpsertGroupAIB
|
||||
return i, err
|
||||
}
|
||||
|
||||
const upsertUserAIBudgetOverride = `-- name: UpsertUserAIBudgetOverride :one
|
||||
INSERT INTO user_ai_budget_overrides (user_id, group_id, spend_limit_micros)
|
||||
VALUES ($1, $2, $3)
|
||||
ON CONFLICT (user_id) DO UPDATE SET
|
||||
group_id = EXCLUDED.group_id,
|
||||
spend_limit_micros = EXCLUDED.spend_limit_micros,
|
||||
updated_at = NOW()
|
||||
RETURNING user_id, group_id, spend_limit_micros, created_at, updated_at
|
||||
`
|
||||
|
||||
type UpsertUserAIBudgetOverrideParams struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
GroupID uuid.UUID `db:"group_id" json:"group_id"`
|
||||
SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertUserAIBudgetOverrideParams) (UserAiBudgetOverride, error) {
|
||||
row := q.db.QueryRowContext(ctx, upsertUserAIBudgetOverride, arg.UserID, arg.GroupID, arg.SpendLimitMicros)
|
||||
var i UserAiBudgetOverride
|
||||
err := row.Scan(
|
||||
&i.UserID,
|
||||
&i.GroupID,
|
||||
&i.SpendLimitMicros,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getActiveAISeatCount = `-- name: GetActiveAISeatCount :one
|
||||
SELECT
|
||||
COUNT(*)
|
||||
|
||||
Reference in New Issue
Block a user