feat(home): implement count for home auth dispatch requests and enable usage statistics

- Added `count` attribute to `homeAuthCount` requests to improve home message batching.
- Enabled usage statistics for home mode by default and added config-level enforcement.
- Adjusted failure logging to include detailed metadata in `UsageReporter`.
- Updated multiple executors to pass error details to `PublishFailure` for better debugging.
- Enhanced unit tests to validate `count` behavior and usage statistics enforcement across components.
This commit is contained in:
Luis Pater
2026-05-10 01:30:43 +08:00
parent 1abf8625d8
commit 66c3dae06b
21 changed files with 281 additions and 52 deletions
+73 -10
View File
@@ -51,6 +51,7 @@ type ExecutionSessionCloser interface {
}
const (
homeAuthCountMetadataKey = "__cliproxy_home_auth_count"
// CloseAllExecutionSessionsID asks an executor to release all active execution sessions.
// Executors that do not support this marker may ignore it.
CloseAllExecutionSessionsID = "__all_execution_sessions__"
@@ -1316,19 +1317,25 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req
}
routeModel := req.Model
opts = ensureRequestedModelMetadata(opts, routeModel)
homeMode := m.HomeEnabled()
homeAuthCount := 1
tried := make(map[string]struct{})
attempted := make(map[string]struct{})
var lastErr error
for {
if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if lastErr != nil {
return cliproxyexecutor.Response{}, lastErr
}
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
pickOpts := opts
if homeMode {
pickOpts = withHomeAuthCount(opts, homeAuthCount)
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
if errPick != nil {
if lastErr != nil {
if lastErr != nil && !homeMode {
return cliproxyexecutor.Response{}, lastErr
}
return cliproxyexecutor.Response{}, errPick
@@ -1384,6 +1391,9 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req
return cliproxyexecutor.Response{}, authErr
}
lastErr = authErr
if homeMode {
homeAuthCount++
}
continue
}
}
@@ -1395,19 +1405,25 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string,
}
routeModel := req.Model
opts = ensureRequestedModelMetadata(opts, routeModel)
homeMode := m.HomeEnabled()
homeAuthCount := 1
tried := make(map[string]struct{})
attempted := make(map[string]struct{})
var lastErr error
for {
if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if lastErr != nil {
return cliproxyexecutor.Response{}, lastErr
}
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
pickOpts := opts
if homeMode {
pickOpts = withHomeAuthCount(opts, homeAuthCount)
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
if errPick != nil {
if lastErr != nil {
if lastErr != nil && !homeMode {
return cliproxyexecutor.Response{}, lastErr
}
return cliproxyexecutor.Response{}, errPick
@@ -1463,6 +1479,9 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string,
return cliproxyexecutor.Response{}, authErr
}
lastErr = authErr
if homeMode {
homeAuthCount++
}
continue
}
}
@@ -1474,19 +1493,25 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string
}
routeModel := req.Model
opts = ensureRequestedModelMetadata(opts, routeModel)
homeMode := m.HomeEnabled()
homeAuthCount := 1
tried := make(map[string]struct{})
attempted := make(map[string]struct{})
var lastErr error
for {
if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if lastErr != nil {
return nil, lastErr
}
return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
pickOpts := opts
if homeMode {
pickOpts = withHomeAuthCount(opts, homeAuthCount)
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
if errPick != nil {
if lastErr != nil {
if lastErr != nil && !homeMode {
return nil, lastErr
}
return nil, errPick
@@ -1516,6 +1541,9 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string
return nil, errStream
}
lastErr = errStream
if homeMode {
homeAuthCount++
}
continue
}
return streamResult, nil
@@ -1543,6 +1571,40 @@ func ensureRequestedModelMetadata(opts cliproxyexecutor.Options, requestedModel
return opts
}
func withHomeAuthCount(opts cliproxyexecutor.Options, count int) cliproxyexecutor.Options {
if count <= 0 {
count = 1
}
meta := make(map[string]any, len(opts.Metadata)+1)
for k, v := range opts.Metadata {
meta[k] = v
}
meta[homeAuthCountMetadataKey] = count
opts.Metadata = meta
return opts
}
func homeAuthCountFromMetadata(meta map[string]any) int {
if len(meta) == 0 {
return 1
}
switch value := meta[homeAuthCountMetadataKey].(type) {
case int:
if value > 0 {
return value
}
case int64:
if value > 0 {
return int(value)
}
case float64:
if value > 0 {
return int(value)
}
}
return 1
}
func hasRequestedModelMetadata(meta map[string]any) bool {
if len(meta) == 0 {
return false
@@ -3099,8 +3161,9 @@ func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts clipro
requestedModel := requestedModelFromMetadata(opts.Metadata, model)
sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata)
count := homeAuthCountFromMetadata(opts.Metadata)
raw, err := client.RPopAuth(ctx, requestedModel, sessionID, opts.Headers)
raw, err := client.RPopAuth(ctx, requestedModel, sessionID, opts.Headers, count)
if err != nil {
return nil, nil, "", &Error{Code: "auth_not_found", Message: err.Error(), HTTPStatus: http.StatusServiceUnavailable}
}
+2
View File
@@ -561,6 +561,7 @@ func forceHomeRuntimeConfig(cfg *config.Config) {
return
}
cfg.APIKeys = nil
cfg.UsageStatisticsEnabled = true
cfg.DisableCooling = true
cfg.WebsocketAuth = false
cfg.EnableGeminiCLIEndpoint = false
@@ -732,6 +733,7 @@ func (s *Service) Run(ctx context.Context) error {
homeEnabled := s.cfg != nil && s.cfg.Home.Enabled
if homeEnabled {
forceHomeRuntimeConfig(s.cfg)
redisqueue.SetUsageStatisticsEnabled(true)
}
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
+29
View File
@@ -99,3 +99,32 @@ func TestServiceApplyCoreAuthAddOrUpdate_DeleteReAddDoesNotInheritStaleRuntimeSt
t.Fatalf("expected re-added auth to re-register models in global registry")
}
}
func TestForceHomeRuntimeConfigEnablesUsageStatistics(t *testing.T) {
cfg := &config.Config{
UsageStatisticsEnabled: false,
}
forceHomeRuntimeConfig(cfg)
if !cfg.UsageStatisticsEnabled {
t.Fatal("expected home runtime config to force usage statistics enabled")
}
}
func TestApplyHomeOverlayForcesUsageStatisticsEnabled(t *testing.T) {
baseCfg := &config.Config{}
baseCfg.Home.Enabled = true
service := &Service{cfg: baseCfg}
service.applyHomeOverlay(&config.Config{
UsageStatisticsEnabled: false,
})
if service.cfg == nil || !service.cfg.UsageStatisticsEnabled {
t.Fatal("expected home overlay to force usage statistics enabled")
}
if !service.cfg.Home.Enabled {
t.Fatal("expected home overlay to preserve local home settings")
}
}
+7
View File
@@ -22,9 +22,16 @@ type Record struct {
RequestedAt time.Time
Latency time.Duration
Failed bool
Fail Failure
Detail Detail
}
// Failure holds HTTP failure metadata for an upstream request attempt.
type Failure struct {
StatusCode int
Body string
}
// Detail holds the token usage breakdown.
type Detail struct {
InputTokens int64