[EE] Key limits, and user and team budgets

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)）。其余已发布功能仍不
+> 门控。
 
 
 | 阶段 | 内容 | 估时 |

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% 提醒停止。

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）：复用这里的会话和邀请。

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.

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
+}

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")
+	}
+}

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)