From 25a9d49c550ba4ed2b67db666c1d763f9aa7c113 Mon Sep 17 00:00:00 2001
From: vaimdevs <vaimdevs@v-aim.com>
Date: Fri, 10 Jul 2026 11:54:45 +0800
Subject: [PATCH 01/11] Fix

---
 graphistry/arrow_uploader.py            | 45 +++++++++++++--
 graphistry/tests/test_arrow_uploader.py | 73 +++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 5 deletions(-)

diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py
index 91cb6b5f63..b8b24aa61c 100644
--- a/graphistry/arrow_uploader.py
+++ b/graphistry/arrow_uploader.py
@@ -1,6 +1,6 @@
 from typing import List, Optional, Dict, Any
 
-import base64, io, json, pyarrow as pa, requests, sys
+import base64, io, json, pyarrow as pa, re, requests, sys, unicodedata
 
 from graphistry.privacy import Mode, Privacy, ModeAction
 from graphistry.otel import inject_trace_headers
@@ -22,8 +22,21 @@
 logger = setup_logger(__name__)
 
 
+def _django_slugify(value: str) -> str:
+    """Reimplementation of django.utils.text.slugify(value, allow_unicode=False).
+
+    Server contract: ``Organization.save()`` (common/models.py) sets a personal
+    org's slug via this exact transform on the owner's username — not the raw
+    username. Must stay byte-for-byte consistent with Django's algorithm or
+    slug lookups derived from it will 404.
+    """
+    value = unicodedata.normalize('NFKD', value).encode('ascii', 'ignore').decode('ascii')
+    value = re.sub(r'[^\w\s-]', '', value.lower())
+    return re.sub(r'[-\s]+', '-', value).strip('-_')
+
+
 def _personal_org_from_jwt(token: str) -> Optional[str]:
-    """Extract the username claim from a JWT payload, used as a personal-org slug.
+    """Derive the personal-org slug from the username claim in a JWT payload.
 
     Trust chain: callers pass a JWT just received from the authenticated
     /api/v2/o/sso/oidc/jwt/{state}/ endpoint in the same exchange, so we decode
@@ -33,8 +46,11 @@ def _personal_org_from_jwt(token: str) -> Optional[str]:
     rejects), not grant unauthorized access. Do NOT reuse this helper for
     tokens received from outside that trust chain.
 
-    Server contract: ``personal_org.slug == jwt_payload.username`` for users
-    auto-provisioned via SSO.
+    Server contract: ``personal_org.slug == django_slugify(jwt_payload.username)``
+    for users auto-provisioned via SSO (see _django_slugify). Note the server
+    also appends a numeric suffix on slug collision (common/models.py:315-317),
+    which this best-effort reconstruction cannot replicate — collision cases
+    still fall through to the server's 404, same as before this derivation existed.
 
     Returns None on any decode/parse failure or missing/non-string username.
     """
@@ -50,7 +66,10 @@ def _personal_org_from_jwt(token: str) -> Optional[str]:
         if not isinstance(payload, dict):
             return None
         username = payload.get('username')
-        return username if isinstance(username, str) and username else None
+        if not isinstance(username, str) or not username:
+            return None
+        slug = _django_slugify(username)
+        return slug or None
     except Exception as exc:
         logger.debug("@_personal_org_from_jwt: failed to extract username: %s", exc)
         return None
@@ -464,6 +483,22 @@ def sso_get_token(self, state):
             active_org = data.get('active_organization')
             slug = active_org.get('slug') if isinstance(active_org, dict) else None
 
+            # Defense-in-depth: nexus's SSO claim-to-org resolver can (bug)
+            # bind a user to the reserved site-wide sentinel org (slug
+            # 'SITE') when an IdP org claim happens to case-insensitively
+            # match it. No dataset can ever be created under that org (server
+            # rejects with 403), so treat it the same as "no org bound" and
+            # fall through to the JWT-derived personal-org fallback below.
+            if isinstance(slug, str) and slug.upper() == 'SITE':
+                logger.warning(
+                    "SSO returned active_organization=%r, the reserved "
+                    "site-wide org; treating as unbound and falling back to "
+                    "personal org. This usually means an IdP org claim was "
+                    "mis-mapped server-side — verify per-org SSO config.",
+                    slug
+                )
+                slug = None
+
             if slug:
                 # Layer 1: server-bound active_organization. Caller's intent
                 # (self.org_name from register(org_name=...) or session) must
diff --git a/graphistry/tests/test_arrow_uploader.py b/graphistry/tests/test_arrow_uploader.py
index 867ea66024..8ba4b5bdb7 100644
--- a/graphistry/tests/test_arrow_uploader.py
+++ b/graphistry/tests/test_arrow_uploader.py
@@ -671,6 +671,79 @@ def test_sso_get_token_layer1_caller_server_match_proceeds(self, mock_get):
         assert au.org_name == 'mock-org'
         mock_switch.assert_called_once_with('mock-org', '123')
 
+    @mock.patch('requests.get')
+    def test_sso_get_token_site_wide_slug_falls_back_to_personal_org(self, mock_get):
+        # Defense-in-depth for the nexus SSO claim-to-org resolver bug: server
+        # binds active_organization to the reserved site-wide sentinel org
+        # (slug 'SITE') instead of a real org or nothing. No dataset can ever
+        # be created under that org (server-side 403), so the client must NOT
+        # accept it — treat as unbound and fall back to the JWT-derived
+        # personal org, same as the Layer-3 no-org-at-all path.
+
+        payload_dict = {'username': 'site-user', 'exp': 9999999999}
+        payload_b64 = base64.urlsafe_b64encode(
+            json.dumps(payload_dict).encode()
+        ).rstrip(b'=').decode()
+        fake_jwt = f"eyJhbGciOiJIUzI1NiJ9.{payload_b64}.placeholder-signature"
+
+        mock_resp = self._mock_response(
+            json_data={
+                'status': 'OK',
+                'message': 'State is valid',
+                'data': {
+                    'token': fake_jwt,
+                    'active_organization': {'slug': 'SITE', 'is_found': True, 'is_member': True},
+                }
+        })
+        mock_get.return_value = mock_resp
+
+        client = PyGraphistry.client()
+        client.session.org_name = None
+        PyGraphistry.session.org_name = None
+
+        au = ArrowUploader(client_session=client.session)
+
+        with mock.patch.object(ArrowUploader, "_switch_org") as mock_switch:
+            with self.assertLogs('graphistry.arrow_uploader', level='WARNING') as log_ctx:
+                au.sso_get_token(state='ignored-valid')
+
+        assert au.token == fake_jwt
+        assert au.org_name == 'site-user'
+        mock_switch.assert_called_once_with('site-user', fake_jwt)
+        assert any(
+            record.levelname == 'WARNING' and 'SITE' in record.getMessage()
+            for record in log_ctx.records
+        ), "Expected a WARNING log surfacing the rejected site-wide org slug"
+
+    @mock.patch('requests.get')
+    def test_sso_get_token_site_wide_slug_case_insensitive(self, mock_get):
+        # Server-side matching is case-insensitive (slug__iexact), so the
+        # client-side guard must be too.
+
+        mock_resp = self._mock_response(
+            json_data={
+                'status': 'OK',
+                'message': 'State is valid',
+                'data': {
+                    'token': '123',
+                    'active_organization': {'slug': 'site', 'is_found': True, 'is_member': True},
+                }
+        })
+        mock_get.return_value = mock_resp
+
+        client = PyGraphistry.client()
+        client.session.org_name = None
+        PyGraphistry.session.org_name = None
+
+        au = ArrowUploader(client_session=client.session)
+
+        with mock.patch.object(ArrowUploader, "_switch_org") as mock_switch:
+            au.sso_get_token(state='ignored-valid')
+
+        assert au.token == '123'
+        assert au.org_name is None
+        mock_switch.assert_not_called()
+
     @mock.patch('requests.get')
     def test_sso_get_token_layer3_jwt_username_fallback(self, mock_get):
         # Layer 3: caller passed nothing, server didn't bind — JWT-derived

From 105a79b8a034e8b1e157471257afe310d252d8ad Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 10 Jul 2026 21:26:58 +0800
Subject: [PATCH 02/11] Fix

---
 graphistry/arrow_uploader.py            | 103 +++++-------
 graphistry/tests/test_arrow_uploader.py | 203 ++++++++++++------------
 2 files changed, 142 insertions(+), 164 deletions(-)

diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py
index b8b24aa61c..22f7774792 100644
--- a/graphistry/arrow_uploader.py
+++ b/graphistry/arrow_uploader.py
@@ -1,6 +1,6 @@
 from typing import List, Optional, Dict, Any
 
-import base64, io, json, pyarrow as pa, re, requests, sys, unicodedata
+import io, json, pyarrow as pa, requests, sys
 
 from graphistry.privacy import Mode, Privacy, ModeAction
 from graphistry.otel import inject_trace_headers
@@ -22,59 +22,6 @@
 logger = setup_logger(__name__)
 
 
-def _django_slugify(value: str) -> str:
-    """Reimplementation of django.utils.text.slugify(value, allow_unicode=False).
-
-    Server contract: ``Organization.save()`` (common/models.py) sets a personal
-    org's slug via this exact transform on the owner's username — not the raw
-    username. Must stay byte-for-byte consistent with Django's algorithm or
-    slug lookups derived from it will 404.
-    """
-    value = unicodedata.normalize('NFKD', value).encode('ascii', 'ignore').decode('ascii')
-    value = re.sub(r'[^\w\s-]', '', value.lower())
-    return re.sub(r'[-\s]+', '-', value).strip('-_')
-
-
-def _personal_org_from_jwt(token: str) -> Optional[str]:
-    """Derive the personal-org slug from the username claim in a JWT payload.
-
-    Trust chain: callers pass a JWT just received from the authenticated
-    /api/v2/o/sso/oidc/jwt/{state}/ endpoint in the same exchange, so we decode
-    the payload without local signature verification. The server re-validates
-    the token signature on every subsequent request — an incorrect username
-    can at worst route the session to an org the user isn't a member of (server
-    rejects), not grant unauthorized access. Do NOT reuse this helper for
-    tokens received from outside that trust chain.
-
-    Server contract: ``personal_org.slug == django_slugify(jwt_payload.username)``
-    for users auto-provisioned via SSO (see _django_slugify). Note the server
-    also appends a numeric suffix on slug collision (common/models.py:315-317),
-    which this best-effort reconstruction cannot replicate — collision cases
-    still fall through to the server's 404, same as before this derivation existed.
-
-    Returns None on any decode/parse failure or missing/non-string username.
-    """
-    try:
-        parts = token.split('.')
-        if len(parts) < 2:
-            return None
-        # base64 padding: only the missing chars (0, 2, or 3); never extra.
-        # Inputs whose stripped length is 1 mod 4 are invalid b64; the decode
-        # call below will raise and the caller-level except returns None.
-        segment = parts[1] + '=' * (-len(parts[1]) % 4)
-        payload = json.loads(base64.urlsafe_b64decode(segment))
-        if not isinstance(payload, dict):
-            return None
-        username = payload.get('username')
-        if not isinstance(username, str) or not username:
-            return None
-        slug = _django_slugify(username)
-        return slug or None
-    except Exception as exc:
-        logger.debug("@_personal_org_from_jwt: failed to extract username: %s", exc)
-        return None
-
-
 class ArrowUploader:
 
     def __init__(
@@ -450,6 +397,40 @@ def sso_login(self, org_name: Optional[str] = None, idp_name: Optional[str] = No
             
         return self
 
+    def _personal_org_from_api(self, token: str) -> Optional[str]:
+        """Look up the caller's personal org via GET /api/v2/my/organizations/.
+
+        Trust chain: token is the one just returned by the SSO exchange in
+        this same flow. Replaces the old JWT-username-slugify reconstruction
+        (unreliable: couldn't replicate server-side slug-collision suffixing
+        and depended on byte-for-byte parity with Django's slugify), with a
+        real server-confirmed lookup, filtered on the ``is_personal`` flag.
+
+        Returns None on any request/parse failure or if no personal org is found.
+        """
+        try:
+            out = requests.get(
+                f'{self.server_base_path}/api/v2/my/organizations/',
+                params={'limit': 1000},
+                verify=self.certificate_validation,
+                headers=inject_trace_headers({'Authorization': f'Bearer {token}'})
+            )
+            if not (200 <= out.status_code < 300):
+                return None
+            payload = out.json()
+            orgs = payload.get('results', payload) if isinstance(payload, dict) else payload
+            if not isinstance(orgs, list):
+                return None
+            for org in orgs:
+                if isinstance(org, dict) and org.get('is_personal'):
+                    slug = org.get('slug')
+                    if isinstance(slug, str) and slug:
+                        return slug
+            return None
+        except Exception as exc:
+            logger.debug("@_personal_org_from_api: failed to fetch personal org: %s", exc)
+            return None
+
     def sso_get_token(self, state):
         """
         Koa, 04 May 2022    Use state to get token
@@ -530,22 +511,22 @@ def sso_get_token(self, state):
                 )
             else:
                 # Layer 3: caller didn't ask, server didn't bind. Try
-                # JWT-derived personal-org fallback for first-login UX. See
-                # _personal_org_from_jwt for the trust-chain rationale.
-                fallback = _personal_org_from_jwt(token_value)
+                # server-confirmed personal-org fallback for first-login UX.
+                # See _personal_org_from_api for the trust-chain rationale.
+                fallback = self._personal_org_from_api(token_value)
                 if fallback:
                     logger.info(
                         "SSO did not bind active_organization; falling back to "
-                        "JWT-derived personal org=%s", fallback
+                        "personal org=%s", fallback
                     )
                     self.org_name = fallback
                     self._switch_org(fallback, token_value)
                 else:
                     # Layer 4: nothing claimed, nothing bound, nothing inferable.
                     logger.info(
-                        "SSO did not bind active_organization and no JWT "
-                        "username present; site-wide SSO login completes with "
-                        "no org binding."
+                        "SSO did not bind active_organization and no personal "
+                        "org found; site-wide SSO login completes with no org "
+                        "binding."
                     )
 
         except Exception as e:
diff --git a/graphistry/tests/test_arrow_uploader.py b/graphistry/tests/test_arrow_uploader.py
index 8ba4b5bdb7..b975e9b6d8 100644
--- a/graphistry/tests/test_arrow_uploader.py
+++ b/graphistry/tests/test_arrow_uploader.py
@@ -677,25 +677,29 @@ def test_sso_get_token_site_wide_slug_falls_back_to_personal_org(self, mock_get)
         # binds active_organization to the reserved site-wide sentinel org
         # (slug 'SITE') instead of a real org or nothing. No dataset can ever
         # be created under that org (server-side 403), so the client must NOT
-        # accept it — treat as unbound and fall back to the JWT-derived
-        # personal org, same as the Layer-3 no-org-at-all path.
+        # accept it — treat as unbound and fall back to the server-confirmed
+        # personal org (via GET /api/v2/my/organizations/), same as the
+        # Layer-3 no-org-at-all path.
 
-        payload_dict = {'username': 'site-user', 'exp': 9999999999}
-        payload_b64 = base64.urlsafe_b64encode(
-            json.dumps(payload_dict).encode()
-        ).rstrip(b'=').decode()
-        fake_jwt = f"eyJhbGciOiJIUzI1NiJ9.{payload_b64}.placeholder-signature"
-
-        mock_resp = self._mock_response(
+        sso_resp = self._mock_response(
             json_data={
                 'status': 'OK',
                 'message': 'State is valid',
                 'data': {
-                    'token': fake_jwt,
+                    'token': '123',
                     'active_organization': {'slug': 'SITE', 'is_found': True, 'is_member': True},
                 }
         })
-        mock_get.return_value = mock_resp
+        orgs_resp = self._mock_response(json_data={
+            'results': [
+                {'slug': 'acme-team', 'is_personal': False},
+                {'slug': 'site-user-personal', 'is_personal': True},
+            ]
+        })
+
+        def _route(url, *args, **kwargs):
+            return orgs_resp if '/api/v2/my/organizations/' in url else sso_resp
+        mock_get.side_effect = _route
 
         client = PyGraphistry.client()
         client.session.org_name = None
@@ -707,9 +711,9 @@ def test_sso_get_token_site_wide_slug_falls_back_to_personal_org(self, mock_get)
             with self.assertLogs('graphistry.arrow_uploader', level='WARNING') as log_ctx:
                 au.sso_get_token(state='ignored-valid')
 
-        assert au.token == fake_jwt
-        assert au.org_name == 'site-user'
-        mock_switch.assert_called_once_with('site-user', fake_jwt)
+        assert au.token == '123'
+        assert au.org_name == 'site-user-personal'
+        mock_switch.assert_called_once_with('site-user-personal', '123')
         assert any(
             record.levelname == 'WARNING' and 'SITE' in record.getMessage()
             for record in log_ctx.records
@@ -745,26 +749,28 @@ def test_sso_get_token_site_wide_slug_case_insensitive(self, mock_get):
         mock_switch.assert_not_called()
 
     @mock.patch('requests.get')
-    def test_sso_get_token_layer3_jwt_username_fallback(self, mock_get):
-        # Layer 3: caller passed nothing, server didn't bind — JWT-derived
-        # personal-org slug is used. Pins the first-login-UX path that #1230
-        # contributed to the unified design.
-
-        # Construct a JWT with username="newuser". No signature verification
-        # happens client-side, so the signature segment is just a placeholder.
-        payload_dict = {'user_id': 42, 'username': 'newuser', 'exp': 9999999999}
-        payload_b64 = base64.urlsafe_b64encode(
-            json.dumps(payload_dict).encode()
-        ).rstrip(b'=').decode()
-        fake_jwt = f"eyJhbGciOiJIUzI1NiJ9.{payload_b64}.placeholder-signature"
+    def test_sso_get_token_layer3_personal_org_api_fallback(self, mock_get):
+        # Layer 3: caller passed nothing, server didn't bind — server-confirmed
+        # personal-org slug (via GET /api/v2/my/organizations/) is used.
+        # Replaces the old JWT-username-slugify reconstruction with a real
+        # lookup filtered on is_personal.
 
-        mock_resp = self._mock_response(
+        sso_resp = self._mock_response(
             json_data={
                 'status': 'OK',
                 'message': 'State is valid',
-                'data': {'token': fake_jwt},
+                'data': {'token': '123'},
         })
-        mock_get.return_value = mock_resp
+        orgs_resp = self._mock_response(json_data={
+            'results': [
+                {'slug': 'other-team', 'is_personal': False},
+                {'slug': 'newuser', 'is_personal': True},
+            ]
+        })
+
+        def _route(url, *args, **kwargs):
+            return orgs_resp if '/api/v2/my/organizations/' in url else sso_resp
+        mock_get.side_effect = _route
 
         client = PyGraphistry.client()
         client.session.org_name = None
@@ -775,9 +781,9 @@ def test_sso_get_token_layer3_jwt_username_fallback(self, mock_get):
         with mock.patch.object(ArrowUploader, "_switch_org") as mock_switch:
             au.sso_get_token(state='ignored-valid')
 
-        assert au.token == fake_jwt
+        assert au.token == '123'
         assert au.org_name == 'newuser'
-        mock_switch.assert_called_once_with('newuser', fake_jwt)
+        mock_switch.assert_called_once_with('newuser', '123')
 
     @mock.patch('requests.get')
     def test_sso_get_token_layer2_takes_precedence_over_layer3(self, mock_get):
@@ -810,75 +816,66 @@ def test_sso_get_token_layer2_takes_precedence_over_layer3(self, mock_get):
         mock_switch.assert_not_called()
 
 
-def _make_test_jwt(payload_dict: dict) -> str:
-    payload_b64 = base64.urlsafe_b64encode(
-        json.dumps(payload_dict).encode()
-    ).rstrip(b'=').decode()
-    return f"eyJhbGciOiJIUzI1NiJ9.{payload_b64}.placeholder-signature"
-
-
-class TestPersonalOrgFromJwt(unittest.TestCase):
-    """Unit tests for _personal_org_from_jwt — exercises decode/parse failure
-    modes in isolation, separately from the sso_get_token integration tests."""
-
-    def test_valid_jwt_with_username_returns_username(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        token = _make_test_jwt({'username': 'alice', 'exp': 9999999999})
-        assert _personal_org_from_jwt(token) == 'alice'
-
-    def test_jwt_with_no_dot_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        assert _personal_org_from_jwt('not-a-jwt') is None
-
-    def test_empty_string_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        assert _personal_org_from_jwt('') is None
-
-    def test_jwt_with_only_one_segment_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        assert _personal_org_from_jwt('header-only') is None
-
-    def test_malformed_base64_payload_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        # Use chars that are not valid base64 (e.g. "!!!")
-        assert _personal_org_from_jwt('header.!!!invalid_b64!!!.sig') is None
-
-    def test_non_json_payload_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        not_json = base64.urlsafe_b64encode(b'not even close to json').rstrip(b'=').decode()
-        assert _personal_org_from_jwt(f"header.{not_json}.sig") is None
-
-    def test_non_dict_payload_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        # JWT with a JSON ARRAY (not object) as payload — uncommon but possible.
-        arr_payload = base64.urlsafe_b64encode(json.dumps(['not', 'a', 'dict']).encode()).rstrip(b'=').decode()
-        assert _personal_org_from_jwt(f"header.{arr_payload}.sig") is None
-
-    def test_payload_without_username_field_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        token = _make_test_jwt({'user_id': 1, 'email': 'a@b.com', 'exp': 9999999999})
-        assert _personal_org_from_jwt(token) is None
-
-    def test_payload_with_non_string_username_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        token = _make_test_jwt({'username': 12345, 'exp': 9999999999})
-        assert _personal_org_from_jwt(token) is None
-
-    def test_payload_with_empty_string_username_returns_none(self):
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        token = _make_test_jwt({'username': '', 'exp': 9999999999})
-        assert _personal_org_from_jwt(token) is None
-
-    def test_payload_length_multiple_of_4_after_rstrip_decodes_correctly(self):
-        # Pin the corrected b64 padding formula: when stripped len % 4 == 0,
-        # we must add ZERO pad chars (not 4). Use a payload that produces a
-        # rstripped b64 of length multiple of 4.
-        from graphistry.arrow_uploader import _personal_org_from_jwt
-        # 9-byte body → b64 length 12 → rstrip removes 0 pads → len 12, mod4 0
-        payload_dict = {'username': 'u'}  # tiny payload
-        # Pad the payload until rstripped length is a multiple of 4. A 12-char
-        # rstripped output corresponds to a 9-byte payload. {"username":"u"} is
-        # 16 bytes — let's just trust the formula and exercise it explicitly.
-        token = _make_test_jwt(payload_dict)
-        # The exact byte length will vary; the corrected formula is robust to all.
-        assert _personal_org_from_jwt(token) == 'u'
+class TestPersonalOrgFromApi(unittest.TestCase):
+    """Unit tests for ArrowUploader._personal_org_from_api — exercises the
+    GET /api/v2/my/organizations/ lookup in isolation, separately from the
+    sso_get_token integration tests."""
+
+    def _mock_response(self, json_data, status_code=200):
+        mock_resp = mock.Mock()
+        mock_resp.status_code = status_code
+        mock_resp.json.return_value = json_data
+        mock_resp.text = json.dumps(json_data)
+        return mock_resp
+
+    @mock.patch('requests.get')
+    def test_returns_slug_of_is_personal_org(self, mock_get):
+        mock_get.return_value = self._mock_response({
+            'results': [
+                {'slug': 'acme-team', 'is_personal': False},
+                {'slug': 'alice-personal', 'is_personal': True},
+            ]
+        })
+        au = ArrowUploader()
+        assert au._personal_org_from_api('tok') == 'alice-personal'
+
+    @mock.patch('requests.get')
+    def test_non_paginated_list_response(self, mock_get):
+        # Defensive shape guard: tolerate a plain list too (no 'results' key).
+        mock_get.return_value = self._mock_response([
+            {'slug': 'alice-personal', 'is_personal': True},
+        ])
+        au = ArrowUploader()
+        assert au._personal_org_from_api('tok') == 'alice-personal'
+
+    @mock.patch('requests.get')
+    def test_no_personal_org_in_results_returns_none(self, mock_get):
+        mock_get.return_value = self._mock_response({
+            'results': [{'slug': 'acme-team', 'is_personal': False}]
+        })
+        au = ArrowUploader()
+        assert au._personal_org_from_api('tok') is None
+
+    @mock.patch('requests.get')
+    def test_empty_results_returns_none(self, mock_get):
+        mock_get.return_value = self._mock_response({'results': []})
+        au = ArrowUploader()
+        assert au._personal_org_from_api('tok') is None
+
+    @mock.patch('requests.get')
+    def test_non_200_status_returns_none(self, mock_get):
+        mock_get.return_value = self._mock_response({'detail': 'unauthorized'}, status_code=401)
+        au = ArrowUploader()
+        assert au._personal_org_from_api('tok') is None
+
+    @mock.patch('requests.get')
+    def test_unexpected_shape_returns_none(self, mock_get):
+        mock_get.return_value = self._mock_response({'results': 'not-a-list'})
+        au = ArrowUploader()
+        assert au._personal_org_from_api('tok') is None
+
+    @mock.patch('requests.get')
+    def test_request_exception_returns_none(self, mock_get):
+        mock_get.side_effect = Exception('connection error')
+        au = ArrowUploader()
+        assert au._personal_org_from_api('tok') is None

From 1e7672b3f09594ac46c4745cbfd93c6e3d7b93f2 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 15 Jul 2026 17:19:16 +0800
Subject: [PATCH 03/11] fix(sso): detect SSO-challenge response in
 switch_org(), don't claim false success

switch_org() treated any status=='OK' response as a completed switch, but
the server overloads that status for two different outcomes: a real switch
(data has pk/org_name/organization_slug) and an SSO re-authentication
challenge (data['idp'] populated) for members of SSO-protected orgs who
aren't the owner/admin. The challenge case never actually updates the
server-side active org, so treating it as success silently left callers on
their previous org while locally claiming to have switched.

Now raises with the idp names and a pointer to sso_login(org_name=...) as
the alternative, instead of logging a false "Switched to organization" and
setting self.session.org_name to a value the server never accepted.
---
 graphistry/pygraphistry.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 6416069369..ce4f903ae3 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2710,6 +2710,27 @@ def switch_org(self, value: str):
         result = self._handle_api_response(response)
 
         if result is True:
+            # Server-side quirk: for an org the caller is a plain member of
+            # (not owner/admin) that has its own SSO IDP configured, the
+            # switch endpoint still returns status=='OK' but does NOT
+            # actually switch -- it instead returns an SSO challenge
+            # (data['idp']) requiring a fresh org-scoped SSO login before the
+            # switch is honored. Treating that as success would silently
+            # leave the caller on their previous org.
+            try:
+                data = response.json().get('data', {})
+            except Exception:
+                data = {}
+            if isinstance(data, dict) and data.get('idp'):
+                idp_names = list(data['idp'].keys())
+                raise Exception(
+                    "Switching to organization '{}' requires SSO re-authentication "
+                    "(idp: {}) -- the server did not actually switch. Complete the "
+                    "auth_url in the response's data['idp'], or use "
+                    "sso_login(org_name='{}') instead of switch_org().".format(
+                        value.strip(), idp_names, value.strip()
+                    )
+                )
             self.session.org_name = value.strip()
             logger.info("Switched to organization: {}".format(value.strip()))
         else:  # print the error message

From e4d4e89a692abf5a984be96e2d49bd287a8345c5 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 22 Jul 2026 17:54:37 +0800
Subject: [PATCH 04/11] fix(sso): don't memoize failed/challenged org switch in
 _switch_org

_switch_org wrote _last_switched_org_token unconditionally after the
POST, so a 403/500 or an SSO-challenge response (200 + data['idp'])
got recorded as a completed switch. That poisoned memo made
_maybe_switch_org short-circuit later, skipping switch_org() -- the
only implementation that validates the response and raises on an SSO
challenge. Now _switch_org returns without writing the memo on a
non-2xx status or an idp challenge.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 graphistry/arrow_uploader.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py
index 22f7774792..b3b6ba5eab 100644
--- a/graphistry/arrow_uploader.py
+++ b/graphistry/arrow_uploader.py
@@ -247,6 +247,22 @@ def _switch_org(self, org_name: Optional[str], token: Optional[str]) -> None:
                 verify=self.certificate_validation,
             )
             log_requests_error(response)
+            if not (200 <= response.status_code < 300):
+                logger.warning(
+                    "Org switch to %s failed with HTTP %s; not recording switch",
+                    org_name, response.status_code
+                )
+                return
+            try:
+                body = response.json().get('data', {})
+            except Exception:
+                body = {}
+            if isinstance(body, dict) and body.get('idp'):
+                logger.warning(
+                    "Org switch to %s returned an SSO challenge; not recording switch",
+                    org_name
+                )
+                return
             self._client_session._last_switched_org_token = (org_name, token)
             from .pygraphistry import PyGraphistry
             if PyGraphistry.session is self._client_session:

From 17be0e36388050dd65af90bc5e9a24d94b66f80a Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 28 Aug 2026 11:23:54 +0800
Subject: [PATCH 05/11] refactor(org-switch): unify duplicate switch-org
 request/validation logic

ArrowUploader._switch_org and GraphistryClient.switch_org each POSTed to
the same org-switch endpoint with their own copy of the HTTP-status,
body-status, and SSO-idp-challenge checks. Extract the shared logic into
switch_org_request() in graphistry/utils/requests.py, raising
OrgSwitchError / OrgSwitchIdpChallenge. Each caller keeps its own
success/failure policy (ArrowUploader logs+swallows, GraphistryClient
raises) since that divergence is intentional API contract, not
duplication.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TNGYwRNwAFQdYDucbwgtUm
---
 graphistry/arrow_uploader.py | 43 +++++++---------------
 graphistry/pygraphistry.py   | 52 ++++++--------------------
 graphistry/utils/requests.py | 71 ++++++++++++++++++++++++++++++++++++
 3 files changed, 96 insertions(+), 70 deletions(-)

diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py
index b3b6ba5eab..9d763d7cb1 100644
--- a/graphistry/arrow_uploader.py
+++ b/graphistry/arrow_uploader.py
@@ -16,7 +16,7 @@
 
 from .exceptions import TokenExpireException
 from .validate.validate_encodings import validate_encodings
-from .utils.requests import log_requests_error
+from .utils.requests import log_requests_error, switch_org_request, OrgSwitchError, OrgSwitchIdpChallenge
 from .util import setup_logger
 from graphistry.models.types import ValidationParam
 logger = setup_logger(__name__)
@@ -239,36 +239,21 @@ def _switch_org(self, org_name: Optional[str], token: Optional[str]) -> None:
         if last == (org_name, token):
             return
         try:
-            switch_url = f"{self.server_base_path}/api/v2/o/{org_name}/switch/"
-            response = requests.post(
-                switch_url,
-                data={'slug': org_name},
-                headers=inject_trace_headers({'Authorization': f'Bearer {token}'}),
-                verify=self.certificate_validation,
-            )
-            log_requests_error(response)
-            if not (200 <= response.status_code < 300):
-                logger.warning(
-                    "Org switch to %s failed with HTTP %s; not recording switch",
-                    org_name, response.status_code
-                )
-                return
-            try:
-                body = response.json().get('data', {})
-            except Exception:
-                body = {}
-            if isinstance(body, dict) and body.get('idp'):
-                logger.warning(
-                    "Org switch to %s returned an SSO challenge; not recording switch",
-                    org_name
-                )
-                return
-            self._client_session._last_switched_org_token = (org_name, token)
-            from .pygraphistry import PyGraphistry
-            if PyGraphistry.session is self._client_session:
-                PyGraphistry.session._last_switched_org_token = (org_name, token)
+            switch_org_request(self.server_base_path, org_name, token, self.certificate_validation)
+        except OrgSwitchError as exc:
+            logger.warning("%s; not recording switch", exc)
+            return
+        except OrgSwitchIdpChallenge as exc:
+            logger.warning("%s; not recording switch", exc)
+            return
         except Exception as exc:
             logger.warning("Failed to switch organization %s: %s", org_name, exc)
+            return
+
+        self._client_session._last_switched_org_token = (org_name, token)
+        from .pygraphistry import PyGraphistry
+        if PyGraphistry.session is self._client_session:
+            PyGraphistry.session._last_switched_org_token = (org_name, token)
 
 
     def login(self, username, password, org_name=None):
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index ce4f903ae3..3609a53369 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -1,7 +1,7 @@
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union, cast, overload
 from typing_extensions import Literal
 from graphistry.privacy import Mode, ModeAction
-from graphistry.utils.requests import log_requests_error
+from graphistry.utils.requests import log_requests_error, switch_org_request, OrgSwitchError, OrgSwitchIdpChallenge
 from graphistry.plugins_types.hypergraph import HypergraphResult
 from graphistry.plugins_types.gexf_types import GexfEdgeViz, GexfNodeViz, GexfParseEngine
 from graphistry.client_session import (
@@ -2551,10 +2551,8 @@ def _viz_url(self, info: DatasetInfo, url_params: URLParamsDict) -> str:
             extra,
         )
 
-    def _switch_org_url(self, org_name):
-        hostname = self.session.hostname
-        protocol = self.session.protocol
-        return "{}://{}/api/v2/o/{}/switch/".format(protocol, hostname, org_name)
+    def _base_url(self) -> str:
+        return "{}://{}".format(self.session.protocol, self.session.hostname)
 
 
     def layout_settings(self, 
@@ -2699,42 +2697,14 @@ def personal_key_secret(self, value: Optional[str] = None):
         self.session.personal_key_secret = value.strip()
 
     def switch_org(self, value: str):
-        # print(self._switch_org_url(value))
-        response = requests.post(
-            self._switch_org_url(value),
-            data={'slug': value},
-            headers=inject_trace_headers({'Authorization': f'Bearer {self.api_token()}'}),
-            verify=self.session.certificate_validation,
-        )
-        log_requests_error(response)
-        result = self._handle_api_response(response)
-
-        if result is True:
-            # Server-side quirk: for an org the caller is a plain member of
-            # (not owner/admin) that has its own SSO IDP configured, the
-            # switch endpoint still returns status=='OK' but does NOT
-            # actually switch -- it instead returns an SSO challenge
-            # (data['idp']) requiring a fresh org-scoped SSO login before the
-            # switch is honored. Treating that as success would silently
-            # leave the caller on their previous org.
-            try:
-                data = response.json().get('data', {})
-            except Exception:
-                data = {}
-            if isinstance(data, dict) and data.get('idp'):
-                idp_names = list(data['idp'].keys())
-                raise Exception(
-                    "Switching to organization '{}' requires SSO re-authentication "
-                    "(idp: {}) -- the server did not actually switch. Complete the "
-                    "auth_url in the response's data['idp'], or use "
-                    "sso_login(org_name='{}') instead of switch_org().".format(
-                        value.strip(), idp_names, value.strip()
-                    )
-                )
-            self.session.org_name = value.strip()
-            logger.info("Switched to organization: {}".format(value.strip()))
-        else:  # print the error message
-            raise Exception(result)
+        org_name = value.strip()
+        try:
+            switch_org_request(self._base_url(), org_name, self.api_token(), self.session.certificate_validation)
+        except (OrgSwitchError, OrgSwitchIdpChallenge) as exc:
+            raise Exception(str(exc))
+
+        self.session.org_name = org_name
+        logger.info("Switched to organization: {}".format(org_name))
 
     def _handle_api_response(self, response):
         try:
diff --git a/graphistry/utils/requests.py b/graphistry/utils/requests.py
index 1ec9912512..b9269966c1 100644
--- a/graphistry/utils/requests.py
+++ b/graphistry/utils/requests.py
@@ -1,6 +1,7 @@
 import json
 import requests
 
+from graphistry.otel import inject_trace_headers
 from graphistry.util import setup_logger
 logger = setup_logger(__name__)
 
@@ -14,3 +15,73 @@ def log_requests_error(resp: requests.Response) -> None:
             logger.error("HTTP %s error - response content (JSON): %s", resp.status_code, json.dumps(error_content, indent=2))
         except json.JSONDecodeError:
             logger.error("HTTP %s error - response content (text): %s", resp.status_code, resp.text)
+
+
+class OrgSwitchError(Exception):
+    """Org switch endpoint responded with a non-2xx status or a non-OK body status."""
+
+    def __init__(self, org_name: str, status_code: int, detail: str = ''):
+        self.org_name = org_name
+        self.status_code = status_code
+        self.detail = detail
+        msg = f"Org switch to '{org_name}' failed with HTTP {status_code}"
+        if detail:
+            msg += f": {detail}"
+        super().__init__(msg)
+
+
+class OrgSwitchIdpChallenge(Exception):
+    """Server-side quirk: for an org the caller is a plain member of (not owner/admin)
+    that has its own SSO IDP configured, the switch endpoint still returns
+    status=='OK' but does NOT actually switch -- it instead returns an SSO
+    challenge (data['idp']) requiring a fresh org-scoped SSO login before the
+    switch is honored. Treating that as success would silently leave the
+    caller on their previous org.
+    """
+
+    def __init__(self, org_name: str, idp: dict):
+        self.org_name = org_name
+        self.idp = idp
+        super().__init__(
+            "Switching to organization '{}' requires SSO re-authentication "
+            "(idp: {}) -- the server did not actually switch. Complete the "
+            "auth_url in the response's data['idp'], or use "
+            "sso_login(org_name='{}') instead of switch_org().".format(
+                org_name, list(idp.keys()), org_name
+            )
+        )
+
+
+def switch_org_request(base_url: str, org_name: str, token: str, verify: bool) -> None:
+    """POST to the org-switch endpoint and validate the response.
+
+    Single source of truth for both GraphistryClient.switch_org and
+    ArrowUploader._switch_org: same URL shape, same HTTP/body status check,
+    same SSO-idp-challenge detection. Raises OrgSwitchError or
+    OrgSwitchIdpChallenge on any failure; returns None on a confirmed switch.
+    Callers decide their own success/failure side effects (raising vs.
+    logging-and-swallowing, which session state to update).
+    """
+    switch_url = f"{base_url}/api/v2/o/{org_name}/switch/"
+    response = requests.post(
+        switch_url,
+        data={'slug': org_name},
+        headers=inject_trace_headers({'Authorization': f'Bearer {token}'}),
+        verify=verify,
+    )
+    log_requests_error(response)
+
+    if not (200 <= response.status_code < 300):
+        raise OrgSwitchError(org_name, response.status_code)
+
+    try:
+        body = response.json()
+    except Exception:
+        body = {}
+
+    if body.get('status') != 'OK':
+        raise OrgSwitchError(org_name, response.status_code, detail=body.get('message', ''))
+
+    data = body.get('data', {})
+    if isinstance(data, dict) and data.get('idp'):
+        raise OrgSwitchIdpChallenge(org_name, data['idp'])

From 11d25b84c593b0796e847fa31fb063e914813b71 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 28 Aug 2026 14:33:09 +0800
Subject: [PATCH 06/11] fix(hygiene): add return type annotations and
 parameterize dict in requests.py

Type-hygiene guard flagged missing __init__ return annotations and a bare
dict for OrgSwitchIdpChallenge's idp param.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GgDZr3AgU9uKz5EWcVZZVY
---
 graphistry/utils/requests.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/graphistry/utils/requests.py b/graphistry/utils/requests.py
index b9269966c1..8d93766e0d 100644
--- a/graphistry/utils/requests.py
+++ b/graphistry/utils/requests.py
@@ -1,4 +1,6 @@
 import json
+from typing import Any, Dict
+
 import requests
 
 from graphistry.otel import inject_trace_headers
@@ -20,7 +22,7 @@ def log_requests_error(resp: requests.Response) -> None:
 class OrgSwitchError(Exception):
     """Org switch endpoint responded with a non-2xx status or a non-OK body status."""
 
-    def __init__(self, org_name: str, status_code: int, detail: str = ''):
+    def __init__(self, org_name: str, status_code: int, detail: str = '') -> None:
         self.org_name = org_name
         self.status_code = status_code
         self.detail = detail
@@ -39,7 +41,7 @@ class OrgSwitchIdpChallenge(Exception):
     caller on their previous org.
     """
 
-    def __init__(self, org_name: str, idp: dict):
+    def __init__(self, org_name: str, idp: Dict[str, Any]) -> None:  # hygiene-ok: explicit-any -- idp is arbitrary server-provided JSON
         self.org_name = org_name
         self.idp = idp
         super().__init__(

From 4a9d63e23cf9265fd40f88f1cec94ad1e809f131 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 28 Aug 2026 15:47:04 +0800
Subject: [PATCH 07/11] fix(mypy): silence Optional[str] arg-type in
 switch_org, not cast

Guard baseline already sat at max cast() count; cast() would've pushed it
over. Use a localized type: ignore instead, per DEVELOP.md guidance.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GgDZr3AgU9uKz5EWcVZZVY
---
 graphistry/pygraphistry.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 3609a53369..a8a7de893d 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2699,7 +2699,9 @@ def personal_key_secret(self, value: Optional[str] = None):
     def switch_org(self, value: str):
         org_name = value.strip()
         try:
-            switch_org_request(self._base_url(), org_name, self.api_token(), self.session.certificate_validation)
+            switch_org_request(
+                self._base_url(), org_name, self.api_token(), self.session.certificate_validation  # type: ignore[arg-type]
+            )
         except (OrgSwitchError, OrgSwitchIdpChallenge) as exc:
             raise Exception(str(exc))
 

From cb14becf08093ebca4e1fbac2f7fdcd315fd08ed Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 28 Aug 2026 17:55:09 +0800
Subject: [PATCH 08/11] test: cover switch_org_request and
 ArrowUploader._switch_org branches

Adds direct unit tests for graphistry/utils/requests.py's switch_org_request
(success, HTTP error, non-OK body w/ detail, unparseable JSON, SSO idp
challenge) and ArrowUploader._switch_org (no-op guards, dedup cache,
OrgSwitchError/OrgSwitchIdpChallenge/generic-exception swallow paths),
raising changed-line coverage on the switch-org refactor.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GgDZr3AgU9uKz5EWcVZZVY
---
 graphistry/tests/utils/test_requests.py | 140 ++++++++++++++++++++++++
 1 file changed, 140 insertions(+)
 create mode 100644 graphistry/tests/utils/test_requests.py

diff --git a/graphistry/tests/utils/test_requests.py b/graphistry/tests/utils/test_requests.py
new file mode 100644
index 0000000000..5c7e36aff5
--- /dev/null
+++ b/graphistry/tests/utils/test_requests.py
@@ -0,0 +1,140 @@
+# -*- coding: utf-8 -*-
+import json
+import unittest
+
+try:
+    import mock  # type: ignore
+except ImportError:  # pragma: no cover - fallback for stdlib-only envs
+    from unittest import mock
+
+from graphistry.arrow_uploader import ArrowUploader
+from graphistry.client_session import ClientSession
+from graphistry.utils.requests import (
+    OrgSwitchError,
+    OrgSwitchIdpChallenge,
+    switch_org_request,
+)
+
+
+class FakeResponse(object):
+    """Minimal requests.Response stand-in for switch_org_request tests."""
+
+    def __init__(self, status_code: int, body=None, raise_on_json: bool = False):
+        self.status_code = status_code
+        self._body = body
+        self._raise_on_json = raise_on_json
+        self.text = json.dumps(body) if body is not None else ""
+
+    def json(self):
+        if self._raise_on_json:
+            raise json.JSONDecodeError("no json", "", 0)
+        return self._body
+
+
+class TestSwitchOrgRequest(unittest.TestCase):
+    @mock.patch("graphistry.utils.requests.requests.post")
+    def test_success(self, mock_post):
+        mock_post.return_value = FakeResponse(200, {"status": "OK", "data": []})
+        # Should not raise
+        switch_org_request("https://hub.graphistry.com", "my-org", "tok", True)
+        args, kwargs = mock_post.call_args
+        assert args[0] == "https://hub.graphistry.com/api/v2/o/my-org/switch/"
+        assert kwargs["data"] == {"slug": "my-org"}
+        assert kwargs["headers"]["Authorization"] == "Bearer tok"
+        assert kwargs["verify"] is True
+
+    @mock.patch("graphistry.utils.requests.requests.post")
+    def test_http_error_status(self, mock_post):
+        mock_post.return_value = FakeResponse(404, {"status": "Failed", "message": "not found"})
+        with self.assertRaises(OrgSwitchError) as cm:
+            switch_org_request("https://hub.graphistry.com", "no-such-org", "tok", True)
+        assert cm.exception.status_code == 404
+        assert cm.exception.org_name == "no-such-org"
+        assert "404" in str(cm.exception)
+
+    @mock.patch("graphistry.utils.requests.requests.post")
+    def test_body_status_not_ok_with_detail(self, mock_post):
+        mock_post.return_value = FakeResponse(200, {"status": "Failed", "message": "not permitted"})
+        with self.assertRaises(OrgSwitchError) as cm:
+            switch_org_request("https://hub.graphistry.com", "not-permitted-org", "tok", True)
+        assert cm.exception.detail == "not permitted"
+        assert "not permitted" in str(cm.exception)
+
+    @mock.patch("graphistry.utils.requests.requests.post")
+    def test_unparseable_json_body_treated_as_failure(self, mock_post):
+        mock_post.return_value = FakeResponse(200, None, raise_on_json=True)
+        with self.assertRaises(OrgSwitchError) as cm:
+            switch_org_request("https://hub.graphistry.com", "my-org", "tok", True)
+        # falls back to body={} -> status missing -> not 'OK'
+        assert cm.exception.detail == ""
+
+    @mock.patch("graphistry.utils.requests.requests.post")
+    def test_idp_challenge(self, mock_post):
+        idp = {"auth_url": "https://sso-idp-host/authorize?state=xxx"}
+        mock_post.return_value = FakeResponse(
+            200, {"status": "OK", "data": {"idp": idp}}
+        )
+        with self.assertRaises(OrgSwitchIdpChallenge) as cm:
+            switch_org_request("https://hub.graphistry.com", "sso-org", "tok", True)
+        assert cm.exception.org_name == "sso-org"
+        assert cm.exception.idp == idp
+        assert "sso-org" in str(cm.exception)
+        assert "auth_url" in str(cm.exception)
+
+
+class TestArrowUploaderSwitchOrg(unittest.TestCase):
+    def _uploader(self):
+        return ArrowUploader(client_session=ClientSession())
+
+    def test_noop_without_org_or_token(self):
+        au = self._uploader()
+        with mock.patch("graphistry.arrow_uploader.switch_org_request") as mock_req:
+            au._switch_org(None, "tok")
+            au._switch_org("org", None)
+            au._switch_org(None, None)
+        mock_req.assert_not_called()
+
+    def test_skips_when_already_switched(self):
+        au = self._uploader()
+        au._client_session._last_switched_org_token = ("my-org", "tok")
+        with mock.patch("graphistry.arrow_uploader.switch_org_request") as mock_req:
+            au._switch_org("my-org", "tok")
+        mock_req.assert_not_called()
+
+    def test_success_records_last_switched(self):
+        au = self._uploader()
+        with mock.patch("graphistry.arrow_uploader.switch_org_request") as mock_req:
+            au._switch_org("my-org", "tok")
+        mock_req.assert_called_once_with(au.server_base_path, "my-org", "tok", au.certificate_validation)
+        assert au._client_session._last_switched_org_token == ("my-org", "tok")
+
+    def test_org_switch_error_does_not_record(self):
+        au = self._uploader()
+        with mock.patch(
+            "graphistry.arrow_uploader.switch_org_request",
+            side_effect=OrgSwitchError("my-org", 404),
+        ):
+            au._switch_org("my-org", "tok")
+        assert au._client_session._last_switched_org_token is None
+
+    def test_idp_challenge_does_not_record(self):
+        au = self._uploader()
+        with mock.patch(
+            "graphistry.arrow_uploader.switch_org_request",
+            side_effect=OrgSwitchIdpChallenge("my-org", {"auth_url": "https://x"}),
+        ):
+            au._switch_org("my-org", "tok")
+        assert au._client_session._last_switched_org_token is None
+
+    def test_generic_exception_does_not_record(self):
+        au = self._uploader()
+        with mock.patch(
+            "graphistry.arrow_uploader.switch_org_request",
+            side_effect=RuntimeError("boom"),
+        ):
+            au._switch_org("my-org", "tok")
+        assert au._client_session._last_switched_org_token is None
+
+
+if __name__ == "__main__":
+    unittest.main()

From 2761b1cf6e8a6e310ee7a7a8cddbc928bbabbade Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 28 Aug 2026 18:11:54 +0800
Subject: [PATCH 09/11] fix(test): patch requests.post via module object, not
 dotted string

mock.patch("graphistry.utils.requests.requests.post") failed under the
full suite (AttributeError: module 'graphistry' has no attribute 'utils')
because some other test's import juggling leaves the parent package's
'utils' attribute unset. mock.patch.object against the imported module
sidesteps the string-resolution path entirely.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GgDZr3AgU9uKz5EWcVZZVY
---
 graphistry/tests/utils/test_requests.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/graphistry/tests/utils/test_requests.py b/graphistry/tests/utils/test_requests.py
index 5c7e36aff5..2d9cd5e2b4 100644
--- a/graphistry/tests/utils/test_requests.py
+++ b/graphistry/tests/utils/test_requests.py
@@ -9,6 +9,7 @@
 
 from graphistry.arrow_uploader import ArrowUploader
 from graphistry.client_session import ClientSession
+from graphistry.utils import requests as switch_org_requests_module
 from graphistry.utils.requests import (
     OrgSwitchError,
     OrgSwitchIdpChallenge,
@@ -32,7 +33,7 @@ def json(self):
 
 
 class TestSwitchOrgRequest(unittest.TestCase):
-    @mock.patch("graphistry.utils.requests.requests.post")
+    @mock.patch.object(switch_org_requests_module.requests, "post")
     def test_success(self, mock_post):
         mock_post.return_value = FakeResponse(200, {"status": "OK", "data": []})
         # Should not raise
@@ -43,7 +44,7 @@ def test_success(self, mock_post):
         assert kwargs["headers"]["Authorization"] == "Bearer tok"
         assert kwargs["verify"] is True
 
-    @mock.patch("graphistry.utils.requests.requests.post")
+    @mock.patch.object(switch_org_requests_module.requests, "post")
     def test_http_error_status(self, mock_post):
         mock_post.return_value = FakeResponse(404, {"status": "Failed", "message": "not found"})
         with self.assertRaises(OrgSwitchError) as cm:
@@ -52,7 +53,7 @@ def test_http_error_status(self, mock_post):
         assert cm.exception.org_name == "no-such-org"
         assert "404" in str(cm.exception)
 
-    @mock.patch("graphistry.utils.requests.requests.post")
+    @mock.patch.object(switch_org_requests_module.requests, "post")
     def test_body_status_not_ok_with_detail(self, mock_post):
         mock_post.return_value = FakeResponse(200, {"status": "Failed", "message": "not permitted"})
         with self.assertRaises(OrgSwitchError) as cm:
@@ -60,7 +61,7 @@ def test_body_status_not_ok_with_detail(self, mock_post):
         assert cm.exception.detail == "not permitted"
         assert "not permitted" in str(cm.exception)
 
-    @mock.patch("graphistry.utils.requests.requests.post")
+    @mock.patch.object(switch_org_requests_module.requests, "post")
     def test_unparseable_json_body_treated_as_failure(self, mock_post):
         mock_post.return_value = FakeResponse(200, None, raise_on_json=True)
         with self.assertRaises(OrgSwitchError) as cm:
@@ -68,7 +69,7 @@ def test_unparseable_json_body_treated_as_failure(self, mock_post):
         # falls back to body={} -> status missing -> not 'OK'
         assert cm.exception.detail == ""
 
-    @mock.patch("graphistry.utils.requests.requests.post")
+    @mock.patch.object(switch_org_requests_module.requests, "post")
     def test_idp_challenge(self, mock_post):
         idp = {"auth_url": "https://sso-idp-host/authorize?state=xxx"}
         mock_post.return_value = FakeResponse(

From f24fb4531eec2a55c376c9928ea7482abaebf8a2 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 2 Sep 2026 16:50:25 +0800
Subject: [PATCH 10/11] fix(org-switch): skip SSO re-challenge for
 already-verified orgs (client-only)

Cache per-org verified JWTs (exp-aware) in ClientSession so switch_org()
can no-op or swap in a cached token instead of re-triggering SSO, without
any server-side changes. Multi-org: each org's token cached independently,
so switching back to an org verified under an earlier (now-inactive)
token still skips re-challenge as long as that token hasn't expired.

- client_session.py: _verified_org_tokens dict (org -> (token, exp)),
  is_org_verified / mark_org_verified / get_verified_token, JWT exp
  read (no signature verification, token already trusted).
- pygraphistry.py: switch_org() fast paths (same-token skip, cached-token
  swap-in); sso_login()/_handle_auth_url() record verification immediately
  on token issuance instead of waiting on a follow-up switch call.
- arrow_uploader.py: _switch_org() uses the same cache API.
- Tests added/migrated in test_pygraphistry.py, test_client_session.py,
  tests/utils/test_requests.py.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UnsBpuxHHTvC2bAUNaCfED
---
 graphistry/arrow_uploader.py            | 50 +++-----------
 graphistry/client_session.py            | 58 +++++++++++++++-
 graphistry/pygraphistry.py              | 89 ++++++++++++++-----------
 graphistry/tests/test_client_session.py | 56 +++++++++++++++-
 graphistry/tests/test_pygraphistry.py   | 61 +++++++++++++++--
 graphistry/tests/utils/test_requests.py | 10 +--
 6 files changed, 230 insertions(+), 94 deletions(-)

diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py
index 9d763d7cb1..d84aef0526 100644
--- a/graphistry/arrow_uploader.py
+++ b/graphistry/arrow_uploader.py
@@ -217,26 +217,10 @@ def certificate_validation(self, certificate_validation):
         self.__certificate_validation = certificate_validation
 
 
-    ########################################################################3
-
-    # @property
-    # def sso_state(self) -> str:
-    #     return getattr(self, '__sso_state', "")
-
-    ########################################################################3
-
-    # @property
-    # def sso_auth_url(self) -> str:
-    #     return getattr(self, '__sso_auth_url')
-
-    ########################################################################3
-
-
     def _switch_org(self, org_name: Optional[str], token: Optional[str]) -> None:
         if not org_name or not token:
             return
-        last = self._client_session._last_switched_org_token
-        if last == (org_name, token):
+        if self._client_session.is_org_verified(token, org_name):
             return
         try:
             switch_org_request(self.server_base_path, org_name, token, self.certificate_validation)
@@ -250,10 +234,10 @@ def _switch_org(self, org_name: Optional[str], token: Optional[str]) -> None:
             logger.warning("Failed to switch organization %s: %s", org_name, exc)
             return
 
-        self._client_session._last_switched_org_token = (org_name, token)
+        self._client_session.mark_org_verified(token, org_name)
         from .pygraphistry import PyGraphistry
         if PyGraphistry.session is self._client_session:
-            PyGraphistry.session._last_switched_org_token = (org_name, token)
+            PyGraphistry.session.mark_org_verified(token, org_name)
 
 
     def login(self, username, password, org_name=None):
@@ -311,8 +295,7 @@ def _finalize_login(self, out: requests.Response, org_name: Optional[str]) -> 'A
                 if not logged_in_org_name:  # no active_organization in JWT payload
                     raise Exception("You are not authorized to the organization '{}', or server does not support organization, please omit org_name parameter".format(org_name))
                 else:
-                    # if JWT response with org_name different than the pass in org_name
-                    # => org_name not found and return default organization (currently is personal org)
+                    # JWT's active org differs from the requested org_name -> treat as not found.
                     if logged_in_org_name != org_name:
                         raise Exception("Login Organization is not found in your organization")
 
@@ -465,12 +448,7 @@ def sso_get_token(self, state):
             active_org = data.get('active_organization')
             slug = active_org.get('slug') if isinstance(active_org, dict) else None
 
-            # Defense-in-depth: nexus's SSO claim-to-org resolver can (bug)
-            # bind a user to the reserved site-wide sentinel org (slug
-            # 'SITE') when an IdP org claim happens to case-insensitively
-            # match it. No dataset can ever be created under that org (server
-            # rejects with 403), so treat it the same as "no org bound" and
-            # fall through to the JWT-derived personal-org fallback below.
+            # Defense-in-depth: treat the reserved site-wide sentinel org ('SITE') as unbound -- no dataset can ever be created under it.
             if isinstance(slug, str) and slug.upper() == 'SITE':
                 logger.warning(
                     "SSO returned active_organization=%r, the reserved "
@@ -482,10 +460,7 @@ def sso_get_token(self, state):
                 slug = None
 
             if slug:
-                # Layer 1: server-bound active_organization. Caller's intent
-                # (self.org_name from register(org_name=...) or session) must
-                # MATCH or be ABSENT. Symmetric with _finalize_login's strict
-                # check for username/password (line ~309-316).
+                # Layer 1: server-bound active_organization -- caller's intent (self.org_name) must MATCH or be ABSENT.
                 if self.org_name and self.org_name != slug:
                     raise Exception(
                         f"SSO returned active_organization={slug!r}, but caller "
@@ -498,11 +473,7 @@ def sso_get_token(self, state):
                 self.org_name = slug
                 self._switch_org(slug, token_value)
             elif self.org_name:
-                # Layer 2: caller-supplied, server silent. Preserve caller
-                # intent — subsequent authenticated requests will validate org
-                # membership lazily. WARNING because the asymmetric outcome
-                # (caller asked, server didn't bind) is operationally
-                # interesting and worth investigating per-org SSO config.
+                # Layer 2: caller-supplied, server silent -- preserve caller intent; membership validates lazily on later requests.
                 logger.warning(
                     "SSO did not bind active_organization but caller requested "
                     "org_name=%s; preserving caller value. Subsequent requests "
@@ -511,9 +482,7 @@ def sso_get_token(self, state):
                     self.org_name
                 )
             else:
-                # Layer 3: caller didn't ask, server didn't bind. Try
-                # server-confirmed personal-org fallback for first-login UX.
-                # See _personal_org_from_api for the trust-chain rationale.
+                # Layer 3: caller didn't ask, server didn't bind -- try the server-confirmed personal-org fallback (see _personal_org_from_api).
                 fallback = self._personal_org_from_api(token_value)
                 if fallback:
                     logger.info(
@@ -531,7 +500,8 @@ def sso_get_token(self, state):
                     )
 
         except Exception as e:
-            logger.error('Unexpected SSO authentication error: %s', out, exc_info=True)
+            # debug, not error: fires on every routine "not ready yet" poll, not just genuine failures -- error level would spam the console.
+            logger.debug('SSO token poll did not complete: %s', out, exc_info=True)
             raise e
             
         return self
diff --git a/graphistry/client_session.py b/graphistry/client_session.py
index 5b7da8aeac..d44b5d7adb 100644
--- a/graphistry/client_session.py
+++ b/graphistry/client_session.py
@@ -1,4 +1,6 @@
+import base64
 import os
+import time
 from dataclasses import is_dataclass, replace
 from typing import Any, Optional, Literal, cast, Protocol, TypedDict, Dict, MutableMapping, Type, TypeVar, Union, overload, Iterator, Tuple
 from functools import lru_cache
@@ -33,6 +35,27 @@
 ]
 
 
+def _jwt_exp(token: str) -> Optional[float]:
+    """Best-effort read of a JWT's "exp" claim, without verifying its signature --
+    the token is already trusted (this client received it from the server over
+    the auth flow it just completed); this only reads a claim out of a token
+    it already holds, never validates one from an untrusted source."""
+    try:
+        payload_b64 = token.split(".")[1]
+        padded = payload_b64 + "=" * (-len(payload_b64) % 4)
+        payload = json.loads(base64.urlsafe_b64decode(padded))
+        exp = payload.get("exp")
+        return float(exp) if exp is not None else None
+    except Exception:
+        return None
+
+
+def _jwt_expired(token: str) -> bool:
+    exp = _jwt_exp(token)
+    # No readable exp claim -> not expired by this check; let the server reject it if it's actually stale.
+    return exp is not None and exp <= time.time()
+
+
 class ClientSession:
     """
     Holds all configuration and authentication state for a Graphistry client.
@@ -79,6 +102,8 @@ def __init__(self) -> None:
 
         self.idp_name: Optional[str] = None
         self.sso_state: Optional[str] = None
+        # org_name the in-flight SSO login was initiated for: self.org_name still holds the previous org until the switch is confirmed.
+        self.sso_requested_org_name: Optional[str] = None
 
         self.personal_key: Optional[str] = None
         self.personal_key_id: Optional[str] = None
@@ -96,8 +121,37 @@ def __init__(self) -> None:
 
         # TODO: Migrate to a pattern like Kusto or Spanner
         self._bolt_driver: Optional[Any] = None
-        # Tracks the last (org_slug, jwt_token) we successfully switched to on the Hub server
-        self._last_switched_org_token: Optional[Tuple[str, str]] = None
+        # Per-org token cache (org_name -> (token, exp)) so re-switching to any previously verified org skips the server round trip until that token's own exp passes.
+        self._verified_org_tokens: Dict[str, Tuple[str, Optional[float]]] = {}
+
+    def is_org_verified(self, token: Optional[str], org_name: str) -> bool:
+        """Has org_name already been switched-to under this exact, still-unexpired token?"""
+        if not token:
+            return False
+        entry = self._verified_org_tokens.get(org_name)
+        return entry is not None and entry[0] == token and not _jwt_expired(token)
+
+    def mark_org_verified(self, token: Optional[str], org_name: str) -> None:
+        """Record org_name as switched-to under this token."""
+        if not token:
+            return
+        self._verified_org_tokens[org_name] = (token, _jwt_exp(token))
+
+    def get_verified_token(self, org_name: str) -> Optional[str]:
+        """Return a still-unexpired token previously verified for org_name, if any.
+
+        May differ from the currently-active token (e.g. a later SSO login
+        for a different org minted a new one) -- callers that use this to
+        skip a switch should swap it in as the active token.
+        """
+        entry = self._verified_org_tokens.get(org_name)
+        if entry is None:
+            return None
+        token, _exp = entry
+        if _jwt_expired(token):
+            del self._verified_org_tokens[org_name]
+            return None
+        return token
 
     def copy(self) -> "ClientSession":
         """
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index a8a7de893d..2e46344a68 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -1,7 +1,7 @@
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union, cast, overload
 from typing_extensions import Literal
 from graphistry.privacy import Mode, ModeAction
-from graphistry.utils.requests import log_requests_error, switch_org_request, OrgSwitchError, OrgSwitchIdpChallenge
+from graphistry.utils.requests import log_requests_error, switch_org_request
 from graphistry.plugins_types.hypergraph import HypergraphResult
 from graphistry.plugins_types.gexf_types import GexfEdgeViz, GexfNodeViz, GexfParseEngine
 from graphistry.client_session import (
@@ -233,11 +233,15 @@ def sso_login(self, org_name: Optional[str] = None, idp_name: Optional[str] = No
             )
 
         self.session._is_authenticated = False
+        # Track which org this SSO flow was initiated for, distinct from
+        # self.session.org_name (the still-previous org until this switch
+        # is confirmed) -- see _sso_get_token().
+        self.session.sso_requested_org_name = org_name
         arrow_uploader = ArrowUploader(
             client_session=self.session,
             server_base_path=self.protocol()
-            + "://"                   
-            + self.server(),  
+            + "://"
+            + self.server(),
             certificate_validation=self.certificate_validation(),
         ).sso_login(org_name, idp_name)
         try:
@@ -247,7 +251,11 @@ def sso_login(self, org_name: Optional[str] = None, idp_name: Optional[str] = No
                 raise ValueError("ArrowUploader.sso_login returned no token")
             self.api_token(token)
             self.session._is_authenticated = True
-            self._maybe_switch_org(org_name or self.session.org_name)
+            resolved_org = org_name or self.session.org_name
+            if resolved_org:
+                # Record verification now: a token just SSO-minted for resolved_org must never re-challenge on the switch_org_request() below.
+                self.session.mark_org_verified(token, resolved_org)
+            self._maybe_switch_org(resolved_org)
             arrow_uploader.token = None  # type: ignore[assignment]
             return token
 
@@ -304,7 +312,14 @@ def _handle_auth_url(self, auth_url: str, sso_timeout: Optional[int], sso_opt_in
             token = None
 
             while True:
-                token, org_name = self._sso_get_token()
+                try:
+                    token, org_name = self._sso_get_token()
+                except SsoStateInvalidException:
+                    raise
+                except Exception as exc:
+                    logger.debug("Transient error polling SSO token, will retry: %s", exc)
+                    token = None
+                    org_name = None
                 try:
                     if not token:
                         if elapsed_time % 10 == 1:
@@ -329,25 +344,18 @@ def _handle_auth_url(self, auth_url: str, sso_timeout: Optional[int], sso_opt_in
                 self.session.org_name = org_name
                 # finish, set back to None
                 self.session.sso_state = None
+                self.session.sso_requested_org_name = None
                 print("Successfully logged in")
+                current_token = self.api_token()
+                if org_name and current_token:
+                    # Same reasoning as the direct-token branch above: don't re-challenge a token just minted for org_name via SSO.
+                    self.session.mark_org_verified(current_token, org_name)
                 self._maybe_switch_org(org_name)
                 return self.api_token()
             else:
                 print("Please run graphistry.sso_get_token() to complete the authentication after you have authenticated via SSO")
                 return None
         else:
-            # print("Start getting token ...")
-            # token = None
-            # for i in range(10):
-            #     token, org_name = self._sso_get_token()
-            #     if token:
-            #         # set org_name to sso org
-            #         self._config['org_name'] = org_name
-            #         print("Successfully logged in")
-            #         return self.api_token()
-            #     print("Keep trying to get token ...")
-            #     time.sleep(5)
-
             print("Please run graphistry.sso_get_token() to complete the authentication")
             return None
 
@@ -367,12 +375,14 @@ def _sso_get_token(self) -> Tuple[Optional[str], Optional[str]]:
             raise SsoStateInvalidException("[SSO] Invalid SSO state: NoneType encountered")
 
         # print("_sso_get_token : {}".format(state))
+        # Pass the SSO-requested org explicitly: self.session.org_name still holds the previous org mid-switch.
         arrow_uploader = ArrowUploader(
             client_session=self.session,
             server_base_path=self.protocol()
-            + "://"                   
-            + self.server(),  
+            + "://"
+            + self.server(),
             certificate_validation=self.certificate_validation(),
+            org_name=self.session.sso_requested_org_name,
         ).sso_get_token(state)
 
         try:
@@ -519,18 +529,6 @@ def api_token(self, value: Optional[str] = None) -> Optional[str]:
             self.session._is_authenticated = False
         return value
 
-    # @staticmethod
-    # def api_token_refresh_ms(value=None):
-    #     """Set or get the API token refresh interval in milliseconds.
-    #     None and 0 interpreted as no refreshing."""
-
-    #     if value is None:
-    #         return PyGraphistry._config["api_token_refresh_ms"]
-
-    #     # setter
-    #     if value is not PyGraphistry._config["api_token_refresh_ms"]:
-    #         PyGraphistry._config["api_token_refresh_ms"] = int(value)
-
     def protocol(self, value: Optional[str] = None) -> str:
         """Set or get the protocol ('http' or 'https').
         Set automatically when using a server alias.
@@ -790,8 +788,7 @@ def _maybe_switch_org(self, org_name: Optional[str]) -> None:
         if not org_name:
             return
         token = self.api_token()
-        last = self.session._last_switched_org_token
-        if token and last == (org_name, token):
+        if self.session.is_org_verified(token, org_name):
             return
         try:
             self.switch_org(org_name)
@@ -2698,14 +2695,28 @@ def personal_key_secret(self, value: Optional[str] = None):
 
     def switch_org(self, value: str):
         org_name = value.strip()
-        try:
-            switch_org_request(
-                self._base_url(), org_name, self.api_token(), self.session.certificate_validation  # type: ignore[arg-type]
-            )
-        except (OrgSwitchError, OrgSwitchIdpChallenge) as exc:
-            raise Exception(str(exc))
+        token = self.api_token()
+
+        # Fast path 1: org_name already verified under the current token (a plain switch doesn't reissue the token) -- skip the round trip.
+        if self.session.is_org_verified(token, org_name):
+            self.session.org_name = org_name
+            logger.debug("Already switched to organization %s with current token; skipping request.", org_name)
+            return
+
+        # Fast path 2: an earlier, still-unexpired token already verified org_name (now superseded as active by a later SSO login elsewhere) -- swap it back in.
+        cached_token = self.session.get_verified_token(org_name)
+        if cached_token:
+            self.api_token(cached_token)
+            self.session.org_name = org_name
+            logger.debug("Reusing previously-verified token for organization %s; skipping request.", org_name)
+            return
+
+        switch_org_request(
+            self._base_url(), org_name, token, self.session.certificate_validation  # type: ignore[arg-type]
+        )
 
         self.session.org_name = org_name
+        self.session.mark_org_verified(token, org_name)
         logger.info("Switched to organization: {}".format(org_name))
 
     def _handle_api_response(self, response):
diff --git a/graphistry/tests/test_client_session.py b/graphistry/tests/test_client_session.py
index 40e42c4327..28725d0664 100644
--- a/graphistry/tests/test_client_session.py
+++ b/graphistry/tests/test_client_session.py
@@ -1,3 +1,7 @@
+import base64
+import json
+import time
+
 import pytest
 import graphistry
 from graphistry.client_session import ClientSession
@@ -7,6 +11,52 @@
 from unittest import mock
 
 
+def _fake_jwt(exp=None) -> str:
+    header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode()
+    payload = {} if exp is None else {"exp": exp}
+    body = base64.urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=").decode()
+    return f"{header}.{body}.sig"
+
+
+class TestVerifiedOrgTokenCache:
+    """Per-org, exp-aware token cache backing switch_org()'s client-only fast path."""
+
+    def test_multiple_orgs_survive_under_distinct_tokens(self):
+        session = ClientSession()
+        token_a = _fake_jwt(exp=time.time() + 3600)
+        token_b = _fake_jwt(exp=time.time() + 3600)
+
+        session.mark_org_verified(token_a, "org-a")
+        session.mark_org_verified(token_b, "org-b")
+
+        # Both remain retrievable even though token_b is "more current" --
+        # a fresh SSO login for org-b doesn't evict org-a's own grant.
+        assert session.get_verified_token("org-a") == token_a
+        assert session.get_verified_token("org-b") == token_b
+
+    def test_expired_entry_is_not_reused_and_is_evicted(self):
+        session = ClientSession()
+        expired = _fake_jwt(exp=time.time() - 1)
+        session.mark_org_verified(expired, "org-a")
+
+        assert session.get_verified_token("org-a") is None
+        # Second call confirms it was actually evicted, not just filtered.
+        assert "org-a" not in session._verified_org_tokens
+
+    def test_is_org_verified_requires_exact_current_token(self):
+        session = ClientSession()
+        token = _fake_jwt(exp=time.time() + 3600)
+        session.mark_org_verified(token, "org-a")
+
+        assert session.is_org_verified(token, "org-a") is True
+        assert session.is_org_verified("some-other-token", "org-a") is False
+
+    def test_unknown_org_returns_none(self):
+        session = ClientSession()
+        assert session.get_verified_token("never-switched") is None
+        assert session.is_org_verified("tok", "never-switched") is False
+
+
 class TestClientSession:
     """
     Tests for Graphistry client()/register() interactions using the *current*
@@ -153,7 +203,7 @@ def test_client_register_updates_last_switched_cache(self, mock_post, mock_switc
         client = graphistry.client()
 
         def fake_switch(org_name, token):
-            client.session._last_switched_org_token = (org_name, token)
+            client.session.mark_org_verified(token, org_name)
 
         mock_switch_org.side_effect = fake_switch
         mock_resp = mock.Mock()
@@ -172,11 +222,11 @@ def fake_switch(org_name, token):
         mock_resp.raise_for_status = mock.Mock()
         mock_post.return_value = mock_resp
 
-        assert client.session._last_switched_org_token is None
+        assert client.session._verified_org_tokens == {}
 
         client.register(api=3, username='u', password='p', org_name='mock-org')
 
-        assert client.session._last_switched_org_token == ('mock-org', 'tok123')
+        assert client.session.is_org_verified('tok123', 'mock-org')
 
     # --------------------------------------------------------------------- #
     # Persistence of arbitrary config                                       #
diff --git a/graphistry/tests/test_pygraphistry.py b/graphistry/tests/test_pygraphistry.py
index c71f5b54aa..b32b0724d2 100644
--- a/graphistry/tests/test_pygraphistry.py
+++ b/graphistry/tests/test_pygraphistry.py
@@ -1,5 +1,9 @@
 # -*- coding: utf-8 -*-
 
+import base64
+import json
+import time
+from typing import Optional
 import unittest, pytest
 try:
     from mock import patch  # type: ignore
@@ -86,7 +90,7 @@ def test_refresh_skips_switch_when_cached(mock_refresh):
     client = graphistry.client()
     client.session.org_name = "mock-org"
     client.api_token("tok123")
-    client.session._last_switched_org_token = ("mock-org", "tok123")
+    client.session.mark_org_verified("tok123", "mock-org")
 
     with patch.object(client, "switch_org") as mock_switch:
         client.refresh()
@@ -103,7 +107,7 @@ def test_refresh_switches_when_org_changes(mock_refresh):
     client = graphistry.client()
     client.session.org_name = "new-org"
     client.api_token("tok123")
-    client.session._last_switched_org_token = ("old-org", "tok123")
+    client.session.mark_org_verified("tok123", "old-org")
 
     with patch.object(client, "switch_org") as mock_switch:
         client.refresh()
@@ -114,7 +118,7 @@ def test_refresh_switches_when_org_changes(mock_refresh):
 def test_maybe_switch_org_cached_pair_skips():
     client = graphistry.client()
     client.api_token("tok123")
-    client.session._last_switched_org_token = ("mock-org", "tok123")
+    client.session.mark_org_verified("tok123", "mock-org")
 
     with patch.object(client, "switch_org") as mock_switch:
         client._maybe_switch_org("mock-org")
@@ -125,7 +129,7 @@ def test_maybe_switch_org_cached_pair_skips():
 def test_maybe_switch_org_new_token_switches():
     client = graphistry.client()
     client.api_token("tok123")
-    client.session._last_switched_org_token = ("mock-org", "old-token")
+    client.session.mark_org_verified("old-token", "mock-org")
 
     with patch.object(client, "switch_org") as mock_switch:
         client._maybe_switch_org("mock-org")
@@ -136,7 +140,7 @@ def test_maybe_switch_org_new_token_switches():
 def test_maybe_switch_org_new_org_switches():
     client = graphistry.client()
     client.api_token("tok123")
-    client.session._last_switched_org_token = ("other-org", "tok123")
+    client.session.mark_org_verified("tok123", "other-org")
 
     with patch.object(client, "switch_org") as mock_switch:
         client._maybe_switch_org("mock-org")
@@ -144,6 +148,53 @@ def test_maybe_switch_org_new_org_switches():
     mock_switch.assert_called_once_with("mock-org")
 
 
+def _fake_jwt(exp: Optional[float] = None) -> str:
+    """Minimal unsigned-looking JWT with an optional exp claim, for exercising
+    ClientSession's exp-aware verified-org cache without a real server."""
+    header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode()
+    payload = {} if exp is None else {"exp": exp}
+    body = base64.urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=").decode()
+    return f"{header}.{body}.sig"
+
+
+def test_switch_org_reuses_earlier_verified_token_for_different_org():
+    """Org A was SSO-verified under token1; a later SSO login to org B mints
+    token2 (now active). Switching back to org A should swap token1 back in
+    locally instead of hitting the server, since token1 hasn't expired."""
+    client = graphistry.client()
+    token1 = _fake_jwt(exp=time.time() + 3600)
+    token2 = _fake_jwt(exp=time.time() + 3600)
+
+    client.api_token(token1)
+    client.session.mark_org_verified(token1, "org-a")
+
+    client.api_token(token2)
+    client.session.mark_org_verified(token2, "org-b")
+
+    with patch("graphistry.pygraphistry.switch_org_request") as mock_req:
+        client.switch_org("org-a")
+
+    mock_req.assert_not_called()
+    assert client.api_token() == token1
+    assert client.session.org_name == "org-a"
+
+
+def test_switch_org_does_not_reuse_expired_cached_token():
+    client = graphistry.client()
+    expired_token = _fake_jwt(exp=time.time() - 10)
+    active_token = _fake_jwt(exp=time.time() + 3600)
+
+    client.api_token(expired_token)
+    client.session.mark_org_verified(expired_token, "org-a")
+    client.api_token(active_token)
+
+    with patch("graphistry.pygraphistry.switch_org_request") as mock_req:
+        client.switch_org("org-a")
+
+    mock_req.assert_called_once()
+    assert client.api_token() == active_token
+
+
 class FakeRequestResponse(object):
     def __init__(self, response, status_code: int):
         self.response = response
diff --git a/graphistry/tests/utils/test_requests.py b/graphistry/tests/utils/test_requests.py
index 2d9cd5e2b4..a9f0fa2d00 100644
--- a/graphistry/tests/utils/test_requests.py
+++ b/graphistry/tests/utils/test_requests.py
@@ -97,7 +97,7 @@ def test_noop_without_org_or_token(self):
 
     def test_skips_when_already_switched(self):
         au = self._uploader()
-        au._client_session._last_switched_org_token = ("my-org", "tok")
+        au._client_session.mark_org_verified("tok", "my-org")
         with mock.patch("graphistry.arrow_uploader.switch_org_request") as mock_req:
             au._switch_org("my-org", "tok")
         mock_req.assert_not_called()
@@ -107,7 +107,7 @@ def test_success_records_last_switched(self):
         with mock.patch("graphistry.arrow_uploader.switch_org_request") as mock_req:
             au._switch_org("my-org", "tok")
         mock_req.assert_called_once_with(au.server_base_path, "my-org", "tok", au.certificate_validation)
-        assert au._client_session._last_switched_org_token == ("my-org", "tok")
+        assert au._client_session.is_org_verified("tok", "my-org")
 
     def test_org_switch_error_does_not_record(self):
         au = self._uploader()
@@ -116,7 +116,7 @@ def test_org_switch_error_does_not_record(self):
             side_effect=OrgSwitchError("my-org", 404),
         ):
             au._switch_org("my-org", "tok")
-        assert au._client_session._last_switched_org_token is None
+        assert au._client_session._verified_org_tokens == {}
 
     def test_idp_challenge_does_not_record(self):
         au = self._uploader()
@@ -125,7 +125,7 @@ def test_idp_challenge_does_not_record(self):
             side_effect=OrgSwitchIdpChallenge("my-org", {"auth_url": "https://x"}),
         ):
             au._switch_org("my-org", "tok")
-        assert au._client_session._last_switched_org_token is None
+        assert au._client_session._verified_org_tokens == {}
 
     def test_generic_exception_does_not_record(self):
         au = self._uploader()
@@ -134,7 +134,7 @@ def test_generic_exception_does_not_record(self):
             side_effect=RuntimeError("boom"),
         ):
             au._switch_org("my-org", "tok")
-        assert au._client_session._last_switched_org_token is None
+        assert au._client_session._verified_org_tokens == {}
 
 
 if __name__ == "__main__":

From 56a92976467ff4720d6ee15159210231d33383ae Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Thu, 3 Sep 2026 17:21:59 +0800
Subject: [PATCH 11/11] fix(sso): classify poll failures by type instead of
 swallowing all

The SSO token endpoint answers with a non-OK body status ("State is
invalid") on every poll until the browser login completes, so the poll
loop's blanket `except Exception` was load-bearing. But it also swallowed
permanent failures -- an org mismatch burned the full sso_timeout and then
reported a misleading "run sso_get_token()" message.

Type the retryable cases instead:

- New SsoPendingException, raised by ArrowUploader.sso_get_token for a
  non-OK body status and for an unrecognized body with no 'status' key
  (proxy error page, DRF detail). Nothing before the token is parsed can
  be classified as permanent.
- The poll loop waits on SsoPendingException uncapped until sso_timeout,
  retries RequestException/ValueError up to SSO_POLL_MAX_CONSECUTIVE_ERRORS
  (ValueError covers both json and simplejson decode errors), and lets
  everything else -- org mismatch, missing JWT in an OK response --
  propagate on the first poll.

Also fix switch_org's cached-token fast path clearing _is_authenticated:
it went through the api_token() setter, whose side effect forced a
needless refresh() on the next authenticate(); that refresh minted a token
not verified for the org, reinstating the round trip the fast path exists
to skip. Assign the session field directly, and carry the org grant onto
the refreshed token in refresh() so the fast path stays hit.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AvLMs6S5UosfKf2YB2sCRS
---
 graphistry/arrow_uploader.py            |   8 +-
 graphistry/exceptions.py                |   9 ++
 graphistry/pygraphistry.py              |  33 +++++-
 graphistry/tests/test_arrow_uploader.py |  26 +++++
 graphistry/tests/test_pygraphistry.py   | 128 +++++++++++++++++++++++-
 5 files changed, 196 insertions(+), 8 deletions(-)

diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py
index d84aef0526..0d75a9b839 100644
--- a/graphistry/arrow_uploader.py
+++ b/graphistry/arrow_uploader.py
@@ -14,7 +14,7 @@
     serialize_edge_encodings
 )
 
-from .exceptions import TokenExpireException
+from .exceptions import TokenExpireException, SsoPendingException
 from .validate.validate_encodings import validate_encodings
 from .utils.requests import log_requests_error, switch_org_request, OrgSwitchError, OrgSwitchIdpChallenge
 from .util import setup_logger
@@ -434,10 +434,12 @@ def sso_get_token(self, state):
             json_response = out.json()
             self.token = None
             if 'status' not in json_response:
-                raise Exception(out.text)
+                # Unrecognized body (proxy error page, DRF detail, ...): retryable, not fatal.
+                raise SsoPendingException(out.text)
 
             if json_response['status'] != 'OK':
-                raise Exception(json_response.get('message', out.text))
+                # Normal "not ready yet" answer while the browser login is in flight.
+                raise SsoPendingException(json_response.get('message', out.text))
 
             data = json_response.get('data', {})
             token_value = data.get('token')
diff --git a/graphistry/exceptions.py b/graphistry/exceptions.py
index 7e7287c7ca..9697600286 100644
--- a/graphistry/exceptions.py
+++ b/graphistry/exceptions.py
@@ -18,6 +18,15 @@ class SsoStateInvalidException(SsoException):
     pass
 
 
+class SsoPendingException(SsoException):
+    """
+    Raised while the SSO token is not ready yet: the browser login is still in
+    flight, so the token endpoint answers with a non-OK body status (typically
+    "State is invalid"). Retryable -- the poll loop keeps waiting on it.
+    """
+    pass
+
+
 
 class TokenExpireException(Exception):
     """
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 2e46344a68..f617874795 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -33,7 +33,7 @@
 from . import bolt_util
 from .plotter import Plotter
 from .util import in_databricks, setup_logger, in_ipython
-from .exceptions import SsoRetrieveTokenTimeoutException, TokenExpireException, SsoStateInvalidException
+from .exceptions import SsoRetrieveTokenTimeoutException, TokenExpireException, SsoStateInvalidException, SsoPendingException
 
 from .messages import (
     MSG_REGISTER_MISSING_PASSWORD,
@@ -51,6 +51,10 @@
 
 SSO_GET_TOKEN_ELAPSE_SECONDS = 50
 
+# Consecutive network-level failures tolerated while polling for the SSO token
+# before giving up; a sustained outage should not wait out the whole sso_timeout.
+SSO_POLL_MAX_CONSECUTIVE_ERRORS = 3
+
 _client_mode_enabled = False
 _is_client_mode_warned = False
 PyGraphistry: "GraphistryClient" = None  # type: ignore[assignment]
@@ -310,13 +314,29 @@ def _handle_auth_url(self, auth_url: str, sso_timeout: Optional[int], sso_opt_in
             time.sleep(1)
             elapsed_time = 1
             token = None
+            consecutive_poll_errors = 0
 
             while True:
                 try:
                     token, org_name = self._sso_get_token()
+                    consecutive_poll_errors = 0
                 except SsoStateInvalidException:
                     raise
-                except Exception as exc:
+                except SsoPendingException as exc:
+                    # Browser login still in flight; keep waiting until sso_timeout.
+                    logger.debug("SSO token not ready yet: %s", exc)
+                    token = None
+                    org_name = None
+                except (requests.exceptions.RequestException, ValueError) as exc:
+                    # ValueError covers both json and simplejson decode errors.
+                    # Retry network/parse blips; permanent errors (org mismatch) propagate.
+                    consecutive_poll_errors += 1
+                    if consecutive_poll_errors > SSO_POLL_MAX_CONSECUTIVE_ERRORS:
+                        logger.error(
+                            "Repeated network errors polling SSO token (%s consecutive): %s",
+                            consecutive_poll_errors, exc
+                        )
+                        raise
                     logger.debug("Transient error polling SSO token, will retry: %s", exc)
                     token = None
                     org_name = None
@@ -425,7 +445,11 @@ def refresh(self, token: Optional[str] = None, fail_silent: bool = False) -> Opt
             )
             self.api_token(token)
             self.session._is_authenticated = True
-            self._maybe_switch_org(self.session.org_name)
+            refreshed_org = self.session.org_name
+            if token and refreshed_org and self.session.get_verified_token(refreshed_org):
+                # Refresh reissues under the same org context; carry the grant to the new token.
+                self.session.mark_org_verified(token, refreshed_org)
+            self._maybe_switch_org(refreshed_org)
             return self.api_token()
         except Exception as e:
 
@@ -2706,7 +2730,8 @@ def switch_org(self, value: str):
         # Fast path 2: an earlier, still-unexpired token already verified org_name (now superseded as active by a later SSO login elsewhere) -- swap it back in.
         cached_token = self.session.get_verified_token(org_name)
         if cached_token:
-            self.api_token(cached_token)
+            # Direct assign: api_token()'s setter would clear _is_authenticated and force a needless refresh().
+            self.session.api_token = cached_token
             self.session.org_name = org_name
             logger.debug("Reusing previously-verified token for organization %s; skipping request.", org_name)
             return
diff --git a/graphistry/tests/test_arrow_uploader.py b/graphistry/tests/test_arrow_uploader.py
index b975e9b6d8..c0cd8d9cb9 100644
--- a/graphistry/tests/test_arrow_uploader.py
+++ b/graphistry/tests/test_arrow_uploader.py
@@ -8,6 +8,7 @@
 
 from graphistry import ArrowUploader
 from graphistry.pygraphistry import PyGraphistry
+from graphistry.exceptions import SsoPendingException
 
 # TODO mock requests for testing actual effectful code
 
@@ -472,6 +473,31 @@ def test_sso_login_get_sso_token_ok(self, mock_get):
         assert au.org_name == 'mock-org'
         mock_switch.assert_called_once_with('mock-org', '123')
 
+    @mock.patch('requests.get')
+    def test_sso_get_token_state_not_ready_raises_pending(self, mock_get):
+        # The endpoint answers status != OK on every poll until the browser
+        # login completes. Must be typed retryable, not a generic failure --
+        # the poll loop keys off SsoPendingException to keep waiting.
+        mock_get.return_value = self._mock_response(
+            json_data={'status': 'Failed', 'message': 'State is invalid'}
+        )
+
+        au = ArrowUploader()
+
+        with self.assertRaises(SsoPendingException) as cm:
+            au.sso_get_token(state='not-ready-yet')
+        assert 'State is invalid' in str(cm.exception)
+
+    @mock.patch('requests.get')
+    def test_sso_get_token_status_less_body_raises_pending(self, mock_get):
+        # Proxy/DRF error body with no 'status' key: also retryable.
+        mock_get.return_value = self._mock_response(json_data={'detail': 'Not found.'})
+
+        au = ArrowUploader()
+
+        with self.assertRaises(SsoPendingException):
+            au.sso_get_token(state='ignored')
+
     @mock.patch('requests.get')
     def test_sso_get_token_missing_active_organization_no_caller_org(self, mock_get):
         # Site-wide SSO: server omits active_organization, caller passed no
diff --git a/graphistry/tests/test_pygraphistry.py b/graphistry/tests/test_pygraphistry.py
index b32b0724d2..fcc8ca0e58 100644
--- a/graphistry/tests/test_pygraphistry.py
+++ b/graphistry/tests/test_pygraphistry.py
@@ -11,7 +11,8 @@
     from unittest.mock import patch
 import graphistry
 
-from graphistry.pygraphistry import PyGraphistry, GraphistryClient
+from graphistry.pygraphistry import PyGraphistry, GraphistryClient, SSO_POLL_MAX_CONSECUTIVE_ERRORS
+from graphistry.exceptions import SsoPendingException
 from graphistry.messages import (
     MSG_REGISTER_MISSING_PASSWORD,
     MSG_REGISTER_MISSING_USERNAME,
@@ -170,6 +171,7 @@ def test_switch_org_reuses_earlier_verified_token_for_different_org():
 
     client.api_token(token2)
     client.session.mark_org_verified(token2, "org-b")
+    client.session._is_authenticated = True
 
     with patch("graphistry.pygraphistry.switch_org_request") as mock_req:
         client.switch_org("org-a")
@@ -177,6 +179,9 @@ def test_switch_org_reuses_earlier_verified_token_for_different_org():
     mock_req.assert_not_called()
     assert client.api_token() == token1
     assert client.session.org_name == "org-a"
+    # Swapping in an already-verified token must not de-authenticate the session,
+    # else the next authenticate() refreshes and undoes the skipped round trip.
+    assert client.session._is_authenticated is True
 
 
 def test_switch_org_does_not_reuse_expired_cached_token():
@@ -195,6 +200,127 @@ def test_switch_org_does_not_reuse_expired_cached_token():
     assert client.api_token() == active_token
 
 
+class TestSsoPollErrorHandling:
+    """_handle_auth_url's polling loop: keep waiting on SsoPendingException and
+    network blips, propagate permanent errors."""
+
+    def _client(self):
+        client = graphistry.client()
+        client.session.sso_state = "state-123"
+        return client
+
+    def test_pending_state_keeps_polling_until_login_completes(self):
+        # The token endpoint answers "State is invalid" on every poll until the
+        # user finishes the browser login -- that is the normal wait, not a failure.
+        outcomes = [
+            SsoPendingException("State is invalid"),
+            SsoPendingException("State is invalid"),
+            SsoPendingException("State is invalid"),
+            SsoPendingException("State is invalid"),
+            ("tok-final", "org-a"),
+        ]
+
+        def _poll():
+            outcome = outcomes.pop(0)
+            if isinstance(outcome, Exception):
+                raise outcome
+            return outcome
+
+        client = self._client()
+        with patch.object(client, "_sso_get_token", side_effect=_poll) as mock_get:
+            with patch("graphistry.pygraphistry.time.sleep"):
+                with patch.object(client, "api_token", return_value="tok-final"):
+                    with patch.object(client, "_maybe_switch_org"):
+                        out = client._handle_auth_url("http://auth", 30, None)
+
+        # More pending polls than the network-error budget: pending must not be capped.
+        assert mock_get.call_count == 5
+        assert out == "tok-final"
+
+    def test_pending_state_stops_at_timeout(self):
+        client = self._client()
+        with patch.object(
+            client, "_sso_get_token",
+            side_effect=SsoPendingException("State is invalid")
+        ):
+            with patch("graphistry.pygraphistry.time.sleep"):
+                out = client._handle_auth_url("http://auth", 3, None)
+
+        assert out is None
+
+    def test_status_less_body_is_treated_as_pending(self):
+        # A proxy/DRF error body with no 'status' key must not abort the login.
+        outcomes = [
+            SsoPendingException('{"detail": "Not found."}'),
+            ("tok-final", "org-a"),
+        ]
+
+        def _poll():
+            outcome = outcomes.pop(0)
+            if isinstance(outcome, Exception):
+                raise outcome
+            return outcome
+
+        client = self._client()
+        with patch.object(client, "_sso_get_token", side_effect=_poll):
+            with patch("graphistry.pygraphistry.time.sleep"):
+                with patch.object(client, "api_token", return_value="tok-final"):
+                    with patch.object(client, "_maybe_switch_org"):
+                        out = client._handle_auth_url("http://auth", 30, None)
+
+        assert out == "tok-final"
+
+    def test_permanent_error_propagates_instead_of_timing_out(self):
+        client = self._client()
+        boom = Exception("SSO returned active_organization='org-b', but caller requested org_name='org-a'")
+
+        with patch.object(client, "_sso_get_token", side_effect=boom) as mock_get:
+            with patch("graphistry.pygraphistry.time.sleep"):
+                with pytest.raises(Exception) as excinfo:
+                    client._handle_auth_url("http://auth", 30, None)
+
+        assert "active_organization" in str(excinfo.value)
+        # Failed on the first poll rather than spinning out the full timeout.
+        assert mock_get.call_count == 1
+
+    def test_transient_network_error_is_retried_then_succeeds(self):
+        client = self._client()
+        import requests as _requests
+        outcomes = [
+            _requests.exceptions.ConnectionError("reset"),
+            ("tok-final", "org-a"),
+        ]
+
+        def _poll():
+            outcome = outcomes.pop(0)
+            if isinstance(outcome, Exception):
+                raise outcome
+            return outcome
+
+        with patch.object(client, "_sso_get_token", side_effect=_poll):
+            with patch("graphistry.pygraphistry.time.sleep"):
+                with patch.object(client, "api_token", return_value="tok-final"):
+                    with patch.object(client, "_maybe_switch_org"):
+                        out = client._handle_auth_url("http://auth", 30, None)
+
+        assert out == "tok-final"
+        assert client.session.org_name == "org-a"
+
+    def test_sustained_network_errors_give_up_before_timeout(self):
+        client = self._client()
+        import requests as _requests
+
+        with patch.object(
+            client, "_sso_get_token",
+            side_effect=_requests.exceptions.ConnectionError("down")
+        ) as mock_get:
+            with patch("graphistry.pygraphistry.time.sleep"):
+                with pytest.raises(_requests.exceptions.ConnectionError):
+                    client._handle_auth_url("http://auth", 300, None)
+
+        assert mock_get.call_count == SSO_POLL_MAX_CONSECUTIVE_ERRORS + 1
+
+
 class FakeRequestResponse(object):
     def __init__(self, response, status_code: int):
         self.response = response