From 6b8f84553c647b3a66ce33788aded1e21b3ece76 Mon Sep 17 00:00:00 2001 From: RaresKeY <158580472+RaresKeY@users.noreply.github.com> Date: Sat, 11 Jul 2026 16:06:15 +0200 Subject: [PATCH] fix(llm): avoid blocking Kimi Code async header probes (#5231) --- src/llm_core.py | 34 +++++- tests/test_kimi_code_user_agent.py | 165 +++++++++++++++++++++++++++-- 2 files changed, 186 insertions(+), 13 deletions(-) diff --git a/src/llm_core.py b/src/llm_core.py index d8f94dfb7..af1958f16 100644 --- a/src/llm_core.py +++ b/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: diff --git a/tests/test_kimi_code_user_agent.py b/tests/test_kimi_code_user_agent.py index 0d9f1cb01..ed182306c 100644 --- a/tests/test_kimi_code_user_agent.py +++ b/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(url, "Kilo-Code/1.0") - candidates = _kimi_code_ua_candidates(url) + _remember_kimi_code_user_agent(KIMI_CHAT_URL, "Kilo-Code/1.0") + candidates = _kimi_code_ua_candidates(KIMI_CHAT_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(url, {"Authorization": "Bearer x"}, json={}) + r = httpx_post_kimi_aware(KIMI_CHAT_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()