fix(okta): retry transient errors once on the pooled client path

api/services/okta_service.py

@@ -32,6 +32,14 @@
 # Page size for cursor-paginated list endpoints. Okta caps most list endpoints
 # at 200 per page; the facade drives the ``after`` cursor to walk pages.
 LIST_PAGE_LIMIT = 200
+# Extra attempts for a call that raises ``OktaTransientError``. Transient Okta
+# failures (a 429/5xx or timeout, and — via the reclassification in ``_call`` — a
+# pooled keep-alive connection Okta dropped server-side, which the SDK surfaces as
+# a None-response error) are worth one immediate retry: it gets a fresh connection
+# and usually succeeds, recovering within the run instead of skipping the resource
+# until the next reconcile. Kept at 1 so a genuinely-degraded endpoint, or a
+# lingering 429, still gives up quickly rather than amplifying load.
+OKTA_TRANSIENT_RETRIES = 1
 
 
 logger = logging.getLogger(__name__)
@@ -113,7 +121,17 @@ def __getattr__(self, name: str) -> Any:
 
         @functools.wraps(attr)
         async def wrapper(*args: Any, **kwargs: Any) -> Any:
-            return await OktaService._call(attr(*args, **kwargs))
+            # Rebuild the coroutine each attempt — a spent coroutine can't be
+            # re-awaited — and retry a bounded number of times on a transient
+            # failure, so a stale pooled connection is recovered within the run
+            # rather than deferred to the next reconcile.
+            for attempt in range(OKTA_TRANSIENT_RETRIES + 1):
+                try:
+                    return await OktaService._call(attr(*args, **kwargs))
+                except OktaTransientError:
+                    if attempt == OKTA_TRANSIENT_RETRIES:
+                        raise
+            raise AssertionError("unreachable")  # pragma: no cover
 
         return wrapper
 

tests/test_okta_service.py

@@ -3,12 +3,20 @@
 from typing import Any, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import pytest
+
 from okta.models.add_group_request import AddGroupRequest
 from okta.models.group import Group as OktaGroupType
 from okta.models.group_rule import GroupRule as OktaGroupRuleType
 from okta.models.user_schema import UserSchema as OktaUserSchemaType
 
-from api.services.okta_service import OktaService, UserSchema, is_managed_group
+from api.services.okta_service import (
+    OKTA_TRANSIENT_RETRIES,
+    OktaService,
+    OktaTransientError,
+    UserSchema,
+    is_managed_group,
+)
 from tests.factories import UserFactory
 
 
@@ -215,3 +223,40 @@ def tracking_build_client() -> Any:
     # Each concurrent call built its own client; nothing shared.
     assert len(clients) == call_count
     assert len({id(client) for client in clients}) == call_count
+
+
+async def test_transient_error_is_retried_then_succeeds() -> None:
+    """A transient Okta failure is retried within the same call and can recover.
+
+    The SDK surfaces a transient outcome (here a 503) that ``_call`` turns into an
+    ``OktaTransientError``; ``_WrapperClient`` retries it, and the next attempt —
+    on a fresh connection in production — succeeds.
+    """
+    service = OktaService()
+    service.initialize("fake.domain", "fake.token")
+
+    transient = (None, MagicMock(status_code=503, headers={}), MagicMock(status=503))
+    success = (UserFactory(), MagicMock(status_code=200, headers={}), None)
+    get_user = AsyncMock(side_effect=[transient, success])
+
+    with patch("okta.client.Client.get_user", get_user):
+        user = await service.get_user("okta_id")
+
+    assert user is not None
+    # One transient failure, retried once, then success.
+    assert get_user.call_count == 2
+
+
+async def test_transient_error_gives_up_after_bounded_retries() -> None:
+    """A persistently transient endpoint stops after the bounded retries."""
+    service = OktaService()
+    service.initialize("fake.domain", "fake.token")
+
+    transient = (None, MagicMock(status_code=503, headers={}), MagicMock(status=503))
+    get_user = AsyncMock(return_value=transient)
+
+    with patch("okta.client.Client.get_user", get_user):
+        with pytest.raises(OktaTransientError):
+            await service.get_user("okta_id")
+
+    assert get_user.call_count == OKTA_TRANSIENT_RETRIES + 1