Add async TEE resolution and public refresh-loop control

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

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

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

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) ────────────────────