Skip to content

Commit 6187919

Browse files
luispaterluispater
committed
feat: add support for recent request tracking in auth records
- Implemented `RecentRequestsSnapshot` in `Auth` to capture bucketed recent request data. - Added new fields and methods to `Auth` for tracking request success and failure counts over time. - Updated `/v0/management/auth-files` response to include recent request data for each auth record. - Introduced unit tests to validate request tracking and snapshot generation logic.
1 parent 4035abc commit 6187919

6 files changed

Lines changed: 294 additions & 2 deletions

File tree

‎internal/api/handlers/management/auth_files.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,7 @@ func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth) gin.H {
388388
"source": "memory",
389389
"size": int64(0),
390390
}
391+
entry["recent_requests"] = auth.RecentRequestsSnapshot(time.Now())
391392
if email := authEmail(auth); email != "" {
392393
entry["email"] = email
393394
}

‎internal/api/handlers/management/auth_files_recent_requests_test.go‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
package management
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"net/http"
7+
"net/http/httptest"
8+
"testing"
9+
10+
"github.com/gin-gonic/gin"
11+
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
12+
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
13+
)
14+
15+
func TestListAuthFiles_IncludesRecentRequestsBuckets(t *testing.T) {
16+
t.Setenv("MANAGEMENT_PASSWORD", "")
17+
gin.SetMode(gin.TestMode)
18+
19+
manager := coreauth.NewManager(nil, nil, nil)
20+
record := &coreauth.Auth{
21+
ID: "runtime-only-auth-1",
22+
Provider: "codex",
23+
Attributes: map[string]string{
24+
"runtime_only": "true",
25+
},
26+
Metadata: map[string]any{
27+
"type": "codex",
28+
},
29+
}
30+
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
31+
t.Fatalf("failed to register auth record: %v", errRegister)
32+
}
33+
34+
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
35+
h.tokenStore = &memoryAuthStore{}
36+
37+
rec := httptest.NewRecorder()
38+
ginCtx, _ := gin.CreateTestContext(rec)
39+
req := httptest.NewRequest(http.MethodGet, "/v0/management/auth-files", nil)
40+
ginCtx.Request = req
41+
42+
h.ListAuthFiles(ginCtx)
43+
44+
if rec.Code != http.StatusOK {
45+
t.Fatalf("expected list status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
46+
}
47+
48+
var payload map[string]any
49+
if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &payload); errUnmarshal != nil {
50+
t.Fatalf("failed to decode list payload: %v", errUnmarshal)
51+
}
52+
filesRaw, ok := payload["files"].([]any)
53+
if !ok {
54+
t.Fatalf("expected files array, payload: %#v", payload)
55+
}
56+
if len(filesRaw) != 1 {
57+
t.Fatalf("expected 1 auth entry, got %d", len(filesRaw))
58+
}
59+
60+
fileEntry, ok := filesRaw[0].(map[string]any)
61+
if !ok {
62+
t.Fatalf("expected file entry object, got %#v", filesRaw[0])
63+
}
64+
65+
recentRaw, ok := fileEntry["recent_requests"].([]any)
66+
if !ok {
67+
t.Fatalf("expected recent_requests array, got %#v", fileEntry["recent_requests"])
68+
}
69+
if len(recentRaw) != 20 {
70+
t.Fatalf("expected 20 recent_requests buckets, got %d", len(recentRaw))
71+
}
72+
for idx, item := range recentRaw {
73+
bucket, ok := item.(map[string]any)
74+
if !ok {
75+
t.Fatalf("expected bucket object at %d, got %#v", idx, item)
76+
}
77+
if _, ok := bucket["time"].(string); !ok {
78+
t.Fatalf("expected bucket time string at %d, got %#v", idx, bucket["time"])
79+
}
80+
if _, ok := bucket["success"].(float64); !ok {
81+
t.Fatalf("expected bucket success number at %d, got %#v", idx, bucket["success"])
82+
}
83+
if _, ok := bucket["failed"].(float64); !ok {
84+
t.Fatalf("expected bucket failed number at %d, got %#v", idx, bucket["failed"])
85+
}
86+
}
87+
}

‎sdk/cliproxy/auth/conductor.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2021,6 +2021,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) {
20212021
m.mu.Lock()
20222022
if auth, ok := m.auths[result.AuthID]; ok && auth != nil {
20232023
now := time.Now()
2024+
auth.recordRecentRequest(now, result.Success)
20242025

20252026
if result.Success {
20262027
if result.Model != "" {

‎sdk/cliproxy/auth/conductor_recent_requests_test.go‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package auth
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
)
8+
9+
func TestManagerMarkResultRecordsRecentRequests(t *testing.T) {
10+
mgr := NewManager(nil, nil, nil)
11+
auth := &Auth{
12+
ID: "auth-1",
13+
Provider: "antigravity",
14+
Attributes: map[string]string{
15+
"runtime_only": "true",
16+
},
17+
Metadata: map[string]any{
18+
"type": "antigravity",
19+
},
20+
}
21+
22+
if _, err := mgr.Register(WithSkipPersist(context.Background()), auth); err != nil {
23+
t.Fatalf("Register returned error: %v", err)
24+
}
25+
26+
mgr.MarkResult(context.Background(), Result{AuthID: "auth-1", Provider: "antigravity", Model: "gpt-5", Success: true})
27+
mgr.MarkResult(context.Background(), Result{AuthID: "auth-1", Provider: "antigravity", Model: "gpt-5", Success: false})
28+
29+
gotAuth, ok := mgr.GetByID("auth-1")
30+
if !ok || gotAuth == nil {
31+
t.Fatalf("GetByID returned ok=%v auth=%v", ok, gotAuth)
32+
}
33+
34+
snapshot := gotAuth.RecentRequestsSnapshot(time.Now())
35+
var successTotal int64
36+
var failedTotal int64
37+
for _, bucket := range snapshot {
38+
successTotal += bucket.Success
39+
failedTotal += bucket.Failed
40+
}
41+
if successTotal != 1 || failedTotal != 1 {
42+
t.Fatalf("totals = success=%d failed=%d, want 1/1", successTotal, failedTotal)
43+
}
44+
}

‎sdk/cliproxy/auth/types.go‎

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,29 @@ type Auth struct {
9292
// Runtime carries non-serialisable data used during execution (in-memory only).
9393
Runtime any `json:"-"`
9494

95-
indexAssigned bool `json:"-"`
95+
recentRequests recentRequestRing `json:"-"`
96+
indexAssigned bool `json:"-"`
97+
}
98+
99+
const (
100+
recentRequestBucketSeconds int64 = 10 * 60
101+
recentRequestBucketCount = 20
102+
)
103+
104+
type recentRequestBucket struct {
105+
bucketID int64
106+
success int64
107+
failed int64
108+
}
109+
110+
type recentRequestRing struct {
111+
buckets [recentRequestBucketCount]recentRequestBucket
112+
}
113+
114+
type RecentRequestBucket struct {
115+
Time string `json:"time"`
116+
Success int64 `json:"success"`
117+
Failed int64 `json:"failed"`
96118
}
97119

98120
// QuotaState contains limiter tracking data for a credential.
@@ -125,6 +147,70 @@ type ModelState struct {
125147
UpdatedAt time.Time `json:"updated_at"`
126148
}
127149

150+
func recentRequestBucketID(now time.Time) int64 {
151+
if now.IsZero() {
152+
return 0
153+
}
154+
return now.Unix() / recentRequestBucketSeconds
155+
}
156+
157+
func recentRequestBucketIndex(bucketID int64) int {
158+
mod := bucketID % int64(recentRequestBucketCount)
159+
if mod < 0 {
160+
mod += int64(recentRequestBucketCount)
161+
}
162+
return int(mod)
163+
}
164+
165+
func formatRecentRequestBucketLabel(bucketID int64) string {
166+
start := time.Unix(bucketID*recentRequestBucketSeconds, 0).In(time.Local)
167+
end := start.Add(10 * time.Minute)
168+
return start.Format("15:04") + "-" + end.Format("15:04")
169+
}
170+
171+
func (a *Auth) recordRecentRequest(now time.Time, success bool) {
172+
if a == nil {
173+
return
174+
}
175+
bucketID := recentRequestBucketID(now)
176+
idx := recentRequestBucketIndex(bucketID)
177+
bucket := &a.recentRequests.buckets[idx]
178+
if bucket.bucketID != bucketID {
179+
bucket.bucketID = bucketID
180+
bucket.success = 0
181+
bucket.failed = 0
182+
}
183+
if success {
184+
bucket.success++
185+
return
186+
}
187+
bucket.failed++
188+
}
189+
190+
func (a *Auth) RecentRequestsSnapshot(now time.Time) []RecentRequestBucket {
191+
out := make([]RecentRequestBucket, 0, recentRequestBucketCount)
192+
if a == nil {
193+
return out
194+
}
195+
196+
currentBucketID := recentRequestBucketID(now)
197+
for i := recentRequestBucketCount - 1; i >= 0; i-- {
198+
bucketID := currentBucketID - int64(i)
199+
idx := recentRequestBucketIndex(bucketID)
200+
bucket := a.recentRequests.buckets[idx]
201+
entry := RecentRequestBucket{
202+
Time: formatRecentRequestBucketLabel(bucketID),
203+
}
204+
if bucket.bucketID == bucketID {
205+
entry.Success = bucket.success
206+
entry.Failed = bucket.failed
207+
}
208+
out = append(out, entry)
209+
}
210+
211+
return out
212+
}
213+
128214
// Clone shallow copies the Auth structure, duplicating maps to avoid accidental mutation.
129215
func (a *Auth) Clone() *Auth {
130216
if a == nil {

‎sdk/cliproxy/auth/types_test.go‎

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
package auth
22

3-
import "testing"
3+
import (
4+
"strings"
5+
"testing"
6+
"time"
7+
)
48

59
func TestToolPrefixDisabled(t *testing.T) {
610
var a *Auth
@@ -96,3 +100,72 @@ func TestEnsureIndexUsesCredentialIdentity(t *testing.T) {
96100
t.Fatalf("duplicate config entries should be separated by source-derived seed, got %q", geminiIndex)
97101
}
98102
}
103+
104+
func TestRecentRequestsSnapshotEmptyReturnsTwentyBuckets(t *testing.T) {
105+
now := time.Unix(1_700_000_000, 0).In(time.Local)
106+
a := &Auth{}
107+
108+
got := a.RecentRequestsSnapshot(now)
109+
if len(got) != recentRequestBucketCount {
110+
t.Fatalf("len = %d, want %d", len(got), recentRequestBucketCount)
111+
}
112+
113+
currentBucketID := now.Unix() / recentRequestBucketSeconds
114+
baseBucketID := currentBucketID - int64(recentRequestBucketCount-1)
115+
for i, bucket := range got {
116+
if bucket.Success != 0 || bucket.Failed != 0 {
117+
t.Fatalf("bucket[%d] counts = %d/%d, want 0/0", i, bucket.Success, bucket.Failed)
118+
}
119+
if strings.TrimSpace(bucket.Time) == "" {
120+
t.Fatalf("bucket[%d] time label is empty", i)
121+
}
122+
expectedBucketID := baseBucketID + int64(i)
123+
start := time.Unix(expectedBucketID*recentRequestBucketSeconds, 0).In(time.Local)
124+
end := start.Add(10 * time.Minute)
125+
expected := start.Format("15:04") + "-" + end.Format("15:04")
126+
if bucket.Time != expected {
127+
t.Fatalf("bucket[%d] time = %q, want %q", i, bucket.Time, expected)
128+
}
129+
}
130+
}
131+
132+
func TestRecentRequestsSnapshotIncludesCounts(t *testing.T) {
133+
now := time.Unix(1_700_000_000, 0).In(time.Local)
134+
a := &Auth{}
135+
136+
a.recordRecentRequest(now, true)
137+
a.recordRecentRequest(now, false)
138+
139+
got := a.RecentRequestsSnapshot(now)
140+
if len(got) != recentRequestBucketCount {
141+
t.Fatalf("len = %d, want %d", len(got), recentRequestBucketCount)
142+
}
143+
144+
newest := got[len(got)-1]
145+
if newest.Success != 1 || newest.Failed != 1 {
146+
t.Fatalf("newest bucket = success=%d failed=%d, want 1/1", newest.Success, newest.Failed)
147+
}
148+
}
149+
150+
func TestRecentRequestsSnapshotBucketAdvanceMovesCounts(t *testing.T) {
151+
now := time.Unix(1_700_000_000, 0).In(time.Local)
152+
next := now.Add(10 * time.Minute)
153+
a := &Auth{}
154+
155+
a.recordRecentRequest(now, true)
156+
a.recordRecentRequest(next, false)
157+
158+
got := a.RecentRequestsSnapshot(next)
159+
if len(got) != recentRequestBucketCount {
160+
t.Fatalf("len = %d, want %d", len(got), recentRequestBucketCount)
161+
}
162+
163+
secondNewest := got[len(got)-2]
164+
newest := got[len(got)-1]
165+
if secondNewest.Success != 1 || secondNewest.Failed != 0 {
166+
t.Fatalf("second newest bucket = success=%d failed=%d, want 1/0", secondNewest.Success, secondNewest.Failed)
167+
}
168+
if newest.Success != 0 || newest.Failed != 1 {
169+
t.Fatalf("newest bucket = success=%d failed=%d, want 0/1", newest.Success, newest.Failed)
170+
}
171+
}

0 commit comments

Comments
 (0)