diff --git a/internal/home/client.go b/internal/home/client.go index b5115d43a..f4a295c6f 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -1327,7 +1327,7 @@ func isAmbiguousIssuedRPopAuthError(err error) bool { return !errors.As(err, &redisErr) } -func (c *Client) GetRefreshAuth(ctx context.Context, authIndex string, lastRefreshedAt time.Time, accessTokenSHA256 string) ([]byte, error) { +func (c *Client) GetRefreshAuth(ctx context.Context, authIndex string) ([]byte, error) { cmd, errClient := c.commandClient() if errClient != nil { return nil, errClient @@ -1340,10 +1340,6 @@ func (c *Client) GetRefreshAuth(ctx context.Context, authIndex string, lastRefre Type: "refresh", AuthIndex: authIndex, } - if !lastRefreshedAt.IsZero() { - req.LastRefreshedAt = lastRefreshedAt.UTC().Format(time.RFC3339Nano) - } - req.ObservedAccessTokenSHA256 = strings.TrimSpace(accessTokenSHA256) keyBytes, err := json.Marshal(&req) if err != nil { return nil, err diff --git a/internal/home/requests.go b/internal/home/requests.go index c0ce7a755..eca63742f 100644 --- a/internal/home/requests.go +++ b/internal/home/requests.go @@ -18,10 +18,8 @@ type modelsRequest struct { } type refreshRequest struct { - Type string `json:"type"` - AuthIndex string `json:"auth_index"` - LastRefreshedAt string `json:"last_refreshed_at,omitempty"` - ObservedAccessTokenSHA256 string `json:"access_token_sha256,omitempty"` + Type string `json:"type"` + AuthIndex string `json:"auth_index"` } type InFlightFrameKind string diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index d91c8a282..915f88947 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -90,7 +90,6 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec TTFTMs: record.TTFT.Milliseconds(), Source: record.Source, AuthIndex: record.AuthIndex, - AccessTokenHash: record.AccessTokenSHA256, ClientIP: clientRequestMetadata.ClientIP, XForwardedFor: clientRequestMetadata.XForwardedFor, UserAgent: clientRequestMetadata.UserAgent, @@ -146,7 +145,6 @@ type requestDetail struct { TTFTMs int64 `json:"ttft_ms"` Source string `json:"source"` AuthIndex string `json:"auth_index"` - AccessTokenHash string `json:"access_token_sha256,omitempty"` ClientIP string `json:"client_ip"` XForwardedFor string `json:"x_forwarded_for"` UserAgent string `json:"user_agent"` diff --git a/internal/runtime/executor/helps/home_refresh.go b/internal/runtime/executor/helps/home_refresh.go index 020d5f4fe..7c9719927 100644 --- a/internal/runtime/executor/helps/home_refresh.go +++ b/internal/runtime/executor/helps/home_refresh.go @@ -2,14 +2,10 @@ package helps import ( "context" - "crypto/sha256" - "encoding/hex" "encoding/json" - "errors" "fmt" "net/http" "strings" - "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/home" @@ -47,7 +43,7 @@ type homeErrorDetail struct { type homeRefreshClient interface { HeartbeatOK() bool - GetRefreshAuth(ctx context.Context, authIndex string, lastRefreshedAt time.Time, accessTokenSHA256 string) ([]byte, error) + GetRefreshAuth(ctx context.Context, authIndex string) ([]byte, error) } var currentHomeRefreshClient = func() homeRefreshClient { @@ -81,11 +77,8 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya return nil, true, homeStatusErr{code: http.StatusBadGateway, msg: "home refresh: auth_index is empty"} } - raw, err := client.GetRefreshAuth(ctx, authIndex, auth.LastRefreshedAt, authAccessTokenSHA256(auth)) + raw, err := client.GetRefreshAuth(ctx, authIndex) if err != nil { - if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { - return nil, true, err - } return nil, true, homeStatusErr{code: http.StatusBadGateway, msg: err.Error()} } @@ -114,43 +107,6 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya return updated, true, nil } -func authAccessTokenSHA256(auth *cliproxyauth.Auth) string { - accessToken := authAccessTokenForFingerprint(auth) - if accessToken == "" { - return "" - } - digest := sha256.Sum256([]byte(accessToken)) - return hex.EncodeToString(digest[:]) -} - -func authAccessTokenForFingerprint(auth *cliproxyauth.Auth) string { - if auth == nil || auth.Metadata == nil { - return "" - } - for _, key := range []string{"access_token", "accessToken"} { - if value, ok := auth.Metadata[key].(string); ok && strings.TrimSpace(value) != "" { - return strings.TrimSpace(value) - } - } - for _, key := range []string{"token", "Token"} { - switch token := auth.Metadata[key].(type) { - case map[string]any: - for _, tokenKey := range []string{"access_token", "accessToken"} { - if value, ok := token[tokenKey].(string); ok && strings.TrimSpace(value) != "" { - return strings.TrimSpace(value) - } - } - case map[string]string: - for _, tokenKey := range []string{"access_token", "accessToken"} { - if value := strings.TrimSpace(token[tokenKey]); value != "" { - return value - } - } - } - } - return "" -} - func parseHomeRefreshAuth(raw []byte) (*cliproxyauth.Auth, string, error) { var rawObject map[string]json.RawMessage if errUnmarshal := json.Unmarshal(raw, &rawObject); errUnmarshal != nil { @@ -176,8 +132,6 @@ func statusFromHomeErrorCode(code string) int { return http.StatusUnauthorized case "model_not_found": return http.StatusNotFound - case "refresh_temporarily_unavailable", "home_unavailable": - return http.StatusServiceUnavailable default: return http.StatusBadGateway } diff --git a/internal/runtime/executor/helps/home_refresh_test.go b/internal/runtime/executor/helps/home_refresh_test.go index 26cc51afc..ca7582732 100644 --- a/internal/runtime/executor/helps/home_refresh_test.go +++ b/internal/runtime/executor/helps/home_refresh_test.go @@ -3,11 +3,9 @@ package helps import ( "context" "encoding/json" - "errors" "net/http" "sync/atomic" "testing" - "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -20,60 +18,22 @@ func TestStatusFromHomeErrorCodeMapsAuthenticationErrorToUnauthorized(t *testing if got := statusFromHomeErrorCode("unauthorized"); got != http.StatusUnauthorized { t.Fatalf("statusFromHomeErrorCode(unauthorized) = %d, want %d", got, http.StatusUnauthorized) } - if got := statusFromHomeErrorCode("refresh_temporarily_unavailable"); got != http.StatusServiceUnavailable { - t.Fatalf("statusFromHomeErrorCode(refresh_temporarily_unavailable) = %d, want %d", got, http.StatusServiceUnavailable) - } } type fakeHomeRefreshClient struct { - calls atomic.Int32 - authIndex string - lastRefreshedAt time.Time - accessTokenHash string - raw []byte - err error + calls atomic.Int32 + authIndex string + raw []byte } func (c *fakeHomeRefreshClient) HeartbeatOK() bool { return true } -func (c *fakeHomeRefreshClient) GetRefreshAuth(_ context.Context, authIndex string, lastRefreshedAt time.Time, accessTokenHash string) ([]byte, error) { +func (c *fakeHomeRefreshClient) GetRefreshAuth(_ context.Context, authIndex string) ([]byte, error) { c.calls.Add(1) c.authIndex = authIndex - c.lastRefreshedAt = lastRefreshedAt - c.accessTokenHash = accessTokenHash - return c.raw, c.err -} - -func TestRefreshAuthViaHomePreservesContextErrors(t *testing.T) { - client := &fakeHomeRefreshClient{err: context.DeadlineExceeded} - oldCurrentHomeRefreshClient := currentHomeRefreshClient - currentHomeRefreshClient = func() homeRefreshClient { return client } - t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient }) - - cfg := &config.Config{Home: config.HomeConfig{Enabled: true}} - auth := &cliproxyauth.Auth{ID: "home-auth", Index: "home-auth", Provider: "codex"} - _, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth) - if !handled || !errors.Is(errRefresh, context.DeadlineExceeded) { - t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want true/context.DeadlineExceeded", handled, errRefresh) - } -} - -func TestAuthAccessTokenSHA256SupportsKnownMetadataShapes(t *testing.T) { - want := authAccessTokenSHA256(&cliproxyauth.Auth{Metadata: map[string]any{"access_token": "same-token"}}) - cases := map[string]*cliproxyauth.Auth{ - "camel case": {Metadata: map[string]any{"accessToken": "same-token"}}, - "nested any map": {Metadata: map[string]any{"token": map[string]any{"access_token": "same-token"}}}, - "nested string map": {Metadata: map[string]any{"Token": map[string]string{"accessToken": "same-token"}}}, - } - for name, auth := range cases { - t.Run(name, func(t *testing.T) { - if got := authAccessTokenSHA256(auth); got == "" || got != want { - t.Fatalf("token hash = %q, want %q", got, want) - } - }) - } + return c.raw, nil } func TestRefreshAuthViaHomeAcceptsAuthEnvelope(t *testing.T) { @@ -104,14 +64,11 @@ func TestRefreshAuthViaHomeAcceptsAuthEnvelope(t *testing.T) { }) cfg := &config.Config{Home: config.HomeConfig{Enabled: true}} - observedRefreshAt := time.Now().UTC() auth := &cliproxyauth.Auth{ - ID: "home-auth-1", - Provider: "antigravity", - Index: "home-index-1", - LastRefreshedAt: observedRefreshAt, + ID: "home-auth-1", + Provider: "antigravity", + Index: "home-index-1", Metadata: map[string]any{ - "access_token": "old-access-token", "refresh_token": "refresh-token", }, } @@ -129,12 +86,6 @@ func TestRefreshAuthViaHomeAcceptsAuthEnvelope(t *testing.T) { if client.authIndex != "home-index-1" { t.Fatalf("home refresh auth_index = %q, want home-index-1", client.authIndex) } - if !client.lastRefreshedAt.Equal(observedRefreshAt) { - t.Fatalf("home refresh last_refreshed_at = %v, want %v", client.lastRefreshedAt, observedRefreshAt) - } - if client.accessTokenHash != authAccessTokenSHA256(auth) { - t.Fatalf("home refresh access token hash = %q, want %q", client.accessTokenHash, authAccessTokenSHA256(auth)) - } if updated == nil { t.Fatal("updated auth = nil") } diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 39a320fe8..52e1687f3 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -22,25 +22,24 @@ import ( ) type UsageReporter struct { - provider string - executorType string - model string - alias string - authID string - authIndex string - accessTokenHash string - authType string - apiKey string - source string - reasoning string - serviceTier string - generate bool - requestedAt time.Time - ttftMu sync.RWMutex - ttft time.Duration - ttftStart time.Time - ttftSet bool - once sync.Once + provider string + executorType string + model string + alias string + authID string + authIndex string + authType string + apiKey string + source string + reasoning string + serviceTier string + generate bool + requestedAt time.Time + ttftMu sync.RWMutex + ttft time.Duration + ttftStart time.Time + ttftSet bool + once sync.Once } type usageExecutor interface { @@ -78,7 +77,6 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox if auth != nil { reporter.authID = auth.ID reporter.authIndex = auth.EnsureIndex() - reporter.accessTokenHash = authAccessTokenSHA256(auth) } return reporter } @@ -266,7 +264,6 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f APIKey: r.apiKey, AuthID: r.authID, AuthIndex: r.authIndex, - AccessTokenSHA256: r.accessTokenHash, AuthType: r.authType, ReasoningEffort: r.reasoning, ServiceTier: r.serviceTier, diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 68dca5a5e..a9ca51654 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -636,7 +636,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string models = models[:1] pooled = false } - streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, routing, true, selection != nil) + streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, routing, !homeMode, selection != nil) if errStream != nil { if selection != nil { releaseAttempt() diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index a4cfed50d..324d2f7e4 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -431,18 +431,14 @@ func (m *Manager) endHomeSelectionBeforeRedispatch(ctx context.Context, selectio } func (m *Manager) retainHomeWebsocketSelection(ctx context.Context, opts cliproxyexecutor.Options, model string, selection *HomeDispatchSelection) bool { - if m == nil || selection == nil || !selection.Retained() || !cliproxyexecutor.DownstreamWebsocket(ctx) { - return false - } - selectionAuth := selection.CloneAuth() - if selectionAuth == nil { + if m == nil || selection == nil || !selection.Retained() || !cliproxyexecutor.DownstreamWebsocket(ctx) || selection.Auth == nil { return false } sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) - credentialID := strings.TrimSpace(selectionAuth.ID) + credentialID := strings.TrimSpace(selection.Auth.ID) routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) if selection.accountedModel == "" { - selection.accountedModel, _ = m.predictedHomeConcurrencyModel(selectionAuth, model) + selection.accountedModel, _ = m.predictedHomeConcurrencyModel(selection.Auth, model) } if sessionID == "" || credentialID == "" || !validRouteModel || selection.accountedModel == "" { return false @@ -461,7 +457,7 @@ func (m *Manager) retainHomeWebsocketSelection(ctx context.Context, opts cliprox previous := selections[key] selections[key] = selection m.mu.Unlock() - m.rememberHomeRuntimeAuth(sessionID, selectionAuth) + m.rememberHomeRuntimeAuth(sessionID, selection.Auth) if previous != nil && previous != selection { previous.End("target_replaced") } @@ -537,15 +533,11 @@ func (m *Manager) clearHomeRuntimeAuthsForSessionLocked(sessionID string) { } func (m *Manager) bindHomeSelectionRuntimeAuth(ctx context.Context, opts cliproxyexecutor.Options, selection *HomeDispatchSelection) error { - if m == nil || selection == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) { - return nil - } - selectionAuth := selection.CloneAuth() - if selectionAuth == nil || !authWebsocketsEnabled(selectionAuth) { + if m == nil || selection == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) || selection.Auth == nil || !authWebsocketsEnabled(selection.Auth) { return nil } sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) - authID := strings.TrimSpace(selectionAuth.ID) + authID := strings.TrimSpace(selection.Auth.ID) if sessionID == "" || authID == "" || !selection.runtimeAuthBound.CompareAndSwap(false, true) { return nil } @@ -562,15 +554,11 @@ func (m *Manager) bindHomeSelectionRuntimeAuth(ctx context.Context, opts cliprox } func (m *Manager) rememberHomeSelectionRuntimeAuth(sessionID string, selection *HomeDispatchSelection) { - if m == nil || selection == nil { - return - } - selectionAuth := selection.CloneAuth() - if selectionAuth == nil { + if m == nil || selection == nil || selection.Auth == nil { return } sessionID = strings.TrimSpace(sessionID) - authID := strings.TrimSpace(selectionAuth.ID) + authID := strings.TrimSpace(selection.Auth.ID) if sessionID == "" || authID == "" { return } @@ -587,33 +575,11 @@ func (m *Manager) rememberHomeSelectionRuntimeAuth(sessionID string, selection * if m.homeRuntimeAuthOwners[sessionID] == nil { m.homeRuntimeAuthOwners[sessionID] = make(map[string]*HomeDispatchSelection) } - m.homeRuntimeAuths[sessionID][authID] = selectionAuth + m.homeRuntimeAuths[sessionID][authID] = selection.Auth.Clone() m.homeRuntimeAuthOwners[sessionID][authID] = selection m.mu.Unlock() } -func (m *Manager) replaceHomeSelectionAuth(selection *HomeDispatchSelection, auth *Auth) { - if m == nil || selection == nil || auth == nil { - return - } - m.mu.Lock() - selection.ReplaceAuth(auth) - updated := selection.CloneAuth() - if updated == nil { - m.mu.Unlock() - return - } - for sessionID, owners := range m.homeRuntimeAuthOwners { - for authID, owner := range owners { - if owner != selection || m.homeRuntimeAuths[sessionID] == nil { - continue - } - m.homeRuntimeAuths[sessionID][authID] = updated.Clone() - } - } - m.mu.Unlock() -} - func (m *Manager) forgetHomeRuntimeAuth(sessionID string, authID string, owner *HomeDispatchSelection) { sessionID = strings.TrimSpace(sessionID) authID = strings.TrimSpace(authID) @@ -699,8 +665,7 @@ func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts clipro if errSelection != nil { return nil, nil, "", errSelection } - selectionAuth := selection.CloneAuth() - if selectionAuth == nil || homeAuthAlreadyTried(tried, selectionAuth.ID) { + if selection.Auth == nil || homeAuthAlreadyTried(tried, selection.Auth.ID) { selection.End("repeated_auth") return nil, nil, "", repeatedHomeAuthError() } diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 709b874f7..dfd14ee06 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -81,7 +81,6 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr lastErr = errPrepare continue } - didRefreshOnUnauthorized := false for _, upstreamModel := range models { resultModel := m.stateModelForExecution(preparedAuth, routeModel, upstreamModel, pooled) execReq := req @@ -108,23 +107,10 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr } var response cliproxyexecutor.Response var errExecute error - execute := func() (cliproxyexecutor.Response, error) { - if countTokens { - return selection.Executor.CountTokens(execCtx, preparedAuth, execReq, execOpts) - } - return selection.Executor.Execute(execCtx, preparedAuth, execReq, execOpts) - } - response, errExecute = execute() - if errExecute != nil { - if refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(execCtx, selection.Executor, preparedAuth, errExecute, didRefreshOnUnauthorized, true); errRefresh != nil { - errExecute = errRefresh - } else if okRefresh { - preparedAuth = refreshed - m.replaceHomeSelectionAuth(selection, preparedAuth) - didRefreshOnUnauthorized = true - publishSelectedAuthMetadata(opts.Metadata, preparedAuth) - response, errExecute = execute() - } + if countTokens { + response, errExecute = selection.Executor.CountTokens(execCtx, preparedAuth, execReq, execOpts) + } else { + response, errExecute = selection.Executor.Execute(execCtx, preparedAuth, execReq, execOpts) } result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil} if errExecute == nil { diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index e7ed65cb3..4d9385d4b 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -377,47 +377,8 @@ func clearUnauthorizedModelStates(auth *Auth, now time.Time) []string { return resumed } -// tryRefreshExecutionAuthAfterUnauthorized refreshes OAuth credentials once for -// either a local auth or an ephemeral Home dispatch auth. -func (m *Manager) tryRefreshExecutionAuthAfterUnauthorized(ctx context.Context, executor ProviderExecutor, auth *Auth, execErr error, alreadyTried bool, homeDispatch bool) (*Auth, bool, error) { - if !homeDispatch { - refreshed, ok := m.tryRefreshAfterUnauthorized(ctx, auth, execErr, alreadyTried) - return refreshed, ok, nil - } - if m == nil || executor == nil || auth == nil || alreadyTried || execErr == nil { - return auth, false, nil - } - if !isUnauthorizedError(execErr) || auth.AuthKind() != AuthKindOAuth { - return auth, false, nil - } - - log.Debugf("unauthorized Home response for %s (%s), refreshing credentials before redispatch", auth.Provider, auth.ID) - target := auth.Clone() - updated, errRefresh := executor.Refresh(ctx, target) - if errRefresh != nil { - log.Debugf("Home credential refresh before redispatch failed for %s (%s): %v", auth.Provider, auth.ID, errRefresh) - return auth, false, errRefresh - } - if updated == nil { - updated = target - } - if updated.ID == "" { - updated.ID = auth.ID - } - if updated.Index == "" { - updated.Index = auth.Index - } - if updated.Provider == "" { - updated.Provider = auth.Provider - } - if updated.Runtime == nil { - updated.Runtime = auth.Runtime - } - return updated, true, nil -} - -// tryRefreshAfterUnauthorized refreshes local OAuth credentials once after a -// 401 so the current auth can be retried before fallback/suspend. +// tryRefreshAfterUnauthorized refreshes OAuth credentials once after a 401 so the +// current auth can be retried before fallback/suspend. func (m *Manager) tryRefreshAfterUnauthorized(ctx context.Context, auth *Auth, execErr error, alreadyTried bool) (*Auth, bool) { if m == nil || auth == nil || alreadyTried || execErr == nil { return auth, false diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 30cbb4ca8..81e41b38c 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -1087,15 +1087,14 @@ func (m *Manager) SelectHomeAuthByKind(ctx context.Context, provider string, mod return nil, errSelection } providerMatches := strings.TrimSpace(provider) == "" || strings.EqualFold(strings.TrimSpace(selection.Provider), strings.TrimSpace(provider)) - selectionAuth := selection.CloneAuth() - kindMatches := selectionAuth != nil && selectionAuth.AuthKind() == requiredKind + kindMatches := selection.Auth != nil && selection.Auth.AuthKind() == requiredKind if providerMatches && kindMatches { return selection, nil } authID := "" - if selectionAuth != nil { - authID = strings.TrimSpace(selectionAuth.ID) + if selection.Auth != nil { + authID = strings.TrimSpace(selection.Auth.ID) } reason := "auth_kind_mismatch" if !providerMatches { diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 0d83a9871..be6784af3 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -180,14 +180,6 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} } -func (m *Manager) replaceHomeExecutionLifecycleAuth(lifecycle cliproxyexecutor.ExecutionLifecycle, auth *Auth) { - selection, ok := lifecycle.(*HomeDispatchSelection) - if !ok || selection == nil { - return - } - m.replaceHomeSelectionAuth(selection, auth) -} - func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult, routing *apiKeyModelRoutingSnapshot, allowRetry bool, ephemeralResult bool) (*cliproxyexecutor.StreamResult, error) { if executor == nil { return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} @@ -220,11 +212,8 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi return nil, errCtx } if allowRetry { - if refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, executor, auth, errStream, didRefreshOnUnauthorized, ephemeralResult); errRefresh != nil { - errStream = errRefresh - } else if okRefresh { + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, errStream, didRefreshOnUnauthorized); okRefresh { auth = refreshed - m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) didRefreshOnUnauthorized = true streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) if errStream != nil { @@ -257,14 +246,9 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi return nil, errCtx } if allowRetry { - if refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, executor, auth, bootstrapErr, didRefreshOnUnauthorized, ephemeralResult); errRefresh != nil { - discardStreamChunks(streamResult.Chunks) - bootstrapErr = errRefresh - streamResult = &cliproxyexecutor.StreamResult{} - } else if okRefresh { + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, bootstrapErr, didRefreshOnUnauthorized); okRefresh { discardStreamChunks(streamResult.Chunks) auth = refreshed - m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) didRefreshOnUnauthorized = true retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) if retryErr != nil { diff --git a/sdk/cliproxy/auth/home_selection.go b/sdk/cliproxy/auth/home_selection.go index a9326a804..01a39b322 100644 --- a/sdk/cliproxy/auth/home_selection.go +++ b/sdk/cliproxy/auth/home_selection.go @@ -141,7 +141,6 @@ type HomeDispatchSelection struct { Executor ProviderExecutor Provider string - authMu sync.RWMutex scope *executionregistry.Scope accountedModel string resources *executionResources @@ -250,35 +249,9 @@ func (s *HomeDispatchSelection) EndWithRelease(reason string) *executionregistry return s.scope.EndWithRelease("") } -// ReplaceAuth updates the selection after Home returns refreshed credentials. -func (s *HomeDispatchSelection) ReplaceAuth(auth *Auth) { - if s == nil || auth == nil { - return - } - updated := auth.Clone() - s.authMu.Lock() - defer s.authMu.Unlock() - if s.Auth != nil { - if updated.Attributes == nil { - updated.Attributes = make(map[string]string) - } - for _, key := range []string{homeUpstreamModelAttributeKey, homeForceMappingAttributeKey, homeOriginalAliasAttributeKey} { - if value := strings.TrimSpace(s.Auth.Attributes[key]); value != "" { - updated.Attributes[key] = value - } - } - } - s.Auth = updated -} - // CloneAuth returns a standalone auth copy without the selection handle. func (s *HomeDispatchSelection) CloneAuth() *Auth { - if s == nil { - return nil - } - s.authMu.RLock() - defer s.authMu.RUnlock() - if s.Auth == nil { + if s == nil || s.Auth == nil { return nil } return s.Auth.Clone() diff --git a/sdk/cliproxy/auth/home_selection_test.go b/sdk/cliproxy/auth/home_selection_test.go index 56cbe29d2..1f02fc5cc 100644 --- a/sdk/cliproxy/auth/home_selection_test.go +++ b/sdk/cliproxy/auth/home_selection_test.go @@ -42,70 +42,6 @@ func TestHomeDispatchSelectionOwnsScopeOutsideAuth(t *testing.T) { } } -func TestHomeDispatchSelectionReplaceAuthPreservesRoutingAttributes(t *testing.T) { - selection := &HomeDispatchSelection{Auth: &Auth{ - ID: "cred-1", - Provider: "codex", - Attributes: map[string]string{ - homeUpstreamModelAttributeKey: "gpt-5-upstream", - homeForceMappingAttributeKey: "true", - homeOriginalAliasAttributeKey: "team/gpt-5", - }, - Metadata: map[string]any{"access_token": "old"}, - }} - - selection.ReplaceAuth(&Auth{ - ID: "cred-1", - Provider: "codex", - Attributes: map[string]string{AttributeAuthKind: AuthKindOAuth}, - Metadata: map[string]any{"access_token": "fresh"}, - }) - - updated := selection.CloneAuth() - if updated == nil || updated.Metadata["access_token"] != "fresh" { - t.Fatalf("updated auth = %#v", updated) - } - if updated.Attributes[homeUpstreamModelAttributeKey] != "gpt-5-upstream" || updated.Attributes[homeForceMappingAttributeKey] != "true" || updated.Attributes[homeOriginalAliasAttributeKey] != "team/gpt-5" { - t.Fatalf("routing attributes were not preserved: %#v", updated.Attributes) - } -} - -func TestHomeDispatchSelectionReplaceAuthConcurrentClone(t *testing.T) { - selection := &HomeDispatchSelection{Auth: &Auth{ID: "cred-1", Metadata: map[string]any{"access_token": "old"}}} - done := make(chan struct{}) - go func() { - defer close(done) - for i := 0; i < 1000; i++ { - selection.ReplaceAuth(&Auth{ID: "cred-1", Metadata: map[string]any{"access_token": "fresh"}}) - } - }() - for i := 0; i < 1000; i++ { - if auth := selection.CloneAuth(); auth == nil || auth.ID != "cred-1" { - t.Fatalf("CloneAuth() = %#v", auth) - } - } - <-done -} - -func TestReplaceHomeSelectionAuthUpdatesRetainedRuntimeAuth(t *testing.T) { - selection := &HomeDispatchSelection{Auth: &Auth{ID: "cred-1", Provider: "codex", Metadata: map[string]any{"access_token": "old"}}} - manager := &Manager{ - homeRuntimeAuths: map[string]map[string]*Auth{ - "session-1": {"cred-1": selection.Auth.Clone()}, - }, - homeRuntimeAuthOwners: map[string]map[string]*HomeDispatchSelection{ - "session-1": {"cred-1": selection}, - }, - } - - manager.replaceHomeSelectionAuth(selection, &Auth{ID: "cred-1", Provider: "codex", Metadata: map[string]any{"access_token": "fresh"}}) - - retained := manager.homeRuntimeAuths["session-1"]["cred-1"] - if retained == nil || retained.Metadata["access_token"] != "fresh" { - t.Fatalf("retained runtime auth = %#v, want fresh token", retained) - } -} - func TestHomeDispatchSelectionDrainsResourcesAddedDuringEnd(t *testing.T) { registry := executionregistry.New() pending, errBegin := registry.BeginDispatch() diff --git a/sdk/cliproxy/auth/home_unauthorized_refresh_test.go b/sdk/cliproxy/auth/home_unauthorized_refresh_test.go deleted file mode 100644 index 538046ce8..000000000 --- a/sdk/cliproxy/auth/home_unauthorized_refresh_test.go +++ /dev/null @@ -1,275 +0,0 @@ -package auth - -import ( - "context" - "encoding/json" - "net/http" - "sync/atomic" - "testing" - - internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" - cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" -) - -const homeUnauthorizedRefreshProvider = "home-unauthorized-refresh" - -type homeUnauthorizedRefreshDispatcher struct { - calls atomic.Int32 -} - -func (*homeUnauthorizedRefreshDispatcher) HeartbeatOK() bool { return true } - -func (d *homeUnauthorizedRefreshDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { - d.calls.Add(1) - return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ - ID: "home-refresh-auth", - Provider: homeUnauthorizedRefreshProvider, - Status: StatusActive, - Attributes: map[string]string{ - AttributeAuthKind: AuthKindOAuth, - "websockets": "true", - }, - Metadata: map[string]any{ - "access_token": "stale-access-token", - }, - }}) -} - -func (*homeUnauthorizedRefreshDispatcher) AbortAmbiguousDispatch() {} - -type homeUnauthorizedRefreshExecutor struct { - streamMode string - refreshErr error - retainSelection bool - executeCalls atomic.Int32 - countCalls atomic.Int32 - streamCalls atomic.Int32 - refreshCalls atomic.Int32 -} - -func (*homeUnauthorizedRefreshExecutor) Identifier() string { return homeUnauthorizedRefreshProvider } - -func (e *homeUnauthorizedRefreshExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - e.executeCalls.Add(1) - if e.retainSelection { - if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { - lifecycle.Retain() - } - } - if authAccessToken(auth) == "stale-access-token" { - return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"} - } - return cliproxyexecutor.Response{Payload: []byte("ok")}, nil -} - -func (e *homeUnauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { - e.streamCalls.Add(1) - if authAccessToken(auth) == "stale-access-token" { - switch e.streamMode { - case "bootstrap": - chunks := make(chan cliproxyexecutor.StreamChunk, 1) - chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}} - close(chunks) - return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil - case "started": - chunks := make(chan cliproxyexecutor.StreamChunk, 2) - chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("started")} - chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}} - close(chunks) - return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil - default: - return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"} - } - } - chunks := make(chan cliproxyexecutor.StreamChunk, 1) - chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("ok")} - close(chunks) - return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil -} - -func (e *homeUnauthorizedRefreshExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { - e.refreshCalls.Add(1) - if e.refreshErr != nil { - return nil, e.refreshErr - } - updated := auth.Clone() - if updated.Metadata == nil { - updated.Metadata = make(map[string]any) - } - updated.Metadata["access_token"] = "fresh-access-token" - return updated, nil -} - -func (e *homeUnauthorizedRefreshExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - e.countCalls.Add(1) - if authAccessToken(auth) == "stale-access-token" { - return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"} - } - return cliproxyexecutor.Response{Payload: []byte("ok")}, nil -} - -func (*homeUnauthorizedRefreshExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { - return nil, nil -} - -func newHomeUnauthorizedRefreshManager(dispatcher *homeUnauthorizedRefreshDispatcher, executor *homeUnauthorizedRefreshExecutor) *Manager { - manager := NewManager(nil, nil, nil) - manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) - manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) - manager.RegisterExecutor(executor) - return manager -} - -func TestHomeUnauthorizedRefreshesSameSelectionBeforeRedispatch(t *testing.T) { - for _, test := range []struct { - name string - run func(*Manager) error - }{ - { - name: "execute", - run: func(manager *Manager) error { - _, errExecute := manager.Execute(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}) - return errExecute - }, - }, - { - name: "count_tokens", - run: func(manager *Manager) error { - _, errCount := manager.ExecuteCount(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}) - return errCount - }, - }, - } { - t.Run(test.name, func(t *testing.T) { - dispatcher := &homeUnauthorizedRefreshDispatcher{} - executor := &homeUnauthorizedRefreshExecutor{} - manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) - - if errRun := test.run(manager); errRun != nil { - t.Fatalf("execution error = %v", errRun) - } - if got := dispatcher.calls.Load(); got != 1 { - t.Fatalf("Home dispatch calls = %d, want 1", got) - } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want 1", got) - } - if test.name == "execute" && executor.executeCalls.Load() != 2 { - t.Fatalf("execute calls = %d, want 2", executor.executeCalls.Load()) - } - if test.name == "count_tokens" && executor.countCalls.Load() != 2 { - t.Fatalf("count calls = %d, want 2", executor.countCalls.Load()) - } - }) - } -} - -func TestHomeUnauthorizedRefreshUpdatesRetainedSelection(t *testing.T) { - dispatcher := &homeUnauthorizedRefreshDispatcher{} - executor := &homeUnauthorizedRefreshExecutor{retainSelection: true} - manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) - ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) - opts := cliproxyexecutor.Options{Metadata: map[string]any{ - cliproxyexecutor.ExecutionSessionMetadataKey: "refresh-session", - cliproxyexecutor.PinnedAuthMetadataKey: "home-refresh-auth", - }} - - for range 2 { - if _, errExecute := manager.Execute(ctx, []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil { - t.Fatalf("Execute() error = %v", errExecute) - } - } - if got := dispatcher.calls.Load(); got != 1 { - t.Fatalf("Home dispatch calls = %d, want one retained selection", got) - } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want refreshed token reused by retained selection", got) - } - if got := executor.executeCalls.Load(); got != 3 { - t.Fatalf("execute calls = %d, want stale attempt, retry, and retained reuse", got) - } -} - -func TestHomeUnauthorizedTransientRefreshFailureIsReturned(t *testing.T) { - dispatcher := &homeUnauthorizedRefreshDispatcher{} - executor := &homeUnauthorizedRefreshExecutor{ - refreshErr: &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "Home refresh temporarily unavailable"}, - } - manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) - - _, errExecute := manager.Execute(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}) - if statusCodeFromError(errExecute) != http.StatusServiceUnavailable { - t.Fatalf("Execute() error = %v, want transient 503", errExecute) - } - if got := executor.executeCalls.Load(); got != 1 { - t.Fatalf("execute calls = %d, want 1", got) - } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want 1", got) - } -} - -func TestHomeUnauthorizedStartedStreamDoesNotReplay(t *testing.T) { - dispatcher := &homeUnauthorizedRefreshDispatcher{} - executor := &homeUnauthorizedRefreshExecutor{streamMode: "started"} - manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) - - result, errStream := manager.ExecuteStream(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) - if errStream != nil { - t.Fatalf("ExecuteStream() error = %v", errStream) - } - sawPayload := false - sawUnauthorized := false - for chunk := range result.Chunks { - if string(chunk.Payload) == "started" { - sawPayload = true - } - if statusCodeFromError(chunk.Err) == http.StatusUnauthorized { - sawUnauthorized = true - } - } - if !sawPayload || !sawUnauthorized { - t.Fatalf("stream results = payload %v unauthorized %v, want both", sawPayload, sawUnauthorized) - } - if got := executor.refreshCalls.Load(); got != 0 { - t.Fatalf("refresh calls = %d, want 0 after stream started", got) - } - if got := executor.streamCalls.Load(); got != 1 { - t.Fatalf("stream calls = %d, want 1", got) - } -} - -func TestHomeUnauthorizedStreamRefreshesBeforeRedispatch(t *testing.T) { - for _, mode := range []string{"synchronous", "bootstrap"} { - t.Run(mode, func(t *testing.T) { - dispatcher := &homeUnauthorizedRefreshDispatcher{} - executor := &homeUnauthorizedRefreshExecutor{streamMode: mode} - manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) - - result, errStream := manager.ExecuteStream(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) - if errStream != nil { - t.Fatalf("ExecuteStream() error = %v", errStream) - } - var payload string - for chunk := range result.Chunks { - if chunk.Err != nil { - t.Fatalf("stream chunk error = %v", chunk.Err) - } - payload += string(chunk.Payload) - } - if payload != "ok" { - t.Fatalf("stream payload = %q, want ok", payload) - } - if got := dispatcher.calls.Load(); got != 1 { - t.Fatalf("Home dispatch calls = %d, want 1", got) - } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want 1", got) - } - if got := executor.streamCalls.Load(); got != 2 { - t.Fatalf("stream calls = %d, want 2", got) - } - }) - } -} diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index ca36dc558..7fa604168 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -28,10 +28,8 @@ type Record struct { APIKey string AuthID string AuthIndex string - // AccessTokenSHA256 identifies the OAuth token version without exposing the token. - AccessTokenSHA256 string - AuthType string - Source string + AuthType string + Source string // ReasoningEffort stores the translated upstream thinking level for request event logs. ReasoningEffort string // ServiceTier stores the client-requested service tier.