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.
This commit is contained in:
Luis Pater
@@ -2021,6 +2021,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) {
|
||||
m.mu.Lock()
|
||||
if auth, ok := m.auths[result.AuthID]; ok && auth != nil {
|
||||
now := time.Now()
|
||||
auth.recordRecentRequest(now, result.Success)
|
||||
|
||||
if result.Success {
|
||||
if result.Model != "" {
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestManagerMarkResultRecordsRecentRequests(t *testing.T) {
|
||||
mgr := NewManager(nil, nil, nil)
|
||||
auth := &Auth{
|
||||
ID: "auth-1",
|
||||
Provider: "antigravity",
|
||||
Attributes: map[string]string{
|
||||
"runtime_only": "true",
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "antigravity",
|
||||
},
|
||||
}
|
||||
|
||||
if _, err := mgr.Register(WithSkipPersist(context.Background()), auth); err != nil {
|
||||
t.Fatalf("Register returned error: %v", err)
|
||||
}
|
||||
|
||||
mgr.MarkResult(context.Background(), Result{AuthID: "auth-1", Provider: "antigravity", Model: "gpt-5", Success: true})
|
||||
mgr.MarkResult(context.Background(), Result{AuthID: "auth-1", Provider: "antigravity", Model: "gpt-5", Success: false})
|
||||
|
||||
gotAuth, ok := mgr.GetByID("auth-1")
|
||||
if !ok || gotAuth == nil {
|
||||
t.Fatalf("GetByID returned ok=%v auth=%v", ok, gotAuth)
|
||||
}
|
||||
|
||||
snapshot := gotAuth.RecentRequestsSnapshot(time.Now())
|
||||
var successTotal int64
|
||||
var failedTotal int64
|
||||
for _, bucket := range snapshot {
|
||||
successTotal += bucket.Success
|
||||
failedTotal += bucket.Failed
|
||||
}
|
||||
if successTotal != 1 || failedTotal != 1 {
|
||||
t.Fatalf("totals = success=%d failed=%d, want 1/1", successTotal, failedTotal)
|
||||
}
|
||||
}
|
||||
@@ -92,7 +92,29 @@ type Auth struct {
|
||||
// Runtime carries non-serialisable data used during execution (in-memory only).
|
||||
Runtime any `json:"-"`
|
||||
|
||||
indexAssigned bool `json:"-"`
|
||||
recentRequests recentRequestRing `json:"-"`
|
||||
indexAssigned bool `json:"-"`
|
||||
}
|
||||
|
||||
const (
|
||||
recentRequestBucketSeconds int64 = 10 * 60
|
||||
recentRequestBucketCount = 20
|
||||
)
|
||||
|
||||
type recentRequestBucket struct {
|
||||
bucketID int64
|
||||
success int64
|
||||
failed int64
|
||||
}
|
||||
|
||||
type recentRequestRing struct {
|
||||
buckets [recentRequestBucketCount]recentRequestBucket
|
||||
}
|
||||
|
||||
type RecentRequestBucket struct {
|
||||
Time string `json:"time"`
|
||||
Success int64 `json:"success"`
|
||||
Failed int64 `json:"failed"`
|
||||
}
|
||||
|
||||
// QuotaState contains limiter tracking data for a credential.
|
||||
@@ -125,6 +147,70 @@ type ModelState struct {
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
func recentRequestBucketID(now time.Time) int64 {
|
||||
if now.IsZero() {
|
||||
return 0
|
||||
}
|
||||
return now.Unix() / recentRequestBucketSeconds
|
||||
}
|
||||
|
||||
func recentRequestBucketIndex(bucketID int64) int {
|
||||
mod := bucketID % int64(recentRequestBucketCount)
|
||||
if mod < 0 {
|
||||
mod += int64(recentRequestBucketCount)
|
||||
}
|
||||
return int(mod)
|
||||
}
|
||||
|
||||
func formatRecentRequestBucketLabel(bucketID int64) string {
|
||||
start := time.Unix(bucketID*recentRequestBucketSeconds, 0).In(time.Local)
|
||||
end := start.Add(10 * time.Minute)
|
||||
return start.Format("15:04") + "-" + end.Format("15:04")
|
||||
}
|
||||
|
||||
func (a *Auth) recordRecentRequest(now time.Time, success bool) {
|
||||
if a == nil {
|
||||
return
|
||||
}
|
||||
bucketID := recentRequestBucketID(now)
|
||||
idx := recentRequestBucketIndex(bucketID)
|
||||
bucket := &a.recentRequests.buckets[idx]
|
||||
if bucket.bucketID != bucketID {
|
||||
bucket.bucketID = bucketID
|
||||
bucket.success = 0
|
||||
bucket.failed = 0
|
||||
}
|
||||
if success {
|
||||
bucket.success++
|
||||
return
|
||||
}
|
||||
bucket.failed++
|
||||
}
|
||||
|
||||
func (a *Auth) RecentRequestsSnapshot(now time.Time) []RecentRequestBucket {
|
||||
out := make([]RecentRequestBucket, 0, recentRequestBucketCount)
|
||||
if a == nil {
|
||||
return out
|
||||
}
|
||||
|
||||
currentBucketID := recentRequestBucketID(now)
|
||||
for i := recentRequestBucketCount - 1; i >= 0; i-- {
|
||||
bucketID := currentBucketID - int64(i)
|
||||
idx := recentRequestBucketIndex(bucketID)
|
||||
bucket := a.recentRequests.buckets[idx]
|
||||
entry := RecentRequestBucket{
|
||||
Time: formatRecentRequestBucketLabel(bucketID),
|
||||
}
|
||||
if bucket.bucketID == bucketID {
|
||||
entry.Success = bucket.success
|
||||
entry.Failed = bucket.failed
|
||||
}
|
||||
out = append(out, entry)
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
// Clone shallow copies the Auth structure, duplicating maps to avoid accidental mutation.
|
||||
func (a *Auth) Clone() *Auth {
|
||||
if a == nil {
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
package auth
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestToolPrefixDisabled(t *testing.T) {
|
||||
var a *Auth
|
||||
@@ -96,3 +100,72 @@ func TestEnsureIndexUsesCredentialIdentity(t *testing.T) {
|
||||
t.Fatalf("duplicate config entries should be separated by source-derived seed, got %q", geminiIndex)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecentRequestsSnapshotEmptyReturnsTwentyBuckets(t *testing.T) {
|
||||
now := time.Unix(1_700_000_000, 0).In(time.Local)
|
||||
a := &Auth{}
|
||||
|
||||
got := a.RecentRequestsSnapshot(now)
|
||||
if len(got) != recentRequestBucketCount {
|
||||
t.Fatalf("len = %d, want %d", len(got), recentRequestBucketCount)
|
||||
}
|
||||
|
||||
currentBucketID := now.Unix() / recentRequestBucketSeconds
|
||||
baseBucketID := currentBucketID - int64(recentRequestBucketCount-1)
|
||||
for i, bucket := range got {
|
||||
if bucket.Success != 0 || bucket.Failed != 0 {
|
||||
t.Fatalf("bucket[%d] counts = %d/%d, want 0/0", i, bucket.Success, bucket.Failed)
|
||||
}
|
||||
if strings.TrimSpace(bucket.Time) == "" {
|
||||
t.Fatalf("bucket[%d] time label is empty", i)
|
||||
}
|
||||
expectedBucketID := baseBucketID + int64(i)
|
||||
start := time.Unix(expectedBucketID*recentRequestBucketSeconds, 0).In(time.Local)
|
||||
end := start.Add(10 * time.Minute)
|
||||
expected := start.Format("15:04") + "-" + end.Format("15:04")
|
||||
if bucket.Time != expected {
|
||||
t.Fatalf("bucket[%d] time = %q, want %q", i, bucket.Time, expected)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecentRequestsSnapshotIncludesCounts(t *testing.T) {
|
||||
now := time.Unix(1_700_000_000, 0).In(time.Local)
|
||||
a := &Auth{}
|
||||
|
||||
a.recordRecentRequest(now, true)
|
||||
a.recordRecentRequest(now, false)
|
||||
|
||||
got := a.RecentRequestsSnapshot(now)
|
||||
if len(got) != recentRequestBucketCount {
|
||||
t.Fatalf("len = %d, want %d", len(got), recentRequestBucketCount)
|
||||
}
|
||||
|
||||
newest := got[len(got)-1]
|
||||
if newest.Success != 1 || newest.Failed != 1 {
|
||||
t.Fatalf("newest bucket = success=%d failed=%d, want 1/1", newest.Success, newest.Failed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecentRequestsSnapshotBucketAdvanceMovesCounts(t *testing.T) {
|
||||
now := time.Unix(1_700_000_000, 0).In(time.Local)
|
||||
next := now.Add(10 * time.Minute)
|
||||
a := &Auth{}
|
||||
|
||||
a.recordRecentRequest(now, true)
|
||||
a.recordRecentRequest(next, false)
|
||||
|
||||
got := a.RecentRequestsSnapshot(next)
|
||||
if len(got) != recentRequestBucketCount {
|
||||
t.Fatalf("len = %d, want %d", len(got), recentRequestBucketCount)
|
||||
}
|
||||
|
||||
secondNewest := got[len(got)-2]
|
||||
newest := got[len(got)-1]
|
||||
if secondNewest.Success != 1 || secondNewest.Failed != 0 {
|
||||
t.Fatalf("second newest bucket = success=%d failed=%d, want 1/0", secondNewest.Success, secondNewest.Failed)
|
||||
}
|
||||
if newest.Success != 0 || newest.Failed != 1 {
|
||||
t.Fatalf("newest bucket = success=%d failed=%d, want 0/1", newest.Success, newest.Failed)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user