358 lines
14 KiB
Go
358 lines
14 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
|
"github.com/tidwall/sjson"
|
|
)
|
|
|
|
func (m *Manager) executeHome(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool) (cliproxyexecutor.Response, error) {
|
|
if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil {
|
|
defer unlockSession()
|
|
}
|
|
defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings()
|
|
retryModel := authSelectionModelFromOptions(opts, req.Model)
|
|
homeRetryLimit := -1
|
|
attempt := 0
|
|
retryRoundPending := false
|
|
retryRoundWaited := false
|
|
for {
|
|
response, errExecute := m.executeHomeOnce(ctx, providers, req, opts, countTokens, maxRetryCredentials, &homeRetryLimit, attempt)
|
|
if errExecute == nil {
|
|
return response, nil
|
|
}
|
|
if retryRoundPending {
|
|
if wait, okWait := pendingHomeRetryRoundDelay(errExecute, maxWait, &homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == ""); okWait && m.homeRetryAllowed(attempt-1, homeRetryLimit) {
|
|
if retryRoundWaited {
|
|
return cliproxyexecutor.Response{}, errExecute
|
|
}
|
|
if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil {
|
|
return cliproxyexecutor.Response{}, errWait
|
|
}
|
|
retryRoundWaited = true
|
|
continue
|
|
}
|
|
}
|
|
retryRoundPending = false
|
|
retryRoundWaited = false
|
|
if isRequestTerminatedError(errExecute) || isRequestStopError(errExecute) {
|
|
return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute)
|
|
}
|
|
wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExecute, attempt, providers, retryModel, maxWait, homeRetryLimit, defaultRequestRetry)
|
|
if !shouldRetry {
|
|
return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute)
|
|
}
|
|
if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil {
|
|
return cliproxyexecutor.Response{}, errWait
|
|
}
|
|
attempt++
|
|
retryRoundPending = true
|
|
retryRoundWaited = false
|
|
}
|
|
}
|
|
|
|
func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool, maxRetryCredentials int, homeRetryLimit *int, retryRounds ...int) (cliproxyexecutor.Response, error) {
|
|
retryRound := 0
|
|
if len(retryRounds) > 0 {
|
|
retryRound = retryRounds[0]
|
|
}
|
|
routeModel := authSelectionModelFromOptions(opts, req.Model)
|
|
responseAlias := requestedModelAliasFromOptions(opts, routeModel)
|
|
executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model)
|
|
opts = ensureRequestedModelMetadata(opts, routeModel)
|
|
tried := make(map[string]struct{})
|
|
attempted := make(map[string]struct{})
|
|
var lastErr error
|
|
var roundTiming homeRetryRoundTiming
|
|
for homeAuthCount := 1; ; homeAuthCount++ {
|
|
if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
|
|
if lastErr != nil {
|
|
return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), true)
|
|
}
|
|
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
pickOpts := withHomeRetryRound(opts, retryRound)
|
|
pickOpts = withHomeAuthCount(pickOpts, homeAuthCount)
|
|
pickOpts = withHomeExcludedAuthIDs(pickOpts, tried)
|
|
selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, pickOpts)
|
|
if errSelection != nil {
|
|
var homeCooldown *homeDispatchRetryAfterError
|
|
if lastErr != nil && errors.As(errSelection, &homeCooldown) && homeCooldown != nil {
|
|
observeHomeCooldownRetryLimit(homeCooldown, homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == "")
|
|
return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, homeCooldown.RetryAfter(), false)
|
|
}
|
|
if shouldReturnLastErrorOnPickFailure(true, lastErr, errSelection) {
|
|
return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), isHomeNextRoundImmediatelyAvailable(errSelection))
|
|
}
|
|
return cliproxyexecutor.Response{}, errSelection
|
|
}
|
|
auth := selection.CloneAuthForRoute(routeModel)
|
|
if auth == nil || selection.Executor == nil {
|
|
selection.End("missing_execution_target")
|
|
return cliproxyexecutor.Response{}, &Error{Code: "executor_not_found", Message: "executor not registered"}
|
|
}
|
|
m.observeHomeRetryLimit(auth, selection, homeRetryLimit)
|
|
if _, seen := tried[auth.ID]; seen {
|
|
if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "repeated_auth"); errEnd != nil {
|
|
return cliproxyexecutor.Response{}, errEnd
|
|
}
|
|
if lastErr != nil {
|
|
return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), false)
|
|
}
|
|
return cliproxyexecutor.Response{}, repeatedHomeAuthError()
|
|
}
|
|
tried[auth.ID] = struct{}{}
|
|
attempted[auth.ID] = struct{}{}
|
|
entry := logEntryWithRequestID(ctx)
|
|
debugLogAuthSelection(entry, auth, selection.Provider, routeModel)
|
|
if errRuntimeAuth := m.bindHomeSelectionRuntimeAuth(ctx, opts, selection); errRuntimeAuth != nil {
|
|
selection.End("runtime_auth_bind_failed")
|
|
return cliproxyexecutor.Response{}, errRuntimeAuth
|
|
}
|
|
publishSelectedAuthMetadata(opts.Metadata, auth)
|
|
execCtx, releaseAttempt, errBind := homeExecutionAttemptContext(ctx, selection)
|
|
if errBind != nil {
|
|
selection.End("attempt_bind_failed")
|
|
return cliproxyexecutor.Response{}, errBind
|
|
}
|
|
// Enrich before auth preparation so prepare-stage usage records observe the client request.
|
|
execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel)
|
|
if rt := m.roundTripperFor(auth); rt != nil {
|
|
execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
|
|
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
|
|
}
|
|
models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAlias(auth, routeModel)
|
|
if aliasResult.ForceMapping && responseAlias != "" {
|
|
aliasResult.OriginalAlias = responseAlias
|
|
}
|
|
if len(models) > 1 {
|
|
models = models[:1]
|
|
pooled = false
|
|
}
|
|
if len(models) == 0 {
|
|
releaseAttempt()
|
|
if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "no_execution_models"); errEnd != nil {
|
|
return cliproxyexecutor.Response{}, errEnd
|
|
}
|
|
lastErr = &Error{Code: "auth_not_found", Message: "no execution models available"}
|
|
roundTiming.Observe(lastErr)
|
|
continue
|
|
}
|
|
preparedAuth, errPrepare := m.prepareHomeRequestAuth(execCtx, selection.Executor, selection)
|
|
if errPrepare != nil {
|
|
m.reportHomeResult(execCtx, Result{AuthID: auth.ID, Provider: selection.Provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: opts}, auth)
|
|
releaseAttempt()
|
|
if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil {
|
|
return cliproxyexecutor.Response{}, errEnd
|
|
}
|
|
lastErr = errPrepare
|
|
roundTiming.Observe(lastErr)
|
|
continue
|
|
}
|
|
didRefreshOnUnauthorized := false
|
|
for _, upstreamModel := range models {
|
|
resultModel := m.stateModelForExecution(preparedAuth, routeModel, upstreamModel, pooled)
|
|
execReq := req
|
|
execReq.Model = upstreamModel
|
|
if restoreExecutionModel {
|
|
execReq.Model = executionModel
|
|
}
|
|
execOpts := opts
|
|
execOpts.ExecutionLifecycle = selection
|
|
var errIntercept error
|
|
execReq, execOpts, errIntercept = applyRequestAfterAuthInterceptor(execCtx, selection.Executor, selection.Provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel))
|
|
if errIntercept != nil {
|
|
releaseAttempt()
|
|
selection.End("request_intercepted")
|
|
return cliproxyexecutor.Response{}, errIntercept
|
|
}
|
|
if !restoreExecutionModel {
|
|
execReq = attachResolvedAPIKeyModelInfo(routing, execReq, preparedAuth, routeModel, upstreamModel)
|
|
}
|
|
if errCtx := execCtx.Err(); errCtx != nil {
|
|
releaseAttempt()
|
|
selection.End("attempt_canceled")
|
|
return cliproxyexecutor.Response{}, errCtx
|
|
}
|
|
var response cliproxyexecutor.Response
|
|
var errExecute error
|
|
var effectiveAuthMu sync.RWMutex
|
|
effectiveAuth := preparedAuth.Clone()
|
|
setEffectiveAuth := func(auth *Auth) {
|
|
if auth == nil || AccessTokenSHA256(auth) == "" {
|
|
return
|
|
}
|
|
effectiveAuthMu.Lock()
|
|
effectiveAuth = auth.Clone()
|
|
effectiveAuthMu.Unlock()
|
|
}
|
|
getEffectiveAuth := func() (*Auth, string) {
|
|
effectiveAuthMu.RLock()
|
|
defer effectiveAuthMu.RUnlock()
|
|
if effectiveAuth == nil {
|
|
return nil, ""
|
|
}
|
|
return effectiveAuth.Clone(), AccessTokenSHA256(effectiveAuth)
|
|
}
|
|
executorCtx := execCtx
|
|
if countTokens {
|
|
executorCtx = withAccessTokenFingerprintObserver(execCtx, setEffectiveAuth)
|
|
}
|
|
execute := func() (cliproxyexecutor.Response, error) {
|
|
if countTokens {
|
|
return selection.Executor.CountTokens(executorCtx, preparedAuth, execReq, execOpts)
|
|
}
|
|
return selection.Executor.Execute(execCtx, preparedAuth, execReq, execOpts)
|
|
}
|
|
startHomeExec := time.Now()
|
|
response, errExecute = execute()
|
|
durationHomeExec := time.Since(startHomeExec)
|
|
refreshAuth := preparedAuth
|
|
if countTokens {
|
|
if observedAuth, fingerprint := getEffectiveAuth(); isUnauthorizedError(errExecute) {
|
|
m.reportHomeUnauthorized(execCtx, preparedAuth, selection.Provider, resultModel, fingerprint)
|
|
if observedAuth != nil {
|
|
refreshAuth = observedAuth
|
|
}
|
|
}
|
|
}
|
|
if errExecute != nil {
|
|
if refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(execCtx, selection.Executor, refreshAuth, errExecute, didRefreshOnUnauthorized, true); errRefresh != nil {
|
|
errExecute = errRefresh
|
|
warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeExec, errExecute)
|
|
} else if okRefresh {
|
|
preparedAuth = refreshed
|
|
m.replaceHomeSelectionAuth(selection, preparedAuth)
|
|
didRefreshOnUnauthorized = true
|
|
publishSelectedAuthMetadata(opts.Metadata, preparedAuth)
|
|
setEffectiveAuth(preparedAuth)
|
|
startHomeRetry := time.Now()
|
|
response, errExecute = execute()
|
|
durationHomeRetry := time.Since(startHomeRetry)
|
|
if errExecute != nil {
|
|
warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeRetry, errExecute)
|
|
if countTokens && isUnauthorizedError(errExecute) {
|
|
_, fingerprint := getEffectiveAuth()
|
|
m.reportHomeUnauthorized(execCtx, preparedAuth, selection.Provider, resultModel, fingerprint)
|
|
}
|
|
}
|
|
} else {
|
|
warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeExec, errExecute)
|
|
}
|
|
}
|
|
result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil, Options: execOpts}
|
|
if errExecute == nil {
|
|
m.reportHomeResult(execCtx, result, preparedAuth)
|
|
releaseAttempt()
|
|
attemptAliasResult := resolveAttemptAliasResult(routing, preparedAuth, routeModel, upstreamModel, aliasResult)
|
|
rewriteForceMappedResponse(&response, attemptAliasResult)
|
|
if !m.retainHomeWebsocketSelection(ctx, opts, routeModel, selection) {
|
|
selection.End("completed")
|
|
}
|
|
return response, nil
|
|
}
|
|
result.Error = resultErrorFromError(errExecute)
|
|
result.RetryAfter = retryAfterFromError(errExecute)
|
|
if isCredentialScopedError(errExecute) {
|
|
result.CredentialScope = true
|
|
}
|
|
action, okAction := matchRequestScopedErrorAction(preparedAuth, errExecute, m.runtimeConfigSnapshot())
|
|
applyRequestScopedActionToResult(action, okAction, &result)
|
|
m.reportHomeResult(execCtx, result, preparedAuth)
|
|
lastErr = errExecute
|
|
if okAction {
|
|
if isRequestScopedStop(action, okAction) {
|
|
releaseAttempt()
|
|
selection.End("request_stopped")
|
|
return cliproxyexecutor.Response{}, wrapRequestStopError(errExecute)
|
|
}
|
|
if result.CredentialScope {
|
|
break
|
|
}
|
|
continue
|
|
}
|
|
if isRequestInvalidError(errExecute) {
|
|
releaseAttempt()
|
|
selection.End("request_invalid")
|
|
return cliproxyexecutor.Response{}, errExecute
|
|
}
|
|
if result.CredentialScope {
|
|
break
|
|
}
|
|
}
|
|
roundTiming.Observe(lastErr)
|
|
releaseAttempt()
|
|
if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "execution_failed"); errEnd != nil {
|
|
return cliproxyexecutor.Response{}, errEnd
|
|
}
|
|
if errCtx := execCtx.Err(); errCtx != nil && ctx != nil && ctx.Err() != nil {
|
|
return cliproxyexecutor.Response{}, errCtx
|
|
}
|
|
}
|
|
}
|
|
|
|
func homeExecutionAttemptContext(ctx context.Context, selection *HomeDispatchSelection) (context.Context, func(), error) {
|
|
if selection == nil {
|
|
return nil, func() {}, fmt.Errorf("Home dispatch selection is nil")
|
|
}
|
|
return selection.AttemptContext(ctx)
|
|
}
|
|
|
|
func wrapHomeStream(ctx context.Context, result *cliproxyexecutor.StreamResult, selection *HomeDispatchSelection, releaseAttempt func()) *cliproxyexecutor.StreamResult {
|
|
if result == nil || result.Chunks == nil {
|
|
if releaseAttempt != nil {
|
|
releaseAttempt()
|
|
}
|
|
return result
|
|
}
|
|
out := make(chan cliproxyexecutor.StreamChunk)
|
|
go func() {
|
|
defer close(out)
|
|
if releaseAttempt != nil {
|
|
defer releaseAttempt()
|
|
}
|
|
if selection != nil {
|
|
defer selection.End("stream_closed")
|
|
}
|
|
forward := true
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case chunk, ok := <-result.Chunks:
|
|
if !ok {
|
|
return
|
|
}
|
|
if !forward {
|
|
continue
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case out <- chunk:
|
|
}
|
|
if chunk.Err != nil && selection != nil {
|
|
forward = false
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
return &cliproxyexecutor.StreamResult{Headers: result.Headers, Chunks: out}
|
|
}
|
|
|
|
func sanitizeDownstreamWebsocketFallbackRequest(ctx context.Context, auth *Auth, req cliproxyexecutor.Request) cliproxyexecutor.Request {
|
|
if !cliproxyexecutor.DownstreamWebsocket(ctx) || authWebsocketsEnabled(auth) || len(req.Payload) == 0 {
|
|
return req
|
|
}
|
|
updated, errDelete := sjson.DeleteBytes(req.Payload, "generate")
|
|
if errDelete != nil {
|
|
return req
|
|
}
|
|
req.Payload = updated
|
|
return req
|
|
}
|