From 0e4daade1463de4e6c7c2688a8ad6638c3e15629 Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Mon, 3 Aug 2026 17:30:28 +0530
Subject: [PATCH] perf(zai): early-exit when the highest-priority endpoint wins
 (simplify finding)

The as_completed drain + `with` join made the parallel version WORSE than
sequential main in the common case (first endpoint succeeds fast, others
slow/unreachable): main returned at first success, the parallel version
waited for every straggler. Now: after each completion, walk endpoints in
priority order and return as soon as a success is unbeatable (all
higher-priority probes already finished); pool uses shutdown(wait=False) so
losers drain in the background. Mutation-checked: removing the early exit
makes the new timing test fail (8.2s vs <1.5s).
---
 hermes_cli/auth.py                         | 30 +++++++++++++++++-----
 tests/hermes_cli/test_api_key_providers.py | 22 ++++++++++++++++
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/hermes_cli/auth.py b/hermes_cli/auth.py
index 2ab6227b39961..266f1f42b3c3f 100644
--- a/hermes_cli/auth.py
+++ b/hermes_cli/auth.py
@@ -698,11 +698,17 @@ def detect_zai_endpoint(api_key: str, timeout: float = 8.0) -> Optional[Dict[str
     """
     from concurrent.futures import ThreadPoolExecutor, as_completed
 
-    with ThreadPoolExecutor(max_workers=len(ZAI_ENDPOINTS)) as pool:
+    # No `with` block: a context manager would join ALL probe threads on
+    # exit, defeating the early return below. shutdown(wait=False) lets the
+    # surviving daemon-style probes drain in the background instead of
+    # blocking the caller on slow/unreachable endpoints.
+    pool = ThreadPoolExecutor(max_workers=len(ZAI_ENDPOINTS))
+    try:
         futures = {
             pool.submit(_probe_single_zai_endpoint, api_key, ep, timeout): ep[0]
             for ep in ZAI_ENDPOINTS
         }
+        by_id = {ep_id: f for f, ep_id in futures.items()}
         results: Dict[str, Dict[str, str]] = {}
         for future in as_completed(futures):
             ep_id = futures[future]
@@ -712,12 +718,24 @@ def detect_zai_endpoint(api_key: str, timeout: float = 8.0) -> Optional[Dict[str
                     results[ep_id] = result
             except Exception:
                 pass
+            # Early exit in PRIORITY order: walk endpoints highest-priority
+            # first; if one has succeeded and every higher-priority probe
+            # has already finished (without success), no later completion
+            # can win — return now instead of waiting out slow endpoints
+            # (main's sequential loop also stopped at first success).
+            for ep in ZAI_ENDPOINTS:
+                if not by_id[ep[0]].done():
+                    break  # a higher-priority probe is still in flight
+                if ep[0] in results:
+                    return results[ep[0]]
 
-    # Return first match in priority order (ZAI_ENDPOINTS list order)
-    for ep in ZAI_ENDPOINTS:
-        if ep[0] in results:
-            return results[ep[0]]
-    return None
+        # All probes finished: first match in priority order, if any.
+        for ep in ZAI_ENDPOINTS:
+            if ep[0] in results:
+                return results[ep[0]]
+        return None
+    finally:
+        pool.shutdown(wait=False)
 
 
 def _resolve_zai_base_url(api_key: str, default_url: str, env_override: str) -> str:
diff --git a/tests/hermes_cli/test_api_key_providers.py b/tests/hermes_cli/test_api_key_providers.py
index 1b27186882eb3..775b5ae84e3a0 100644
--- a/tests/hermes_cli/test_api_key_providers.py
+++ b/tests/hermes_cli/test_api_key_providers.py
@@ -712,6 +712,28 @@ class TestZaiParallelProbe:
         monkeypatch.setattr("hermes_cli.auth.httpx.post", self._mock_post({}))
         assert detect_zai_endpoint("bad-key", timeout=1.0) is None
 
+    def test_early_exit_does_not_wait_for_slow_losers(self, monkeypatch):
+        """When the highest-priority endpoint succeeds fast, the caller must
+        return without waiting for slow lower-priority probes to finish."""
+        import time as _time
+
+        from hermes_cli.auth import ZAI_ENDPOINTS, detect_zai_endpoint
+
+        first = ZAI_ENDPOINTS[0]
+        inner = self._mock_post({(first[1], first[2][0]): True})
+
+        def _slow_losers(url, headers=None, json=None, timeout=None):
+            if not url.startswith(first[1]):
+                _time.sleep(2.0)  # slow lower-priority endpoints
+            return inner(url, headers=headers, json=json, timeout=timeout)
+
+        monkeypatch.setattr("hermes_cli.auth.httpx.post", _slow_losers)
+        t0 = _time.perf_counter()
+        result = detect_zai_endpoint("test-key", timeout=5.0)
+        elapsed = _time.perf_counter() - t0
+        assert result is not None and result["id"] == first[0]
+        assert elapsed < 1.5, f"early exit failed: waited {elapsed:.2f}s for losers"
+
 
 # =============================================================================
 # Kimi / Moonshot model list isolation tests