Jake/odysseus
1
0
Fork 0
mirror of https://github.com/pewdiepie-archdaemon/odysseus.git synced 2026-07-12 12:37:32 +00:00
Code Packages Projects Releases Activity

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

Browse Source
This commit is contained in:
RaresKeY
2026-07-11 16:06:15 +02:00
committed by GitHub
parent b571c7ddc1
commit 6b8f84553c
2 changed files with 186 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llm_core.py
+32 -2
View File
@@ -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
+154 -11
View File
@@ -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()