mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
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
296
tests/api/public/test_gemini_files_rust.py
Normal file
296
tests/api/public/test_gemini_files_rust.py
Normal 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
|
||||
142
tests/api/public/test_system_catalog_rust.py
Normal file
142
tests/api/public/test_system_catalog_rust.py
Normal 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"},
|
||||
)
|
||||
Reference in New Issue
Block a user