Skip to content

Commit 32a8102

Browse files
committed
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.
1 parent f2710c0 commit 32a8102

File tree

3 files changed

+38
-2
lines changed

3 files changed

+38
-2
lines changed

internal/runtime/executor/usage_helpers.go

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@ import (
44
"bytes"
55
"context"
66
"fmt"
7+
"strings"
78
"sync"
89
"time"
910

1011
"github.com/gin-gonic/gin"
12+
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
1113
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
1214
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
1315
"github.com/tidwall/gjson"
@@ -18,20 +20,23 @@ type usageReporter struct {
1820
model string
1921
authID string
2022
apiKey string
23+
source string
2124
requestedAt time.Time
2225
once sync.Once
2326
}
2427

2528
func newUsageReporter(ctx context.Context, provider, model string, auth *cliproxyauth.Auth) *usageReporter {
29+
apiKey := apiKeyFromContext(ctx)
2630
reporter := &usageReporter{
2731
provider: provider,
2832
model: model,
2933
requestedAt: time.Now(),
34+
apiKey: apiKey,
35+
source: util.HideAPIKey(resolveUsageSource(auth, apiKey)),
3036
}
3137
if auth != nil {
3238
reporter.authID = auth.ID
3339
}
34-
reporter.apiKey = apiKeyFromContext(ctx)
3540
return reporter
3641
}
3742

@@ -52,6 +57,7 @@ func (r *usageReporter) publish(ctx context.Context, detail usage.Detail) {
5257
usage.PublishRecord(ctx, usage.Record{
5358
Provider: r.provider,
5459
Model: r.model,
60+
Source: r.source,
5561
APIKey: r.apiKey,
5662
AuthID: r.authID,
5763
RequestedAt: r.requestedAt,
@@ -81,6 +87,30 @@ func apiKeyFromContext(ctx context.Context) string {
8187
return ""
8288
}
8389

90+
func resolveUsageSource(auth *cliproxyauth.Auth, ctxAPIKey string) string {
91+
if auth != nil {
92+
if _, value := auth.AccountInfo(); value != "" {
93+
return strings.TrimSpace(value)
94+
}
95+
if auth.Metadata != nil {
96+
if email, ok := auth.Metadata["email"].(string); ok {
97+
if trimmed := strings.TrimSpace(email); trimmed != "" {
98+
return trimmed
99+
}
100+
}
101+
}
102+
if auth.Attributes != nil {
103+
if key := strings.TrimSpace(auth.Attributes["api_key"]); key != "" {
104+
return key
105+
}
106+
}
107+
}
108+
if trimmed := strings.TrimSpace(ctxAPIKey); trimmed != "" {
109+
return trimmed
110+
}
111+
return ""
112+
}
113+
84114
func parseCodexUsage(data []byte) (usage.Detail, bool) {
85115
usageNode := gjson.ParseBytes(data).Get("response.usage")
86116
if !usageNode.Exists() {

internal/usage/logger_plugin.go

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ type modelStats struct {
8989
// RequestDetail stores the timestamp and token usage for a single request.
9090
type RequestDetail struct {
9191
Timestamp time.Time `json:"timestamp"`
92+
Source string `json:"source"`
9293
Tokens TokenStats `json:"tokens"`
9394
}
9495

@@ -188,7 +189,11 @@ func (s *RequestStatistics) Record(ctx context.Context, record coreusage.Record)
188189
stats = &apiStats{Models: make(map[string]*modelStats)}
189190
s.apis[statsKey] = stats
190191
}
191-
s.updateAPIStats(stats, modelName, RequestDetail{Timestamp: timestamp, Tokens: detail})
192+
s.updateAPIStats(stats, modelName, RequestDetail{
193+
Timestamp: timestamp,
194+
Source: record.Source,
195+
Tokens: detail,
196+
})
192197

193198
s.requestsByDay[dayKey]++
194199
s.requestsByHour[hourKey]++

sdk/cliproxy/usage/manager.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ type Record struct {
1414
Model string
1515
APIKey string
1616
AuthID string
17+
Source string
1718
RequestedAt time.Time
1819
Detail Detail
1920
}

0 commit comments

Comments
 (0)