Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
SSH key fingerprint:
SHA256:GBzT2DXvAuGV9XIV5W3WrzVpjU54FThmxHXdbz95J24
1802 changed files with 503346 additions and 2 deletions
358
backend/sdk/cliproxy/auth/conductor_home_execution.go
Normal file
358
backend/sdk/cliproxy/auth/conductor_home_execution.go
Normal file
|
|
@ -0,0 +1,358 @@
|
|||
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
|
||||
}
|
||||
Loading…
Reference in a new issue