feat: 引入 Rust executor/gateway sidecar 及 Python 侧双后端适配

- 新增 Rust workspace crates: aether-contracts, aether-executor, aether-gateway
- aether-executor: 支持 Unix Socket/TCP 双传输模式,处理同步/流式上游请求
- aether-gateway: 作为本地主入口代理,集成 /api/internal/gateway/resolve 认证预解析
- Python 侧新增 ExecutionPlan 契约和 RustExecutorClient,各 handler 支持
  executor_backend=rust 时将可序列化请求转发给 Rust executor 执行
- 重构 dev.sh 支持 executor/gateway 进程编排与生命周期管理
- 新增 internal gateway 路由,提供 resolve/passthrough 端点
- handler 层(chat/cli/video/endpoint_checker 等)全面适配 Rust executor 回退逻辑
- pipeline 层支持 trusted auth context 跳过重复认证
- 新增 Rust CI workflow 及对应测试用例
This commit is contained in:
fawney19
2026-03-21 12:57:09 +08:00
parent 46737d32f8
commit d735b6316f
79 changed files with 19032 additions and 522 deletions

View File

@@ -0,0 +1,565 @@
from __future__ import annotations
from types import SimpleNamespace
import httpx
import pytest
import src.api.handlers.base.chat_sync_executor as chat_sync_mod
from src.api.handlers.base.chat_sync_executor import ChatSyncExecutor
from src.core.exceptions import EmbeddedErrorException
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
PreparedExecutionPlan,
)
from src.services.request.rust_executor_client import (
RustExecutorClientError,
RustExecutorSyncResult,
)
class _FakeEnvelope:
name = "fake-envelope"
def __init__(self) -> None:
self.status_codes: list[int] = []
self.postprocessed_payloads: list[dict[str, object]] = []
def on_http_status(self, *, base_url: str | None, status_code: int) -> None:
self.status_codes.append(status_code)
def on_connection_error(self, *, base_url: str | None, exc: Exception) -> None:
raise AssertionError("connection error hook should not be used in this test")
def unwrap_response(self, data: dict[str, object]) -> dict[str, object]:
return dict(data["payload"]) # type: ignore[index]
def postprocess_unwrapped_response(self, *, model: str, data: dict[str, object]) -> None:
self.postprocessed_payloads.append(dict(data))
class _FakeNormalizer:
def response_from_internal(
self,
internal_resp: object,
*,
requested_model: str,
) -> dict[str, object]:
return {
"aggregated": True,
"requested_model": requested_model,
"internal_id": getattr(internal_resp, "id", "missing"),
}
def _make_prepared_plan() -> PreparedExecutionPlan:
return PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-test",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="POST",
url="https://example.com/v1/chat/completions",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(json_body={"model": "gpt-4.1"}),
stream=False,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
),
payload={"model": "gpt-4.1"},
headers={"content-type": "application/json"},
upstream_is_stream=False,
needs_conversion=False,
provider_type="openai",
request_timeout=30.0,
)
def _make_proxy_prepared_plan() -> PreparedExecutionPlan:
prepared = _make_prepared_plan()
prepared.contract.proxy = chat_sync_mod.ExecutionProxySnapshot(
enabled=True,
mode="http",
label="proxy.internal",
url="http://proxy.internal:8080",
)
prepared.proxy_config = {"url": "http://proxy.internal:8080"}
return prepared
def _make_tunnel_prepared_plan() -> PreparedExecutionPlan:
prepared = _make_prepared_plan()
prepared.contract.proxy = chat_sync_mod.ExecutionProxySnapshot(
enabled=True,
mode="tunnel",
node_id="node-1",
label="relay-node",
)
prepared.delegate_config = {"tunnel": True, "node_id": "node-1"}
prepared.proxy_config = {"node_id": "node-1"}
return prepared
def _make_upstream_stream_prepared_plan() -> PreparedExecutionPlan:
prepared = _make_prepared_plan()
prepared.contract.stream = True
prepared.upstream_is_stream = True
return prepared
def _make_tls_prepared_plan() -> PreparedExecutionPlan:
prepared = _make_prepared_plan()
prepared.contract.tls_profile = "claude_code_nodejs"
prepared.provider_type = "claude_code"
return prepared
def _make_executor() -> ChatSyncExecutor:
handler = SimpleNamespace(request_id="req-test")
executor = ChatSyncExecutor(handler)
executor._ctx.provider_api_format_for_error = "openai:chat"
executor._ctx.client_api_format_for_error = "openai:chat"
executor._ctx.needs_conversion_for_error = False
return executor
@pytest.mark.asyncio
async def test_execute_sync_plan_uses_rust_executor_when_available(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_prepared_plan()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.request_id == "req-test"
return RustExecutorSyncResult(
status_code=200,
response_json={"id": "chatcmpl-1"},
headers={"content-type": "application/json"},
)
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
raise AssertionError("local execution should not be used")
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
response = await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert response == {"id": "chatcmpl-1"}
assert executor._ctx.status_code == 200
assert executor._ctx.response_json == {"id": "chatcmpl-1"}
@pytest.mark.asyncio
async def test_execute_sync_plan_allows_supported_proxy_urls_for_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_proxy_prepared_plan()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.proxy is not None
assert plan.proxy.url == "http://proxy.internal:8080"
return RustExecutorSyncResult(
status_code=200,
response_json={"id": "chatcmpl-proxy"},
headers={"content-type": "application/json"},
)
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
raise AssertionError("local execution should not be used")
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
response = await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert response == {"id": "chatcmpl-proxy"}
@pytest.mark.asyncio
async def test_execute_sync_plan_allows_tunnel_delegate_for_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_tunnel_prepared_plan()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.proxy is not None
assert plan.proxy.mode == "tunnel"
assert plan.proxy.node_id == "node-1"
return RustExecutorSyncResult(
status_code=200,
response_json={"id": "chatcmpl-tunnel"},
headers={"content-type": "application/json"},
)
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
raise AssertionError("local execution should not be used")
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
response = await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert response == {"id": "chatcmpl-tunnel"}
@pytest.mark.asyncio
async def test_execute_sync_plan_allows_tls_profile_for_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_tls_prepared_plan()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.tls_profile == "claude_code_nodejs"
return RustExecutorSyncResult(
status_code=200,
response_json={"id": "chatcmpl-tls"},
headers={"content-type": "application/json"},
)
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
raise AssertionError("local execution should not be used")
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
response = await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="claude-3.7-sonnet",
)
assert response == {"id": "chatcmpl-tls"}
@pytest.mark.asyncio
async def test_execute_sync_plan_applies_envelope_postprocessing_after_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_prepared_plan()
prepared_plan.envelope = _FakeEnvelope()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
return RustExecutorSyncResult(
status_code=200,
response_json={"payload": {"id": "wrapped-1", "message": "ok"}},
headers={"content-type": "application/json"},
)
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
raise AssertionError("local execution should not be used")
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
response = await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert response == {"id": "wrapped-1", "message": "ok"}
assert prepared_plan.envelope.status_codes == [200]
assert prepared_plan.envelope.postprocessed_payloads == [
{"id": "wrapped-1", "message": "ok"}
]
@pytest.mark.asyncio
async def test_execute_sync_plan_applies_format_conversion_after_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_prepared_plan()
prepared_plan.needs_conversion = True
prepared_plan.contract.provider_api_format = "gemini:chat"
prepared_plan.contract.client_api_format = "openai:chat"
executor._ctx.provider_api_format_for_error = "gemini:chat"
executor._ctx.client_api_format_for_error = "openai:chat"
executor._ctx.needs_conversion_for_error = True
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
class _FakeRegistry:
def convert_response(
self,
response_json: dict[str, object],
source_format: str,
target_format: str,
*,
requested_model: str,
) -> dict[str, object]:
assert source_format == "gemini:chat"
assert target_format == "openai:chat"
assert requested_model == "gpt-4.1"
return {
"converted": True,
"source_id": response_json["provider_id"],
}
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.provider_api_format == "gemini:chat"
return RustExecutorSyncResult(
status_code=200,
response_json={"provider_id": "gemini-1"},
headers={"content-type": "application/json"},
)
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
raise AssertionError("local execution should not be used")
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(chat_sync_mod, "get_format_converter_registry", lambda: _FakeRegistry())
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
response = await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert response == {"converted": True, "source_id": "gemini-1"}
assert executor._ctx.provider_response_json == {"provider_id": "gemini-1"}
@pytest.mark.asyncio
async def test_execute_sync_plan_aggregates_upstream_stream_after_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_upstream_stream_prepared_plan()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
class _FakeRegistry:
def get_normalizer(self, format_id: str) -> _FakeNormalizer:
assert format_id == "openai:chat"
return _FakeNormalizer()
captured_chunks: list[bytes] = []
async def _fake_aggregate(
byte_iter: object,
*,
provider_api_format: str,
provider_name: str,
model: str,
request_id: str,
envelope: object = None,
provider_parser: object = None,
) -> object:
async for chunk in byte_iter: # type: ignore[attr-defined]
captured_chunks.append(chunk)
assert provider_api_format == "openai:chat"
assert provider_name == "provider"
assert model == "gpt-4.1"
assert request_id == "req-test"
return SimpleNamespace(id="agg-1")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.stream is True
return RustExecutorSyncResult(
status_code=200,
response_body_bytes=b"data: {\"id\":\"chunk-1\"}\n\ndata: [DONE]\n\n",
headers={"content-type": "text/event-stream"},
)
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
raise AssertionError("local execution should not be used")
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(chat_sync_mod, "get_format_converter_registry", lambda: _FakeRegistry())
monkeypatch.setattr(
"src.api.handlers.base.upstream_stream_bridge.aggregate_upstream_stream_to_internal_response",
_fake_aggregate,
)
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
response = await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert response == {
"aggregated": True,
"requested_model": "gpt-4.1",
"internal_id": "agg-1",
}
assert captured_chunks == [b"data: {\"id\":\"chunk-1\"}\n\ndata: [DONE]\n\n"]
assert executor._ctx.status_code == 200
@pytest.mark.asyncio
async def test_execute_sync_plan_turns_rust_http_error_into_httpx_status_error(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_prepared_plan()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.url.endswith("/chat/completions")
return RustExecutorSyncResult(
status_code=429,
response_json={"error": {"message": "slow down"}},
headers={"content-type": "application/json"},
)
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
raise AssertionError("local execution should not be used")
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
with pytest.raises(httpx.HTTPStatusError) as exc_info:
await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert exc_info.value.response.status_code == 429
assert '"message": "slow down"' in exc_info.value.upstream_response # type: ignore[attr-defined]
@pytest.mark.asyncio
async def test_execute_sync_plan_preserves_embedded_error_semantics_from_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_prepared_plan()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.provider_api_format == "openai:chat"
return RustExecutorSyncResult(
status_code=200,
response_json={
"error": {
"message": "bad request",
"type": "invalid_request_error",
"code": 400,
}
},
headers={"content-type": "application/json"},
)
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
with pytest.raises(EmbeddedErrorException) as exc_info:
await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert exc_info.value.error_message == "bad request"
assert exc_info.value.error_code == 400
@pytest.mark.asyncio
async def test_execute_sync_plan_falls_back_to_local_when_rust_unavailable(
monkeypatch: pytest.MonkeyPatch,
) -> None:
executor = _make_executor()
prepared_plan = _make_prepared_plan()
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
assert plan.request_id == "req-test"
raise RustExecutorClientError("executor down")
fallback_called = False
async def _fake_local_execute(**kwargs: object) -> dict[str, object]:
nonlocal fallback_called
fallback_called = True
return {"id": "local-fallback"}
monkeypatch.setattr(
chat_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _fake_local_execute)
response = await executor._execute_sync_plan(
prepared_plan=prepared_plan,
provider=SimpleNamespace(name="provider"),
model="gpt-4.1",
)
assert fallback_called is True
assert response == {"id": "local-fallback"}

View File

@@ -0,0 +1,553 @@
from __future__ import annotations
from collections.abc import AsyncGenerator
from types import SimpleNamespace
from typing import Any
import httpx
import pytest
import src.api.handlers.base.chat_handler_base as chatmod
import src.services.proxy_node.resolver as proxymod
from src.api.handlers.base.chat_handler_base import ChatHandlerBase
from src.api.handlers.base.stream_context import StreamContext
from src.services.request.rust_executor_client import (
RustExecutorClientError,
RustExecutorStreamResult,
)
class _DummyAuthInfo:
auth_header = "authorization"
auth_value = "Bearer test"
decrypted_auth_config = None
def as_tuple(self) -> tuple[str, str]:
return self.auth_header, self.auth_value
class _PassBuilder:
def build(self, request_body: dict[str, Any], *args: Any, **kwargs: Any) -> Any:
return request_body, {"content-type": "application/json"}
class _DummyStreamResponseCtx:
def __init__(self) -> None:
self.closed = False
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
self.closed = True
class _FakeStreamProcessor:
def __init__(self) -> None:
self.prefetched_chunks: list[bytes] | None = None
self.response_ctx: _DummyStreamResponseCtx | None = None
async def prefetch_and_check_error(
self,
byte_iterator: Any,
provider: Any,
endpoint: Any,
ctx: Any,
max_prefetch_lines: int = 5,
max_prefetch_bytes: int = 65536,
) -> list[bytes]:
del provider, endpoint, ctx, max_prefetch_lines, max_prefetch_bytes
first = await anext(byte_iterator)
self.prefetched_chunks = [first]
return self.prefetched_chunks
async def create_response_stream(
self,
ctx: Any,
byte_iterator: Any,
response_ctx: _DummyStreamResponseCtx,
prefetched_chunks: list[bytes] | None = None,
*,
start_time: float | None = None,
) -> AsyncGenerator[bytes]:
del ctx, start_time
self.response_ctx = response_ctx
try:
for chunk in prefetched_chunks or []:
yield chunk
async for chunk in byte_iterator:
yield chunk
finally:
await response_ctx.__aexit__(None, None, None)
class _DummyChatHandler(ChatHandlerBase):
FORMAT_ID = "openai:chat"
def __init__(self) -> None:
self.request_id = "req-test"
self.api_key = SimpleNamespace(id="user-key-1")
self._request_builder = _PassBuilder()
self.allowed_api_formats = ["openai:chat"]
self.api_family = None
self.endpoint_kind = None
self.start_time = 0.0
async def _convert_request(self, request: Any) -> Any:
return request
def _extract_usage(self, response: dict) -> dict[str, int]:
return {}
async def _get_mapped_model(
self,
source_model: str,
provider_id: str,
api_format: str | None = None,
) -> str | None:
del source_model, provider_id, api_format
return None
def apply_mapped_model(self, request_body: dict[str, Any], mapped_model: str) -> dict[str, Any]:
out = dict(request_body)
out["model"] = mapped_model
return out
def prepare_provider_request_body(self, request_body: dict[str, Any]) -> dict[str, Any]:
return dict(request_body)
def finalize_provider_request(
self,
request_body: dict[str, Any],
*,
mapped_model: str | None,
provider_api_format: str | None,
) -> dict[str, Any]:
del mapped_model, provider_api_format
return dict(request_body)
def get_model_for_url(
self,
request_body: dict[str, Any],
mapped_model: str | None,
) -> str | None:
return mapped_model or str(request_body.get("model") or "")
def _patch_stream_setup(
monkeypatch: pytest.MonkeyPatch,
*,
proxy_info: dict[str, Any] | None = None,
delegate_config: dict[str, Any] | None = None,
) -> None:
async def _fake_get_provider_auth(endpoint: Any, key: Any) -> _DummyAuthInfo:
del endpoint, key
return _DummyAuthInfo()
async def _fake_resolve_proxy_info(proxy_config: Any) -> Any:
del proxy_config
return proxy_info
async def _fake_resolve_delegate(proxy_config: Any) -> Any:
del proxy_config
return delegate_config
async def _fake_get_system_proxy() -> None:
return None
monkeypatch.setattr(chatmod, "get_provider_auth", _fake_get_provider_auth)
monkeypatch.setattr(
chatmod,
"get_provider_behavior",
lambda **kwargs: SimpleNamespace(
envelope=None,
same_format_variant=None,
cross_format_variant=None,
),
)
monkeypatch.setattr(chatmod, "build_provider_url", lambda *args, **kwargs: "https://upstream.test/v1/chat/completions")
monkeypatch.setattr(chatmod, "get_upstream_stream_policy", lambda *args, **kwargs: None)
monkeypatch.setattr(
chatmod,
"resolve_upstream_is_stream",
lambda *, client_is_stream, policy: client_is_stream,
)
monkeypatch.setattr(chatmod, "enforce_stream_mode_for_upstream", lambda *args, **kwargs: None)
monkeypatch.setattr(
chatmod,
"maybe_patch_request_with_prompt_cache_key",
lambda request_body, **kwargs: request_body,
)
monkeypatch.setattr(proxymod, "resolve_effective_proxy", lambda provider_proxy, key_proxy=None: None)
monkeypatch.setattr(proxymod, "resolve_proxy_info_async", _fake_resolve_proxy_info)
monkeypatch.setattr(proxymod, "get_proxy_label", lambda proxy_info: "direct")
monkeypatch.setattr(proxymod, "resolve_delegate_config_async", _fake_resolve_delegate)
monkeypatch.setattr(proxymod, "get_system_proxy_config_async", _fake_get_system_proxy)
monkeypatch.setattr(proxymod, "build_proxy_url_async", _fake_get_system_proxy)
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
for chunk in chunks:
yield chunk
@pytest.mark.asyncio
async def test_execute_stream_request_uses_rust_executor_when_available(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_patch_stream_setup(monkeypatch)
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
handler = _DummyChatHandler()
stream_processor = _FakeStreamProcessor()
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
ctx.client_api_format = "openai:chat"
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
key = SimpleNamespace(id="key-1", proxy=None)
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
dummy_ctx = _DummyStreamResponseCtx()
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "stream") is True
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "text/event-stream", "x-upstream-test": "true"},
byte_iterator=_iter_chunks(
[
b"data: {\"id\":\"chunk-1\"}\n\n",
b"data: [DONE]\n\n",
]
),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
stream = await handler._execute_stream_request(
ctx,
stream_processor,
provider,
endpoint,
key,
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
{},
candidate=candidate,
)
received = [chunk async for chunk in stream]
assert received == [
b"data: {\"id\":\"chunk-1\"}\n\n",
b"data: [DONE]\n\n",
]
assert ctx.status_code == 200
assert ctx.response_headers["x-upstream-test"] == "true"
assert stream_processor.prefetched_chunks == [b"data: {\"id\":\"chunk-1\"}\n\n"]
assert dummy_ctx.closed is True
@pytest.mark.asyncio
async def test_execute_stream_request_accepts_async_generator_stream_processor(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_patch_stream_setup(monkeypatch)
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
handler = _DummyChatHandler()
stream_processor = _FakeStreamProcessor()
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
ctx.client_api_format = "openai:chat"
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
key = SimpleNamespace(id="key-1", proxy=None)
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
dummy_ctx = _DummyStreamResponseCtx()
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "stream") is True
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "text/event-stream"},
byte_iterator=_iter_chunks(
[
b"data: {\"id\":\"chunk-1\"}\n\n",
b"data: [DONE]\n\n",
]
),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
stream = await handler._execute_stream_request(
ctx,
stream_processor,
provider,
endpoint,
key,
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
{},
candidate=candidate,
)
received = [chunk async for chunk in stream]
assert received == [
b"data: {\"id\":\"chunk-1\"}\n\n",
b"data: [DONE]\n\n",
]
assert dummy_ctx.closed is True
@pytest.mark.asyncio
async def test_execute_stream_request_allows_tunnel_delegate_for_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_patch_stream_setup(
monkeypatch,
proxy_info={"node_id": "node-1", "node_name": "relay-node", "mode": "tunnel"},
delegate_config={"tunnel": True, "node_id": "node-1"},
)
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
handler = _DummyChatHandler()
stream_processor = _FakeStreamProcessor()
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
ctx.client_api_format = "openai:chat"
provider = SimpleNamespace(
name="provider",
id="provider-1",
provider_type="",
proxy={"enabled": True, "node_id": "node-1"},
)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
key = SimpleNamespace(id="key-1", proxy=None)
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
dummy_ctx = _DummyStreamResponseCtx()
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "proxy") is not None
assert getattr(plan.proxy, "mode") == "tunnel"
assert getattr(plan.proxy, "node_id") == "node-1"
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "text/event-stream"},
byte_iterator=_iter_chunks([b"data: [DONE]\n\n"]),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
stream = await handler._execute_stream_request(
ctx,
stream_processor,
provider,
endpoint,
key,
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
{},
candidate=candidate,
)
received = [chunk async for chunk in stream]
assert received == [b"data: [DONE]\n\n"]
assert dummy_ctx.closed is True
@pytest.mark.asyncio
async def test_execute_stream_request_allows_tls_profile_for_rust(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_patch_stream_setup(monkeypatch)
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
handler = _DummyChatHandler()
stream_processor = _FakeStreamProcessor()
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
ctx.client_api_format = "openai:chat"
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
key = SimpleNamespace(id="key-1", proxy=None)
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
dummy_ctx = _DummyStreamResponseCtx()
async def _fake_prepare_provider_request(self: object, **kwargs: Any) -> object:
del self, kwargs
return chatmod.ProviderRequestResult(
request_body={"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
url_model="gpt-test",
mapped_model=None,
envelope=None,
extra_headers={},
upstream_is_stream=True,
needs_conversion=False,
provider_api_format="openai:chat",
client_api_format="openai:chat",
auth_info=_DummyAuthInfo(),
tls_profile="claude_code_nodejs",
)
monkeypatch.setattr(
_DummyChatHandler,
"_prepare_provider_request",
_fake_prepare_provider_request,
)
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "tls_profile") == "claude_code_nodejs"
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "text/event-stream"},
byte_iterator=_iter_chunks([b"data: [DONE]\n\n"]),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
stream = await handler._execute_stream_request(
ctx,
stream_processor,
provider,
endpoint,
key,
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
{},
candidate=candidate,
)
received = [chunk async for chunk in stream]
assert received == [b"data: [DONE]\n\n"]
assert dummy_ctx.closed is True
@pytest.mark.asyncio
async def test_execute_stream_request_turns_rust_upstream_error_into_http_status_error(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_patch_stream_setup(monkeypatch)
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
handler = _DummyChatHandler()
stream_processor = _FakeStreamProcessor()
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
ctx.client_api_format = "openai:chat"
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
key = SimpleNamespace(id="key-1", proxy=None)
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
dummy_ctx = _DummyStreamResponseCtx()
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "stream") is True
return RustExecutorStreamResult(
status_code=429,
headers={"content-type": "application/json"},
byte_iterator=_iter_chunks([b'{"error":{"message":"slow down"}}']),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
with pytest.raises(httpx.HTTPStatusError) as exc_info:
await handler._execute_stream_request(
ctx,
stream_processor,
provider,
endpoint,
key,
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
{},
candidate=candidate,
)
assert exc_info.value.response.status_code == 429
assert "slow down" in exc_info.value.upstream_response # type: ignore[attr-defined]
assert dummy_ctx.closed is True
@pytest.mark.asyncio
async def test_execute_stream_request_falls_back_to_python_when_rust_unavailable(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_patch_stream_setup(monkeypatch)
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
handler = _DummyChatHandler()
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
ctx.client_api_format = "openai:chat"
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
key = SimpleNamespace(id="key-1", proxy=None)
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
del plan
raise RustExecutorClientError("executor down")
class _FakeHTTPClient:
def stream(self, **kwargs: Any) -> Any:
raise RuntimeError("local-http-client-used")
async def _fake_get_upstream_client(*args: Any, **kwargs: Any) -> _FakeHTTPClient:
return _FakeHTTPClient()
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
monkeypatch.setattr(
"src.clients.http_client.HTTPClientPool.get_upstream_client",
_fake_get_upstream_client,
)
with pytest.raises(RuntimeError) as exc_info:
await handler._execute_stream_request(
ctx,
object(),
provider,
endpoint,
key,
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
{},
candidate=candidate,
)
assert "local-http-client-used" in str(exc_info.value)

View File

@@ -0,0 +1,427 @@
from __future__ import annotations
import json
from collections.abc import AsyncGenerator
from types import SimpleNamespace
from typing import Any
import pytest
import src.api.handlers.base.cli_stream_mixin as cli_stream_mod
import src.api.handlers.base.cli_sync_mixin as cli_sync_mod
import src.services.proxy_node.resolver as proxymod
import src.services.task as taskmod
from src.api.handlers.base.cli_stream_mixin import CliStreamMixin
from src.api.handlers.base.cli_sync_mixin import CliSyncMixin
from src.api.handlers.base.stream_context import StreamContext
from src.services.request.rust_executor_client import (
RustExecutorStreamResult,
RustExecutorSyncResult,
)
class _DummyParser:
def extract_usage_from_response(self, response: dict[str, Any]) -> dict[str, int]:
del response
return {
"input_tokens": 0,
"output_tokens": 0,
"cache_read_tokens": 0,
"cache_creation_tokens": 0,
}
def extract_text_content(self, response: dict[str, Any]) -> str:
return str(response.get("id") or "")
class _DummyTelemetry:
async def record_success(self, **kwargs: Any) -> int:
del kwargs
return 0
async def record_failure(self, **kwargs: Any) -> None:
del kwargs
class _DummySyncHandler(CliSyncMixin):
FORMAT_ID = "openai:cli"
def __init__(self) -> None:
self.db = None
self.redis = None
self.user = SimpleNamespace(id="user-1")
self.api_key = SimpleNamespace(id="user-key-1")
self.request_id = "req-cli-sync"
self.client_ip = "127.0.0.1"
self.user_agent = "pytest"
self.start_time = 0.0
self.allowed_api_formats = ["openai:cli"]
self.primary_api_format = "openai:cli"
self.api_family = None
self.endpoint_kind = None
self.telemetry = _DummyTelemetry()
self.perf_metrics = None
self._parser = _DummyParser()
@property
def parser(self) -> _DummyParser:
return self._parser
def _create_pending_usage(self, **kwargs: object) -> bool:
del kwargs
return True
def _build_request_metadata(self, http_request: Any | None = None) -> dict[str, Any]:
del http_request
return {}
def _merge_scheduling_metadata(
self,
request_metadata: dict[str, Any] | None,
**kwargs: Any,
) -> dict[str, Any]:
del kwargs
return dict(request_metadata or {})
def _resolve_capability_requirements(
self,
model_name: str,
request_headers: dict[str, str] | None = None,
request_body: dict[str, Any] | None = None,
) -> dict[str, bool]:
del model_name, request_headers, request_body
return {}
async def _resolve_preferred_key_ids(
self,
model_name: str,
request_body: dict[str, Any] | None = None,
) -> list[str] | None:
del model_name, request_body
return None
def extract_model_from_request(
self,
request_body: dict[str, Any],
path_params: dict[str, Any] | None = None,
) -> str:
del path_params
return str(request_body.get("model") or "unknown")
async def _get_mapped_model(self, source_model: str, provider_id: str) -> str | None:
del source_model, provider_id
return None
async def _build_upstream_request(self, **kwargs: Any) -> Any:
payload = dict(kwargs["request_body"])
return SimpleNamespace(
payload=payload,
headers={"content-type": "application/json"},
url="https://upstream.test/v1/responses",
url_model=str(payload.get("model") or ""),
envelope=None,
upstream_is_stream=False,
tls_profile=None,
selected_base_url=None,
)
def _extract_response_metadata(self, response_json: dict[str, Any]) -> dict[str, Any]:
return {"id": response_json.get("id")}
class _DummyStreamResponseCtx:
def __init__(self) -> None:
self.closed = False
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
self.closed = True
class _DummyCliStreamHandler(CliStreamMixin):
FORMAT_ID = "openai:cli"
def __init__(self, *, upstream_is_stream: bool) -> None:
self.request_id = "req-cli-stream"
self.api_key = SimpleNamespace(id="user-key-1")
self._upstream_is_stream = upstream_is_stream
async def _get_mapped_model(self, source_model: str, provider_id: str) -> str | None:
del source_model, provider_id
return None
async def _build_upstream_request(self, **kwargs: Any) -> Any:
payload = dict(kwargs["request_body"])
return SimpleNamespace(
payload=payload,
headers={"content-type": "application/json"},
url="https://upstream.test/v1/responses",
url_model=str(payload.get("model") or ""),
envelope=None,
upstream_is_stream=self._upstream_is_stream,
tls_profile=None,
selected_base_url=None,
)
def apply_mapped_model(self, request_body: dict[str, Any], mapped_model: str) -> dict[str, Any]:
out = dict(request_body)
out["model"] = mapped_model
return out
def _extract_response_metadata(self, response_json: dict[str, Any]) -> dict[str, Any]:
return {"id": response_json.get("id")}
def _record_converted_chunks(self, ctx: Any, converted_events: Any) -> None:
del ctx, converted_events
def _mark_first_output(self, ctx: Any, output_state: dict[str, Any]) -> None:
del ctx
output_state["first_yield"] = False
async def _prefetch_and_check_embedded_error(
self,
byte_iterator: Any,
provider: Any,
endpoint: Any,
ctx: Any,
) -> list[bytes]:
del provider, endpoint, ctx
first = await anext(byte_iterator)
return [first]
async def _create_response_stream_with_prefetch(
self,
ctx: Any,
byte_iterator: Any,
response_ctx: _DummyStreamResponseCtx,
prefetched_chunks: list[bytes],
) -> AsyncGenerator[bytes]:
del ctx
async def _gen() -> AsyncGenerator[bytes]:
try:
for chunk in prefetched_chunks:
yield chunk
async for chunk in byte_iterator:
yield chunk
finally:
await response_ctx.__aexit__(None, None, None)
return _gen()
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
for chunk in chunks:
yield chunk
def _patch_proxy_resolver(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(proxymod, "resolve_effective_proxy", lambda provider_proxy, key_proxy=None: None)
monkeypatch.setattr(proxymod, "get_proxy_label", lambda proxy_info: "direct")
async def _fake_resolve_proxy_info(proxy_config: Any) -> Any:
del proxy_config
return None
async def _fake_resolve_delegate(proxy_config: Any) -> Any:
del proxy_config
return None
async def _fake_build_proxy_url(proxy_config: Any) -> Any:
del proxy_config
return None
monkeypatch.setattr(proxymod, "resolve_proxy_info_async", _fake_resolve_proxy_info)
monkeypatch.setattr(proxymod, "resolve_delegate_config_async", _fake_resolve_delegate)
monkeypatch.setattr(proxymod, "build_proxy_url_async", _fake_build_proxy_url)
@pytest.mark.asyncio
async def test_cli_process_sync_uses_rust_executor_when_available(
monkeypatch: pytest.MonkeyPatch,
) -> None:
handler = _DummySyncHandler()
monkeypatch.setattr(cli_sync_mod.config, "executor_backend", "rust")
_patch_proxy_resolver(monkeypatch)
class _FakeTaskService:
def __init__(self, db: Any, redis: Any) -> None:
del db, redis
async def execute(self, **kwargs: Any) -> Any:
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
provider = SimpleNamespace(
name="provider",
id="provider-1",
provider_type="",
proxy=None,
request_timeout=None,
stream_first_byte_timeout=None,
)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:cli")
key = SimpleNamespace(id="key-1", api_key="sk-test", proxy=None)
response = await kwargs["request_func"](provider, endpoint, key, candidate)
return SimpleNamespace(
response=response,
provider_name="provider",
request_candidate_id="cand-1",
provider_id="provider-1",
endpoint_id="endpoint-1",
key_id="key-1",
pool_summary=None,
)
async def _fake_execute_sync_json(self: object, plan: object) -> RustExecutorSyncResult:
assert getattr(plan, "provider_api_format") == "openai:cli"
return RustExecutorSyncResult(
status_code=200,
response_json={"id": "resp-rust-cli"},
headers={"content-type": "application/json"},
)
monkeypatch.setattr(taskmod, "TaskService", _FakeTaskService)
monkeypatch.setattr(
cli_sync_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
response = await handler.process_sync(
original_request_body={"model": "gpt-4.1", "input": "hello"},
original_headers={},
)
assert response.status_code == 200
assert json.loads(response.body) == {"id": "resp-rust-cli"}
@pytest.mark.asyncio
async def test_cli_execute_stream_request_uses_rust_sync_bridge(
monkeypatch: pytest.MonkeyPatch,
) -> None:
handler = _DummyCliStreamHandler(upstream_is_stream=False)
ctx = StreamContext(model="gpt-test", api_format="openai:cli")
ctx.client_api_format = "openai:cli"
monkeypatch.setattr(cli_stream_mod.config, "executor_backend", "rust")
_patch_proxy_resolver(monkeypatch)
async def _fake_execute_sync_json(self: object, plan: object) -> RustExecutorSyncResult:
assert getattr(plan, "stream") is False
return RustExecutorSyncResult(
status_code=200,
response_json={"id": "sync-bridge-rust"},
headers={"content-type": "application/json"},
)
async def _fake_streamify(**kwargs: Any) -> AsyncGenerator[bytes]:
assert kwargs["response_json"] == {"id": "sync-bridge-rust"}
yield b"data: cli-bridge\n\n"
monkeypatch.setattr(
cli_stream_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(handler, "_streamify_sync_response", _fake_streamify)
provider = SimpleNamespace(
name="provider",
id="provider-1",
provider_type="",
proxy=None,
request_timeout=None,
stream_first_byte_timeout=None,
)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:cli", base_url="https://x")
key = SimpleNamespace(id="key-1", proxy=None, auth_type="", api_key="sk-test")
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
stream = await handler._execute_stream_request(
ctx,
provider,
endpoint,
key,
{"model": "gpt-test", "input": "hello"},
{},
candidate=candidate,
)
chunks = [chunk async for chunk in stream]
assert chunks == [b"data: cli-bridge\n\n"]
@pytest.mark.asyncio
async def test_cli_execute_stream_request_uses_rust_native_stream(
monkeypatch: pytest.MonkeyPatch,
) -> None:
handler = _DummyCliStreamHandler(upstream_is_stream=True)
ctx = StreamContext(model="gpt-test", api_format="openai:cli")
ctx.client_api_format = "openai:cli"
monkeypatch.setattr(cli_stream_mod.config, "executor_backend", "rust")
_patch_proxy_resolver(monkeypatch)
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "stream") is True
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "text/event-stream", "x-upstream-test": "true"},
byte_iterator=_iter_chunks(
[
b"data: {\"id\":\"chunk-1\"}\n\n",
b"data: [DONE]\n\n",
]
),
response_ctx=_DummyStreamResponseCtx(),
)
monkeypatch.setattr(
cli_stream_mod.RustExecutorClient,
"execute_stream",
_fake_execute_stream,
)
provider = SimpleNamespace(
name="provider",
id="provider-1",
provider_type="",
proxy=None,
request_timeout=None,
stream_first_byte_timeout=None,
)
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:cli", base_url="https://x")
key = SimpleNamespace(id="key-1", proxy=None, api_key="sk-test")
candidate = SimpleNamespace(
request_candidate_id="cand-1",
mapping_matched_model=None,
needs_conversion=False,
output_limit=None,
)
stream = await handler._execute_stream_request(
ctx,
provider,
endpoint,
key,
{"model": "gpt-test", "input": "hello"},
{},
candidate=candidate,
)
chunks = [chunk async for chunk in stream]
assert chunks == [
b"data: {\"id\":\"chunk-1\"}\n\n",
b"data: [DONE]\n\n",
]
assert ctx.status_code == 200
assert ctx.response_headers["x-upstream-test"] == "true"

View File

@@ -0,0 +1,168 @@
from __future__ import annotations
from typing import Any
from unittest.mock import AsyncMock
import pytest
from src.api.handlers.base.endpoint_checker import EndpointCheckRequest, HttpRequestExecutor
from src.services.request.rust_executor_client import (
RustExecutorStreamResult,
RustExecutorSyncResult,
)
class _DummyStreamContext:
def __init__(self) -> None:
self.closed = False
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
self.closed = True
@pytest.mark.asyncio
async def test_endpoint_checker_sync_prefers_rust_executor(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.api.handlers.base import endpoint_checker as mod
from src.services.request import rust_executor_client as rust_mod
monkeypatch.setattr(mod.config, "executor_backend", "rust")
executor = HttpRequestExecutor(timeout=15.0)
proxy_snapshot = object()
monkeypatch.setattr(
executor,
"_build_rust_proxy_snapshot",
AsyncMock(return_value=proxy_snapshot),
)
captured: dict[str, Any] = {}
async def _fake_execute_sync_json(
self: object,
plan: Any,
) -> RustExecutorSyncResult:
captured["plan"] = plan
return RustExecutorSyncResult(
status_code=200,
response_json={"id": "resp_1", "usage": {"prompt_tokens": 1, "completion_tokens": 2}},
headers={"content-type": "application/json"},
)
monkeypatch.setattr(rust_mod.RustExecutorClient, "execute_sync_json", _fake_execute_sync_json)
result = await executor.execute(
EndpointCheckRequest(
url="https://upstream.test/v1/chat/completions",
headers={"authorization": "Bearer test"},
json_body={"model": "gpt-test", "messages": [{"role": "user", "content": "hi"}]},
api_format="openai:chat",
provider_name="openai",
model_name="gpt-test",
api_key_id="key_1",
provider_id="provider_1",
)
)
assert result.status_code == 200
assert result.response_data == {
"id": "resp_1",
"usage": {"prompt_tokens": 1, "completion_tokens": 2},
}
assert captured["plan"].proxy is proxy_snapshot
assert captured["plan"].method == "POST"
assert captured["plan"].url == "https://upstream.test/v1/chat/completions"
@pytest.mark.asyncio
async def test_endpoint_checker_stream_prefers_rust_executor(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.api.handlers.base import endpoint_checker as mod
from src.services.request import rust_executor_client as rust_mod
monkeypatch.setattr(mod.config, "executor_backend", "rust")
executor = HttpRequestExecutor(timeout=15.0)
monkeypatch.setattr(
executor,
"_build_rust_proxy_snapshot",
AsyncMock(return_value=None),
)
async def _byte_iter() -> Any:
yield b'data: {"choices":[{"delta":{"content":"Hel'
yield b'lo"}}]}\n\n'
yield b'data: {"choices":[{"delta":{"content":" world"},"finish_reason":"stop"}]}\n\n'
stream_ctx = _DummyStreamContext()
async def _fake_execute_stream(self: object, plan: Any) -> RustExecutorStreamResult:
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "text/event-stream"},
byte_iterator=_byte_iter(),
response_ctx=stream_ctx,
)
monkeypatch.setattr(rust_mod.RustExecutorClient, "execute_stream", _fake_execute_stream)
result = await executor.execute(
EndpointCheckRequest(
url="https://upstream.test/v1/chat/completions",
headers={"authorization": "Bearer test"},
json_body={
"model": "gpt-test",
"messages": [{"role": "user", "content": "hi"}],
"stream": True,
},
api_format="openai:chat",
provider_name="openai",
model_name="gpt-test",
is_stream=True,
)
)
assert result.status_code == 200
assert result.response_data == {
"choices": [{"delta": {"content": " world"}, "finish_reason": "stop"}]
}
assert stream_ctx.closed is True
@pytest.mark.asyncio
async def test_endpoint_checker_proxy_snapshot_falls_back_to_system_proxy(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.services.proxy_node import resolver as resolver_mod
executor = HttpRequestExecutor()
monkeypatch.setattr(
resolver_mod,
"get_system_proxy_config_async",
AsyncMock(return_value={"enabled": True, "url": "http://system-proxy.test:8080"}),
)
monkeypatch.setattr(
resolver_mod,
"resolve_delegate_config_async",
AsyncMock(return_value=None),
)
monkeypatch.setattr(
resolver_mod,
"build_proxy_url_async",
AsyncMock(return_value="http://system-proxy.test:8080"),
)
monkeypatch.setattr(
resolver_mod,
"resolve_proxy_info_async",
AsyncMock(return_value={"mode": "http", "label": "system-proxy"}),
)
snapshot = await executor._build_rust_proxy_snapshot(None)
assert snapshot is not None
assert snapshot.enabled is True
assert snapshot.url == "http://system-proxy.test:8080"
assert snapshot.mode == "http"