Code review fixes

Author: mateuszkuprowski

README.md

@@ -504,7 +504,7 @@ s = UnstructuredClient(debug_logger=logging.getLogger("unstructured_client"))
 
 ## Authentication with Client Secrets
 
-> Available from SDK version **X.Y.Z** (first release carrying the
+> Available from SDK version **0.44.0** (first release carrying the
 > `unstructured_client.auth` module).
 
 If you are running against an Unstructured deployment that issues **client
@@ -586,6 +586,14 @@ client = UnstructuredClient(
   `token_exchange_enabled=False`, the call raises
   `TokenExchangeDisabledError` — that deployment expects the plain
   `api_key_auth="..."` string form instead.
+- **Custom HTTP client:** pass `http_client=httpx.Client(...)` (or
+  `httpx.AsyncClient(...)` for the async variant) to share a connection
+  pool, route through a corporate proxy, pin a custom CA bundle for mTLS,
+  or otherwise control how the SDK reaches account-service. When the
+  argument is omitted, the auth callable lazily creates and owns a private
+  `httpx.Client`; that client is closed automatically when the auth
+  instance is garbage-collected, or you can call `close()` /
+  `aclose()` explicitly.
 
 ### Backward compatibility
 

_test_unstructured_client/integration/test_client_credentials_e2e.py

@@ -2,15 +2,15 @@
 :class:`LegacyKeyExchange`.
 
 This test is **opt-in**: it only runs when every required env var is set.
-Point it at a deployment (e.g. the ``unsightly-koala`` dedicated-instance
-test cluster) that has ``DEPLOYMENT_MODE=dedicated`` and a valid client
-secret provisioned via account-service.
+Point it at any deployment that has ``DEPLOYMENT_MODE=dedicated`` (or any
+other configuration that accepts ``/auth/token-exchange``) and a valid
+client secret provisioned via account-service.
 
 Required env vars
 -----------------
 
 - ``UNS_ACCOUNTS_URL``   base URL of account-service (e.g.
-  ``https://accounts.unsightly-koala.example``)
+  ``https://accounts.<your-deployment>.example``)
 - ``UNS_CLIENT_SECRET``  ``uns_sk_...`` client secret
 - ``UNS_PLATFORM_API_URL`` platform-api base URL to hit after exchange
 
@@ -48,8 +48,8 @@
 
 _REASON = (
     "Opt-in E2E: set UNS_ACCOUNTS_URL, UNS_CLIENT_SECRET, and "
-    "UNS_PLATFORM_API_URL to run against a real dedicated-instance "
-    "deployment (e.g. unsightly-koala)."
+    "UNS_PLATFORM_API_URL to run against a real deployment that supports "
+    "/auth/token-exchange."
 )
 
 

_test_unstructured_client/unit/auth/test_async_client_credentials.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import threading
 from typing import List
 
 import httpx
@@ -168,3 +169,94 @@ async def test_sync_call_works_inside_running_loop(fake_clock):
 
     token = await asyncio.to_thread(acc)
     assert token == "jwt-1"
+
+
+def test_sync_call_re_exchanges_after_cache_lapse_across_loops(fake_clock):
+    """Regression: ``__call__`` must not crash on the second exchange.
+
+    Each invocation outside a running loop spins a fresh event loop via
+    ``asyncio.run``. The internal ``asyncio.Lock`` is created lazily inside
+    ``acquire`` so it binds to whatever loop is current; if it were created
+    in ``__init__`` it would stay bound to the first loop and ``async with``
+    would raise ``RuntimeError`` on the second exchange.
+    """
+    transport = AsyncScriptedTransport(
+        [
+            exchange_response(access_token="jwt-1", expires_in=900),
+            exchange_response(access_token="jwt-2", expires_in=900),
+        ]
+    )
+    http_client = httpx.AsyncClient(transport=transport)
+    acc = AsyncClientCredentials(
+        client_secret=SECRET,
+        server_url=SERVER_URL,
+        http_client=http_client,
+        refresh_buffer_seconds=60,
+    )
+
+    assert acc() == "jwt-1"
+    fake_clock["now"] += 900 - 30  # past refresh buffer
+    assert acc() == "jwt-2"
+    assert len(transport.requests) == 2
+
+
+@pytest.mark.asyncio
+async def test_acquire_re_uses_correct_async_lock_after_loop_change(fake_clock):
+    """The lazy per-loop lock must refresh when the running loop changes."""
+    transport = AsyncScriptedTransport(
+        [
+            exchange_response(access_token="jwt-1", expires_in=900),
+            exchange_response(access_token="jwt-2", expires_in=900),
+        ]
+    )
+    http_client = httpx.AsyncClient(transport=transport)
+    acc = AsyncClientCredentials(
+        client_secret=SECRET,
+        server_url=SERVER_URL,
+        http_client=http_client,
+        refresh_buffer_seconds=60,
+    )
+
+    first = await acc.acquire()
+    assert first == "jwt-1"
+    first_loop = acc._async_lock_loop  # type: ignore[attr-defined]
+
+    fake_clock["now"] += 900 - 30  # past refresh buffer
+
+    # Run the next exchange on a *different* event loop driven from a worker
+    # thread; the lazy-init code path must re-bind the lock to that loop.
+    second = await asyncio.to_thread(lambda: asyncio.run(acc.acquire()))
+    assert second == "jwt-2"
+
+    second_loop = acc._async_lock_loop  # type: ignore[attr-defined]
+    assert second_loop is not None
+    assert second_loop is not first_loop
+
+
+def test_sync_call_coalesces_concurrent_threads_into_one_exchange(fake_clock):
+    """Two OS threads racing into ``__call__`` must share one exchange."""
+    barrier = threading.Barrier(2)
+    transport = AsyncScriptedTransport(
+        [exchange_response(access_token="jwt-1", expires_in=900)]
+    )
+    http_client = httpx.AsyncClient(transport=transport)
+    acc = AsyncClientCredentials(
+        client_secret=SECRET,
+        server_url=SERVER_URL,
+        http_client=http_client,
+    )
+
+    results = []
+
+    def _worker():
+        barrier.wait()
+        results.append(acc())
+
+    threads = [threading.Thread(target=_worker) for _ in range(2)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert results == ["jwt-1", "jwt-1"]
+    assert len(transport.requests) == 1

_test_unstructured_client/unit/auth/test_auth_header_hook.py

@@ -175,6 +175,40 @@ def _mock(request: httpx.Request) -> httpx.Response:
     assert "unstructured-api-key" not in {k.lower() for k in headers}
 
 
+def test_security_factory_exposes_wrapped_callable_for_hook_detection():
+    """Speakeasy-regen guard: the security factory built by ``UnstructuredClient``
+    must expose ``__wrapped_callable__`` so :class:`AuthHeaderBeforeRequestHook`
+    can recognize our token-exchange callables.
+
+    If a future Speakeasy regeneration strips the hand-edited block in
+    ``sdk.py``, this test will fail loudly instead of letting the auth-header
+    rewrite silently break for every user.
+    """
+    transport = ScriptedTransport([exchange_response()])
+    cc = ClientCredentials(
+        client_secret=SECRET,
+        server_url=ACCOUNTS_URL,
+        http_client=httpx.Client(transport=transport),
+    )
+
+    session = UnstructuredClient(
+        api_key_auth=cc,
+        client=httpx.Client(transport=httpx.MockTransport(lambda r: httpx.Response(200))),
+        server_url=SERVER_URL,
+    )
+
+    security_factory = session.sdk_configuration.security
+    assert callable(security_factory), (
+        "When api_key_auth is callable, sdk.py must wrap it in a factory."
+    )
+    wrapped = getattr(security_factory, "__wrapped_callable__", None)
+    assert wrapped is cc, (
+        "sdk.py must attach __wrapped_callable__ pointing at the original "
+        "user callable; without it AuthHeaderBeforeRequestHook cannot detect "
+        "ClientCredentials / LegacyKeyExchange instances."
+    )
+
+
 def test_integration_leaves_legacy_path_unchanged_for_plain_string():
     captured: dict = {}
 

_test_unstructured_client/unit/auth/test_legacy_key_exchange.py

@@ -95,3 +95,53 @@ async def test_async_sends_grant_type_api_key_and_api_key_field():
 def test_rejects_empty_api_key(bad):
     with pytest.raises(ValueError):
         LegacyKeyExchange(api_key=bad, server_url=SERVER_URL)  # type: ignore[arg-type]
+
+
+def test_caches_jwt_within_ttl_for_legacy_path(monkeypatch):
+    """Caching applies to the legacy api-key path just like client-secrets."""
+    state = {"now": 1_000_000.0}
+    monkeypatch.setattr("unstructured_client.auth._base.time.monotonic", lambda: state["now"])
+    monkeypatch.setattr(
+        "unstructured_client.auth.client_credentials.time.monotonic", lambda: state["now"]
+    )
+
+    transport = ScriptedTransport([exchange_response(access_token="jwt-1", expires_in=900)])
+    http_client = httpx.Client(transport=transport)
+    lke = LegacyKeyExchange(
+        api_key=LEGACY_KEY,
+        server_url=SERVER_URL,
+        http_client=http_client,
+        refresh_buffer_seconds=60,
+    )
+
+    assert lke() == "jwt-1"
+    assert lke() == "jwt-1"
+    assert lke() == "jwt-1"
+    assert len(transport.requests) == 1
+
+
+def test_retries_5xx_on_legacy_path_then_succeeds(monkeypatch):
+    """5xx exponential-backoff retries also apply to the legacy api-key path."""
+    state = {"now": 1_000_000.0}
+    monkeypatch.setattr("unstructured_client.auth._base.time.monotonic", lambda: state["now"])
+    monkeypatch.setattr(
+        "unstructured_client.auth.client_credentials.time.monotonic", lambda: state["now"]
+    )
+
+    transport = ScriptedTransport(
+        [
+            httpx.Response(500),
+            httpx.Response(503),
+            exchange_response(access_token="jwt-1", expires_in=900),
+        ]
+    )
+    http_client = httpx.Client(transport=transport)
+    lke = LegacyKeyExchange(
+        api_key=LEGACY_KEY,
+        server_url=SERVER_URL,
+        http_client=http_client,
+        max_retries=3,
+    )
+
+    assert lke() == "jwt-1"
+    assert len(transport.requests) == 3

src/unstructured_client/_hooks/custom/auth_header_hook.py

@@ -1,11 +1,12 @@
 """Before-request hook that promotes exchanged JWTs to ``Authorization: Bearer``.
 
-Speakeasy's generated ``Security`` model places ``api_key_auth`` in the
+The generated ``Security`` model places ``api_key_auth`` in the
 ``unstructured-api-key`` header. When the user supplies a token-exchange
-callable (``ClientCredentials`` or ``LegacyKeyExchange``) the value is a JWT
-and must be sent as ``Authorization: Bearer <jwt>`` so the service-side
-``utic-jwt-auth`` validator picks it up (see ``core-product`` auth_context
-and ``platform-api`` public_api/dependencies).
+callable from :mod:`unstructured_client.auth` (such as
+:class:`~unstructured_client.auth.ClientCredentials` or
+:class:`~unstructured_client.auth.LegacyKeyExchange`) the value is a JWT
+and must be sent as ``Authorization: Bearer <jwt>`` so the server-side
+JWT validator accepts it.
 
 Plain-string ``api_key_auth`` is untouched.
 """
@@ -43,8 +44,10 @@ def _is_exchange_callable(security_source: object) -> bool:
         """Return True when ``security_source`` was built from one of our
         token-exchange callables.
 
-        The SDK wraps a user-supplied callable into an internal factory and
-        attaches ``__wrapped_callable__`` to it (see ``sdk.py``).
+        ``UnstructuredClient.__init__`` wraps a user-supplied callable into an
+        internal security factory and attaches ``__wrapped_callable__`` to it
+        so this hook can detect token-exchange instances without reaching
+        into the lambda's closure.
         """
         if security_source is None:
             return False

src/unstructured_client/_version.py

@@ -3,10 +3,10 @@
 import importlib.metadata
 
 __title__: str = "unstructured-client"
-__version__: str = "0.43.3"
+__version__: str = "0.44.0"
 __openapi_doc_version__: str = "1.2.31"
 __gen_version__: str = "2.680.0"
-__user_agent__: str = "speakeasy-sdk/python 0.43.3 2.680.0 1.2.31 unstructured-client"
+__user_agent__: str = "speakeasy-sdk/python 0.44.0 2.680.0 1.2.31 unstructured-client"
 
 try:
     if __package__ is not None:

src/unstructured_client/auth/__init__.py

@@ -19,7 +19,6 @@
 as the ``unstructured-api-key`` header.
 """
 
-from ._base import _ExchangeCallableBase
 from ._exceptions import (
     InvalidCredentialError,
     TokenExchangeDisabledError,
@@ -36,5 +35,4 @@
     "LegacyKeyExchange",
     "TokenExchangeDisabledError",
     "TokenExchangeError",
-    "_ExchangeCallableBase",
 ]

src/unstructured_client/auth/_base.py

@@ -165,12 +165,16 @@ def _raise_for_status(self, response: httpx.Response) -> None:
         """Map HTTP status to auth-specific exceptions before retry decisions."""
         if response.status_code == 401:
             raise InvalidCredentialError(
-                "Account-service rejected the credential (401). Check that the "
-                "client secret / API key is correct and not revoked.",
+                f"Account-service rejected the credential (401) at "
+                f"{self._exchange_url}. Check that the client secret / API "
+                f"key is correct and not revoked.",
             )
         if response.status_code == 400:
             raise TokenExchangeError(
-                f"Account-service rejected the token-exchange request (400): "
+                f"Account-service rejected the token-exchange request (400) "
+                f"at {self._exchange_url}. This commonly indicates the "
+                f"server_url is pointing somewhere other than account-service "
+                f"(for example, at platform-api). Response body: "
                 f"{response.text[:500]}",
             )