Skip to content

feat(NEX-584): First iteration of changes to SDK handling new client secrets#339

Open
mateuszkuprowski wants to merge 4 commits intomainfrom
nex-584-ad-client-credentials-support-and-exchange
Open

feat(NEX-584): First iteration of changes to SDK handling new client secrets#339
mateuszkuprowski wants to merge 4 commits intomainfrom
nex-584-ad-client-credentials-support-and-exchange

Conversation

@mateuszkuprowski
Copy link
Copy Markdown

@mateuszkuprowski mateuszkuprowski commented Apr 24, 2026

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.

Comment thread src/unstructured_client/auth/client_credentials.py Outdated
Comment thread src/unstructured_client/auth/client_credentials.py Outdated
Comment thread src/unstructured_client/auth/client_credentials.py Outdated
Copy link
Copy Markdown

@imjoshholloway imjoshholloway left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

  1. 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.

  2. 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.

  1. 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.

Copy link
Copy Markdown

@cursor cursor Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Fix All in Cursor

❌ 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.

_close_httpx_client(owner._http_client)
owner._http_client = None

self._finalizer = weakref.finalize(self, _finalize)
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 4fb09ac. Configure here.

@mateuszkuprowski
Copy link
Copy Markdown
Author

@imjoshholloway thanks for the thorough pass. All fixes pushed in the latest commits. Walkthrough below.

Critical

1 + 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:

  • test_sync_call_re_exchanges_after_cache_lapse_across_loops — calls acc() twice across two asyncio.run invocations; would have raised RuntimeError: ... bound to a different event loop before the fix.
  • test_acquire_re_uses_correct_async_lock_after_loop_change — asserts _async_lock_loop actually re-binds to a new loop after the cache lapses.
  • test_sync_call_coalesces_concurrent_threads_into_one_exchange — two threads racing into __call__, only one HTTP exchange.

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

  • README X.Y.Z → bumped _version.py to 0.44.0 and replaced the placeholder.
  • _ExchangeCallableBase removed from auth.__all__ — only the public callables and exception classes are exported now. The hook still imports it directly from auth._base.
  • Speakeasy regen guard — added test_security_factory_exposes_wrapped_callable_for_hook_detection. It boots a real UnstructuredClient with a ClientCredentials instance, asserts sdk_configuration.security is callable and that __wrapped_callable__ points back at the original ClientCredentials. If a future Speakeasy regen wipes the hand-edited block in sdk.py this test fails loudly. (Long-term I think the better fix is a Speakeasy overlay; happy to follow up in a separate ticket.)
  • weakref.finalize for the private httpx.Client — both ClientCredentials and AsyncClientCredentials now register a finalizer that closes the lazily-created client. The finalizer is detached on explicit close() / aclose(). The async finalizer runs aclose() on a fresh loop in a worker thread so it doesn't touch whatever loop (if any) the user is on at GC time. README documents that explicit close is optional.

Suggestions

  • _get_http_client atomicity — moved off self._lock onto a dedicated _http_client_init_lock (sync + async) and added a docstring explaining the invariant.
  • Hook docstring codenames — stripped core-product / platform-api/public_api/dependencies from auth_header_hook.py. The docstring now describes only what's true at the SDK boundary (rewriting unstructured-api-keyAuthorization: Bearer for token-exchange callables).
  • LegacyKeyExchange field collapse + tests — removed the duplicate _api_key; the legacy key now lives only on the inherited _client_secret slot, and _build_request_body reads it from there. Added test_caches_jwt_within_ttl_for_legacy_path and test_retries_5xx_on_legacy_path_then_succeeds so caching + 5xx backoff are explicitly covered on the legacy grant.
  • 400 error includes the exchange URL_raise_for_status now embeds self._exchange_url and a hint that pointing server_url at platform-api instead of account-service is the most likely cause. Same treatment for 401 (URL only, no platform-api hint).
  • unsightly-koala — replaced with generic phrasing in both the module docstring and the skip reason of test_client_credentials_e2e.py.
  • http_client= in README — added a 'Custom HTTP client' bullet under 'Behavior and tuning' covering the use cases (proxy/mTLS/shared pool), the lazy ownership semantics, and the close() / aclose() interaction.

Lint is clean (10.00/10) and the unit suite is green (201 passed, 1 pre-existing xfail unrelated to this change).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants