From 06db295d9e285dc35d6fbed5ab61cccd0c7f2cd6 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@localai.io>
Date: Wed, 20 May 2026 22:11:51 +0000
Subject: [PATCH] feat(usage): add Source, APIKeyID, APIKeyName columns to
 UsageRecord

Adds three additive columns plus UsageSource* constants. The columns
are auto-migrated by InitDB. APIKeyID is a nullable foreign reference
to UserAPIKey.ID; APIKeyName is snapshotted on each row so revoked
keys keep showing their name in history.

Refs: #9862
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
---
 core/http/auth/usage.go      | 27 ++++++++++++++++++----
 core/http/auth/usage_test.go | 43 ++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/core/http/auth/usage.go b/core/http/auth/usage.go
index 31c3202b2..ad56ee309 100644
--- a/core/http/auth/usage.go
+++ b/core/http/auth/usage.go
@@ -8,11 +8,27 @@ import (
 	"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 +46,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"`
diff --git a/core/http/auth/usage_test.go b/core/http/auth/usage_test.go
index 8782ac095..c3ad1e43f 100644
--- a/core/http/auth/usage_test.go
+++ b/core/http/auth/usage_test.go
@@ -158,4 +158,47 @@ var _ = Describe("Usage", func() {
 			}
 		})
 	})
+
+	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())
+		})
+	})
 })