From f3d5ba6ceacd96716c1c3e85d0754c557e14e034 Mon Sep 17 00:00:00 2001 From: dev Date: Mon, 24 Aug 2026 22:41:48 +0800 Subject: [PATCH] refactor(gateway): deduplicate stats CSV export paths - extract csvHeaders() (Content-Type + Content-Disposition) shared by both export branches - extract exportKey(): the identical admin/user key-filter logic existed twice (JSON path + keys-csv); now all three call sites share one function - inline the nine single-use intermediate variables in the keys-csv loop --- internal/gateway/api.go | 68 +++++++++++++++++++---------------------- 1 file changed, 32 insertions(+), 36 deletions(-) diff --git a/internal/gateway/api.go b/internal/gateway/api.go index c8da5ed..417f396 100644 --- a/internal/gateway/api.go +++ b/internal/gateway/api.go @@ -128,6 +128,23 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) { } } +// csvHeaders stamps the shared download headers for CSV exports. +func csvHeaders(w http.ResponseWriter, filename string) { + w.Header().Set("Content-Type", "text/csv; charset=utf-8") + w.Header().Set("Content-Disposition", "attachment; filename="+filename) +} + +// exportKey resolves which masked key id an export covers: admins may pass +// any key filter, user keys are always scoped to themselves (records and +// aggregates are keyed by the masked keyID form). +func exportKey(r *http.Request) string { + k := r.URL.Query().Get("key") + if reqRole(r.Context()) != "admin" { + k = keyID(reqKey(r.Context())) + } + return k +} + // handleStatsAPI returns per-key / per-model / per-source usage aggregates and // the recent request audit trail. func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) { @@ -141,20 +158,14 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) { limit = n } } - key := r.URL.Query().Get("key") - if reqRole(r.Context()) != "admin" { - // user keys may only see their own usage - records and aggregates are - // keyed by the masked id (keyID), so filter on that masked form. - key = keyID(reqKey(r.Context())) - } + key := exportKey(r) if r.URL.Query().Get("export") == "csv" { from, _ := strconv.ParseInt(r.URL.Query().Get("from"), 10, 64) to, _ := strconv.ParseInt(r.URL.Query().Get("to"), 10, 64) if to == 0 { to = time.Now().UnixMilli() } - w.Header().Set("Content-Type", "text/csv; charset=utf-8") - w.Header().Set("Content-Disposition", "attachment; filename=llmsproxy-requests.csv") + csvHeaders(w, "llmsproxy-requests.csv") cw := csv.NewWriter(w) names := map[string]string{} for _, k := range g.core.ListKeys() { @@ -181,31 +192,24 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) { return } if r.URL.Query().Get("export") == "keys-csv" { - w.Header().Set("Content-Type", "text/csv; charset=utf-8") - w.Header().Set("Content-Disposition", "attachment; filename=llmsproxy-keys.csv") + csvHeaders(w, "llmsproxy-keys.csv") cw := csv.NewWriter(w) _ = cw.Write([]string{"key", "key_name", "role", "models", "total_requests", "success_requests", "failed_requests", "prompt_tokens", "completion_tokens", "total_tokens", "avg_latency_ms", "max_latency_ms", "created_at"}) - // Use the same key filtering as the JSON API - exportKey := r.URL.Query().Get("key") - if reqRole(r.Context()) != "admin" { - exportKey = keyID(reqKey(r.Context())) - } + keyFilter := exportKey(r) // by_key rows are keyed by the masked id (keyID); build a masked-id -> // record map so names/roles/models resolve for the export. byMasked := map[string]config.GWKey{} for _, k := range g.core.ListKeys() { byMasked[keyID(k.Key)] = k } - snap := g.stats.Snapshot(0, exportKey) + snap := g.stats.Snapshot(0, keyFilter) byKeyRaw, ok := snap["by_key"].([]StatsRow) if !ok { byKeyRaw = []StatsRow{} } for _, row := range byKeyRaw { keyInfo, found := byMasked[row.Name] - name := "" - role := "" - models := "" + name, role, models := "", "", "" if found { name = keyInfo.Name role = keyInfo.Role @@ -215,17 +219,9 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) { } models = strings.Join(modelNames, ",") } - reqs := row.Stat.Reqs - success := row.Stat.OK - errCount := row.Stat.Err - prompt := row.Stat.Prompt - compl := row.Stat.Compl - tokens := row.Stat.Prompt + row.Stat.Compl - latSum := row.Stat.LatSum - latMax := row.Stat.LatMax avgLat := int64(0) - if reqs > 0 { - avgLat = latSum / reqs + if row.Stat.Reqs > 0 { + avgLat = row.Stat.LatSum / row.Stat.Reqs } createdAt := "" if found && keyInfo.CreatedAt > 0 { @@ -233,14 +229,14 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) { } _ = cw.Write([]string{ row.Name, name, role, models, - strconv.FormatInt(reqs, 10), - strconv.FormatInt(success, 10), - strconv.FormatInt(errCount, 10), - strconv.FormatInt(prompt, 10), - strconv.FormatInt(compl, 10), - strconv.FormatInt(tokens, 10), + strconv.FormatInt(row.Stat.Reqs, 10), + strconv.FormatInt(row.Stat.OK, 10), + strconv.FormatInt(row.Stat.Err, 10), + strconv.FormatInt(row.Stat.Prompt, 10), + strconv.FormatInt(row.Stat.Compl, 10), + strconv.FormatInt(row.Stat.Prompt+row.Stat.Compl, 10), strconv.FormatInt(avgLat, 10), - strconv.FormatInt(latMax, 10), + strconv.FormatInt(row.Stat.LatMax, 10), createdAt, }) }