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:
@@ -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}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user