mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-09 01:46:48 +08:00
feat(auth): add error event publishing and Redis queue integration
- Introduced `publishErrorEvent` in `Manager` to publish error events to Redis. - Implemented error event structure to capture authentication errors with detailed metadata. - Added test cases for error event publishing, subscription, and Redis protocol handling. - Enhanced error and usage queue handling with `SubscribeErrors` and `EnqueueError`. Closes: #3701
This commit is contained in:
@@ -14,7 +14,10 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const redisUsageChannel = "usage"
|
||||
const (
|
||||
redisUsageChannel = "usage"
|
||||
redisErrorsChannel = "errors"
|
||||
)
|
||||
|
||||
type redisSubscriptionCommand struct {
|
||||
args []string
|
||||
@@ -150,15 +153,15 @@ func (s *Server) handleRedisConnection(conn net.Conn, reader *bufio.Reader) {
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !strings.EqualFold(channel, redisUsageChannel) {
|
||||
messages, unsubscribe, ok := subscribeRedisChannel(channel)
|
||||
if !ok {
|
||||
_ = writeRedisError(writer, fmt.Sprintf("ERR unsupported channel '%s'", channel))
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
messages, unsubscribe := redisqueue.SubscribeUsage()
|
||||
if errWrite := writeRedisPubSubSubscribe(writer, redisUsageChannel, 1); errWrite != nil {
|
||||
if errWrite := writeRedisPubSubSubscribe(writer, channel, 1); errWrite != nil {
|
||||
unsubscribe()
|
||||
log.Errorf("redis protocol subscribe response error: %v", errWrite)
|
||||
return
|
||||
@@ -167,7 +170,7 @@ func (s *Server) handleRedisConnection(conn net.Conn, reader *bufio.Reader) {
|
||||
unsubscribe()
|
||||
return
|
||||
}
|
||||
s.streamRedisUsageSubscription(reader, writer, messages, unsubscribe)
|
||||
s.streamRedisSubscription(reader, writer, channel, messages, unsubscribe)
|
||||
return
|
||||
case "LPOP", "RPOP":
|
||||
count, hasCount, ok := parsePopCount(args)
|
||||
@@ -185,7 +188,14 @@ func (s *Server) handleRedisConnection(conn net.Conn, reader *bufio.Reader) {
|
||||
}
|
||||
continue
|
||||
}
|
||||
items := redisqueue.PopOldest(count)
|
||||
items, ok := popRedisQueueItems(args[1], count)
|
||||
if !ok {
|
||||
_ = writeRedisError(writer, fmt.Sprintf("ERR unsupported channel '%s'", strings.TrimSpace(args[1])))
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if hasCount {
|
||||
_ = writeRedisArrayOfBulkStrings(writer, items)
|
||||
if !flush() {
|
||||
@@ -213,7 +223,29 @@ func (s *Server) handleRedisConnection(conn net.Conn, reader *bufio.Reader) {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) streamRedisUsageSubscription(reader *bufio.Reader, writer *bufio.Writer, messages <-chan []byte, unsubscribe func()) {
|
||||
func subscribeRedisChannel(channel string) (<-chan []byte, func(), bool) {
|
||||
switch strings.ToLower(strings.TrimSpace(channel)) {
|
||||
case redisUsageChannel:
|
||||
messages, unsubscribe := redisqueue.SubscribeUsage()
|
||||
return messages, unsubscribe, true
|
||||
case redisErrorsChannel:
|
||||
messages, unsubscribe := redisqueue.SubscribeErrors()
|
||||
return messages, unsubscribe, true
|
||||
default:
|
||||
return nil, nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func popRedisQueueItems(channel string, count int) ([][]byte, bool) {
|
||||
switch strings.ToLower(strings.TrimSpace(channel)) {
|
||||
case redisUsageChannel:
|
||||
return redisqueue.PopOldest(count), true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) streamRedisSubscription(reader *bufio.Reader, writer *bufio.Writer, channel string, messages <-chan []byte, unsubscribe func()) {
|
||||
if unsubscribe == nil {
|
||||
return
|
||||
}
|
||||
@@ -231,7 +263,7 @@ func (s *Server) streamRedisUsageSubscription(reader *bufio.Reader, writer *bufi
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if errWrite := writeRedisPubSubMessage(writer, redisUsageChannel, msg); errWrite != nil {
|
||||
if errWrite := writeRedisPubSubMessage(writer, channel, msg); errWrite != nil {
|
||||
log.Errorf("redis protocol publish message error: %v", errWrite)
|
||||
return
|
||||
}
|
||||
@@ -243,7 +275,7 @@ func (s *Server) streamRedisUsageSubscription(reader *bufio.Reader, writer *bufi
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
keepOpen := handleRedisSubscriptionCommand(writer, command)
|
||||
keepOpen := handleRedisSubscriptionCommand(writer, channel, command)
|
||||
if errFlush := writer.Flush(); errFlush != nil {
|
||||
log.Errorf("redis protocol flush error: %v", errFlush)
|
||||
return
|
||||
@@ -277,7 +309,7 @@ func readRedisSubscriptionCommands(reader *bufio.Reader, commands chan<- redisSu
|
||||
}
|
||||
}
|
||||
|
||||
func handleRedisSubscriptionCommand(writer *bufio.Writer, command redisSubscriptionCommand) bool {
|
||||
func handleRedisSubscriptionCommand(writer *bufio.Writer, channel string, command redisSubscriptionCommand) bool {
|
||||
if command.err != nil {
|
||||
_ = writeRedisError(writer, "ERR "+command.err.Error())
|
||||
return false
|
||||
@@ -297,7 +329,7 @@ func handleRedisSubscriptionCommand(writer *bufio.Writer, command redisSubscript
|
||||
_ = writeRedisPubSubPong(writer, payload)
|
||||
return true
|
||||
case "UNSUBSCRIBE":
|
||||
_ = writeRedisPubSubUnsubscribe(writer, redisUsageChannel, 0)
|
||||
_ = writeRedisPubSubUnsubscribe(writer, channel, 0)
|
||||
return false
|
||||
case "QUIT":
|
||||
_ = writeRedisSimpleString(writer, "OK")
|
||||
|
||||
@@ -359,6 +359,60 @@ func TestRedisProtocol_SUBSCRIBE_UsageSendsSupportRefresh(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisProtocol_SUBSCRIBE_ErrorsReceivesErrorEvents(t *testing.T) {
|
||||
const managementPassword = "test-management-password"
|
||||
|
||||
t.Setenv("MANAGEMENT_PASSWORD", managementPassword)
|
||||
redisqueue.SetEnabled(false)
|
||||
t.Cleanup(func() { redisqueue.SetEnabled(false) })
|
||||
|
||||
server := newTestServer(t)
|
||||
if !server.managementRoutesEnabled.Load() {
|
||||
t.Fatalf("expected managementRoutesEnabled to be true")
|
||||
}
|
||||
|
||||
addr, stop := startRedisMuxListener(t, server)
|
||||
t.Cleanup(stop)
|
||||
|
||||
conn, errDial := net.DialTimeout("tcp", addr, time.Second)
|
||||
if errDial != nil {
|
||||
t.Fatalf("failed to dial redis listener: %v", errDial)
|
||||
}
|
||||
t.Cleanup(func() { _ = conn.Close() })
|
||||
|
||||
reader := bufio.NewReader(conn)
|
||||
_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "AUTH", managementPassword); errWrite != nil {
|
||||
t.Fatalf("failed to write AUTH command: %v", errWrite)
|
||||
}
|
||||
if msg, errRead := readTestRESPSimpleString(reader); errRead != nil {
|
||||
t.Fatalf("failed to read AUTH response: %v", errRead)
|
||||
} else if msg != "OK" {
|
||||
t.Fatalf("unexpected AUTH response: %q", msg)
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "SUBSCRIBE", "errors"); errWrite != nil {
|
||||
t.Fatalf("failed to write SUBSCRIBE command: %v", errWrite)
|
||||
}
|
||||
channel, subscriptions, errSubscribe := readTestRESPPubSubSubscribe(reader)
|
||||
if errSubscribe != nil {
|
||||
t.Fatalf("failed to read subscribe response: %v", errSubscribe)
|
||||
}
|
||||
if channel != "errors" || subscriptions != 1 {
|
||||
t.Fatalf("unexpected subscribe response channel=%q subscriptions=%d", channel, subscriptions)
|
||||
}
|
||||
|
||||
redisqueue.EnqueueError([]byte(`{"auth_index":"auth-1","status_code":401}`))
|
||||
channel, payload, errMessage := readTestRESPPubSubMessage(reader)
|
||||
if errMessage != nil {
|
||||
t.Fatalf("failed to read error message: %v", errMessage)
|
||||
}
|
||||
if channel != "errors" || string(payload) != `{"auth_index":"auth-1","status_code":401}` {
|
||||
t.Fatalf("unexpected error message channel=%q payload=%q", channel, string(payload))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisProtocol_AUTH_And_PopContracts(t *testing.T) {
|
||||
const managementPassword = "test-management-password"
|
||||
|
||||
@@ -450,4 +504,13 @@ func TestRedisProtocol_AUTH_And_PopContracts(t *testing.T) {
|
||||
if len(emptyItems) != 0 {
|
||||
t.Fatalf("expected empty array for empty queue with count, got %#v", emptyItems)
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "RPOP", "errors", "2"); errWrite != nil {
|
||||
t.Fatalf("failed to write RPOP errors count command: %v", errWrite)
|
||||
}
|
||||
if msg, errRead := readTestRESPError(reader); errRead != nil {
|
||||
t.Fatalf("failed to read RPOP errors response: %v", errRead)
|
||||
} else if msg != "ERR unsupported channel 'errors'" {
|
||||
t.Fatalf("unexpected RPOP errors response: %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ const (
|
||||
defaultRetentionSeconds int64 = 60
|
||||
maxRetentionSeconds int64 = 3600
|
||||
usageSubscriberBuffer = 256
|
||||
errorSubscriberBuffer = 256
|
||||
|
||||
usageSupportRefreshPayload = `{"support_refresh":true}`
|
||||
usageRefreshPayload = `{"refresh":true}`
|
||||
@@ -32,6 +33,7 @@ var (
|
||||
enabled atomic.Bool
|
||||
retentionSeconds atomic.Int64
|
||||
global queue
|
||||
errorGlobal queue
|
||||
)
|
||||
|
||||
func init() {
|
||||
@@ -42,6 +44,7 @@ func SetEnabled(value bool) {
|
||||
enabled.Store(value)
|
||||
if !value {
|
||||
global.clear()
|
||||
errorGlobal.clear()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,6 +75,16 @@ func Enqueue(payload []byte) {
|
||||
global.enqueue(payload)
|
||||
}
|
||||
|
||||
func EnqueueError(payload []byte) {
|
||||
if !Enabled() {
|
||||
return
|
||||
}
|
||||
if len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
errorGlobal.publishToSubscribers(payload)
|
||||
}
|
||||
|
||||
func PopOldest(count int) [][]byte {
|
||||
if !Enabled() {
|
||||
return nil
|
||||
@@ -83,7 +96,11 @@ func PopOldest(count int) [][]byte {
|
||||
}
|
||||
|
||||
func SubscribeUsage() (<-chan []byte, func()) {
|
||||
return global.subscribeUsage()
|
||||
return global.subscribe(usageSubscriberBuffer, []byte(usageSupportRefreshPayload))
|
||||
}
|
||||
|
||||
func SubscribeErrors() (<-chan []byte, func()) {
|
||||
return errorGlobal.subscribe(errorSubscriberBuffer, nil)
|
||||
}
|
||||
|
||||
func NotifyUsageRefresh() {
|
||||
@@ -142,9 +159,11 @@ func (q *queue) publishToSubscribers(payload []byte) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (q *queue) subscribeUsage() (<-chan []byte, func()) {
|
||||
subscriber := make(chan []byte, usageSubscriberBuffer)
|
||||
subscriber <- []byte(usageSupportRefreshPayload)
|
||||
func (q *queue) subscribe(buffer int, initialPayload []byte) (<-chan []byte, func()) {
|
||||
subscriber := make(chan []byte, buffer)
|
||||
if len(initialPayload) > 0 {
|
||||
subscriber <- append([]byte(nil), initialPayload...)
|
||||
}
|
||||
|
||||
q.mu.Lock()
|
||||
if q.subscribers == nil {
|
||||
@@ -158,13 +177,13 @@ func (q *queue) subscribeUsage() (<-chan []byte, func()) {
|
||||
var once sync.Once
|
||||
unsubscribe := func() {
|
||||
once.Do(func() {
|
||||
q.unsubscribeUsage(id)
|
||||
q.unsubscribe(id)
|
||||
})
|
||||
}
|
||||
return subscriber, unsubscribe
|
||||
}
|
||||
|
||||
func (q *queue) unsubscribeUsage(id uint64) {
|
||||
func (q *queue) unsubscribe(id uint64) {
|
||||
q.mu.Lock()
|
||||
subscriber, ok := q.subscribers[id]
|
||||
if ok {
|
||||
|
||||
@@ -39,6 +39,8 @@ func TestSetEnabledFalseClosesUsageSubscribers(t *testing.T) {
|
||||
withEnabledQueue(t, func() {
|
||||
subscriber, unsubscribe := SubscribeUsage()
|
||||
defer unsubscribe()
|
||||
errorSubscriber, unsubscribeErrors := SubscribeErrors()
|
||||
defer unsubscribeErrors()
|
||||
|
||||
requireUsageSubscriberPayload(t, subscriber, usageSupportRefreshPayload)
|
||||
|
||||
@@ -52,6 +54,30 @@ func TestSetEnabledFalseClosesUsageSubscribers(t *testing.T) {
|
||||
case <-time.After(time.Second):
|
||||
t.Fatalf("timeout waiting for subscriber close")
|
||||
}
|
||||
|
||||
select {
|
||||
case _, ok := <-errorSubscriber:
|
||||
if ok {
|
||||
t.Fatalf("error subscriber channel remained open after SetEnabled(false)")
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatalf("timeout waiting for error subscriber close")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestEnqueueErrorBroadcastsToErrorSubscribersAndDiscardsWithoutSubscribers(t *testing.T) {
|
||||
withEnabledQueue(t, func() {
|
||||
subscriber, unsubscribe := SubscribeErrors()
|
||||
defer unsubscribe()
|
||||
|
||||
EnqueueError([]byte("error-record"))
|
||||
requireUsageSubscriberPayload(t, subscriber, "error-record")
|
||||
|
||||
unsubscribe()
|
||||
|
||||
EnqueueError([]byte("discarded-error"))
|
||||
requireErrorQueueEmpty(t)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -59,12 +85,20 @@ func TestNotifyUsageRefreshBroadcastsOnlyToUsageSubscribers(t *testing.T) {
|
||||
withEnabledQueue(t, func() {
|
||||
subscriber, unsubscribe := SubscribeUsage()
|
||||
defer unsubscribe()
|
||||
errorSubscriber, unsubscribeErrors := SubscribeErrors()
|
||||
defer unsubscribeErrors()
|
||||
|
||||
requireUsageSubscriberPayload(t, subscriber, usageSupportRefreshPayload)
|
||||
|
||||
NotifyUsageRefresh()
|
||||
requireUsageSubscriberPayload(t, subscriber, usageRefreshPayload)
|
||||
|
||||
select {
|
||||
case got := <-errorSubscriber:
|
||||
t.Fatalf("error subscriber received usage refresh payload %q", string(got))
|
||||
default:
|
||||
}
|
||||
|
||||
unsubscribe()
|
||||
NotifyUsageRefresh()
|
||||
if items := PopOldest(1); len(items) != 0 {
|
||||
@@ -88,3 +122,14 @@ func requireUsageSubscriberPayload(t *testing.T, subscriber <-chan []byte, want
|
||||
t.Fatalf("timeout waiting for subscriber payload %q", want)
|
||||
}
|
||||
}
|
||||
|
||||
func requireErrorQueueEmpty(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
errorGlobal.mu.Lock()
|
||||
defer errorGlobal.mu.Unlock()
|
||||
|
||||
if len(errorGlobal.items)-errorGlobal.head != 0 {
|
||||
t.Fatalf("error queue retained %d item(s), want none", len(errorGlobal.items)-errorGlobal.head)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2523,6 +2523,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) {
|
||||
}
|
||||
|
||||
m.hook.OnResult(ctx, result)
|
||||
m.publishErrorEvent(result, authSnapshot)
|
||||
}
|
||||
|
||||
func ensureModelState(auth *Auth, model string) *ModelState {
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"
|
||||
)
|
||||
|
||||
type errorEvent struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
AuthID string `json:"auth_id,omitempty"`
|
||||
AuthIndex string `json:"auth_index"`
|
||||
StatusCode int `json:"status_code"`
|
||||
Body string `json:"body"`
|
||||
Code string `json:"code,omitempty"`
|
||||
Retryable bool `json:"retryable,omitempty"`
|
||||
AuthStatus errorEventAuthStatus `json:"auth_status"`
|
||||
}
|
||||
|
||||
type errorEventAuthStatus struct {
|
||||
Status Status `json:"status"`
|
||||
StatusMessage string `json:"status_message,omitempty"`
|
||||
Disabled bool `json:"disabled"`
|
||||
Unavailable bool `json:"unavailable"`
|
||||
NextRetryAfter *time.Time `json:"next_retry_after,omitempty"`
|
||||
Quota *errorEventQuotaStatus `json:"quota,omitempty"`
|
||||
Model *errorEventModelStatus `json:"model,omitempty"`
|
||||
}
|
||||
|
||||
type errorEventQuotaStatus struct {
|
||||
Exceeded bool `json:"exceeded"`
|
||||
Reason string `json:"reason,omitempty"`
|
||||
NextRecoverAt *time.Time `json:"next_recover_at,omitempty"`
|
||||
BackoffLevel int `json:"backoff_level,omitempty"`
|
||||
}
|
||||
|
||||
type errorEventModelStatus struct {
|
||||
Name string `json:"name"`
|
||||
Status Status `json:"status"`
|
||||
StatusMessage string `json:"status_message,omitempty"`
|
||||
Unavailable bool `json:"unavailable"`
|
||||
NextRetryAfter *time.Time `json:"next_retry_after,omitempty"`
|
||||
Quota *errorEventQuotaStatus `json:"quota,omitempty"`
|
||||
}
|
||||
|
||||
func (m *Manager) publishErrorEvent(result Result, authSnapshot *Auth) {
|
||||
if m == nil || result.Success || authSnapshot == nil || m.HomeEnabled() {
|
||||
return
|
||||
}
|
||||
payload, ok := buildErrorEventPayload(result, authSnapshot)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
redisqueue.EnqueueError(payload)
|
||||
}
|
||||
|
||||
func buildErrorEventPayload(result Result, authSnapshot *Auth) ([]byte, bool) {
|
||||
if authSnapshot == nil || result.Success {
|
||||
return nil, false
|
||||
}
|
||||
authSnapshot.EnsureIndex()
|
||||
event := errorEvent{
|
||||
Timestamp: time.Now(),
|
||||
Provider: strings.TrimSpace(result.Provider),
|
||||
Model: strings.TrimSpace(result.Model),
|
||||
AuthID: strings.TrimSpace(result.AuthID),
|
||||
AuthIndex: strings.TrimSpace(authSnapshot.Index),
|
||||
StatusCode: errorEventStatusCode(result.Error),
|
||||
Body: errorEventBody(result.Error),
|
||||
AuthStatus: buildErrorEventAuthStatus(result.Model, authSnapshot),
|
||||
}
|
||||
if result.Error != nil {
|
||||
event.Code = strings.TrimSpace(result.Error.Code)
|
||||
event.Retryable = result.Error.Retryable
|
||||
}
|
||||
payload, errMarshal := json.Marshal(event)
|
||||
if errMarshal != nil {
|
||||
return nil, false
|
||||
}
|
||||
return payload, true
|
||||
}
|
||||
|
||||
func buildErrorEventAuthStatus(model string, authSnapshot *Auth) errorEventAuthStatus {
|
||||
status := errorEventAuthStatus{
|
||||
Status: authSnapshot.Status,
|
||||
StatusMessage: strings.TrimSpace(authSnapshot.StatusMessage),
|
||||
Disabled: authSnapshot.Disabled,
|
||||
Unavailable: authSnapshot.Unavailable,
|
||||
NextRetryAfter: timePtrIfSet(authSnapshot.NextRetryAfter),
|
||||
Quota: errorEventQuotaStatusFrom(authSnapshot.Quota),
|
||||
}
|
||||
if modelState := errorEventModelStatusFrom(model, authSnapshot); modelState != nil {
|
||||
status.Model = modelState
|
||||
}
|
||||
return status
|
||||
}
|
||||
|
||||
func errorEventModelStatusFrom(model string, authSnapshot *Auth) *errorEventModelStatus {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" || authSnapshot == nil || authSnapshot.ModelStates == nil {
|
||||
return nil
|
||||
}
|
||||
state := authSnapshot.ModelStates[model]
|
||||
if state == nil {
|
||||
return nil
|
||||
}
|
||||
return &errorEventModelStatus{
|
||||
Name: model,
|
||||
Status: state.Status,
|
||||
StatusMessage: strings.TrimSpace(state.StatusMessage),
|
||||
Unavailable: state.Unavailable,
|
||||
NextRetryAfter: timePtrIfSet(state.NextRetryAfter),
|
||||
Quota: errorEventQuotaStatusFrom(state.Quota),
|
||||
}
|
||||
}
|
||||
|
||||
func errorEventQuotaStatusFrom(quota QuotaState) *errorEventQuotaStatus {
|
||||
if !quota.Exceeded && strings.TrimSpace(quota.Reason) == "" && quota.NextRecoverAt.IsZero() && quota.BackoffLevel == 0 {
|
||||
return nil
|
||||
}
|
||||
return &errorEventQuotaStatus{
|
||||
Exceeded: quota.Exceeded,
|
||||
Reason: strings.TrimSpace(quota.Reason),
|
||||
NextRecoverAt: timePtrIfSet(quota.NextRecoverAt),
|
||||
BackoffLevel: quota.BackoffLevel,
|
||||
}
|
||||
}
|
||||
|
||||
func errorEventStatusCode(err *Error) int {
|
||||
if err != nil && err.HTTPStatus > 0 {
|
||||
return err.HTTPStatus
|
||||
}
|
||||
return 500
|
||||
}
|
||||
|
||||
func errorEventBody(err *Error) string {
|
||||
if err == nil {
|
||||
return "request failed"
|
||||
}
|
||||
if msg := strings.TrimSpace(err.Message); msg != "" {
|
||||
return msg
|
||||
}
|
||||
if msg := strings.TrimSpace(err.Error()); msg != "" {
|
||||
return msg
|
||||
}
|
||||
return "request failed"
|
||||
}
|
||||
|
||||
func timePtrIfSet(value time.Time) *time.Time {
|
||||
if value.IsZero() {
|
||||
return nil
|
||||
}
|
||||
copyValue := value
|
||||
return ©Value
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"
|
||||
)
|
||||
|
||||
func TestManagerMarkResultPublishesErrorEventAfterAuthStateUpdate(t *testing.T) {
|
||||
withEnabledErrorQueue(t)
|
||||
subscriber, unsubscribe := redisqueue.SubscribeErrors()
|
||||
defer unsubscribe()
|
||||
|
||||
manager := NewManager(nil, nil, nil)
|
||||
auth := &Auth{
|
||||
ID: "auth-error-event",
|
||||
Provider: "codex",
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
|
||||
t.Fatalf("Register returned error: %v", errRegister)
|
||||
}
|
||||
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: auth.ID,
|
||||
Provider: "codex",
|
||||
Model: "gpt-5",
|
||||
Success: false,
|
||||
Error: &Error{
|
||||
Code: "rate_limit",
|
||||
Message: `{"error":"quota"}`,
|
||||
Retryable: true,
|
||||
HTTPStatus: http.StatusTooManyRequests,
|
||||
},
|
||||
})
|
||||
|
||||
payload := requireErrorSubscriberPayload(t, subscriber)
|
||||
|
||||
var event struct {
|
||||
Provider string `json:"provider"`
|
||||
Model string `json:"model"`
|
||||
AuthID string `json:"auth_id"`
|
||||
AuthIndex string `json:"auth_index"`
|
||||
StatusCode int `json:"status_code"`
|
||||
Body string `json:"body"`
|
||||
Code string `json:"code"`
|
||||
Retryable bool `json:"retryable"`
|
||||
AuthStatus struct {
|
||||
Status Status `json:"status"`
|
||||
StatusMessage string `json:"status_message"`
|
||||
Unavailable bool `json:"unavailable"`
|
||||
Quota *struct {
|
||||
Exceeded bool `json:"exceeded"`
|
||||
Reason string `json:"reason"`
|
||||
} `json:"quota"`
|
||||
Model *struct {
|
||||
Name string `json:"name"`
|
||||
Status Status `json:"status"`
|
||||
Unavailable bool `json:"unavailable"`
|
||||
Quota *struct {
|
||||
Exceeded bool `json:"exceeded"`
|
||||
Reason string `json:"reason"`
|
||||
} `json:"quota"`
|
||||
} `json:"model"`
|
||||
} `json:"auth_status"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(payload, &event); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal error event: %v body=%s", errUnmarshal, string(payload))
|
||||
}
|
||||
if event.Provider != "codex" || event.Model != "gpt-5" || event.AuthID != auth.ID {
|
||||
t.Fatalf("unexpected event routing fields: %+v", event)
|
||||
}
|
||||
if event.AuthIndex == "" {
|
||||
t.Fatalf("auth_index is empty in event: %s", string(payload))
|
||||
}
|
||||
if event.StatusCode != http.StatusTooManyRequests || event.Body != `{"error":"quota"}` {
|
||||
t.Fatalf("unexpected error fields: status=%d body=%q", event.StatusCode, event.Body)
|
||||
}
|
||||
if event.Code != "rate_limit" || !event.Retryable {
|
||||
t.Fatalf("unexpected error code fields: code=%q retryable=%t", event.Code, event.Retryable)
|
||||
}
|
||||
if event.AuthStatus.Status != StatusError || !event.AuthStatus.Unavailable {
|
||||
t.Fatalf("unexpected auth status: %+v", event.AuthStatus)
|
||||
}
|
||||
if event.AuthStatus.Model == nil || event.AuthStatus.Model.Name != "gpt-5" || event.AuthStatus.Model.Status != StatusError || !event.AuthStatus.Model.Unavailable {
|
||||
t.Fatalf("unexpected model status: %+v", event.AuthStatus.Model)
|
||||
}
|
||||
if event.AuthStatus.Quota == nil || !event.AuthStatus.Quota.Exceeded || event.AuthStatus.Quota.Reason != "quota" {
|
||||
t.Fatalf("unexpected auth quota: %+v", event.AuthStatus.Quota)
|
||||
}
|
||||
if event.AuthStatus.Model.Quota == nil || !event.AuthStatus.Model.Quota.Exceeded || event.AuthStatus.Model.Quota.Reason != "quota" {
|
||||
t.Fatalf("unexpected model quota: %+v", event.AuthStatus.Model.Quota)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerMarkResultSkipsErrorEventInHomeMode(t *testing.T) {
|
||||
withEnabledErrorQueue(t)
|
||||
subscriber, unsubscribe := redisqueue.SubscribeErrors()
|
||||
defer unsubscribe()
|
||||
|
||||
manager := NewManager(nil, nil, nil)
|
||||
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
||||
auth := &Auth{
|
||||
ID: "home-auth-error-event",
|
||||
Provider: "codex",
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
|
||||
t.Fatalf("Register returned error: %v", errRegister)
|
||||
}
|
||||
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: auth.ID,
|
||||
Provider: "codex",
|
||||
Model: "gpt-5",
|
||||
Success: false,
|
||||
Error: &Error{
|
||||
Message: "unauthorized",
|
||||
HTTPStatus: http.StatusUnauthorized,
|
||||
},
|
||||
})
|
||||
|
||||
select {
|
||||
case got := <-subscriber:
|
||||
t.Fatalf("received home-mode error event %q, want none", string(got))
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func withEnabledErrorQueue(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
prevQueueEnabled := redisqueue.Enabled()
|
||||
redisqueue.SetEnabled(false)
|
||||
redisqueue.SetEnabled(true)
|
||||
|
||||
t.Cleanup(func() {
|
||||
redisqueue.SetEnabled(false)
|
||||
redisqueue.SetEnabled(prevQueueEnabled)
|
||||
})
|
||||
}
|
||||
|
||||
func requireErrorSubscriberPayload(t *testing.T, subscriber <-chan []byte) []byte {
|
||||
t.Helper()
|
||||
|
||||
select {
|
||||
case got, ok := <-subscriber:
|
||||
if !ok {
|
||||
t.Fatalf("error subscriber closed before receiving payload")
|
||||
}
|
||||
return got
|
||||
case <-time.After(time.Second):
|
||||
t.Fatalf("timeout waiting for error subscriber payload")
|
||||
return nil
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user