From fd62de3e1ae9d20e50ccc2a1ddbbc2841b9b8d21 Mon Sep 17 00:00:00 2001
From: AMATH <116212274+amathxbt@users.noreply.github.com>
Date: Fri, 3 Apr 2026 22:34:41 +0100
Subject: [PATCH] fix: correct misleading log and swallowed exception in
 RegistryTEEConnection.reconnect()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three bugs fixed:
1. Wrong log message — error said 'close previous HTTP client' but the failure
   was actually in _connect(), not in closing anything.
2. Exception was silently swallowed — callers had no way to know reconnect
   failed and self._active was left pointing to a potentially stale/dead TEE.
3. Resource leak — old HTTP client was never closed on a successful reconnect,
   leaking the underlying TCP connection.

Fix: snapshot old_active before attempting _connect(), re-raise on failure
with a correct warning-level log, and close the old client in the else branch
only after a successful reconnect.
---
 src/opengradient/client/tee_connection.py | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/opengradient/client/tee_connection.py b/src/opengradient/client/tee_connection.py
index 2807b88..9787d76 100644
--- a/src/opengradient/client/tee_connection.py
+++ b/src/opengradient/client/tee_connection.py
@@ -151,12 +151,31 @@ def _connect(self) -> 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.
+
+        On success, closes the previous HTTP client to free the underlying
+        TCP connection.  On failure, the existing connection is kept and the
+        exception is re-raised so callers can take appropriate action.
+        """
         async with self._refresh_lock:
+            old_active = self._active
             try:
                 self._active = 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 from registry; keeping existing connection.",
+                    exc_info=True,
+                )
+                raise
+            else:
+                # Close the old HTTP client now that we have a live replacement.
+                try:
+                    await old_active.http_client.aclose()
+                except Exception:
+                    logger.debug(
+                        "Failed to close previous HTTP client after successful TEE refresh.",
+                        exc_info=True,
+                    )
 
     # ── Background health check ─────────────────────────────────────────