fix(llm): avoid blocking Kimi Code async header probes (#5231)

src/llm_core.py

@@ -766,6 +766,36 @@ def apply_kimi_code_headers(headers: Optional[Dict], url: str) -> Dict[str, str]
     return h
 
 
+async def apply_kimi_code_headers_async(client, headers: Optional[Dict], url: str) -> Dict[str, str]:
+    """Pick a Kimi Code User-Agent without blocking the event loop."""
+    h = dict(headers or {})
+    if not _is_kimi_code_url(url):
+        return h
+    base_key = _kimi_code_base_key(url)
+    cached = _kimi_code_ua_cache.get(base_key)
+    if cached:
+        h["User-Agent"] = cached
+        return h
+    models_url = base_key.rstrip("/") + "/models"
+    for ua in KIMI_CODE_USER_AGENTS:
+        trial = dict(h)
+        trial["User-Agent"] = ua
+        try:
+            r = await client.get(models_url, headers=trial, timeout=8)
+        except Exception:
+            continue
+        if _is_kimi_code_access_denied(r.status_code, r.content):
+            logger.debug("Kimi Code rejected User-Agent %s (403), trying next", ua)
+            continue
+        if r.status_code < 400:
+            _remember_kimi_code_user_agent(url, ua)
+            h["User-Agent"] = ua
+            return h
+        break
+    h.setdefault("User-Agent", KIMI_CODE_USER_AGENT)
+    return h
+
+
 def httpx_get_kimi_aware(url: str, headers: Optional[Dict], **kwargs):
     h = apply_kimi_code_headers(headers, url)
     if not _is_kimi_code_url(url):
@@ -799,7 +829,7 @@ def httpx_post_kimi_aware(url: str, headers: Optional[Dict], **kwargs):
 
 
 async def httpx_post_kimi_aware_async(client, url: str, headers: Optional[Dict], **kwargs):
-    h = apply_kimi_code_headers(headers, url)
+    h = await apply_kimi_code_headers_async(client, headers, url)
     if not _is_kimi_code_url(url):
         return await client.post(url, headers=h, **kwargs)
     last = None
@@ -2466,9 +2496,9 @@ async def _stream_llm_inner(url: str, model: str, messages: List[Dict], temperat
             events.append(_stream_delta_event(part))
         return events
 
-    h = apply_kimi_code_headers(h, target_url)
     try:
         client = _get_http_client()
+        h = await apply_kimi_code_headers_async(client, h, target_url)
         async with client.stream('POST', target_url, json=payload, headers=h, timeout=stream_timeout) as r:
             _clear_host_dead(target_url)
             if r.status_code != 200:

tests/test_kimi_code_user_agent.py

@@ -1,4 +1,7 @@
 """Kimi Code User-Agent fallback list and 403 detection."""
+import pytest
+
+from src import llm_core
 from src.llm_core import (
     KIMI_CODE_USER_AGENTS,
     KIMI_CODE_USER_AGENT,
@@ -12,6 +15,35 @@ from src.llm_core import (
 )
 
 
+KIMI_CHAT_URL = "https://api.kimi.com/coding/v1/chat/completions"
+
+
+class _Resp:
+    def __init__(self, status, text="{}"):
+        self.status_code = status
+        self.content = text.encode()
+        self.text = text
+
+
+class _FakeStreamResp(_Resp):
+    async def aiter_lines(self):
+        yield "data: [DONE]"
+
+    async def aread(self):
+        return b""
+
+
+class _FakeStreamCtx:
+    def __init__(self, response):
+        self.response = response
+
+    async def __aenter__(self):
+        return self.response
+
+    async def __aexit__(self, *args):
+        return False
+
+
 class TestKimiCodeUserAgents:
     def test_default_is_first_fallback(self):
         assert KIMI_CODE_USER_AGENT == KIMI_CODE_USER_AGENTS[0]
@@ -29,9 +61,8 @@ class TestKimiCodeUserAgents:
 
     def test_ua_candidates_prefers_cache(self):
         _kimi_code_ua_cache.clear()
-        url = "https://api.kimi.com/coding/v1/chat/completions"
+        _remember_kimi_code_user_agent(KIMI_CHAT_URL, "Kilo-Code/1.0")
-        _remember_kimi_code_user_agent(url, "Kilo-Code/1.0")
+        candidates = _kimi_code_ua_candidates(KIMI_CHAT_URL)
-        candidates = _kimi_code_ua_candidates(url)
         assert candidates[0] == "Kilo-Code/1.0"
         assert len(candidates) == len(KIMI_CODE_USER_AGENTS)
         _kimi_code_ua_cache.clear()
@@ -48,22 +79,134 @@ class TestKimiCodeUserAgents:
         _kimi_code_ua_cache.clear()
         calls = []
 
-        class _Resp:
-            def __init__(self, status, text=""):
-                self.status_code = status
-                self.content = text.encode()
-                self.text = text
-
         def fake_post(url, headers=None, **kwargs):
             calls.append(headers.get("User-Agent"))
             if headers.get("User-Agent") == KIMI_CODE_USER_AGENTS[0]:
                 return _Resp(403, '{"error":{"type":"access_terminated_error"}}')
             return _Resp(200, "{}")
 
+        monkeypatch.setattr(llm_core.httpx, "get", lambda *a, **k: (_ for _ in ()).throw(RuntimeError()))
         monkeypatch.setattr("src.llm_core.httpx.post", fake_post)
-        url = "https://api.kimi.com/coding/v1/chat/completions"
+        r = httpx_post_kimi_aware(KIMI_CHAT_URL, {"Authorization": "Bearer x"}, json={})
-        r = httpx_post_kimi_aware(url, {"Authorization": "Bearer x"}, json={})
         assert r.status_code == 200
         assert calls[0] == KIMI_CODE_USER_AGENTS[0]
         assert calls[1] == KIMI_CODE_USER_AGENTS[1]
         _kimi_code_ua_cache.clear()
+
+    @pytest.mark.asyncio
+    async def test_async_post_uses_async_probe_not_sync_httpx_get(self, monkeypatch):
+        _kimi_code_ua_cache.clear()
+
+        class FakeClient:
+            def __init__(self):
+                self.get_user_agents = []
+                self.post_user_agents = []
+
+            async def get(self, url, headers=None, **kwargs):
+                self.get_user_agents.append(headers.get("User-Agent"))
+                if headers.get("User-Agent") == KIMI_CODE_USER_AGENTS[0]:
+                    return _Resp(403, '{"error":{"type":"access_terminated_error"}}')
+                return _Resp(200)
+
+            async def post(self, url, headers=None, **kwargs):
+                self.post_user_agents.append(headers.get("User-Agent"))
+                return _Resp(200)
+
+        def forbidden_sync_get(*args, **kwargs):
+            raise AssertionError("async Kimi path must not call sync httpx.get")
+
+        client = FakeClient()
+        monkeypatch.setattr(llm_core.httpx, "get", forbidden_sync_get)
+
+        r = await llm_core.httpx_post_kimi_aware_async(
+            client,
+            KIMI_CHAT_URL,
+            {"Authorization": "Bearer x"},
+            json={},
+        )
+
+        assert r.status_code == 200
+        assert client.get_user_agents == [KIMI_CODE_USER_AGENTS[0], KIMI_CODE_USER_AGENTS[1]]
+        assert client.post_user_agents == [KIMI_CODE_USER_AGENTS[1]]
+        assert _kimi_code_ua_cache[_kimi_code_base_key(KIMI_CHAT_URL)] == KIMI_CODE_USER_AGENTS[1]
+        _kimi_code_ua_cache.clear()
+
+    @pytest.mark.asyncio
+    async def test_async_post_preserves_fallback_when_probe_fails(self, monkeypatch):
+        _kimi_code_ua_cache.clear()
+
+        class FakeClient:
+            def __init__(self):
+                self.post_user_agents = []
+
+            async def get(self, url, headers=None, **kwargs):
+                raise RuntimeError("models probe unavailable")
+
+            async def post(self, url, headers=None, **kwargs):
+                self.post_user_agents.append(headers.get("User-Agent"))
+                if headers.get("User-Agent") == KIMI_CODE_USER_AGENTS[0]:
+                    return _Resp(403, '{"error":{"type":"access_terminated_error"}}')
+                return _Resp(200)
+
+        def forbidden_sync_get(*args, **kwargs):
+            raise AssertionError("async Kimi path must not call sync httpx.get")
+
+        client = FakeClient()
+        monkeypatch.setattr(llm_core.httpx, "get", forbidden_sync_get)
+
+        r = await llm_core.httpx_post_kimi_aware_async(
+            client,
+            KIMI_CHAT_URL,
+            {"Authorization": "Bearer x"},
+            json={},
+        )
+
+        assert r.status_code == 200
+        assert client.post_user_agents == [KIMI_CODE_USER_AGENTS[0], KIMI_CODE_USER_AGENTS[1]]
+        assert _kimi_code_ua_cache[_kimi_code_base_key(KIMI_CHAT_URL)] == KIMI_CODE_USER_AGENTS[1]
+        _kimi_code_ua_cache.clear()
+
+    @pytest.mark.asyncio
+    async def test_stream_uses_async_kimi_probe_not_sync_httpx_get(self, monkeypatch):
+        _kimi_code_ua_cache.clear()
+
+        class FakeClient:
+            def __init__(self):
+                self.get_user_agents = []
+                self.stream_headers = []
+
+            async def get(self, url, headers=None, **kwargs):
+                self.get_user_agents.append(headers.get("User-Agent"))
+                if headers.get("User-Agent") == KIMI_CODE_USER_AGENTS[0]:
+                    return _Resp(403, '{"error":{"type":"access_terminated_error"}}')
+                return _Resp(200)
+
+            def stream(self, method, url, **kwargs):
+                self.stream_headers.append(kwargs.get("headers") or {})
+                return _FakeStreamCtx(_FakeStreamResp(200))
+
+        def forbidden_sync_get(*args, **kwargs):
+            raise AssertionError("streaming Kimi path must not call sync httpx.get")
+
+        client = FakeClient()
+        monkeypatch.setattr(llm_core.httpx, "get", forbidden_sync_get)
+        monkeypatch.setattr(llm_core, "_get_http_client", lambda: client)
+        monkeypatch.setattr(llm_core, "_is_host_dead", lambda url: False)
+        monkeypatch.setattr(llm_core, "note_model_activity", lambda *args, **kwargs: None)
+        monkeypatch.setattr(llm_core, "_clear_host_dead", lambda *args, **kwargs: None)
+
+        chunks = [
+            chunk
+            async for chunk in llm_core.stream_llm(
+                KIMI_CHAT_URL,
+                "kimi-for-coding",
+                [{"role": "user", "content": "hi"}],
+                headers={"Authorization": "Bearer x"},
+            )
+        ]
+
+        assert chunks == ["data: [DONE]\n\n"]
+        assert client.get_user_agents == [KIMI_CODE_USER_AGENTS[0], KIMI_CODE_USER_AGENTS[1]]
+        assert client.stream_headers[0]["User-Agent"] == KIMI_CODE_USER_AGENTS[1]
+        assert _kimi_code_ua_cache[_kimi_code_base_key(KIMI_CHAT_URL)] == KIMI_CODE_USER_AGENTS[1]
+        _kimi_code_ua_cache.clear()