From 32a8102d711ce3b821e342a172c5c699cc0bd844 Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Tue, 14 Oct 2025 02:11:43 +0800
Subject: [PATCH] feat(usage): add support for tracking request source in usage
 records

- Introduced `Source` field to usage-related structs for better origin tracking.
- Updated `newUsageReporter` to resolve and populate the `Source` attribute.
- Implemented `resolveUsageSource` to determine source from auth metadata or API key.
---
 internal/runtime/executor/usage_helpers.go | 32 +++++++++++++++++++++-
 internal/usage/logger_plugin.go            |  7 ++++-
 sdk/cliproxy/usage/manager.go              |  1 +
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/internal/runtime/executor/usage_helpers.go b/internal/runtime/executor/usage_helpers.go
index 0bb3c682..d27fd524 100644
--- a/internal/runtime/executor/usage_helpers.go
+++ b/internal/runtime/executor/usage_helpers.go
@@ -4,10 +4,12 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"strings"
 	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
 	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
 	"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
 	"github.com/tidwall/gjson"
@@ -18,20 +20,23 @@ type usageReporter struct {
 	model       string
 	authID      string
 	apiKey      string
+	source      string
 	requestedAt time.Time
 	once        sync.Once
 }
 
 func newUsageReporter(ctx context.Context, provider, model string, auth *cliproxyauth.Auth) *usageReporter {
+	apiKey := apiKeyFromContext(ctx)
 	reporter := &usageReporter{
 		provider:    provider,
 		model:       model,
 		requestedAt: time.Now(),
+		apiKey:      apiKey,
+		source:      util.HideAPIKey(resolveUsageSource(auth, apiKey)),
 	}
 	if auth != nil {
 		reporter.authID = auth.ID
 	}
-	reporter.apiKey = apiKeyFromContext(ctx)
 	return reporter
 }
 
@@ -52,6 +57,7 @@ func (r *usageReporter) publish(ctx context.Context, detail usage.Detail) {
 		usage.PublishRecord(ctx, usage.Record{
 			Provider:    r.provider,
 			Model:       r.model,
+			Source:      r.source,
 			APIKey:      r.apiKey,
 			AuthID:      r.authID,
 			RequestedAt: r.requestedAt,
@@ -81,6 +87,30 @@ func apiKeyFromContext(ctx context.Context) string {
 	return ""
 }
 
+func resolveUsageSource(auth *cliproxyauth.Auth, ctxAPIKey string) string {
+	if auth != nil {
+		if _, value := auth.AccountInfo(); value != "" {
+			return strings.TrimSpace(value)
+		}
+		if auth.Metadata != nil {
+			if email, ok := auth.Metadata["email"].(string); ok {
+				if trimmed := strings.TrimSpace(email); trimmed != "" {
+					return trimmed
+				}
+			}
+		}
+		if auth.Attributes != nil {
+			if key := strings.TrimSpace(auth.Attributes["api_key"]); key != "" {
+				return key
+			}
+		}
+	}
+	if trimmed := strings.TrimSpace(ctxAPIKey); trimmed != "" {
+		return trimmed
+	}
+	return ""
+}
+
 func parseCodexUsage(data []byte) (usage.Detail, bool) {
 	usageNode := gjson.ParseBytes(data).Get("response.usage")
 	if !usageNode.Exists() {
diff --git a/internal/usage/logger_plugin.go b/internal/usage/logger_plugin.go
index a96dbc7c..78bce813 100644
--- a/internal/usage/logger_plugin.go
+++ b/internal/usage/logger_plugin.go
@@ -89,6 +89,7 @@ type modelStats struct {
 // RequestDetail stores the timestamp and token usage for a single request.
 type RequestDetail struct {
 	Timestamp time.Time  `json:"timestamp"`
+	Source    string     `json:"source"`
 	Tokens    TokenStats `json:"tokens"`
 }
 
@@ -188,7 +189,11 @@ func (s *RequestStatistics) Record(ctx context.Context, record coreusage.Record)
 		stats = &apiStats{Models: make(map[string]*modelStats)}
 		s.apis[statsKey] = stats
 	}
-	s.updateAPIStats(stats, modelName, RequestDetail{Timestamp: timestamp, Tokens: detail})
+	s.updateAPIStats(stats, modelName, RequestDetail{
+		Timestamp: timestamp,
+		Source:    record.Source,
+		Tokens:    detail,
+	})
 
 	s.requestsByDay[dayKey]++
 	s.requestsByHour[hourKey]++
diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go
index 48f0c003..25efbdc2 100644
--- a/sdk/cliproxy/usage/manager.go
+++ b/sdk/cliproxy/usage/manager.go
@@ -14,6 +14,7 @@ type Record struct {
 	Model       string
 	APIKey      string
 	AuthID      string
+	Source      string
 	RequestedAt time.Time
 	Detail      Detail
 }