feat(NEX-584): First iteration of changes to SDK handling new client secrets - #339
feat(NEX-584): First iteration of changes to SDK handling new client secrets#339mateuszkuprowski wants to merge 4 commits into
Conversation
Sorry, something went wrong.
imjoshholloway
left a comment
There was a problem hiding this comment.
Nice shape overall — opt-in design, clean exception hierarchy, hook only fires when it sees one of our callables, plain-string api_key_auth untouched. A few fixes needed before this goes to PyPI.
Critical
-
AsyncClientCredentials.__call__ will crash on the second token exchange. self._async_lock is created in __init__; once asyncio.run(self.acquire()) runs the lock binds to that loop. The next exchange (after the cached JWT lapses) creates a new loop via asyncio.run and async with self._async_lock raises RuntimeError: ... bound to a different event loop. Existing tests only force a single exchange so they miss this. Please add a regression test that exhausts the cache and re-acquires.
-
AsyncClientCredentials has no cross-thread protection. The inherited _lock (threading.Lock) isn't held in __call__, so concurrent threads each spin a fresh loop, each call asyncio.run, and the asyncio-only lock provides zero serialization. Wrap the cache check + asyncio.run in __call__ under self._lock.
Suggested shape for #1 + #2 — lazy per-loop async lock, threading lock around the sync entry point:
def __init__(self, ...):
...
self._async_lock: Optional[asyncio.Lock] = None
self._async_lock_loop: Optional[asyncio.AbstractEventLoop] = None
# self._lock (threading.Lock) is inherited from _ExchangeCallableBase
def _get_async_lock(self) -> asyncio.Lock:
"""Return an asyncio.Lock bound to the *currently running* loop."""
loop = asyncio.get_running_loop()
with self._lock: # threading lock guards lazy init
if self._async_lock is None or self._async_lock_loop is not loop:
self._async_lock = asyncio.Lock()
self._async_lock_loop = loop
return self._async_lock
async def acquire(self) -> str:
now = time.monotonic()
cached = self._cached_token_if_fresh(now)
if cached is not None:
return cached
async with self._get_async_lock(): # always bound to current loop
now = time.monotonic()
cached = self._cached_token_if_fresh(now)
if cached is not None:
return cached
return await self._exchange()
def __call__(self) -> str:
try:
asyncio.get_running_loop()
except RuntimeError:
# No loop running — coalesce threads via the threading lock so
# multiple sync callers don't each spin their own asyncio.run().
now = time.monotonic()
cached = self._cached_token_if_fresh(now)
if cached is not None:
return cached
with self._lock:
now = time.monotonic()
cached = self._cached_token_if_fresh(now)
if cached is not None:
return cached
return asyncio.run(self.acquire())
# Inside a running loop — see Critical #3 below.
...Why this works: the async lock is created lazily inside a coroutine, so get_running_loop() returns the loop the coroutine is on and asyncio.Lock() binds correctly. The is not loop check refreshes the lock when asyncio.run opens a new loop. The threading lock guards both the lazy field assignment and the sync-entry cache-then-exchange path so concurrent OS threads collapse onto a single exchange.
- future.result() in the running-loop branch blocks the caller's event loop while the worker thread runs the exchange. The "doesn't block the caller's loop" comment is misleading. Either let async callers await acquire() directly, route through asyncio.run_coroutine_threadsafe, or fix the comment.
Important
- README still says Available from SDK version **X.Y.Z** — fill in before merge.
- Drop _ExchangeCallableBase from unstructured_client.auth.__all__; it's an underscore-private detail used only by the hook.
- The hook depends on a hand-edited block in the Speakeasy-generated sdk.py (_security_factory + __wrapped_callable__). The next regen will wipe this and silently break auth. Add either a CI guard, a Speakeasy overlay, or move the wrapping out of the generated file. At minimum add an end-to-end test that asserts the hook actually rewrites the header through real UnstructuredClient init.
- ClientCredentials / AsyncClientCredentials leak their private httpx.Client when close() / aclose() isn't called — and most users won't call it. Add a weakref.finalize like UnstructuredClient does.
Suggestions
- _get_http_client lazy-init relies on the call-site lock being held. Document or atomicize.
- Hook docstring leaks internal codenames (core-product, platform-api/public_api/dependencies) — strip from a public SDK.
- LegacyKeyExchange stores the key on both _client_secret (via super) and _api_key — collapse to one. Add a small test asserting retry/caching applies to the legacy path too.
- Include server_url / exchange URL in the 400 error message — misconfigurations (pointing at platform-api instead of accounts-service) will be common.
- Replace "unsightly-koala" reference in the integration test with generic phrasing.
- Document the http_client= injection knob in the README — useful for proxies/mTLS.
Nice work
- Outage fallback (serve still-unexpired cache after retry exhaustion with a warning log) is exactly what we want for a public SDK.
- Exception classes give callers something to branch on.
- Sync double-checked locking with time.monotonic() is correct.
- Test coverage on the sync path is good.
Sorry, something went wrong.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 4fb09ac. Configure here.
Sorry, something went wrong.
| _close_httpx_client(owner._http_client) | ||
| owner._http_client = None | ||
|
|
||
| self._finalizer = weakref.finalize(self, _finalize) |
There was a problem hiding this comment.
Weak reference returns None inside finalizer callback
Medium Severity
The weakref.finalize callback uses owner_ref = weakref.ref(self) to access the HTTP client, but during garbage collection CPython clears all weak references to an object before invoking finalize callbacks. So owner_ref() always returns None when _finalize runs, causing the function to return immediately without closing the lazily-created httpx.Client / httpx.AsyncClient. This means the finalizer safety net never actually works, leaking HTTP connections when users don't call close() / aclose() explicitly. The same issue exists in both ClientCredentials and AsyncClientCredentials.
Additional Locations (1)Reviewed by Cursor Bugbot for commit 4fb09ac. Configure here.
Sorry, something went wrong.
|
@imjoshholloway thanks for the thorough pass. All fixes pushed in the latest commits. Walkthrough below. Critical1 + 2. AsyncClientCredentials event-loop binding + cross-thread protection Adopted the lazy-per-loop async lock pattern you suggested, with one tweak: the lazy field assignment is guarded by a dedicated threading.Lock (_async_lock_init_lock) rather than the inherited self._lock. Reason: self._lock is also held in __call__ around the cache-check + asyncio.run, and acquire() runs inside that critical section when called from the sync entry point. Sharing one non-reentrant lock for both jobs would deadlock on the first re-entry. With the dedicated init lock the pattern is otherwise identical to your sketch — _get_async_lock() reads the running loop, swaps the asyncio.Lock whenever it's missing or bound to a stale loop, and acquire() always uses self._get_async_lock() instead of a stored field. __call__ now wraps the cache-check + dispatch in self._lock, both for the no-loop and inside-running-loop branches, so concurrent OS threads coalesce onto a single exchange. Regression coverage:
Same _http_client_init_lock treatment was applied to _get_http_client for the same deadlock reason. 3. future.result() blocks the caller's loop You're right, the comment was misleading. I left the worker-thread offload in place (Speakeasy's security factory still calls us synchronously, so we have no way to suspend back into the caller's loop here) but rewrote the docstring + inline comment to say plainly that __call__ from inside a running loop does block the loop and that async-native code should await acquire() directly. Happy to escalate further (e.g. raise instead of block when called from an async context) if you'd rather force users onto acquire() from day one. Important
Suggestions
Lint is clean (10.00/10) and the unit suite is green (201 passed, 1 pre-existing xfail unrelated to this change). |
Sorry, something went wrong.
Reviewers
Requested changes must be addressed to merge this pull request.
This PR lets our SDK handle new type client secrets. Exchanges them for JWT and manages refresh cycle. This change is fully transparent for all the existing external apps using our SDK. Switching to clinet-secrets is an opt in.
Note
Medium Risk
Adds new authentication flow that exchanges client secrets/legacy keys for JWTs and mutates outgoing auth headers, which could affect request authentication if misconfigured. Includes new caching/retry/concurrency logic and hook ordering changes that need careful review across sync/async usage.
Overview
Introduces a new unstructured_client.auth module that can exchange client secrets (and legacy API keys) for short-lived JWTs, with in-memory caching, refresh-before-expiry, retry/backoff on transient failures, and a fallback to still-valid cached tokens during account-service outages.
Updates request handling to send exchanged JWTs as Authorization: Bearer via a new AuthHeaderBeforeRequestHook, and adjusts sdk.py to preserve the original auth callable (__wrapped_callable__) so the hook can detect exchange-based auth.
Bumps SDK version to 0.44.0, adds extensive unit coverage plus an opt-in E2E integration test, and documents the new client-secret auth flow and tuning knobs in the README.
Reviewed by Cursor Bugbot for commit 4fb09ac. Bugbot is set up for automated code reviews on this repo. Configure here.