airkjw/CLIProxyAPI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: remove usage tracking and logging functionality

Browse Source 
- Deleted the `LoggerPlugin` along with associated usage tracking and in-memory statistics logic.
- Removed all related tests (`logger_plugin_test.go`, `usage_tab_test.go`) and external-facing handler (`usage.go`) for usage statistics export/import.
- Cleaned up TUI integration by deleting `usage_tab.go`.
This commit is contained in:
Luis Pater
2026-05-02 04:50:58 +08:00
parent b8bba053fc
commit 18bb9c315f
18 changed files with 116 additions and 1470 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/server/main.go
+2 -2
View File
@@ -24,11 +24,11 @@ import (
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging" "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset" "github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc" "github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
"github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry" "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
"github.com/router-for-me/CLIProxyAPI/v6/internal/store" "github.com/router-for-me/CLIProxyAPI/v6/internal/store"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator"
"github.com/router-for-me/CLIProxyAPI/v6/internal/tui" "github.com/router-for-me/CLIProxyAPI/v6/internal/tui"
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util" "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth" sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
@@ -417,7 +417,7 @@ func main() {
configFileExists = true configFileExists = true
} }
} }
usage.SetStatisticsEnabled(cfg.UsageStatisticsEnabled) redisqueue.SetUsageStatisticsEnabled(cfg.UsageStatisticsEnabled)
coreauth.SetQuotaCooldownDisabled(cfg.DisableCooling) coreauth.SetQuotaCooldownDisabled(cfg.DisableCooling)
if err = logging.ConfigureLogOutput(cfg); err != nil { if err = logging.ConfigureLogOutput(cfg); err != nil {
docker-build.sh
+3 -129
View File
@@ -5,123 +5,13 @@
# This script automates the process of building and running the Docker container # This script automates the process of building and running the Docker container
# with version information dynamically injected at build time. # with version information dynamically injected at build time.
# Hidden feature: Preserve usage statistics across rebuilds
# Usage: ./docker-build.sh --with-usage
# First run prompts for management API key, saved to temp/stats/.api_secret
set -euo pipefail set -euo pipefail
STATS_DIR="temp/stats" if [[ "${1:-}" != "" ]]; then
STATS_FILE="${STATS_DIR}/.usage_backup.json" echo "Error: unknown option '${1}'."
SECRET_FILE="${STATS_DIR}/.api_secret" echo "Usage: ./docker-build.sh"
WITH_USAGE=false
get_port() {
if [[ -f "config.yaml" ]]; then
grep -E "^port:" config.yaml | sed -E 's/^port: *["'"'"']?([0-9]+)["'"'"']?.*$/\1/'
else
echo "8317"
fi
}
export_stats_api_secret() {
if [[ -f "${SECRET_FILE}" ]]; then
API_SECRET=$(cat "${SECRET_FILE}")
else
if [[ ! -d "${STATS_DIR}" ]]; then
mkdir -p "${STATS_DIR}"
fi
echo "First time using --with-usage. Management API key required."
read -r -p "Enter management key: " -s API_SECRET
echo
echo "${API_SECRET}" > "${SECRET_FILE}"
chmod 600 "${SECRET_FILE}"
fi
}
check_container_running() {
local port
port=$(get_port)
if ! curl -s -o /dev/null -w "%{http_code}" "http://localhost:${port}/" | grep -q "200"; then
echo "Error: cli-proxy-api service is not responding at localhost:${port}"
echo "Please start the container first or use without --with-usage flag."
exit 1 exit 1
fi fi
}
export_stats() {
local port
port=$(get_port)
if [[ ! -d "${STATS_DIR}" ]]; then
mkdir -p "${STATS_DIR}"
fi
check_container_running
echo "Exporting usage statistics..."
EXPORT_RESPONSE=$(curl -s -w "\n%{http_code}" -H "X-Management-Key: ${API_SECRET}" \
"http://localhost:${port}/v0/management/usage/export")
HTTP_CODE=$(echo "${EXPORT_RESPONSE}" | tail -n1)
RESPONSE_BODY=$(echo "${EXPORT_RESPONSE}" | sed '$d')
if [[ "${HTTP_CODE}" != "200" ]]; then
echo "Export failed (HTTP ${HTTP_CODE}): ${RESPONSE_BODY}"
exit 1
fi
echo "${RESPONSE_BODY}" > "${STATS_FILE}"
echo "Statistics exported to ${STATS_FILE}"
}
import_stats() {
local port
port=$(get_port)
echo "Importing usage statistics..."
IMPORT_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST \
-H "X-Management-Key: ${API_SECRET}" \
-H "Content-Type: application/json" \
-d @"${STATS_FILE}" \
"http://localhost:${port}/v0/management/usage/import")
IMPORT_CODE=$(echo "${IMPORT_RESPONSE}" | tail -n1)
IMPORT_BODY=$(echo "${IMPORT_RESPONSE}" | sed '$d')
if [[ "${IMPORT_CODE}" == "200" ]]; then
echo "Statistics imported successfully"
else
echo "Import failed (HTTP ${IMPORT_CODE}): ${IMPORT_BODY}"
fi
rm -f "${STATS_FILE}"
}
wait_for_service() {
local port
port=$(get_port)
echo "Waiting for service to be ready..."
for i in {1..30}; do
if curl -s -o /dev/null -w "%{http_code}" "http://localhost:${port}/" | grep -q "200"; then
break
fi
sleep 1
done
sleep 2
}
case "${1:-}" in
"")
;;
"--with-usage")
WITH_USAGE=true
export_stats_api_secret
;;
*)
echo "Error: unknown option '${1}'. Did you mean '--with-usage'?"
echo "Usage: ./docker-build.sh [--with-usage]"
exit 1
;;
esac
# --- Step 1: Choose Environment --- # --- Step 1: Choose Environment ---
echo "Please select an option:" echo "Please select an option:"
@@ -133,14 +23,7 @@ read -r -p "Enter choice [1-2]: " choice
case "$choice" in case "$choice" in
1) 1)
echo "--- Running with Pre-built Image ---" echo "--- Running with Pre-built Image ---"
if [[ "${WITH_USAGE}" == "true" ]]; then
export_stats
fi
docker compose up -d --remove-orphans --no-build docker compose up -d --remove-orphans --no-build
if [[ "${WITH_USAGE}" == "true" ]]; then
wait_for_service
import_stats
fi
echo "Services are starting from remote image." echo "Services are starting from remote image."
echo "Run 'docker compose logs -f' to see the logs." echo "Run 'docker compose logs -f' to see the logs."
;; ;;
@@ -167,18 +50,9 @@ case "$choice" in
--build-arg COMMIT="${COMMIT}" \ --build-arg COMMIT="${COMMIT}" \
--build-arg BUILD_DATE="${BUILD_DATE}" --build-arg BUILD_DATE="${BUILD_DATE}"
if [[ "${WITH_USAGE}" == "true" ]]; then
export_stats
fi
echo "Starting the services..." echo "Starting the services..."
docker compose up -d --remove-orphans --pull never docker compose up -d --remove-orphans --pull never
if [[ "${WITH_USAGE}" == "true" ]]; then
wait_for_service
import_stats
fi
echo "Build complete. Services are starting." echo "Build complete. Services are starting."
echo "Run 'docker compose logs -f' to see the logs." echo "Run 'docker compose logs -f' to see the logs."
;; ;;
internal/api/handlers/management/handler.go
-6
View File
@@ -15,7 +15,6 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/buildinfo" "github.com/router-for-me/CLIProxyAPI/v6/internal/buildinfo"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth" sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
"golang.org/x/crypto/bcrypt" "golang.org/x/crypto/bcrypt"
@@ -41,7 +40,6 @@ type Handler struct {
attemptsMu sync.Mutex attemptsMu sync.Mutex
failedAttempts map[string]*attemptInfo // keyed by client IP failedAttempts map[string]*attemptInfo // keyed by client IP
authManager *coreauth.Manager authManager *coreauth.Manager
usageStats *usage.RequestStatistics
tokenStore coreauth.Store tokenStore coreauth.Store
localPassword string localPassword string
allowRemoteOverride bool allowRemoteOverride bool
@@ -60,7 +58,6 @@ func NewHandler(cfg *config.Config, configFilePath string, manager *coreauth.Man
configFilePath: configFilePath, configFilePath: configFilePath,
failedAttempts: make(map[string]*attemptInfo), failedAttempts: make(map[string]*attemptInfo),
authManager: manager, authManager: manager,
usageStats: usage.GetRequestStatistics(),
tokenStore: sdkAuth.GetTokenStore(), tokenStore: sdkAuth.GetTokenStore(),
allowRemoteOverride: envSecret != "", allowRemoteOverride: envSecret != "",
envSecret: envSecret, envSecret: envSecret,
@@ -124,9 +121,6 @@ func (h *Handler) SetAuthManager(manager *coreauth.Manager) {
h.mu.Unlock() h.mu.Unlock()
} }
// SetUsageStatistics allows replacing the usage statistics reference.
func (h *Handler) SetUsageStatistics(stats *usage.RequestStatistics) { h.usageStats = stats }
// SetLocalPassword configures the runtime-local password accepted for localhost requests. // SetLocalPassword configures the runtime-local password accepted for localhost requests.
func (h *Handler) SetLocalPassword(password string) { h.localPassword = password } func (h *Handler) SetLocalPassword(password string) { h.localPassword = password }
internal/api/handlers/management/usage.go
-79
View File
@@ -1,79 +0,0 @@
package management
import (
"encoding/json"
"net/http"
"time"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
)
type usageExportPayload struct {
Version int `json:"version"`
ExportedAt time.Time `json:"exported_at"`
Usage usage.StatisticsSnapshot `json:"usage"`
}
type usageImportPayload struct {
Version int `json:"version"`
Usage usage.StatisticsSnapshot `json:"usage"`
}
// GetUsageStatistics returns the in-memory request statistics snapshot.
func (h *Handler) GetUsageStatistics(c *gin.Context) {
var snapshot usage.StatisticsSnapshot
if h != nil && h.usageStats != nil {
snapshot = h.usageStats.Snapshot()
}
c.JSON(http.StatusOK, gin.H{
"usage": snapshot,
"failed_requests": snapshot.FailureCount,
})
}
// ExportUsageStatistics returns a complete usage snapshot for backup/migration.
func (h *Handler) ExportUsageStatistics(c *gin.Context) {
var snapshot usage.StatisticsSnapshot
if h != nil && h.usageStats != nil {
snapshot = h.usageStats.Snapshot()
}
c.JSON(http.StatusOK, usageExportPayload{
Version: 1,
ExportedAt: time.Now().UTC(),
Usage: snapshot,
})
}
// ImportUsageStatistics merges a previously exported usage snapshot into memory.
func (h *Handler) ImportUsageStatistics(c *gin.Context) {
if h == nil || h.usageStats == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "usage statistics unavailable"})
return
}
data, err := c.GetRawData()
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "failed to read request body"})
return
}
var payload usageImportPayload
if err := json.Unmarshal(data, &payload); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json"})
return
}
if payload.Version != 0 && payload.Version != 1 {
c.JSON(http.StatusBadRequest, gin.H{"error": "unsupported version"})
return
}
result := h.usageStats.MergeSnapshot(payload.Usage)
snapshot := h.usageStats.Snapshot()
c.JSON(http.StatusOK, gin.H{
"added": result.Added,
"skipped": result.Skipped,
"total_requests": snapshot.TotalRequests,
"failed_requests": snapshot.FailureCount,
})
}
internal/api/server.go
+1 -5
View File
@@ -31,7 +31,6 @@ import (
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging" "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset" "github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
"github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue" "github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util" "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access" sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers" "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
@@ -507,9 +506,6 @@ func (s *Server) registerManagementRoutes() {
mgmt := s.engine.Group("/v0/management") mgmt := s.engine.Group("/v0/management")
mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware()) mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware())
{ {
mgmt.GET("/usage", s.mgmt.GetUsageStatistics)
mgmt.GET("/usage/export", s.mgmt.ExportUsageStatistics)
mgmt.POST("/usage/import", s.mgmt.ImportUsageStatistics)
mgmt.GET("/config", s.mgmt.GetConfig) mgmt.GET("/config", s.mgmt.GetConfig)
mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML) mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML)
mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML) mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML)
@@ -1001,7 +997,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
} }
if oldCfg == nil || oldCfg.UsageStatisticsEnabled != cfg.UsageStatisticsEnabled { if oldCfg == nil || oldCfg.UsageStatisticsEnabled != cfg.UsageStatisticsEnabled {
usage.SetStatisticsEnabled(cfg.UsageStatisticsEnabled) redisqueue.SetUsageStatisticsEnabled(cfg.UsageStatisticsEnabled)
} }
if s.requestLogger != nil && (oldCfg == nil || oldCfg.ErrorLogsMaxFiles != cfg.ErrorLogsMaxFiles) { if s.requestLogger != nil && (oldCfg == nil || oldCfg.ErrorLogsMaxFiles != cfg.ErrorLogsMaxFiles) {
internal/redisqueue/plugin.go
+22 -6
View File
@@ -7,7 +7,6 @@ import (
"time" "time"
internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging" internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
internalusage "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage" coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
) )
@@ -21,7 +20,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
if p == nil { if p == nil {
return return
} }
if !Enabled() || !internalusage.StatisticsEnabled() { if !Enabled() || !UsageStatisticsEnabled() {
return return
} }
@@ -45,7 +44,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
apiKey := strings.TrimSpace(record.APIKey) apiKey := strings.TrimSpace(record.APIKey)
requestID := strings.TrimSpace(internallogging.GetRequestID(ctx)) requestID := strings.TrimSpace(internallogging.GetRequestID(ctx))
tokens := internalusage.TokenStats{ tokens := tokenStats{
InputTokens: record.Detail.InputTokens, InputTokens: record.Detail.InputTokens,
OutputTokens: record.Detail.OutputTokens, OutputTokens: record.Detail.OutputTokens,
ReasoningTokens: record.Detail.ReasoningTokens, ReasoningTokens: record.Detail.ReasoningTokens,
@@ -64,7 +63,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
failed = !resolveSuccess(ctx) failed = !resolveSuccess(ctx)
} }
detail := internalusage.RequestDetail{ detail := requestDetail{
Timestamp: timestamp, Timestamp: timestamp,
LatencyMs: record.Latency.Milliseconds(), LatencyMs: record.Latency.Milliseconds(),
Source: record.Source, Source: record.Source,
@@ -74,7 +73,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
} }
payload, err := json.Marshal(queuedUsageDetail{ payload, err := json.Marshal(queuedUsageDetail{
RequestDetail: detail, requestDetail: detail,
Provider: provider, Provider: provider,
Model: modelName, Model: modelName,
Endpoint: resolveEndpoint(ctx), Endpoint: resolveEndpoint(ctx),
@@ -89,7 +88,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
} }
type queuedUsageDetail struct { type queuedUsageDetail struct {
internalusage.RequestDetail requestDetail
Provider string `json:"provider"` Provider string `json:"provider"`
Model string `json:"model"` Model string `json:"model"`
Endpoint string `json:"endpoint"` Endpoint string `json:"endpoint"`
@@ -98,6 +97,23 @@ type queuedUsageDetail struct {
RequestID string `json:"request_id"` RequestID string `json:"request_id"`
} }
type requestDetail struct {
Timestamp time.Time `json:"timestamp"`
LatencyMs int64 `json:"latency_ms"`
Source string `json:"source"`
AuthIndex string `json:"auth_index"`
Tokens tokenStats `json:"tokens"`
Failed bool `json:"failed"`
}
type tokenStats struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
ReasoningTokens int64 `json:"reasoning_tokens"`
CachedTokens int64 `json:"cached_tokens"`
TotalTokens int64 `json:"total_tokens"`
}
func resolveSuccess(ctx context.Context) bool { func resolveSuccess(ctx context.Context) bool {
status := internallogging.GetResponseStatus(ctx) status := internallogging.GetResponseStatus(ctx)
if status == 0 { if status == 0 {
internal/redisqueue/plugin_test.go
+3 -4
View File
@@ -10,7 +10,6 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging" internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
internalusage "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage" coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
) )
@@ -127,16 +126,16 @@ func withEnabledQueue(t *testing.T, fn func()) {
t.Helper() t.Helper()
prevQueueEnabled := Enabled() prevQueueEnabled := Enabled()
prevStatsEnabled := internalusage.StatisticsEnabled() prevUsageEnabled := UsageStatisticsEnabled()
SetEnabled(false) SetEnabled(false)
SetEnabled(true) SetEnabled(true)
internalusage.SetStatisticsEnabled(true) SetUsageStatisticsEnabled(true)
defer func() { defer func() {
SetEnabled(false) SetEnabled(false)
SetEnabled(prevQueueEnabled) SetEnabled(prevQueueEnabled)
internalusage.SetStatisticsEnabled(prevStatsEnabled) SetUsageStatisticsEnabled(prevUsageEnabled)
}() }()
fn() fn()
internal/redisqueue/usage_toggle.go
+16
View File
@@ -0,0 +1,16 @@
package redisqueue
import "sync/atomic"
var usageStatisticsEnabled atomic.Bool
func init() {
usageStatisticsEnabled.Store(true)
}
// SetUsageStatisticsEnabled toggles whether usage records are enqueued into the redisqueue payload buffer.
// This is controlled by the config field `usage-statistics-enabled` and the corresponding management API.
func SetUsageStatisticsEnabled(enabled bool) { usageStatisticsEnabled.Store(enabled) }
// UsageStatisticsEnabled reports whether the usage queue plugin should publish records.
func UsageStatisticsEnabled() bool { return usageStatisticsEnabled.Load() }
internal/tui/app.go
+4 -18
View File
@@ -18,7 +18,6 @@ const (
tabAuthFiles tabAuthFiles
tabAPIKeys tabAPIKeys
tabOAuth tabOAuth
tabUsage
tabLogs tabLogs
) )
@@ -40,7 +39,6 @@ type App struct {
auth authTabModel auth authTabModel
keys keysTabModel keys keysTabModel
oauth oauthTabModel oauth oauthTabModel
usage usageTabModel
logs logsTabModel logs logsTabModel
client *Client client *Client
@@ -50,7 +48,7 @@ type App struct {
ready bool ready bool
// Track which tabs have been initialized (fetched data) // Track which tabs have been initialized (fetched data)
initialized [7]bool initialized [6]bool
} }
type authConnectMsg struct { type authConnectMsg struct {
@@ -81,10 +79,9 @@ func NewApp(port int, secretKey string, hook *LogHook) App {
auth: newAuthTabModel(client), auth: newAuthTabModel(client),
keys: newKeysTabModel(client), keys: newKeysTabModel(client),
oauth: newOAuthTabModel(client), oauth: newOAuthTabModel(client),
usage: newUsageTabModel(client),
logs: newLogsTabModel(client, hook), logs: newLogsTabModel(client, hook),
client: client, client: client,
initialized: [7]bool{ initialized: [6]bool{
tabDashboard: true, tabDashboard: true,
tabLogs: true, tabLogs: true,
}, },
@@ -92,7 +89,7 @@ func NewApp(port int, secretKey string, hook *LogHook) App {
app.refreshTabs() app.refreshTabs()
if authRequired { if authRequired {
app.initialized = [7]bool{} app.initialized = [6]bool{}
} }
app.setAuthInputPrompt() app.setAuthInputPrompt()
return app return app
@@ -128,7 +125,6 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
a.auth.SetSize(contentW, contentH) a.auth.SetSize(contentW, contentH)
a.keys.SetSize(contentW, contentH) a.keys.SetSize(contentW, contentH)
a.oauth.SetSize(contentW, contentH) a.oauth.SetSize(contentW, contentH)
a.usage.SetSize(contentW, contentH)
a.logs.SetSize(contentW, contentH) a.logs.SetSize(contentW, contentH)
return a, nil return a, nil
@@ -142,7 +138,7 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
a.authenticated = true a.authenticated = true
a.logsEnabled = a.standalone || isLogsEnabledFromConfig(msg.cfg) a.logsEnabled = a.standalone || isLogsEnabledFromConfig(msg.cfg)
a.refreshTabs() a.refreshTabs()
a.initialized = [7]bool{} a.initialized = [6]bool{}
a.initialized[tabDashboard] = true a.initialized[tabDashboard] = true
cmds := []tea.Cmd{a.dashboard.Init()} cmds := []tea.Cmd{a.dashboard.Init()}
if a.logsEnabled { if a.logsEnabled {
@@ -258,8 +254,6 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
a.keys, cmd = a.keys.Update(msg) a.keys, cmd = a.keys.Update(msg)
case tabOAuth: case tabOAuth:
a.oauth, cmd = a.oauth.Update(msg) a.oauth, cmd = a.oauth.Update(msg)
case tabUsage:
a.usage, cmd = a.usage.Update(msg)
case tabLogs: case tabLogs:
a.logs, cmd = a.logs.Update(msg) a.logs, cmd = a.logs.Update(msg)
} }
@@ -322,8 +316,6 @@ func (a *App) initTabIfNeeded(_ int) tea.Cmd {
return a.keys.Init() return a.keys.Init()
case tabOAuth: case tabOAuth:
return a.oauth.Init() return a.oauth.Init()
case tabUsage:
return a.usage.Init()
case tabLogs: case tabLogs:
if !a.logsEnabled { if !a.logsEnabled {
return nil return nil
@@ -360,8 +352,6 @@ func (a App) View() string {
sb.WriteString(a.keys.View()) sb.WriteString(a.keys.View())
case tabOAuth: case tabOAuth:
sb.WriteString(a.oauth.View()) sb.WriteString(a.oauth.View())
case tabUsage:
sb.WriteString(a.usage.View())
case tabLogs: case tabLogs:
if a.logsEnabled { if a.logsEnabled {
sb.WriteString(a.logs.View()) sb.WriteString(a.logs.View())
@@ -529,10 +519,6 @@ func (a App) broadcastToAllTabs(msg tea.Msg) (tea.Model, tea.Cmd) {
if cmd != nil { if cmd != nil {
cmds = append(cmds, cmd) cmds = append(cmds, cmd)
} }
a.usage, cmd = a.usage.Update(msg)
if cmd != nil {
cmds = append(cmds, cmd)
}
a.logs, cmd = a.logs.Update(msg) a.logs, cmd = a.logs.Update(msg)
if cmd != nil { if cmd != nil {
cmds = append(cmds, cmd) cmds = append(cmds, cmd)
internal/tui/client.go
-5
View File
@@ -140,11 +140,6 @@ func (c *Client) PutConfigYAML(yamlContent string) error {
return err return err
} }
// GetUsage fetches usage statistics.
func (c *Client) GetUsage() (map[string]any, error) {
return c.getJSON("/v0/management/usage")
}
// GetAuthFiles lists auth credential files. // GetAuthFiles lists auth credential files.
// API returns {"files": [...]}. // API returns {"files": [...]}.
func (c *Client) GetAuthFiles() ([]map[string]any, error) { func (c *Client) GetAuthFiles() ([]map[string]any, error) {
internal/tui/dashboard.go
+7 -70
View File
@@ -22,14 +22,12 @@ type dashboardModel struct {
// Cached data for re-rendering on locale change // Cached data for re-rendering on locale change
lastConfig map[string]any lastConfig map[string]any
lastUsage map[string]any
lastAuthFiles []map[string]any lastAuthFiles []map[string]any
lastAPIKeys []string lastAPIKeys []string
} }
type dashboardDataMsg struct { type dashboardDataMsg struct {
config map[string]any config map[string]any
usage map[string]any
authFiles []map[string]any authFiles []map[string]any
apiKeys []string apiKeys []string
err error err error
@@ -47,25 +45,24 @@ func (m dashboardModel) Init() tea.Cmd {
func (m dashboardModel) fetchData() tea.Msg { func (m dashboardModel) fetchData() tea.Msg {
cfg, cfgErr := m.client.GetConfig() cfg, cfgErr := m.client.GetConfig()
usage, usageErr := m.client.GetUsage()
authFiles, authErr := m.client.GetAuthFiles() authFiles, authErr := m.client.GetAuthFiles()
apiKeys, keysErr := m.client.GetAPIKeys() apiKeys, keysErr := m.client.GetAPIKeys()
var err error var err error
for _, e := range []error{cfgErr, usageErr, authErr, keysErr} { for _, e := range []error{cfgErr, authErr, keysErr} {
if e != nil { if e != nil {
err = e err = e
break break
} }
} }
return dashboardDataMsg{config: cfg, usage: usage, authFiles: authFiles, apiKeys: apiKeys, err: err} return dashboardDataMsg{config: cfg, authFiles: authFiles, apiKeys: apiKeys, err: err}
} }
func (m dashboardModel) Update(msg tea.Msg) (dashboardModel, tea.Cmd) { func (m dashboardModel) Update(msg tea.Msg) (dashboardModel, tea.Cmd) {
switch msg := msg.(type) { switch msg := msg.(type) {
case localeChangedMsg: case localeChangedMsg:
// Re-render immediately with cached data using new locale // Re-render immediately with cached data using new locale
m.content = m.renderDashboard(m.lastConfig, m.lastUsage, m.lastAuthFiles, m.lastAPIKeys) m.content = m.renderDashboard(m.lastConfig, m.lastAuthFiles, m.lastAPIKeys)
m.viewport.SetContent(m.content) m.viewport.SetContent(m.content)
// Also fetch fresh data in background // Also fetch fresh data in background
return m, m.fetchData return m, m.fetchData
@@ -78,11 +75,10 @@ func (m dashboardModel) Update(msg tea.Msg) (dashboardModel, tea.Cmd) {
m.err = nil m.err = nil
// Cache data for locale switching // Cache data for locale switching
m.lastConfig = msg.config m.lastConfig = msg.config
m.lastUsage = msg.usage
m.lastAuthFiles = msg.authFiles m.lastAuthFiles = msg.authFiles
m.lastAPIKeys = msg.apiKeys m.lastAPIKeys = msg.apiKeys
m.content = m.renderDashboard(msg.config, msg.usage, msg.authFiles, msg.apiKeys) m.content = m.renderDashboard(msg.config, msg.authFiles, msg.apiKeys)
} }
m.viewport.SetContent(m.content) m.viewport.SetContent(m.content)
return m, nil return m, nil
@@ -121,7 +117,7 @@ func (m dashboardModel) View() string {
return m.viewport.View() return m.viewport.View()
} }
func (m dashboardModel) renderDashboard(cfg, usage map[string]any, authFiles []map[string]any, apiKeys []string) string { func (m dashboardModel) renderDashboard(cfg map[string]any, authFiles []map[string]any, apiKeys []string) string {
var sb strings.Builder var sb strings.Builder
sb.WriteString(titleStyle.Render(T("dashboard_title"))) sb.WriteString(titleStyle.Render(T("dashboard_title")))
@@ -138,7 +134,7 @@ func (m dashboardModel) renderDashboard(cfg, usage map[string]any, authFiles []m
// ━━━ Stats Cards ━━━ // ━━━ Stats Cards ━━━
cardWidth := 25 cardWidth := 25
if m.width > 0 { if m.width > 0 {
cardWidth = (m.width - 6) / 4 cardWidth = (m.width - 2) / 2
if cardWidth < 18 { if cardWidth < 18 {
cardWidth = 18 cardWidth = 18
} }
@@ -173,34 +169,7 @@ func (m dashboardModel) renderDashboard(cfg, usage map[string]any, authFiles []m
lipgloss.NewStyle().Foreground(colorMuted).Render(fmt.Sprintf("%s (%d %s)", T("auth_files_label"), activeAuth, T("active_suffix"))), lipgloss.NewStyle().Foreground(colorMuted).Render(fmt.Sprintf("%s (%d %s)", T("auth_files_label"), activeAuth, T("active_suffix"))),
)) ))
// Card 3: Total Requests sb.WriteString(lipgloss.JoinHorizontal(lipgloss.Top, card1, " ", card2))
totalReqs := int64(0)
successReqs := int64(0)
failedReqs := int64(0)
totalTokens := int64(0)
if usage != nil {
if usageMap, ok := usage["usage"].(map[string]any); ok {
totalReqs = int64(getFloat(usageMap, "total_requests"))
successReqs = int64(getFloat(usageMap, "success_count"))
failedReqs = int64(getFloat(usageMap, "failure_count"))
totalTokens = int64(getFloat(usageMap, "total_tokens"))
}
}
card3 := cardStyle.Render(fmt.Sprintf(
"%s\n%s",
lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("214")).Render(fmt.Sprintf("📈 %d", totalReqs)),
lipgloss.NewStyle().Foreground(colorMuted).Render(fmt.Sprintf("%s (✓%d ✗%d)", T("total_requests"), successReqs, failedReqs)),
))
// Card 4: Total Tokens
tokenStr := formatLargeNumber(totalTokens)
card4 := cardStyle.Render(fmt.Sprintf(
"%s\n%s",
lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("170")).Render(fmt.Sprintf("🔤 %s", tokenStr)),
lipgloss.NewStyle().Foreground(colorMuted).Render(T("total_tokens")),
))
sb.WriteString(lipgloss.JoinHorizontal(lipgloss.Top, card1, " ", card2, " ", card3, " ", card4))
sb.WriteString("\n\n") sb.WriteString("\n\n")
// ━━━ Current Config ━━━ // ━━━ Current Config ━━━
@@ -258,38 +227,6 @@ func (m dashboardModel) renderDashboard(cfg, usage map[string]any, authFiles []m
sb.WriteString("\n") sb.WriteString("\n")
// ━━━ Per-Model Usage ━━━
if usage != nil {
if usageMap, ok := usage["usage"].(map[string]any); ok {
if apis, ok := usageMap["apis"].(map[string]any); ok && len(apis) > 0 {
sb.WriteString(lipgloss.NewStyle().Bold(true).Foreground(colorHighlight).Render(T("model_stats")))
sb.WriteString("\n")
sb.WriteString(strings.Repeat("─", minInt(m.width, 60)))
sb.WriteString("\n")
header := fmt.Sprintf(" %-40s %10s %12s", T("model"), T("requests"), T("tokens"))
sb.WriteString(tableHeaderStyle.Render(header))
sb.WriteString("\n")
for _, apiSnap := range apis {
if apiMap, ok := apiSnap.(map[string]any); ok {
if models, ok := apiMap["models"].(map[string]any); ok {
for model, v := range models {
if stats, ok := v.(map[string]any); ok {
reqs := int64(getFloat(stats, "total_requests"))
toks := int64(getFloat(stats, "total_tokens"))
row := fmt.Sprintf(" %-40s %10d %12s", truncate(model, 40), reqs, formatLargeNumber(toks))
sb.WriteString(tableCellStyle.Render(row))
sb.WriteString("\n")
}
}
}
}
}
}
}
}
return sb.String() return sb.String()
} }
internal/tui/i18n.go
+2 -2
View File
@@ -50,8 +50,8 @@ var locales = map[string]map[string]string{
// ────────────────────────────────────────── // ──────────────────────────────────────────
// Tab names // Tab names
// ────────────────────────────────────────── // ──────────────────────────────────────────
var zhTabNames = []string{"仪表盘", "配置", "认证文件", "API 密钥", "OAuth", "使用统计", "日志"} var zhTabNames = []string{"仪表盘", "配置", "认证文件", "API 密钥", "OAuth", "日志"}
var enTabNames = []string{"Dashboard", "Config", "Auth Files", "API Keys", "OAuth", "Usage", "Logs"} var enTabNames = []string{"Dashboard", "Config", "Auth Files", "API Keys", "OAuth", "Logs"}
// TabNames returns tab names in the current locale. // TabNames returns tab names in the current locale.
func TabNames() []string { func TabNames() []string {
internal/tui/usage_tab.go
-418
View File
@@ -1,418 +0,0 @@
package tui
import (
"fmt"
"sort"
"strings"
"github.com/charmbracelet/bubbles/viewport"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
// usageTabModel displays usage statistics with charts and breakdowns.
type usageTabModel struct {
client *Client
viewport viewport.Model
usage map[string]any
err error
width int
height int
ready bool
}
type usageDataMsg struct {
usage map[string]any
err error
}
func newUsageTabModel(client *Client) usageTabModel {
return usageTabModel{
client: client,
}
}
func (m usageTabModel) Init() tea.Cmd {
return m.fetchData
}
func (m usageTabModel) fetchData() tea.Msg {
usage, err := m.client.GetUsage()
return usageDataMsg{usage: usage, err: err}
}
func (m usageTabModel) Update(msg tea.Msg) (usageTabModel, tea.Cmd) {
switch msg := msg.(type) {
case localeChangedMsg:
m.viewport.SetContent(m.renderContent())
return m, nil
case usageDataMsg:
if msg.err != nil {
m.err = msg.err
} else {
m.err = nil
m.usage = msg.usage
}
m.viewport.SetContent(m.renderContent())
return m, nil
case tea.KeyMsg:
if msg.String() == "r" {
return m, m.fetchData
}
var cmd tea.Cmd
m.viewport, cmd = m.viewport.Update(msg)
return m, cmd
}
var cmd tea.Cmd
m.viewport, cmd = m.viewport.Update(msg)
return m, cmd
}
func (m *usageTabModel) SetSize(w, h int) {
m.width = w
m.height = h
if !m.ready {
m.viewport = viewport.New(w, h)
m.viewport.SetContent(m.renderContent())
m.ready = true
} else {
m.viewport.Width = w
m.viewport.Height = h
}
}
func (m usageTabModel) View() string {
if !m.ready {
return T("loading")
}
return m.viewport.View()
}
func (m usageTabModel) renderContent() string {
var sb strings.Builder
sb.WriteString(titleStyle.Render(T("usage_title")))
sb.WriteString("\n")
sb.WriteString(helpStyle.Render(T("usage_help")))
sb.WriteString("\n\n")
if m.err != nil {
sb.WriteString(errorStyle.Render("⚠ Error: " + m.err.Error()))
sb.WriteString("\n")
return sb.String()
}
if m.usage == nil {
sb.WriteString(subtitleStyle.Render(T("usage_no_data")))
sb.WriteString("\n")
return sb.String()
}
usageMap, _ := m.usage["usage"].(map[string]any)
if usageMap == nil {
sb.WriteString(subtitleStyle.Render(T("usage_no_data")))
sb.WriteString("\n")
return sb.String()
}
totalReqs := int64(getFloat(usageMap, "total_requests"))
successCnt := int64(getFloat(usageMap, "success_count"))
failureCnt := int64(getFloat(usageMap, "failure_count"))
totalTokens := int64(getFloat(usageMap, "total_tokens"))
// ━━━ Overview Cards ━━━
cardWidth := 20
if m.width > 0 {
cardWidth = (m.width - 6) / 4
if cardWidth < 16 {
cardWidth = 16
}
}
cardStyle := lipgloss.NewStyle().
Border(lipgloss.RoundedBorder()).
BorderForeground(lipgloss.Color("240")).
Padding(0, 1).
Width(cardWidth).
Height(3)
// Total Requests
card1 := cardStyle.Copy().BorderForeground(lipgloss.Color("111")).Render(fmt.Sprintf(
"%s\n%s\n%s",
lipgloss.NewStyle().Foreground(colorMuted).Render(T("usage_total_reqs")),
lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("111")).Render(fmt.Sprintf("%d", totalReqs)),
lipgloss.NewStyle().Foreground(colorMuted).Render(fmt.Sprintf("● %s: %d ● %s: %d", T("usage_success"), successCnt, T("usage_failure"), failureCnt)),
))
// Total Tokens
card2 := cardStyle.Copy().BorderForeground(lipgloss.Color("214")).Render(fmt.Sprintf(
"%s\n%s\n%s",
lipgloss.NewStyle().Foreground(colorMuted).Render(T("usage_total_tokens")),
lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("214")).Render(formatLargeNumber(totalTokens)),
lipgloss.NewStyle().Foreground(colorMuted).Render(fmt.Sprintf("%s: %s", T("usage_total_token_l"), formatLargeNumber(totalTokens))),
))
// RPM
rpm := float64(0)
if totalReqs > 0 {
if rByH, ok := usageMap["requests_by_hour"].(map[string]any); ok && len(rByH) > 0 {
rpm = float64(totalReqs) / float64(len(rByH)) / 60.0
}
}
card3 := cardStyle.Copy().BorderForeground(lipgloss.Color("76")).Render(fmt.Sprintf(
"%s\n%s\n%s",
lipgloss.NewStyle().Foreground(colorMuted).Render(T("usage_rpm")),
lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("76")).Render(fmt.Sprintf("%.2f", rpm)),
lipgloss.NewStyle().Foreground(colorMuted).Render(fmt.Sprintf("%s: %d", T("usage_total_reqs"), totalReqs)),
))
// TPM
tpm := float64(0)
if totalTokens > 0 {
if tByH, ok := usageMap["tokens_by_hour"].(map[string]any); ok && len(tByH) > 0 {
tpm = float64(totalTokens) / float64(len(tByH)) / 60.0
}
}
card4 := cardStyle.Copy().BorderForeground(lipgloss.Color("170")).Render(fmt.Sprintf(
"%s\n%s\n%s",
lipgloss.NewStyle().Foreground(colorMuted).Render(T("usage_tpm")),
lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("170")).Render(fmt.Sprintf("%.2f", tpm)),
lipgloss.NewStyle().Foreground(colorMuted).Render(fmt.Sprintf("%s: %s", T("usage_total_tokens"), formatLargeNumber(totalTokens))),
))
sb.WriteString(lipgloss.JoinHorizontal(lipgloss.Top, card1, " ", card2, " ", card3, " ", card4))
sb.WriteString("\n\n")
// ━━━ Requests by Hour (ASCII bar chart) ━━━
if rByH, ok := usageMap["requests_by_hour"].(map[string]any); ok && len(rByH) > 0 {
sb.WriteString(lipgloss.NewStyle().Bold(true).Foreground(colorHighlight).Render(T("usage_req_by_hour")))
sb.WriteString("\n")
sb.WriteString(strings.Repeat("─", minInt(m.width, 60)))
sb.WriteString("\n")
sb.WriteString(renderBarChart(rByH, m.width-6, lipgloss.Color("111")))
sb.WriteString("\n")
}
// ━━━ Tokens by Hour ━━━
if tByH, ok := usageMap["tokens_by_hour"].(map[string]any); ok && len(tByH) > 0 {
sb.WriteString(lipgloss.NewStyle().Bold(true).Foreground(colorHighlight).Render(T("usage_tok_by_hour")))
sb.WriteString("\n")
sb.WriteString(strings.Repeat("─", minInt(m.width, 60)))
sb.WriteString("\n")
sb.WriteString(renderBarChart(tByH, m.width-6, lipgloss.Color("214")))
sb.WriteString("\n")
}
// ━━━ Requests by Day ━━━
if rByD, ok := usageMap["requests_by_day"].(map[string]any); ok && len(rByD) > 0 {
sb.WriteString(lipgloss.NewStyle().Bold(true).Foreground(colorHighlight).Render(T("usage_req_by_day")))
sb.WriteString("\n")
sb.WriteString(strings.Repeat("─", minInt(m.width, 60)))
sb.WriteString("\n")
sb.WriteString(renderBarChart(rByD, m.width-6, lipgloss.Color("76")))
sb.WriteString("\n")
}
// ━━━ API Detail Stats ━━━
if apis, ok := usageMap["apis"].(map[string]any); ok && len(apis) > 0 {
sb.WriteString(lipgloss.NewStyle().Bold(true).Foreground(colorHighlight).Render(T("usage_api_detail")))
sb.WriteString("\n")
sb.WriteString(strings.Repeat("─", minInt(m.width, 80)))
sb.WriteString("\n")
header := fmt.Sprintf(" %-30s %10s %12s", "API", T("requests"), T("tokens"))
sb.WriteString(tableHeaderStyle.Render(header))
sb.WriteString("\n")
for apiName, apiSnap := range apis {
if apiMap, ok := apiSnap.(map[string]any); ok {
apiReqs := int64(getFloat(apiMap, "total_requests"))
apiToks := int64(getFloat(apiMap, "total_tokens"))
row := fmt.Sprintf(" %-30s %10d %12s",
truncate(maskKey(apiName), 30), apiReqs, formatLargeNumber(apiToks))
sb.WriteString(lipgloss.NewStyle().Bold(true).Render(row))
sb.WriteString("\n")
// Per-model breakdown
if models, ok := apiMap["models"].(map[string]any); ok {
for model, v := range models {
if stats, ok := v.(map[string]any); ok {
mReqs := int64(getFloat(stats, "total_requests"))
mToks := int64(getFloat(stats, "total_tokens"))
mRow := fmt.Sprintf(" ├─ %-28s %10d %12s",
truncate(model, 28), mReqs, formatLargeNumber(mToks))
sb.WriteString(tableCellStyle.Render(mRow))
sb.WriteString("\n")
// Token type breakdown from details
sb.WriteString(m.renderTokenBreakdown(stats))
// Latency breakdown from details
sb.WriteString(m.renderLatencyBreakdown(stats))
}
}
}
}
}
}
sb.WriteString("\n")
return sb.String()
}
// renderTokenBreakdown aggregates input/output/cached/reasoning tokens from model details.
func (m usageTabModel) renderTokenBreakdown(modelStats map[string]any) string {
details, ok := modelStats["details"]
if !ok {
return ""
}
detailList, ok := details.([]any)
if !ok || len(detailList) == 0 {
return ""
}
var inputTotal, outputTotal, cachedTotal, reasoningTotal int64
for _, d := range detailList {
dm, ok := d.(map[string]any)
if !ok {
continue
}
tokens, ok := dm["tokens"].(map[string]any)
if !ok {
continue
}
inputTotal += int64(getFloat(tokens, "input_tokens"))
outputTotal += int64(getFloat(tokens, "output_tokens"))
cachedTotal += int64(getFloat(tokens, "cached_tokens"))
reasoningTotal += int64(getFloat(tokens, "reasoning_tokens"))
}
if inputTotal == 0 && outputTotal == 0 && cachedTotal == 0 && reasoningTotal == 0 {
return ""
}
parts := []string{}
if inputTotal > 0 {
parts = append(parts, fmt.Sprintf("%s:%s", T("usage_input"), formatLargeNumber(inputTotal)))
}
if outputTotal > 0 {
parts = append(parts, fmt.Sprintf("%s:%s", T("usage_output"), formatLargeNumber(outputTotal)))
}
if cachedTotal > 0 {
parts = append(parts, fmt.Sprintf("%s:%s", T("usage_cached"), formatLargeNumber(cachedTotal)))
}
if reasoningTotal > 0 {
parts = append(parts, fmt.Sprintf("%s:%s", T("usage_reasoning"), formatLargeNumber(reasoningTotal)))
}
return fmt.Sprintf(" │ %s\n",
lipgloss.NewStyle().Foreground(colorMuted).Render(strings.Join(parts, " ")))
}
// renderLatencyBreakdown aggregates latency_ms from model details and displays avg/min/max.
func (m usageTabModel) renderLatencyBreakdown(modelStats map[string]any) string {
details, ok := modelStats["details"]
if !ok {
return ""
}
detailList, ok := details.([]any)
if !ok || len(detailList) == 0 {
return ""
}
var totalLatency int64
var count int
var minLatency, maxLatency int64
first := true
for _, d := range detailList {
dm, ok := d.(map[string]any)
if !ok {
continue
}
latencyMs := int64(getFloat(dm, "latency_ms"))
if latencyMs <= 0 {
continue
}
totalLatency += latencyMs
count++
if first {
minLatency = latencyMs
maxLatency = latencyMs
first = false
} else {
if latencyMs < minLatency {
minLatency = latencyMs
}
if latencyMs > maxLatency {
maxLatency = latencyMs
}
}
}
if count == 0 {
return ""
}
avgLatency := totalLatency / int64(count)
return fmt.Sprintf(" │ %s: avg %dms min %dms max %dms\n",
lipgloss.NewStyle().Foreground(colorMuted).Render(T("usage_time")),
avgLatency, minLatency, maxLatency)
}
// renderBarChart renders a simple ASCII horizontal bar chart.
func renderBarChart(data map[string]any, maxBarWidth int, barColor lipgloss.Color) string {
if maxBarWidth < 10 {
maxBarWidth = 10
}
// Sort keys
keys := make([]string, 0, len(data))
for k := range data {
keys = append(keys, k)
}
sort.Strings(keys)
// Find max value
maxVal := float64(0)
for _, k := range keys {
v := getFloat(data, k)
if v > maxVal {
maxVal = v
}
}
if maxVal == 0 {
return ""
}
barStyle := lipgloss.NewStyle().Foreground(barColor)
var sb strings.Builder
labelWidth := 12
barAvail := maxBarWidth - labelWidth - 12
if barAvail < 5 {
barAvail = 5
}
for _, k := range keys {
v := getFloat(data, k)
barLen := int(v / maxVal * float64(barAvail))
if barLen < 1 && v > 0 {
barLen = 1
}
bar := strings.Repeat("█", barLen)
label := k
if len(label) > labelWidth {
label = label[:labelWidth]
}
sb.WriteString(fmt.Sprintf(" %-*s %s %s\n",
labelWidth, label,
barStyle.Render(bar),
lipgloss.NewStyle().Foreground(colorMuted).Render(fmt.Sprintf("%.0f", v)),
))
}
return sb.String()
}
internal/tui/usage_tab_test.go
-134
View File
@@ -1,134 +0,0 @@
package tui
import (
"strings"
"testing"
)
func TestRenderLatencyBreakdown(t *testing.T) {
tests := []struct {
name string
modelStats map[string]any
wantEmpty bool
wantContains string
}{
{
name: "no details",
modelStats: map[string]any{},
wantEmpty: true,
},
{
name: "empty details",
modelStats: map[string]any{
"details": []any{},
},
wantEmpty: true,
},
{
name: "details with zero latency",
modelStats: map[string]any{
"details": []any{
map[string]any{
"latency_ms": float64(0),
},
},
},
wantEmpty: true,
},
{
name: "single request with latency",
modelStats: map[string]any{
"details": []any{
map[string]any{
"latency_ms": float64(1500),
},
},
},
wantEmpty: false,
wantContains: "avg 1500ms min 1500ms max 1500ms",
},
{
name: "multiple requests with varying latency",
modelStats: map[string]any{
"details": []any{
map[string]any{
"latency_ms": float64(100),
},
map[string]any{
"latency_ms": float64(200),
},
map[string]any{
"latency_ms": float64(300),
},
},
},
wantEmpty: false,
wantContains: "avg 200ms min 100ms max 300ms",
},
{
name: "mixed valid and invalid latency values",
modelStats: map[string]any{
"details": []any{
map[string]any{
"latency_ms": float64(500),
},
map[string]any{
"latency_ms": float64(0),
},
map[string]any{
"latency_ms": float64(1500),
},
},
},
wantEmpty: false,
wantContains: "avg 1000ms min 500ms max 1500ms",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
m := usageTabModel{}
result := m.renderLatencyBreakdown(tt.modelStats)
if tt.wantEmpty {
if result != "" {
t.Errorf("renderLatencyBreakdown() = %q, want empty string", result)
}
return
}
if result == "" {
t.Errorf("renderLatencyBreakdown() = empty, want non-empty string")
return
}
if tt.wantContains != "" && !strings.Contains(result, tt.wantContains) {
t.Errorf("renderLatencyBreakdown() = %q, want to contain %q", result, tt.wantContains)
}
})
}
}
func TestUsageTimeTranslations(t *testing.T) {
prevLocale := CurrentLocale()
t.Cleanup(func() {
SetLocale(prevLocale)
})
tests := []struct {
locale string
want string
}{
{locale: "en", want: "Time"},
{locale: "zh", want: "时间"},
}
for _, tt := range tests {
t.Run(tt.locale, func(t *testing.T) {
SetLocale(tt.locale)
if got := T("usage_time"); got != tt.want {
t.Fatalf("T(usage_time) = %q, want %q", got, tt.want)
}
})
}
}
internal/usage/logger_plugin.go
-464
View File
@@ -1,464 +0,0 @@
// Package usage provides usage tracking and logging functionality for the CLI Proxy API server.
// It includes plugins for monitoring API usage, token consumption, and other metrics
// to help with observability and billing purposes.
package usage
import (
"context"
"fmt"
"strings"
"sync"
"sync/atomic"
"time"
internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
)
var statisticsEnabled atomic.Bool
func init() {
statisticsEnabled.Store(true)
coreusage.RegisterPlugin(NewLoggerPlugin())
}
// LoggerPlugin collects in-memory request statistics for usage analysis.
// It implements coreusage.Plugin to receive usage records emitted by the runtime.
type LoggerPlugin struct {
stats *RequestStatistics
}
// NewLoggerPlugin constructs a new logger plugin instance.
//
// Returns:
// - *LoggerPlugin: A new logger plugin instance wired to the shared statistics store.
func NewLoggerPlugin() *LoggerPlugin { return &LoggerPlugin{stats: defaultRequestStatistics} }
// HandleUsage implements coreusage.Plugin.
// It updates the in-memory statistics store whenever a usage record is received.
//
// Parameters:
// - ctx: The context for the usage record
// - record: The usage record to aggregate
func (p *LoggerPlugin) HandleUsage(ctx context.Context, record coreusage.Record) {
if !statisticsEnabled.Load() {
return
}
if p == nil || p.stats == nil {
return
}
p.stats.Record(ctx, record)
}
// SetStatisticsEnabled toggles whether in-memory statistics are recorded.
func SetStatisticsEnabled(enabled bool) { statisticsEnabled.Store(enabled) }
// StatisticsEnabled reports the current recording state.
func StatisticsEnabled() bool { return statisticsEnabled.Load() }
// RequestStatistics maintains aggregated request metrics in memory.
type RequestStatistics struct {
mu sync.RWMutex
totalRequests int64
successCount int64
failureCount int64
totalTokens int64
apis map[string]*apiStats
requestsByDay map[string]int64
requestsByHour map[int]int64
tokensByDay map[string]int64
tokensByHour map[int]int64
}
// apiStats holds aggregated metrics for a single API key.
type apiStats struct {
TotalRequests int64
TotalTokens int64
Models map[string]*modelStats
}
// modelStats holds aggregated metrics for a specific model within an API.
type modelStats struct {
TotalRequests int64
TotalTokens int64
Details []RequestDetail
}
// RequestDetail stores the timestamp, latency, and token usage for a single request.
type RequestDetail struct {
Timestamp time.Time `json:"timestamp"`
LatencyMs int64 `json:"latency_ms"`
Source string `json:"source"`
AuthIndex string `json:"auth_index"`
Tokens TokenStats `json:"tokens"`
Failed bool `json:"failed"`
}
// TokenStats captures the token usage breakdown for a request.
type TokenStats struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
ReasoningTokens int64 `json:"reasoning_tokens"`
CachedTokens int64 `json:"cached_tokens"`
TotalTokens int64 `json:"total_tokens"`
}
// StatisticsSnapshot represents an immutable view of the aggregated metrics.
type StatisticsSnapshot struct {
TotalRequests int64 `json:"total_requests"`
SuccessCount int64 `json:"success_count"`
FailureCount int64 `json:"failure_count"`
TotalTokens int64 `json:"total_tokens"`
APIs map[string]APISnapshot `json:"apis"`
RequestsByDay map[string]int64 `json:"requests_by_day"`
RequestsByHour map[string]int64 `json:"requests_by_hour"`
TokensByDay map[string]int64 `json:"tokens_by_day"`
TokensByHour map[string]int64 `json:"tokens_by_hour"`
}
// APISnapshot summarises metrics for a single API key.
type APISnapshot struct {
TotalRequests int64 `json:"total_requests"`
TotalTokens int64 `json:"total_tokens"`
Models map[string]ModelSnapshot `json:"models"`
}
// ModelSnapshot summarises metrics for a specific model.
type ModelSnapshot struct {
TotalRequests int64 `json:"total_requests"`
TotalTokens int64 `json:"total_tokens"`
Details []RequestDetail `json:"details"`
}
var defaultRequestStatistics = NewRequestStatistics()
// GetRequestStatistics returns the shared statistics store.
func GetRequestStatistics() *RequestStatistics { return defaultRequestStatistics }
// NewRequestStatistics constructs an empty statistics store.
func NewRequestStatistics() *RequestStatistics {
return &RequestStatistics{
apis: make(map[string]*apiStats),
requestsByDay: make(map[string]int64),
requestsByHour: make(map[int]int64),
tokensByDay: make(map[string]int64),
tokensByHour: make(map[int]int64),
}
}
// Record ingests a new usage record and updates the aggregates.
func (s *RequestStatistics) Record(ctx context.Context, record coreusage.Record) {
if s == nil {
return
}
if !statisticsEnabled.Load() {
return
}
timestamp := record.RequestedAt
if timestamp.IsZero() {
timestamp = time.Now()
}
detail := normaliseDetail(record.Detail)
totalTokens := detail.TotalTokens
statsKey := record.APIKey
if statsKey == "" {
statsKey = resolveAPIIdentifier(ctx, record)
}
failed := record.Failed
if !failed {
failed = !resolveSuccess(ctx)
}
success := !failed
modelName := record.Model
if modelName == "" {
modelName = "unknown"
}
dayKey := timestamp.Format("2006-01-02")
hourKey := timestamp.Hour()
s.mu.Lock()
defer s.mu.Unlock()
s.totalRequests++
if success {
s.successCount++
} else {
s.failureCount++
}
s.totalTokens += totalTokens
stats, ok := s.apis[statsKey]
if !ok {
stats = &apiStats{Models: make(map[string]*modelStats)}
s.apis[statsKey] = stats
}
s.updateAPIStats(stats, modelName, RequestDetail{
Timestamp: timestamp,
LatencyMs: normaliseLatency(record.Latency),
Source: record.Source,
AuthIndex: record.AuthIndex,
Tokens: detail,
Failed: failed,
})
s.requestsByDay[dayKey]++
s.requestsByHour[hourKey]++
s.tokensByDay[dayKey] += totalTokens
s.tokensByHour[hourKey] += totalTokens
}
func (s *RequestStatistics) updateAPIStats(stats *apiStats, model string, detail RequestDetail) {
stats.TotalRequests++
stats.TotalTokens += detail.Tokens.TotalTokens
modelStatsValue, ok := stats.Models[model]
if !ok {
modelStatsValue = &modelStats{}
stats.Models[model] = modelStatsValue
}
modelStatsValue.TotalRequests++
modelStatsValue.TotalTokens += detail.Tokens.TotalTokens
modelStatsValue.Details = append(modelStatsValue.Details, detail)
}
// Snapshot returns a copy of the aggregated metrics for external consumption.
func (s *RequestStatistics) Snapshot() StatisticsSnapshot {
result := StatisticsSnapshot{}
if s == nil {
return result
}
s.mu.RLock()
defer s.mu.RUnlock()
result.TotalRequests = s.totalRequests
result.SuccessCount = s.successCount
result.FailureCount = s.failureCount
result.TotalTokens = s.totalTokens
result.APIs = make(map[string]APISnapshot, len(s.apis))
for apiName, stats := range s.apis {
apiSnapshot := APISnapshot{
TotalRequests: stats.TotalRequests,
TotalTokens: stats.TotalTokens,
Models: make(map[string]ModelSnapshot, len(stats.Models)),
}
for modelName, modelStatsValue := range stats.Models {
requestDetails := make([]RequestDetail, len(modelStatsValue.Details))
copy(requestDetails, modelStatsValue.Details)
apiSnapshot.Models[modelName] = ModelSnapshot{
TotalRequests: modelStatsValue.TotalRequests,
TotalTokens: modelStatsValue.TotalTokens,
Details: requestDetails,
}
}
result.APIs[apiName] = apiSnapshot
}
result.RequestsByDay = make(map[string]int64, len(s.requestsByDay))
for k, v := range s.requestsByDay {
result.RequestsByDay[k] = v
}
result.RequestsByHour = make(map[string]int64, len(s.requestsByHour))
for hour, v := range s.requestsByHour {
key := formatHour(hour)
result.RequestsByHour[key] = v
}
result.TokensByDay = make(map[string]int64, len(s.tokensByDay))
for k, v := range s.tokensByDay {
result.TokensByDay[k] = v
}
result.TokensByHour = make(map[string]int64, len(s.tokensByHour))
for hour, v := range s.tokensByHour {
key := formatHour(hour)
result.TokensByHour[key] = v
}
return result
}
type MergeResult struct {
Added int64 `json:"added"`
Skipped int64 `json:"skipped"`
}
// MergeSnapshot merges an exported statistics snapshot into the current store.
// Existing data is preserved and duplicate request details are skipped.
func (s *RequestStatistics) MergeSnapshot(snapshot StatisticsSnapshot) MergeResult {
result := MergeResult{}
if s == nil {
return result
}
s.mu.Lock()
defer s.mu.Unlock()
seen := make(map[string]struct{})
for apiName, stats := range s.apis {
if stats == nil {
continue
}
for modelName, modelStatsValue := range stats.Models {
if modelStatsValue == nil {
continue
}
for _, detail := range modelStatsValue.Details {
seen[dedupKey(apiName, modelName, detail)] = struct{}{}
}
}
}
for apiName, apiSnapshot := range snapshot.APIs {
apiName = strings.TrimSpace(apiName)
if apiName == "" {
continue
}
stats, ok := s.apis[apiName]
if !ok || stats == nil {
stats = &apiStats{Models: make(map[string]*modelStats)}
s.apis[apiName] = stats
} else if stats.Models == nil {
stats.Models = make(map[string]*modelStats)
}
for modelName, modelSnapshot := range apiSnapshot.Models {
modelName = strings.TrimSpace(modelName)
if modelName == "" {
modelName = "unknown"
}
for _, detail := range modelSnapshot.Details {
detail.Tokens = normaliseTokenStats(detail.Tokens)
if detail.LatencyMs < 0 {
detail.LatencyMs = 0
}
if detail.Timestamp.IsZero() {
detail.Timestamp = time.Now()
}
key := dedupKey(apiName, modelName, detail)
if _, exists := seen[key]; exists {
result.Skipped++
continue
}
seen[key] = struct{}{}
s.recordImported(apiName, modelName, stats, detail)
result.Added++
}
}
}
return result
}
func (s *RequestStatistics) recordImported(apiName, modelName string, stats *apiStats, detail RequestDetail) {
totalTokens := detail.Tokens.TotalTokens
if totalTokens < 0 {
totalTokens = 0
}
s.totalRequests++
if detail.Failed {
s.failureCount++
} else {
s.successCount++
}
s.totalTokens += totalTokens
s.updateAPIStats(stats, modelName, detail)
dayKey := detail.Timestamp.Format("2006-01-02")
hourKey := detail.Timestamp.Hour()
s.requestsByDay[dayKey]++
s.requestsByHour[hourKey]++
s.tokensByDay[dayKey] += totalTokens
s.tokensByHour[hourKey] += totalTokens
}
func dedupKey(apiName, modelName string, detail RequestDetail) string {
timestamp := detail.Timestamp.UTC().Format(time.RFC3339Nano)
tokens := normaliseTokenStats(detail.Tokens)
return fmt.Sprintf(
"%s|%s|%s|%s|%s|%t|%d|%d|%d|%d|%d",
apiName,
modelName,
timestamp,
detail.Source,
detail.AuthIndex,
detail.Failed,
tokens.InputTokens,
tokens.OutputTokens,
tokens.ReasoningTokens,
tokens.CachedTokens,
tokens.TotalTokens,
)
}
func resolveAPIIdentifier(ctx context.Context, record coreusage.Record) string {
if ctx != nil {
if endpoint := strings.TrimSpace(internallogging.GetEndpoint(ctx)); endpoint != "" {
return endpoint
}
}
if record.Provider != "" {
return record.Provider
}
return "unknown"
}
func resolveSuccess(ctx context.Context) bool {
status := internallogging.GetResponseStatus(ctx)
if status == 0 {
return true
}
return status < httpStatusBadRequest
}
const httpStatusBadRequest = 400
func normaliseDetail(detail coreusage.Detail) TokenStats {
tokens := TokenStats{
InputTokens: detail.InputTokens,
OutputTokens: detail.OutputTokens,
ReasoningTokens: detail.ReasoningTokens,
CachedTokens: detail.CachedTokens,
TotalTokens: detail.TotalTokens,
}
if tokens.TotalTokens == 0 {
tokens.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
if tokens.TotalTokens == 0 {
tokens.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens + detail.CachedTokens
}
return tokens
}
func normaliseTokenStats(tokens TokenStats) TokenStats {
if tokens.TotalTokens == 0 {
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
}
if tokens.TotalTokens == 0 {
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens + tokens.CachedTokens
}
return tokens
}
func normaliseLatency(latency time.Duration) int64 {
if latency <= 0 {
return 0
}
return latency.Milliseconds()
}
func formatHour(hour int) string {
if hour < 0 {
hour = 0
}
hour = hour % 24
return fmt.Sprintf("%02d", hour)
}
internal/usage/logger_plugin_test.go
-96
View File
@@ -1,96 +0,0 @@
package usage
import (
"context"
"testing"
"time"
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
)
func TestRequestStatisticsRecordIncludesLatency(t *testing.T) {
stats := NewRequestStatistics()
stats.Record(context.Background(), coreusage.Record{
APIKey: "test-key",
Model: "gpt-5.4",
RequestedAt: time.Date(2026, 3, 20, 12, 0, 0, 0, time.UTC),
Latency: 1500 * time.Millisecond,
Detail: coreusage.Detail{
InputTokens: 10,
OutputTokens: 20,
TotalTokens: 30,
},
})
snapshot := stats.Snapshot()
details := snapshot.APIs["test-key"].Models["gpt-5.4"].Details
if len(details) != 1 {
t.Fatalf("details len = %d, want 1", len(details))
}
if details[0].LatencyMs != 1500 {
t.Fatalf("latency_ms = %d, want 1500", details[0].LatencyMs)
}
}
func TestRequestStatisticsMergeSnapshotDedupIgnoresLatency(t *testing.T) {
stats := NewRequestStatistics()
timestamp := time.Date(2026, 3, 20, 12, 0, 0, 0, time.UTC)
first := StatisticsSnapshot{
APIs: map[string]APISnapshot{
"test-key": {
Models: map[string]ModelSnapshot{
"gpt-5.4": {
Details: []RequestDetail{{
Timestamp: timestamp,
LatencyMs: 0,
Source: "user@example.com",
AuthIndex: "0",
Tokens: TokenStats{
InputTokens: 10,
OutputTokens: 20,
TotalTokens: 30,
},
}},
},
},
},
},
}
second := StatisticsSnapshot{
APIs: map[string]APISnapshot{
"test-key": {
Models: map[string]ModelSnapshot{
"gpt-5.4": {
Details: []RequestDetail{{
Timestamp: timestamp,
LatencyMs: 2500,
Source: "user@example.com",
AuthIndex: "0",
Tokens: TokenStats{
InputTokens: 10,
OutputTokens: 20,
TotalTokens: 30,
},
}},
},
},
},
},
}
result := stats.MergeSnapshot(first)
if result.Added != 1 || result.Skipped != 0 {
t.Fatalf("first merge = %+v, want added=1 skipped=0", result)
}
result = stats.MergeSnapshot(second)
if result.Added != 0 || result.Skipped != 1 {
t.Fatalf("second merge = %+v, want added=0 skipped=1", result)
}
snapshot := stats.Snapshot()
details := snapshot.APIs["test-key"].Models["gpt-5.4"].Details
if len(details) != 1 {
t.Fatalf("details len = %d, want 1", len(details))
}
}
sdk/cliproxy/service.go
-1
View File
@@ -16,7 +16,6 @@ import (
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry" "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor" "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
"github.com/router-for-me/CLIProxyAPI/v6/internal/watcher" "github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
"github.com/router-for-me/CLIProxyAPI/v6/internal/wsrelay" "github.com/router-for-me/CLIProxyAPI/v6/internal/wsrelay"
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access" sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
test/usage_logging_test.go
+49 -24
View File
@@ -2,6 +2,7 @@ package test
import ( import (
"context" "context"
"encoding/json"
"fmt" "fmt"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
@@ -9,14 +10,14 @@ import (
"time" "time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"
runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor" runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor"
internalusage "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator" sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
) )
func TestGeminiExecutorRecordsSuccessfulZeroUsageInStatistics(t *testing.T) { func TestGeminiExecutorRecordsSuccessfulZeroUsageInQueue(t *testing.T) {
model := fmt.Sprintf("gemini-2.5-flash-zero-usage-%d", time.Now().UnixNano()) model := fmt.Sprintf("gemini-2.5-flash-zero-usage-%d", time.Now().UnixNano())
source := fmt.Sprintf("zero-usage-%d@example.com", time.Now().UnixNano()) source := fmt.Sprintf("zero-usage-%d@example.com", time.Now().UnixNano())
@@ -42,10 +43,15 @@ func TestGeminiExecutorRecordsSuccessfulZeroUsageInStatistics(t *testing.T) {
}, },
} }
prevStatsEnabled := internalusage.StatisticsEnabled() prevQueueEnabled := redisqueue.Enabled()
internalusage.SetStatisticsEnabled(true) prevUsageEnabled := redisqueue.UsageStatisticsEnabled()
redisqueue.SetEnabled(false)
redisqueue.SetEnabled(true)
redisqueue.SetUsageStatisticsEnabled(true)
t.Cleanup(func() { t.Cleanup(func() {
internalusage.SetStatisticsEnabled(prevStatsEnabled) redisqueue.SetEnabled(false)
redisqueue.SetEnabled(prevQueueEnabled)
redisqueue.SetUsageStatisticsEnabled(prevUsageEnabled)
}) })
_, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{
@@ -59,39 +65,58 @@ func TestGeminiExecutorRecordsSuccessfulZeroUsageInStatistics(t *testing.T) {
t.Fatalf("Execute error: %v", err) t.Fatalf("Execute error: %v", err)
} }
detail := waitForStatisticsDetail(t, "gemini", model, source) waitForQueuedUsageModelTotalTokens(t, "gemini", model, 0)
if detail.Failed {
t.Fatalf("detail failed = true, want false")
}
if detail.Tokens.TotalTokens != 0 {
t.Fatalf("total tokens = %d, want 0", detail.Tokens.TotalTokens)
}
} }
func waitForStatisticsDetail(t *testing.T, apiName, model, source string) internalusage.RequestDetail { func waitForQueuedUsageModelTotalTokens(t *testing.T, wantProvider, wantModel string, wantTokens int64) {
t.Helper() t.Helper()
deadline := time.Now().Add(2 * time.Second) deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) { for time.Now().Before(deadline) {
snapshot := internalusage.GetRequestStatistics().Snapshot() items := redisqueue.PopOldest(10)
apiSnapshot, ok := snapshot.APIs[apiName] for _, item := range items {
got, ok := parseQueuedUsagePayload(t, item)
if !ok { if !ok {
time.Sleep(10 * time.Millisecond)
continue continue
} }
modelSnapshot, ok := apiSnapshot.Models[model] if got.Provider != wantProvider || got.Model != wantModel {
if !ok {
time.Sleep(10 * time.Millisecond)
continue continue
} }
for _, detail := range modelSnapshot.Details { if got.Failed {
if detail.Source == source { t.Fatalf("payload failed = true, want false")
return detail
} }
if got.Tokens.TotalTokens != wantTokens {
t.Fatalf("payload total tokens = %d, want %d", got.Tokens.TotalTokens, wantTokens)
}
return
} }
time.Sleep(10 * time.Millisecond) time.Sleep(10 * time.Millisecond)
} }
t.Fatalf("timed out waiting for statistics detail for api=%q model=%q source=%q", apiName, model, source) t.Fatalf("timed out waiting for queued usage payload for provider=%q model=%q", wantProvider, wantModel)
return internalusage.RequestDetail{} }
type queuedUsagePayload struct {
Provider string `json:"provider"`
Model string `json:"model"`
Failed bool `json:"failed"`
Tokens struct {
TotalTokens int64 `json:"total_tokens"`
} `json:"tokens"`
}
func parseQueuedUsagePayload(t *testing.T, payload []byte) (queuedUsagePayload, bool) {
t.Helper()
var parsed queuedUsagePayload
if len(payload) == 0 {
return parsed, false
}
if err := json.Unmarshal(payload, &parsed); err != nil {
return parsed, false
}
if parsed.Provider == "" || parsed.Model == "" {
return parsed, false
}
return parsed, true
} }