Skip to content

Commit ccf15f2

Browse files
Ilanlidoclaude
andauthored
CM-72432: Retry token mint when a concurrent login is refused (#546)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent e0db88f commit ccf15f2

2 files changed

Lines changed: 70 additions & 1 deletion

File tree

cycode/cyclient/base_token_auth_client.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
1+
import secrets
2+
import time
13
from abc import ABC, abstractmethod
24
from threading import Lock
35
from typing import Any, Optional
46

57
import arrow
68
from requests import Response
79

10+
from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError
811
from cycode.cli.user_settings.credentials_manager import CredentialsManager
912
from cycode.cli.user_settings.jwt_creator import JwtCreator
1013
from cycode.cyclient.cycode_client import CycodeClient
@@ -15,6 +18,11 @@
1518
b'JWT Token validation failed',
1619
]
1720

21+
# Identity provider brute-force protection rejects logins for the same user that land within
22+
# milliseconds of each other, so when several processes mint at once all but one are refused.
23+
_MINT_CONFLICT_RETRY_MIN_MS = 50
24+
_MINT_CONFLICT_RETRY_SPREAD_MS = 100
25+
1826

1927
class BaseTokenAuthClient(CycodeClient, ABC):
2028
"""Base client for token-based authentication flows with cached JWTs."""
@@ -49,7 +57,20 @@ def refresh_access_token_if_needed(self) -> None:
4957
self._load_token_from_disk()
5058
if self._has_valid_token():
5159
return
52-
self.refresh_access_token()
60+
61+
try:
62+
self.refresh_access_token()
63+
except HttpUnauthorizedError:
64+
# Processes sharing one cached token all expire at the same instant, so a burst of
65+
# them mints together and the identity provider refuses all but the first as a
66+
# too-fast login. The winner persists a usable token, so prefer re-reading it over
67+
# minting again. The wait is randomized to keep the losers from colliding a second
68+
# time. A genuinely invalid token still raises on the retry.
69+
time.sleep((_MINT_CONFLICT_RETRY_MIN_MS + secrets.randbelow(_MINT_CONFLICT_RETRY_SPREAD_MS)) / 1000)
70+
self._load_token_from_disk()
71+
if self._has_valid_token():
72+
return
73+
self.refresh_access_token()
5374

5475
def _has_valid_token(self) -> bool:
5576
return self._access_token is not None and self._expires_in is not None and arrow.utcnow() < self._expires_in

tests/cyclient/test_token_based_client.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
import arrow
2+
import pytest
23
import responses
4+
from pyfakefs.fake_filesystem import FakeFilesystem
35

6+
from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError
47
from cycode.cyclient.cycode_token_based_client import CycodeTokenBasedClient
58
from tests.conftest import _EXPECTED_API_TOKEN, create_token_based_client
69

@@ -66,6 +69,51 @@ def test_access_token_cached_creator_changed(
6669
assert client2._expires_in is None
6770

6871

72+
@responses.activate
73+
def test_access_token_mint_conflict_prefers_token_persisted_by_another_process(
74+
api_token_url: str, fs: FakeFilesystem
75+
) -> None:
76+
client = create_token_based_client()
77+
78+
def _refuse_while_another_process_wins(_request: object) -> tuple:
79+
# the process that won the race persists its token while this one is being refused
80+
client._credentials_manager.update_access_token(
81+
_EXPECTED_API_TOKEN, arrow.utcnow().shift(hours=1).timestamp(), client._create_jwt_creator()
82+
)
83+
return 401, {}, ''
84+
85+
responses.add_callback(responses.POST, api_token_url, callback=_refuse_while_another_process_wins)
86+
87+
assert client.get_access_token() == _EXPECTED_API_TOKEN
88+
assert len(responses.calls) == 1
89+
90+
91+
@responses.activate
92+
def test_access_token_mint_conflict_retries_when_no_other_process_won(
93+
api_token_url: str, api_token_response: responses.Response, fs: FakeFilesystem
94+
) -> None:
95+
client = create_token_based_client()
96+
97+
responses.add(responses.Response(method=responses.POST, url=api_token_url, status=401))
98+
responses.add(api_token_response)
99+
100+
assert client.get_access_token() == _EXPECTED_API_TOKEN
101+
assert len(responses.calls) == 2
102+
103+
104+
@responses.activate
105+
def test_access_token_mint_conflict_raises_when_retry_is_refused_too(api_token_url: str, fs: FakeFilesystem) -> None:
106+
client = create_token_based_client()
107+
108+
responses.add(responses.Response(method=responses.POST, url=api_token_url, status=401))
109+
responses.add(responses.Response(method=responses.POST, url=api_token_url, status=401))
110+
111+
with pytest.raises(HttpUnauthorizedError):
112+
client.get_access_token()
113+
114+
assert len(responses.calls) == 2
115+
116+
69117
@responses.activate
70118
def test_access_token_invalidation(
71119
token_based_client: CycodeTokenBasedClient, api_token_response: responses.Response

0 commit comments

Comments
 (0)