diff --git a/docs/guides/ee-license-design.md b/docs/guides/ee-license-design.md index 882538e..e949e94 100644 --- a/docs/guides/ee-license-design.md +++ b/docs/guides/ee-license-design.md @@ -498,8 +498,10 @@ csghub-lite license features # 列出 Catalog 及当前是否可用 > 阶段三的 License 管理部分（Settings 页导入/校验/删除、侧边栏 “CSGLite EE” > 标识）已完成；导航锁标与 `FeatureLocked` 组件随第一条被门控的路由一起做。 > 第一条被门控的功能是账号的企业版部分（`feature.lite.team_management`：审计员 -> 角色和登录安全策略；账号本身免费），导航锁标与锁定页随它一起完成；其余已发布 -> 功能仍不门控。 +> 角色和登录安全策略；账号本身免费），导航锁标与锁定页随它一起完成。第二条是 +> 预算的企业版部分（`feature.lite.budgets`：用户和团队预算、80% 提醒；单个 Key +> 的限制免费，见 [`key-governance.md`](key-governance.md)）。其余已发布功能仍不 +> 门控。 | 阶段 | 内容 | 估时 | diff --git a/docs/guides/key-governance.md b/docs/guides/key-governance.md new file mode 100644 index 0000000..56853cb --- /dev/null +++ b/docs/guides/key-governance.md @@ -0,0 +1,68 @@ +# Key 治理与预算 + +和 LiteLLM 开源版一样，单个 API Key 的限制**免费**：允许的模型、到期时间、 +RPM/TPM 限流，以及 Key 自己的 Token 或金额预算。企业版功能键 +`feature.lite.budgets`（布尔）只解锁两项： + +- **用户和团队预算**：一个用户在所有 Key 和网页端的累计用量，以及团队所有成员的 + 累计用量。没有 License 时，`PUT /api/team/users/{id}/budget` 和 + `PUT /api/team/teams/{id}/budget` 返回 `403 feature_not_licensed`；读取不门控。 +- **80% 提醒**：任何预算用到 80% 时，推理响应带 + `X-CSGLite-Budget-Warning` 头（例如 `key=85%, team=92%`），列表接口的 + `spend.warning` 为 `true`，页面上显示黄色标记。 + +代码位置：Key 限制、用量账本和通用预算检查在 `internal/keygov`（Apache-2.0）， +HTTP 接入在 `internal/server/key_governance.go`；用户和团队预算的范围与提醒头在 +`ee/budget`。 + +## 1. 参考 + +| 方面 | LiteLLM 开源版 | CSGLite | +|---|---|---| +| 模型白名单 | `models` | `limits.models`，`*` 结尾按前缀匹配；`/v1/models` 只列出允许的模型 | +| 到期 | `duration` / `expires` | `limits.expires_at`，页面上按日期选，当天 23:59:59 到期 | +| 限流 | `rpm_limit` / `tpm_limit` | `limits.rpm` / `limits.tpm`，最近一分钟的滑动窗口 | +| 预算 | `max_budget` + `budget_duration`（美元） | `budget.tokens` 和/或 `budget.cost` + `currency`，周期为每天、每月或总计 | +| 用户 / 团队预算 | 免费 | 企业版 | +| 软预算告警 | 企业版（Slack/邮件） | 企业版（响应头和页面标记） | + +## 2. 拒绝方式 + +只对推理路由生效（`/v1/*`、`/anthropic/*`、`/api/chat`、`/api/generate`， +以及 `/providers/{id}/v1/*`）。本机操作者不带 Key 时不受限制。 + +| 情况 | 状态码 | 代码 | +|---|---|---| +| Key 已到期 | 401 | `api_key_expired` | +| 模型不在白名单 | 403 | `model_not_allowed` | +| Key、用户或团队预算用完 | 429 | `budget_exceeded` | +| 超过 RPM/TPM | 429 | `rate_limited`，带 `Retry-After` | + +OpenAI 兼容路由和 Anthropic 路由把代码放在 `error.type`，Ollama 路由放在 +`code`。检查顺序是到期、模型、预算、限流，被拒绝的请求不占用限流额度。 + +## 3. 用量怎么算 + +- 每个计量过的请求结束时，把 Token 数和估算金额记到存储根目录下的 + `spend.db`（SQLite），按「对象、本地日期、币种」累加。对象是 Key、Key 的 + 所有者，以及所有者当时所在的团队。网页端登录用户的对话也记到用户和团队上。 +- 账本从本版本开始记录，之前的用量不计入预算。删除再新建 Key 不能绕过用户和 + 团队预算，因为用户和团队有自己的记录。 +- 金额只统计模型有该币种定价的请求；没有定价的请求只计 Token。 +- 「每天」从本地零点开始，「每月」从本地每月 1 日开始，「总计」不重置。 +- 预算在请求开始前检查，所以最后一个请求可能让用量略超预算，和 LiteLLM 一样。 +- RPM/TPM 计数在内存里，重启后清零。TPM 在请求结束时计入，所以窗口内的 + 第一个大请求不会被拦。 + +## 4. 谁能改 + +- 任何人创建 Key 时都可以给自己的 Key 加限制。 +- 创建后只有管理员能改限制（`PUT /api/api-keys/{id}` 带 `limits`），成员只能 + 改名，避免成员解除管理员给的限制。`limits: {}` 删除全部限制。 +- 用户和团队预算只有管理员能设（企业版）。 + +## 5. License 失效后 + +- Key 限制不受影响。 +- 已设置的用户和团队预算继续生效，保证花费仍有上限；但在续期前不能修改。 +- 80% 提醒停止。 diff --git a/docs/guides/team-management.md b/docs/guides/team-management.md index babcc28..b3aa785 100644 --- a/docs/guides/team-management.md +++ b/docs/guides/team-management.md @@ -23,7 +23,7 @@ LiteLLM Proxy 和 OpenRouter Organizations： | 方面 | LiteLLM | OpenRouter | CSGLite | |---|---|---|---| | 初始化 | `UI_USERNAME`/`UI_PASSWORD` 环境变量登录，仅用于初始化，之后关掉并提示 | 账号注册 | 环境变量**只创建一次真实账号**，之后忽略并在页面提示移除；或使用本机访问、初始化令牌、CLI | -| 免费与收费 | 开源版含团队、内部用户、Virtual Key、基础预算；组织、团队管理员与细粒度角色、SSO（超过 5 人）、审计日志、高级预算收费 | — | 免费：用户、邀请、团队、管理员和成员角色；收费：审计员角色、登录安全策略，以及后续的 Key 治理、花费报表、审计日志、SSO | +| 免费与收费 | 开源版含团队、内部用户、Virtual Key、基础预算；组织、团队管理员与细粒度角色、SSO（超过 5 人）、审计日志、高级预算收费 | — | 免费：用户、邀请、团队、管理员和成员角色、Key 限制（见 [`key-governance.md`](key-governance.md)）；收费：审计员角色、登录安全策略、用户和团队预算，以及后续的花费报表、审计日志、SSO | | 加入方式 | Invite User 生成链接；管理员设的密码首登必须修改 | 邀请加入组织 | 默认邀请链接（7 天、一次性）；直接创建时首登强制改密 | | 角色 | proxy_admin / proxy_admin_viewer / internal_user | Admin / Member | 管理员 / 审计员（企业版） / 成员 | | 应急 | break-glass 账号 | — | `csghub-lite user reset-password` 在主机上直接重置，不受 License 影响 | @@ -136,7 +136,7 @@ LiteLLM Proxy 和 OpenRouter Organizations： ## 9. 后续 -- E2（Key 治理）：按用户和团队设置预算和模型白名单。 +- E2（Key 治理）：已完成，见 [`key-governance.md`](key-governance.md)。 - E3（花费报表）：按用户和团队统计用量。 - E4（审计日志）：记录账号的增删改和登录事件。 - E5（SSO）：复用这里的会话和邀请。 diff --git a/ee/README.md b/ee/README.md index a8657a7..56b6af5 100644 --- a/ee/README.md +++ b/ee/README.md @@ -243,5 +243,19 @@ inference. `feature.lite.team_management` gates giving someone that role and changing the sign-in policy; both checks sit in `internal/server`. Design: `docs/guides/team-management.md`. +### Budgets for users and teams (`ee/budget`) + +A key's own limits (allowed models, expiry, RPM/TPM and a budget) are free, as +in LiteLLM's open-source proxy, and live under Apache-2.0 in `internal/keygov` +together with the spend ledger and the generic budget check. This package holds +only the Enterprise part: `OwnerScopes`, which turns a user's and a team's +budget into scopes the check counts every request against, and +`WarningHeader`, the `X-CSGLite-Budget-Warning` response header at 80%. +`feature.lite.budgets` gates setting those budgets +(`PUT /api/team/{users,teams}/{id}/budget`) and reporting the warning. Budgets +already set keep being enforced after a license lapses, so spending stays +capped; they cannot be changed until it is renewed. Design: +`docs/guides/key-governance.md`. + See `docs/agent-guidelines/ee-features.md` for the full rules and `docs/guides/ee-license-design.md` for the design. diff --git a/ee/budget/budget.go b/ee/budget/budget.go new file mode 100644 index 0000000..b877445 --- /dev/null +++ b/ee/budget/budget.go @@ -0,0 +1,59 @@ +// Copyright (c) OpenCSG. Licensed under the CSGLite Enterprise Edition License. +// See ee/LICENSE for details. + +// Package budget holds the Enterprise part of spend control: budgets for a +// user and for a team, counted across every key and the web UI, and the +// warning at 80% of any budget. Key limits themselves (models, expiry, rate +// limits, a key's own budget) are free and live in internal/keygov. +package budget + +import ( + "fmt" + "math" + "strings" + + "github.com/opencsgs/csglite/internal/keygov" + accounts "github.com/opencsgs/csglite/internal/team" + "github.com/opencsgs/csglite/pkg/api" +) + +// OwnerScopes lists the budgets of user and of team, which may be nil, that +// are set. +func OwnerScopes(user accounts.User, team *accounts.Team) []keygov.Scope { + var scopes []keygov.Scope + if user.ID != "" && user.Budget != nil { + scopes = append(scopes, keygov.Scope{Kind: "user", Subject: keygov.UserSubject(user.ID), Budget: user.Budget}) + } + if team != nil && team.Budget != nil { + scopes = append(scopes, keygov.Scope{Kind: "team", Subject: keygov.TeamSubject(team.ID), Budget: team.Budget}) + } + return scopes +} + +// WarningHeader is the value of the X-CSGLite-Budget-Warning response header +// for the statuses past 80%, such as "key=85%, team=92%", or empty. +func WarningHeader(statuses []keygov.Status) string { + var parts []string + for _, status := range statuses { + if !status.Warning { + continue + } + parts = append(parts, fmt.Sprintf("%s=%d%%", status.Scope.Kind, int(math.Floor(100*Used(status.Scope.Budget, status.Spend))))) + } + return strings.Join(parts, ", ") +} + +// Used is the largest share of a cap that spend has reached. +func Used(budget *api.Budget, spend keygov.Spend) float64 { + if budget == nil { + return 0 + } + used := 0.0 + if budget.Tokens > 0 { + used = math.Max(used, float64(spend.Tokens)/float64(budget.Tokens)) + } + if budget.Cost > 0 { + used = math.Max(used, spend.Cost/budget.Cost) + } + return used +} diff --git a/ee/budget/budget_test.go b/ee/budget/budget_test.go new file mode 100644 index 0000000..f04b862 --- /dev/null +++ b/ee/budget/budget_test.go @@ -0,0 +1,50 @@ +// Copyright (c) OpenCSG. Licensed under the CSGLite Enterprise Edition License. +// See ee/LICENSE for details. + +package budget + +import ( + "testing" + "time" + + "github.com/opencsgs/csglite/internal/keygov" + accounts "github.com/opencsgs/csglite/internal/team" + "github.com/opencsgs/csglite/pkg/api" +) + +func TestCheckUserAndTeamBudgets(t *testing.T) { + ledger := keygov.NewLedger(t.TempDir()) + defer ledger.Close() + now := time.Now() + user := accounts.User{ID: "u1", Budget: &api.Budget{Tokens: 1000, Period: keygov.PeriodDay}} + team := &accounts.Team{ID: "t1", Budget: &api.Budget{Cost: 10, Currency: "CNY", Period: keygov.PeriodMonth}} + scopes := OwnerScopes(user, team) + if len(scopes) != 2 { + t.Fatalf("scopes = %+v", scopes) + } + if err := ledger.Record([]string{keygov.UserSubject("u1"), keygov.TeamSubject("t1")}, now, 850, 10, "CNY"); err != nil { + t.Fatal(err) + } + statuses, err := keygov.Check(ledger, scopes, now) + if err != nil { + t.Fatal(err) + } + if statuses[0].Exhausted || !statuses[0].Warning { + t.Errorf("user at 85%% = %+v; want a warning only", statuses[0]) + } + if !statuses[1].Exhausted { + t.Errorf("team at 100%% = %+v; want exhausted", statuses[1]) + } + if got := WarningHeader(statuses); got != "user=85%, team=100%" { + t.Errorf("header = %q", got) + } +} + +func TestOwnerScopesSkipsUnsetBudgets(t *testing.T) { + if scopes := OwnerScopes(accounts.User{ID: "u1"}, &accounts.Team{ID: "t1"}); len(scopes) != 0 { + t.Fatalf("scopes = %+v; want none", scopes) + } + if scopes := OwnerScopes(accounts.User{Budget: &api.Budget{Tokens: 1}}, nil); len(scopes) != 0 { + t.Fatal("a key that predates accounts has no user to budget") + } +} diff --git a/internal/config/api_auth.go b/internal/config/api_auth.go index a41e725..aa1d04f 100644 --- a/internal/config/api_auth.go +++ b/internal/config/api_auth.go @@ -14,6 +14,8 @@ import ( "strings" "sync" "time" + + "github.com/opencsgs/csglite/pkg/api" ) const ( @@ -31,6 +33,8 @@ type APIKeyRecord struct { // OwnerID is the user account that owns the key when accounts are on. // Keys created before accounts have none. OwnerID string `json:"owner_id,omitempty"` + // Limits restrict the key; nil means none. + Limits *api.KeyLimits `json:"limits,omitempty"` } type APIKeyState struct { @@ -69,7 +73,8 @@ func (s *APIKeyStore) SetAuthEnabled(enabled bool) (APIKeyState, error) { } // Create adds a key owned by ownerID, which is empty when accounts are off. -func (s *APIKeyStore) Create(name, ownerID string) (APIKeyRecord, string, error) { +// Limits must already be normalized. +func (s *APIKeyStore) Create(name, ownerID string, limits *api.KeyLimits) (APIKeyRecord, string, error) { s.mu.Lock() defer s.mu.Unlock() @@ -99,6 +104,7 @@ func (s *APIKeyStore) Create(name, ownerID string) (APIKeyRecord, string, error) Hash: hashAPIKey(plain), CreatedAt: now, OwnerID: strings.TrimSpace(ownerID), + Limits: limits, } state.Keys = append(state.Keys, record) if err := s.saveLocked(state); err != nil { @@ -107,6 +113,36 @@ func (s *APIKeyStore) Create(name, ownerID string) (APIKeyRecord, string, error) return record, plain, nil } +// Update renames a key when name is set and replaces its limits when +// setLimits is true. Limits must already be normalized. +func (s *APIKeyStore) Update(id string, name *string, limits *api.KeyLimits, setLimits bool) (APIKeyRecord, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + + state, err := s.loadLocked() + if err != nil { + return APIKeyRecord{}, false, err + } + for i := range state.Keys { + if state.Keys[i].ID != id { + continue + } + if name != nil { + if trimmed := strings.TrimSpace(*name); trimmed != "" { + state.Keys[i].Name = trimmed + } + } + if setLimits { + state.Keys[i].Limits = limits + } + if err := s.saveLocked(state); err != nil { + return APIKeyRecord{}, false, err + } + return state.Keys[i], true, nil + } + return APIKeyRecord{}, false, nil +} + func (s *APIKeyStore) Delete(id string) (bool, error) { s.mu.Lock() defer s.mu.Unlock() @@ -189,6 +225,27 @@ func (s *APIKeyStore) Validate(plain string) (APIKeyRecord, bool, error) { return APIKeyRecord{}, false, nil } +// Find looks a key up like Validate without recording that it was used. +func (s *APIKeyStore) Find(plain string) (APIKeyRecord, bool, error) { + plain = strings.TrimSpace(plain) + if plain == "" { + return APIKeyRecord{}, false, nil + } + s.mu.RLock() + defer s.mu.RUnlock() + state, err := s.loadLocked() + if err != nil { + return APIKeyRecord{}, false, err + } + hash := hashAPIKey(plain) + for _, record := range state.Keys { + if subtle.ConstantTimeCompare([]byte(record.Hash), []byte(hash)) == 1 { + return record, true, nil + } + } + return APIKeyRecord{}, false, nil +} + func (s *APIKeyStore) loadLocked() (APIKeyState, error) { var state APIKeyState data, err := os.ReadFile(s.path) diff --git a/internal/keygov/keygov_test.go b/internal/keygov/keygov_test.go new file mode 100644 index 0000000..8810582 --- /dev/null +++ b/internal/keygov/keygov_test.go @@ -0,0 +1,132 @@ +package keygov + +import ( + "testing" + "time" + + "github.com/opencsgs/csglite/pkg/api" +) + +func TestNormalizeLimits(t *testing.T) { + got, err := NormalizeLimits(&api.KeyLimits{Models: []string{" qwen/* ", "Qwen/*", "", "gpt-4o"}, RPM: 10}) + if err != nil { + t.Fatal(err) + } + if len(got.Models) != 2 || got.Models[0] != "qwen/*" || got.Models[1] != "gpt-4o" { + t.Fatalf("models = %v", got.Models) + } + if got, err := NormalizeLimits(&api.KeyLimits{Budget: &api.Budget{}}); err != nil || got != nil { + t.Fatalf("empty limits = %v, %v; want nil", got, err) + } + for _, bad := range []api.KeyLimits{ + {RPM: -1}, + {Models: []string{"a*b"}}, + {Budget: &api.Budget{Cost: 5}}, + {Budget: &api.Budget{Tokens: 1, Period: "week"}}, + } { + if _, err := NormalizeLimits(&bad); err == nil { + t.Errorf("NormalizeLimits(%+v) accepted", bad) + } + } + got, err = NormalizeLimits(&api.KeyLimits{Budget: &api.Budget{Cost: 5, Currency: "cny", Period: "Month"}}) + if err != nil || got.Budget.Currency != "CNY" || got.Budget.Period != PeriodMonth { + t.Fatalf("budget = %+v, %v", got, err) + } +} + +func TestAllowsModel(t *testing.T) { + limits := &api.KeyLimits{Models: []string{"Qwen/*", "gpt-4o"}} + for model, want := range map[string]bool{ + "qwen/Qwen3-8B": true, + "GPT-4o": true, + "gpt-4o-mini": false, + "deepseek-r1": false, + } { + if got := AllowsModel(limits, model); got != want { + t.Errorf("AllowsModel(%q) = %v, want %v", model, got, want) + } + } + if !AllowsModel(nil, "anything") || !AllowsModel(&api.KeyLimits{Models: []string{"*"}}, "x") { + t.Fatal("no allowlist, or *, must allow every model") + } +} + +func TestPeriodStartAndBudgets(t *testing.T) { + now := time.Date(2026, 10, 22, 15, 4, 5, 0, time.Local) + if got := PeriodStart(PeriodDay, now); !got.Equal(time.Date(2026, 10, 22, 0, 0, 0, 0, time.Local)) { + t.Errorf("day start = %v", got) + } + if got := PeriodStart(PeriodMonth, now); !got.Equal(time.Date(2026, 10, 1, 0, 0, 0, 0, time.Local)) { + t.Errorf("month start = %v", got) + } + if !PeriodStart("", now).IsZero() { + t.Error("a budget without a period never restarts") + } + budget := &api.Budget{Tokens: 100, Cost: 10, Currency: "CNY"} + if Exhausted(budget, Spend{Tokens: 99, Cost: 9.99}) || !Exhausted(budget, Spend{Tokens: 100}) || !Exhausted(budget, Spend{Cost: 10}) { + t.Error("Exhausted is wrong at the caps") + } + if Warning(budget, Spend{Tokens: 79}) || !Warning(budget, Spend{Tokens: 80}) || !Warning(budget, Spend{Cost: 8}) { + t.Error("Warning is wrong at 80%") + } +} + +func TestLedger(t *testing.T) { + ledger := NewLedger(t.TempDir()) + defer ledger.Close() + today := time.Now() + yesterday := today.AddDate(0, 0, -1) + if err := ledger.Record([]string{KeySubject("k"), UserSubject("u")}, yesterday, 100, 1, "cny"); err != nil { + t.Fatal(err) + } + if err := ledger.Record([]string{KeySubject("k")}, today, 50, 2, "USD"); err != nil { + t.Fatal(err) + } + if err := ledger.Record([]string{KeySubject("k")}, today, 7, 0, ""); err != nil { + t.Fatal(err) + } + all, err := ledger.Spend(KeySubject("k"), time.Time{}, "CNY") + if err != nil || all.Tokens != 157 || all.Cost != 1 { + t.Fatalf("lifetime spend = %+v, %v", all, err) + } + day, err := ledger.Spend(KeySubject("k"), PeriodStart(PeriodDay, today), "USD") + if err != nil || day.Tokens != 57 || day.Cost != 2 { + t.Fatalf("today's spend = %+v, %v", day, err) + } + user, err := ledger.Spend(UserSubject("u"), time.Time{}, "CNY") + if err != nil || user.Tokens != 100 { + t.Fatalf("user spend = %+v, %v", user, err) + } +} + +func TestLimiter(t *testing.T) { + limiter := NewLimiter() + now := time.Now() + for i := 0; i < 2; i++ { + if ok, _ := limiter.Allow("k", 2, 0, now); !ok { + t.Fatalf("request %d refused", i) + } + } + ok, wait := limiter.Allow("k", 2, 0, now.Add(10*time.Second)) + if ok || wait != 50*time.Second { + t.Fatalf("third request = %v, wait %v; want refused for 50s", ok, wait) + } + if ok, _ := limiter.Allow("k", 2, 0, now.Add(61*time.Second)); !ok { + t.Fatal("the window should have moved on") + } + + if ok, _ := limiter.Allow("t", 0, 100, now); !ok { + t.Fatal("first request refused") + } + limiter.AddTokens("t", 120, now) + if ok, _ := limiter.Allow("t", 0, 100, now.Add(time.Second)); ok { + t.Fatal("tokens over the TPM cap should refuse") + } + if ok, _ := limiter.Allow("t", 0, 100, now.Add(61*time.Second)); !ok { + t.Fatal("tokens should leave the window after a minute") + } + limiter.AddTokens("untracked", 10, now) + if len(limiter.windows) != 2 { + t.Fatal("a subject without a rate limit must not be tracked") + } +} diff --git a/internal/keygov/ledger.go b/internal/keygov/ledger.go new file mode 100644 index 0000000..e994746 --- /dev/null +++ b/internal/keygov/ledger.go @@ -0,0 +1,161 @@ +package keygov + +import ( + "database/sql" + "fmt" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/opencsgs/csglite/pkg/api" + _ "modernc.org/sqlite" +) + +// LedgerFile holds what each key, user and team used per day. +const LedgerFile = "spend.db" + +const ledgerSchema = ` +PRAGMA journal_mode=WAL; +PRAGMA busy_timeout=5000; +CREATE TABLE IF NOT EXISTS spend ( + subject TEXT NOT NULL, + day TEXT NOT NULL, + currency TEXT NOT NULL, + requests INTEGER NOT NULL DEFAULT 0, + tokens INTEGER NOT NULL DEFAULT 0, + cost REAL NOT NULL DEFAULT 0, + PRIMARY KEY (subject, day, currency) +); +` + +const ledgerUpsert = ` +INSERT INTO spend (subject, day, currency, requests, tokens, cost) VALUES (?, ?, ?, 1, ?, ?) +ON CONFLICT (subject, day, currency) DO UPDATE SET + requests = spend.requests + 1, + tokens = spend.tokens + excluded.tokens, + cost = spend.cost + excluded.cost +` + +// Subjects name who spending is counted against. +func KeySubject(id string) string { return "key:" + id } +func UserSubject(id string) string { return "user:" + id } +func TeamSubject(id string) string { return "team:" + id } + +// Spend is what a subject used since some time. +type Spend struct { + Tokens int64 + // Cost is in the currency asked for; requests priced in another + // currency, or not priced, add only tokens. + Cost float64 +} + +// Ledger records spending per subject and local day. Days are local so that +// daily and monthly budgets restart at the operator's midnight. +type Ledger struct { + mu sync.Mutex + path string + db *sql.DB +} + +func NewLedger(appHome string) *Ledger { + return &Ledger{path: filepath.Join(appHome, LedgerFile)} +} + +func (l *Ledger) Close() error { + if l == nil { + return nil + } + l.mu.Lock() + defer l.mu.Unlock() + if l.db == nil { + return nil + } + err := l.db.Close() + l.db = nil + return err +} + +// Record adds one request to every subject. +func (l *Ledger) Record(subjects []string, at time.Time, tokens int64, cost float64, currency string) error { + if l == nil || len(subjects) == 0 { + return nil + } + if cost <= 0 { + cost, currency = 0, "" + } + day := at.Local().Format(time.DateOnly) + l.mu.Lock() + defer l.mu.Unlock() + db, err := l.openLocked() + if err != nil { + return err + } + tx, err := db.Begin() + if err != nil { + return err + } + for _, subject := range subjects { + if _, err := tx.Exec(ledgerUpsert, subject, day, strings.ToUpper(currency), tokens, cost); err != nil { + _ = tx.Rollback() + return fmt.Errorf("recording spend: %w", err) + } + } + return tx.Commit() +} + +// Spend sums what subject used from the local day of since onwards, with +// cost counted in currency. A zero since sums everything. +func (l *Ledger) Spend(subject string, since time.Time, currency string) (Spend, error) { + if l == nil { + return Spend{}, nil + } + from := "" + if !since.IsZero() { + from = since.Local().Format(time.DateOnly) + } + l.mu.Lock() + defer l.mu.Unlock() + db, err := l.openLocked() + if err != nil { + return Spend{}, err + } + var out Spend + err = db.QueryRow( + `SELECT COALESCE(SUM(tokens), 0), COALESCE(SUM(CASE WHEN currency = ? THEN cost ELSE 0 END), 0) + FROM spend WHERE subject = ? AND day >= ?`, + strings.ToUpper(currency), subject, from, + ).Scan(&out.Tokens, &out.Cost) + if err != nil { + return Spend{}, fmt.Errorf("reading spend: %w", err) + } + return out, nil +} + +func (l *Ledger) openLocked() (*sql.DB, error) { + if l.db != nil { + return l.db, nil + } + db, err := sql.Open("sqlite", l.path) + if err != nil { + return nil, err + } + db.SetMaxOpenConns(1) + if _, err := db.Exec(ledgerSchema); err != nil { + _ = db.Close() + return nil, fmt.Errorf("opening spend ledger: %w", err) + } + l.db = db + return db, nil +} + +// Current reads what subject spent against budget in the period containing +// now, and when that period began. +func (l *Ledger) Current(subject string, budget *api.Budget, now time.Time) (Spend, time.Time, error) { + if budget == nil { + return Spend{}, time.Time{}, nil + } + since := PeriodStart(budget.Period, now) + spend, err := l.Spend(subject, since, budget.Currency) + return spend, since, err +} diff --git a/internal/keygov/limiter.go b/internal/keygov/limiter.go new file mode 100644 index 0000000..915e779 --- /dev/null +++ b/internal/keygov/limiter.go @@ -0,0 +1,92 @@ +package keygov + +import ( + "sync" + "time" +) + +const rateWindow = time.Minute + +// Limiter counts requests and tokens per subject over the last minute. It +// lives in memory, so a restart forgets the window, which errs on the side +// of letting requests through. +type Limiter struct { + mu sync.Mutex + windows map[string]*window +} + +type window struct { + requests []time.Time + tokens []tokenMark +} + +type tokenMark struct { + at time.Time + tokens int64 +} + +func NewLimiter() *Limiter { + return &Limiter{windows: map[string]*window{}} +} + +// Allow admits one request for subject unless that would pass rpm, or the +// tokens already used in the window reached tpm. A zero cap is not checked. +// When refused it says how long until the window has room. +func (l *Limiter) Allow(subject string, rpm, tpm int, now time.Time) (bool, time.Duration) { + if l == nil || (rpm <= 0 && tpm <= 0) { + return true, 0 + } + l.mu.Lock() + defer l.mu.Unlock() + w := l.windows[subject] + if w == nil { + w = &window{} + l.windows[subject] = w + } + w.trim(now) + if rpm > 0 && len(w.requests) >= rpm { + return false, w.requests[len(w.requests)-rpm].Add(rateWindow).Sub(now) + } + if tpm > 0 { + var used int64 + for _, mark := range w.tokens { + used += mark.tokens + } + if used >= int64(tpm) { + // Room comes back as the oldest tokens leave the window. + return false, w.tokens[0].at.Add(rateWindow).Sub(now) + } + } + w.requests = append(w.requests, now) + return true, 0 +} + +// AddTokens counts tokens a finished request used. Subjects without a rate +// limit are not tracked. +func (l *Limiter) AddTokens(subject string, tokens int64, now time.Time) { + if l == nil || tokens <= 0 { + return + } + l.mu.Lock() + defer l.mu.Unlock() + w := l.windows[subject] + if w == nil { + return + } + w.trim(now) + w.tokens = append(w.tokens, tokenMark{at: now, tokens: tokens}) +} + +func (w *window) trim(now time.Time) { + cutoff := now.Add(-rateWindow) + i := 0 + for i < len(w.requests) && !w.requests[i].After(cutoff) { + i++ + } + w.requests = w.requests[i:] + j := 0 + for j < len(w.tokens) && !w.tokens[j].at.After(cutoff) { + j++ + } + w.tokens = w.tokens[j:] +} diff --git a/internal/keygov/limits.go b/internal/keygov/limits.go new file mode 100644 index 0000000..b30de9e --- /dev/null +++ b/internal/keygov/limits.go @@ -0,0 +1,198 @@ +// Package keygov enforces what an API key may do: which models it calls, +// until when, how fast, and how much it spends. Key limits are free, as in +// LiteLLM's open-source proxy; budgets for users and teams are in ee/budget. +package keygov + +import ( + "errors" + "fmt" + "math" + "regexp" + "strings" + "time" + + "github.com/opencsgs/csglite/pkg/api" +) + +// Budget periods. An empty period never restarts. +const ( + PeriodDay = "day" + PeriodMonth = "month" +) + +// WarnFraction is the share of a cap at which a budget warns. +const WarnFraction = 0.8 + +const maxModels = 200 + +var currencyPattern = regexp.MustCompile(`^[A-Z]{3}$`) + +// NormalizeBudget trims a budget and checks it. A budget that caps nothing +// comes back nil. +func NormalizeBudget(b *api.Budget) (*api.Budget, error) { + if b == nil { + return nil, nil + } + out := *b + out.Currency = strings.ToUpper(strings.TrimSpace(out.Currency)) + out.Period = strings.ToLower(strings.TrimSpace(out.Period)) + if out.Tokens < 0 { + return nil, errors.New("token budget must not be negative") + } + if out.Cost < 0 || math.IsNaN(out.Cost) || math.IsInf(out.Cost, 0) { + return nil, errors.New("cost budget must be a positive number") + } + switch out.Period { + case "", PeriodDay, PeriodMonth: + default: + return nil, fmt.Errorf("budget period must be %q, %q or empty", PeriodDay, PeriodMonth) + } + if out.Cost > 0 && !currencyPattern.MatchString(out.Currency) { + return nil, errors.New("a cost budget needs a three-letter currency such as CNY or USD") + } + if out.Cost == 0 { + out.Currency = "" + } + if out.Tokens == 0 && out.Cost == 0 { + return nil, nil + } + return &out, nil +} + +// NormalizeLimits trims key limits and checks them. Limits that restrict +// nothing come back nil. +func NormalizeLimits(l *api.KeyLimits) (*api.KeyLimits, error) { + if l == nil { + return nil, nil + } + out := api.KeyLimits{RPM: l.RPM, TPM: l.TPM, ExpiresAt: l.ExpiresAt} + if out.RPM < 0 || out.TPM < 0 { + return nil, errors.New("rate limits must not be negative") + } + seen := map[string]bool{} + for _, model := range l.Models { + model = strings.TrimSpace(model) + if model == "" || seen[strings.ToLower(model)] { + continue + } + if model != "*" && strings.Contains(strings.TrimSuffix(model, "*"), "*") { + return nil, fmt.Errorf("model %q may only end in *", model) + } + seen[strings.ToLower(model)] = true + out.Models = append(out.Models, model) + } + if len(out.Models) > maxModels { + return nil, fmt.Errorf("a key may list at most %d models", maxModels) + } + if out.ExpiresAt != nil { + if out.ExpiresAt.IsZero() { + out.ExpiresAt = nil + } else { + at := out.ExpiresAt.UTC() + out.ExpiresAt = &at + } + } + budget, err := NormalizeBudget(l.Budget) + if err != nil { + return nil, err + } + out.Budget = budget + if len(out.Models) == 0 && out.ExpiresAt == nil && out.RPM == 0 && out.TPM == 0 && out.Budget == nil { + return nil, nil + } + return &out, nil +} + +// AllowsModel reports whether limits let a key call model. Names compare +// without case; an entry ending in "*" matches by prefix. +func AllowsModel(l *api.KeyLimits, model string) bool { + if l == nil || len(l.Models) == 0 { + return true + } + model = strings.ToLower(strings.TrimSpace(model)) + for _, entry := range l.Models { + entry = strings.ToLower(entry) + if prefix, ok := strings.CutSuffix(entry, "*"); ok { + if strings.HasPrefix(model, prefix) { + return true + } + continue + } + if entry == model { + return true + } + } + return false +} + +// Expired reports whether limits have run out at now. +func Expired(l *api.KeyLimits, now time.Time) bool { + return l != nil && l.ExpiresAt != nil && !now.Before(*l.ExpiresAt) +} + +// PeriodStart is when the budget period containing now began, in now's +// location. It is the zero time for a budget that never restarts. +func PeriodStart(period string, now time.Time) time.Time { + switch period { + case PeriodDay: + return time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()) + case PeriodMonth: + return time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, now.Location()) + default: + return time.Time{} + } +} + +// Exhausted reports whether spend has reached a cap of b. +func Exhausted(b *api.Budget, spend Spend) bool { + if b == nil { + return false + } + return (b.Tokens > 0 && spend.Tokens >= b.Tokens) || (b.Cost > 0 && spend.Cost >= b.Cost) +} + +// Warning reports whether spend has reached WarnFraction of a cap of b. +func Warning(b *api.Budget, spend Spend) bool { + if b == nil { + return false + } + return (b.Tokens > 0 && float64(spend.Tokens) >= WarnFraction*float64(b.Tokens)) || + (b.Cost > 0 && spend.Cost >= WarnFraction*b.Cost) +} + +// Scope is one budget a request counts against: a key's, or with an +// Enterprise license a user's or a team's. +type Scope struct { + // Kind is "key", "user" or "team". + Kind string + Subject string + Budget *api.Budget +} + +// Status is where one scope stands in its current period. +type Status struct { + Scope Scope + Spend Spend + Since time.Time + Exhausted bool + Warning bool +} + +// Check reads each scope's spend in its current period. +func Check(ledger *Ledger, scopes []Scope, now time.Time) ([]Status, error) { + out := make([]Status, 0, len(scopes)) + for _, scope := range scopes { + spend, since, err := ledger.Current(scope.Subject, scope.Budget, now) + if err != nil { + return nil, err + } + out = append(out, Status{ + Scope: scope, + Spend: spend, + Since: since, + Exhausted: Exhausted(scope.Budget, spend), + Warning: Warning(scope.Budget, spend), + }) + } + return out, nil +} diff --git a/internal/license/features.go b/internal/license/features.go index 188d0f4..d06554f 100644 --- a/internal/license/features.go +++ b/internal/license/features.go @@ -51,7 +51,7 @@ type FeatureDefinition struct { Since string } -// Catalog entries. FeatureTeamManagement and the quota marked Gated consult +// Catalog entries. FeatureTeamManagement, FeatureBudgets and the quota marked Gated consult // the license; every other boolean is enabled regardless of license. The ungated entries are declared so that the // issuer registry and this binary agree on the names before a feature moves // to EE. @@ -110,6 +110,14 @@ var ( Key: featurePrefix + "team_management", Type: FeatureTypeBoolean, Gated: true, DefaultValue: true, Since: "0.9.92", } + // FeatureBudgets is the Enterprise part of spend control: budgets for a + // user and for a team, and the warning at 80% of any budget. A key's own + // limits (models, expiry, RPM/TPM, budget) are free, as in LiteLLM's + // open-source proxy. See docs/guides/key-governance.md. + FeatureBudgets = FeatureDefinition{ + Key: featurePrefix + "budgets", Type: FeatureTypeBoolean, Gated: true, DefaultValue: true, + Since: "0.9.92", + } ) var catalog = []FeatureDefinition{ @@ -123,6 +131,7 @@ var catalog = []FeatureDefinition{ FeatureLANCluster, QuotaMaxClusterNodes, FeatureTeamManagement, + FeatureBudgets, } // Catalog returns a copy of every registered feature definition. diff --git a/internal/license/manager_test.go b/internal/license/manager_test.go index 9f0b45a..887005a 100644 --- a/internal/license/manager_test.go +++ b/internal/license/manager_test.go @@ -103,7 +103,7 @@ func TestUngatedFeaturesIgnoreTheLicense(t *testing.T) { key := licensetest.NewKey(t) // Shipped catalog: the cluster node cap and team management are gated. m, _ := newManager(t, key, func(o *license.Options) { o.Catalog = license.Catalog() }) - wantGated := []string{license.QuotaMaxClusterNodes.Key, license.FeatureTeamManagement.Key} + wantGated := []string{license.QuotaMaxClusterNodes.Key, license.FeatureTeamManagement.Key, license.FeatureBudgets.Key} gated := license.GatedCatalog() if len(gated) != len(wantGated) { t.Fatalf("shipped catalog gates %v; only %v are expected. Update this test and the docs deliberately", gated, wantGated) @@ -432,7 +432,7 @@ func TestAPIViewsMirrorState(t *testing.T) { } } full := st.APIState("") - if full.GraceUntil == nil || len(full.Features) != 7 { + if full.GraceUntil == nil || len(full.Features) != 8 { t.Fatalf("state view %+v", full) } verify.Limits["quota.lite.max_provider_pools"] = 99 diff --git a/internal/server/auth_test.go b/internal/server/auth_test.go index a6743f4..4c26d77 100644 --- a/internal/server/auth_test.go +++ b/internal/server/auth_test.go @@ -77,7 +77,7 @@ func TestRemoteAPIAuthProtectsAnthropicAlias(t *testing.T) { func TestRemoteAPIAuthAcceptsBearerAndXAPIKey(t *testing.T) { s := newTestServer(t) - _, plain, err := s.apiKeys.Create("client", "") + _, plain, err := s.apiKeys.Create("client", "", nil) if err != nil { t.Fatalf("create key: %v", err) } @@ -231,7 +231,7 @@ func TestDesktopExternalDownloadAuthFollowsRemoteAccessSetting(t *testing.T) { } } - _, plain, err := s.apiKeys.Create("downloads", "") + _, plain, err := s.apiKeys.Create("downloads", "", nil) if err != nil { t.Fatalf("create key: %v", err) } @@ -296,7 +296,7 @@ func TestDesktopExternalRuntimeAuthFollowsLocalAPIKeySetting(t *testing.T) { t.Fatalf("%s status with auth disabled = %d body=%q, want request to reach handler", path, status, body) } } - _, plain, err := s.apiKeys.Create("desktop-runtime", "") + _, plain, err := s.apiKeys.Create("desktop-runtime", "", nil) if err != nil { t.Fatalf("create key: %v", err) } @@ -400,7 +400,7 @@ func TestAPIKeyCreateDoesNotExposeHash(t *testing.T) { func TestAPIUsageAggregatesByKeyAndModel(t *testing.T) { s := newTestServer(t) - record, _, err := s.apiKeys.Create("client", "") + record, _, err := s.apiKeys.Create("client", "", nil) if err != nil { t.Fatalf("create key: %v", err) } @@ -474,7 +474,7 @@ func TestAPIUsageRecordsBuiltinClientWithoutAPIKey(t *testing.T) { func TestAPIUsageBackfillsLegacyMissingSourceFromModelHistory(t *testing.T) { s := newTestServer(t) - record, _, err := s.apiKeys.Create("client", "") + record, _, err := s.apiKeys.Create("client", "", nil) if err != nil { t.Fatalf("create key: %v", err) } @@ -565,7 +565,7 @@ func TestAPIUsageResolvesLegacyLocalModelID(t *testing.T) { func TestAPIUsageSourceTotalsKeepThirdPartyProvidersSeparate(t *testing.T) { s := newTestServer(t) - record, _, err := s.apiKeys.Create("client", "") + record, _, err := s.apiKeys.Create("client", "", nil) if err != nil { t.Fatalf("create key: %v", err) } @@ -619,7 +619,7 @@ func TestAPIUsageSourceTotalsKeepThirdPartyProvidersSeparate(t *testing.T) { func TestAPIUsageFiltersByProvider(t *testing.T) { s := newTestServer(t) - record, _, err := s.apiKeys.Create("client", "") + record, _, err := s.apiKeys.Create("client", "", nil) if err != nil { t.Fatalf("create key: %v", err) } @@ -725,7 +725,7 @@ func TestAPIUsageReportsPoolTotalsMemberSplitAndPoolFilter(t *testing.T) { func TestAPIUsageFiltersByPeriod(t *testing.T) { s := newTestServer(t) - record, _, err := s.apiKeys.Create("client", "") + record, _, err := s.apiKeys.Create("client", "", nil) if err != nil { t.Fatalf("create key: %v", err) } @@ -916,7 +916,7 @@ func TestAPIUsageKeyTotalsAggregatePerKeyAndFollowFilters(t *testing.T) { func TestAPIKeyIsIdentifiedForUsageWhenAuthIsNotEnforced(t *testing.T) { s := newTestServer(t) - record, plain, err := s.apiKeys.Create("reporting", "") + record, plain, err := s.apiKeys.Create("reporting", "", nil) if err != nil { t.Fatalf("create key: %v", err) } diff --git a/internal/server/handlers_api_keys.go b/internal/server/handlers_api_keys.go index 3f7c754..2e16852 100644 --- a/internal/server/handlers_api_keys.go +++ b/internal/server/handlers_api_keys.go @@ -10,6 +10,7 @@ import ( "time" "github.com/opencsgs/csglite/internal/config" + "github.com/opencsgs/csglite/internal/keygov" "github.com/opencsgs/csglite/pkg/api" ) @@ -62,17 +63,76 @@ func (s *Server) handleAPIKeyCreate(w http.ResponseWriter, r *http.Request) { httpjson.WriteCodedError(w, http.StatusForbidden, "sign in to create API keys", permissionDeniedCode) return } - record, plain, err := s.apiKeys.Create(strings.TrimSpace(req.Name), id.User.ID) + // Anyone may restrict a key they create; only an administrator may + // change limits later, so a member cannot lift what was set for them. + limits, err := keygov.NormalizeLimits(req.Limits) + if err != nil { + writeError(w, http.StatusBadRequest, err.Error()) + return + } + record, plain, err := s.apiKeys.Create(strings.TrimSpace(req.Name), id.User.ID, limits) if err != nil { writeError(w, http.StatusInternalServerError, "failed to create API key") return } writeJSON(w, http.StatusCreated, api.APIKeyCreateResponse{ - Key: apiKeyInfo(record, s.apiKeyOwnerNames()), + Key: s.apiKeyInfo(record, s.apiKeyOwnerNames()), APIKey: plain, }) } +// PUT /api/api-keys/{id}. The owner may rename a key; changing its limits +// takes an administrator. +func (s *Server) handleAPIKeyUpdate(w http.ResponseWriter, r *http.Request) { + if s.apiKeys == nil { + writeError(w, http.StatusInternalServerError, "API key store is unavailable") + return + } + var req api.APIKeyUpdateRequest + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 64<<10)).Decode(&req); err != nil { + writeError(w, http.StatusBadRequest, "invalid request body") + return + } + id := strings.TrimSpace(r.PathValue("id")) + if !requestIsAdmin(r) { + if req.Limits != nil { + httpjson.WriteCodedError(w, http.StatusForbidden, "only an administrator may change key limits", permissionDeniedCode) + return + } + state, err := s.apiKeys.State() + if err != nil { + writeError(w, http.StatusInternalServerError, "failed to load API keys") + return + } + owned := false + for _, record := range state.Keys { + if record.ID == id && deletableAPIKey(r, record) { + owned = true + break + } + } + if !owned { + writeError(w, http.StatusNotFound, "API key not found") + return + } + } + limits, err := keygov.NormalizeLimits(req.Limits) + if err != nil { + writeError(w, http.StatusBadRequest, err.Error()) + return + } + record, ok, err := s.apiKeys.Update(id, req.Name, limits, req.Limits != nil) + if err != nil { + writeError(w, http.StatusInternalServerError, "failed to update API key") + return + } + if !ok { + writeError(w, http.StatusNotFound, "API key not found") + return + } + writeJSON(w, http.StatusOK, s.apiKeyInfo(record, s.apiKeyOwnerNames())) +} + func (s *Server) handleAPIKeyDelete(w http.ResponseWriter, r *http.Request) { if s.apiKeys == nil { writeError(w, http.StatusInternalServerError, "API key store is unavailable") @@ -514,24 +574,8 @@ func (s *Server) apiKeysResponseFor(r *http.Request, state config.APIKeyState) a owners := s.apiKeyOwnerNames() for _, record := range state.Keys { if visibleAPIKey(r, record) { - resp.Keys = append(resp.Keys, apiKeyInfo(record, owners)) + resp.Keys = append(resp.Keys, s.apiKeyInfo(record, owners)) } } return resp } - -func apiKeyInfo(record config.APIKeyRecord, owners map[string]string) api.APIKeyInfo { - var lastUsedAt *time.Time - if !record.LastUsedAt.IsZero() { - lastUsedAt = &record.LastUsedAt - } - return api.APIKeyInfo{ - ID: record.ID, - Name: record.Name, - Prefix: record.Prefix, - CreatedAt: record.CreatedAt, - LastUsedAt: lastUsedAt, - OwnerID: record.OwnerID, - OwnerName: owners[record.OwnerID], - } -} diff --git a/internal/server/handlers_models.go b/internal/server/handlers_models.go index 5f8bf01..07847cb 100644 --- a/internal/server/handlers_models.go +++ b/internal/server/handlers_models.go @@ -48,6 +48,7 @@ func (s *Server) handleAnthropicModels(w http.ResponseWriter, r *http.Request) { return } models = filterModelsByProviderRoute(models, providerRouteSourceFromContext(r.Context())) + models = s.filterModelsForKey(r, models) data := make([]api.AnthropicModelInfo, 0, len(models)) for _, item := range models { diff --git a/internal/server/handlers_openai.go b/internal/server/handlers_openai.go index 395dd3f..a95b235 100644 --- a/internal/server/handlers_openai.go +++ b/internal/server/handlers_openai.go @@ -564,6 +564,7 @@ func (s *Server) handleOpenAIModels(w http.ResponseWriter, r *http.Request) { return } models = filterModelsByProviderRoute(models, providerRouteSourceFromContext(r.Context())) + models = s.filterModelsForKey(r, models) seen := make(map[string]struct{}, len(models)) data := make([]api.OpenAIModel, 0, len(models)+4) diff --git a/internal/server/key_governance.go b/internal/server/key_governance.go new file mode 100644 index 0000000..bcd8148 --- /dev/null +++ b/internal/server/key_governance.go @@ -0,0 +1,269 @@ +package server + +import ( + "bytes" + "encoding/json" + "io" + "math" + "mime" + "mime/multipart" + "net/http" + "strconv" + "strings" + "time" + + eebudget "github.com/opencsgs/csglite/ee/budget" + "github.com/opencsgs/csglite/internal/config" + "github.com/opencsgs/csglite/internal/httpjson" + "github.com/opencsgs/csglite/internal/keygov" + "github.com/opencsgs/csglite/internal/license" + "github.com/opencsgs/csglite/internal/team" + "github.com/opencsgs/csglite/pkg/api" +) + +// Codes inference callers and the web UI switch on. +const ( + apiKeyExpiredCode = "api_key_expired" + modelNotAllowedCode = "model_not_allowed" + rateLimitedCode = "rate_limited" + budgetExceededCode = "budget_exceeded" +) + +// budgetWarningHeader reports budgets past 80% on a response, with an +// Enterprise license. +const budgetWarningHeader = "X-CSGLite-Budget-Warning" + +// Key statuses in API responses. +const ( + apiKeyStatusActive = "active" + apiKeyStatusExpired = "expired" + apiKeyStatusExhausted = "exhausted" +) + +// keyGovernanceMiddleware applies a key's limits, and the budgets of the user +// and team behind a request, to inference calls. Requests without a key or a +// signed-in user, such as the local operator's, are not limited. +func (s *Server) keyGovernanceMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !identifiesAPIKeyForUsage(r) { + next.ServeHTTP(w, r) + return + } + now := time.Now() + key, hasKey := authenticatedAPIKey(r) + var limits *api.KeyLimits + if hasKey { + limits = key.Limits + } + if keygov.Expired(limits, now) { + writeGovernanceError(w, r, http.StatusUnauthorized, apiKeyExpiredCode, "this API key has expired") + return + } + if limits != nil && len(limits.Models) > 0 { + model := peekRequestModel(r) + if model != "" && !keygov.AllowsModel(limits, model) { + writeGovernanceError(w, r, http.StatusForbidden, modelNotAllowedCode, "this API key may not call model "+model) + return + } + } + + statuses, err := keygov.Check(s.spend, s.requestBudgetScopes(r), now) + if err != nil { + writeGovernanceError(w, r, http.StatusInternalServerError, "server_error", "failed to read budgets") + return + } + for _, status := range statuses { + if status.Exhausted { + writeGovernanceError(w, r, http.StatusTooManyRequests, budgetExceededCode, "the "+status.Scope.Kind+" budget is used up") + return + } + } + if limits != nil { + if ok, wait := s.rateLimits.Allow(keygov.KeySubject(key.ID), limits.RPM, limits.TPM, now); !ok { + w.Header().Set("Retry-After", strconv.Itoa(int(math.Ceil(wait.Seconds())))) + writeGovernanceError(w, r, http.StatusTooManyRequests, rateLimitedCode, "this API key is over its rate limit") + return + } + } + if s.license.Enabled(license.FeatureBudgets) { + if warning := eebudget.WarningHeader(statuses); warning != "" { + w.Header().Set(budgetWarningHeader, warning) + } + } + next.ServeHTTP(w, r) + }) +} + +// requestBudgetScopes lists the budgets a request counts against: its key's, +// and its user's and its user's team's, which only a license sets. +func (s *Server) requestBudgetScopes(r *http.Request) []keygov.Scope { + var scopes []keygov.Scope + if key, ok := authenticatedAPIKey(r); ok && key.Limits != nil && key.Limits.Budget != nil { + scopes = append(scopes, keygov.Scope{Kind: "key", Subject: keygov.KeySubject(key.ID), Budget: key.Limits.Budget}) + } + if id, ok := identityFromRequest(r); ok && !id.Local && id.User.ID != "" { + scopes = append(scopes, eebudget.OwnerScopes(id.User, s.teamByID(id.User.TeamID))...) + } + return scopes +} + +func (s *Server) teamByID(id string) *team.Team { + if id == "" || s.team == nil { + return nil + } + teams, err := s.team.Teams() + if err != nil { + return nil + } + for i := range teams { + if teams[i].ID == id { + return &teams[i] + } + } + return nil +} + +// recordSpend counts a finished request against its key, user and team, and +// against the key's rate limit window. +func (s *Server) recordSpend(r *http.Request, tokens int64, cost float64, currency string) { + now := time.Now() + var subjects []string + if key, ok := authenticatedAPIKey(r); ok { + subjects = append(subjects, keygov.KeySubject(key.ID)) + if s.rateLimits != nil { + s.rateLimits.AddTokens(keygov.KeySubject(key.ID), tokens, now) + } + } + if id, ok := identityFromRequest(r); ok && !id.Local && id.User.ID != "" { + subjects = append(subjects, keygov.UserSubject(id.User.ID)) + if id.User.TeamID != "" { + subjects = append(subjects, keygov.TeamSubject(id.User.TeamID)) + } + } + _ = s.spend.Record(subjects, now, tokens, cost, currency) +} + +// budgetSpend reports where subject stands against budget, or nil without a +// budget. Warnings are an Enterprise feature. +func (s *Server) budgetSpend(subject string, budget *api.Budget) *api.BudgetSpend { + if budget == nil { + return nil + } + spend, since, err := s.spend.Current(subject, budget, time.Now()) + if err != nil { + return nil + } + out := &api.BudgetSpend{ + Tokens: spend.Tokens, + Cost: spend.Cost, + Exhausted: keygov.Exhausted(budget, spend), + Warning: keygov.Warning(budget, spend) && s.license.Enabled(license.FeatureBudgets), + } + if !since.IsZero() { + out.Since = &since + } + return out +} + +func (s *Server) apiKeyInfo(record config.APIKeyRecord, owners map[string]string) api.APIKeyInfo { + var lastUsedAt *time.Time + if !record.LastUsedAt.IsZero() { + lastUsedAt = &record.LastUsedAt + } + info := api.APIKeyInfo{ + ID: record.ID, + Name: record.Name, + Prefix: record.Prefix, + CreatedAt: record.CreatedAt, + LastUsedAt: lastUsedAt, + OwnerID: record.OwnerID, + OwnerName: owners[record.OwnerID], + Limits: record.Limits, + Status: apiKeyStatusActive, + } + if record.Limits != nil { + info.Spend = s.budgetSpend(keygov.KeySubject(record.ID), record.Limits.Budget) + } + switch { + case keygov.Expired(record.Limits, time.Now()): + info.Status = apiKeyStatusExpired + case info.Spend != nil && info.Spend.Exhausted: + info.Status = apiKeyStatusExhausted + } + return info +} + +// filterModelsForKey drops the models a key's allowlist leaves out of a +// model list. /v1/models is not metered, so the key is looked up here +// without marking it used. +func (s *Server) filterModelsForKey(r *http.Request, models []api.ModelInfo) []api.ModelInfo { + key, ok := authenticatedAPIKey(r) + if !ok && s.apiKeys != nil { + key, ok, _ = s.apiKeys.Find(requestAPIKey(r)) + } + if !ok || key.Limits == nil || len(key.Limits.Models) == 0 { + return models + } + out := models[:0:0] + for _, model := range models { + if keygov.AllowsModel(key.Limits, model.Model) { + out = append(out, model) + } + } + return out +} + +// peekRequestModel reads the model a request names without consuming its +// body: a JSON body is read and put back, and a multipart form is read only +// up to its model field. +func peekRequestModel(r *http.Request) string { + if r.Body == nil { + return "" + } + mediaType, params, _ := mime.ParseMediaType(r.Header.Get("Content-Type")) + if mediaType == "multipart/form-data" && params["boundary"] != "" { + var seen bytes.Buffer + reader := multipart.NewReader(io.TeeReader(r.Body, &seen), params["boundary"]) + model := "" + for { + part, err := reader.NextPart() + if err != nil { + break + } + if part.FormName() == "model" { + value, _ := io.ReadAll(io.LimitReader(part, 1024)) + model = strings.TrimSpace(string(value)) + break + } + } + r.Body = struct { + io.Reader + io.Closer + }{io.MultiReader(&seen, r.Body), r.Body} + return model + } + body, err := io.ReadAll(r.Body) + _ = r.Body.Close() + r.Body = io.NopCloser(bytes.NewReader(body)) + if err != nil { + return "" + } + var probe struct { + Model string `json:"model"` + } + _ = json.Unmarshal(body, &probe) + return strings.TrimSpace(probe.Model) +} + +// writeGovernanceError answers in the shape the caller's protocol expects. +func writeGovernanceError(w http.ResponseWriter, r *http.Request, status int, code, msg string) { + path := providerRouteLegacyPath(r.URL.Path) + switch { + case strings.HasPrefix(path, "/anthropic/") || strings.HasPrefix(path, "/v1/messages"): + writeAnthropicErrorWithType(w, status, code, msg) + case strings.HasPrefix(path, "/v1/"): + writeOpenAIError(w, status, code, msg) + default: + httpjson.WriteCodedError(w, status, msg, code) + } +} diff --git a/internal/server/key_governance_test.go b/internal/server/key_governance_test.go new file mode 100644 index 0000000..e1aba58 --- /dev/null +++ b/internal/server/key_governance_test.go @@ -0,0 +1,245 @@ +package server + +import ( + "bytes" + "context" + "encoding/json" + "io" + "mime/multipart" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/opencsgs/csglite/internal/config" + "github.com/opencsgs/csglite/internal/keygov" + "github.com/opencsgs/csglite/internal/team" + "github.com/opencsgs/csglite/pkg/api" +) + +// expectOpenAIError checks the error type an OpenAI-compatible route returns. +func expectOpenAIError(t *testing.T, rec *httptest.ResponseRecorder, status int, errType, what string) { + t.Helper() + expectStatus(t, rec, status, what) + var body struct { + Error struct { + Type string `json:"type"` + } `json:"error"` + } + _ = json.Unmarshal(rec.Body.Bytes(), &body) + if body.Error.Type != errType { + t.Fatalf("%s: error type %q, want %q: %s", what, body.Error.Type, errType, rec.Body.String()) + } +} + +func chatRequest(model string) api.OpenAIChatRequest { + return api.OpenAIChatRequest{Model: model, Messages: []api.Message{{Role: "user", Content: "hi"}}} +} + +// newKeyClient enables API key authentication, has the local operator +// create a key with limits, and returns a remote client using it. +func newKeyClient(t *testing.T, s *Server, limits *api.KeyLimits) (*teamClient, api.APIKeyInfo) { + t.Helper() + local := newTeamClient(t, s) + local.remote = "127.0.0.1:50000" + expectStatus(t, local.do(http.MethodPost, "/api/api-keys/settings", api.APIKeySettingsUpdateRequest{AuthEnabled: true}), http.StatusOK, "enable API keys") + rec := local.do(http.MethodPost, "/api/api-keys", api.APIKeyCreateRequest{Name: "client", Limits: limits}) + expectStatus(t, rec, http.StatusCreated, "create a key with limits") + var created api.APIKeyCreateResponse + if err := json.Unmarshal(rec.Body.Bytes(), &created); err != nil { + t.Fatal(err) + } + client := newTeamClient(t, s) + client.apiKey = created.APIKey + return client, created.Key +} + +func TestKeyLimitsWorkWithoutALicense(t *testing.T) { + s := newTestServer(t) + client, key := newKeyClient(t, s, &api.KeyLimits{Models: []string{"qwen/*"}, RPM: 1}) + if key.Limits == nil || key.Limits.RPM != 1 || key.Status != apiKeyStatusActive { + t.Fatalf("created key = %+v", key) + } + + expectOpenAIError(t, client.do(http.MethodPost, "/v1/chat/completions", chatRequest("gpt-4o")), + http.StatusForbidden, modelNotAllowedCode, "a model outside the allowlist") + rec := client.do(http.MethodPost, "/v1/chat/completions", chatRequest("qwen/Qwen3-8B")) + if rec.Code == http.StatusForbidden || rec.Code == http.StatusTooManyRequests { + t.Fatalf("an allowed model was refused: %d %s", rec.Code, rec.Body.String()) + } + rec = client.do(http.MethodPost, "/v1/chat/completions", chatRequest("qwen/Qwen3-8B")) + expectOpenAIError(t, rec, http.StatusTooManyRequests, rateLimitedCode, "a second request within the minute") + if rec.Header().Get("Retry-After") == "" { + t.Fatal("a rate-limited response must say when to retry") + } +} + +func TestAnExpiredKeyIsRefused(t *testing.T) { + s := newTestServer(t) + past := time.Now().Add(-time.Hour) + client, key := newKeyClient(t, s, &api.KeyLimits{ExpiresAt: &past}) + if key.Status != apiKeyStatusExpired { + t.Fatalf("status = %q, want expired", key.Status) + } + expectCode(t, client.do(http.MethodPost, "/api/chat", chatRequest("any")), http.StatusUnauthorized, apiKeyExpiredCode, "an expired key") +} + +func TestAKeyStopsAtItsBudget(t *testing.T) { + s := newTestServer(t) + client, key := newKeyClient(t, s, &api.KeyLimits{Budget: &api.Budget{Tokens: 100, Period: keygov.PeriodDay}}) + if err := s.spend.Record([]string{keygov.KeySubject(key.ID)}, time.Now(), 100, 0, ""); err != nil { + t.Fatal(err) + } + rec := client.do(http.MethodPost, "/anthropic/v1/messages", api.AnthropicMessageRequest{Model: "any", MaxTokens: 10}) + expectStatus(t, rec, http.StatusTooManyRequests, "a key past its budget") + if !bytes.Contains(rec.Body.Bytes(), []byte(budgetExceededCode)) { + t.Fatalf("body = %s, want the %s type", rec.Body.String(), budgetExceededCode) + } + + local := newTeamClient(t, s) + local.remote = "127.0.0.1:50000" + var list api.APIKeysResponse + _ = json.Unmarshal(local.do(http.MethodGet, "/api/api-keys", nil).Body.Bytes(), &list) + got := list.Keys[0] + if got.Status != apiKeyStatusExhausted || got.Spend == nil || got.Spend.Tokens != 100 || got.Spend.Since == nil { + t.Fatalf("listed key = %+v", got) + } + if got.Spend.Warning { + t.Fatal("budget warnings are Enterprise") + } +} + +func TestOnlyAnAdministratorChangesKeyLimits(t *testing.T) { + s, admin := newAccountsServer(t) + _, member := inviteUser(t, s, admin, "bob", "member") + rec := member.do(http.MethodPost, "/api/api-keys", api.APIKeyCreateRequest{Name: "mine", Limits: &api.KeyLimits{RPM: 5}}) + expectStatus(t, rec, http.StatusCreated, "a member restricts their own new key") + var created api.APIKeyCreateResponse + _ = json.Unmarshal(rec.Body.Bytes(), &created) + path := "/api/api-keys/" + created.Key.ID + + expectCode(t, member.do(http.MethodPut, path, api.APIKeyUpdateRequest{Limits: &api.KeyLimits{}}), + http.StatusForbidden, permissionDeniedCode, "a member lifting their own limits") + expectStatus(t, member.do(http.MethodPut, path, api.APIKeyUpdateRequest{Name: ptr("renamed")}), http.StatusOK, "a member renaming their key") + + rec = admin.do(http.MethodPut, path, api.APIKeyUpdateRequest{Limits: &api.KeyLimits{TPM: 1000}}) + expectStatus(t, rec, http.StatusOK, "an administrator changing limits") + var updated api.APIKeyInfo + _ = json.Unmarshal(rec.Body.Bytes(), &updated) + if updated.Name != "renamed" || updated.Limits == nil || updated.Limits.TPM != 1000 || updated.Limits.RPM != 0 { + t.Fatalf("updated key = %+v", updated) + } + expectStatus(t, admin.do(http.MethodPut, path, api.APIKeyUpdateRequest{Limits: &api.KeyLimits{RPM: -1}}), http.StatusBadRequest, "a negative limit") +} + +func TestUserAndTeamBudgetsNeedALicense(t *testing.T) { + s := newTestServer(t) + local := newTeamClient(t, s) + local.remote = "127.0.0.1:50000" + admin := createTeamUser(t, local, api.TeamUserCreateRequest{Username: "admin", Password: "admin-password"}) + expectCode(t, local.do(http.MethodPut, "/api/team/users/"+admin.ID+"/budget", api.BudgetRequest{Budget: &api.Budget{Tokens: 10}}), + http.StatusForbidden, licenseErrorCode, "a user budget without a license") +} + +func TestUserAndTeamBudgetsWarnAndStop(t *testing.T) { + s, admin := newAccountsServer(t) + rec := admin.do(http.MethodPost, "/api/team/teams", api.TeamGroupRequest{Name: ptr("research")}) + var group api.TeamGroup + _ = json.Unmarshal(rec.Body.Bytes(), &group) + bob, member := inviteUser(t, s, admin, "bob", "member") + expectStatus(t, admin.do(http.MethodPut, "/api/team/users/"+bob.ID, api.TeamUserUpdateRequest{TeamID: &group.ID}), http.StatusOK, "join the team") + + rec = admin.do(http.MethodPut, "/api/team/users/"+bob.ID+"/budget", api.BudgetRequest{Budget: &api.Budget{Tokens: 1000, Period: "month"}}) + expectStatus(t, rec, http.StatusOK, "set a user budget") + expectStatus(t, admin.do(http.MethodPut, "/api/team/teams/"+group.ID+"/budget", api.BudgetRequest{Budget: &api.Budget{Cost: 5, Currency: "cny"}}), + http.StatusOK, "set a team budget") + expectStatus(t, admin.do(http.MethodPut, "/api/team/teams/"+group.ID+"/budget", api.BudgetRequest{Budget: &api.Budget{Cost: 5}}), + http.StatusBadRequest, "a cost budget without a currency") + + if err := s.spend.Record([]string{keygov.UserSubject(bob.ID)}, time.Now(), 850, 0, ""); err != nil { + t.Fatal(err) + } + rec = member.do(http.MethodPost, "/v1/chat/completions", chatRequest("any")) + if rec.Code == http.StatusTooManyRequests { + t.Fatalf("refused at 85%%: %s", rec.Body.String()) + } + if got := rec.Header().Get(budgetWarningHeader); got != "user=85%" { + t.Fatalf("warning header = %q", got) + } + var users api.TeamUsersResponse + _ = json.Unmarshal(admin.do(http.MethodGet, "/api/team/users", nil).Body.Bytes(), &users) + for _, user := range users.Users { + if user.ID == bob.ID && (user.Spend == nil || !user.Spend.Warning || user.Spend.Exhausted) { + t.Fatalf("listed user = %+v", user) + } + } + + if err := s.spend.Record([]string{keygov.TeamSubject(group.ID)}, time.Now(), 0, 5, "CNY"); err != nil { + t.Fatal(err) + } + expectOpenAIError(t, member.do(http.MethodPost, "/v1/chat/completions", chatRequest("any")), + http.StatusTooManyRequests, budgetExceededCode, "a member of a team past its budget") + + expectStatus(t, admin.do(http.MethodPut, "/api/team/teams/"+group.ID+"/budget", api.BudgetRequest{}), http.StatusOK, "remove the team budget") + if rec := member.do(http.MethodPost, "/v1/chat/completions", chatRequest("any")); rec.Code == http.StatusTooManyRequests { + t.Fatalf("still refused after the budget was removed: %s", rec.Body.String()) + } +} + +func TestSpendIsCountedAgainstKeyUserAndTeam(t *testing.T) { + s := newTestServer(t) + key := config.APIKeyRecord{ID: "k1", Limits: &api.KeyLimits{TPM: 10}} + ctx := context.WithValue(context.Background(), apiKeyContextKey{}, key) + ctx = context.WithValue(ctx, requestIdentityKey{}, requestIdentity{User: team.User{ID: "u1", TeamID: "t1"}, APIKey: &key}) + r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(ctx) + + s.rateLimits.Allow(keygov.KeySubject("k1"), 0, 10, time.Now()) + s.recordSpend(r, 42, 0.5, "USD") + for _, subject := range []string{keygov.KeySubject("k1"), keygov.UserSubject("u1"), keygov.TeamSubject("t1")} { + spend, err := s.spend.Spend(subject, time.Time{}, "USD") + if err != nil || spend.Tokens != 42 || spend.Cost != 0.5 { + t.Fatalf("%s spend = %+v, %v", subject, spend, err) + } + } + if ok, _ := s.rateLimits.Allow(keygov.KeySubject("k1"), 0, 10, time.Now()); ok { + t.Fatal("the tokens should count against the key's TPM") + } +} + +func TestModelListFollowsTheKeyAllowlist(t *testing.T) { + s := newTestServer(t) + key := config.APIKeyRecord{ID: "k1", Limits: &api.KeyLimits{Models: []string{"qwen/*"}}} + r := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + r = r.WithContext(context.WithValue(r.Context(), apiKeyContextKey{}, key)) + got := s.filterModelsForKey(r, []api.ModelInfo{{Model: "Qwen/Qwen3-8B"}, {Model: "gpt-4o"}}) + if len(got) != 1 || got[0].Model != "Qwen/Qwen3-8B" { + t.Fatalf("models = %+v", got) + } +} + +func TestPeekRequestModelLeavesTheBodyIntact(t *testing.T) { + var form bytes.Buffer + writer := multipart.NewWriter(&form) + file, _ := writer.CreateFormFile("file", "a.wav") + _, _ = file.Write(bytes.Repeat([]byte("x"), 10_000)) + _ = writer.WriteField("model", "whisper-large") + _ = writer.Close() + original := form.Bytes() + r := httptest.NewRequest(http.MethodPost, "/v1/audio/transcriptions", bytes.NewReader(append([]byte(nil), original...))) + r.Header.Set("Content-Type", writer.FormDataContentType()) + if got := peekRequestModel(r); got != "whisper-large" { + t.Fatalf("multipart model = %q", got) + } + if body, _ := io.ReadAll(r.Body); !bytes.Equal(body, original) { + t.Fatal("the multipart body changed") + } + + r = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader([]byte(`{"model":" gpt-4o ","messages":[]}`))) + r.Header.Set("Content-Type", "application/json") + if got := peekRequestModel(r); got != "gpt-4o" { + t.Fatalf("JSON model = %q", got) + } + if body, _ := io.ReadAll(r.Body); string(body) != `{"model":" gpt-4o ","messages":[]}` { + t.Fatalf("the JSON body changed: %s", body) + } +} diff --git a/internal/server/license_test.go b/internal/server/license_test.go index d54ad97..778a74e 100644 --- a/internal/server/license_test.go +++ b/internal/server/license_test.go @@ -108,9 +108,9 @@ func TestRequireFeatureUngatedAlwaysPasses(t *testing.T) { t.Fatalf("%s should be gated: %+v", entry.Key, entry) } continue - case license.FeatureTeamManagement.Key: + case license.FeatureTeamManagement.Key, license.FeatureBudgets.Key: if !entry.Gated || entry.Enabled { - t.Fatalf("team management should be gated and off without a license: %+v", entry) + t.Fatalf("%s should be gated and off without a license: %+v", entry.Key, entry) } continue } @@ -205,7 +205,7 @@ func TestLicenseLifecycleOverHTTP(t *testing.T) { t.Fatal(err) } if settings.Edition != license.EditionEnterprise || settings.LicenseStatus != "valid" || settings.License == nil || - len(settings.Features) != 8 || len(settings.FeatureCatalog) != len(license.Catalog()) { + len(settings.Features) != 9 || len(settings.FeatureCatalog) != len(license.Catalog()) { t.Fatalf("settings after import: edition=%s status=%s features=%v catalog=%d", settings.Edition, settings.LicenseStatus, settings.Features, len(settings.FeatureCatalog)) } diff --git a/internal/server/routes.go b/internal/server/routes.go index 9085e1a..404ad17 100644 --- a/internal/server/routes.go +++ b/internal/server/routes.go @@ -132,10 +132,14 @@ func (s *Server) routes() http.Handler { mux.HandleFunc("POST /api/team/teams", s.handleTeamGroupCreate) mux.HandleFunc("PUT /api/team/teams/{id}", s.handleTeamGroupUpdate) mux.HandleFunc("DELETE /api/team/teams/{id}", s.handleTeamGroupDelete) + // Budgets for users and teams are Enterprise; reading them is not. + mux.HandleFunc("PUT /api/team/users/{id}/budget", s.requireFeature(license.FeatureBudgets)(s.handleTeamUserBudget)) + mux.HandleFunc("PUT /api/team/teams/{id}/budget", s.requireFeature(license.FeatureBudgets)(s.handleTeamGroupBudget)) mux.HandleFunc("GET /api/api-keys", s.handleAPIKeysList) mux.HandleFunc("POST /api/api-keys/settings", s.handleAPIKeysSettingsUpdate) mux.HandleFunc("POST /api/api-keys", s.handleAPIKeyCreate) + mux.HandleFunc("PUT /api/api-keys/{id}", s.handleAPIKeyUpdate) mux.HandleFunc("DELETE /api/api-keys/{id}", s.handleAPIKeyDelete) mux.HandleFunc("GET /api/api-usage", s.handleAPIUsage) mux.HandleFunc("GET /api/observability/requests", s.handleObservabilityRequests) @@ -228,7 +232,7 @@ func (s *Server) routes() http.Handler { return correlationMiddleware(LogMiddleware( s.desktopAuthMiddleware(s.corsMiddleware(sensitiveOriginGuard(s.accountsMiddleware(s.apiAuthMiddleware( - s.observabilityMiddleware(providerPoolUsageMiddleware(mux)), + s.observabilityMiddleware(s.keyGovernanceMiddleware(providerPoolUsageMiddleware(mux))), ))))), )) } @@ -279,7 +283,7 @@ func (s *Server) externalAPIRoutes() http.Handler { return correlationMiddleware(LogMiddleware( desktopExternalAPIMiddleware(s.apiAuthMiddleware( - s.observabilityMiddleware(providerPoolUsageMiddleware(mux)), + s.observabilityMiddleware(s.keyGovernanceMiddleware(providerPoolUsageMiddleware(mux))), )), )) } diff --git a/internal/server/server.go b/internal/server/server.go index 88ab3e5..efc9416 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -30,6 +30,7 @@ import ( "github.com/opencsgs/csglite/internal/dataset" "github.com/opencsgs/csglite/internal/imagegen" "github.com/opencsgs/csglite/internal/inference" + "github.com/opencsgs/csglite/internal/keygov" "github.com/opencsgs/csglite/internal/license" "github.com/opencsgs/csglite/internal/model" "github.com/opencsgs/csglite/internal/modelmetadata" @@ -257,6 +258,8 @@ type Server struct { loginLimiter loginLimiter setupTokenMu sync.Mutex apiUsage *config.APIUsageStore + spend *keygov.Ledger + rateLimits *keygov.Limiter observabilityMu sync.RWMutex observability *observability.Store observabilityCleanupAt atomic.Int64 @@ -360,7 +363,9 @@ func New(cfg *config.Config, version string) *Server { s.apiKeys = config.NewAPIKeyStore(appHome) s.team = team.NewStore(appHome) s.apiUsage = config.NewAPIUsageStore(appHome) + s.spend = keygov.NewLedger(appHome) } + s.rateLimits = keygov.NewLimiter() s.bootstrapAccounts() handler := s.routes() @@ -597,6 +602,7 @@ func (s *Server) shutdownRuntime() { if s.apiUsage != nil { _ = s.apiUsage.Close() } + _ = s.spend.Close() s.modelMetadataMu.Lock() if s.modelMetadata != nil { _ = s.modelMetadata.Close() diff --git a/internal/server/static/openapi/local-api.json b/internal/server/static/openapi/local-api.json index 8a5f174..9cdc742 100644 --- a/internal/server/static/openapi/local-api.json +++ b/internal/server/static/openapi/local-api.json @@ -3780,6 +3780,122 @@ } } }, + "/api/team/users/{id}/budget": { + "put": { + "tags": [ + "Accounts" + ], + "summary": "Set a user's budget", + "description": "Sets or, with a null `budget`, removes the user's budget. It caps what the user spends across every key and the web UI, per day, per month or in total; requests past it are refused with 429 (`budget_exceeded`). Enterprise (`feature.lite.budgets`); reading budgets is not gated. Administrators only.", + "operationId": "setTeamUserBudget", + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/BudgetRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Updated user.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TeamUser" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest" + }, + "403": { + "description": "Not licensed (`code: feature_not_licensed`), the caller's role does not allow this (`code: permission_denied`), or the caller must sign in (`code: login_required`, 401).", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccountsErrorResponse" + } + } + } + }, + "404": { + "$ref": "#/components/responses/NotFound" + } + } + } + }, + "/api/team/teams/{id}/budget": { + "put": { + "tags": [ + "Accounts" + ], + "summary": "Set a team's budget", + "description": "Sets or, with a null `budget`, removes the team's budget. It caps what the team spends across every key and the web UI, per day, per month or in total; requests past it are refused with 429 (`budget_exceeded`). Enterprise (`feature.lite.budgets`); reading budgets is not gated. Administrators only.", + "operationId": "setTeamBudget", + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/BudgetRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Updated team.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TeamGroup" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest" + }, + "403": { + "description": "Not licensed (`code: feature_not_licensed`), the caller's role does not allow this (`code: permission_denied`), or the caller must sign in (`code: login_required`, 401).", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccountsErrorResponse" + } + } + } + }, + "404": { + "$ref": "#/components/responses/NotFound" + } + } + } + }, "/api/api-keys": { "get": { "tags": [ @@ -3877,6 +3993,62 @@ } }, "/api/api-keys/{id}": { + "put": { + "tags": [ + "System" + ], + "summary": "Update a local API key", + "description": "Renames a key or replaces its limits. The owner may rename a key; changing `limits` takes an administrator, and `{}` removes them. Limits are free in every edition.", + "operationId": "updateAPIKey", + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/APIKeyUpdateRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Updated key.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/APIKeyInfo" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest" + }, + "403": { + "description": "Only an administrator may change limits (`code: permission_denied`).", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccountsErrorResponse" + } + } + } + }, + "404": { + "$ref": "#/components/responses/NotFound" + } + } + }, "delete": { "tags": [ "System" @@ -10011,6 +10183,12 @@ "invite_token": { "type": "string", "description": "Returned once when an invitation is issued; the link is /invite/." + }, + "budget": { + "$ref": "#/components/schemas/Budget" + }, + "spend": { + "$ref": "#/components/schemas/BudgetSpend" } } }, @@ -10204,6 +10382,12 @@ "member_count": { "type": "integer" }, + "budget": { + "$ref": "#/components/schemas/Budget" + }, + "spend": { + "$ref": "#/components/schemas/BudgetSpend" + }, "created_at": { "type": "string", "format": "date-time" @@ -10248,7 +10432,8 @@ "id", "name", "prefix", - "created_at" + "created_at", + "status" ], "properties": { "owner_id": { @@ -10274,6 +10459,21 @@ "last_used_at": { "type": "string", "format": "date-time" + }, + "limits": { + "$ref": "#/components/schemas/KeyLimits" + }, + "spend": { + "$ref": "#/components/schemas/BudgetSpend" + }, + "status": { + "type": "string", + "enum": [ + "active", + "expired", + "exhausted" + ], + "description": "`exhausted` means the key's budget is used up for the current period." } } }, @@ -10302,6 +10502,130 @@ "properties": { "name": { "type": "string" + }, + "limits": { + "$ref": "#/components/schemas/KeyLimits" + } + } + }, + "APIKeyUpdateRequest": { + "type": "object", + "additionalProperties": false, + "properties": { + "name": { + "type": "string" + }, + "limits": { + "$ref": "#/components/schemas/KeyLimits" + } + } + }, + "KeyLimits": { + "type": "object", + "additionalProperties": false, + "description": "Restrictions on one API key, enforced on inference routes. A key past them is refused: an expired key with 401 (`api_key_expired`), a model outside `models` with 403 (`model_not_allowed`), over `rpm`/`tpm` with 429 (`rate_limited`, with `Retry-After`), past its budget with 429 (`budget_exceeded`). OpenAI-compatible routes put the code in `error.type`, Anthropic routes in `error.type`, Ollama routes in `code`. With an Enterprise license, responses carry `X-CSGLite-Budget-Warning` (for example `key=85%, team=92%`) once a budget passes 80%.", + "properties": { + "models": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Models the key may call; an entry ending in `*` matches by prefix. Empty allows every model. `/v1/models` lists only these." + }, + "expires_at": { + "type": "string", + "format": "date-time" + }, + "rpm": { + "type": "integer", + "minimum": 0, + "description": "Requests over the last minute; 0 is unlimited." + }, + "tpm": { + "type": "integer", + "minimum": 0, + "description": "Tokens over the last minute; 0 is unlimited." + }, + "budget": { + "$ref": "#/components/schemas/Budget" + } + } + }, + "Budget": { + "type": "object", + "additionalProperties": false, + "description": "Caps what a key, a user or a team uses per period. A zero cap is not checked.", + "properties": { + "tokens": { + "type": "integer", + "minimum": 0 + }, + "cost": { + "type": "number", + "minimum": 0, + "description": "Counted only from requests priced in `currency`." + }, + "currency": { + "type": "string", + "pattern": "^[A-Z]{3}$", + "description": "Required with `cost`, for example CNY or USD." + }, + "period": { + "type": "string", + "enum": [ + "", + "day", + "month" + ], + "description": "`day` and `month` restart at local midnight and on the first of the month; empty never restarts." + } + } + }, + "BudgetSpend": { + "type": "object", + "additionalProperties": false, + "required": [ + "tokens", + "cost", + "exhausted" + ], + "properties": { + "tokens": { + "type": "integer" + }, + "cost": { + "type": "number" + }, + "since": { + "type": "string", + "format": "date-time", + "description": "Start of the current period; absent for a budget that never restarts." + }, + "exhausted": { + "type": "boolean" + }, + "warning": { + "type": "boolean", + "description": "At 80% of a cap. Reported only with an Enterprise license." + } + } + }, + "BudgetRequest": { + "type": "object", + "additionalProperties": false, + "required": [ + "budget" + ], + "properties": { + "budget": { + "oneOf": [ + { + "$ref": "#/components/schemas/Budget" + }, + { + "type": "null" + } + ] } } }, diff --git a/internal/server/team_handlers.go b/internal/server/team_handlers.go index defa6cd..3d87470 100644 --- a/internal/server/team_handlers.go +++ b/internal/server/team_handlers.go @@ -13,6 +13,7 @@ import ( "github.com/opencsgs/csglite/internal/config" "github.com/opencsgs/csglite/internal/httpjson" + "github.com/opencsgs/csglite/internal/keygov" "github.com/opencsgs/csglite/internal/license" "github.com/opencsgs/csglite/internal/team" "github.com/opencsgs/csglite/pkg/api" @@ -347,7 +348,7 @@ func (s *Server) handleTeamGroupsList(w http.ResponseWriter, r *http.Request) { } out := make([]api.TeamGroup, 0, len(teams)) for _, group := range teams { - out = append(out, teamGroupAPI(group, members[group.ID])) + out = append(out, s.teamGroupAPI(group, members[group.ID])) } writeJSON(w, http.StatusOK, api.TeamGroupsResponse{Teams: out}) } @@ -371,7 +372,7 @@ func (s *Server) handleTeamGroupCreate(w http.ResponseWriter, r *http.Request) { writeTeamError(w, err) return } - writeJSON(w, http.StatusCreated, teamGroupAPI(group, 0)) + writeJSON(w, http.StatusCreated, s.teamGroupAPI(group, 0)) } // PUT /api/team/teams/{id} @@ -386,7 +387,7 @@ func (s *Server) handleTeamGroupUpdate(w http.ResponseWriter, r *http.Request) { writeTeamError(w, err) return } - writeJSON(w, http.StatusOK, teamGroupAPI(group, s.teamMemberCount(group.ID))) + writeJSON(w, http.StatusOK, s.teamGroupAPI(group, s.teamMemberCount(group.ID))) } // DELETE /api/team/teams/{id} @@ -522,15 +523,19 @@ func (s *Server) teamUserAPI(user team.User, teams map[string]string, invites ma UpdatedAt: user.UpdatedAt, LastLoginAt: lastLogin, InviteExpiresAt: inviteExpires, + Budget: user.Budget, + Spend: s.budgetSpend(keygov.UserSubject(user.ID), user.Budget), } } -func teamGroupAPI(group team.Team, members int) api.TeamGroup { +func (s *Server) teamGroupAPI(group team.Team, members int) api.TeamGroup { return api.TeamGroup{ ID: group.ID, Name: group.Name, Description: group.Description, MemberCount: members, + Budget: group.Budget, + Spend: s.budgetSpend(keygov.TeamSubject(group.ID), group.Budget), CreatedAt: group.CreatedAt, UpdatedAt: group.UpdatedAt, } @@ -579,3 +584,45 @@ func ownsAPIKey(r *http.Request, record config.APIKeyRecord) bool { func secureEqual(a, b string) bool { return a != "" && b != "" && subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1 } + +// PUT /api/team/users/{id}/budget (Enterprise) +func (s *Server) handleTeamUserBudget(w http.ResponseWriter, r *http.Request) { + budget, ok := decodeBudget(w, r) + if !ok { + return + } + user, err := s.team.SetUserBudget(r.PathValue("id"), budget) + if err != nil { + writeTeamError(w, err) + return + } + writeJSON(w, http.StatusOK, s.teamUserAPI(user, nil, s.inviteExpiries())) +} + +// PUT /api/team/teams/{id}/budget (Enterprise) +func (s *Server) handleTeamGroupBudget(w http.ResponseWriter, r *http.Request) { + budget, ok := decodeBudget(w, r) + if !ok { + return + } + group, err := s.team.SetTeamBudget(r.PathValue("id"), budget) + if err != nil { + writeTeamError(w, err) + return + } + writeJSON(w, http.StatusOK, s.teamGroupAPI(group, s.teamMemberCount(group.ID))) +} + +func decodeBudget(w http.ResponseWriter, r *http.Request) (*api.Budget, bool) { + var req api.BudgetRequest + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 64<<10)).Decode(&req); err != nil { + writeError(w, http.StatusBadRequest, "invalid request body") + return nil, false + } + budget, err := keygov.NormalizeBudget(req.Budget) + if err != nil { + writeError(w, http.StatusBadRequest, err.Error()) + return nil, false + } + return budget, true +} diff --git a/internal/server/team_test.go b/internal/server/team_test.go index a1048e0..d58de67 100644 --- a/internal/server/team_test.go +++ b/internal/server/team_test.go @@ -350,7 +350,7 @@ func TestConversationsAndImagesArePrivateToTheirOwner(t *testing.T) { func TestAPIKeysBelongToTheirCreator(t *testing.T) { s, admin := newAccountsServer(t) member, mc := inviteUser(t, s, admin, "member", "member") - legacy, legacyPlain, err := s.apiKeys.Create("legacy", "") + legacy, legacyPlain, err := s.apiKeys.Create("legacy", "", nil) if err != nil { t.Fatal(err) } @@ -439,7 +439,7 @@ func ptr[T any](v T) *T { return &v } func TestAKeyFromBeforeAccountsCannotMintKeys(t *testing.T) { s, _ := newAccountsServer(t) - _, plain, err := s.apiKeys.Create("legacy", "") + _, plain, err := s.apiKeys.Create("legacy", "", nil) if err != nil { t.Fatal(err) } diff --git a/internal/server/usage.go b/internal/server/usage.go index 4cbb15e..0dff271 100644 --- a/internal/server/usage.go +++ b/internal/server/usage.go @@ -260,6 +260,7 @@ func (s *Server) recordAPIUsageWithPool(r *http.Request, model, source string, i costSnapshot, estimatedCost := s.requestPricingSnapshot( resolvedSource, pricingModel, int64(inputTokens), int64(outputTokens), ) + s.recordSpend(r, int64(inputTokens+outputTokens), estimatedCost, costSnapshot.Currency) observationFromContext(r.Context()).setUsage( model, resolvedSource, diff --git a/internal/team/policy.go b/internal/team/policy.go index 4dc4f44..b43a24a 100644 --- a/internal/team/policy.go +++ b/internal/team/policy.go @@ -101,11 +101,12 @@ func adminOnlyRead(path string) bool { return strings.HasPrefix(path, "/api/observability") || strings.HasPrefix(path, "/api/team/") } -// OwnKeyWrite reports whether a request creates or deletes one of the +// OwnKeyWrite reports whether a request creates, renames or deletes one of the // caller's own API keys. func OwnKeyWrite(method, path string) bool { return (path == "/api/api-keys" && method == http.MethodPost) || - (strings.HasPrefix(path, "/api/api-keys/") && path != "/api/api-keys/settings" && method == http.MethodDelete) + (strings.HasPrefix(path, "/api/api-keys/") && path != "/api/api-keys/settings" && + (method == http.MethodDelete || method == http.MethodPut)) } // memberWrite lists the changes a member may make besides running inference. diff --git a/internal/team/store.go b/internal/team/store.go index 0985922..1626366 100644 --- a/internal/team/store.go +++ b/internal/team/store.go @@ -21,6 +21,7 @@ import ( "time" "github.com/opencsgs/csglite/internal/safefile" + "github.com/opencsgs/csglite/pkg/api" ) // File is the runtime file under the app home that holds users, teams, @@ -75,10 +76,13 @@ type User struct { PasswordHash string `json:"password_hash,omitempty"` // MustChangePassword is set when someone other than the user chose the // password, so the user replaces it before doing anything else. - MustChangePassword bool `json:"must_change_password,omitempty"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` - LastLoginAt time.Time `json:"last_login_at,omitempty"` + MustChangePassword bool `json:"must_change_password,omitempty"` + // Budget caps what the user spends across all their keys and the web + // UI. It is an Enterprise feature; see ee/budget. + Budget *api.Budget `json:"budget,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + LastLoginAt time.Time `json:"last_login_at,omitempty"` } // Pending reports whether the user was invited and has not set a password. @@ -87,11 +91,14 @@ func (u User) Pending() bool { return u.PasswordHash == "" } // Team groups users. Usage and budgets are attributed to the team of the user // who owns a request. type Team struct { - ID string `json:"id"` - Name string `json:"name"` - Description string `json:"description,omitempty"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + // Budget caps what the team's members spend together. It is an + // Enterprise feature; see ee/budget. + Budget *api.Budget `json:"budget,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } // Invite is an outstanding invitation. Only the SHA-256 of its token is kept. @@ -805,6 +812,48 @@ func (s *Store) UpdateTeam(id string, name, description *string) (Team, error) { return team, nil } +// SetUserBudget replaces a user's budget; nil removes it. The budget must +// already be normalized. +func (s *Store) SetUserBudget(id string, budget *api.Budget) (User, error) { + s.mu.Lock() + defer s.mu.Unlock() + st, err := s.loadLocked() + if err != nil { + return User{}, err + } + i := st.userIndex(id) + if i < 0 { + return User{}, ErrNotFound + } + st.Users[i].Budget = budget + st.Users[i].UpdatedAt = s.now().UTC() + if err := s.saveLocked(st); err != nil { + return User{}, err + } + return st.Users[i], nil +} + +// SetTeamBudget replaces a team's budget; nil removes it. The budget must +// already be normalized. +func (s *Store) SetTeamBudget(id string, budget *api.Budget) (Team, error) { + s.mu.Lock() + defer s.mu.Unlock() + st, err := s.loadLocked() + if err != nil { + return Team{}, err + } + i := st.teamIndex(id) + if i < 0 { + return Team{}, ErrNotFound + } + st.Teams[i].Budget = budget + st.Teams[i].UpdatedAt = s.now().UTC() + if err := s.saveLocked(st); err != nil { + return Team{}, err + } + return st.Teams[i], nil +} + // DeleteTeam removes a team. Its members stay, without a team. func (s *Store) DeleteTeam(id string) error { s.mu.Lock() diff --git a/openapi/local-api.json b/openapi/local-api.json index 8a5f174..9cdc742 100644 --- a/openapi/local-api.json +++ b/openapi/local-api.json @@ -3780,6 +3780,122 @@ } } }, + "/api/team/users/{id}/budget": { + "put": { + "tags": [ + "Accounts" + ], + "summary": "Set a user's budget", + "description": "Sets or, with a null `budget`, removes the user's budget. It caps what the user spends across every key and the web UI, per day, per month or in total; requests past it are refused with 429 (`budget_exceeded`). Enterprise (`feature.lite.budgets`); reading budgets is not gated. Administrators only.", + "operationId": "setTeamUserBudget", + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/BudgetRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Updated user.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TeamUser" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest" + }, + "403": { + "description": "Not licensed (`code: feature_not_licensed`), the caller's role does not allow this (`code: permission_denied`), or the caller must sign in (`code: login_required`, 401).", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccountsErrorResponse" + } + } + } + }, + "404": { + "$ref": "#/components/responses/NotFound" + } + } + } + }, + "/api/team/teams/{id}/budget": { + "put": { + "tags": [ + "Accounts" + ], + "summary": "Set a team's budget", + "description": "Sets or, with a null `budget`, removes the team's budget. It caps what the team spends across every key and the web UI, per day, per month or in total; requests past it are refused with 429 (`budget_exceeded`). Enterprise (`feature.lite.budgets`); reading budgets is not gated. Administrators only.", + "operationId": "setTeamBudget", + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/BudgetRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Updated team.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TeamGroup" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest" + }, + "403": { + "description": "Not licensed (`code: feature_not_licensed`), the caller's role does not allow this (`code: permission_denied`), or the caller must sign in (`code: login_required`, 401).", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccountsErrorResponse" + } + } + } + }, + "404": { + "$ref": "#/components/responses/NotFound" + } + } + } + }, "/api/api-keys": { "get": { "tags": [ @@ -3877,6 +3993,62 @@ } }, "/api/api-keys/{id}": { + "put": { + "tags": [ + "System" + ], + "summary": "Update a local API key", + "description": "Renames a key or replaces its limits. The owner may rename a key; changing `limits` takes an administrator, and `{}` removes them. Limits are free in every edition.", + "operationId": "updateAPIKey", + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/APIKeyUpdateRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Updated key.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/APIKeyInfo" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest" + }, + "403": { + "description": "Only an administrator may change limits (`code: permission_denied`).", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccountsErrorResponse" + } + } + } + }, + "404": { + "$ref": "#/components/responses/NotFound" + } + } + }, "delete": { "tags": [ "System" @@ -10011,6 +10183,12 @@ "invite_token": { "type": "string", "description": "Returned once when an invitation is issued; the link is /invite/." + }, + "budget": { + "$ref": "#/components/schemas/Budget" + }, + "spend": { + "$ref": "#/components/schemas/BudgetSpend" } } }, @@ -10204,6 +10382,12 @@ "member_count": { "type": "integer" }, + "budget": { + "$ref": "#/components/schemas/Budget" + }, + "spend": { + "$ref": "#/components/schemas/BudgetSpend" + }, "created_at": { "type": "string", "format": "date-time" @@ -10248,7 +10432,8 @@ "id", "name", "prefix", - "created_at" + "created_at", + "status" ], "properties": { "owner_id": { @@ -10274,6 +10459,21 @@ "last_used_at": { "type": "string", "format": "date-time" + }, + "limits": { + "$ref": "#/components/schemas/KeyLimits" + }, + "spend": { + "$ref": "#/components/schemas/BudgetSpend" + }, + "status": { + "type": "string", + "enum": [ + "active", + "expired", + "exhausted" + ], + "description": "`exhausted` means the key's budget is used up for the current period." } } }, @@ -10302,6 +10502,130 @@ "properties": { "name": { "type": "string" + }, + "limits": { + "$ref": "#/components/schemas/KeyLimits" + } + } + }, + "APIKeyUpdateRequest": { + "type": "object", + "additionalProperties": false, + "properties": { + "name": { + "type": "string" + }, + "limits": { + "$ref": "#/components/schemas/KeyLimits" + } + } + }, + "KeyLimits": { + "type": "object", + "additionalProperties": false, + "description": "Restrictions on one API key, enforced on inference routes. A key past them is refused: an expired key with 401 (`api_key_expired`), a model outside `models` with 403 (`model_not_allowed`), over `rpm`/`tpm` with 429 (`rate_limited`, with `Retry-After`), past its budget with 429 (`budget_exceeded`). OpenAI-compatible routes put the code in `error.type`, Anthropic routes in `error.type`, Ollama routes in `code`. With an Enterprise license, responses carry `X-CSGLite-Budget-Warning` (for example `key=85%, team=92%`) once a budget passes 80%.", + "properties": { + "models": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Models the key may call; an entry ending in `*` matches by prefix. Empty allows every model. `/v1/models` lists only these." + }, + "expires_at": { + "type": "string", + "format": "date-time" + }, + "rpm": { + "type": "integer", + "minimum": 0, + "description": "Requests over the last minute; 0 is unlimited." + }, + "tpm": { + "type": "integer", + "minimum": 0, + "description": "Tokens over the last minute; 0 is unlimited." + }, + "budget": { + "$ref": "#/components/schemas/Budget" + } + } + }, + "Budget": { + "type": "object", + "additionalProperties": false, + "description": "Caps what a key, a user or a team uses per period. A zero cap is not checked.", + "properties": { + "tokens": { + "type": "integer", + "minimum": 0 + }, + "cost": { + "type": "number", + "minimum": 0, + "description": "Counted only from requests priced in `currency`." + }, + "currency": { + "type": "string", + "pattern": "^[A-Z]{3}$", + "description": "Required with `cost`, for example CNY or USD." + }, + "period": { + "type": "string", + "enum": [ + "", + "day", + "month" + ], + "description": "`day` and `month` restart at local midnight and on the first of the month; empty never restarts." + } + } + }, + "BudgetSpend": { + "type": "object", + "additionalProperties": false, + "required": [ + "tokens", + "cost", + "exhausted" + ], + "properties": { + "tokens": { + "type": "integer" + }, + "cost": { + "type": "number" + }, + "since": { + "type": "string", + "format": "date-time", + "description": "Start of the current period; absent for a budget that never restarts." + }, + "exhausted": { + "type": "boolean" + }, + "warning": { + "type": "boolean", + "description": "At 80% of a cap. Reported only with an Enterprise license." + } + } + }, + "BudgetRequest": { + "type": "object", + "additionalProperties": false, + "required": [ + "budget" + ], + "properties": { + "budget": { + "oneOf": [ + { + "$ref": "#/components/schemas/Budget" + }, + { + "type": "null" + } + ] } } }, diff --git a/pkg/api/types.go b/pkg/api/types.go index 6d81c9c..0b3a635 100644 --- a/pkg/api/types.go +++ b/pkg/api/types.go @@ -875,6 +875,64 @@ type APIKeyInfo struct { // A key without an owner predates accounts. OwnerID string `json:"owner_id,omitempty"` OwnerName string `json:"owner_name,omitempty"` + // Limits restrict what the key may do; nil means none. + Limits *KeyLimits `json:"limits,omitempty"` + // Spend is what the key used in its budget period, present when the key + // has a budget. + Spend *BudgetSpend `json:"spend,omitempty"` + // Status is "active", "expired" or "exhausted" (its budget is used up). + Status string `json:"status"` +} + +// Budget caps what a key, a user or a team may use in each period. A zero +// Tokens or Cost is not capped. +type Budget struct { + Tokens int64 `json:"tokens,omitempty"` + // Cost is counted only from requests whose price is known in Currency. + Cost float64 `json:"cost,omitempty"` + Currency string `json:"currency,omitempty"` + // Period is "day" or "month", which restart at local midnight and on the + // first of the month, or empty for a budget that never restarts. + Period string `json:"period,omitempty"` +} + +// BudgetSpend is what was used against a budget in its current period. +type BudgetSpend struct { + Tokens int64 `json:"tokens"` + Cost float64 `json:"cost"` + // Since is when the current period began; absent for a budget that never + // restarts. + Since *time.Time `json:"since,omitempty"` + // Exhausted is true once a cap is reached and requests are refused. + Exhausted bool `json:"exhausted"` + // Warning is true at 80% of a cap. It is reported only with an + // Enterprise license. + Warning bool `json:"warning,omitempty"` +} + +// KeyLimits restrict one API key. Zero values are not limited. +type KeyLimits struct { + // Models the key may call. An entry ending in "*" matches every model + // starting with the rest. Empty allows every model. + Models []string `json:"models,omitempty"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` + // RPM and TPM cap requests and tokens over the last minute. + RPM int `json:"rpm,omitempty"` + TPM int `json:"tpm,omitempty"` + Budget *Budget `json:"budget,omitempty"` +} + +// APIKeyUpdateRequest renames a key or replaces its limits. Only an +// administrator may change limits. +type APIKeyUpdateRequest struct { + Name *string `json:"name,omitempty"` + Limits *KeyLimits `json:"limits,omitempty"` +} + +// BudgetRequest sets or, with a null budget, removes a user's or a team's +// budget. +type BudgetRequest struct { + Budget *Budget `json:"budget"` } type APIKeysResponse struct { @@ -883,7 +941,8 @@ type APIKeysResponse struct { } type APIKeyCreateRequest struct { - Name string `json:"name,omitempty"` + Name string `json:"name,omitempty"` + Limits *KeyLimits `json:"limits,omitempty"` } type APIKeyCreateResponse struct { @@ -1556,7 +1615,9 @@ type TeamUser struct { InviteExpiresAt *time.Time `json:"invite_expires_at,omitempty"` // InviteToken is returned once, when an invitation is issued. The link to // send is /invite/ on this instance. - InviteToken string `json:"invite_token,omitempty"` + InviteToken string `json:"invite_token,omitempty"` + Budget *Budget `json:"budget,omitempty"` + Spend *BudgetSpend `json:"spend,omitempty"` } type TeamUsersResponse struct { @@ -1591,12 +1652,14 @@ type TeamUserUpdateRequest struct { } type TeamGroup struct { - ID string `json:"id"` - Name string `json:"name"` - Description string `json:"description,omitempty"` - MemberCount int `json:"member_count"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + MemberCount int `json:"member_count"` + Budget *Budget `json:"budget,omitempty"` + Spend *BudgetSpend `json:"spend,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } type TeamGroupsResponse struct { diff --git a/web/src/api/client.ts b/web/src/api/client.ts index a2dcd67..e8974a5 100644 --- a/web/src/api/client.ts +++ b/web/src/api/client.ts @@ -591,8 +591,44 @@ export interface LocalAPIKeyInfo { /** User account that owns the key when accounts are on. */ owner_id?: string; owner_name?: string; + limits?: KeyLimits; + /** Present when the key has a budget. */ + spend?: BudgetSpend; + status: "active" | "expired" | "exhausted"; } +/** Caps what a key, a user or a team uses per period; 0 is not capped. */ +export interface Budget { + tokens?: number; + /** Counted only from requests priced in `currency`. */ + cost?: number; + currency?: string; + /** "day", "month", or empty for a budget that never restarts. */ + period?: "" | "day" | "month"; +} + +export interface BudgetSpend { + tokens: number; + cost: number; + since?: string; + exhausted: boolean; + /** At 80% of a cap; reported only with an Enterprise license. */ + warning?: boolean; +} + +/** Restrictions on one API key. They are free in every edition. */ +export interface KeyLimits { + /** An entry ending in "*" matches by prefix; empty allows every model. */ + models?: string[]; + expires_at?: string; + rpm?: number; + tpm?: number; + budget?: Budget; +} + +/** Enterprise: budgets for users and teams, and the 80% warning. */ +export const BUDGETS_FEATURE = "feature.lite.budgets"; + export interface LocalAPIKeysResponse { auth_enabled: boolean; keys: LocalAPIKeyInfo[]; @@ -1809,11 +1845,20 @@ export async function updateLocalAPIKeySettings(authEnabled: boolean): Promise { +export async function createLocalAPIKey(name: string, limits?: KeyLimits): Promise { return fetchJSON("/api/api-keys", { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ name }), + body: JSON.stringify({ name, limits }), + }); +} + +/** Only an administrator may change limits; `{}` removes them. */ +export async function updateLocalAPIKey(id: string, req: { name?: string; limits?: KeyLimits }): Promise { + return fetchJSON(`/api/api-keys/${encodeURIComponent(id)}`, { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(req), }); } @@ -3486,6 +3531,8 @@ export interface TeamUser { invite_expires_at?: string; /** Returned once when an invitation is issued; the link is /invite/. */ invite_token?: string; + budget?: Budget; + spend?: BudgetSpend; } export interface TeamSettings { @@ -3521,6 +3568,8 @@ export interface TeamGroup { name: string; description?: string; member_count: number; + budget?: Budget; + spend?: BudgetSpend; created_at: string; updated_at: string; } @@ -3609,6 +3658,16 @@ export async function deleteTeamUser(id: string): Promise { await fetchAccountsJSON<{ status: string }>(`/api/team/users/${encodeURIComponent(id)}`, { method: "DELETE" }); } +/** Enterprise. A null budget removes it. */ +export function setTeamUserBudget(id: string, budget: Budget | null): Promise { + return fetchAccountsJSON(`/api/team/users/${encodeURIComponent(id)}/budget`, { method: "PUT", body: JSON.stringify({ budget }) }); +} + +/** Enterprise. A null budget removes it. */ +export function setTeamBudget(id: string, budget: Budget | null): Promise { + return fetchAccountsJSON(`/api/team/teams/${encodeURIComponent(id)}/budget`, { method: "PUT", body: JSON.stringify({ budget }) }); +} + export async function getTeams(): Promise { const data = await fetchAccountsJSON<{ teams: TeamGroup[] }>("/api/team/teams", { cache: "no-store" }); return data.teams || []; diff --git a/web/src/components/KeyLimits.tsx b/web/src/components/KeyLimits.tsx new file mode 100644 index 0000000..43f7d1a --- /dev/null +++ b/web/src/components/KeyLimits.tsx @@ -0,0 +1,411 @@ +import { useEffect, useState } from "preact/hooks"; +import { createLocalAPIKey, updateLocalAPIKey } from "../api/client"; +import type { Budget, BudgetSpend, KeyLimits, LocalAPIKeyInfo } from "../api/client"; +import { locale, t } from "../i18n"; + +const fieldClass = + "mt-1 w-full rounded-lg border border-gray-200 px-3 py-2 text-sm focus:outline-none focus:ring-2 focus:ring-indigo-500 disabled:bg-gray-50"; + +/** Editable form of KeyLimits: numbers and lists kept as typed text. */ +export interface LimitsDraft { + models: string; + expiresAt: string; + rpm: string; + tpm: string; + budget: BudgetDraft; +} + +export interface BudgetDraft { + tokens: string; + cost: string; + currency: string; + period: "" | "day" | "month"; +} + +export function budgetDraft(budget?: Budget | null): BudgetDraft { + return { + tokens: budget?.tokens ? String(budget.tokens) : "", + cost: budget?.cost ? String(budget.cost) : "", + currency: budget?.currency || "CNY", + period: budget?.period || "month", + }; +} + +export function limitsDraft(limits?: KeyLimits | null): LimitsDraft { + return { + models: (limits?.models || []).join("\n"), + expiresAt: limits?.expires_at ? limits.expires_at.slice(0, 10) : "", + rpm: limits?.rpm ? String(limits.rpm) : "", + tpm: limits?.tpm ? String(limits.tpm) : "", + budget: budgetDraft(limits?.budget), + }; +} + +function wholeNumber(value: string): number { + const n = Math.floor(Number(value.trim())); + return Number.isFinite(n) && n > 0 ? n : 0; +} + +/** The budget a draft describes, or null when it caps nothing. */ +export function draftBudget(draft: BudgetDraft): Budget | null { + const tokens = wholeNumber(draft.tokens); + const cost = Number(draft.cost.trim()); + const budget: Budget = { period: draft.period }; + if (tokens > 0) budget.tokens = tokens; + if (Number.isFinite(cost) && cost > 0) { + budget.cost = cost; + budget.currency = draft.currency.trim().toUpperCase(); + } + return budget.tokens || budget.cost ? budget : null; +} + +/** The limits a draft describes; `{}` means none. A date expires at the end of that local day. */ +export function draftLimits(draft: LimitsDraft): KeyLimits { + const limits: KeyLimits = {}; + const models = draft.models + .split(/[\n,]/) + .map((m) => m.trim()) + .filter(Boolean); + if (models.length) limits.models = models; + if (draft.expiresAt) limits.expires_at = new Date(`${draft.expiresAt}T23:59:59`).toISOString(); + const rpm = wholeNumber(draft.rpm); + const tpm = wholeNumber(draft.tpm); + if (rpm) limits.rpm = rpm; + if (tpm) limits.tpm = tpm; + const budget = draftBudget(draft.budget); + if (budget) limits.budget = budget; + return limits; +} + +function compact(n: number): string { + return new Intl.NumberFormat(locale.value === "zh" ? "zh-CN" : "en-US", { notation: "compact", maximumFractionDigits: 1 }).format(n); +} + +function money(n: number, currency?: string): string { + try { + return new Intl.NumberFormat(locale.value === "zh" ? "zh-CN" : "en-US", { style: "currency", currency: currency || "CNY", maximumFractionDigits: 2 }).format(n); + } catch { + return `${currency || ""} ${n.toFixed(2)}`; + } +} + +/** "1.2K / 10K tokens · ¥3.20 / ¥50.00 · per month" */ +export function budgetText(budget: Budget, spend?: BudgetSpend): string { + const parts: string[] = []; + if (budget.tokens) parts.push(t("limits.tokensOf", compact(spend?.tokens ?? 0), compact(budget.tokens))); + if (budget.cost) parts.push(`${money(spend?.cost ?? 0, budget.currency)} / ${money(budget.cost, budget.currency)}`); + parts.push(t(`limits.period.${budget.period || "total"}`)); + return parts.join(" · "); +} + +/** Largest share of a cap used, 0–1. */ +export function budgetUsed(budget: Budget, spend?: BudgetSpend): number { + if (!spend) return 0; + let used = 0; + if (budget.tokens) used = Math.max(used, spend.tokens / budget.tokens); + if (budget.cost) used = Math.max(used, spend.cost / budget.cost); + return Math.min(used, 1); +} + +/** A budget as a line of text over a thin bar. */ +export function BudgetMeter({ budget, spend }: { budget: Budget; spend?: BudgetSpend }) { + const used = budgetUsed(budget, spend); + const tone = spend?.exhausted ? "bg-red-500" : spend?.warning ? "bg-amber-500" : "bg-indigo-500"; + return ( +
+

{budgetText(budget, spend)}

+
+
+
+
+ ); +} + +const chip = "rounded-full px-2.5 py-1"; + +/** The limits on a key as chips, for a key list. */ +export function KeyLimitChips({ apiKey }: { apiKey: LocalAPIKeyInfo }) { + const limits = apiKey.limits; + return ( + <> + {apiKey.status === "expired" && {t("limits.statusExpired")}} + {apiKey.status === "exhausted" && {t("limits.statusExhausted")}} + {apiKey.spend?.warning && apiKey.status === "active" && {t("limits.statusWarning")}} + {limits?.models?.length ? ( + + {limits.models.length === 1 ? t("limits.oneModel", limits.models[0]) : t("limits.models", limits.models.length)} + + ) : null} + {limits?.expires_at && apiKey.status !== "expired" && ( + {t("limits.expires", new Date(limits.expires_at).toLocaleDateString(locale.value === "zh" ? "zh-CN" : "en-US"))} + )} + {limits?.rpm ? {t("limits.rpmChip", limits.rpm)} : null} + {limits?.tpm ? {t("limits.tpmChip", compact(limits.tpm))} : null} + {limits?.budget && {budgetText(limits.budget, apiKey.spend)}} + + ); +} + +/** Token and cost caps with their period. */ +export function BudgetFields({ draft, onChange, disabled }: { draft: BudgetDraft; onChange: (draft: BudgetDraft) => void; disabled?: boolean }) { + return ( +
+ + + + +

{t("limits.costHint")}

+
+ ); +} + +/** Every key limit. */ +export function LimitsFields({ draft, onChange, disabled }: { draft: LimitsDraft; onChange: (draft: LimitsDraft) => void; disabled?: boolean }) { + return ( +
+