|
|
|
@@ -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}
|
|
|
|
|
}
|
|
|
|
|