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