diff --git a/src/opengradient/client/llm.py b/src/opengradient/client/llm.py index 9155d36..a568bae 100644 --- a/src/opengradient/client/llm.py +++ b/src/opengradient/client/llm.py @@ -152,9 +152,37 @@ def resolve_tee_connection(self, tee_id: Optional[str] = None) -> ActiveTEE: This is primarily for backend relays that need SDK-managed TEE routing, TLS pinning, and x402 clients without using the chat/completion helpers directly, for example when forwarding OHTTP ciphertext. + + Warning: + Resolving a TEE id other than the active one scans the on-chain + registry with a blocking web3 call. Async servers resolving a TEE + per request should use ``aresolve_tee_connection`` instead. """ return self._tee.resolve(tee_id) + async def aresolve_tee_connection(self, tee_id: Optional[str] = None) -> ActiveTEE: + """Async, event-loop-safe variant of ``resolve_tee_connection``. + + Built for backend relays that resolve a TEE on every request: registry + scans run in a worker thread and each pinned id's outcome (found or + not-active) is cached briefly, so this never blocks the event loop and + doesn't hit the chain RPC per request. Also starts the background TEE + refresh loop, so the active TEE fails over when it is retired from the + registry. + """ + return await self._tee.aresolve(tee_id) + + def ensure_tee_refresh_loop(self) -> None: + """Start the background TEE health-check/failover loop if not running. + + The loop starts lazily from the SDK's own request helpers and from + ``aresolve_tee_connection``. Call this explicitly from server startup + when neither is used on every code path and you still want the active + TEE to fail over once it is retired from the registry. Requires a + running event loop. No-op for static/dev connections. + """ + self._tee.ensure_refresh_loop() + # ── Image helpers ──────────────────────────────────────────────────── @staticmethod diff --git a/src/opengradient/client/tee_connection.py b/src/opengradient/client/tee_connection.py index d3470e9..5ade6b5 100644 --- a/src/opengradient/client/tee_connection.py +++ b/src/opengradient/client/tee_connection.py @@ -3,6 +3,7 @@ import asyncio import logging import ssl +import time from dataclasses import dataclass from typing import Dict, Optional, Protocol, Union @@ -19,6 +20,8 @@ logger = logging.getLogger(__name__) _TEE_REFRESH_INTERVAL = 300 # Re-resolve TEE from registry every 5 minutes +_TEE_RESOLVE_TTL = 30 # Seconds a pinned aresolve() outcome (found or not) stays cached +_ARESOLVE_CACHE_SWEEP_THRESHOLD = 256 # Sweep expired aresolve entries past this size @dataclass(frozen=True) @@ -44,6 +47,7 @@ class TEEConnectionInterface(Protocol): def get(self) -> ActiveTEE: ... def resolve(self, tee_id: Optional[str] = None) -> ActiveTEE: ... + async def aresolve(self, tee_id: Optional[str] = None) -> ActiveTEE: ... def ensure_refresh_loop(self) -> None: ... async def reconnect(self) -> None: ... async def close(self) -> None: ... @@ -98,6 +102,10 @@ def resolve(self, tee_id: Optional[str] = None) -> ActiveTEE: """ return self._active + async def aresolve(self, tee_id: Optional[str] = None) -> ActiveTEE: + """Async variant of ``resolve``; static connections never do I/O.""" + return self._active + def _connect(self) -> ActiveTEE: return ActiveTEE( endpoint=self._endpoint, @@ -143,6 +151,9 @@ def __init__(self, x402_client: x402Client, registry: TEERegistry): self._refresh_task: Optional[asyncio.Task] = None self._active_by_tee_id: dict[str, ActiveTEE] = {} + self._aresolve_lock = asyncio.Lock() + self._aresolve_cache: dict[str, tuple[float, Optional[ActiveTEE]]] = {} + self._active: ActiveTEE = self._connect() # ── Public API ────────────────────────────────────────────────────── @@ -184,6 +195,60 @@ def resolve(self, tee_id: Optional[str] = None) -> ActiveTEE: raise ValueError(f"Selected TEE is not active in the registry: {normalized_tee_id}") + async def aresolve(self, tee_id: Optional[str] = None) -> ActiveTEE: + """Event-loop-safe ``resolve`` for per-request use in async backends. + + ``resolve`` scans the registry with a blocking web3 call whenever the + requested TEE is not the active one, which stalls the event loop when + called per request. This variant runs the scan in a worker thread and + caches each pinned id's outcome — found or not-active — for + ``_TEE_RESOLVE_TTL`` seconds, so steady traffic costs at most one + chain RPC per TTL window per TEE id, and concurrent cold lookups + collapse into a single scan. It also starts the background refresh + loop, so long-running relays fail over when the active TEE is retired + from the registry. + + Raises: + ValueError: If the requested TEE id is not active in the registry + (the miss may be cached for up to ``_TEE_RESOLVE_TTL`` seconds). + """ + self.ensure_refresh_loop() + + normalized_tee_id = _normalize_tee_id(tee_id) + if normalized_tee_id is None or normalized_tee_id == _normalize_tee_id(self._active.tee_id): + return self._active + + cached = self._aresolve_cache.get(normalized_tee_id) + if cached is not None and time.monotonic() - cached[0] < _TEE_RESOLVE_TTL: + return self._require_resolved(normalized_tee_id, cached[1]) + + async with self._aresolve_lock: + cached = self._aresolve_cache.get(normalized_tee_id) + if cached is not None and time.monotonic() - cached[0] < _TEE_RESOLVE_TTL: + return self._require_resolved(normalized_tee_id, cached[1]) + + try: + resolved: Optional[ActiveTEE] = await asyncio.to_thread(self.resolve, normalized_tee_id) + except ValueError: + resolved = None + + # Bogus ids each add a (negative) entry; sweep expired ones once + # the cache outgrows any plausible number of real TEEs so callers + # relaying untrusted ids can't grow it without bound. + if len(self._aresolve_cache) > _ARESOLVE_CACHE_SWEEP_THRESHOLD: + now = time.monotonic() + self._aresolve_cache = {key: entry for key, entry in self._aresolve_cache.items() if now - entry[0] < _TEE_RESOLVE_TTL} + + self._aresolve_cache[normalized_tee_id] = (time.monotonic(), resolved) + + return self._require_resolved(normalized_tee_id, resolved) + + @staticmethod + def _require_resolved(normalized_tee_id: str, resolved: Optional[ActiveTEE]) -> ActiveTEE: + if resolved is None: + raise ValueError(f"Selected TEE is not active in the registry: {normalized_tee_id}") + return resolved + # ── Connection management ─────────────────────────────────────────── def _resolve_tee(self) -> TEEEndpoint: @@ -223,12 +288,17 @@ def _connect_to_tee(self, tee: TEEEndpoint) -> ActiveTEE: ) async def reconnect(self) -> None: - """Connect to a new TEE from the registry and rebuild the HTTP client.""" + """Connect to a new TEE from the registry and rebuild the HTTP client. + + The registry lookup is a blocking web3 call, so it runs in a worker + thread rather than on the event loop. A failed reconnect keeps the + previous connection. + """ async with self._refresh_lock: try: - self._active = self._connect() + self._active = await asyncio.to_thread(self._connect) except Exception: - logger.debug("Failed to close previous HTTP client during TEE refresh.", exc_info=True) + logger.warning("Failed to reconnect to a new TEE; keeping the previous connection.", exc_info=True) # ── Background health check ───────────────────────────────────────── @@ -250,7 +320,8 @@ async def _tee_refresh_loop(self) -> None: while True: await asyncio.sleep(_TEE_REFRESH_INTERVAL) try: - active_tees = self._registry.get_active_tees_by_type(TEE_TYPE_LLM_PROXY) + # Blocking web3 call — keep it off the event loop. + active_tees = await asyncio.to_thread(self._registry.get_active_tees_by_type, TEE_TYPE_LLM_PROXY) if any(t.tee_id == self._active.tee_id for t in active_tees): logger.debug("Current TEE %s still active; no refresh needed.", self._active.tee_id) continue diff --git a/tests/llm_test.py b/tests/llm_test.py index fdc93bf..8d935c1 100644 --- a/tests/llm_test.py +++ b/tests/llm_test.py @@ -852,3 +852,28 @@ async def test_ssl_verification_failure_triggers_tee_refresh_streaming(self, fak assert len(chunks) == 1 assert chunks[0].choices[0].delta.content == "ok" assert len(fake_http.post_calls) == 2 + + +@pytest.mark.asyncio +class TestTeeConnectionApi: + """Public TEE-connection helpers used by backend relays.""" + + async def test_aresolve_tee_connection_returns_active(self, fake_http): + llm = _make_llm() + + resolved = await llm.aresolve_tee_connection() + + assert resolved is llm.resolve_tee_connection() + assert resolved.endpoint == "https://test.tee.server" + + async def test_aresolve_tee_connection_static_ignores_pin(self, fake_http): + llm = _make_llm() + + resolved = await llm.aresolve_tee_connection("0xsome-tee-id") + + assert resolved is llm.resolve_tee_connection() + + async def test_ensure_tee_refresh_loop_is_noop_for_static(self, fake_http): + llm = _make_llm() + + llm.ensure_tee_refresh_loop() # should not raise diff --git a/tests/tee_connection_test.py b/tests/tee_connection_test.py index 20d02f7..6705821 100644 --- a/tests/tee_connection_test.py +++ b/tests/tee_connection_test.py @@ -18,6 +18,7 @@ from opengradient.client.tee_connection import ( ActiveTEE, RegistryTEEConnection, + StaticTEEConnection, ) from opengradient.client.tee_registry import build_ssl_context_from_der @@ -360,6 +361,110 @@ async def test_close_without_refresh_task(self): await conn.close() # should not raise +def _mock_registry_tee(endpoint="https://pinned.tee", tee_id="0xbb", tls_cert_der=b"fake-der", payment_address="0xPay2"): + """A TEE entry as returned by registry.get_active_tees_by_type().""" + tee = MagicMock() + tee.endpoint = endpoint + tee.tls_cert_der = tls_cert_der + tee.tee_id = tee_id + tee.payment_address = payment_address + return tee + + +def _patch_connection_externals(): + """Patch client/ssl construction for resolve calls made after init.""" + return ( + patch( + "opengradient.client.tee_connection.x402HttpxClient", + side_effect=FakeHTTPClient, + ), + patch( + "opengradient.client.tee_connection.build_ssl_context_from_der", + return_value=MagicMock(spec=ssl.SSLContext), + ), + ) + + +@pytest.mark.asyncio +class TestAresolve: + async def test_none_returns_active_without_scan(self): + mock_reg = _mock_registry_with_tee(tee_id="0xaa") + conn = _make_registry_connection(registry=mock_reg) + + assert await conn.aresolve(None) is conn.get() + mock_reg.get_active_tees_by_type.assert_not_called() + await conn.close() + + async def test_matching_active_id_returns_active_without_scan(self): + mock_reg = _mock_registry_with_tee(tee_id="0xaa") + conn = _make_registry_connection(registry=mock_reg) + + # Case/prefix-insensitive match against the active TEE. + assert await conn.aresolve("AA") is conn.get() + mock_reg.get_active_tees_by_type.assert_not_called() + await conn.close() + + async def test_starts_refresh_loop(self): + mock_reg = _mock_registry_with_tee(tee_id="0xaa") + conn = _make_registry_connection(registry=mock_reg) + + await conn.aresolve(None) + + assert conn._refresh_task is not None + assert not conn._refresh_task.done() + await conn.close() + + async def test_pinned_id_scans_once_then_caches(self): + mock_reg = _mock_registry_with_tee(tee_id="0xaa") + mock_reg.get_active_tees_by_type.return_value = [_mock_registry_tee(tee_id="0xbb")] + conn = _make_registry_connection(registry=mock_reg) + + p1, p2 = _patch_connection_externals() + with p1, p2: + results = [await conn.aresolve("0xBB") for _ in range(5)] + + assert all(r.endpoint == "https://pinned.tee" for r in results) + assert mock_reg.get_active_tees_by_type.call_count == 1 + await conn.close() + + async def test_unknown_id_raises_and_caches_miss(self): + mock_reg = _mock_registry_with_tee(tee_id="0xaa") + mock_reg.get_active_tees_by_type.return_value = [] + conn = _make_registry_connection(registry=mock_reg) + + p1, p2 = _patch_connection_externals() + with p1, p2: + for _ in range(3): + with pytest.raises(ValueError, match="not active in the registry"): + await conn.aresolve("0xdead") + + assert mock_reg.get_active_tees_by_type.call_count == 1 + await conn.close() + + async def test_concurrent_cold_lookups_collapse_to_one_scan(self): + mock_reg = _mock_registry_with_tee(tee_id="0xaa") + mock_reg.get_active_tees_by_type.return_value = [_mock_registry_tee(tee_id="0xbb")] + conn = _make_registry_connection(registry=mock_reg) + + p1, p2 = _patch_connection_externals() + with p1, p2: + results = await asyncio.gather(*[conn.aresolve("0xbb") for _ in range(10)]) + + assert all(r.endpoint == "https://pinned.tee" for r in results) + assert mock_reg.get_active_tees_by_type.call_count == 1 + await conn.close() + + async def test_static_connection_resolves_without_io(self): + with patch( + "opengradient.client.tee_connection.x402HttpxClient", + side_effect=FakeHTTPClient, + ): + conn = StaticTEEConnection(x402_client=_mock_x402_client(), endpoint="https://static.tee") + + assert await conn.aresolve("0xanything") is conn.get() + await conn.close() + + # ── TLS certificate verification (real handshake) ────────────────────