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"

View File

@@ -0,0 +1,226 @@
from __future__ import annotations
from collections.abc import AsyncGenerator
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock
import httpx
import pytest
from fastapi.responses import JSONResponse, StreamingResponse
import src.api.handlers.gemini.video_handler as video_mod
import src.services.proxy_node.resolver as resolver_mod
import src.services.request.rust_executor_client as rust_client_mod
from src.api.handlers.gemini.video_handler import GeminiVeoHandler
from src.core.api_format.conversion.internal_video import VideoStatus
from src.services.request.rust_executor_client import RustExecutorStreamResult
class _DummyStreamResponseCtx:
def __init__(self) -> None:
self.closed = False
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
self.closed = True
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
for chunk in chunks:
yield chunk
def _make_handler() -> GeminiVeoHandler:
return GeminiVeoHandler(
db=SimpleNamespace(),
user=SimpleNamespace(id="user-1"),
api_key=SimpleNamespace(id="api-key-1"),
request_id="req-gemini-video-test",
client_ip="127.0.0.1",
user_agent="pytest",
start_time=0.0,
)
@pytest.mark.asyncio
async def test_handle_create_task_uses_rust_sync_helper(
monkeypatch: pytest.MonkeyPatch,
) -> None:
handler = _make_handler()
monkeypatch.setattr(video_mod.UsageService, "create_pending_usage", lambda *args, **kwargs: None)
monkeypatch.setattr(
handler._normalizer,
"video_request_to_internal",
lambda body: SimpleNamespace(model=str(body.get("model") or "veo-3")),
)
candidate = SimpleNamespace(provider=SimpleNamespace(name="provider-1", id="prov-1"))
endpoint = SimpleNamespace(
id="ep-1",
provider_id="prov-1",
api_family="gemini",
endpoint_kind="video",
base_url="https://generativelanguage.googleapis.com",
body_rules=None,
)
provider_key = SimpleNamespace(id="key-1")
monkeypatch.setattr(
handler,
"_resolve_upstream_key",
AsyncMock(return_value=("upstream-key", endpoint, provider_key, None)),
)
monkeypatch.setattr(
handler,
"_build_upstream_url",
lambda base_url, model: f"https://generativelanguage.googleapis.com/v1beta/models/{model}:predictLongRunning",
)
monkeypatch.setattr(
handler,
"_build_upstream_headers",
lambda original_headers, upstream_key, endpoint, auth_info, **kwargs: {
"x-goog-api-key": upstream_key
},
)
monkeypatch.setattr(
video_mod.HTTPClientPool,
"get_default_client_async",
AsyncMock(side_effect=AssertionError("python fallback should not run")),
)
async def _fake_rust_sync(**kwargs: object) -> httpx.Response:
assert kwargs["method"] == "POST"
assert kwargs["provider_id"] == "prov-1"
assert kwargs["endpoint_id"] == "ep-1"
assert kwargs["key_id"] == "key-1"
assert kwargs["body"] == {"model": "veo-3", "prompt": "hello"}
return httpx.Response(
200,
request=httpx.Request("POST", str(kwargs["url"])),
json={"name": "operations/ext-1"},
)
monkeypatch.setattr(handler, "_try_rust_sync_http_response", _fake_rust_sync)
async def _fake_submit_with_failover(**kwargs: object) -> JSONResponse:
response = await kwargs["submit_func"](candidate)
assert response.status_code == 200
assert response.json()["name"] == "operations/ext-1"
return JSONResponse(status_code=400, content={"error": {"message": "stop"}})
monkeypatch.setattr(handler, "_submit_with_failover", _fake_submit_with_failover)
response = await handler.handle_create_task(
http_request=SimpleNamespace(
headers={},
url=SimpleNamespace(scheme="https", netloc="example.com"),
),
original_headers={},
original_request_body={"model": "veo-3", "prompt": "hello"},
)
assert response.status_code == 400
@pytest.mark.asyncio
async def test_handle_download_content_uses_rust_executor_with_proxy_snapshot(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(video_mod.config, "executor_backend", "rust")
handler = _make_handler()
dummy_ctx = _DummyStreamResponseCtx()
monkeypatch.setattr(
handler,
"_get_task_by_external_id",
lambda task_id: SimpleNamespace(
id=task_id,
status=VideoStatus.COMPLETED.value,
video_url="https://storage.example.com/video.mp4",
video_expires_at=datetime.now(timezone.utc).replace(year=2099),
model="veo-3",
),
)
endpoint = SimpleNamespace(
id="ep-1",
provider_id="prov-1",
api_family="gemini",
endpoint_kind="video",
base_url="https://generativelanguage.googleapis.com",
proxy={"enabled": True, "url": "http://proxy.local:8080"},
)
key = SimpleNamespace(id="key-1", api_key="encrypted", proxy=None)
monkeypatch.setattr(handler, "_get_endpoint_and_key", lambda task: (endpoint, key))
monkeypatch.setattr(video_mod.crypto_service, "decrypt", lambda _: "upstream-key")
monkeypatch.setattr(
video_mod,
"get_provider_auth",
AsyncMock(return_value=None),
)
monkeypatch.setattr(
video_mod,
"resolve_provider_proxy",
lambda endpoint, key: {"enabled": True, "url": "http://proxy.local:8080"},
)
monkeypatch.setattr(
resolver_mod,
"resolve_effective_proxy",
lambda provider_proxy, key_proxy=None: provider_proxy,
)
monkeypatch.setattr(
resolver_mod,
"get_system_proxy_config_async",
AsyncMock(return_value=None),
)
monkeypatch.setattr(
resolver_mod,
"resolve_delegate_config_async",
AsyncMock(return_value=None),
)
async def _fake_build_proxy_url_async(proxy_config: object) -> str:
assert proxy_config == {"enabled": True, "url": "http://proxy.local:8080"}
return "http://proxy.local:8080"
async def _fake_resolve_proxy_info_async(proxy_config: object) -> dict[str, str]:
assert proxy_config == {"enabled": True, "url": "http://proxy.local:8080"}
return {"url": "http://proxy.local:8080"}
monkeypatch.setattr(resolver_mod, "build_proxy_url_async", _fake_build_proxy_url_async)
monkeypatch.setattr(
resolver_mod,
"resolve_proxy_info_async",
_fake_resolve_proxy_info_async,
)
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "method") == "GET"
assert getattr(plan, "url") == "https://storage.example.com/video.mp4"
assert getattr(plan, "headers") == {"x-goog-api-key": "upstream-key"}
assert getattr(plan, "proxy").url == "http://proxy.local:8080"
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "video/mp4", "x-rust-download": "true"},
byte_iterator=_iter_chunks([b"gemini-", b"video"]),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(rust_client_mod.RustExecutorClient, "execute_stream", _fake_execute_stream)
monkeypatch.setattr(
video_mod.HTTPClientPool,
"get_default_client_async",
AsyncMock(side_effect=AssertionError("python fallback should not run")),
)
response = await handler.handle_download_content(
task_id="operations/ext-1",
http_request=SimpleNamespace(),
original_headers={},
query_params=None,
)
assert isinstance(response, StreamingResponse)
body = b"".join([chunk async for chunk in response.body_iterator])
assert body == b"gemini-video"
assert dummy_ctx.closed is True

View File

@@ -0,0 +1,198 @@
from __future__ import annotations
from collections.abc import AsyncGenerator
from types import SimpleNamespace
from typing import Any
import pytest
from fastapi.responses import Response, StreamingResponse
import src.api.handlers.openai.video_handler as video_mod
from src.api.handlers.openai.video_handler import OpenAIVideoHandler
from src.core.api_format.conversion.internal_video import VideoStatus
from src.services.request.rust_executor_client import (
RustExecutorClientError,
RustExecutorStreamResult,
)
class _DummyStreamResponseCtx:
def __init__(self) -> None:
self.closed = False
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
self.closed = True
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
for chunk in chunks:
yield chunk
def _make_handler() -> OpenAIVideoHandler:
return OpenAIVideoHandler(
db=SimpleNamespace(),
user=SimpleNamespace(id="user-1"),
api_key=SimpleNamespace(id="api-key-1"),
request_id="req-video-test",
client_ip="127.0.0.1",
user_agent="pytest",
start_time=0.0,
)
@pytest.mark.asyncio
async def test_handle_download_content_uses_rust_executor_for_direct_video_url(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(video_mod.config, "executor_backend", "rust")
handler = _make_handler()
dummy_ctx = _DummyStreamResponseCtx()
monkeypatch.setattr(
handler,
"_get_task",
lambda task_id: SimpleNamespace(
id=task_id,
status=VideoStatus.COMPLETED.value,
video_url="https://cdn.example.com/video.mp4",
model="sora-2",
),
)
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "method") == "GET"
assert getattr(plan, "url") == "https://cdn.example.com/video.mp4"
assert getattr(plan, "body").json_body is None
assert getattr(plan, "body").body_bytes_b64 is None
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "video/mp4", "x-rust-download": "true"},
byte_iterator=_iter_chunks([b"video-", b"bytes"]),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(video_mod.RustExecutorClient, "execute_stream", _fake_execute_stream)
response = await handler.handle_download_content(
task_id="task-1",
http_request=SimpleNamespace(),
original_headers={},
query_params={"variant": "video"},
)
assert isinstance(response, StreamingResponse)
assert response.headers["x-rust-download"] == "true"
body = b"".join([chunk async for chunk in response.body_iterator])
assert body == b"video-bytes"
assert dummy_ctx.closed is True
@pytest.mark.asyncio
async def test_handle_download_content_uses_rust_executor_for_upstream_content_endpoint(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(video_mod.config, "executor_backend", "rust")
handler = _make_handler()
dummy_ctx = _DummyStreamResponseCtx()
monkeypatch.setattr(
handler,
"_get_task",
lambda task_id: SimpleNamespace(
id=task_id,
status=VideoStatus.COMPLETED.value,
video_url=None,
external_task_id="ext-1",
model="sora-2",
),
)
monkeypatch.setattr(
handler,
"_get_endpoint_and_key",
lambda task: (
SimpleNamespace(id="ep-1", provider_id="prov-1", base_url="https://api.openai.com"),
SimpleNamespace(id="key-1", api_key="encrypted"),
),
)
monkeypatch.setattr(video_mod.crypto_service, "decrypt", lambda _: "upstream-key")
monkeypatch.setattr(
handler,
"_build_upstream_url",
lambda base_url, suffix=None: "https://api.openai.com/v1/videos/ext-1/content",
)
monkeypatch.setattr(
handler,
"_build_upstream_headers",
lambda original_headers, upstream_key, endpoint: {"authorization": f"Bearer {upstream_key}"},
)
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "method") == "GET"
assert getattr(plan, "url") == "https://api.openai.com/v1/videos/ext-1/content"
assert getattr(plan, "headers") == {"authorization": "Bearer upstream-key"}
assert getattr(plan, "provider_id") == "prov-1"
assert getattr(plan, "endpoint_id") == "ep-1"
assert getattr(plan, "key_id") == "key-1"
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "video/mp4", "x-rust-download": "true"},
byte_iterator=_iter_chunks([b"upstream-", b"video"]),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(video_mod.RustExecutorClient, "execute_stream", _fake_execute_stream)
response = await handler.handle_download_content(
task_id="task-1",
http_request=SimpleNamespace(),
original_headers={},
query_params=None,
)
assert isinstance(response, StreamingResponse)
body = b"".join([chunk async for chunk in response.body_iterator])
assert body == b"upstream-video"
assert dummy_ctx.closed is True
@pytest.mark.asyncio
async def test_handle_download_content_falls_back_to_python_proxy_when_rust_unavailable(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(video_mod.config, "executor_backend", "rust")
handler = _make_handler()
monkeypatch.setattr(
handler,
"_get_task",
lambda task_id: SimpleNamespace(
id=task_id,
status=VideoStatus.COMPLETED.value,
video_url="https://cdn.example.com/video.mp4",
model="sora-2",
),
)
async def _failing_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
del self, plan
raise RustExecutorClientError("executor down")
fallback_response = Response(content=b"python-fallback", media_type="video/mp4")
async def _fake_proxy_direct_url(url: str, task_id: str) -> Response:
assert url == "https://cdn.example.com/video.mp4"
assert task_id == "task-1"
return fallback_response
monkeypatch.setattr(video_mod.RustExecutorClient, "execute_stream", _failing_execute_stream)
monkeypatch.setattr(handler, "_proxy_direct_url", _fake_proxy_direct_url)
response = await handler.handle_download_content(
task_id="task-1",
http_request=SimpleNamespace(),
original_headers={},
query_params={"variant": "video"},
)
assert response is fallback_response

View File

@@ -0,0 +1,103 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock
import httpx
import pytest
from fastapi.responses import JSONResponse
import src.api.handlers.openai.video_handler as video_mod
from src.api.handlers.openai.video_handler import OpenAIVideoHandler
def _make_handler() -> OpenAIVideoHandler:
return OpenAIVideoHandler(
db=SimpleNamespace(),
user=SimpleNamespace(id="user-1"),
api_key=SimpleNamespace(id="api-key-1"),
request_id="req-video-sync-test",
client_ip="127.0.0.1",
user_agent="pytest",
start_time=0.0,
)
@pytest.mark.asyncio
async def test_handle_create_task_uses_rust_sync_helper(
monkeypatch: pytest.MonkeyPatch,
) -> None:
handler = _make_handler()
monkeypatch.setattr(video_mod.UsageService, "create_pending_usage", lambda *args, **kwargs: None)
monkeypatch.setattr(
handler._normalizer,
"video_request_to_internal",
lambda body: SimpleNamespace(model=str(body.get("model") or "sora-2")),
)
candidate = SimpleNamespace(provider=SimpleNamespace(name="provider-1", id="prov-1"))
endpoint = SimpleNamespace(
id="ep-1",
api_family="openai",
endpoint_kind="video",
base_url="https://api.openai.com",
body_rules=None,
)
provider_key = SimpleNamespace(id="key-1")
monkeypatch.setattr(
handler,
"_resolve_upstream_key",
AsyncMock(return_value=("upstream-key", endpoint, provider_key)),
)
monkeypatch.setattr(
handler,
"_build_upstream_url",
lambda base_url: "https://api.openai.com/v1/videos",
)
monkeypatch.setattr(
handler,
"_build_upstream_headers",
lambda original_headers, upstream_key, endpoint, **kwargs: {
"authorization": f"Bearer {upstream_key}"
},
)
monkeypatch.setattr(
video_mod.HTTPClientPool,
"get_default_client_async",
AsyncMock(side_effect=AssertionError("python fallback should not run")),
)
async def _fake_rust_sync(**kwargs: object) -> httpx.Response:
assert kwargs["method"] == "POST"
assert kwargs["url"] == "https://api.openai.com/v1/videos"
assert kwargs["provider_id"] == "prov-1"
assert kwargs["endpoint_id"] == "ep-1"
assert kwargs["key_id"] == "key-1"
assert kwargs["body"] == {"model": "sora-2", "prompt": "hello"}
return httpx.Response(
200,
request=httpx.Request("POST", str(kwargs["url"])),
json={"id": "ext-1"},
)
create_failed = AsyncMock()
monkeypatch.setattr(handler, "_try_rust_sync_http_response", _fake_rust_sync)
monkeypatch.setattr(handler, "_create_failed_task_and_usage", create_failed)
async def _fake_submit_with_failover(**kwargs: object) -> JSONResponse:
response = await kwargs["submit_func"](candidate)
assert response.status_code == 200
assert response.json()["id"] == "ext-1"
return JSONResponse(status_code=400, content={"error": {"message": "stop"}})
monkeypatch.setattr(handler, "_submit_with_failover", _fake_submit_with_failover)
response = await handler.handle_create_task(
http_request=SimpleNamespace(),
original_headers={},
original_request_body={"model": "sora-2", "prompt": "hello"},
)
assert response.status_code == 400
create_failed.assert_awaited_once()

View File

@@ -0,0 +1,296 @@
from __future__ import annotations
import json
from collections.abc import AsyncGenerator
from types import SimpleNamespace
from unittest.mock import AsyncMock
import httpx
import pytest
from fastapi.responses import StreamingResponse
import src.api.public.gemini_files as gemini_files_mod
import src.services.proxy_node.resolver as resolver_mod
import src.services.request.rust_executor_client as rust_client_mod
from src.api.public.gemini_files import UpstreamContext
from src.config.settings import config
from src.services.request.executor_plan import ExecutionProxySnapshot
from src.services.request.rust_executor_client import (
RustExecutorStreamResult,
RustExecutorSyncResult,
)
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 _FakeDBContext:
def __init__(self, db: object) -> None:
self._db = db
def __enter__(self) -> object:
return self._db
def __exit__(self, exc_type: object, exc: object, tb: object) -> bool:
return False
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
for chunk in chunks:
yield chunk
@pytest.mark.asyncio
async def test_enrich_upstream_context_proxy_builds_tunnel_snapshot(
monkeypatch: pytest.MonkeyPatch,
) -> None:
ctx = UpstreamContext(
upstream_key="upstream-key",
base_url="https://generativelanguage.googleapis.com",
file_key_id="key-1",
user_id="user-1",
provider_id="prov-1",
endpoint_id="ep-1",
provider_proxy={"enabled": True, "node_id": "node-1"},
key_proxy=None,
)
monkeypatch.setattr(
resolver_mod,
"resolve_effective_proxy",
lambda provider_proxy, key_proxy=None: provider_proxy,
)
monkeypatch.setattr(
resolver_mod,
"get_system_proxy_config_async",
AsyncMock(return_value=None),
)
monkeypatch.setattr(
resolver_mod,
"resolve_delegate_config_async",
AsyncMock(return_value={"tunnel": True, "node_id": "node-1"}),
)
monkeypatch.setattr(
resolver_mod,
"resolve_proxy_info_async",
AsyncMock(return_value={"mode": "tunnel", "node_id": "node-1", "label": "Node 1"}),
)
async def _unexpected_build_proxy_url_async(proxy_config: object) -> str:
raise AssertionError(f"proxy url should not be built for tunnel: {proxy_config!r}")
monkeypatch.setattr(
resolver_mod,
"build_proxy_url_async",
_unexpected_build_proxy_url_async,
)
enriched = await gemini_files_mod._enrich_upstream_context_proxy(ctx)
assert enriched.proxy_config == {"enabled": True, "node_id": "node-1"}
assert enriched.delegate_config == {"tunnel": True, "node_id": "node-1"}
assert enriched.proxy_snapshot is not None
assert enriched.proxy_snapshot.mode == "tunnel"
assert enriched.proxy_snapshot.node_id == "node-1"
assert enriched.proxy_snapshot.url is None
@pytest.mark.asyncio
async def test_proxy_request_passes_proxy_snapshot_to_rust_executor(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(config, "executor_backend", "rust")
proxy_snapshot = ExecutionProxySnapshot(
enabled=True,
mode="http",
url="http://proxy.local:8080",
)
async def _fake_execute_sync_json(self: object, plan: object) -> RustExecutorSyncResult:
assert getattr(plan, "method") == "GET"
assert getattr(plan, "provider_id") == "prov-1"
assert getattr(plan, "endpoint_id") == "ep-1"
assert getattr(plan, "proxy").url == "http://proxy.local:8080"
return RustExecutorSyncResult(
status_code=200,
headers={"content-type": "application/json", "x-rust-files": "true"},
response_json={"files": [{"name": "files/abc"}]},
)
monkeypatch.setattr(
rust_client_mod.RustExecutorClient,
"execute_sync_json",
_fake_execute_sync_json,
)
monkeypatch.setattr(
gemini_files_mod.HTTPClientPool,
"get_upstream_client",
AsyncMock(side_effect=AssertionError("python fallback should not run")),
)
response = await gemini_files_mod._proxy_request(
"GET",
"https://generativelanguage.googleapis.com/v1beta/files",
{"x-goog-api-key": "upstream-key"},
file_key_id="key-1",
user_id="user-1",
provider_id="prov-1",
endpoint_id="ep-1",
proxy=proxy_snapshot,
)
assert response.status_code == 200
assert response.headers["x-rust-files"] == "true"
assert json.loads(response.body) == {"files": [{"name": "files/abc"}]}
@pytest.mark.asyncio
async def test_proxy_request_fallback_uses_upstream_client_proxy_context(
monkeypatch: pytest.MonkeyPatch,
) -> None:
class _FakeClient:
async def get(self, url: str, headers: dict[str, str]) -> httpx.Response:
assert url == "https://generativelanguage.googleapis.com/v1beta/files"
assert headers == {"x-goog-api-key": "upstream-key"}
return httpx.Response(
200,
request=httpx.Request("GET", url),
json={"files": []},
)
async def _fake_try_rust_sync_proxy_request(*args: object, **kwargs: object) -> None:
del args, kwargs
return None
async def _fake_get_upstream_client(
delegate_cfg: dict[str, object] | None,
*,
proxy_config: dict[str, object] | None = None,
tls_profile: str | None = None,
) -> _FakeClient:
assert delegate_cfg == {"tunnel": True, "node_id": "node-1"}
assert proxy_config == {"enabled": True, "node_id": "node-1"}
assert tls_profile is None
return _FakeClient()
monkeypatch.setattr(
gemini_files_mod,
"_try_rust_sync_proxy_request",
_fake_try_rust_sync_proxy_request,
)
monkeypatch.setattr(
gemini_files_mod.HTTPClientPool,
"get_upstream_client",
_fake_get_upstream_client,
)
response = await gemini_files_mod._proxy_request(
"GET",
"https://generativelanguage.googleapis.com/v1beta/files",
{"x-goog-api-key": "upstream-key"},
file_key_id="key-1",
user_id="user-1",
proxy_config={"enabled": True, "node_id": "node-1"},
delegate_config={"tunnel": True, "node_id": "node-1"},
)
assert response.status_code == 200
assert json.loads(response.body) == {"files": []}
@pytest.mark.asyncio
async def test_download_file_uses_enriched_proxy_snapshot_for_regular_files(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(config, "executor_backend", "rust")
dummy_ctx = _DummyStreamResponseCtx()
raw_ctx = UpstreamContext(
upstream_key="upstream-key",
base_url="https://generativelanguage.googleapis.com",
file_key_id="key-1",
user_id="user-1",
provider_id="prov-1",
endpoint_id="ep-1",
)
enriched_ctx = UpstreamContext(
upstream_key="upstream-key",
base_url="https://generativelanguage.googleapis.com",
file_key_id="key-1",
user_id="user-1",
provider_id="prov-1",
endpoint_id="ep-1",
proxy_config={"enabled": True, "url": "http://proxy.local:8080"},
delegate_config=None,
proxy_snapshot=ExecutionProxySnapshot(
enabled=True,
mode="http",
url="http://proxy.local:8080",
),
)
monkeypatch.setattr(gemini_files_mod, "_extract_gemini_api_key", lambda request: "client-key")
monkeypatch.setattr(
gemini_files_mod,
"create_session",
lambda: _FakeDBContext(SimpleNamespace()),
)
monkeypatch.setattr(
gemini_files_mod.AuthService,
"authenticate_api_key",
lambda db, key: (SimpleNamespace(id="user-1"), SimpleNamespace(id="user-api-key")),
)
monkeypatch.setattr(gemini_files_mod, "_ensure_balance_access", lambda db, user, api_key: None)
monkeypatch.setattr(
gemini_files_mod,
"_resolve_upstream_context",
AsyncMock(return_value=raw_ctx),
)
monkeypatch.setattr(
gemini_files_mod,
"_enrich_upstream_context_proxy",
AsyncMock(return_value=enriched_ctx),
)
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
assert getattr(plan, "method") == "GET"
assert getattr(plan, "url") == (
"https://generativelanguage.googleapis.com/v1beta/files/file-1:download?alt=media"
)
assert getattr(plan, "headers") == {"x-goog-api-key": "upstream-key"}
assert getattr(plan, "proxy").url == "http://proxy.local:8080"
return RustExecutorStreamResult(
status_code=200,
headers={"content-type": "application/octet-stream", "x-rust-files": "true"},
byte_iterator=_iter_chunks([b"file-", b"bytes"]),
response_ctx=dummy_ctx,
)
monkeypatch.setattr(
rust_client_mod.RustExecutorClient,
"execute_stream",
_fake_execute_stream,
)
monkeypatch.setattr(
gemini_files_mod.HTTPClientPool,
"get_upstream_client",
AsyncMock(side_effect=AssertionError("python fallback should not run")),
)
response = await gemini_files_mod.download_file(
"file-1",
SimpleNamespace(
headers={},
query_params={"alt": "media"},
),
)
assert isinstance(response, StreamingResponse)
assert response.headers["x-rust-files"] == "true"
body = b"".join([chunk async for chunk in response.body_iterator])
assert body == b"file-bytes"
assert dummy_ctx.closed is True

View File

@@ -0,0 +1,142 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
def _build_provider_fixture() -> SimpleNamespace:
endpoint = SimpleNamespace(
id="endpoint_1",
base_url="https://upstream.test",
api_format="openai:chat",
is_active=True,
)
key = SimpleNamespace(
id="key_1",
is_active=True,
api_formats=None,
)
return SimpleNamespace(
id="provider_1",
name="openai",
endpoints=[endpoint],
api_keys=[key],
)
@pytest.mark.asyncio
async def test_test_connection_prefers_rust_executor(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.api.public import system_catalog as mod
provider = _build_provider_fixture()
monkeypatch.setattr(mod, "_select_provider", lambda _db, _provider_name: provider)
monkeypatch.setattr(mod, "build_test_request_body", lambda *_args, **_kwargs: {"model": "gpt-test"})
monkeypatch.setattr(mod, "get_provider_auth", AsyncMock(return_value=None))
class _DummyBuilder:
def build(self, *_args: object, **_kwargs: object) -> tuple[dict[str, str], dict[str, str]]:
return {"model": "gpt-test"}, {"authorization": "Bearer test"}
monkeypatch.setattr(mod, "PassthroughRequestBuilder", lambda: _DummyBuilder())
monkeypatch.setattr(
mod,
"build_provider_url",
lambda *_args, **_kwargs: "https://upstream.test/v1/chat/completions",
)
proxy_snapshot = object()
monkeypatch.setattr(
mod,
"_build_test_connection_transport_context",
AsyncMock(return_value=({"enabled": True}, {"node_id": "node-1", "tunnel": True}, proxy_snapshot)),
)
rust_response = httpx.Response(
200,
request=httpx.Request("POST", "https://upstream.test/v1/chat/completions"),
json={"id": "resp_rust"},
)
rust_call = AsyncMock(return_value=rust_response)
monkeypatch.setattr(mod, "_try_rust_test_connection_response", rust_call)
get_upstream_client = AsyncMock(side_effect=AssertionError("python upstream client should not be used"))
monkeypatch.setattr(mod.HTTPClientPool, "get_upstream_client", get_upstream_client)
result = await mod.test_connection(
request=SimpleNamespace(query_params={}),
db=MagicMock(),
provider=None,
model="gpt-test",
api_format=None,
)
assert result["status"] == "success"
assert result["response_id"] == "resp_rust"
rust_call.assert_awaited_once()
get_upstream_client.assert_not_awaited()
@pytest.mark.asyncio
async def test_test_connection_fallback_uses_transport_aware_client(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.api.public import system_catalog as mod
provider = _build_provider_fixture()
monkeypatch.setattr(mod, "_select_provider", lambda _db, _provider_name: provider)
monkeypatch.setattr(mod, "build_test_request_body", lambda *_args, **_kwargs: {"model": "gpt-test"})
monkeypatch.setattr(mod, "get_provider_auth", AsyncMock(return_value=None))
class _DummyBuilder:
def build(self, *_args: object, **_kwargs: object) -> tuple[dict[str, str], dict[str, str]]:
return {"model": "gpt-test"}, {"authorization": "Bearer test"}
monkeypatch.setattr(mod, "PassthroughRequestBuilder", lambda: _DummyBuilder())
monkeypatch.setattr(
mod,
"build_provider_url",
lambda *_args, **_kwargs: "https://upstream.test/v1/chat/completions",
)
proxy_config = {"enabled": True, "url": "http://proxy.test:8080"}
delegate_cfg = {"node_id": "node-1", "tunnel": True}
monkeypatch.setattr(
mod,
"_build_test_connection_transport_context",
AsyncMock(return_value=(proxy_config, delegate_cfg, None)),
)
monkeypatch.setattr(mod, "_try_rust_test_connection_response", AsyncMock(return_value=None))
upstream_response = httpx.Response(
200,
request=httpx.Request("POST", "https://upstream.test/v1/chat/completions"),
json={"id": "resp_python"},
)
upstream_client = MagicMock()
upstream_client.post = AsyncMock(return_value=upstream_response)
get_upstream_client = AsyncMock(return_value=upstream_client)
monkeypatch.setattr(mod.HTTPClientPool, "get_upstream_client", get_upstream_client)
result = await mod.test_connection(
request=SimpleNamespace(query_params={}),
db=MagicMock(),
provider=None,
model="gpt-test",
api_format=None,
)
assert result["status"] == "success"
assert result["response_id"] == "resp_python"
get_upstream_client.assert_awaited_once_with(delegate_cfg, proxy_config=proxy_config)
upstream_client.post.assert_awaited_once_with(
"https://upstream.test/v1/chat/completions",
json={"model": "gpt-test"},
headers={"authorization": "Bearer test"},
)

View File

@@ -0,0 +1,633 @@
import base64
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from fastapi.responses import JSONResponse, StreamingResponse
from src.api.internal.gateway import (
CONTROL_ACTION_HEADER,
CONTROL_ACTION_PROXY_PUBLIC,
CONTROL_EXECUTED_HEADER,
GatewayResolveRequest,
_is_streaming_sync_payload,
_resolve_auth_context,
_resolve_gateway_sync_adapter,
classify_gateway_route,
router,
)
from src.database import get_db
def test_classify_openai_chat_route_as_ai_public() -> None:
decision = classify_gateway_route("POST", "/v1/chat/completions")
assert decision.route_class == "ai_public"
assert decision.route_family == "openai"
assert decision.route_kind == "chat"
assert decision.auth_endpoint_signature == "openai:chat"
assert decision.executor_candidate is True
assert decision.action == "proxy_public"
def test_classify_gemini_files_download_route_as_ai_public() -> None:
decision = classify_gateway_route("GET", "/v1beta/files/file-123:download")
assert decision.route_class == "ai_public"
assert decision.route_family == "gemini"
assert decision.route_kind == "files"
assert decision.auth_endpoint_signature == "gemini:chat"
assert decision.executor_candidate is True
def test_classify_gemini_files_nested_metadata_route_as_ai_public() -> None:
decision = classify_gateway_route("GET", "/v1beta/files/files/abc-123")
assert decision.route_class == "ai_public"
assert decision.route_family == "gemini"
assert decision.route_kind == "files"
assert decision.executor_candidate is True
def test_classify_gemini_video_operation_route_as_ai_public() -> None:
decision = classify_gateway_route("GET", "/v1beta/models/veo-3/operations/op-123")
assert decision.route_class == "ai_public"
assert decision.route_family == "gemini"
assert decision.route_kind == "video"
assert decision.auth_endpoint_signature == "gemini:video"
assert decision.executor_candidate is True
def test_classify_non_ai_route_as_passthrough() -> None:
decision = classify_gateway_route("GET", "/api/admin/system/info")
assert decision.route_class == "passthrough"
assert decision.route_family is None
assert decision.route_kind is None
assert decision.executor_candidate is False
def test_classify_claude_cli_route_from_bearer_header() -> None:
decision = classify_gateway_route(
"POST",
"/v1/messages",
{"authorization": "Bearer sk-cli"},
)
assert decision.route_class == "ai_public"
assert decision.route_family == "claude"
assert decision.route_kind == "cli"
assert decision.auth_endpoint_signature == "claude:cli"
def test_classify_gemini_cli_route_from_user_agent() -> None:
decision = classify_gateway_route(
"POST",
"/v1beta/models/gemini-2.5-pro:generateContent",
{"user-agent": "GeminiCLI/1.2.3"},
)
assert decision.route_class == "ai_public"
assert decision.route_family == "gemini"
assert decision.route_kind == "cli"
assert decision.auth_endpoint_signature == "gemini:cli"
def test_resolve_sync_adapter_for_openai_chat_route() -> None:
decision = classify_gateway_route("POST", "/v1/chat/completions")
adapter, path_params = _resolve_gateway_sync_adapter(decision, "/v1/chat/completions")
assert adapter is not None
assert adapter.name == "openai.chat"
assert path_params == {}
def test_resolve_sync_adapter_for_gemini_route_extracts_model_path_params() -> None:
decision = classify_gateway_route("POST", "/v1beta/models/gemini-2.5-pro:generateContent")
adapter, path_params = _resolve_gateway_sync_adapter(
decision,
"/v1beta/models/gemini-2.5-pro:generateContent",
)
assert adapter is not None
assert adapter.name == "gemini.chat"
assert path_params == {"model": "gemini-2.5-pro", "stream": False}
def test_resolve_sync_adapter_rejects_non_sync_files_route() -> None:
decision = classify_gateway_route("GET", "/v1beta/files/file-123:download")
adapter, path_params = _resolve_gateway_sync_adapter(decision, "/v1beta/files/file-123:download")
assert adapter is None
assert path_params == {}
def test_is_streaming_sync_payload_detects_body_and_path_stream_flags() -> None:
assert _is_streaming_sync_payload({"stream": True}, {}) is True
assert _is_streaming_sync_payload({}, {"stream": True}) is True
assert _is_streaming_sync_payload({"stream": False}, {"stream": False}) is False
@pytest.mark.asyncio
async def test_resolve_auth_context_from_openai_bearer_header(
monkeypatch: pytest.MonkeyPatch,
) -> None:
payload = GatewayResolveRequest(
method="POST",
path="/v1/chat/completions",
headers={"Authorization": "Bearer sk-test"},
)
decision = classify_gateway_route(payload.method, payload.path, payload.headers)
monkeypatch.setattr(
"src.api.internal.gateway.AuthService.authenticate_api_key_threadsafe",
AsyncMock(
return_value=SimpleNamespace(
user=SimpleNamespace(id="user-123"),
api_key=SimpleNamespace(id="key-123"),
balance_remaining=42.5,
access_allowed=True,
)
),
)
auth_context = await _resolve_auth_context(payload, decision)
assert auth_context == {
"user_id": "user-123",
"api_key_id": "key-123",
"balance_remaining": 42.5,
"access_allowed": True,
}
def test_execute_sync_route_returns_controlled_response(monkeypatch: pytest.MonkeyPatch) -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
class FakeAdapter:
mode = SimpleNamespace(value="standard")
allowed_api_formats = ["openai:chat"]
def authorize(self, context: object) -> None:
self.authorized_context = context
async def handle(self, context: object) -> JSONResponse:
assert getattr(context, "path_params", {}) == {}
return JSONResponse(status_code=201, content={"ok": True, "request_id": context.request_id})
fake_adapter = FakeAdapter()
fake_pipeline = SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None))
monkeypatch.setattr(
"src.api.internal.gateway._resolve_gateway_sync_adapter",
lambda decision, path: (fake_adapter, {}),
)
monkeypatch.setattr(
"src.api.internal.gateway._load_gateway_auth_models",
lambda db, auth_context: (
SimpleNamespace(id="user-123"),
SimpleNamespace(id="key-123"),
),
)
monkeypatch.setattr("src.api.internal.gateway.get_pipeline", lambda: fake_pipeline)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/execute-sync",
json={
"trace_id": "trace-sync-123",
"method": "POST",
"path": "/v1/chat/completions",
"headers": {"user-agent": "pytest"},
"body_json": {"model": "gpt-5", "messages": []},
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"balance_remaining": 12.5,
"access_allowed": True,
},
},
)
assert response.status_code == 201
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
assert response.json()["ok"] is True
assert response.json()["request_id"] == "trace-sync-123"
def test_execute_sync_route_falls_back_for_stream_payload() -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch = pytest.MonkeyPatch()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/execute-sync",
json={
"trace_id": "trace-stream-123",
"method": "POST",
"path": "/v1/chat/completions",
"headers": {"user-agent": "pytest"},
"body_json": {"model": "gpt-5", "messages": [], "stream": True},
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"access_allowed": True,
},
},
)
assert response.status_code == 409
assert response.headers[CONTROL_ACTION_HEADER] == CONTROL_ACTION_PROXY_PUBLIC
assert response.json() == {"action": CONTROL_ACTION_PROXY_PUBLIC}
monkeypatch.undo()
def test_execute_stream_route_returns_controlled_stream(monkeypatch: pytest.MonkeyPatch) -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
class FakeAdapter:
mode = SimpleNamespace(value="standard")
allowed_api_formats = ["openai:chat"]
def authorize(self, context: object) -> None:
self.authorized_context = context
async def handle(self, context: object) -> StreamingResponse:
async def _iter() -> object:
yield b"data: one\n\n"
yield b"data: [DONE]\n\n"
return StreamingResponse(_iter(), media_type="text/event-stream")
fake_pipeline = SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None))
monkeypatch.setattr(
"src.api.internal.gateway._resolve_gateway_sync_adapter",
lambda decision, path: (FakeAdapter(), {}),
)
monkeypatch.setattr(
"src.api.internal.gateway._load_gateway_auth_models",
lambda db, auth_context: (
SimpleNamespace(id="user-123"),
SimpleNamespace(id="key-123"),
),
)
monkeypatch.setattr("src.api.internal.gateway.get_pipeline", lambda: fake_pipeline)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/execute-stream",
json={
"trace_id": "trace-stream-123",
"method": "POST",
"path": "/v1/chat/completions",
"headers": {"user-agent": "pytest"},
"body_json": {"model": "gpt-5", "messages": [], "stream": True},
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"balance_remaining": 12.5,
"access_allowed": True,
},
},
)
assert response.status_code == 200
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
assert response.text == "data: one\n\ndata: [DONE]\n\n"
def test_execute_stream_route_falls_back_for_sync_payload(monkeypatch: pytest.MonkeyPatch) -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/execute-stream",
json={
"trace_id": "trace-sync-456",
"method": "POST",
"path": "/v1/chat/completions",
"headers": {"user-agent": "pytest"},
"body_json": {"model": "gpt-5", "messages": []},
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"access_allowed": True,
},
},
)
assert response.status_code == 409
assert response.headers[CONTROL_ACTION_HEADER] == CONTROL_ACTION_PROXY_PUBLIC
assert response.json() == {"action": CONTROL_ACTION_PROXY_PUBLIC}
def test_execute_sync_route_handles_gemini_files_list(monkeypatch: pytest.MonkeyPatch) -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
monkeypatch.setattr(
"src.api.internal.gateway._load_gateway_auth_models",
lambda db, auth_context: (
SimpleNamespace(id="user-123"),
SimpleNamespace(id="key-123"),
),
)
monkeypatch.setattr(
"src.api.internal.gateway.get_pipeline",
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
)
async def _fake_list_files(
request: object,
pageSize: int | None = None,
pageToken: str | None = None,
) -> JSONResponse:
del request
assert pageSize == 25
assert pageToken == "page-2"
return JSONResponse(status_code=200, content={"files": [{"name": "files/abc"}]})
monkeypatch.setattr("src.api.public.gemini_files.list_files", _fake_list_files)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/execute-sync",
json={
"trace_id": "trace-files-list",
"method": "GET",
"path": "/v1beta/files",
"query_string": "pageSize=25&pageToken=page-2",
"headers": {"x-goog-api-key": "client-key"},
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"access_allowed": True,
},
},
)
assert response.status_code == 200
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
assert response.json() == {"files": [{"name": "files/abc"}]}
def test_execute_sync_route_handles_gemini_files_upload_raw_body(
monkeypatch: pytest.MonkeyPatch,
) -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
monkeypatch.setattr(
"src.api.internal.gateway._load_gateway_auth_models",
lambda db, auth_context: (
SimpleNamespace(id="user-123"),
SimpleNamespace(id="key-123"),
),
)
monkeypatch.setattr(
"src.api.internal.gateway.get_pipeline",
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
)
async def _fake_upload_file(request: object) -> JSONResponse:
body = await request.body()
assert body == b"upload-bytes"
assert request.headers["content-type"] == "application/octet-stream"
return JSONResponse(status_code=201, content={"uploaded": True})
monkeypatch.setattr("src.api.public.gemini_files.upload_file", _fake_upload_file)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/execute-sync",
json={
"trace_id": "trace-files-upload",
"method": "POST",
"path": "/upload/v1beta/files",
"query_string": "uploadType=resumable",
"headers": {
"x-goog-api-key": "client-key",
"content-type": "application/octet-stream",
},
"body_base64": base64.b64encode(b"upload-bytes").decode("ascii"),
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"access_allowed": True,
},
},
)
assert response.status_code == 201
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
assert response.json() == {"uploaded": True}
def test_execute_sync_route_handles_openai_video_remix_with_original_request(
monkeypatch: pytest.MonkeyPatch,
) -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
class FakeVideoAdapter:
mode = SimpleNamespace(value="standard")
allowed_api_formats = ["openai:video"]
def authorize(self, context: object) -> None:
self.authorized_context = context
async def handle(self, context: object) -> JSONResponse:
assert context.request.method == "POST"
assert context.request.url.path == "/v1/videos/task-123/remix"
assert context.path_params == {"task_id": "task-123"}
assert await context.ensure_json_body_async() == {"prompt": "remix this"}
return JSONResponse(status_code=200, content={"video": True})
monkeypatch.setattr(
"src.api.internal.gateway._resolve_gateway_sync_adapter",
lambda decision, path: (FakeVideoAdapter(), {"task_id": "task-123"}),
)
monkeypatch.setattr(
"src.api.internal.gateway._load_gateway_auth_models",
lambda db, auth_context: (
SimpleNamespace(id="user-123"),
SimpleNamespace(id="key-123"),
),
)
monkeypatch.setattr(
"src.api.internal.gateway.get_pipeline",
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/execute-sync",
json={
"trace_id": "trace-video-remix",
"method": "POST",
"path": "/v1/videos/task-123/remix",
"headers": {
"content-type": "application/json",
"user-agent": "pytest",
},
"body_json": {"prompt": "remix this"},
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"access_allowed": True,
},
},
)
assert response.status_code == 200
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
assert response.json() == {"video": True}
def test_plan_stream_route_returns_executor_plan_for_gemini_files_download(
monkeypatch: pytest.MonkeyPatch,
) -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
fake_plan = {
"request_id": "plan-123",
"provider_id": "provider-123",
"endpoint_id": "endpoint-123",
"key_id": "key-123",
"method": "GET",
"url": "https://example.com/v1beta/files/file-123:download",
"headers": {"x-goog-api-key": "upstream-key"},
"body": {},
"stream": True,
"provider_api_format": "gemini:files",
"client_api_format": "gemini:files",
"model_name": "gemini-files",
}
monkeypatch.setattr(
"src.api.internal.gateway._load_gateway_auth_models",
lambda db, auth_context: (
SimpleNamespace(id="user-123"),
SimpleNamespace(id="key-123"),
),
)
monkeypatch.setattr(
"src.api.internal.gateway.get_pipeline",
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
)
monkeypatch.setattr(
"src.api.internal.gateway._build_gemini_files_download_stream_plan",
AsyncMock(return_value=fake_plan),
)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/plan-stream",
json={
"trace_id": "trace-files-plan",
"method": "GET",
"path": "/v1beta/files/file-123:download",
"query_string": "alt=media",
"headers": {"x-goog-api-key": "client-key"},
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"access_allowed": True,
},
},
)
assert response.status_code == 200
assert response.json() == {
"action": "executor_stream",
"plan_kind": "gemini_files_download",
"plan": fake_plan,
}
def test_plan_stream_route_returns_executor_plan_for_openai_video_content(
monkeypatch: pytest.MonkeyPatch,
) -> None:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[get_db] = lambda: object()
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
fake_plan = {
"request_id": "plan-video-123",
"provider_id": "provider-video-123",
"endpoint_id": "endpoint-video-123",
"key_id": "key-video-123",
"method": "GET",
"url": "https://api.openai.com/v1/videos/ext-123/content",
"headers": {"authorization": "Bearer upstream-key"},
"body": {},
"stream": True,
"provider_api_format": "openai:video",
"client_api_format": "openai:video",
"model_name": "sora-2",
}
monkeypatch.setattr(
"src.api.internal.gateway._load_gateway_auth_models",
lambda db, auth_context: (
SimpleNamespace(id="user-123"),
SimpleNamespace(id="key-123"),
),
)
monkeypatch.setattr(
"src.api.internal.gateway.get_pipeline",
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
)
monkeypatch.setattr(
"src.api.internal.gateway._build_openai_video_content_stream_plan",
AsyncMock(return_value=fake_plan),
)
client = TestClient(app, base_url="http://127.0.0.1")
response = client.post(
"/api/internal/gateway/plan-stream",
json={
"trace_id": "trace-video-plan",
"method": "GET",
"path": "/v1/videos/task-123/content",
"query_string": "variant=video",
"headers": {"authorization": "Bearer client-key"},
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"access_allowed": True,
},
},
)
assert response.status_code == 200
assert response.json() == {
"action": "executor_stream",
"plan_kind": "openai_video_content",
"plan": fake_plan,
}

View File

@@ -529,6 +529,101 @@ class TestPipelineAuthentication:
assert exc_info.value.status_code == 403
assert "锁定" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_authenticate_client_reuses_trusted_gateway_auth_context(
self, pipeline: ApiRequestPipeline
) -> None:
mock_request = MagicMock()
mock_request.headers = {
"x-aether-gateway": "rust-phase3b",
"x-aether-auth-user-id": "user-123",
"x-aether-auth-api-key-id": "key-123",
"x-aether-auth-balance-remaining": "42.5",
"x-aether-auth-access-allowed": "true",
}
mock_request.client = SimpleNamespace(host="127.0.0.1")
mock_request.url.path = "/v1/chat/completions"
mock_request.state = MagicMock()
db_user = MagicMock()
db_user.id = "user-123"
db_user.is_active = True
db_user.is_deleted = False
db_api_key = MagicMock()
db_api_key.id = "key-123"
db_api_key.user_id = "user-123"
db_api_key.is_active = True
db_api_key.is_locked = False
db_api_key.is_standalone = False
db_api_key.expires_at = None
mock_db = MagicMock()
user_query = MagicMock()
user_query.filter.return_value.first.return_value = db_user
api_key_query = MagicMock()
api_key_query.filter.return_value.first.return_value = db_api_key
mock_db.query.side_effect = [user_query, api_key_query]
mock_adapter = MagicMock()
mock_adapter.extract_api_key.side_effect = AssertionError("trusted auth should short-circuit")
with patch.object(
pipeline.auth_service,
"authenticate_api_key_threadsafe",
new_callable=AsyncMock,
side_effect=AssertionError("threadsafe auth should not be called"),
):
user, api_key = await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
assert user == db_user
assert api_key == db_api_key
assert mock_request.state.user_id == "user-123"
assert mock_request.state.api_key_id == "key-123"
assert mock_request.state.prefetched_balance_remaining == 42.5
@pytest.mark.asyncio
async def test_authenticate_client_trusted_gateway_balance_denied(
self, pipeline: ApiRequestPipeline
) -> None:
mock_request = MagicMock()
mock_request.headers = {
"x-aether-gateway": "rust-phase3b",
"x-aether-auth-user-id": "user-123",
"x-aether-auth-api-key-id": "key-123",
"x-aether-auth-balance-remaining": "0",
"x-aether-auth-access-allowed": "false",
}
mock_request.client = SimpleNamespace(host="127.0.0.1")
mock_request.url.path = "/v1/chat/completions"
mock_request.state = MagicMock()
db_user = MagicMock()
db_user.id = "user-123"
db_user.is_active = True
db_user.is_deleted = False
db_api_key = MagicMock()
db_api_key.id = "key-123"
db_api_key.user_id = "user-123"
db_api_key.is_active = True
db_api_key.is_locked = False
db_api_key.is_standalone = False
db_api_key.expires_at = None
mock_db = MagicMock()
user_query = MagicMock()
user_query.filter.return_value.first.return_value = db_user
api_key_query = MagicMock()
api_key_query.filter.return_value.first.return_value = db_api_key
mock_db.query.side_effect = [user_query, api_key_query]
mock_adapter = MagicMock()
mock_adapter.extract_api_key.side_effect = AssertionError("trusted auth should short-circuit")
from src.core.exceptions import BalanceInsufficientException
with pytest.raises(BalanceInsufficientException):
await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
class TestPipelineUserRateLimit:
@pytest.fixture

899
tests/e2e_rust_executor.py Normal file
View File

@@ -0,0 +1,899 @@
"""
aether-executor 本地端到端测试
测试流程:
1. 启动本地假上游 HTTP 服务
2. 启动 aether-executorUnix Socket
3. 用 Python RustExecutorClient 发送 ExecutionPlan
4. 验证执行结果与上游响应一致
运行:
cargo build -p aether-executor
uv run python tests/e2e_rust_executor.py
"""
from __future__ import annotations
import asyncio
import base64
import contextlib
import gzip
import json
import signal
import subprocess
import sys
from pathlib import Path
import httpx
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
ExecutionProxySnapshot,
ExecutionPlanTimeouts,
)
from src.services.request.rust_executor_client import RustExecutorClient
async def _upstream_app(scope, receive, send) -> None: # type: ignore[no-untyped-def]
assert scope["type"] == "http"
body = b""
while True:
message = await receive()
if message["type"] != "http.request":
continue
body += message.get("body", b"")
if not message.get("more_body", False):
break
header_map = {
key.decode("latin1").lower(): value.decode("latin1")
for key, value in scope.get("headers", [])
}
path = str(scope.get("path") or "")
content_encoding = header_map.get("content-encoding")
decoded_body = gzip.decompress(body) if content_encoding == "gzip" else body
if path == "/v1/raw":
response_body = json.dumps(
{
"ok": True,
"path": path,
"received_body_text": decoded_body.decode("utf-8", errors="replace"),
"received_content_type": header_map.get("content-type"),
"received_content_encoding": content_encoding,
},
ensure_ascii=False,
).encode("utf-8")
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [
[b"content-type", b"application/json"],
[b"x-upstream-test", b"true"],
],
}
)
await send(
{
"type": "http.response.body",
"body": response_body,
"more_body": False,
}
)
return
if path == "/v1/download":
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [
[b"content-type", b"video/mp4"],
[b"x-download-test", b"true"],
],
}
)
await send(
{
"type": "http.response.body",
"body": b"video-download-bytes",
"more_body": False,
}
)
return
if path == "/v1/delete" and scope.get("method") == "DELETE":
await send(
{
"type": "http.response.start",
"status": 204,
"headers": [[b"x-delete-test", b"true"]],
}
)
await send(
{
"type": "http.response.body",
"body": b"",
"more_body": False,
}
)
return
payload = json.loads(decoded_body.decode("utf-8"))
if payload.get("stream") is True:
stream_body = (
"data: "
+ json.dumps(
{
"id": "chatcmpl-stream-1",
"object": "chat.completion.chunk",
"model": payload.get("model"),
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": "hello"},
"finish_reason": None,
}
],
},
ensure_ascii=False,
)
+ "\n\n"
+ "data: "
+ json.dumps(
{
"id": "chatcmpl-stream-1",
"object": "chat.completion.chunk",
"model": payload.get("model"),
"choices": [
{
"index": 0,
"delta": {},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 3,
"completion_tokens": 1,
"total_tokens": 4,
},
},
ensure_ascii=False,
)
+ "\n\n"
+ "data: [DONE]\n\n"
).encode("utf-8")
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [
[b"content-type", b"text/event-stream"],
[b"x-upstream-test", b"true"],
],
}
)
await send(
{
"type": "http.response.body",
"body": stream_body,
"more_body": False,
}
)
return
response = {
"ok": True,
"received_model": payload.get("model"),
"received_messages": payload.get("messages", []),
"path": path,
"received_content_encoding": content_encoding,
}
response_body = json.dumps(response, ensure_ascii=False).encode("utf-8")
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [
[b"content-type", b"application/json"],
[b"x-upstream-test", b"true"],
],
}
)
await send(
{
"type": "http.response.body",
"body": response_body,
"more_body": False,
}
)
async def _relay_app(scope, receive, send) -> None: # type: ignore[no-untyped-def]
assert scope["type"] == "http"
body = b""
while True:
message = await receive()
if message["type"] != "http.request":
continue
body += message.get("body", b"")
if not message.get("more_body", False):
break
if len(body) < 4:
await send(
{
"type": "http.response.start",
"status": 400,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
}
)
await send({"type": "http.response.body", "body": b"invalid relay envelope", "more_body": False})
return
path = str(scope.get("path") or "")
node_id = path.rsplit("/", 1)[-1]
meta_len = int.from_bytes(body[:4], "big")
if meta_len <= 0 or len(body) < 4 + meta_len:
await send(
{
"type": "http.response.start",
"status": 400,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
}
)
await send({"type": "http.response.body", "body": b"invalid relay metadata", "more_body": False})
return
meta = json.loads(body[4 : 4 + meta_len].decode("utf-8"))
request_json = json.loads(body[4 + meta_len :].decode("utf-8"))
if request_json.get("stream") is True:
stream_body = (
"data: "
+ json.dumps(
{
"id": "relay-stream-1",
"object": "chat.completion.chunk",
"model": request_json.get("model"),
"choices": [{"index": 0, "delta": {"content": "relay"}, "finish_reason": None}],
},
ensure_ascii=False,
)
+ "\n\n"
+ "data: [DONE]\n\n"
).encode("utf-8")
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [
[b"content-type", b"text/event-stream"],
[b"x-relay-node", node_id.encode("utf-8")],
],
}
)
await send({"type": "http.response.body", "body": stream_body, "more_body": False})
return
response_body = json.dumps(
{
"ok": True,
"via_tunnel": True,
"node_id": node_id,
"meta_url": meta.get("url"),
"received_model": request_json.get("model"),
},
ensure_ascii=False,
).encode("utf-8")
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [
[b"content-type", b"application/json"],
[b"x-relay-node", node_id.encode("utf-8")],
],
}
)
await send(
{
"type": "http.response.body",
"body": response_body,
"more_body": False,
}
)
async def _start_proxy_server() -> tuple[asyncio.AbstractServer, int, asyncio.Future[str]]:
loop = asyncio.get_running_loop()
request_line_future: asyncio.Future[str] = loop.create_future()
async def _handle_proxy(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
try:
header_bytes = await reader.readuntil(b"\r\n\r\n")
header_text = header_bytes.decode("latin-1")
lines = header_text.split("\r\n")
if lines and not request_line_future.done():
request_line_future.set_result(lines[0])
content_length = 0
for line in lines[1:]:
if not line:
continue
if line.lower().startswith("content-length:"):
content_length = int(line.split(":", 1)[1].strip() or "0")
break
if content_length:
await reader.readexactly(content_length)
body = json.dumps(
{
"ok": True,
"via_proxy": True,
"request_line": lines[0] if lines else "",
},
ensure_ascii=False,
).encode("utf-8")
writer.write(
b"HTTP/1.1 200 OK\r\n"
b"content-type: application/json\r\n"
b"x-proxy-test: true\r\n"
+ f"content-length: {len(body)}\r\n\r\n".encode("ascii")
+ body
)
await writer.drain()
finally:
writer.close()
with contextlib.suppress(Exception):
await writer.wait_closed()
server = await asyncio.start_server(_handle_proxy, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
return server, port, request_line_future
async def run_test() -> bool:
repo_root = Path(__file__).resolve().parents[1]
executor_binary = repo_root / "target" / "debug" / "aether-executor"
if not executor_binary.is_file():
print(f"FAIL: executor binary not found at {executor_binary}")
print(" run: cargo build -p aether-executor")
return False
upstream_listener = await asyncio.start_server(lambda r, w: None, "127.0.0.1", 0)
upstream_port = upstream_listener.sockets[0].getsockname()[1]
upstream_listener.close()
await upstream_listener.wait_closed()
import uvicorn
upstream_config = uvicorn.Config(
_upstream_app,
host="127.0.0.1",
port=upstream_port,
log_level="warning",
)
upstream_server = uvicorn.Server(upstream_config)
upstream_task = asyncio.create_task(upstream_server.serve())
relay_listener = await asyncio.start_server(lambda r, w: None, "127.0.0.1", 0)
relay_port = relay_listener.sockets[0].getsockname()[1]
relay_listener.close()
await relay_listener.wait_closed()
relay_config = uvicorn.Config(
_relay_app,
host="127.0.0.1",
port=relay_port,
log_level="warning",
)
relay_server = uvicorn.Server(relay_config)
relay_task = asyncio.create_task(relay_server.serve())
executor_socket = Path("/tmp/aether-executor-e2e.sock")
executor_socket.unlink(missing_ok=True)
executor_proc = subprocess.Popen(
[
str(executor_binary),
"--transport",
"unix_socket",
"--unix-socket",
str(executor_socket),
],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
)
try:
print("[1/4] Waiting for upstream server ...")
async with httpx.AsyncClient() as client:
for _ in range(50):
try:
resp = await client.post(
f"http://127.0.0.1:{upstream_port}/healthz",
json={"model": "probe"},
timeout=0.2,
)
if resp.status_code == 200:
break
except Exception:
await asyncio.sleep(0.1)
else:
print("FAIL: upstream server did not start in time")
return False
for _ in range(50):
try:
resp = await client.post(
f"http://127.0.0.1:{relay_port}/local/relay/probe",
content=(0).to_bytes(4, "big"),
timeout=0.2,
)
if resp.status_code in {400, 500}:
break
except Exception:
await asyncio.sleep(0.1)
else:
print("FAIL: relay server did not start in time")
return False
print("[2/4] Waiting for aether-executor ...")
for _ in range(100):
if executor_socket.exists():
try:
async with httpx.AsyncClient(
transport=httpx.AsyncHTTPTransport(uds=str(executor_socket)),
base_url="http://127.0.0.1:5219",
) as client:
resp = await client.get("/health", timeout=0.2)
if resp.status_code == 200:
break
except Exception:
pass
await asyncio.sleep(0.1)
else:
print("FAIL: executor did not start in time")
if executor_proc.stdout is not None:
print(executor_proc.stdout.read())
return False
print("[3/4] Sending ExecutionPlan via RustExecutorClient ...")
client = RustExecutorClient(
transport="unix_socket",
socket_path=str(executor_socket),
base_url="http://127.0.0.1:5219",
request_timeout=5.0,
)
result = await client.execute_sync_json(
ExecutionPlan(
request_id="req-e2e-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="POST",
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(
json_body={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "hello"}],
}
),
stream=False,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_type="application/json",
timeouts=ExecutionPlanTimeouts(
connect_ms=5_000,
read_ms=5_000,
write_ms=5_000,
pool_ms=5_000,
total_ms=5_000,
),
)
)
print("[4/4] Verifying response ...")
assert result.status_code == 200
assert result.headers.get("x-upstream-test") == "true"
assert result.response_json["ok"] is True
assert result.response_json["received_model"] == "gpt-4.1"
assert result.response_json["path"] == "/v1/chat/completions"
print("[extra] Verifying tunnel relay sync execution ...")
tunnel_result = await client.execute_sync_json(
ExecutionPlan(
request_id="req-e2e-tunnel-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="POST",
url="https://upstream-via-relay.test/v1/chat/completions",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(
json_body={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "hello tunnel"}],
}
),
stream=False,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_type="application/json",
proxy=ExecutionProxySnapshot(
enabled=True,
mode="tunnel",
node_id="node-1",
label="relay-node",
extra={"hub_base_url": f"http://127.0.0.1:{relay_port}"},
),
timeouts=ExecutionPlanTimeouts(
connect_ms=5_000,
read_ms=5_000,
write_ms=5_000,
pool_ms=5_000,
total_ms=5_000,
),
)
)
assert tunnel_result.status_code == 200
assert tunnel_result.headers.get("x-relay-node") == "node-1"
assert tunnel_result.response_json["via_tunnel"] is True
assert tunnel_result.response_json["meta_url"] == "https://upstream-via-relay.test/v1/chat/completions"
print("[extra] Verifying HTTP proxy execution ...")
proxy_server, proxy_port, proxy_request_line = await _start_proxy_server()
try:
proxied_result = await client.execute_sync_json(
ExecutionPlan(
request_id="req-e2e-proxy-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="POST",
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(
json_body={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "proxy hello"}],
}
),
stream=False,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_type="application/json",
proxy=ExecutionProxySnapshot(
enabled=True,
mode="http",
label="local-proxy",
url=f"http://127.0.0.1:{proxy_port}",
),
timeouts=ExecutionPlanTimeouts(
connect_ms=5_000,
read_ms=5_000,
write_ms=5_000,
pool_ms=5_000,
total_ms=5_000,
),
)
)
request_line = await asyncio.wait_for(proxy_request_line, timeout=5.0)
finally:
proxy_server.close()
await proxy_server.wait_closed()
assert proxied_result.status_code == 200
assert proxied_result.headers.get("x-proxy-test") == "true"
assert proxied_result.response_json["via_proxy"] is True
assert request_line.startswith("POST http://127.0.0.1:")
assert "/v1/chat/completions HTTP/1.1" in request_line
print("[extra] Verifying upstream-stream raw byte path ...")
stream_result = await client.execute_sync_json(
ExecutionPlan(
request_id="req-e2e-stream-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="POST",
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(
json_body={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "hello stream"}],
"stream": True,
}
),
stream=True,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_type="application/json",
timeouts=ExecutionPlanTimeouts(
connect_ms=5_000,
read_ms=5_000,
write_ms=5_000,
pool_ms=5_000,
total_ms=5_000,
),
)
)
assert stream_result.status_code == 200
assert stream_result.response_json is None
assert stream_result.response_body_bytes is not None
assert b"chat.completion.chunk" in stream_result.response_body_bytes
assert b"[DONE]" in stream_result.response_body_bytes
print("[extra] Verifying gzip request body path ...")
gzip_result = await client.execute_sync_json(
ExecutionPlan(
request_id="req-e2e-gzip-1",
candidate_id="cand-e2e-gzip-1",
provider_name="openai",
provider_id="prov-e2e",
endpoint_id="ep-e2e",
key_id="key-e2e",
method="POST",
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(
json_body={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "hello gzip"}],
}
),
stream=False,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_type="application/json",
content_encoding="gzip",
timeouts=ExecutionPlanTimeouts(connect_ms=5_000, total_ms=30_000),
)
)
assert gzip_result.status_code == 200
assert gzip_result.response_json is not None
assert gzip_result.response_json["received_content_encoding"] == "gzip"
assert gzip_result.response_json["received_model"] == "gpt-4.1"
print("[extra] Verifying raw body bytes path ...")
raw_result = await client.execute_sync_json(
ExecutionPlan(
request_id="req-e2e-raw-1",
candidate_id="cand-e2e-raw-1",
provider_name="openai",
provider_id="prov-e2e",
endpoint_id="ep-e2e",
key_id="key-e2e",
method="POST",
url=f"http://127.0.0.1:{upstream_port}/v1/raw",
headers={"content-type": "text/plain"},
body=ExecutionPlanBody(
body_bytes_b64=base64.b64encode(b"hello raw body").decode("ascii")
),
stream=False,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_type="text/plain",
timeouts=ExecutionPlanTimeouts(connect_ms=5_000, total_ms=30_000),
)
)
assert raw_result.status_code == 200
assert raw_result.response_json is not None
assert raw_result.response_json["received_body_text"] == "hello raw body"
assert raw_result.response_json["received_content_type"] == "text/plain"
print("[extra] Verifying sync empty-body delete path ...")
delete_result = await client.execute_sync_json(
ExecutionPlan(
request_id="req-e2e-delete-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="DELETE",
url=f"http://127.0.0.1:{upstream_port}/v1/delete",
headers={},
body=ExecutionPlanBody(),
stream=False,
provider_api_format="openai:video",
client_api_format="openai:video",
model_name="sora-2",
timeouts=ExecutionPlanTimeouts(
connect_ms=5_000,
read_ms=5_000,
write_ms=5_000,
pool_ms=5_000,
total_ms=5_000,
),
)
)
assert delete_result.status_code == 204
assert delete_result.response_json is None
assert delete_result.response_body_bytes is None
print("[extra] Verifying GET/no-body stream path ...")
download_stream = await client.execute_stream(
ExecutionPlan(
request_id="req-e2e-download-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="GET",
url=f"http://127.0.0.1:{upstream_port}/v1/download",
headers={},
body=ExecutionPlanBody(),
stream=True,
provider_api_format="openai:video",
client_api_format="openai:video",
model_name="sora-2",
timeouts=ExecutionPlanTimeouts(
connect_ms=5_000,
read_ms=5_000,
write_ms=5_000,
pool_ms=5_000,
total_ms=5_000,
),
)
)
try:
download_chunks = [chunk async for chunk in download_stream.byte_iterator]
finally:
await download_stream.response_ctx.__aexit__(None, None, None)
assert download_stream.status_code == 200
assert download_stream.headers.get("x-download-test") == "true"
assert b"".join(download_chunks) == b"video-download-bytes"
print("[extra] Verifying native stream path ...")
live_stream = await client.execute_stream(
ExecutionPlan(
request_id="req-e2e-live-stream-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="POST",
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(
json_body={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "hello live stream"}],
"stream": True,
}
),
stream=True,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_type="application/json",
timeouts=ExecutionPlanTimeouts(
connect_ms=5_000,
read_ms=5_000,
write_ms=5_000,
pool_ms=5_000,
total_ms=5_000,
),
)
)
try:
live_chunks = [chunk async for chunk in live_stream.byte_iterator]
finally:
await live_stream.response_ctx.__aexit__(None, None, None)
assert live_stream.status_code == 200
assert live_stream.headers.get("x-upstream-test") == "true"
live_body = b"".join(live_chunks)
assert live_body
assert b"chat.completion.chunk" in live_body
assert b"[DONE]" in live_body
print("[extra] Verifying native tunnel stream path ...")
tunnel_stream = await client.execute_stream(
ExecutionPlan(
request_id="req-e2e-tunnel-stream-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="POST",
url="https://upstream-via-relay.test/v1/chat/completions",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(
json_body={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "hello relay stream"}],
"stream": True,
}
),
stream=True,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_type="application/json",
proxy=ExecutionProxySnapshot(
enabled=True,
mode="tunnel",
node_id="node-1",
label="relay-node",
extra={"hub_base_url": f"http://127.0.0.1:{relay_port}"},
),
timeouts=ExecutionPlanTimeouts(
connect_ms=5_000,
read_ms=5_000,
write_ms=5_000,
pool_ms=5_000,
total_ms=5_000,
),
)
)
try:
tunnel_stream_chunks = [chunk async for chunk in tunnel_stream.byte_iterator]
finally:
await tunnel_stream.response_ctx.__aexit__(None, None, None)
tunnel_stream_body = b"".join(tunnel_stream_chunks)
assert tunnel_stream.status_code == 200
assert tunnel_stream.headers.get("x-relay-node") == "node-1"
assert b"relay-stream-1" in tunnel_stream_body
assert b"[DONE]" in tunnel_stream_body
print("ALL TESTS PASSED")
return True
finally:
upstream_server.should_exit = True
try:
await asyncio.wait_for(upstream_task, timeout=5.0)
except Exception:
upstream_task.cancel()
with contextlib.suppress(Exception):
await upstream_task
relay_server.should_exit = True
try:
await asyncio.wait_for(relay_task, timeout=5.0)
except Exception:
relay_task.cancel()
with contextlib.suppress(Exception):
await relay_task
executor_proc.send_signal(signal.SIGTERM)
try:
executor_proc.wait(timeout=5)
except subprocess.TimeoutExpired:
executor_proc.kill()
executor_proc.wait()
executor_socket.unlink(missing_ok=True)
if __name__ == "__main__":
success = asyncio.run(run_test())
sys.exit(0 if success else 1)

View File

@@ -5,9 +5,11 @@ import sys
import types
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock
import pytest
from src.config import config as app_config
from src.services.provider_keys.codex_usage_parser import (
CodexUsageParseError,
parse_codex_wham_usage_response,
@@ -17,6 +19,7 @@ from src.services.provider_keys.quota_refresh.antigravity_refresher import (
)
from src.services.provider_keys.quota_refresh.codex_refresher import refresh_codex_key_quota
from src.services.provider_keys.quota_refresh.kiro_refresher import refresh_kiro_key_quota
from src.services.request.rust_executor_client import RustExecutorSyncResult
class _FakeDB:
@@ -149,6 +152,73 @@ async def test_codex_refresher_http_non_200_returns_error(
assert result["status_code"] == 503
@pytest.mark.asyncio
async def test_codex_refresher_prefers_rust_executor(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.services.provider_keys.quota_refresh import codex_refresher as module
from src.services.request import rust_executor_client as rust_module
monkeypatch.setattr(app_config, "executor_backend", "rust")
key = SimpleNamespace(
id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
)
provider = SimpleNamespace(id="provider-1", proxy=None)
endpoint = SimpleNamespace(id="endpoint-1")
metadata_updates: dict[str, dict[str, Any]] = {}
state_updates: dict[str, dict[str, Any]] = {}
captured: dict[str, Any] = {}
async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
return None
_install_module(
monkeypatch,
"src.services.proxy_node.resolver",
{
"resolve_effective_proxy": lambda provider_proxy, key_proxy: {"enabled": True, "url": "http://proxy.test:8080"},
"build_proxy_client_kwargs": lambda proxy, timeout: {"timeout": timeout},
"get_system_proxy_config_async": AsyncMock(return_value=None),
"resolve_delegate_config_async": AsyncMock(return_value=None),
"build_proxy_url_async": AsyncMock(return_value="http://proxy.test:8080"),
"resolve_proxy_info_async": AsyncMock(return_value={"mode": "http", "label": "proxy"}),
},
)
monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
monkeypatch.setattr(
module, "parse_codex_wham_usage_response", lambda _data: {"used_percent": 12.5}
)
async def _fake_execute_sync_json(self: object, plan: Any) -> RustExecutorSyncResult:
captured["plan"] = plan
return RustExecutorSyncResult(
status_code=200,
response_json={"ok": True},
headers={"content-type": "application/json"},
)
monkeypatch.setattr(rust_module.RustExecutorClient, "execute_sync_json", _fake_execute_sync_json)
result = await refresh_codex_key_quota(
db=cast(Any, _FakeDB()),
provider=cast(Any, provider),
key=cast(Any, key),
endpoint=cast(Any, endpoint),
codex_wham_usage_url="https://example.test/wham/usage",
metadata_updates=metadata_updates,
state_updates=state_updates,
)
assert result["status"] == "success"
assert captured["plan"].method == "GET"
assert captured["plan"].url == "https://example.test/wham/usage"
assert captured["plan"].proxy is not None
assert captured["plan"].proxy.url == "http://proxy.test:8080"
assert metadata_updates["k1"]["codex"]["used_percent"] == 12.5
@pytest.mark.asyncio
async def test_codex_refresher_http_401_marks_auth_invalid_without_disabling_key(
monkeypatch: pytest.MonkeyPatch,

View File

@@ -1,11 +1,13 @@
from __future__ import annotations
from typing import Any
from unittest.mock import AsyncMock
import pytest
from src.services.provider_ops.service import ProviderOpsService
from src.services.provider_ops.types import ConnectorAuthType
from src.services.provider_ops.types import ConnectorAuthType, ProviderActionType
from src.services.request.rust_executor_client import RustExecutorSyncResult
class _FakeDB:
@@ -51,6 +53,43 @@ class _FakeRegistry:
return self._architecture
class _SuccessResult:
def __init__(self) -> None:
self.success = True
self.quota = 200.0
self.extra = {"window": "day"}
def to_dict(self) -> dict[str, Any]:
return {"success": True, "quota": self.quota}
class _SuccessArchitecture:
default_action_configs = {ProviderActionType.QUERY_BALANCE: {"quota_divisor": 100}}
def get_verify_endpoint(self) -> str:
return "/verify"
async def prepare_verify_config(
self,
_base_url: str,
_config: dict[str, Any],
_credentials: dict[str, Any],
) -> dict[str, Any]:
return {}
def build_verify_headers(
self,
_config: dict[str, Any],
_credentials: dict[str, Any],
) -> dict[str, str]:
return {"authorization": "Bearer test"}
def parse_verify_response(self, status_code: int, data: dict[str, Any]) -> _SuccessResult:
assert status_code == 200
assert data == {"ok": True}
return _SuccessResult()
@pytest.mark.asyncio
async def test_verify_auth_returns_failure_when_prepare_verify_config_raises_value_error(
monkeypatch: pytest.MonkeyPatch,
@@ -72,3 +111,57 @@ async def test_verify_auth_returns_failure_when_prepare_verify_config_raises_val
)
assert result == {"success": False, "message": "invalid refresh token"}
@pytest.mark.asyncio
async def test_verify_auth_prefers_rust_executor(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.services.provider_ops import service as module
from src.services.request import rust_executor_client as rust_module
service = ProviderOpsService(_FakeDB())
architecture = _SuccessArchitecture()
monkeypatch.setattr(
module,
"get_registry",
lambda: _FakeRegistry(architecture),
)
monkeypatch.setattr(module.config, "executor_backend", "rust")
monkeypatch.setattr(
"src.services.proxy_node.resolver.resolve_ops_proxy_config_async",
AsyncMock(return_value=(None, "node-1")),
)
cache_balance = AsyncMock()
monkeypatch.setattr(service, "_cache_balance_from_verify", cache_balance)
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={"ok": True},
headers={"content-type": "application/json"},
)
monkeypatch.setattr(rust_module.RustExecutorClient, "execute_sync_json", _fake_execute_sync_json)
result = await service.verify_auth(
base_url="https://example.com",
architecture_id="sub2api",
auth_type=ConnectorAuthType.SESSION_LOGIN,
config={},
credentials={"access_token": "token"},
provider_id="provider-1",
)
assert result == {"success": True, "quota": 200.0}
assert captured["plan"].method == "GET"
assert captured["plan"].url == "https://example.com/verify"
assert captured["plan"].proxy is not None
assert captured["plan"].proxy.mode == "tunnel"
assert captured["plan"].proxy.node_id == "node-1"
cache_balance.assert_awaited_once_with("provider-1", 2.0, {"window": "day"})

View File

@@ -0,0 +1,173 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import src.services.task.video.cancel as cancel_mod
import src.services.request.rust_executor_client as rust_client_mod
from src.core.api_format.conversion.internal_video import VideoStatus
from src.services.request.rust_executor_client import RustExecutorSyncResult
from src.services.task.video.cancel import VideoTaskCancelService
class _Query:
def __init__(self, result: object) -> None:
self._result = result
def filter(self, *args: object, **kwargs: object) -> "_Query":
del args, kwargs
return self
def first(self) -> object:
return self._result
class _FakeDB:
def __init__(self, *results: object) -> None:
self._results = list(results)
self.committed = False
def query(self, model: object) -> _Query:
del model
return _Query(self._results.pop(0))
def commit(self) -> None:
self.committed = True
@pytest.mark.asyncio
async def test_video_cancel_service_uses_rust_for_openai_delete(
monkeypatch: pytest.MonkeyPatch,
) -> None:
endpoint = SimpleNamespace(
id="ep-1",
provider_id="prov-1",
api_family="openai",
endpoint_kind="video",
base_url="https://api.openai.com",
)
key = SimpleNamespace(id="key-1", api_key="encrypted")
db = _FakeDB(endpoint, key)
service = VideoTaskCancelService(db)
task = SimpleNamespace(
id="task-1",
status=VideoStatus.PROCESSING.value,
external_task_id="ext-1",
endpoint_id="ep-1",
key_id="key-1",
request_id="req-1",
model="sora-2",
completed_at=None,
updated_at=None,
)
monkeypatch.setattr(cancel_mod.config, "executor_backend", "rust")
monkeypatch.setattr(
rust_client_mod.RustExecutorClient,
"execute_sync_json",
AsyncMock(
return_value=RustExecutorSyncResult(
status_code=204,
headers={},
response_json=None,
response_body_bytes=None,
)
),
)
monkeypatch.setattr(
"src.core.crypto.crypto_service.decrypt",
lambda value: "upstream-key",
)
monkeypatch.setattr(
"src.core.api_format.build_upstream_headers_for_endpoint",
lambda *args, **kwargs: {"authorization": "Bearer upstream-key"},
)
monkeypatch.setattr(
"src.services.provider.transport.build_provider_url",
lambda endpoint, is_stream=False, key=None: "https://api.openai.com/v1/videos",
)
monkeypatch.setattr(
"src.clients.http_client.HTTPClientPool.get_default_client_async",
AsyncMock(side_effect=AssertionError("python fallback should not run")),
)
monkeypatch.setattr(cancel_mod.UsageService, "finalize_void", lambda *args, **kwargs: True)
result = await service.cancel_task(
task=task,
task_id="task-1",
original_headers={"x-test": "1"},
)
assert result is None
assert db.committed is True
assert task.status == VideoStatus.CANCELLED.value
@pytest.mark.asyncio
async def test_video_cancel_service_uses_rust_for_gemini_cancel(
monkeypatch: pytest.MonkeyPatch,
) -> None:
endpoint = SimpleNamespace(
id="ep-1",
provider_id="prov-1",
api_family="gemini",
endpoint_kind="video",
base_url="https://generativelanguage.googleapis.com",
)
key = SimpleNamespace(id="key-1", api_key="encrypted")
db = _FakeDB(endpoint, key)
service = VideoTaskCancelService(db)
task = SimpleNamespace(
id="task-1",
status=VideoStatus.PROCESSING.value,
external_task_id="operations/ext-1",
endpoint_id="ep-1",
key_id="key-1",
request_id="req-1",
model="veo-3",
completed_at=None,
updated_at=None,
)
monkeypatch.setattr(cancel_mod.config, "executor_backend", "rust")
execute_sync = AsyncMock(
return_value=RustExecutorSyncResult(
status_code=200,
headers={"content-type": "application/json"},
response_json={"ok": True},
response_body_bytes=None,
)
)
monkeypatch.setattr(rust_client_mod.RustExecutorClient, "execute_sync_json", execute_sync)
monkeypatch.setattr(
"src.core.crypto.crypto_service.decrypt",
lambda value: "upstream-key",
)
monkeypatch.setattr(
"src.core.api_format.build_upstream_headers_for_endpoint",
lambda *args, **kwargs: {"x-goog-api-key": "upstream-key"},
)
monkeypatch.setattr(
"src.services.provider.auth.get_provider_auth",
AsyncMock(return_value=SimpleNamespace(auth_header="authorization", auth_value="Bearer token")),
)
monkeypatch.setattr(
"src.clients.http_client.HTTPClientPool.get_default_client_async",
AsyncMock(side_effect=AssertionError("python fallback should not run")),
)
monkeypatch.setattr(cancel_mod.UsageService, "finalize_void", lambda *args, **kwargs: True)
result = await service.cancel_task(
task=task,
task_id="task-1",
original_headers={"x-test": "1"},
)
assert result is None
assert db.committed is True
assert task.status == VideoStatus.CANCELLED.value
plan = execute_sync.await_args.args[0]
assert plan.method == "POST"
assert plan.body.json_body == {}

View File

@@ -6,7 +6,8 @@ import pytest
from src.core.api_format.conversion.internal_video import InternalVideoPollResult, VideoStatus
from src.models.database import VideoTask
from src.services.task.video.poller_adapter import VideoTaskPollerAdapter
from src.services.request.rust_executor_client import RustExecutorSyncResult
from src.services.task.video.poller_adapter import VideoPollContext, VideoTaskPollerAdapter
@pytest.mark.asyncio
@@ -78,3 +79,140 @@ async def test_update_task_after_poll_skips_terminal_cancelled_task(
finalize.assert_not_awaited()
session.commit.assert_not_called()
@pytest.mark.asyncio
async def test_video_poller_try_rust_payload_passes_proxy_snapshot(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.services.request import rust_executor_client as rust_mod
from src.services.task.video import poller_adapter as mod
adapter = VideoTaskPollerAdapter()
monkeypatch.setattr(mod.config, "executor_backend", "rust")
proxy_snapshot = object()
captured: dict[str, object] = {}
async def _fake_execute_sync_json(self: object, plan: object) -> RustExecutorSyncResult:
captured["plan"] = plan
return RustExecutorSyncResult(status_code=200, response_json={"id": "op_1", "done": False})
monkeypatch.setattr(rust_mod.RustExecutorClient, "execute_sync_json", _fake_execute_sync_json)
ctx = VideoPollContext(
task_id="task-1",
external_task_id="vid_1",
provider_api_format="openai:video",
base_url="https://api.openai.com",
upstream_key="upstream",
headers={"authorization": "Bearer test"},
poll_count=0,
retry_count=0,
poll_interval_seconds=15,
max_poll_count=10,
current_status=VideoStatus.PROCESSING.value,
proxy_snapshot=proxy_snapshot,
)
payload = await adapter._try_rust_poll_payload(
ctx=ctx,
url="https://api.openai.com/v1/videos/vid_1",
)
assert payload == {"id": "op_1", "done": False}
assert getattr(captured["plan"], "method") == "GET"
assert getattr(captured["plan"], "proxy") is proxy_snapshot
@pytest.mark.asyncio
async def test_video_poller_openai_poll_prefers_rust_payload(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.services.task.video import poller_adapter as mod
adapter = VideoTaskPollerAdapter()
rust_poll = AsyncMock(return_value={"id": "vid_1", "status": "processing"})
monkeypatch.setattr(adapter, "_try_rust_poll_payload", rust_poll)
normalized = InternalVideoPollResult(status=VideoStatus.PROCESSING, progress_percent=42)
normalizer = MagicMock(return_value=normalized)
monkeypatch.setattr(adapter._openai_normalizer, "video_poll_to_internal", normalizer)
get_upstream_client = AsyncMock(side_effect=AssertionError("python upstream client should not be used"))
monkeypatch.setattr(mod.HTTPClientPool, "get_upstream_client", get_upstream_client)
ctx = VideoPollContext(
task_id="task-1",
external_task_id="vid_1",
provider_api_format="openai:video",
base_url="https://api.openai.com",
upstream_key="upstream",
headers={"authorization": "Bearer test"},
poll_count=0,
retry_count=0,
poll_interval_seconds=15,
max_poll_count=10,
current_status=VideoStatus.PROCESSING.value,
)
result = await adapter._poll_openai_with_context(ctx)
assert result is normalized
normalizer.assert_called_once_with({"id": "vid_1", "status": "processing"})
get_upstream_client.assert_not_awaited()
rust_poll.assert_awaited_once()
@pytest.mark.asyncio
async def test_video_poller_openai_poll_fallback_uses_transport_aware_client(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.services.task.video import poller_adapter as mod
adapter = VideoTaskPollerAdapter()
monkeypatch.setattr(adapter, "_try_rust_poll_payload", AsyncMock(return_value=None))
response = MagicMock()
response.status_code = 200
response.json.return_value = {"id": "vid_2", "status": "processing"}
client = MagicMock()
client.get = AsyncMock(return_value=response)
get_upstream_client = AsyncMock(return_value=client)
monkeypatch.setattr(mod.HTTPClientPool, "get_upstream_client", get_upstream_client)
normalized = InternalVideoPollResult(status=VideoStatus.PROCESSING, progress_percent=7)
normalizer = MagicMock(return_value=normalized)
monkeypatch.setattr(adapter._openai_normalizer, "video_poll_to_internal", normalizer)
ctx = VideoPollContext(
task_id="task-2",
external_task_id="vid_2",
provider_api_format="openai:video",
base_url="https://api.openai.com",
upstream_key="upstream",
headers={"authorization": "Bearer test"},
poll_count=0,
retry_count=0,
poll_interval_seconds=15,
max_poll_count=10,
current_status=VideoStatus.PROCESSING.value,
proxy_config={"enabled": True, "url": "http://proxy.test:8080"},
delegate_config={"node_id": "node-1", "tunnel": True},
)
result = await adapter._poll_openai_with_context(ctx)
assert result is normalized
get_upstream_client.assert_awaited_once_with(
{"node_id": "node-1", "tunnel": True},
proxy_config={"enabled": True, "url": "http://proxy.test:8080"},
)
client.get.assert_awaited_once_with(
"https://api.openai.com/v1/videos/vid_2",
headers={"authorization": "Bearer test"},
)

View File

@@ -0,0 +1,43 @@
from __future__ import annotations
import pytest
from src.config.settings import Config
def test_executor_config_defaults_to_rust_and_unix_socket(
monkeypatch: pytest.MonkeyPatch,
) -> None:
for key in (
"EXECUTOR_BACKEND",
"EXECUTOR_TRANSPORT",
"EXECUTOR_SOCKET_PATH",
"EXECUTOR_BASE_URL",
"EXECUTOR_REQUEST_TIMEOUT",
"HTTP_REQUEST_TIMEOUT",
):
monkeypatch.delenv(key, raising=False)
cfg = Config()
assert cfg.executor_backend == "rust"
assert cfg.executor_transport == "unix_socket"
assert cfg.executor_socket_path == "/tmp/aether-executor.sock"
assert cfg.executor_base_url == "http://127.0.0.1:5219"
assert cfg.executor_request_timeout == cfg.http_request_timeout
def test_executor_config_accepts_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("EXECUTOR_BACKEND", "RUST")
monkeypatch.setenv("EXECUTOR_TRANSPORT", "TCP")
monkeypatch.setenv("EXECUTOR_SOCKET_PATH", "/var/run/aether.sock")
monkeypatch.setenv("EXECUTOR_BASE_URL", "http://127.0.0.1:9311")
monkeypatch.setenv("EXECUTOR_REQUEST_TIMEOUT", "12.5")
cfg = Config()
assert cfg.executor_backend == "rust"
assert cfg.executor_transport == "tcp"
assert cfg.executor_socket_path == "/var/run/aether.sock"
assert cfg.executor_base_url == "http://127.0.0.1:9311"
assert cfg.executor_request_timeout == 12.5

View File

@@ -0,0 +1,344 @@
from __future__ import annotations
import base64
from src.services.request.executor_plan import PreparedExecutionPlan
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
ExecutionPlanTimeouts,
ExecutionProxySnapshot,
build_execution_plan_body,
)
def test_execution_plan_to_payload_drops_none_fields() -> None:
plan = ExecutionPlan(
request_id="req-1",
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",
proxy=ExecutionProxySnapshot(
enabled=True,
mode="direct",
label="no-proxy",
),
timeouts=ExecutionPlanTimeouts(
connect_ms=10_000,
total_ms=300_000,
),
)
payload = plan.to_payload()
assert "candidate_id" not in payload
assert payload["body"] == {"json_body": {"model": "gpt-4.1"}}
assert payload["proxy"] == {
"enabled": True,
"mode": "direct",
"label": "no-proxy",
}
assert payload["timeouts"] == {
"connect_ms": 10_000,
"total_ms": 300_000,
}
def test_prepared_execution_plan_remote_eligible_for_non_stream_json() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
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=300.0,
)
assert prepared.remote_eligible is True
def test_prepared_execution_plan_remote_eligible_allows_upstream_stream() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
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": True}),
stream=True,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
),
payload={"model": "gpt-4.1", "stream": True},
headers={"content-type": "application/json"},
upstream_is_stream=True,
needs_conversion=False,
provider_type="openai",
request_timeout=300.0,
)
assert prepared.remote_eligible is True
def test_prepared_execution_plan_remote_eligible_allows_tunnel_delegate_proxy() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
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="gemini:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
proxy=ExecutionProxySnapshot(
enabled=True,
mode="tunnel",
node_id="node-1",
label="relay-node",
),
),
payload={"model": "gpt-4.1"},
headers={"content-type": "application/json"},
upstream_is_stream=False,
needs_conversion=True,
provider_type="gemini",
request_timeout=300.0,
delegate_config={"tunnel": True, "node_id": "node-1"},
proxy_config={"node_id": "node-1"},
)
assert prepared.remote_eligible is True
def test_prepared_execution_plan_remote_eligible_allows_url_proxy_without_delegate() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
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",
proxy=ExecutionProxySnapshot(
enabled=True,
mode="http",
label="proxy.internal",
url="http://proxy.internal:8080",
),
),
payload={"model": "gpt-4.1"},
headers={"content-type": "application/json"},
upstream_is_stream=False,
needs_conversion=False,
provider_type="openai",
request_timeout=300.0,
proxy_config={"url": "http://proxy.internal:8080"},
)
assert prepared.remote_eligible is True
def test_prepared_execution_plan_remote_eligible_allows_conversion_and_envelope() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
candidate_id=None,
provider_name="gemini",
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="gemini: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=True,
provider_type="gemini",
request_timeout=300.0,
envelope=object(),
)
assert prepared.remote_eligible is True
def test_prepared_execution_plan_remote_eligible_allows_tls_profile() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
candidate_id=None,
provider_name="claude",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="POST",
url="https://example.com/v1/messages",
headers={"content-type": "application/json"},
body=ExecutionPlanBody(json_body={"model": "claude-3.7-sonnet"}),
stream=False,
provider_api_format="claude:chat",
client_api_format="claude:chat",
model_name="claude-3.7-sonnet",
tls_profile="claude_code_nodejs",
),
payload={"model": "claude-3.7-sonnet"},
headers={"content-type": "application/json"},
upstream_is_stream=False,
needs_conversion=False,
provider_type="claude_code",
request_timeout=300.0,
)
assert prepared.remote_eligible is True
def test_build_execution_plan_body_encodes_raw_bytes_payload() -> None:
body = build_execution_plan_body(
b"raw-payload",
content_type="text/plain",
)
assert body.json_body is None
assert base64.b64decode(body.body_bytes_b64 or "") == b"raw-payload"
def test_prepared_execution_plan_remote_eligible_allows_gzip_json_body() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
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",
content_encoding="gzip",
),
payload={"model": "gpt-4.1"},
headers={"content-type": "application/json"},
upstream_is_stream=False,
needs_conversion=False,
provider_type="openai",
request_timeout=300.0,
)
assert prepared.remote_eligible is True
def test_prepared_execution_plan_remote_eligible_allows_raw_body_with_passthrough_encoding() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
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": "text/plain"},
body=ExecutionPlanBody(body_bytes_b64=base64.b64encode(b"raw").decode("ascii")),
stream=False,
provider_api_format="openai:chat",
client_api_format="openai:chat",
model_name="gpt-4.1",
content_encoding="br",
),
payload={"model": "gpt-4.1"},
headers={"content-type": "text/plain"},
upstream_is_stream=False,
needs_conversion=False,
provider_type="openai",
request_timeout=300.0,
)
assert prepared.remote_eligible is True
def test_prepared_execution_plan_remote_eligible_allows_empty_get_body() -> None:
prepared = PreparedExecutionPlan(
contract=ExecutionPlan(
request_id="req-1",
candidate_id=None,
provider_name="openai",
provider_id="prov-1",
endpoint_id="ep-1",
key_id="key-1",
method="GET",
url="https://example.com/v1/videos/video-1/content",
headers={},
body=ExecutionPlanBody(),
stream=True,
provider_api_format="openai:video",
client_api_format="openai:video",
model_name="sora-2",
),
payload={},
headers={},
upstream_is_stream=True,
needs_conversion=False,
provider_type="openai",
request_timeout=300.0,
)
assert prepared.remote_eligible is True