test handling

Author: balogh.adam@icloud.com

src/opengradient/client/llm.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from typing import AsyncGenerator, Awaitable, Callable, Dict, List, Optional, TypeVar, Union
 import httpx
+import asyncio
 
 from eth_account import Account
 from eth_account.account import LocalAccount
@@ -170,6 +171,8 @@ async def _call_with_tee_retry(
             return await call()
         except httpx.HTTPStatusError:
             raise
+        except asyncio.CancelledError:
+            raise
         except Exception as exc:
             logger.warning(
                 "Connection failure during %s; refreshing TEE and retrying once: %s",
@@ -424,6 +427,8 @@ async def _chat_stream(self, params: _ChatParams, messages: List[Dict]) -> Async
             return
         except httpx.HTTPStatusError:
             raise
+        except asyncio.CancelledError:
+            raise
         except Exception as exc:
             if chunks_yielded:
                 raise

src/opengradient/client/tee_connection.py

@@ -133,7 +133,7 @@ def _resolve_tee(self):
             raise RuntimeError(f"Failed to fetch LLM TEE endpoint from registry: {e}") from e
 
         if tee is None:
-            raise ValueError("No active LLM proxy TEE found in the registry. Pass llm_server_url explicitly to override.")
+            raise ValueError("No active LLM proxy TEE found in the registry.")
 
         logger.info("Using TEE endpoint from registry: %s (teeId=%s)", tee.endpoint, tee.tee_id)
         return tee
@@ -155,10 +155,8 @@ def _connect(self) -> ActiveTEE:
     async def reconnect(self) -> None:
         """Connect to a new TEE from the registry and rebuild the HTTP client."""
         async with self._refresh_lock:
-            old_client = self._active.http_client
-            self._active = self._connect()
             try:
-                await old_client.aclose()
+                self._active = self._connect()
             except Exception:
                 logger.debug("Failed to close previous HTTP client during TEE refresh.", exc_info=True)
 
@@ -188,6 +186,9 @@ async def _tee_refresh_loop(self) -> None:
                     continue
                 logger.info("Current TEE %s no longer active; switching to a new one.", self._active.tee_id)
                 await self.reconnect()
+            except asyncio.CancelledError:
+                logger.debug("Background TEE health check cancelled; exiting loop.")
+                raise
             except Exception:
                 logger.warning("Background TEE health check failed; will retry next cycle.", exc_info=True)