fix(sso): bound SSO poll loop without IPython; bind cached-token swap to user_id (#1721 follow-up) - #2083
Draft
vaimdev wants to merge 3 commits into
Draft
fix(sso): bound SSO poll loop without IPython; bind cached-token swap to user_id (#1721 follow-up)#2083vaimdev wants to merge 3 commits into
vaimdev wants to merge 3 commits into
Conversation
…swap to user_id Two follow-ups to #1721 (fix/site-wide-org-dataset-error): - sso_login() blocking poll: the notebook countdown imported IPython before the per-iteration sleep/timeout, and the surrounding `except Exception` swallowed the ImportError, so with IPython absent (it is only in the `jupyter` extra) the loop re-polled the token endpoint with zero delay and `sso_timeout` never fired. Newly reachable since SsoPendingException is now retried. Sleep/increment/timeout now run unconditionally; the display is gated on in_ipython() and wrapped in its own try/except. - switch_org() fast path 2 swapped in any unexpired token cached for the org, but the cache is keyed by org only and login()/pkey_login()/ api_token() change the active principal without resetting the session, so a switch could silently reinstate a previous user's token. get_verified_token() gains an opt-in `same_user_as` guard comparing the Hub `user_id` claim of both tokens; tokens without a readable claim never qualify. The same-user multi-org swap-in still works. Tests: no-IPython poll loop terminates at sso_timeout (hangs on the old code); cross-user and claim-less tokens fall through to a real switch; ClientSession guard semantics. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TyjxnzaqeB3SeirofUQhSj
B1 follow-up to #1721: refresh() carried the org grant from the pre-refresh token onto the newly reissued one, but the server resets the active org on every reissue -- so the carried grant let _maybe_switch_org() skip the required /switch/ POST, and plot() silently uploaded to the default org while session.org_name still showed the requested one. Deleting the carry-grant lines alone regresses further: switch_org()'s same-user fast path (added by 5cdd1d6 for B3) would then find the stale cache entry left over from before the rotation and swap that superseded, pre-refresh token back in as active -- no real switch request, and refresh()'s own return value reverted to the old token. Fix: ClientSession.forget_verified_org() drops the cache entry for the active org once refresh() confirms the token actually rotated (skipped on a no-op refresh, so it doesn't discard a still-valid verification), forcing a real switch_org() call with the new token. Tests: test_refresh_rotated_token_does_not_skip_switch pins the mocked-switch_org_request behavior; test_refresh_rotated_token_hits_real_switch_endpoint is the HTTP-level version (only requests.post mocked, real ArrowUploader.refresh()/switch_org_request() code paths run) -- both fail against the pre-fix code. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01B1Pvaw1aZQvtraEWkEvdyt
…n rotation" This reverts commit 602a260. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Follow-up to #1721 (targets
fix/site-wide-org-dataset-error, intended to merge into it). Fixes two blockers found while reviewing that branch. Both are client-only and independent of Hub version.B2 — Blocking SSO poll loop spins with no sleep and no timeout when IPython is not installed
_handle_auth_url()'s countdown didfrom IPython.core.display import HTMLbefore the per-iterationtime.sleep(1)/elapsed_time += 1, inside atrywhoseexcept Exception: token = Noneswallowed theImportError. With IPython absent (it's only in thejupyterextra),elapsed_timestayed at 1,1 % 10 == 1was true every iteration, and the loop re-polledGET /api/v2/o/sso/oidc/jwt/<state>/with zero delay forever —sso_timeoutnever fired. Measured: 201 polls in 3 ms.Newly reachable in #1721: on master
_sso_get_token()raised a plainExceptionon the first "State is invalid" poll, so the countdown branch was effectively dead; the newSsoPendingExceptionretry makes it live.Fix: sleep / increment / timeout now run unconditionally each iteration; the notebook HTML countdown is gated on
in_ipython()and wrapped in its owntry/except.SsoRetrieveTokenTimeoutExceptionwas only used as local control flow here, so the import is dropped.B3 —
switch_org()fast path 2 reinstated a token belonging to a different user_verified_org_tokensis keyed by org only and is reset only byregister()(freshClientSession). Publiclogin(),pkey_login(),sso_login()and theapi_token()setter never touch it. So:Subsequent uploads ran as user A. Same-user variant: a rotated/revoked old token got reinstated.
Fix:
ClientSession.get_verified_token()gains an opt-insame_user_asguard that returns the cached token only when both it and the current token carry a readable Hubuser_idclaim and they match; tokens without the claim (personal keys, opaque tokens) never qualify.switch_org()passes the current token. The multi-org same-user swap-in that #1721 introduced still works (existing test updated to give both tokens the sameuser_id).refresh()'s unguardedget_verified_token()call is left as-is — that's B1 (server resets the active org on refresh), which is a separate decision.Tests
TestSsoPollErrorHandling::test_pending_state_stops_at_timeout_without_ipython— blocksIPython*insys.modules; asserts 3 polls and 4 sleeps forsso_timeout=3. Hangs on the pre-fix code.test_switch_org_does_not_reuse_cached_token_of_another_user,test_switch_org_does_not_reuse_cached_token_without_user_claim— fall through to a realswitch_org_request, active token untouched.TestVerifiedOrgTokenCache::test_same_user_guard_*— guard semantics incl. int/struser_idequivalence and opt-in behavior.ruff / mypy / type-hygiene / comment-density guards clean; 109 tests pass across the four touched test files.
🤖 Generated with Claude Code
https://claude.ai/code/session_01TyjxnzaqeB3SeirofUQhSj