Commit Graph

3 Commits

Author SHA1 Message Date
Yevhenii Shcherbina 1a91d31793 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.
2026-05-29 10:08:25 -04:00
Yevhenii Shcherbina 238968cfa0 feat: add per-group AI budget table and endpoints (#25203)
Closes
https://linear.app/codercom/issue/AIGOV-284/add-group-budgets-table-and-crud-api

## Summary

Adds the `group_ai_budgets` table and the following endpoints:

- `GET /api/v2/groups/{group}/ai/budget`
- `PUT /api/v2/groups/{group}/ai/budget`
- `DELETE /api/v2/groups/{group}/ai/budget`

Each group may have at most one budget row. If no row exists, no budget
is enforced.

### Feature gate
  
Added `RequireFeatureMW(FeatureAIBridge)` on the `/ai/budget` sub-route.

## RBAC

Authorization reuses `rbac.ResourceGroup` with the existing
`.InOrganization(...).WithID(...)` scoping model.

The `dbauthz` wrappers load the parent `groups` row and authorize
against it.

No new resource type is introduced. As a result, anyone with
`group:update` permissions (Owner, OrgAdmin, or UserAdmin within the
organization) can manage AI budgets for that group.

## Read access for group members

`database.Group.RBACObject()` grants `policy.ActionRead` to all members
of the group through the group ACL:

```go
func (g Group) RBACObject() rbac.Object {
	return rbac.ResourceGroup.WithID(g.ID).
		InOrg(g.OrganizationID).
		// Group members can read the group.
		WithGroupACL(map[string][]policy.Action{
			g.ID.String(): {
				policy.ActionRead,
			},
		})
}
```

Because the `GET` endpoint authorizes against the same loaded `Group`
object, any group member can call:

```text
GET /api/v2/groups/{group}/ai/budget
```

`PUT` and `DELETE` remain admin-only. The group ACL grants only
`ActionRead`, so write operations continue to require role-based
`group:update` permissions.

## Alternative considered

A dedicated `rbac.ResourceGroupAiBudget` resource would allow budget
management to be separated from general group administration.

We decided not to add that complexity for now.
2026-05-14 15:54:37 -04:00
Yevhenii Shcherbina 4124d1137d feat: add ai_model_prices table (#24932)
# Summary

Implements
https://linear.app/codercom/issue/AIGOV-282/add-ai-model-price-table-and-seed-generator

This PR lays the groundwork for AI Bridge cost controls (per the AI
Governance RFC). It adds the foundation needed for future cost tracking:
a place to store per-model token prices, a way to keep those prices in
sync with upstream pricing data, and a startup mechanism that ensures
every deployment has prices loaded before AI Bridge starts processing
requests.

The price data comes from [models.dev](https://models.dev/), a
community-maintained catalogue of AI provider pricing. A generator
script fetches the latest prices, filters to Anthropic and OpenAI for
now, and produces a seed file checked into the repository.

On every server startup the seed is applied to the database, so new
releases automatically pick up any price corrections that landed since
the previous one. Existing rows are overwritten with the latest prices;
rows for models no longer in the seed are left untouched.

# Batching the AI model price seed: three approaches

Context: at server startup we seed the `ai_model_prices` table from an
embedded JSON price book (~70 rows today, will grow as we add providers,
potentially 4000+).

Each row is:

```text
(provider, model, input_price, output_price, cache_read_price, cache_write_price)
```

Any of the four price columns can be:

- `NULL` → “price unknown for this dimension”
- explicit `0` → “free”

The batch must be an UPSERT so re-running is idempotent and existing
rows pick up new prices.

We considered three implementations.

---

## Approach 1 — Per-row UPSERT in a Go loop

```go
for _, row := range rows {
    if err := db.UpsertAIModelPrice(ctx, database.UpsertAIModelPriceParams{
        Provider:   row.Provider,
        Model:      row.Model,
        InputPrice: nullInt64(row.InputPrice),
        // ...
    }); err != nil {
        return err
    }
}
```

### Pros

- Trivial.
- NULL handling falls out naturally from `sql.NullInt64`.

### Cons

- `N` round-trips per seed.
- With ~70 rows that means ~70 statement executions on every startup,
even inside a transaction.
- Doesn't scale gracefully as the price book grows, potentially 4000+.

---

## Approach 2 — `UNNEST` with parallel arrays

Pass each column as a separate Go slice. Postgres unnests them in
parallel into a virtual table, then `INSERT ... SELECT`.

```sql
INSERT INTO ai_model_prices (
    provider,
    model,
    input_price,
    output_price,
    cache_read_price,
    cache_write_price
)
SELECT
    UNNEST(@providers::text[]),
    UNNEST(@models::text[]),
    NULLIF(UNNEST(@input_prices::bigint[]), -1),
    NULLIF(UNNEST(@output_prices::bigint[]), -1),
    NULLIF(UNNEST(@cache_read_prices::bigint[]), -1),
    NULLIF(UNNEST(@cache_write_prices::bigint[]), -1)
ON CONFLICT (provider, model) DO UPDATE SET
    input_price       = EXCLUDED.input_price,
    output_price      = EXCLUDED.output_price,
    cache_read_price  = EXCLUDED.cache_read_price,
    cache_write_price = EXCLUDED.cache_write_price,
    updated_at        = NOW();
```

Go side: flatten rows into six parallel slices.

Use a sentinel (`-1`) for “missing”, since `lib/pq` can't encode `NULL`
into a `bigint[]` element.

```go
providers := make([]string, len(rows))
models    := make([]string, len(rows))
inputs    := make([]int64,  len(rows))
outputs   := make([]int64,  len(rows))
cacheR    := make([]int64,  len(rows))
cacheW    := make([]int64,  len(rows))

for i, r := range rows {
    providers[i] = r.Provider
    models[i]    = r.Model

    inputs[i] = -1
    if r.InputPrice != nil {
        inputs[i] = *r.InputPrice
    }

    outputs[i] = -1
    if r.OutputPrice != nil {
        outputs[i] = *r.OutputPrice
    }

    cacheR[i] = -1
    if r.CacheReadPrice != nil {
        cacheR[i] = *r.CacheReadPrice
    }

    cacheW[i] = -1
    if r.CacheWritePrice != nil {
        cacheW[i] = *r.CacheWritePrice
    }
}

return db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
    Providers:        providers,
    Models:           models,
    InputPrices:      inputs,
    OutputPrices:     outputs,
    CacheReadPrices:  cacheR,
    CacheWritePrices: cacheW,
})
```

### Pros

- Single round-trip.

### Cons

- The generated `sqlc` params become plain `[]int64`, which can't
represent `NULL`.

---

## Approach 3 — `jsonb_array_elements` over a single `@seed::jsonb`
(chosen)

Pass the raw seed JSON as one parameter; let Postgres expand and parse
it.

```sql
INSERT INTO ai_model_prices (
    provider,
    model,
    input_price,
    output_price,
    cache_read_price,
    cache_write_price
)
SELECT
    elem->>'provider',
    elem->>'model',
    (elem->>'input_price')::bigint,
    (elem->>'output_price')::bigint,
    (elem->>'cache_read_price')::bigint,
    (elem->>'cache_write_price')::bigint
FROM jsonb_array_elements(@seed::jsonb) AS elem
ON CONFLICT (provider, model) DO UPDATE SET
    input_price       = EXCLUDED.input_price,
    output_price      = EXCLUDED.output_price,
    cache_read_price  = EXCLUDED.cache_read_price,
    cache_write_price = EXCLUDED.cache_write_price,
    updated_at        = NOW();
```

Go side reduces to:

```go
return db.UpsertAIModelPrices(ctx, seedJSON)
```

### Pros

- Single round-trip.
- NULLs fall out naturally:
  - `(elem->>'cache_write_price')::bigint` becomes `NULL`
  - no sentinels
- The seed is already JSON:
- Existing precedent:
  - `jsonb_array_elements` is already used elsewhere in the codebase

### Cons

- Less type-safe at the SQL boundary than `UNNEST`
- Slightly less standard than `UNNEST`
- Readers need familiarity with:
  - `jsonb_array_elements`
  - `->>` extraction syntax
- Postgres pays JSON parse cost
  - negligible at our scale

---

---

# Decision

We picked Approach 3.

It collapses the round-trips like `UNNEST` does, but without:

- nullable-array workarounds
- sentinel values
2026-05-08 16:45:14 -04:00