From 908110bbbb4066a7f74371a2c0f039540371946d Mon Sep 17 00:00:00 2001 From: ganisback <370036720@qq.com> Date: Thu, 22 Oct 2026 17:46:21 +0800 Subject: [PATCH 1/3] Add key limits, and user and team budgets as an Enterprise feature An API key can now be limited to a list of models, given an expiry date, rate-limited by requests and tokens per minute, and capped by a token or cost budget per day, per month or in total. As in LiteLLM's open-source proxy these limits are free; they live in internal/keygov with a spend ledger (spend.db) that counts every metered request against its key, its owner and the owner's team. The Enterprise part, gated by feature.lite.budgets, is budgets for users and teams, counted across every key and the web UI, and the warning at 80% of any budget (X-CSGLite-Budget-Warning, and a flag in list responses). Only that lives in ee/budget. Inference requests past a limit are refused before they run: 401 api_key_expired, 403 model_not_allowed, 429 budget_exceeded or 429 rate_limited with Retry-After. /v1/models lists only allowed models. Anyone may restrict a key they create; afterwards only an administrator changes limits, so a member cannot lift what was set for them. Co-Authored-By: Claude Opus 5.5 --- docs/guides/ee-license-design.md | 6 +- docs/guides/key-governance.md | 68 ++++ docs/guides/team-management.md | 4 +- ee/README.md | 14 + ee/budget/budget.go | 59 ++++ ee/budget/budget_test.go | 50 +++ internal/config/api_auth.go | 59 +++- internal/keygov/keygov_test.go | 132 +++++++ internal/keygov/ledger.go | 161 +++++++++ internal/keygov/limiter.go | 92 +++++ internal/keygov/limits.go | 198 +++++++++++ internal/license/features.go | 11 +- internal/license/manager_test.go | 4 +- internal/server/auth_test.go | 18 +- internal/server/handlers_api_keys.go | 82 ++++- internal/server/handlers_models.go | 1 + internal/server/handlers_openai.go | 1 + internal/server/key_governance.go | 269 +++++++++++++++ internal/server/key_governance_test.go | 245 +++++++++++++ internal/server/license_test.go | 6 +- internal/server/routes.go | 8 +- internal/server/server.go | 6 + internal/server/static/openapi/local-api.json | 326 +++++++++++++++++- internal/server/team_handlers.go | 55 ++- internal/server/team_test.go | 4 +- internal/server/usage.go | 1 + internal/team/policy.go | 5 +- internal/team/store.go | 67 +++- openapi/local-api.json | 326 +++++++++++++++++- pkg/api/types.go | 79 ++++- web/src/api/client.ts | 63 +++- web/src/components/KeyLimits.tsx | 302 ++++++++++++++++ web/src/i18n.ts | 84 ++++- web/src/keyLimits.test.ts | 42 +++ web/src/pages/AIGateway.tsx | 55 ++- web/src/pages/AccountKeys.tsx | 6 + web/src/pages/Team.tsx | 96 +++++- 37 files changed, 2914 insertions(+), 91 deletions(-) create mode 100644 docs/guides/key-governance.md create mode 100644 ee/budget/budget.go create mode 100644 ee/budget/budget_test.go create mode 100644 internal/keygov/keygov_test.go create mode 100644 internal/keygov/ledger.go create mode 100644 internal/keygov/limiter.go create mode 100644 internal/keygov/limits.go create mode 100644 internal/server/key_governance.go create mode 100644 internal/server/key_governance_test.go create mode 100644 web/src/components/KeyLimits.tsx create mode 100644 web/src/keyLimits.test.ts diff --git a/docs/guides/ee-license-design.md b/docs/guides/ee-license-design.md index 882538e1..e949e946 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 00000000..56853cb3 --- /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 babcc283..b3aa7850 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 a8657a79..56b6af52 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 00000000..b8774452 --- /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 00000000..f04b8623 --- /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 a41e7254..aa1d04f3 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 00000000..8810582f --- /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 00000000..e9947464 --- /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 00000000..915e779a --- /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 00000000..b30de9e8 --- /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 188d0f44..d06554f9 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 9f0b45a8..887005a8 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 a6743f4f..4c26d77b 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 3f7c754e..2e168529 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 5f8bf016..07847cb9 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 395dd3f7..a95b235d 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 00000000..bcd81480 --- /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 00000000..e1aba580 --- /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 d54ad97c..778a74e6 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 9085e1ac..404ad173 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 88ab3e53..efc9416c 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 8a5f1748..9cdc7420 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 defa6cdb..3d874706 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 a1048e06..d58de679 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 4cbb15e7..0dff271d 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 4dc4f446..b43a24a1 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 09859228..16263667 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 8a5f1748..9cdc7420 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 6d81c9c7..0b3a635f 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 a2dcd674..e8974a53 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 00000000..0099b82c --- /dev/null +++ b/web/src/components/KeyLimits.tsx @@ -0,0 +1,302 @@ +import { useState } from "preact/hooks"; +import { 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 ( +
+