mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-07-12 12:37:32 +00:00
fix(llm): avoid blocking Kimi Code async header probes (#5231)
This commit is contained in:
RaresKeY
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user