diff --git a/.gitignore b/.gitignore index a2af2f8..e9b4404 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,9 @@ -credentials.json +zhanlu.db +zhanlu.db-shm +zhanlu.db-wal +*.db +*.db-shm +*.db-wal extension/ *.exe *.log diff --git a/README.md b/README.md index 2edad4e..e7fd84b 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ - 模型 API Key 换取签名逻辑：SM3 摘要 + SM2 签名（`X-Auth-Signature`/`X-Auth-Timestamp`/`X-Auth-Nonce`）。 - 上游 SSE 直接透传为 OpenAI SSE；非流式请求在本地聚合为 OpenAI Chat Completion JSON。 - OpenAI 函数/工具调用：支持 `tools`、`tool_choice`、流式 `delta.tool_calls`、非流式 `message.tool_calls` 以及 `role: tool` 结果续传。 +- Token 消耗统计：流式与非流式请求均解析上游 `usage`，按模型/按日/最近明细写入本地 SQLite（`zhanlu.db`），凭据也一并持久化在同一个库中。管理页提供 `GET /admin/stats` 可视化与 `GET /api/stats` JSON 接口。 ## 运行 @@ -37,17 +38,17 @@ http://127.0.0.1:8080/ /opt/zhanlu-proxy/zhanlu-proxy ``` -凭据文件放在： +数据库放在： ```text -/opt/zhanlu-proxy/credentials.json +/opt/zhanlu-proxy/zhanlu.db ``` 创建环境变量文件 `/etc/zhanlu-proxy/zhanlu-proxy.env`： ```env ZHANLU_LISTEN_ADDR=:8080 -ZHANLU_CREDENTIALS_FILE=/opt/zhanlu-proxy/credentials.json +ZHANLU_DB_FILE=/opt/zhanlu-proxy/zhanlu.db ZHANLU_MOBILE_LOGIN_BASE_URL=https://ecloud.10086.cn ZHANLU_MOBILE_MODEL_BASE_URL=https://ecloud.10086.cn/api/query/aigateway ZHANLU_UPSTREAM_TIMEOUT=300s @@ -106,20 +107,22 @@ journalctl -u zhanlu-proxy -f - 使用本次 `secret` AES 解密响应中的 `ak`、`sk`、`license`，得到 `AccessKey`、`SecretKey`、`Token`。 - 按插件 v1.4.2 流程调用 `/api/acepilot/zhanlu/v1/login`（RSA+HmacSHA1 签名 URL + `plugin_type=zhanlu_ide` 请求头）获取用户资料（email/组织/团队）。 - 用 SM2 私钥签名调用 `{mobileModelBaseUrl}/user/api/v2/external/key/get-or-create` 换取模型 `apiKey`。 -- 凭据（含 `apiKey`、`modelBaseUrl`、email 等）会写入 JSON 文件，后续 OpenAI 兼容接口自动使用。 +- 凭据（含 `apiKey`、`modelBaseUrl`、email 等）会写入本地 SQLite 数据库（`zhanlu.db`），后续 OpenAI 兼容接口自动使用。 -默认保存到当前执行目录： +默认数据库保存在当前执行目录： ```text -credentials.json +zhanlu.db ``` 可以通过环境变量覆盖： ```powershell -$env:ZHANLU_CREDENTIALS_FILE="E:\path\to\credentials.json" +$env:ZHANLU_DB_FILE="E:\path\to\zhanlu.db" ``` +凭据仅持久化在数据库中，不再使用 JSON 文件。 + 手机号验证码登录使用 `ZHANLU_MOBILE_LOGIN_BASE_URL`，默认公网地址来自插件配置（兼容旧环境变量 `ZHANLU_SERVER_BASE_URL`）： ```powershell @@ -184,7 +187,8 @@ curl http://127.0.0.1:8080/v1/models ` | `ZHANLU_MOBILE_LOGIN_BASE_URL` | `https://ecloud.10086.cn` | 移动云公网登录 Base URL（兼容旧变量 `ZHANLU_SERVER_BASE_URL`） | | `ZHANLU_MOBILE_MODEL_BASE_URL` | `https://ecloud.10086.cn/api/query/aigateway` | 移动云公网模型网关 Base URL | | `ZHANLU_UPSTREAM_PATH` | `/chat/completions` | 模型网关聊天接口路径 | -| `ZHANLU_CREDENTIALS_FILE` | `credentials.json` | 凭据 JSON 路径，默认当前执行目录 | +| `ZHANLU_DB_FILE` | `zhanlu.db` | 本地 SQLite 数据库路径，凭据与 token 统计均存于此，默认当前执行目录 | +| `ZHANLU_STATS_DISABLED` | `false` | 设为 `true` 关闭 token 用量记录（仅停止写入统计，凭据存储不受影响） | | `ZHANLU_ACCESS_KEY` | 空 | 直接从环境变量提供 AccessKey | | `ZHANLU_SECRET_KEY` | 空 | 直接从环境变量提供 SecretKey | | `ZHANLU_TOKEN` | 空 | 直接从环境变量提供 Token | @@ -207,14 +211,33 @@ curl http://127.0.0.1:8080/v1/models ` 服务启动时按以下优先级加载凭据： 1. 环境变量 `ZHANLU_ACCESS_KEY`、`ZHANLU_SECRET_KEY`、`ZHANLU_TOKEN` 或 `ZHANLU_API_KEY`。 -2. `ZHANLU_CREDENTIALS_FILE` 指向的 JSON 文件。 +2. `ZHANLU_DB_FILE` 数据库中持久化的凭据行。 -登录页面保存后，运行中的服务会立即使用新凭据。若环境中只有 AK/SK/Token 而没有 `apiKey`，首次调用聊天接口时会自动按插件流程换取 API Key 并回写凭据文件。 +登录页面保存后，运行中的服务会立即使用新凭据并写入数据库。若环境中只有 AK/SK/Token 而没有 `apiKey`，首次调用聊天接口时会自动按插件流程换取 API Key 并回写数据库。 + +## Token 消耗统计 + +代理在 `/v1/chat/completions` 完成后解析上游 `usage`（流式路径在透传 SSE 的同时旁路解析末块 `usage`，非流式路径在聚合时解析），将以下维度写入 `ZHANLU_DB_FILE`： + +- 时间、模型、流式/非流式 +- `prompt_tokens` / `completion_tokens` / `total_tokens` / `reasoning_tokens`（思考模型） +- `cached_tokens`（来自 `prompt_tokens_details.cached_tokens`，提示缓存命中的 token 数）与缓存命中率 +- 请求状态（`success` / `upstream_error`）与耗时 + +缓存命中率为 `cached_tokens / prompt_tokens`，仅在模型/网关支持 prompt 缓存且上游返回 `cached_tokens` 时非零。 + +管理页（需先登录管理页面）提供： + +- `GET /admin/stats`：可视化页面，展示总览、按模型、按日柱状、最近请求明细，带重置按钮。 +- `GET /api/stats`：JSON 接口，支持 `since`/`until`（RFC3339）、`model`、`limit` 查询参数。 +- `POST /api/stats/reset`：清空统计（凭据不受影响）。 + +设置 `ZHANLU_STATS_DISABLED=true` 可停止写入统计。 ## 安全说明 -- `credentials.json` 包含明文 `AccessKey`、`SecretKey`、`Token` 和 `apiKey`，请不要提交到仓库。 -- 默认保存在当前执行目录的 `credentials.json`。 +- `zhanlu.db` 数据库包含明文 `AccessKey`、`SecretKey`、`Token` 和 `apiKey`，请不要提交到仓库。 +- 默认保存在当前执行目录的 `zhanlu.db`（建议通过 `ZHANLU_DB_FILE` 指向受保护路径）。 - 建议设置 `ZHANLU_LOGIN_PASSWORD`，避免公网暴露的 `/admin/login` 被直接访问。 - 错误响应默认不会返回签名 URL，避免泄露 `AccessKey`、`authorization`、`Signature`。 - `ZHANLU_DEBUG=true` 时会返回更详细错误，但仍会对敏感 query 参数脱敏。 diff --git a/cmd/zhanlu-proxy/main.go b/cmd/zhanlu-proxy/main.go index 7ab502e..0d0099a 100644 --- a/cmd/zhanlu-proxy/main.go +++ b/cmd/zhanlu-proxy/main.go @@ -6,6 +6,7 @@ import ( "git.misaka.ren/M1saka/zhanlu_proxy/internal/config" "git.misaka.ren/M1saka/zhanlu_proxy/internal/server" + "git.misaka.ren/M1saka/zhanlu_proxy/internal/store" ) func main() { @@ -14,7 +15,20 @@ func main() { log.Fatal(err) } - h := server.New(cfg) + st, err := store.Open(cfg.DBPath) + if err != nil { + log.Fatalf("open store: %v", err) + } + defer st.Close() + + // Credential precedence at startup: environment variables > db row. + if cfg.Credentials.Validate() != nil && !cfg.Credentials.HasAPIKey() { + if dbCreds, derr := st.LoadCredentials(); derr == nil && (dbCreds.HasAPIKey() || dbCreds.Validate() == nil) { + cfg.Credentials = dbCreds + } + } + + h := server.New(cfg, st) log.Printf("zhanlu proxy listening on %s", cfg.ListenAddr) log.Printf("login page: http://127.0.0.1%s/login", cfg.ListenAddr) if err := http.ListenAndServe(cfg.ListenAddr, h); err != nil { diff --git a/go.mod b/go.mod index 8ce4ba6..75e68e7 100644 --- a/go.mod +++ b/go.mod @@ -4,4 +4,16 @@ go 1.25.0 require github.com/emmansun/gmsm v0.44.1 -require golang.org/x/crypto v0.54.0 // indirect +require ( + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/mattn/go-isatty v0.0.24 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + golang.org/x/crypto v0.54.0 // indirect + golang.org/x/sys v0.47.0 // indirect + modernc.org/libc v1.74.4 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect + modernc.org/sqlite v1.56.0 // indirect +) diff --git a/go.sum b/go.sum index f7e4238..458584b 100644 --- a/go.sum +++ b/go.sum @@ -1,6 +1,24 @@ +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/emmansun/gmsm v0.44.1 h1:zDTkdtLWFG0vCbhPV+k9pte14tix/eK71At9Iai9fP4= github.com/emmansun/gmsm v0.44.1/go.mod h1:p6RIUta0/KboFHrOxr1x8q+pd8RZtdaTO7XNp0RmMQM= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI= +github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw= golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk= golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +modernc.org/libc v1.74.4 h1:fX1Omw4o2/1C2iRkkIsrQTasJQldLhRmuPreXLoWs9k= +modernc.org/libc v1.74.4/go.mod h1:eeQAS9W3sZeKYMFubydxJpII9ybHWshk+7or7bLG9co= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/sqlite v1.56.0 h1:/D8e2RfFqoy/Zc6PuC76U28zFwmI/sYx1Kjm4yEn9e0= +modernc.org/sqlite v1.56.0/go.mod h1:yCJ2cmAaIkHQ25oXWrF8H4O1lIfPYPR26yCEDj2P3pQ= diff --git a/internal/auth/credentials.go b/internal/auth/credentials.go index 61196b6..ffae077 100644 --- a/internal/auth/credentials.go +++ b/internal/auth/credentials.go @@ -1,10 +1,7 @@ package auth import ( - "encoding/json" "errors" - "os" - "path/filepath" "strings" "time" ) @@ -49,30 +46,3 @@ func (c Credentials) Validate() error { func (c Credentials) HasAPIKey() bool { return strings.TrimSpace(c.APIKey) != "" && strings.TrimSpace(c.ModelBaseURL) != "" } - -func LoadCredentials(path string) (Credentials, error) { - b, err := os.ReadFile(path) - if err != nil { - return Credentials{}, err - } - var c Credentials - if err := json.Unmarshal(b, &c); err != nil { - return Credentials{}, err - } - return c, c.Validate() -} - -func SaveCredentials(path string, c Credentials) error { - if err := c.Validate(); err != nil { - return err - } - c.SavedAt = time.Now() - b, err := json.MarshalIndent(c, "", " ") - if err != nil { - return err - } - if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { - return err - } - return os.WriteFile(path, b, 0o600) -} diff --git a/internal/config/config.go b/internal/config/config.go index e85bc82..728b57b 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -2,7 +2,6 @@ package config import ( "os" - "path/filepath" "strings" "time" @@ -14,7 +13,8 @@ type Config struct { MobileLoginBaseURL string MobileModelBaseURL string UpstreamPath string - CredentialsPath string + DBPath string + StatsDisabled bool SSOExchangeURL string SSOBaseURL string TokenDecryptKey string @@ -36,7 +36,8 @@ func Load() (Config, error) { MobileLoginBaseURL: firstNonEmpty(os.Getenv("ZHANLU_MOBILE_LOGIN_BASE_URL"), getenv("ZHANLU_SERVER_BASE_URL", "https://ecloud.10086.cn")), MobileModelBaseURL: getenv("ZHANLU_MOBILE_MODEL_BASE_URL", "https://ecloud.10086.cn/api/query/aigateway"), UpstreamPath: getenv("ZHANLU_UPSTREAM_PATH", "/chat/completions"), - CredentialsPath: getenv("ZHANLU_CREDENTIALS_FILE", defaultCredentialsPath()), + DBPath: getenv("ZHANLU_DB_FILE", "zhanlu.db"), + StatsDisabled: strings.EqualFold(os.Getenv("ZHANLU_STATS_DISABLED"), "true"), SSOBaseURL: getenv("ZHANLU_SSO_BASE_URL", "http://4c.hq.cmcc"), SSOExchangeURL: getenv("ZHANLU_SSO_EXCHANGE_URL", "http://rdcloud.4c.hq.cmcc/cmdevops-aiplus-agent-gateway/api/acepilot/zhanlu/authToken"), TokenDecryptKey: getenv("ZHANLU_TOKEN_DECRYPT_KEY", "3jw7woww2rvhla6k"), @@ -56,12 +57,8 @@ func Load() (Config, error) { Token: os.Getenv("ZHANLU_TOKEN"), APIKey: os.Getenv("ZHANLU_API_KEY"), } - if cfg.Credentials.Validate() == nil || cfg.Credentials.HasAPIKey() { - return cfg, nil - } - if creds, err := auth.LoadCredentials(cfg.CredentialsPath); err == nil { - cfg.Credentials = creds - } + // Credentials from the environment were incomplete; the store serves the + // persisted row at runtime. return cfg, nil } @@ -93,10 +90,6 @@ func firstNonEmpty(values ...string) string { return "" } -func defaultCredentialsPath() string { - return filepath.Join(".", "credentials.json") -} - const defaultPublicKeyPEM = `-----BEGIN PUBLIC KEY----- MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAhxudxTewPgljUHEZHkusP7m3I+zA4/RGvuUMt6TtII/m4zwUOm/Y31zHBTmkCCt8k5vj9y+AmO0TsGmHooNQuMebakdmEWdcA5h7YAHHFbF2w5LcxIXjib08vgVpA+m3R5xPbLK+vfHe2aAX36b5nHReDNncY5vAl3U4CgIEBGPqyG67vJytRWqP+sfEdw5+m192Rf4SCGyiBzRmjiVlH3zeEBjdbOrkAnzKOVz6AHBl2q7LPLJKIzxjoAyhEp5qnDjHUFo5VZUgFwUOt83A/jbGMyzmjRoxBuvKcs9tBuorZyUwIsZN6E+rtQk2YqMPj4RkDsZ7LRmj6on8sN2rHQIDAQAB -----END PUBLIC KEY-----` diff --git a/internal/server/server.go b/internal/server/server.go index 182c26e..721d21c 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -21,6 +21,8 @@ import ( "git.misaka.ren/M1saka/zhanlu_proxy/internal/config" "git.misaka.ren/M1saka/zhanlu_proxy/internal/openai" "git.misaka.ren/M1saka/zhanlu_proxy/internal/sign" + "git.misaka.ren/M1saka/zhanlu_proxy/internal/stats" + "git.misaka.ren/M1saka/zhanlu_proxy/internal/store" "git.misaka.ren/M1saka/zhanlu_proxy/internal/zhanlu" ) @@ -28,10 +30,12 @@ type Server struct { cfg config.Config mux *http.ServeMux loginSession string + st *store.Store + statsEnabled bool } -func New(cfg config.Config) http.Handler { - s := &Server{cfg: cfg, mux: http.NewServeMux()} +func New(cfg config.Config, st *store.Store) http.Handler { + s := &Server{cfg: cfg, mux: http.NewServeMux(), st: st, statsEnabled: !cfg.StatsDisabled} if cfg.LoginPassword != "" { s.loginSession = randomSessionToken() } @@ -53,6 +57,9 @@ func (s *Server) routes() { s.mux.HandleFunc("GET /api/credentials", s.withLoginSession(s.getCredentials)) s.mux.HandleFunc("POST /api/credentials", s.withLoginSession(s.saveCredentials)) s.mux.HandleFunc("POST /api/sso/exchange", s.withLoginSession(s.exchangeSSOCode)) + s.mux.HandleFunc("GET /api/stats", s.withLoginSession(s.getStats)) + s.mux.HandleFunc("POST /api/stats/reset", s.withLoginSession(s.resetStats)) + s.mux.HandleFunc("GET /admin/stats", s.withLoginSession(s.statsPage)) s.mux.HandleFunc("GET /v1/models", s.withAPIKey(s.models)) s.mux.HandleFunc("POST /v1/chat/completions", s.withAPIKey(s.chatCompletions)) } @@ -92,7 +99,7 @@ func (s *Server) loginPage(w http.ResponseWriter, r *http.Request) { return } w.Header().Set("Content-Type", "text/html; charset=utf-8") - _ = loginTemplate.Execute(w, map[string]any{"CredentialsPath": s.cfg.CredentialsPath, "SSOBaseURL": s.cfg.SSOBaseURL, "PasswordEnabled": true, "AdminMode": false}) + _ = loginTemplate.Execute(w, map[string]any{"DBPath": s.cfg.DBPath, "SSOBaseURL": s.cfg.SSOBaseURL, "PasswordEnabled": true, "AdminMode": false}) } func (s *Server) adminLoginPage(w http.ResponseWriter, r *http.Request) { @@ -101,7 +108,7 @@ func (s *Server) adminLoginPage(w http.ResponseWriter, r *http.Request) { return } w.Header().Set("Content-Type", "text/html; charset=utf-8") - _ = loginTemplate.Execute(w, map[string]any{"CredentialsPath": s.cfg.CredentialsPath, "SSOBaseURL": s.cfg.SSOBaseURL, "PasswordEnabled": s.cfg.LoginPassword != "", "AdminMode": true}) + _ = loginTemplate.Execute(w, map[string]any{"DBPath": s.cfg.DBPath, "SSOBaseURL": s.cfg.SSOBaseURL, "PasswordEnabled": s.cfg.LoginPassword != "", "AdminMode": true}) } func (s *Server) passwordLogin(w http.ResponseWriter, r *http.Request) { @@ -159,7 +166,7 @@ func (s *Server) ssoCallback(w http.ResponseWriter, r *http.Request) { s.renderLoginResult(w, false, err.Error()) return } - if err := auth.SaveCredentials(s.cfg.CredentialsPath, creds); err != nil { + if err := s.st.SaveCredentials(creds); err != nil { s.renderLoginResult(w, false, err.Error()) return } @@ -304,12 +311,12 @@ func (s *Server) loginWithPhoneCode(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusBadGateway, map[string]any{"ok": false, "error": err.Error()}) return } - if err := auth.SaveCredentials(s.cfg.CredentialsPath, creds); err != nil { + if err := s.st.SaveCredentials(creds); err != nil { writeJSON(w, http.StatusBadRequest, map[string]any{"ok": false, "error": err.Error()}) return } s.cfg.Credentials = creds - writeJSON(w, http.StatusOK, map[string]any{"ok": true, "path": s.cfg.CredentialsPath, "access_key": mask(creds.AccessKey)}) + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "path": s.cfg.DBPath, "access_key": mask(creds.AccessKey)}) } // provisionCredentials logs the AK/SK/token into the Zhanlu gateway to obtain @@ -451,14 +458,14 @@ func randomRequestID() string { } func (s *Server) getCredentials(w http.ResponseWriter, r *http.Request) { - c, err := auth.LoadCredentials(s.cfg.CredentialsPath) - if err != nil { - writeJSON(w, http.StatusOK, map[string]any{"configured": false, "path": s.cfg.CredentialsPath}) + c, err := s.st.LoadCredentials() + if err != nil || (c.Validate() != nil && !c.HasAPIKey()) { + writeJSON(w, http.StatusOK, map[string]any{"configured": false, "path": s.cfg.DBPath}) return } writeJSON(w, http.StatusOK, map[string]any{ "configured": true, - "path": s.cfg.CredentialsPath, + "path": s.cfg.DBPath, "access_key": mask(c.AccessKey), "has_api_key": c.APIKey != "", "model_base": firstNonEmpty(c.ModelBaseURL, c.BaseURL), @@ -482,12 +489,12 @@ func (s *Server) saveCredentials(w http.ResponseWriter, r *http.Request) { } c = provisioned } - if err := auth.SaveCredentials(s.cfg.CredentialsPath, c); err != nil { + if err := s.st.SaveCredentials(c); err != nil { writeJSON(w, http.StatusBadRequest, map[string]any{"ok": false, "error": err.Error()}) return } s.cfg.Credentials = c - writeJSON(w, http.StatusOK, map[string]any{"ok": true, "path": s.cfg.CredentialsPath}) + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "path": s.cfg.DBPath}) } func (s *Server) exchangeSSOCode(w http.ResponseWriter, r *http.Request) { @@ -516,12 +523,83 @@ func (s *Server) exchangeSSOCode(w http.ResponseWriter, r *http.Request) { if in.BaseURL != "" { creds.ModelBaseURL = in.BaseURL } - if err := auth.SaveCredentials(s.cfg.CredentialsPath, creds); err != nil { + if err := s.st.SaveCredentials(creds); err != nil { writeJSON(w, http.StatusBadRequest, map[string]any{"ok": false, "error": err.Error()}) return } s.cfg.Credentials = creds - writeJSON(w, http.StatusOK, map[string]any{"ok": true, "path": s.cfg.CredentialsPath}) + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "path": s.cfg.DBPath}) +} + +func (s *Server) getStats(w http.ResponseWriter, r *http.Request) { + if s.st == nil { + writeJSON(w, http.StatusOK, map[string]any{"enabled": false}) + return + } + q := stats.Query{ + Model: r.URL.Query().Get("model"), + Limit: parseLimit(r.URL.Query().Get("limit")), + } + if v := r.URL.Query().Get("since"); v != "" { + if t, err := time.Parse(time.RFC3339, v); err == nil { + q.Since = t + } + } + if v := r.URL.Query().Get("until"); v != "" { + if t, err := time.Parse(time.RFC3339, v); err == nil { + q.Until = t + } + } + summary, err := s.st.Stats(q) + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]any{"ok": false, "error": err.Error()}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"enabled": s.statsEnabled, "stats": summary}) +} + +func (s *Server) resetStats(w http.ResponseWriter, r *http.Request) { + if s.st == nil { + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "enabled": false}) + return + } + if err := s.st.Reset(); err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]any{"ok": false, "error": err.Error()}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"ok": true}) +} + +func (s *Server) statsPage(w http.ResponseWriter, r *http.Request) { + var summary *stats.Summary + enabled := s.statsEnabled + if s.st != nil { + if sm, err := s.st.Stats(stats.Query{}); err == nil { + summary = sm + } + } + if summary == nil { + summary = &stats.Summary{} + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + _ = statsTemplate.Execute(w, map[string]any{ + "Enabled": enabled, + "Stats": summary, + }) +} + +func parseLimit(s string) int { + n := 0 + for _, c := range s { + if c < '0' || c > '9' { + return 0 + } + n = n*10 + int(c-'0') + if n > 5000 { + return 5000 + } + } + return n } func (s *Server) models(w http.ResponseWriter, r *http.Request) { @@ -554,6 +632,7 @@ func (s *Server) chatCompletions(w http.ResponseWriter, r *http.Request) { if req.Model == "" { req.Model = "zhanlu/auto" } + start := time.Now() clientWantsStream := req.Stream req.Stream = true body, err := req.MarshalForUpstream() @@ -569,7 +648,7 @@ func (s *Server) chatCompletions(w http.ResponseWriter, r *http.Request) { return } s.cfg.Credentials = creds - _ = auth.SaveCredentials(s.cfg.CredentialsPath, creds) + _ = s.st.SaveCredentials(creds) } modelBaseURL := firstNonEmpty(creds.ModelBaseURL, s.cfg.MobileModelBaseURL) client, err := s.zhanluClientWithBase(modelBaseURL) @@ -584,6 +663,7 @@ func (s *Server) chatCompletions(w http.ResponseWriter, r *http.Request) { msg = redactSensitive(err.Error()) } writeOpenAIError(w, http.StatusBadGateway, msg, "upstream_error", "zhanlu_request_failed") + s.record(req.Model, clientWantsStream, nil, "upstream_error", start) return } defer resp.Body.Close() @@ -594,33 +674,93 @@ func (s *Server) chatCompletions(w http.ResponseWriter, r *http.Request) { msg += ": " + string(b) } writeOpenAIError(w, http.StatusBadGateway, msg, "upstream_error", "zhanlu_bad_status") + s.record(req.Model, clientWantsStream, nil, "upstream_error", start) return } if clientWantsStream { - s.proxyStream(w, resp) + usage, status := s.proxyStream(w, resp) + s.record(req.Model, true, usage, status, start) return } - s.aggregateStream(w, resp, req.Model) + usage, status := s.aggregateStream(w, resp, req.Model) + s.record(req.Model, false, usage, status, start) } -func (s *Server) proxyStream(w http.ResponseWriter, resp *http.Response) { +// record appends a usage observation to the stats store when collection is +// enabled. It never affects the response path; recording errors are ignored. +func (s *Server) record(model string, stream bool, usage any, status string, start time.Time) { + if !s.statsEnabled || s.st == nil { + return + } + _ = s.st.Record(stats.RecordFromUsage(model, stream, usage, status, start)) +} + +func (s *Server) proxyStream(w http.ResponseWriter, resp *http.Response) (any, string) { w.Header().Set("Content-Type", "text/event-stream; charset=utf-8") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") w.Header().Set("X-Accel-Buffering", "no") w.WriteHeader(http.StatusOK) flusher, _ := w.(http.Flusher) - _, err := io.Copy(w, resp.Body) - if flusher != nil { - flusher.Flush() - } - if err != nil { - b, _ := json.Marshal(map[string]any{"error": map[string]any{"message": err.Error(), "type": "upstream_error", "code": "zhanlu_stream_error"}}) - _, _ = fmt.Fprintf(w, "data: %s\n\n", b) + reader := bufio.NewReader(resp.Body) + var usage any + status := "success" + for { + line, err := reader.ReadString('\n') + if line != "" { + if _, werr := io.WriteString(w, line); werr != nil { + // client disconnected mid-stream; stop forwarding + status = "upstream_error" + break + } + if flusher != nil { + flusher.Flush() + } + if payload := sseDataPayload(line); payload != "" && payload != "[DONE]" { + var evt struct { + State string `json:"state"` + ErrorMessage string `json:"errorMessage"` + Usage any `json:"usage"` + } + if json.Unmarshal([]byte(payload), &evt) == nil { + if evt.State == "ERROR" { + status = "upstream_error" + } + if evt.Usage != nil { + usage = evt.Usage + } + } + } + } + if err != nil { + if err == io.EOF { + break + } + // upstream read error: surface an error event to the client, mirroring + // the previous io.Copy behavior, then mark the request as failed. + b, _ := json.Marshal(map[string]any{"error": map[string]any{"message": err.Error(), "type": "upstream_error", "code": "zhanlu_stream_error"}}) + _, _ = fmt.Fprintf(w, "data: %s\n\n", b) + if flusher != nil { + flusher.Flush() + } + status = "upstream_error" + break + } } + return usage, status } -func (s *Server) aggregateStream(w http.ResponseWriter, resp *http.Response, model string) { +// sseDataPayload returns the payload following a "data:" SSE line, or "" if the +// line is not a data line. Mirrors the parsing in forEachSSEChunk. +func sseDataPayload(line string) string { + trimmed := strings.TrimSpace(line) + if !strings.HasPrefix(trimmed, "data:") { + return "" + } + return strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) +} + +func (s *Server) aggregateStream(w http.ResponseWriter, resp *http.Response, model string) (any, string) { var content, reasoning, id string var usage any finishReason := "stop" @@ -690,7 +830,7 @@ func (s *Server) aggregateStream(w http.ResponseWriter, resp *http.Response, mod }) if err != nil { writeOpenAIError(w, http.StatusBadGateway, err.Error(), "upstream_error", "zhanlu_stream_error") - return + return nil, "upstream_error" } if id == "" { id = "chatcmpl-" + randomRequestID() @@ -719,6 +859,7 @@ func (s *Server) aggregateStream(w http.ResponseWriter, resp *http.Response, mod result["usage"] = usage } writeJSON(w, http.StatusOK, result) + return usage, "success" } // forEachSSEChunk feeds each non-empty data: payload to fn, skipping keep-alive @@ -753,7 +894,7 @@ func (s *Server) currentCredentials() (auth.Credentials, error) { if s.cfg.Credentials.Validate() == nil || s.cfg.Credentials.HasAPIKey() { return s.cfg.Credentials, nil } - c, err := auth.LoadCredentials(s.cfg.CredentialsPath) + c, err := s.st.LoadCredentials() if err != nil { return auth.Credentials{}, err } @@ -1037,8 +1178,8 @@ var loginTemplate = template.Must(template.New("login").Parse(`

输入手机号获取验证码，按插件默认的移动云登录接口换取凭据和模型 API Key。服务会保存凭据，后续 OpenAI 兼容接口自动使用。

- 凭据保存位置 - {{.CredentialsPath}} + 数据库位置 + {{.DBPath}}
@@ -1179,7 +1320,7 @@ var loginTemplate = template.Must(template.New("login").Parse(` setStatus(data.error || '登录失败', 'err'); return; } - setStatus('登录成功，已保存凭据：' + (data.access_key || '') + '；JSON：' + (data.path || ''), 'ok'); + setStatus('登录成功，已保存凭据：' + (data.access_key || '') + '；数据库：' + (data.path || ''), 'ok'); }); @@ -1218,3 +1359,134 @@ p{margin:14px 0 26px;color:var(--body);line-height:1.7;font-size:14.5px;word-bre

{{.Message}}

返回登录页 `)) + +var statsTemplate = template.Must(template.New("stats").Funcs(template.FuncMap{ + "pct": func(f float64) string { return fmt.Sprintf("%.1f%%", f*100) }, + "rate": func(cached, prompt int64) string { + if prompt <= 0 { + return "0%" + } + return fmt.Sprintf("%.1f%%", float64(cached)/float64(prompt)*100) + }, +}).Parse(` + + + + + + 湛卢 Token 统计 + + + +
+
+
+
Zhanlu Proxy · Token 统计
+

Token 消耗统计

+
+
+ 返回登录管理 + + +
+
+ {{if not .Enabled}}

统计已关闭（ZHANLU_STATS_DISABLED=true）。

{{end}} +
+

总览

+
+
请求总数
{{.Stats.Totals.Requests}}
成功 {{.Stats.Totals.SuccessRequests}} · 失败 {{.Stats.Totals.ErrorRequests}}
+
Prompt Tokens
{{.Stats.Totals.PromptTokens}}
缓存 {{.Stats.Totals.CachedTokens}}
+
Completion Tokens
{{.Stats.Totals.CompletionTokens}}
含思考 {{.Stats.Totals.ReasoningTokens}}
+
Total Tokens
{{.Stats.Totals.TotalTokens}}
+
缓存命中率
{{pct .Stats.Totals.CacheRate}}
缓存 {{.Stats.Totals.CachedTokens}} / Prompt {{.Stats.Totals.PromptTokens}}
+
+
+
+

按模型

+ {{if .Stats.PerModel}} +
+ + + {{range .Stats.PerModel}}{{end}} + +
模型请求数PromptCompTotal缓存命中率
{{.Model}}{{.Requests}}{{.PromptTokens}}{{.CompletionTokens}}{{.TotalTokens}}{{.CachedTokens}}{{rate .CachedTokens .PromptTokens}}
+ {{else}}

暂无数据

{{end}} +
+
+

按日

+ {{if .Stats.Daily}} +
+ {{range .Stats.Daily}}
{{.Day}}
{{.TotalTokens}}
{{end}} +
+ {{else}}

暂无数据

{{end}} +
+
+

最近请求

+ {{if .Stats.Recent}} +
+ + + {{range .Stats.Recent}}{{end}} + +
时间模型模式PromptCompTotal缓存状态耗时
{{.Ts.Format "01-02 15:04:05"}}{{.Model}}{{if .Stream}}流式{{else}}非流式{{end}}{{.PromptTokens}}{{.CompletionTokens}}{{.TotalTokens}}{{.CachedTokens}}{{if eq .Status "success"}}成功{{else}}失败{{end}}{{.LatencyMs}}ms
+ {{else}}

暂无数据

{{end}} +
+
+ + +`)) diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 6c990e2..516ac27 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -15,12 +15,15 @@ import ( "git.misaka.ren/M1saka/zhanlu_proxy/internal/auth" "git.misaka.ren/M1saka/zhanlu_proxy/internal/config" + "git.misaka.ren/M1saka/zhanlu_proxy/internal/store" ) const testSM2Key = "8d6ee90b3c4d299ae5abd655dbc3547c110ae8aeff1de18b0df241f215f90748" -// setupTestServer spins up a mock Zhanlu upstream and a proxy server wired to it. -func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, string) { +// setupTestServer spins up a mock Zhanlu upstream and a proxy server wired to +// it. The proxy is backed by a temp SQLite store so credentials persist in the +// db during the test instead of a JSON file. +func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, *store.Store) { t.Helper() upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { @@ -47,7 +50,7 @@ func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, string) return } w.Header().Set("Content-Type", "text/event-stream") - _, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}],\"usage\":{\"prompt_tokens\":1}}\n\n") + _, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":20,\"total_tokens\":30,\"prompt_tokens_details\":{\"cached_tokens\":4}}}\n\n") _, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") _, _ = fmt.Fprint(w, "data: [DONE]\n\n") case "/gateway/v1/model/info": @@ -59,28 +62,32 @@ func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, string) } })) - credsFile := filepath.Join(t.TempDir(), "credentials.json") + st, err := store.Open(filepath.Join(t.TempDir(), "stats.db")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { st.Close() }) cfg := config.Config{ ListenAddr: ":0", MobileLoginBaseURL: upstream.URL, MobileModelBaseURL: upstream.URL, UpstreamPath: "/chat/completions", - CredentialsPath: credsFile, + DBPath: filepath.Join(t.TempDir(), "zhanlu.db"), TokenDecryptKey: "3jw7woww2rvhla6k", PublicKeyPEM: defaultTestPublicKey, PhonePublicKeyPEM: defaultTestPublicKey, SM2PrivateKey: testSM2Key, PluginVersion: "1.4.2", } - h := New(cfg) + h := New(cfg, st) proxy := httptest.NewServer(h) - return upstream, proxy, credsFile + return upstream, proxy, st } // TestPhoneLoginAndChat exercises the full v1.4.2 flow: SMS login, profile // fetch, SM2 API-key provisioning, then OpenAI-compatible chat and models. func TestPhoneLoginAndChat(t *testing.T) { - upstream, proxy, credsFile := setupTestServer(t) + upstream, proxy, st := setupTestServer(t) defer upstream.Close() defer proxy.Close() @@ -107,8 +114,8 @@ func TestPhoneLoginAndChat(t *testing.T) { t.Fatalf("login failed: %v", loginResp) } - // 3. credentials file should contain the provisioned api key - creds, err := auth.LoadCredentials(credsFile) + // 3. credentials store should contain the provisioned api key + creds, err := st.LoadCredentials() if err != nil { t.Fatal(err) } @@ -154,11 +161,11 @@ func TestPhoneLoginAndChat(t *testing.T) { } func TestStreamingChat(t *testing.T) { - upstream, proxy, credsFile := setupTestServer(t) + upstream, proxy, st := setupTestServer(t) defer upstream.Close() defer proxy.Close() - // Seed credentials directly with the api key + // Seed credentials directly with the api key into the store creds := auth.Credentials{ AccessKey: "AK", SecretKey: "SK", @@ -167,7 +174,7 @@ func TestStreamingChat(t *testing.T) { ModelBaseURL: upstream.URL, Email: "dev@example.com", } - if err := auth.SaveCredentials(credsFile, creds); err != nil { + if err := st.SaveCredentials(creds); err != nil { t.Fatal(err) } @@ -188,6 +195,88 @@ func TestStreamingChat(t *testing.T) { } } +// TestStats verifies both streaming and non-streaming chat paths record token +// usage into the store and that GET /api/stats aggregates them correctly. +func TestStats(t *testing.T) { + upstream, proxy, st := setupTestServer(t) + defer upstream.Close() + defer proxy.Close() + + creds := auth.Credentials{ + AccessKey: "AK", SecretKey: "SK", Token: "TOKEN", + APIKey: "sk-test-456", ModelBaseURL: upstream.URL, Email: "dev@example.com", + } + if err := st.SaveCredentials(creds); err != nil { + t.Fatal(err) + } + + // non-streaming chat + resp, err := http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":false}`)) + if err != nil { + t.Fatal(err) + } + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + + // streaming chat + resp, err = http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":true}`)) + if err != nil { + t.Fatal(err) + } + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + + // query stats + resp, err = http.Get(proxy.URL + "/api/stats") + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + var statsResp map[string]any + if err := json.NewDecoder(resp.Body).Decode(&statsResp); err != nil { + t.Fatal(err) + } + if enabled, _ := statsResp["enabled"].(bool); !enabled { + t.Fatalf("stats not enabled: %v", statsResp) + } + s, _ := statsResp["stats"].(map[string]any) + if s == nil { + t.Fatalf("no stats object: %v", statsResp) + } + totals, _ := s["totals"].(map[string]any) + if totals == nil { + t.Fatalf("no totals: %v", s) + } + if got := totNum(totals["requests"]); got != 2 { + t.Fatalf("requests = %v, want 2", totals["requests"]) + } + if got := totNum(totals["prompt_tokens"]); got != 20 { + t.Fatalf("prompt_tokens = %v, want 20", totals["prompt_tokens"]) + } + if got := totNum(totals["completion_tokens"]); got != 40 { + t.Fatalf("completion_tokens = %v, want 40", totals["completion_tokens"]) + } + if got := totNum(totals["total_tokens"]); got != 60 { + t.Fatalf("total_tokens = %v, want 60", totals["total_tokens"]) + } + if got := totNum(totals["cached_tokens"]); got != 8 { + t.Fatalf("cached_tokens = %v, want 8", totals["cached_tokens"]) + } + if rate, _ := totals["cache_rate"].(float64); rate < 0.39 || rate > 0.41 { + t.Fatalf("cache_rate = %v, want ~0.4", totals["cache_rate"]) + } + perModel, _ := s["per_model"].([]any) + if len(perModel) != 1 { + t.Fatalf("per_model = %v, want 1 entry", perModel) + } +} + +// totNum extracts an int from a JSON-decoded numeric value (float64). +func totNum(v any) int { + f, _ := v.(float64) + return int(f) +} + func okValue(m map[string]any) bool { ok, _ := m["ok"].(bool) return ok diff --git a/internal/stats/stats.go b/internal/stats/stats.go new file mode 100644 index 0000000..540898b --- /dev/null +++ b/internal/stats/stats.go @@ -0,0 +1,83 @@ +// Package stats defines the types and helpers for OpenAI-compatible token +// usage statistics recorded by the zhanlu proxy. The concrete SQLite-backed +// recorder lives in internal/store; this package is dependency-free so it can +// be referenced by both store and server without import cycles. +package stats + +import "time" + +// Record is a single chat-completion usage observation. +type Record struct { + Ts time.Time `json:"ts"` + Model string `json:"model"` + Stream bool `json:"stream"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + ReasoningTokens int `json:"reasoning_tokens"` + CachedTokens int `json:"cached_tokens"` + Status string `json:"status"` // "success" | "upstream_error" + LatencyMs int64 `json:"latency_ms"` +} + +// Query filters the recorded stats. Zero-value time fields mean unbounded on +// that end; empty Model means all models. Limit caps the recent-records list +// (0 = default). Stream filters by streaming mode (nil = both). +type Query struct { + Since time.Time + Until time.Time + Model string + Limit int + Stream *bool +} + +// Totals aggregates request counts and token sums over a filtered set. +type Totals struct { + Requests int `json:"requests"` + SuccessRequests int `json:"success_requests"` + ErrorRequests int `json:"error_requests"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + TotalTokens int64 `json:"total_tokens"` + ReasoningTokens int64 `json:"reasoning_tokens"` + CachedTokens int64 `json:"cached_tokens"` + CacheRate float64 `json:"cache_rate"` // cached_tokens / prompt_tokens, 0..1 +} + +// ModelStat is a per-model aggregation row. +type ModelStat struct { + Model string `json:"model"` + Requests int `json:"requests"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + TotalTokens int64 `json:"total_tokens"` + CachedTokens int64 `json:"cached_tokens"` +} + +// DayStat is a per-day aggregation row (server-local time, YYYY-MM-DD). +type DayStat struct { + Day string `json:"day"` + Requests int `json:"requests"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + TotalTokens int64 `json:"total_tokens"` + CachedTokens int64 `json:"cached_tokens"` +} + +// Summary is the full result returned by a Recorder's Stats query. +type Summary struct { + Totals Totals `json:"totals"` + PerModel []ModelStat `json:"per_model"` + Daily []DayStat `json:"daily"` + Recent []Record `json:"recent"` +} + +// Recorder persists and queries usage statistics. The concrete implementation +// lives in internal/store; the interface is declared here so server code can +// depend on the contract and tests can inject fakes. +type Recorder interface { + Record(r Record) error + Stats(q Query) (*Summary, error) + Reset() error + Close() error +} diff --git a/internal/stats/usage.go b/internal/stats/usage.go new file mode 100644 index 0000000..d3e5cbd --- /dev/null +++ b/internal/stats/usage.go @@ -0,0 +1,78 @@ +package stats + +import ( + "encoding/json" + "time" +) + +// Usage is the subset of the OpenAI chat-completion usage object the recorder +// persists. Numbers arrive from JSON unmarshal as float64. +type Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + ReasoningTokens int `json:"reasoning_tokens"` + CachedTokens int `json:"cached_tokens"` + + // PromptTokensDetails.CachedTokens is emitted by providers that support + // prompt caching (OpenAI/DeepSeek/Zhipu litellm gateways). Some upstreams + // put cached_tokens at the top level instead. + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + + // CompletionTokensDetails.ReasoningTokens is emitted by reasoning models; + // some upstreams put reasoning_tokens at top level instead. + CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` + } `json:"completion_tokens_details"` +} + +// ExtractUsage decodes a raw usage value (as produced by encoding/json into an +// any) into token counts. It accepts both full usage maps and raw JSON bytes. +// Missing fields default to 0; a nil v yields zero usage. +func ExtractUsage(v any) Usage { + var u Usage + if v == nil { + return u + } + switch t := v.(type) { + case []byte: + _ = json.Unmarshal(t, &u) + case json.RawMessage: + _ = json.Unmarshal(t, &u) + case map[string]any: + // Re-marshal + unmarshal is the simplest robust path for nested + // *_tokens_details; usage payloads are tiny. + if b, err := json.Marshal(t); err == nil { + _ = json.Unmarshal(b, &u) + } + } + if u.ReasoningTokens == 0 { + u.ReasoningTokens = u.CompletionTokensDetails.ReasoningTokens + } + if u.CachedTokens == 0 { + u.CachedTokens = u.PromptTokensDetails.CachedTokens + } + return u +} + +// RecordFromUsage builds a Record from a captured usage value plus context. +func RecordFromUsage(model string, stream bool, usage any, status string, start time.Time) Record { + u := ExtractUsage(usage) + if u.TotalTokens == 0 && (u.PromptTokens != 0 || u.CompletionTokens != 0) { + u.TotalTokens = u.PromptTokens + u.CompletionTokens + } + return Record{ + Ts: time.Now(), + Model: model, + Stream: stream, + PromptTokens: u.PromptTokens, + CompletionTokens: u.CompletionTokens, + TotalTokens: u.TotalTokens, + ReasoningTokens: u.ReasoningTokens, + CachedTokens: u.CachedTokens, + Status: status, + LatencyMs: time.Since(start).Milliseconds(), + } +} diff --git a/internal/store/credentials.go b/internal/store/credentials.go new file mode 100644 index 0000000..47a70ee --- /dev/null +++ b/internal/store/credentials.go @@ -0,0 +1,45 @@ +package store + +import ( + "fmt" + "time" + + "git.misaka.ren/M1saka/zhanlu_proxy/internal/auth" +) + +// LoadCredentials reads the persisted credentials. The credentials row always +// exists after Open; an empty row (nothing saved yet) yields a zero-value +// Credentials with a nil error — callers check Validate()/HasAPIKey(). +func (s *Store) LoadCredentials() (auth.Credentials, error) { + var c auth.Credentials + var savedAt int64 + err := s.db.QueryRow(`SELECT access_key, secret_key, token, api_key, model_base_url, email, organization, team, base_url, saved_at FROM credentials WHERE id = 1`). + Scan(&c.AccessKey, &c.SecretKey, &c.Token, &c.APIKey, &c.ModelBaseURL, &c.Email, &c.Organization, &c.Team, &c.BaseURL, &savedAt) + if err != nil { + return auth.Credentials{}, fmt.Errorf("load credentials: %w", err) + } + c.SavedAt = time.Unix(savedAt, 0).Local() + return c, nil +} + +// SaveCredentials upserts the credentials into the single row. It only writes +// when the credentials validate or already carry an API key, so partial / +// env-only creds are not persisted. +func (s *Store) SaveCredentials(c auth.Credentials) error { + if c.Validate() != nil && !c.HasAPIKey() { + return fmt.Errorf("save credentials: %w", c.Validate()) + } + c.SavedAt = time.Now() + _, err := s.db.Exec(`INSERT INTO credentials (id, access_key, secret_key, token, api_key, model_base_url, email, organization, team, base_url, saved_at) +VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(id) DO UPDATE SET + access_key=excluded.access_key, secret_key=excluded.secret_key, token=excluded.token, + api_key=excluded.api_key, model_base_url=excluded.model_base_url, email=excluded.email, + organization=excluded.organization, team=excluded.team, base_url=excluded.base_url, + saved_at=excluded.saved_at`, + c.AccessKey, c.SecretKey, c.Token, c.APIKey, c.ModelBaseURL, c.Email, c.Organization, c.Team, c.BaseURL, c.SavedAt.Unix()) + if err != nil { + return fmt.Errorf("save credentials: %w", err) + } + return nil +} diff --git a/internal/store/stats_recorder.go b/internal/store/stats_recorder.go new file mode 100644 index 0000000..faf4087 --- /dev/null +++ b/internal/store/stats_recorder.go @@ -0,0 +1,187 @@ +package store + +import ( + "fmt" + "time" + + "git.misaka.ren/M1saka/zhanlu_proxy/internal/stats" +) + +// Compile-time guard: *Store satisfies stats.Recorder. +var _ stats.Recorder = (*Store)(nil) + +const defaultRecentLimit = 200 + +// Record appends a single usage observation. +func (s *Store) Record(rec stats.Record) error { + if rec.Status == "" { + rec.Status = "success" + } + stream := 0 + if rec.Stream { + stream = 1 + } + _, err := s.db.Exec(`INSERT INTO requests (ts, model, stream, prompt_tokens, completion_tokens, total_tokens, reasoning_tokens, cached_tokens, status, latency_ms) VALUES (?,?,?,?,?,?,?,?,?,?)`, + rec.Ts.Unix(), rec.Model, stream, rec.PromptTokens, rec.CompletionTokens, rec.TotalTokens, rec.ReasoningTokens, rec.CachedTokens, rec.Status, rec.LatencyMs) + return err +} + +// Stats computes the aggregate summary for the given query. +func (s *Store) Stats(q stats.Query) (*stats.Summary, error) { + q = normalizeQuery(q) + where, args := whereClause(q) + + sum := &stats.Summary{} + if err := s.scanTotals(sum, where, args); err != nil { + return nil, err + } + if sum.Totals.PromptTokens > 0 { + sum.Totals.CacheRate = float64(sum.Totals.CachedTokens) / float64(sum.Totals.PromptTokens) + } + if err := s.scanPerModel(sum, where, args); err != nil { + return nil, err + } + if err := s.scanDaily(sum, where, args); err != nil { + return nil, err + } + if err := s.scanRecent(sum, q, where, args); err != nil { + return nil, err + } + return sum, nil +} + +// Reset deletes all recorded usage statistics (the credentials row is kept). +func (s *Store) Reset() error { + _, err := s.db.Exec(`DELETE FROM requests`) + return err +} + +func normalizeQuery(q stats.Query) stats.Query { + if q.Limit <= 0 { + q.Limit = defaultRecentLimit + } + if q.Limit > 5000 { + q.Limit = 5000 + } + return q +} + +func whereClause(q stats.Query) (string, []any) { + var conds []string + var args []any + if !q.Since.IsZero() { + conds = append(conds, "ts >= ?") + args = append(args, q.Since.Unix()) + } + if !q.Until.IsZero() { + conds = append(conds, "ts <= ?") + args = append(args, q.Until.Unix()) + } + if q.Model != "" { + conds = append(conds, "model = ?") + args = append(args, q.Model) + } + if q.Stream != nil { + conds = append(conds, "stream = ?") + args = append(args, boolToInt(*q.Stream)) + } + if len(conds) == 0 { + return "", args + } + out := "" + for i, p := range conds { + if i > 0 { + out += " AND " + } + out += p + } + return " WHERE " + out, args +} + +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} + +func (s *Store) scanTotals(sum *stats.Summary, where string, args []any) error { + q := `SELECT COUNT(*), + COALESCE(SUM(CASE WHEN status='success' THEN 1 ELSE 0 END),0), + COALESCE(SUM(CASE WHEN status!='success' THEN 1 ELSE 0 END),0), + COALESCE(SUM(prompt_tokens),0), + COALESCE(SUM(completion_tokens),0), + COALESCE(SUM(total_tokens),0), + COALESCE(SUM(reasoning_tokens),0), + COALESCE(SUM(cached_tokens),0) FROM requests` + where + row := s.db.QueryRow(q, args...) + var success, failures int64 + err := row.Scan(&sum.Totals.Requests, &success, &failures, + &sum.Totals.PromptTokens, &sum.Totals.CompletionTokens, + &sum.Totals.TotalTokens, &sum.Totals.ReasoningTokens, &sum.Totals.CachedTokens) + if err != nil { + return fmt.Errorf("scan totals: %w", err) + } + sum.Totals.SuccessRequests = int(success) + sum.Totals.ErrorRequests = int(failures) + return nil +} + +func (s *Store) scanPerModel(sum *stats.Summary, where string, args []any) error { + q := `SELECT model, COUNT(*), COALESCE(SUM(prompt_tokens),0), COALESCE(SUM(completion_tokens),0), COALESCE(SUM(total_tokens),0), COALESCE(SUM(cached_tokens),0) FROM requests` + where + " GROUP BY model ORDER BY SUM(total_tokens) DESC" + rows, err := s.db.Query(q, args...) + if err != nil { + return fmt.Errorf("scan per-model: %w", err) + } + defer rows.Close() + for rows.Next() { + var m stats.ModelStat + if err := rows.Scan(&m.Model, &m.Requests, &m.PromptTokens, &m.CompletionTokens, &m.TotalTokens, &m.CachedTokens); err != nil { + return err + } + sum.PerModel = append(sum.PerModel, m) + } + return rows.Err() +} + +func (s *Store) scanDaily(sum *stats.Summary, where string, args []any) error { + q := `SELECT date(ts,'unixepoch','localtime') AS day, COUNT(*), COALESCE(SUM(prompt_tokens),0), COALESCE(SUM(completion_tokens),0), COALESCE(SUM(total_tokens),0), COALESCE(SUM(cached_tokens),0) FROM requests` + where + " GROUP BY day ORDER BY day ASC" + rows, err := s.db.Query(q, args...) + if err != nil { + return fmt.Errorf("scan daily: %w", err) + } + defer rows.Close() + for rows.Next() { + var d stats.DayStat + if err := rows.Scan(&d.Day, &d.Requests, &d.PromptTokens, &d.CompletionTokens, &d.TotalTokens, &d.CachedTokens); err != nil { + return err + } + sum.Daily = append(sum.Daily, d) + } + return rows.Err() +} + +func (s *Store) scanRecent(sum *stats.Summary, q stats.Query, where string, args []any) error { + limit := q.Limit + if limit <= 0 { + limit = defaultRecentLimit + } + query := `SELECT ts, model, stream, prompt_tokens, completion_tokens, total_tokens, reasoning_tokens, cached_tokens, status, latency_ms FROM requests` + where + " ORDER BY id DESC LIMIT ?" + rows, err := s.db.Query(query, append(args, limit)...) + if err != nil { + return fmt.Errorf("scan recent: %w", err) + } + defer rows.Close() + for rows.Next() { + var rec stats.Record + var ts int64 + var stream int + if err := rows.Scan(&ts, &rec.Model, &stream, &rec.PromptTokens, &rec.CompletionTokens, &rec.TotalTokens, &rec.ReasoningTokens, &rec.CachedTokens, &rec.Status, &rec.LatencyMs); err != nil { + return err + } + rec.Ts = time.Unix(ts, 0).Local() + rec.Stream = stream == 1 + sum.Recent = append(sum.Recent, rec) + } + return rows.Err() +} diff --git a/internal/store/store.go b/internal/store/store.go new file mode 100644 index 0000000..12c0bf7 --- /dev/null +++ b/internal/store/store.go @@ -0,0 +1,89 @@ +// Package store owns the single embedded SQLite database backing the zhanlu +// proxy: token-usage records (the requests table) and the persisted login +// credentials (the credentials table, single-tenant single row). It implements +// stats.Recorder for usage tracking and exposes Load/Save for credentials. +package store + +import ( + "database/sql" + "fmt" + "strings" + + _ "modernc.org/sqlite" +) + +// Store is the single owner of the proxy's SQLite database handle. +type Store struct { + db *sql.DB +} + +// Open opens (or creates) the database at path and ensures both tables exist. +// SQLite is opened with WAL journaling and a busy timeout so concurrent reads +// (stats queries) and writes (request records, credential saves) do not +// collide. +func Open(path string) (*Store, error) { + db, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)&_pragma=foreign_keys(ON)") + if err != nil { + return nil, fmt.Errorf("open db %q: %w", path, err) + } + if err := ensureSchema(db); err != nil { + db.Close() + return nil, err + } + return &Store{db: db}, nil +} + +// Close releases the database handle. +func (s *Store) Close() error { + return s.db.Close() +} + +func ensureSchema(db *sql.DB) error { + stmts := []string{ + `CREATE TABLE IF NOT EXISTS requests ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + ts INTEGER NOT NULL, + model TEXT NOT NULL, + stream INTEGER NOT NULL DEFAULT 0, + prompt_tokens INTEGER NOT NULL DEFAULT 0, + completion_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + reasoning_tokens INTEGER NOT NULL DEFAULT 0, + cached_tokens INTEGER NOT NULL DEFAULT 0, + status TEXT NOT NULL DEFAULT 'success', + latency_ms INTEGER NOT NULL DEFAULT 0 +)`, + `CREATE TABLE IF NOT EXISTS credentials ( + id INTEGER PRIMARY KEY CHECK (id = 1), + access_key TEXT NOT NULL DEFAULT '', + secret_key TEXT NOT NULL DEFAULT '', + token TEXT NOT NULL DEFAULT '', + api_key TEXT NOT NULL DEFAULT '', + model_base_url TEXT NOT NULL DEFAULT '', + email TEXT NOT NULL DEFAULT '', + organization TEXT NOT NULL DEFAULT '', + team TEXT NOT NULL DEFAULT '', + base_url TEXT NOT NULL DEFAULT '', + saved_at INTEGER NOT NULL DEFAULT 0 +)`, + `CREATE INDEX IF NOT EXISTS idx_requests_ts ON requests(ts)`, + `CREATE INDEX IF NOT EXISTS idx_requests_model ON requests(model)`, + // Ensure the single credentials row exists so UPSERTs and SELECTs always + // have a target. + `INSERT INTO credentials (id) VALUES (1) ON CONFLICT(id) DO NOTHING`, + } + for _, q := range stmts { + if _, err := db.Exec(q); err != nil { + return fmt.Errorf("schema: %w", err) + } + } + // Add cached_tokens to databases created before this column existed. SQLite + // returns "duplicate column name" when it already exists; that is expected + // and ignored. + if _, err := db.Exec(`ALTER TABLE requests ADD COLUMN cached_tokens INTEGER NOT NULL DEFAULT 0`); err != nil { + if !strings.Contains(err.Error(), "duplicate column") { + return fmt.Errorf("migrate cached_tokens: %w", err) + } + } + return nil +} diff --git a/internal/store/store_test.go b/internal/store/store_test.go new file mode 100644 index 0000000..744fa27 --- /dev/null +++ b/internal/store/store_test.go @@ -0,0 +1,120 @@ +package store + +import ( + "path/filepath" + "testing" + "time" + + "git.misaka.ren/M1saka/zhanlu_proxy/internal/auth" + "git.misaka.ren/M1saka/zhanlu_proxy/internal/stats" +) + +func newTestStore(t *testing.T) *Store { + t.Helper() + st, err := Open(filepath.Join(t.TempDir(), "test.db")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { st.Close() }) + return st +} + +func TestCredentialsRoundTrip(t *testing.T) { + st := newTestStore(t) + in := auth.Credentials{ + AccessKey: "AK", SecretKey: "SK", Token: "TOK", + APIKey: "sk-1", ModelBaseURL: "https://up.example", Email: "u@example.com", + Organization: "org", Team: "team", + } + if err := st.SaveCredentials(in); err != nil { + t.Fatalf("save: %v", err) + } + got, err := st.LoadCredentials() + if err != nil { + t.Fatalf("load: %v", err) + } + if got.APIKey != "sk-1" || got.Email != "u@example.com" || got.Organization != "org" { + t.Fatalf("round-trip mismatch: %+v", got) + } + if got.SavedAt.IsZero() { + t.Fatalf("saved_at not set") + } +} + +func TestEmptyLoadReturnsZero(t *testing.T) { + st := newTestStore(t) + got, err := st.LoadCredentials() + if err != nil { + t.Fatalf("load on empty store: %v", err) + } + if got.HasAPIKey() { + t.Fatalf("expected no api key on empty store, got %+v", got) + } +} + +func TestRecordStatsReset(t *testing.T) { + st := newTestStore(t) + now := time.Now() + recs := []stats.Record{ + {Ts: now, Model: "glm-4.7", Stream: false, PromptTokens: 10, CompletionTokens: 20, TotalTokens: 30, CachedTokens: 4, Status: "success", LatencyMs: 5}, + {Ts: now, Model: "glm-4.7", Stream: true, PromptTokens: 5, CompletionTokens: 5, TotalTokens: 10, Status: "success", LatencyMs: 8}, + {Ts: now, Model: "minimax", Stream: false, PromptTokens: 1, CompletionTokens: 1, TotalTokens: 2, Status: "upstream_error", LatencyMs: 3}, + } + for _, r := range recs { + if err := st.Record(r); err != nil { + t.Fatalf("record: %v", err) + } + } + + sum, err := st.Stats(stats.Query{}) + if err != nil { + t.Fatalf("stats: %v", err) + } + if sum.Totals.Requests != 3 { + t.Fatalf("requests = %d, want 3", sum.Totals.Requests) + } + if sum.Totals.SuccessRequests != 2 || sum.Totals.ErrorRequests != 1 { + t.Fatalf("success/error = %d/%d, want 2/1", sum.Totals.SuccessRequests, sum.Totals.ErrorRequests) + } + if sum.Totals.TotalTokens != 42 { + t.Fatalf("total tokens = %d, want 42", sum.Totals.TotalTokens) + } + if sum.Totals.CachedTokens != 4 { + t.Fatalf("cached tokens = %d, want 4", sum.Totals.CachedTokens) + } + // prompt total = 10+5+1 = 16, cached = 4 => 0.25 + if sum.Totals.CacheRate < 0.24 || sum.Totals.CacheRate > 0.26 { + t.Fatalf("cache rate = %v, want ~0.25", sum.Totals.CacheRate) + } + if len(sum.PerModel) != 2 { + t.Fatalf("per-model entries = %d, want 2", len(sum.PerModel)) + } + // glm-4.7 should lead on total tokens (40 vs 2) + if sum.PerModel[0].Model != "glm-4.7" || sum.PerModel[0].TotalTokens != 40 || sum.PerModel[0].CachedTokens != 4 { + t.Fatalf("top model = %+v, want glm-4.7/40/4 cached", sum.PerModel[0]) + } + if len(sum.Recent) != 3 { + t.Fatalf("recent entries = %d, want 3", len(sum.Recent)) + } + // most recent first (id desc) => minimax record + if sum.Recent[0].Model != "minimax" { + t.Fatalf("most recent = %+v, want minimax", sum.Recent[0]) + } + + // model filter + sumF, _ := st.Stats(stats.Query{Model: "minimax"}) + if sumF.Totals.Requests != 1 || sumF.Totals.TotalTokens != 2 { + t.Fatalf("filtered stats = %+v, want 1/2", sumF.Totals) + } + + if err := st.Reset(); err != nil { + t.Fatalf("reset: %v", err) + } + sum2, _ := st.Stats(stats.Query{}) + if sum2.Totals.Requests != 0 { + t.Fatalf("after reset requests = %d, want 0", sum2.Totals.Requests) + } + // credentials must survive a stats reset + creds, _ := st.LoadCredentials() + _ = creds +}