feat: add process-global HTTPX instrumentation

.changelog/httpx-global-instrumentation.md

@@ -0,0 +1,5 @@
+---
+pympp: minor
+---
+
+Added explicit process-global payment instrumentation for HTTPX clients.

.github/workflows/ci.yml

@@ -120,7 +120,9 @@ jobs:
           HTTPX_VERSION: ${{ matrix.httpx-version }}
 
       - name: Run HTTPX adapter tests
-        run: .venv/bin/python -m pytest --noconftest -q tests/test_httpx_adapters.py
+        run: |
+          .venv/bin/python -m pytest --noconftest -q \
+            tests/test_httpx_adapters.py tests/test_instrumentation.py
 
   integration:
     name: Integration Test

README.md

@@ -105,8 +105,27 @@ with OwnedPaymentRuntime([method]) as runtime:
 
 Use `runtime.wrap_async_client(httpx.AsyncClient())` for asynchronous clients.
 These adapters support HTTPX 0.27.x and 0.28.x and fail before changing a client
-when its version or private send seam is incompatible. Explicit `PaymentTransport`
-and `SyncPaymentTransport` integrations do not depend on those private seams.
+when its send seams are incompatible. Explicit `PaymentTransport` and
+`SyncPaymentTransport` integrations do not depend on those private seams.
+
+To cover HTTPX clients created by other libraries, install one explicit
+process-global binding:
+
+```python
+from mpp.instrumentation import instrument
+
+with OwnedPaymentRuntime(
+    [method],
+    allowed_origins=["https://api.example.com"],
+) as runtime:
+    with instrument(runtime):
+        response = httpx.get("https://api.example.com/paid")
+```
+
+The binding applies across threads, and per-client adapters take precedence.
+Install or remove it outside concurrent HTTPX use. An unrestricted runtime must
+be enabled explicitly with `instrument(runtime, allow_unrestricted=True)`.
+Payment methods should not spawn raw threads that make instrumented HTTPX calls.
 
 ## Examples
 

src/mpp/_httpx.py

@@ -104,6 +104,29 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
         return await self._send(request)
 
 
+def _send_sync(
+    runtime: OwnedPaymentRuntime,
+    send: Callable[[httpx.Request], httpx.Response],
+    request: httpx.Request,
+) -> httpx.Response:
+    from mpp.client import SyncPaymentTransport
+
+    return SyncPaymentTransport(inner=_SyncSend(send), runtime=runtime).handle_request(request)
+
+
+async def _send_async(
+    runtime: OwnedPaymentRuntime,
+    send: Callable[[httpx.Request], Awaitable[httpx.Response]],
+    request: httpx.Request,
+) -> httpx.Response:
+    from mpp.client import PaymentTransport
+
+    return await PaymentTransport(
+        inner=_AsyncSend(send),
+        runtime=runtime,
+    ).handle_async_request(request)
+
+
 def wrap_client(client: httpx.Client, runtime: OwnedPaymentRuntime) -> httpx.Client:
     """Make one synchronous HTTPX client payment-aware."""
     if not isinstance(client, httpx.Client):
@@ -192,29 +215,53 @@ def sync_auth(
 def _validate_client(client: Any, *, asynchronous: bool) -> tuple[Any, Any]:
     _validate_version()
     owner = type(client)
+    _validate_class(owner, asynchronous=asynchronous)
     bound = []
-    for name, expected, bound_expected in (
-        ("_send_single_request", _PRIVATE, _BOUND_PRIVATE),
-        ("_send_handling_auth", _AUTH, _BOUND_AUTH),
+    for name, expected in (
+        ("_send_single_request", _BOUND_PRIVATE),
+        ("_send_handling_auth", _BOUND_AUTH),
     ):
-        try:
-            seam = inspect.getattr_static(owner, name)
-        except AttributeError as error:
-            raise HttpxCompatibilityError(
-                f"HTTPX adapter seam {owner.__name__}.{name} is missing"
-            ) from error
-        _validate_seam(f"{owner.__name__}.{name}", seam, expected, asynchronous)
         value = getattr(client, name)
         _validate_seam(
             f"{owner.__name__} instance.{name}",
             value,
-            bound_expected,
+            expected,
             asynchronous,
         )
         bound.append(value)
     return bound[0], bound[1]
 
 
+def _validate_httpx_classes() -> tuple[Any, Any, Any, Any]:
+    _validate_version()
+    return (
+        *_validate_class(httpx.Client, asynchronous=False),
+        *_validate_class(httpx.AsyncClient, asynchronous=True),
+    )
+
+
+def _validate_class(owner: Any, *, asynchronous: bool) -> tuple[Any, Any]:
+    seams = []
+    for name, expected in (("_send_single_request", _PRIVATE), ("_send_handling_auth", _AUTH)):
+        try:
+            seam = inspect.getattr_static(owner, name)
+        except AttributeError as error:
+            raise HttpxCompatibilityError(
+                f"HTTPX adapter seam {owner.__name__}.{name} is missing"
+            ) from error
+        _validate_seam(f"{owner.__name__}.{name}", seam, expected, asynchronous)
+        seams.append(seam)
+    return seams[0], seams[1]
+
+
+def _client_adapter_active(client: Any) -> bool:
+    adapter = client.__dict__.get(_MARKER)
+    return isinstance(adapter, _Adapter) and (
+        inspect.getattr_static(client, "_send_single_request") is adapter.private
+        and inspect.getattr_static(client, "_send_handling_auth") is adapter.auth
+    )
+
+
 def _validate_version() -> None:
     installed = version("httpx")
     try:

src/mpp/instrumentation.py

@@ -0,0 +1,204 @@
+"""Explicit process-global HTTPX payment instrumentation."""
+
+from __future__ import annotations
+
+import inspect
+import threading
+from dataclasses import dataclass
+from functools import partial, wraps
+from typing import Any
+
+import httpx
+
+from mpp._httpx import (
+    _client_adapter_active,
+    _payment_budget,
+    _send_async,
+    _send_operation,
+    _send_sync,
+    _validate_httpx_classes,
+)
+from mpp.runtime import OwnedPaymentRuntime, _owned_in_scope
+
+
+@dataclass(frozen=True, slots=True)
+class _Patch:
+    owner: type[Any]
+    name: str
+    original: Any
+    wrapper: Any
+
+
+@dataclass(slots=True)
+class _Binding:
+    runtime: OwnedPaymentRuntime
+    patches: tuple[_Patch, ...] = ()
+    references: int = 1
+
+
+_LOCK = threading.RLock()
+_binding: _Binding | None = None
+
+
+class Instrumentation:
+    """A reference-counted handle for global HTTPX instrumentation."""
+
+    def __init__(self, binding: _Binding) -> None:
+        self._binding: _Binding | None = binding
+
+    def disable(self) -> None:
+        """Release this handle and restore pympp-owned patches."""
+        global _binding
+
+        with _LOCK:
+            binding = self._binding
+            if binding is None:
+                return
+            if _binding is not binding:
+                self._binding = None
+                return
+            if binding.references > 1:
+                binding.references -= 1
+            else:
+                owned = tuple(
+                    patch
+                    for patch in binding.patches
+                    if inspect.getattr_static(patch.owner, patch.name) is patch.wrapper
+                )
+                _replace(owned, enable=False)
+                _binding = None
+            self._binding = None
+
+    def __enter__(self) -> Instrumentation:
+        return self
+
+    def __exit__(self, *_args: Any) -> None:
+        self.disable()
+
+
+def instrument(
+    runtime: OwnedPaymentRuntime,
+    *,
+    allow_unrestricted: bool = False,
+) -> Instrumentation:
+    """Make otherwise-unadapted HTTPX clients payment-aware."""
+    global _binding
+
+    if not isinstance(runtime, OwnedPaymentRuntime):
+        raise TypeError("instrument requires an OwnedPaymentRuntime")
+    if runtime._allowed_origins._allow_all and not allow_unrestricted:
+        raise ValueError(
+            "Global instrumentation requires allowed_origins or allow_unrestricted=True"
+        )
+
+    with _LOCK:
+        if _binding is not None:
+            if _binding.runtime is not runtime:
+                raise RuntimeError("Another payment runtime is already globally instrumented")
+            if not all(
+                inspect.getattr_static(patch.owner, patch.name) is patch.wrapper
+                for patch in _binding.patches
+            ):
+                raise RuntimeError("Global HTTPX instrumentation was modified while active")
+            _binding.references += 1
+            return Instrumentation(_binding)
+
+        binding = _Binding(runtime)
+        binding.patches = _make_patches(binding, _validate_httpx_classes())
+        _replace(binding.patches, enable=True)
+        _binding = binding
+        return Instrumentation(binding)
+
+
+def _make_patches(binding: _Binding, seams: tuple[Any, Any, Any, Any]) -> tuple[_Patch, ...]:
+    sync_private, sync_auth, async_private, async_auth = seams
+
+    @wraps(sync_private)
+    def patched_sync_private(client: httpx.Client, request: httpx.Request) -> httpx.Response:
+        if _bypass(binding, client):
+            return sync_private(client, request)
+        with _payment_budget(request):
+            return _send_sync(binding.runtime, partial(sync_private, client), request)
+
+    @wraps(sync_auth)
+    def patched_sync_auth(
+        client: httpx.Client,
+        request: httpx.Request,
+        *args: Any,
+        **kwargs: Any,
+    ) -> httpx.Response:
+        if _bypass(binding, client):
+            return sync_auth(client, request, *args, **kwargs)
+        with _send_operation():
+            return sync_auth(client, request, *args, **kwargs)
+
+    @wraps(async_private)
+    async def patched_async_private(
+        client: httpx.AsyncClient,
+        request: httpx.Request,
+    ) -> httpx.Response:
+        if _bypass(binding, client):
+            return await async_private(client, request)
+        with _payment_budget(request):
+            return await _send_async(
+                binding.runtime,
+                partial(async_private, client),
+                request,
+            )
+
+    @wraps(async_auth)
+    async def patched_async_auth(
+        client: httpx.AsyncClient,
+        request: httpx.Request,
+        *args: Any,
+        **kwargs: Any,
+    ) -> httpx.Response:
+        if _bypass(binding, client):
+            return await async_auth(client, request, *args, **kwargs)
+        with _send_operation():
+            return await async_auth(client, request, *args, **kwargs)
+
+    return (
+        _Patch(httpx.Client, "_send_single_request", sync_private, patched_sync_private),
+        _Patch(httpx.Client, "_send_handling_auth", sync_auth, patched_sync_auth),
+        _Patch(httpx.AsyncClient, "_send_single_request", async_private, patched_async_private),
+        _Patch(httpx.AsyncClient, "_send_handling_auth", async_auth, patched_async_auth),
+    )
+
+
+def _bypass(binding: _Binding, client: Any) -> bool:
+    return (
+        _binding is not binding
+        or _client_adapter_active(client)
+        or threading.get_ident() == binding.runtime._owner_thread_id
+        or _owned_in_scope(binding.runtime._scope_key)
+    )
+
+
+def _replace(patches: tuple[_Patch, ...], *, enable: bool) -> None:
+    changed: list[_Patch] = []
+    try:
+        for patch in patches:
+            expected = patch.original if enable else patch.wrapper
+            replacement = patch.wrapper if enable else patch.original
+            if inspect.getattr_static(patch.owner, patch.name) is not expected:
+                raise RuntimeError("HTTPX changed while instrumentation was being updated")
+            changed.append(patch)
+            _assign(patch.owner, patch.name, replacement)
+    except BaseException as error:
+        rollback_error: BaseException | None = None
+        for patch in reversed(changed):
+            expected = patch.original if enable else patch.wrapper
+            replacement = patch.wrapper if enable else patch.original
+            try:
+                if inspect.getattr_static(patch.owner, patch.name) is replacement:
+                    _assign(patch.owner, patch.name, expected)
+            except BaseException as cause:
+                rollback_error = rollback_error or cause
+        if rollback_error is not None:
+            raise RuntimeError("Failed to roll back HTTPX instrumentation") from error
+        raise
+
+
+def _assign(owner: type[Any], name: str, value: Any) -> None:
+    setattr(owner, name, value)