diff --git a/core/http/app.go b/core/http/app.go index 99d11bd69..464e506db 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -28,6 +28,7 @@ import ( "github.com/mudler/LocalAI/core/services/monitoring" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/quantization" + "github.com/mudler/LocalAI/pkg/signals" "github.com/mudler/xlog" ) @@ -267,9 +268,12 @@ func API(application *application.Application) (*echo.Echo, error) { e.Static("/generated-videos", videoPath) } - // Initialize usage recording when auth DB is available + // Initialize usage recording when auth DB is available, and ensure the + // batcher drains its in-memory queue on graceful shutdown so the last + // few seconds of usage don't disappear when the process exits. if application.AuthDB() != nil { httpMiddleware.InitUsageRecorder(application.AuthDB()) + signals.RegisterGracefulTerminationHandler(httpMiddleware.ShutdownUsageRecorder) } // Auth is applied to _all_ endpoints. Filtering out endpoints to bypass is diff --git a/core/http/auth/db.go b/core/http/auth/db.go index d860e5068..5d94557ea 100644 --- a/core/http/auth/db.go +++ b/core/http/auth/db.go @@ -38,9 +38,15 @@ func InitDB(databaseURL string) (*gorm.DB, error) { } // Backfill: users created before the provider column existed have an empty - // provider — treat them as local accounts so the UI can identify them. + // provider - treat them as local accounts so the UI can identify them. db.Exec("UPDATE users SET provider = ? WHERE provider = '' OR provider IS NULL", ProviderLocal) + // Backfill: pre-feature usage_records have no source column. Classify them so the + // new per-source aggregators include them. + if err := BackfillUsageSource(db); err != nil { + return nil, fmt.Errorf("failed to backfill usage source: %w", err) + } + // Create composite index on users(provider, subject) for fast OAuth lookups if err := db.Exec("CREATE INDEX IF NOT EXISTS idx_users_provider_subject ON users(provider, subject)").Error; err != nil { // Ignore error on postgres if index already exists diff --git a/core/http/auth/middleware.go b/core/http/auth/middleware.go index 01ec33a68..c67954640 100644 --- a/core/http/auth/middleware.go +++ b/core/http/auth/middleware.go @@ -16,8 +16,10 @@ import ( ) const ( - contextKeyUser = "auth_user" - contextKeyRole = "auth_role" + contextKeyUser = "auth_user" + contextKeyRole = "auth_role" + contextKeyAPIKey = "auth_apikey" + contextKeySource = "auth_source" ) // Middleware returns an Echo middleware that handles authentication. @@ -75,6 +77,7 @@ func Middleware(db *gorm.DB, appConfig *config.ApplicationConfig) echo.Middlewar } c.Set(contextKeyUser, syntheticUser) c.Set(contextKeyRole, RoleAdmin) + c.Set(contextKeySource, UsageSourceLegacy) authenticated = true } } @@ -213,6 +216,20 @@ func GetUserRole(c echo.Context) string { return role } +// GetAPIKey returns the resolved API key from the echo context, or nil. +// Nil for session-cookie and legacy-env-key authentication. +func GetAPIKey(c echo.Context) *UserAPIKey { + k, _ := c.Get(contextKeyAPIKey).(*UserAPIKey) + return k +} + +// GetSource returns the request's authentication source: UsageSourceAPIKey, +// UsageSourceWeb, UsageSourceLegacy, or empty if no authentication was performed. +func GetSource(c echo.Context) string { + s, _ := c.Get(contextKeySource).(string) + return s +} + // RequireRouteFeature returns a global middleware that checks the user has access // to the feature required by the matched route. It uses the RouteFeatureRegistry // to look up the required feature for each route pattern + HTTP method. @@ -421,47 +438,67 @@ func RequireQuota(db *gorm.DB) echo.MiddlewareFunc { } // tryAuthenticate attempts to authenticate the request using the database. +// +// On success it returns the user and, as a side effect, sets the following +// values on the Echo context: +// - contextKeySource ("auth_source"): always set, one of UsageSourceWeb / +// UsageSourceAPIKey. UsageSourceLegacy is set elsewhere by the parent +// Middleware when a legacy env key matches. +// - contextKeyAPIKey ("auth_apikey"): set to the resolved *UserAPIKey for +// named-key branches (Bearer, x-api-key, xi-api-key, token cookie). +// - "_auth_session": session record, used by Middleware to drive cookie +// rotation. Only set on the session-cookie branch. +// +// contextKeyUser and contextKeyRole are populated by the parent Middleware +// after this function returns. func tryAuthenticate(c echo.Context, db *gorm.DB, appConfig *config.ApplicationConfig) *User { hmacSecret := appConfig.Auth.APIKeyHMACSecret - // a. Session cookie + // a. Session cookie -> web UI if cookie, err := c.Cookie(sessionCookie); err == nil && cookie.Value != "" { if user, session := ValidateSession(db, cookie.Value, hmacSecret); user != nil { // Store session for rotation check in middleware c.Set("_auth_session", session) + c.Set(contextKeySource, UsageSourceWeb) return user } } - // b. Authorization: Bearer token + // b. Authorization: Bearer authHeader := c.Request().Header.Get("Authorization") if strings.HasPrefix(authHeader, "Bearer ") { token := strings.TrimPrefix(authHeader, "Bearer ") - // Try as session ID first + // b1. Session token via Bearer -> still web UI if user, _ := ValidateSession(db, token, hmacSecret); user != nil { + c.Set(contextKeySource, UsageSourceWeb) return user } - // Try as user API key + // b2. Named API key if key, err := ValidateAPIKey(db, token, hmacSecret); err == nil { + c.Set(contextKeySource, UsageSourceAPIKey) + c.Set(contextKeyAPIKey, key) return &key.User } } - // c. x-api-key / xi-api-key headers + // c. x-api-key / xi-api-key -> named API key for _, header := range []string{"x-api-key", "xi-api-key"} { - if key := c.Request().Header.Get(header); key != "" { - if apiKey, err := ValidateAPIKey(db, key, hmacSecret); err == nil { + if k := c.Request().Header.Get(header); k != "" { + if apiKey, err := ValidateAPIKey(db, k, hmacSecret); err == nil { + c.Set(contextKeySource, UsageSourceAPIKey) + c.Set(contextKeyAPIKey, apiKey) return &apiKey.User } } } - // d. token cookie (legacy) + // d. token cookie -> named API key if cookie, err := c.Cookie("token"); err == nil && cookie.Value != "" { - // Try as user API key if key, err := ValidateAPIKey(db, cookie.Value, hmacSecret); err == nil { + c.Set(contextKeySource, UsageSourceAPIKey) + c.Set(contextKeyAPIKey, key) return &key.User } } diff --git a/core/http/auth/middleware_test.go b/core/http/auth/middleware_test.go index e7b4daa60..5137851e1 100644 --- a/core/http/auth/middleware_test.go +++ b/core/http/auth/middleware_test.go @@ -303,4 +303,122 @@ var _ = Describe("Auth Middleware", func() { } }) }) + + Describe("auth context plumbing for usage source", func() { + // probeApp builds a minimal echo app with the auth middleware and a single + // "/probe" route that captures the user, source, and apikey from context. + type probe struct { + user *auth.User + source string + key *auth.UserAPIKey + } + probeApp := func(db *gorm.DB, appConfig *config.ApplicationConfig, p *probe) *echo.Echo { + e := echo.New() + e.Use(auth.Middleware(db, appConfig)) + e.GET("/probe", func(c echo.Context) error { + p.user = auth.GetUser(c) + p.source = auth.GetSource(c) + p.key = auth.GetAPIKey(c) + return c.NoContent(http.StatusOK) + }) + return e + } + + It("session cookie sets source=web, apikey=nil", func() { + db := testDB() + appConfig := config.NewApplicationConfig() + user := createTestUser(db, "alice@example.com", auth.RoleUser, auth.ProviderLocal) + token := createTestSession(db, user.ID) + + var p probe + app := probeApp(db, appConfig, &p) + rec := doRequest(app, http.MethodGet, "/probe", withSessionCookie(token)) + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(p.user).ToNot(BeNil()) + Expect(p.user.ID).To(Equal(user.ID)) + Expect(p.source).To(Equal(auth.UsageSourceWeb)) + Expect(p.key).To(BeNil()) + }) + + It("Bearer session token sets source=web, apikey=nil", func() { + db := testDB() + appConfig := config.NewApplicationConfig() + user := createTestUser(db, "alice@example.com", auth.RoleUser, auth.ProviderLocal) + token := createTestSession(db, user.ID) + + var p probe + app := probeApp(db, appConfig, &p) + rec := doRequest(app, http.MethodGet, "/probe", withBearerToken(token)) + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(p.user).ToNot(BeNil()) + Expect(p.user.ID).To(Equal(user.ID)) + Expect(p.source).To(Equal(auth.UsageSourceWeb)) + Expect(p.key).To(BeNil()) + }) + + It("Bearer API key sets source=apikey and exposes the resolved *UserAPIKey", func() { + db := testDB() + appConfig := config.NewApplicationConfig() + user := createTestUser(db, "alice@example.com", auth.RoleUser, auth.ProviderLocal) + plaintext, key, err := auth.CreateAPIKey(db, user.ID, "ci", auth.RoleUser, appConfig.Auth.APIKeyHMACSecret, nil) + Expect(err).ToNot(HaveOccurred()) + + var p probe + app := probeApp(db, appConfig, &p) + rec := doRequest(app, http.MethodGet, "/probe", withBearerToken(plaintext)) + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(p.source).To(Equal(auth.UsageSourceAPIKey)) + Expect(p.key).ToNot(BeNil()) + Expect(p.key.ID).To(Equal(key.ID)) + }) + + It("x-api-key header sets source=apikey", func() { + db := testDB() + appConfig := config.NewApplicationConfig() + user := createTestUser(db, "alice@example.com", auth.RoleUser, auth.ProviderLocal) + plaintext, _, err := auth.CreateAPIKey(db, user.ID, "ci", auth.RoleUser, appConfig.Auth.APIKeyHMACSecret, nil) + Expect(err).ToNot(HaveOccurred()) + + var p probe + app := probeApp(db, appConfig, &p) + rec := doRequest(app, http.MethodGet, "/probe", withXApiKey(plaintext)) + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(p.source).To(Equal(auth.UsageSourceAPIKey)) + Expect(p.key).ToNot(BeNil()) + }) + + It("token cookie sets source=apikey", func() { + db := testDB() + appConfig := config.NewApplicationConfig() + user := createTestUser(db, "alice@example.com", auth.RoleUser, auth.ProviderLocal) + plaintext, _, err := auth.CreateAPIKey(db, user.ID, "ci", auth.RoleUser, appConfig.Auth.APIKeyHMACSecret, nil) + Expect(err).ToNot(HaveOccurred()) + + var p probe + app := probeApp(db, appConfig, &p) + rec := doRequest(app, http.MethodGet, "/probe", withTokenCookie(plaintext)) + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(p.source).To(Equal(auth.UsageSourceAPIKey)) + Expect(p.key).ToNot(BeNil()) + }) + + It("legacy env key sets source=legacy, apikey=nil", func() { + db := testDB() + appConfig := config.NewApplicationConfig() + appConfig.ApiKeys = []string{"legacy-secret"} + + var p probe + app := probeApp(db, appConfig, &p) + rec := doRequest(app, http.MethodGet, "/probe", withBearerToken("legacy-secret")) + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(p.source).To(Equal(auth.UsageSourceLegacy)) + Expect(p.key).To(BeNil()) + }) + }) }) diff --git a/core/http/auth/usage.go b/core/http/auth/usage.go index 31c3202b2..98c11093e 100644 --- a/core/http/auth/usage.go +++ b/core/http/auth/usage.go @@ -5,14 +5,31 @@ import ( "strings" "time" + "github.com/mudler/xlog" "gorm.io/gorm" ) +// Source classification for a UsageRecord. +const ( + UsageSourceAPIKey = "apikey" // request authenticated with a named UserAPIKey + UsageSourceWeb = "web" // request authenticated with a session cookie (web UI) + UsageSourceLegacy = "legacy" // request authenticated with an env-configured legacy key +) + // UsageRecord represents a single API request's token usage. type UsageRecord struct { - ID uint `gorm:"primaryKey;autoIncrement"` - UserID string `gorm:"size:36;index:idx_usage_user_time"` - UserName string `gorm:"size:255"` + ID uint `gorm:"primaryKey;autoIncrement"` + UserID string `gorm:"size:36;index:idx_usage_user_time"` + UserName string `gorm:"size:255"` + + // Source classifies how the request authenticated. One of UsageSource* constants. + // Empty for pre-feature rows until the InitDB backfill runs. + Source string `gorm:"size:16;index:idx_usage_source"` + // APIKeyID is the UserAPIKey.ID when Source == UsageSourceAPIKey. Nil otherwise. + APIKeyID *string `gorm:"size:36;index:idx_usage_apikey"` + // APIKeyName is a snapshot of UserAPIKey.Name at write time. Survives key deletion. + APIKeyName string `gorm:"size:255"` + Model string `gorm:"size:255;index"` Endpoint string `gorm:"size:255"` PromptTokens int64 @@ -30,9 +47,12 @@ func RecordUsage(db *gorm.DB, record *UsageRecord) error { // UsageBucket is an aggregated time bucket for the dashboard. type UsageBucket struct { Bucket string `json:"bucket"` - Model string `json:"model"` + Model string `json:"model,omitempty"` UserID string `json:"user_id,omitempty"` UserName string `json:"user_name,omitempty"` + Source string `json:"source,omitempty"` + APIKeyID string `json:"api_key_id,omitempty"` + APIKeyName string `json:"api_key_name,omitempty"` PromptTokens int64 `json:"prompt_tokens"` CompletionTokens int64 `json:"completion_tokens"` TotalTokens int64 `json:"total_tokens"` @@ -119,6 +139,28 @@ func GetUserUsage(db *gorm.DB, userID, period string) ([]UsageBucket, error) { return buckets, nil } +// BackfillUsageSource sets the Source column on pre-feature usage rows. +// Idempotent: only touches rows where source is NULL or empty. +// - rows whose user_id == "legacy-api-key" -> UsageSourceLegacy +// - everything else -> UsageSourceWeb +func BackfillUsageSource(db *gorm.DB) error { + // Legacy first (more specific predicate) + if err := db.Exec( + `UPDATE usage_records SET source = ? WHERE (source IS NULL OR source = '') AND user_id = ?`, + UsageSourceLegacy, "legacy-api-key", + ).Error; err != nil { + return fmt.Errorf("backfill legacy usage source: %w", err) + } + // Everything else -> web + if err := db.Exec( + `UPDATE usage_records SET source = ? WHERE (source IS NULL OR source = '')`, + UsageSourceWeb, + ).Error; err != nil { + return fmt.Errorf("backfill web usage source: %w", err) + } + return nil +} + // GetAllUsage returns aggregated usage for all users (admin). Optional userID filter. func GetAllUsage(db *gorm.DB, period, userID string) ([]UsageBucket, error) { sqlite := isSQLiteDB(db) @@ -149,3 +191,216 @@ func GetAllUsage(db *gorm.DB, period, userID string) ([]UsageBucket, error) { } return buckets, nil } + +// TotalsEntry is a token+request roll-up. +type TotalsEntry struct { + Tokens int64 `json:"tokens"` + Requests int64 `json:"requests"` +} + +// KeyTotal is the per-key roll-up returned by sources endpoints. +type KeyTotal struct { + APIKeyID string `json:"api_key_id"` + APIKeyName string `json:"api_key_name"` + Tokens int64 `json:"tokens"` + Requests int64 `json:"requests"` + LastUsed time.Time `json:"last_used"` +} + +// SourceTotals summarises a per-source breakdown. +type SourceTotals struct { + BySource map[string]TotalsEntry `json:"by_source"` + ByKey []KeyTotal `json:"by_key"` // server-sorted desc by tokens, capped + GrandTotal TotalsEntry `json:"grand_total"` +} + +const maxKeyTotals = 200 + +// GetUserUsageBySource returns per-source aggregated usage for one user. Legacy +// is excluded by design (visible to admins only via the admin variant). +func GetUserUsageBySource(db *gorm.DB, userID, period string) ([]UsageBucket, SourceTotals, error) { + sqlite := isSQLiteDB(db) + since, dateFmt := periodToWindow(period, sqlite) + bucketExpr := fmt.Sprintf("%s as bucket", dateFmt) + + query := db.Model(&UsageRecord{}). + Select(bucketExpr+", source, COALESCE(api_key_id, '') as api_key_id, api_key_name, "+ + "SUM(prompt_tokens) as prompt_tokens, "+ + "SUM(completion_tokens) as completion_tokens, "+ + "SUM(total_tokens) as total_tokens, "+ + "COUNT(*) as request_count"). + Where("user_id = ?", userID). + Where("source <> ?", UsageSourceLegacy). + Group("bucket, source, api_key_id, api_key_name"). + Order("bucket ASC") + + if !since.IsZero() { + query = query.Where("created_at >= ?", since) + } + + var buckets []UsageBucket + if err := query.Find(&buckets).Error; err != nil { + return nil, SourceTotals{}, err + } + + totals := computeSourceTotals(db, userID, "", since, false) + return buckets, totals, nil +} + +// computeSourceTotals rolls up by_source / by_key / grand_total. +// userID/apiKeyID are optional filters. includeLegacy controls whether the +// legacy bucket is exposed (admin-only). +func computeSourceTotals(db *gorm.DB, userID, apiKeyID string, since time.Time, includeLegacy bool) SourceTotals { + totals := SourceTotals{BySource: map[string]TotalsEntry{}} + + bySourceQ := db.Model(&UsageRecord{}). + Select("source, SUM(total_tokens) as tokens, COUNT(*) as requests"). + Group("source") + bySourceQ = applyFilters(bySourceQ, userID, apiKeyID, since, includeLegacy) + + var bySourceRows []struct { + Source string + Tokens int64 + Requests int64 + } + if err := bySourceQ.Scan(&bySourceRows).Error; err != nil { + xlog.Warn("computeSourceTotals: by-source Scan failed", "error", err) + return totals + } + for _, r := range bySourceRows { + totals.BySource[r.Source] = TotalsEntry{Tokens: r.Tokens, Requests: r.Requests} + totals.GrandTotal.Tokens += r.Tokens + totals.GrandTotal.Requests += r.Requests + } + + byKeyQ := db.Model(&UsageRecord{}). + Select("COALESCE(api_key_id, '') as api_key_id, api_key_name, "+ + "SUM(total_tokens) as tokens, COUNT(*) as requests, MAX(created_at) as last_used"). + Where("api_key_id IS NOT NULL AND api_key_id <> ''"). + Group("api_key_id, api_key_name"). + Order("tokens DESC"). + Limit(maxKeyTotals) + byKeyQ = applyFilters(byKeyQ, userID, apiKeyID, since, includeLegacy) + + // Iterate Rows() manually because MAX(created_at) is returned as a string by + // the SQLite driver, and Go's database/sql refuses to scan that into + // *time.Time. Postgres returns a proper timestamp. We accept both shapes + // via a Rows.Scan into a string column, then parse uniformly. + rows, err := byKeyQ.Rows() + if err != nil { + xlog.Warn("computeSourceTotals: by-key Rows() failed", "error", err) + } else { + defer func() { _ = rows.Close() }() + out := make([]KeyTotal, 0) + for rows.Next() { + var ( + apiKeyID, apiKeyName, lastUsedRaw string + tokens, requests int64 + ) + if scanErr := rows.Scan(&apiKeyID, &apiKeyName, &tokens, &requests, &lastUsedRaw); scanErr != nil { + continue + } + out = append(out, KeyTotal{ + APIKeyID: apiKeyID, + APIKeyName: apiKeyName, + Tokens: tokens, + Requests: requests, + LastUsed: parseLastUsedString(lastUsedRaw), + }) + } + if rerr := rows.Err(); rerr != nil { + xlog.Warn("computeSourceTotals: by-key rows iteration failed", "error", rerr) + } + totals.ByKey = out + } + + return totals +} + +// parseLastUsedString converts the textual MAX(created_at) value returned by +// SQLite (or any driver that surfaces the timestamp as a string) into a +// time.Time. Returns the zero time on parse failure. +func parseLastUsedString(s string) time.Time { + if s == "" { + return time.Time{} + } + // GORM's SQLite driver emits Go's default time formatting. Try the formats + // it commonly produces, falling back to RFC3339Nano. + layouts := []string{ + "2006-01-02 15:04:05.999999999 -0700 MST", + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02 15:04:05", + time.RFC3339Nano, + time.RFC3339, + } + for _, layout := range layouts { + if t, err := time.Parse(layout, s); err == nil { + return t + } + } + xlog.Warn("parseLastUsedString: unrecognised format", "value", s) + return time.Time{} +} + +// GetAllUsageBySource is the admin variant of GetUserUsageBySource. +// Optional filters: userID and apiKeyID. Legacy is included. +// truncated == true iff the per-key roll-up was capped at maxKeyTotals. +func GetAllUsageBySource(db *gorm.DB, period, userID, apiKeyID string) ([]UsageBucket, SourceTotals, bool, error) { + sqlite := isSQLiteDB(db) + since, dateFmt := periodToWindow(period, sqlite) + bucketExpr := fmt.Sprintf("%s as bucket", dateFmt) + + query := db.Model(&UsageRecord{}). + Select(bucketExpr+", source, COALESCE(api_key_id, '') as api_key_id, api_key_name, "+ + "user_id, user_name, "+ + "SUM(prompt_tokens) as prompt_tokens, "+ + "SUM(completion_tokens) as completion_tokens, "+ + "SUM(total_tokens) as total_tokens, "+ + "COUNT(*) as request_count"). + Group("bucket, source, api_key_id, api_key_name, user_id, user_name"). + Order("bucket ASC") + + query = applyFilters(query, userID, apiKeyID, since, true) + + var buckets []UsageBucket + if err := query.Find(&buckets).Error; err != nil { + return nil, SourceTotals{}, false, err + } + + totals := computeSourceTotals(db, userID, apiKeyID, since, true) + + // Count distinct api_key_ids matching the filters. If > maxKeyTotals, + // the by_key slice was capped and we signal truncation to the caller. + truncated := false + var distinct int64 + countQ := applyFilters( + db.Model(&UsageRecord{}). + Distinct("api_key_id"). + Where("api_key_id IS NOT NULL AND api_key_id <> ''"), + userID, apiKeyID, since, true, + ) + if err := countQ.Count(&distinct).Error; err != nil { + xlog.Warn("GetAllUsageBySource: distinct api_key_id count failed", "error", err) + } else { + truncated = distinct > maxKeyTotals + } + + return buckets, totals, truncated, nil +} + +func applyFilters(q *gorm.DB, userID, apiKeyID string, since time.Time, includeLegacy bool) *gorm.DB { + if userID != "" { + q = q.Where("user_id = ?", userID) + } + if apiKeyID != "" { + q = q.Where("api_key_id = ?", apiKeyID) + } + if !since.IsZero() { + q = q.Where("created_at >= ?", since) + } + if !includeLegacy { + q = q.Where("source <> ?", UsageSourceLegacy) + } + return q +} diff --git a/core/http/auth/usage_test.go b/core/http/auth/usage_test.go index 8782ac095..7b8a457a2 100644 --- a/core/http/auth/usage_test.go +++ b/core/http/auth/usage_test.go @@ -3,11 +3,13 @@ package auth_test import ( + "fmt" "time" "github.com/mudler/LocalAI/core/http/auth" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "gorm.io/gorm" ) var _ = Describe("Usage", func() { @@ -158,4 +160,194 @@ var _ = Describe("Usage", func() { } }) }) + + Describe("Usage source backfill", func() { + It("backfills 'web' for pre-feature rows", func() { + db := testDB() + + rawDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + _, err = rawDB.Exec( + `INSERT INTO usage_records (user_id, source, model, created_at, total_tokens, prompt_tokens, completion_tokens, duration) VALUES (?, '', ?, ?, 0, 0, 0, 0)`, + "user-x", "gpt-4", time.Now()) + Expect(err).ToNot(HaveOccurred()) + + Expect(auth.BackfillUsageSource(db)).To(Succeed()) + + var loaded auth.UsageRecord + Expect(db.Where("user_id = ?", "user-x").First(&loaded).Error).To(Succeed()) + Expect(loaded.Source).To(Equal(auth.UsageSourceWeb)) + }) + + It("backfills 'legacy' for pre-feature rows with legacy-api-key user_id", func() { + db := testDB() + + rawDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + _, err = rawDB.Exec( + `INSERT INTO usage_records (user_id, source, model, created_at, total_tokens, prompt_tokens, completion_tokens, duration) VALUES (?, '', ?, ?, 0, 0, 0, 0)`, + "legacy-api-key", "gpt-4", time.Now()) + Expect(err).ToNot(HaveOccurred()) + + Expect(auth.BackfillUsageSource(db)).To(Succeed()) + + var loaded auth.UsageRecord + Expect(db.Where("user_id = ?", "legacy-api-key").First(&loaded).Error).To(Succeed()) + Expect(loaded.Source).To(Equal(auth.UsageSourceLegacy)) + }) + + It("is idempotent on re-run", func() { + db := testDB() + Expect(auth.BackfillUsageSource(db)).To(Succeed()) + Expect(auth.BackfillUsageSource(db)).To(Succeed()) + }) + }) + + Describe("UsageRecord with source fields", func() { + It("persists Source, APIKeyID, APIKeyName", func() { + db := testDB() + keyID := "key-uuid-1" + record := &auth.UsageRecord{ + UserID: "user-1", + UserName: "Test User", + Source: auth.UsageSourceAPIKey, + APIKeyID: &keyID, + APIKeyName: "ci-runner", + Model: "gpt-4", + Endpoint: "/v1/chat/completions", + TotalTokens: 150, + CreatedAt: time.Now(), + } + Expect(auth.RecordUsage(db, record)).To(Succeed()) + + var loaded auth.UsageRecord + Expect(db.First(&loaded, record.ID).Error).To(Succeed()) + Expect(loaded.Source).To(Equal(auth.UsageSourceAPIKey)) + Expect(loaded.APIKeyID).ToNot(BeNil()) + Expect(*loaded.APIKeyID).To(Equal("key-uuid-1")) + Expect(loaded.APIKeyName).To(Equal("ci-runner")) + }) + + It("allows nil APIKeyID for web/legacy sources", func() { + db := testDB() + record := &auth.UsageRecord{ + UserID: "user-1", + Source: auth.UsageSourceWeb, + Model: "gpt-4", + CreatedAt: time.Now(), + } + Expect(auth.RecordUsage(db, record)).To(Succeed()) + + var loaded auth.UsageRecord + Expect(db.First(&loaded, record.ID).Error).To(Succeed()) + Expect(loaded.Source).To(Equal(auth.UsageSourceWeb)) + Expect(loaded.APIKeyID).To(BeNil()) + Expect(loaded.APIKeyName).To(BeEmpty()) + }) + }) + + Describe("GetUserUsageBySource", func() { + insert := func(db *gorm.DB, userID, source, keyID, keyName string, tokens int64, when time.Time) { + rec := &auth.UsageRecord{ + UserID: userID, + Source: source, + Model: "gpt-4", + TotalTokens: tokens, + CreatedAt: when, + } + if keyID != "" { + rec.APIKeyID = &keyID + rec.APIKeyName = keyName + } + Expect(auth.RecordUsage(db, rec)).To(Succeed()) + } + + It("returns only the caller's rows, never legacy", func() { + db := testDB() + now := time.Now() + insert(db, "alice", auth.UsageSourceAPIKey, "k1", "ci", 100, now) + insert(db, "alice", auth.UsageSourceWeb, "", "", 50, now) + insert(db, "alice", auth.UsageSourceLegacy, "", "", 30, now) + insert(db, "bob", auth.UsageSourceAPIKey, "k2", "bobk", 90, now) + + buckets, totals, err := auth.GetUserUsageBySource(db, "alice", "month") + Expect(err).ToNot(HaveOccurred()) + + for _, b := range buckets { + Expect(b.UserID).To(Or(BeEmpty(), Equal("alice"))) + Expect(b.Source).ToNot(Equal(auth.UsageSourceLegacy)) + } + + Expect(totals.GrandTotal.Tokens).To(Equal(int64(150))) + Expect(totals.BySource[auth.UsageSourceAPIKey].Tokens).To(Equal(int64(100))) + Expect(totals.BySource[auth.UsageSourceWeb].Tokens).To(Equal(int64(50))) + _, hasLegacy := totals.BySource[auth.UsageSourceLegacy] + Expect(hasLegacy).To(BeFalse()) + }) + + It("snapshots survive key deletion", func() { + db := testDB() + now := time.Now() + insert(db, "alice", auth.UsageSourceAPIKey, "deleted-key", "old-name", 42, now) + _, totals, err := auth.GetUserUsageBySource(db, "alice", "month") + Expect(err).ToNot(HaveOccurred()) + Expect(totals.ByKey).To(HaveLen(1)) + Expect(totals.ByKey[0].APIKeyName).To(Equal("old-name")) + Expect(totals.ByKey[0].APIKeyID).To(Equal("deleted-key")) + Expect(totals.ByKey[0].LastUsed).ToNot(BeZero()) + Expect(totals.ByKey[0].LastUsed).To(BeTemporally("~", now, 2*time.Second)) + }) + }) + + Describe("GetAllUsageBySource", func() { + insert := func(db *gorm.DB, userID, source, keyID string, tokens int64) { + rec := &auth.UsageRecord{ + UserID: userID, + Source: source, + Model: "gpt-4", + TotalTokens: tokens, + CreatedAt: time.Now(), + } + if keyID != "" { + rec.APIKeyID = &keyID + rec.APIKeyName = "name-" + keyID + } + Expect(auth.RecordUsage(db, rec)).To(Succeed()) + } + + It("includes legacy for admins", func() { + db := testDB() + insert(db, "alice", auth.UsageSourceAPIKey, "k1", 10) + insert(db, "legacy-api-key", auth.UsageSourceLegacy, "", 5) + + _, totals, _, err := auth.GetAllUsageBySource(db, "month", "", "") + Expect(err).ToNot(HaveOccurred()) + Expect(totals.BySource).To(HaveKey(auth.UsageSourceLegacy)) + Expect(totals.BySource[auth.UsageSourceLegacy].Tokens).To(Equal(int64(5))) + }) + + It("filters by user_id AND api_key_id", func() { + db := testDB() + insert(db, "alice", auth.UsageSourceAPIKey, "k1", 10) + insert(db, "alice", auth.UsageSourceAPIKey, "k2", 20) + insert(db, "bob", auth.UsageSourceAPIKey, "k3", 30) + + _, totals, _, err := auth.GetAllUsageBySource(db, "month", "alice", "k2") + Expect(err).ToNot(HaveOccurred()) + Expect(totals.GrandTotal.Tokens).To(Equal(int64(20))) + }) + + It("sets truncated=true when by_key exceeds the cap", func() { + db := testDB() + for i := 0; i < 210; i++ { + insert(db, "alice", auth.UsageSourceAPIKey, fmt.Sprintf("key-%03d", i), int64(210-i)) + } + + _, totals, truncated, err := auth.GetAllUsageBySource(db, "month", "", "") + Expect(err).ToNot(HaveOccurred()) + Expect(truncated).To(BeTrue()) + Expect(totals.ByKey).To(HaveLen(200)) + Expect(totals.ByKey[0].Tokens > totals.ByKey[199].Tokens).To(BeTrue()) + }) + }) }) diff --git a/core/http/middleware/usage.go b/core/http/middleware/usage.go index b82c1ee3f..6dc4699b8 100644 --- a/core/http/middleware/usage.go +++ b/core/http/middleware/usage.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/json" "sync" + "sync/atomic" "time" "github.com/labstack/echo/v4" @@ -14,18 +15,37 @@ import ( const ( usageFlushInterval = 5 * time.Second - usageMaxPending = 5000 + // usageMaxPending bounds the in-memory queue. Sized for bursty inference + // traffic on a self-hosted instance with a slow or unavailable DB. + usageMaxPending = 50000 ) // usageBatcher accumulates usage records and flushes them to the DB periodically. type usageBatcher struct { - mu sync.Mutex - pending []*auth.UsageRecord - db *gorm.DB + mu sync.Mutex + pending []*auth.UsageRecord + db *gorm.DB + stop chan struct{} + done chan struct{} + stopOnce sync.Once } +// droppedRecords counts records discarded because the in-memory queue was full. +// Used to rate-limit the warn log so a sustained outage doesn't flood it. +var droppedRecords atomic.Uint64 + func (b *usageBatcher) add(r *auth.UsageRecord) { b.mu.Lock() + if len(b.pending) >= usageMaxPending { + b.mu.Unlock() + // Rate-limit: one warn per 1024 drops keeps the log readable. + n := droppedRecords.Add(1) + if n&1023 == 1 { + xlog.Warn("usage batcher full, dropping record", + "cap", usageMaxPending, "total_dropped", n) + } + return + } b.pending = append(b.pending, r) b.mu.Unlock() } @@ -42,31 +62,102 @@ func (b *usageBatcher) flush() { if err := b.db.Create(&batch).Error; err != nil { xlog.Error("Failed to flush usage batch", "count", len(batch), "error", err) - // Re-queue failed records with a cap to avoid unbounded growth + // Cap-aware re-queue: prepend as much of the failed batch as fits + // alongside any records added concurrently with the failed write. b.mu.Lock() - if len(b.pending) < usageMaxPending { - b.pending = append(batch, b.pending...) + room := usageMaxPending - len(b.pending) + if room > 0 { + if room > len(batch) { + room = len(batch) + } + b.pending = append(batch[:room], b.pending...) } b.mu.Unlock() } } -var batcher *usageBatcher +func (b *usageBatcher) run() { + defer close(b.done) + ticker := time.NewTicker(usageFlushInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + b.flush() + case <-b.stop: + b.flush() // final drain + return + } + } +} + +func (b *usageBatcher) shutdown() { + b.stopOnce.Do(func() { + close(b.stop) + <-b.done + }) +} + +// The package-level batcher is guarded by batcherMu so Init / Shutdown cycles +// (the test pattern) don't race against UsageMiddleware reads. +var ( + batcherMu sync.RWMutex + batcher *usageBatcher +) + +func currentBatcher() *usageBatcher { + batcherMu.RLock() + defer batcherMu.RUnlock() + return batcher +} // InitUsageRecorder starts a background goroutine that periodically flushes -// accumulated usage records to the database. +// accumulated usage records to the database. Calling it more than once +// shuts down the previous batcher first so its goroutine doesn't leak. func InitUsageRecorder(db *gorm.DB) { if db == nil { return } - batcher = &usageBatcher{db: db} - go func() { - ticker := time.NewTicker(usageFlushInterval) - defer ticker.Stop() - for range ticker.C { - batcher.flush() - } - }() + + batcherMu.Lock() + old := batcher + batcher = nil + batcherMu.Unlock() + if old != nil { + old.shutdown() + } + + b := &usageBatcher{ + db: db, + stop: make(chan struct{}), + done: make(chan struct{}), + } + batcherMu.Lock() + batcher = b + batcherMu.Unlock() + + go b.run() +} + +// ShutdownUsageRecorder stops the background flusher and synchronously drains +// pending records once. Safe to call multiple times. Not yet wired into the +// application lifecycle; intended for graceful process exit and tests. +func ShutdownUsageRecorder() { + batcherMu.Lock() + b := batcher + batcher = nil + batcherMu.Unlock() + if b != nil { + b.shutdown() + } +} + +// FlushNow synchronously flushes any pending usage records. Intended for tests +// that need deterministic behaviour without waiting for the ticker. +func FlushNow() { + if b := currentBatcher(); b != nil { + b.flush() + } } // usageResponseBody is the minimal structure we need from the response JSON. @@ -84,7 +175,8 @@ type usageResponseBody struct { func UsageMiddleware(db *gorm.DB) echo.MiddlewareFunc { return func(next echo.HandlerFunc) echo.HandlerFunc { return func(c echo.Context) error { - if db == nil || batcher == nil { + b := currentBatcher() + if db == nil || b == nil { return next(c) } @@ -149,9 +241,17 @@ func UsageMiddleware(db *gorm.DB) echo.MiddlewareFunc { return handlerErr } + source := auth.GetSource(c) + if source == "" { + // Auth disabled or unrecognised path: classify as web so the row is still + // bucketable rather than silently dropped from per-source aggregates. + source = auth.UsageSourceWeb + } + record := &auth.UsageRecord{ UserID: user.ID, UserName: user.Name, + Source: source, Model: resp.Model, Endpoint: c.Request().URL.Path, PromptTokens: resp.Usage.PromptTokens, @@ -161,7 +261,13 @@ func UsageMiddleware(db *gorm.DB) echo.MiddlewareFunc { CreatedAt: startTime, } - batcher.add(record) + if key := auth.GetAPIKey(c); key != nil { + id := key.ID + record.APIKeyID = &id + record.APIKeyName = key.Name + } + + b.add(record) return handlerErr } diff --git a/core/http/middleware/usage_test.go b/core/http/middleware/usage_test.go new file mode 100644 index 000000000..7db03a6ba --- /dev/null +++ b/core/http/middleware/usage_test.go @@ -0,0 +1,140 @@ +//go:build auth + +package middleware_test + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/http/auth" + "github.com/mudler/LocalAI/core/http/middleware" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// testAuthDB returns a fresh in-memory SQLite auth DB. +func testAuthDB() *gorm.DB { + db, err := auth.InitDB(":memory:") + if err != nil { + panic(err) + } + return db +} + +var _ = Describe("UsageMiddleware", func() { + var ( + e *echo.Echo + db *gorm.DB + ) + + BeforeEach(func() { + db = testAuthDB() + e = echo.New() + middleware.InitUsageRecorder(db) + }) + + AfterEach(func() { + middleware.ShutdownUsageRecorder() + }) + + okHandler := func(c echo.Context) error { + body, _ := json.Marshal(map[string]any{ + "model": "gpt-4", + "usage": map[string]int{ + "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, + }, + }) + c.Response().Header().Set("Content-Type", "application/json") + c.Response().WriteHeader(http.StatusOK) + _, _ = c.Response().Write(body) + return nil + } + + // FlushNow drains pending records synchronously, replacing the 6s sleep + // that was previously needed to wait for the batcher's ticker. + flush := middleware.FlushNow + + It("records source=web when auth_source is web", func() { + e.POST("/v1/chat/completions", okHandler, func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + c.Set("auth_user", &auth.User{ID: "alice", Name: "Alice"}) + c.Set("auth_source", auth.UsageSourceWeb) + return next(c) + } + }, middleware.UsageMiddleware(db)) + + req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`))) + e.ServeHTTP(httptest.NewRecorder(), req) + flush() + + var rec auth.UsageRecord + Expect(db.Where("user_id = ?", "alice").First(&rec).Error).To(Succeed()) + Expect(rec.Source).To(Equal(auth.UsageSourceWeb)) + Expect(rec.APIKeyID).To(BeNil()) + Expect(rec.APIKeyName).To(BeEmpty()) + }) + + It("records source=apikey with snapshotted name when auth_apikey is set", func() { + e.POST("/v1/chat/completions", okHandler, func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + c.Set("auth_user", &auth.User{ID: "alice", Name: "Alice"}) + c.Set("auth_source", auth.UsageSourceAPIKey) + c.Set("auth_apikey", &auth.UserAPIKey{ID: "key-1", Name: "ci-runner"}) + return next(c) + } + }, middleware.UsageMiddleware(db)) + + req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`))) + e.ServeHTTP(httptest.NewRecorder(), req) + flush() + + var rec auth.UsageRecord + Expect(db.Where("user_id = ?", "alice").First(&rec).Error).To(Succeed()) + Expect(rec.Source).To(Equal(auth.UsageSourceAPIKey)) + Expect(rec.APIKeyID).ToNot(BeNil()) + Expect(*rec.APIKeyID).To(Equal("key-1")) + Expect(rec.APIKeyName).To(Equal("ci-runner")) + }) + + It("FlushNow drains pending records synchronously", func() { + e.POST("/v1/chat/completions", okHandler, func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + c.Set("auth_user", &auth.User{ID: "carol", Name: "Carol"}) + c.Set("auth_source", auth.UsageSourceWeb) + return next(c) + } + }, middleware.UsageMiddleware(db)) + + req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`))) + e.ServeHTTP(httptest.NewRecorder(), req) + + // No sleep: FlushNow should drain immediately. + middleware.FlushNow() + + var rec auth.UsageRecord + Expect(db.Where("user_id = ?", "carol").First(&rec).Error).To(Succeed()) + Expect(rec.Source).To(Equal(auth.UsageSourceWeb)) + }) + + It("falls back to source=web when auth_source is empty", func() { + e.POST("/v1/chat/completions", okHandler, func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + c.Set("auth_user", &auth.User{ID: "alice", Name: "Alice"}) + // no auth_source set + return next(c) + } + }, middleware.UsageMiddleware(db)) + + req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`))) + e.ServeHTTP(httptest.NewRecorder(), req) + flush() + + var rec auth.UsageRecord + Expect(db.Where("user_id = ?", "alice").First(&rec).Error).To(Succeed()) + Expect(rec.Source).To(Equal(auth.UsageSourceWeb)) + }) +}) diff --git a/core/http/react-ui/public/locales/en/admin.json b/core/http/react-ui/public/locales/en/admin.json index a76ef8df1..4b5ce0bb0 100644 --- a/core/http/react-ui/public/locales/en/admin.json +++ b/core/http/react-ui/public/locales/en/admin.json @@ -53,7 +53,29 @@ }, "usage": { "title": "Usage", - "subtitle": "API token usage statistics" + "subtitle": "API token usage statistics", + "sources": { + "tab": "Sources", + "mixTitle": "Source mix", + "ribbonAria": "{{apikey}}% API keys, {{web}}% Web UI, {{legacy}}% Legacy", + "topSources": "Top sources over time", + "searchPlaceholder": "Search by name or prefix", + "sortBy": "Sort", + "sortTokens": "Tokens", + "sortRequests": "Requests", + "sortLastUsed": "Last used", + "sortName": "Name", + "webUI": "Web UI", + "legacy": "Legacy", + "revoked": "revoked", + "filteredTo": "Filtered to: {{name}}", + "clearFilter": "Clear filter", + "other": "Other ({{count}})", + "noTrafficShort": "No requests in this period.", + "noKeysYet": "Once requests come in, you'll see them broken down here.", + "createKey": "Create your first API key", + "truncatedWarning": "Showing top 200 keys. Apply a filter to narrow further." + } }, "explorer": { "title": "Explorer", diff --git a/core/http/react-ui/src/pages/Usage.jsx b/core/http/react-ui/src/pages/Usage.jsx index 9d5b51f69..468c6dd07 100644 --- a/core/http/react-ui/src/pages/Usage.jsx +++ b/core/http/react-ui/src/pages/Usage.jsx @@ -4,6 +4,7 @@ import { useTranslation } from 'react-i18next' import { useAuth } from '../context/AuthContext' import { apiUrl } from '../utils/basePath' import LoadingSpinner from '../components/LoadingSpinner' +import SourcesTab from './Usage/SourcesTab' const PERIODS = [ { key: 'day', label: 'Day' }, @@ -724,23 +725,27 @@ export default function Usage() { {p.label} ))} +
+ {isAdmin && ( - <> - - - - > + )} +