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
+91
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import ipaddress
import time
from datetime import datetime, timezone
from enum import Enum
@@ -43,6 +44,12 @@ QUIET_POLLING_PATHS: set[str] = {
"/api/wallet/today-cost",
}
TRUSTED_GATEWAY_HEADER = "x-aether-gateway"
TRUSTED_AUTH_USER_ID_HEADER = "x-aether-auth-user-id"
TRUSTED_AUTH_API_KEY_ID_HEADER = "x-aether-auth-api-key-id"
TRUSTED_AUTH_BALANCE_HEADER = "x-aether-auth-balance-remaining"
TRUSTED_AUTH_ACCESS_ALLOWED_HEADER = "x-aether-auth-access-allowed"
class ApiRequestPipeline:
"""负责统一执行认证、余额校验、上下文构建等通用逻辑的管道。"""
@@ -349,6 +356,10 @@ class ApiRequestPipeline:
async def _authenticate_client(
self, request: Request, db: Session, adapter: ApiAdapter, **_kw: object
) -> tuple[User, ApiKey]:
trusted_auth = self._try_trusted_gateway_auth(request, db)
if trusted_auth is not None:
return trusted_auth
client_api_key = adapter.extract_api_key(request)
if not client_api_key:
raise HTTPException(status_code=401, detail="请提供API密钥")
@@ -393,6 +404,86 @@ class ApiRequestPipeline:
return db_user, db_api_key
def _try_trusted_gateway_auth(
self,
request: Request,
db: Session,
) -> tuple[User, ApiKey] | None:
if not self._is_trusted_gateway_request(request):
return None
user_id = str(request.headers.get(TRUSTED_AUTH_USER_ID_HEADER) or "").strip()
api_key_id = str(request.headers.get(TRUSTED_AUTH_API_KEY_ID_HEADER) or "").strip()
if not user_id or not api_key_id:
return None
db_user = db.query(User).filter(User.id == user_id).first()
db_api_key = db.query(ApiKey).filter(ApiKey.id == api_key_id).first()
if not db_user or not db_api_key:
return None
if not db_user.is_active or db_user.is_deleted:
return None
if not db_api_key.is_active:
return None
if db_api_key.is_locked and not db_api_key.is_standalone:
raise HTTPException(status_code=403, detail="该密钥已被管理员锁定,请联系管理员")
if db_api_key.user_id != db_user.id:
return None
if db_api_key.expires_at:
expires_at = db_api_key.expires_at
if expires_at.tzinfo is None:
expires_at = expires_at.replace(tzinfo=timezone.utc)
if expires_at < datetime.now(timezone.utc):
return None
balance_remaining = self._parse_trusted_balance_remaining(
request.headers.get(TRUSTED_AUTH_BALANCE_HEADER)
)
access_allowed = self._parse_trusted_bool_header(
request.headers.get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER),
default=True,
)
request.state.user_id = db_user.id
request.state.api_key_id = db_api_key.id
request.state.prefetched_balance_remaining = balance_remaining
if not access_allowed:
raise BalanceInsufficientException(balance_type="USD", remaining=balance_remaining)
return db_user, db_api_key
@staticmethod
def _is_trusted_gateway_request(request: Request) -> bool:
gateway_marker = str(request.headers.get(TRUSTED_GATEWAY_HEADER) or "").strip().lower()
if not gateway_marker.startswith("rust-phase3"):
return False
host = request.client.host if request.client else ""
try:
return ipaddress.ip_address(host).is_loopback
except ValueError:
return False
@staticmethod
def _parse_trusted_balance_remaining(value: str | None) -> float | None:
raw = str(value or "").strip()
if not raw:
return None
try:
return float(raw)
except ValueError:
return None
@staticmethod
def _parse_trusted_bool_header(value: str | None, *, default: bool) -> bool:
raw = str(value or "").strip().lower()
if raw in {"1", "true", "yes", "on"}:
return True
if raw in {"0", "false", "no", "off"}:
return False
return default
async def _try_token_prefix_auth(
self, token: str, request: Request, db: Session
) -> tuple[User, Any] | None:
+141 -5
View File
@@ -92,6 +92,17 @@ from src.services.provider.transport import (
build_provider_url,
)
from src.services.provider.upstream_headers import build_upstream_extra_headers
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
ExecutionProxySnapshot,
build_execution_plan_body,
is_remote_contract_eligible,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
from src.services.scheduling.aware_scheduler import ProviderCandidate
from src.services.system.config import SystemConfigService
from src.services.task.request_state import MutableRequestBodyState
@@ -958,7 +969,10 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
# 解析有效代理(Key 级别优先于 Provider 级别)
from src.services.proxy_node.resolver import (
build_proxy_url_async,
get_proxy_label,
get_system_proxy_config_async,
resolve_delegate_config_async,
resolve_effective_proxy,
resolve_proxy_info_async,
)
@@ -971,6 +985,7 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
f" [{self.request_id}] 发送流式请求: Provider={provider.name}, "
f"模型={ctx.model} -> {mapped_model or '无映射'}, 代理={proxy_label}"
)
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
# If upstream is forced to non-stream mode, we execute a sync request and then
# simulate streaming to the client (sync -> stream bridge).
@@ -978,11 +993,9 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
from src.clients.http_client import HTTPClientPool
from src.services.proxy_node.resolver import (
build_post_kwargs_async,
resolve_delegate_config_async,
)
request_timeout_sync = provider.request_timeout or config.http_request_timeout
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
http_client = await HTTPClientPool.get_upstream_client(
delegate_cfg,
proxy_config=effective_proxy,
@@ -1179,19 +1192,142 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
return _streamified()
if config.executor_backend == "rust":
effective_proxy_for_contract = effective_proxy
if not effective_proxy_for_contract or not effective_proxy_for_contract.get(
"enabled", True
):
effective_proxy_for_contract = await get_system_proxy_config_async()
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
proxy_url: str | None = None
if effective_proxy_for_contract and not is_tunnel_delegate:
proxy_url = await build_proxy_url_async(effective_proxy_for_contract)
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
ctx.proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if is_tunnel_delegate else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if is_tunnel_delegate
else None
),
)
rust_plan = ExecutionPlan(
request_id=str(self.request_id or ""),
candidate_id=str(
getattr(candidate, "request_candidate_id", "")
or getattr(candidate, "id", "")
or ""
)
or None,
provider_name=str(provider.name),
provider_id=str(provider.id),
endpoint_id=str(endpoint.id),
key_id=str(key.id),
method="POST",
url=url,
headers=dict(provider_headers),
body=build_execution_plan_body(
provider_payload,
content_type=str(provider_headers.get("content-type") or "").strip() or None,
),
stream=True,
provider_api_format=provider_api_format,
client_api_format=client_api_format,
model_name=str(ctx.model or ""),
content_type=str(provider_headers.get("content-type") or "").strip() or None,
content_encoding=client_content_encoding,
proxy=proxy_snapshot,
tls_profile=tls_profile,
timeouts=ExecutionPlanTimeouts(
connect_ms=int(config.http_connect_timeout * 1000),
read_ms=int(config.http_read_timeout * 1000),
write_ms=int(config.http_write_timeout * 1000),
pool_ms=int(config.http_pool_timeout * 1000),
total_ms=None,
),
)
if is_remote_contract_eligible(rust_plan):
try:
rust_stream = await RustExecutorClient().execute_stream(rust_plan)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning(
"[{}] Rust executor stream 不可用,回退 Python 执行: {}",
self.request_id,
exc,
)
else:
ctx.status_code = rust_stream.status_code
ctx.response_headers = dict(rust_stream.headers)
ctx.set_proxy_timing(ctx.response_headers)
if envelope:
envelope.on_http_status(
base_url=ctx.selected_base_url,
status_code=ctx.status_code,
)
try:
if ctx.status_code >= 400:
error_chunks: list[bytes] = []
async for chunk in rust_stream.byte_iterator:
if chunk:
error_chunks.append(chunk)
if sum(len(item) for item in error_chunks) >= 4000:
break
error_body = b"".join(error_chunks)[:4000].decode(
"utf-8",
errors="replace",
)
request = httpx.Request("POST", url, headers=provider_headers)
response = httpx.Response(
ctx.status_code,
request=request,
headers=rust_stream.headers,
content=b"".join(error_chunks),
)
error = httpx.HTTPStatusError(
f"Upstream status error: {ctx.status_code}",
request=request,
response=response,
)
error.upstream_response = error_body # type: ignore[attr-defined]
raise error
prefetched_chunks = await stream_processor.prefetch_and_check_error(
rust_stream.byte_iterator,
provider,
endpoint,
ctx,
max_prefetch_lines=config.stream_prefetch_lines,
)
except Exception:
await rust_stream.response_ctx.__aexit__(None, None, None)
raise
return stream_processor.create_response_stream(
ctx,
rust_stream.byte_iterator,
rust_stream.response_ctx,
prefetched_chunks,
start_time=self.start_time,
)
# 流式请求使用 stream_first_byte_timeout 作为首字节超时
# 优先使用 Provider 配置,否则使用全局配置
request_timeout = provider.stream_first_byte_timeout or config.stream_first_byte_timeout
request_timeout = (
getattr(provider, "stream_first_byte_timeout", None) or config.stream_first_byte_timeout
)
# 获取 HTTP 客户端(支持代理配置,Key 级别优先于 Provider 级别)
# 使用连接池复用客户端,避免每次流式请求都新建 TCP/TLS 连接
from src.clients.http_client import HTTPClientPool
from src.services.proxy_node.resolver import (
build_stream_kwargs_async,
resolve_delegate_config_async,
)
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
http_client = await HTTPClientPool.get_upstream_client(
delegate_cfg,
proxy_config=effective_proxy,
+416 -99
View File
@@ -35,6 +35,7 @@ from src.api.handlers.base.utils import (
resolve_client_accept_encoding,
resolve_client_content_encoding,
)
from src.config.settings import config
from src.core.error_utils import extract_client_error_message
from src.core.exceptions import (
EmbeddedErrorException,
@@ -46,6 +47,17 @@ from src.core.exceptions import (
UpstreamClientException,
)
from src.core.logger import logger
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
ExecutionProxySnapshot,
PreparedExecutionPlan,
build_execution_plan_body,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
from src.services.task.request_state import MutableRequestBodyState
if TYPE_CHECKING:
@@ -442,6 +454,39 @@ class ChatSyncExecutor:
client_content_encoding: str | None = None,
) -> dict[str, Any]:
"""单次同步请求(原 sync_request_func 内嵌函数)"""
prepared_plan = await self._build_sync_execution_plan(
provider,
endpoint,
key,
candidate,
model=model,
api_format=api_format,
original_headers=original_headers,
request_state=request_state,
query_params=query_params,
client_content_encoding=client_content_encoding,
)
return await self._execute_sync_plan(
prepared_plan=prepared_plan,
provider=provider,
model=model,
)
async def _build_sync_execution_plan(
self,
provider: Provider,
endpoint: ProviderEndpoint,
key: ProviderAPIKey,
candidate: ProviderCandidate,
*,
model: str,
api_format: Any,
original_headers: dict[str, Any],
request_state: MutableRequestBodyState,
query_params: dict[str, str] | None = None,
client_content_encoding: str | None = None,
) -> PreparedExecutionPlan:
"""构建可序列化的执行计划,并保留本地执行所需的运行时上下文。"""
handler = self._handler
ctx = self._ctx
@@ -518,7 +563,9 @@ class ChatSyncExecutor:
# 解析有效代理(Key 级别优先于 Provider 级别)
from src.services.proxy_node.resolver import (
build_proxy_url_async,
get_proxy_label,
get_system_proxy_config_async,
resolve_effective_proxy,
resolve_proxy_info_async,
)
@@ -536,13 +583,8 @@ class ChatSyncExecutor:
)
logger.debug(f" [{handler.request_id}] 请求URL: {redact_url_for_log(url)}")
# 获取复用的 HTTP 客户端(支持代理配置,Key 级别优先于 Provider 级别)
# 注意:使用 get_proxy_client 复用连接池,不再每次创建新客户端
from src.clients.http_client import HTTPClientPool
from src.config.settings import config
# 解析 delegate 配置,用于后续本地执行或交给 Rust executor。
from src.services.proxy_node.resolver import (
build_post_kwargs_async,
build_stream_kwargs_async,
resolve_delegate_config_async,
)
@@ -551,49 +593,354 @@ class ChatSyncExecutor:
request_timeout = provider.request_timeout or config.http_request_timeout
delegate_cfg = await resolve_delegate_config_async(_effective_proxy)
http_client = await HTTPClientPool.get_upstream_client(
delegate_cfg,
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
effective_proxy_for_contract = _effective_proxy
if not effective_proxy_for_contract or not effective_proxy_for_contract.get(
"enabled", True
):
effective_proxy_for_contract = await get_system_proxy_config_async()
proxy_url: str | None = None
if effective_proxy_for_contract and not is_tunnel_delegate:
proxy_url = await build_proxy_url_async(effective_proxy_for_contract)
return PreparedExecutionPlan(
contract=ExecutionPlan(
request_id=str(handler.request_id or ""),
candidate_id=str(
getattr(candidate, "request_candidate_id", "")
or getattr(candidate, "id", "")
or ""
)
or None,
provider_name=str(provider.name),
provider_id=str(provider.id),
endpoint_id=str(endpoint.id),
key_id=str(key.id),
method="POST",
url=url,
headers=dict(provider_hdrs),
body=build_execution_plan_body(
provider_payload,
content_type=str(provider_hdrs.get("content-type") or "").strip() or None,
),
stream=upstream_is_stream,
provider_api_format=provider_api_format,
client_api_format=client_api_format,
model_name=str(model or ""),
content_type=str(provider_hdrs.get("content-type") or "").strip() or None,
content_encoding=client_content_encoding,
proxy=ExecutionProxySnapshot.from_proxy_info(
ctx.sync_proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if is_tunnel_delegate else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if is_tunnel_delegate
else None
),
),
tls_profile=tls_profile,
timeouts=ExecutionPlanTimeouts(
connect_ms=int(config.http_connect_timeout * 1000),
read_ms=int(config.http_read_timeout * 1000),
write_ms=int(config.http_write_timeout * 1000),
pool_ms=int(config.http_pool_timeout * 1000),
total_ms=int(request_timeout * 1000),
),
),
payload=provider_payload,
headers=dict(provider_hdrs),
upstream_is_stream=upstream_is_stream,
needs_conversion=needs_conversion,
provider_type=provider_type,
request_timeout=request_timeout,
delegate_config=delegate_cfg,
proxy_config=_effective_proxy,
tls_profile=tls_profile,
envelope=envelope,
selected_base_url=selected_base_url_cached,
client_content_encoding=client_content_encoding,
proxy_info=ctx.sync_proxy_info,
)
async def _execute_sync_plan(
self,
*,
prepared_plan: PreparedExecutionPlan,
provider: Provider,
model: str,
) -> dict[str, Any]:
handler = self._handler
ctx = self._ctx
if config.executor_backend == "rust" and prepared_plan.remote_eligible:
try:
rust_result = await RustExecutorClient().execute_sync_json(prepared_plan.contract)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning(
"[{}] Rust executor 不可用,回退 Python 执行: {}",
handler.request_id,
exc,
)
else:
result = await self._finalize_rust_sync_result(
prepared_plan=prepared_plan,
provider=provider,
model=model,
status_code=rust_result.status_code,
response_headers=rust_result.headers,
response_json=rust_result.response_json,
provider_response_json=rust_result.provider_response_json,
response_body_bytes=rust_result.response_body_bytes,
)
logger.debug(
"[{}] sync chat 请求由 Rust executor 执行完成",
handler.request_id,
)
return result
return await self._execute_sync_plan_locally(
prepared_plan=prepared_plan,
provider=provider,
model=model,
)
async def _finalize_rust_sync_result(
self,
*,
prepared_plan: PreparedExecutionPlan,
provider: Provider,
model: str,
status_code: int,
response_headers: dict[str, str],
response_json: dict[str, Any] | None,
provider_response_json: dict[str, Any] | None = None,
response_body_bytes: bytes | None = None,
) -> dict[str, Any]:
ctx = self._ctx
synthetic_content = response_body_bytes
if synthetic_content is None:
synthetic_content = json.dumps(
response_json or {},
ensure_ascii=False,
).encode("utf-8")
request = httpx.Request(
prepared_plan.contract.method,
prepared_plan.contract.url,
headers=prepared_plan.headers,
)
synthetic_response = httpx.Response(
status_code,
request=request,
headers=response_headers,
content=synthetic_content,
)
ctx.status_code = status_code
ctx.response_headers = dict(synthetic_response.headers)
extract_proxy_timing(ctx.sync_proxy_info, ctx.response_headers)
if prepared_plan.envelope:
prepared_plan.envelope.on_http_status(
base_url=prepared_plan.selected_base_url,
status_code=ctx.status_code,
)
if status_code >= 400:
error = httpx.HTTPStatusError(
f"Upstream status error: {status_code}",
request=request,
response=synthetic_response,
)
error_body = ""
try:
if prepared_plan.envelope and hasattr(prepared_plan.envelope, "extract_error_text"):
error_body = await prepared_plan.envelope.extract_error_text(synthetic_response)
else:
error_body = synthetic_response.text[:4000] if synthetic_response.text else ""
except Exception:
error_body = synthetic_response.text[:4000] if synthetic_response.text else ""
error.upstream_response = error_body[:4000] # type: ignore[attr-defined]
raise error
if prepared_plan.upstream_is_stream:
if response_body_bytes is None:
raise RustExecutorClientError("Rust executor stream result must contain body bytes")
return await self._finalize_rust_stream_sync_result(
prepared_plan=prepared_plan,
provider=provider,
model=model,
response_body_bytes=response_body_bytes,
)
if response_json is None:
raise RustExecutorClientError("Rust executor sync result must contain response_json")
ctx.response_json = dict(response_json)
if provider_response_json is not None:
ctx.provider_response_json = dict(provider_response_json)
if prepared_plan.envelope:
ctx.response_json = prepared_plan.envelope.unwrap_response(ctx.response_json)
prepared_plan.envelope.postprocess_unwrapped_response(
model=model,
data=ctx.response_json,
)
if isinstance(ctx.response_json, dict):
parser = get_parser_for_format(ctx.provider_api_format_for_error or "")
if parser.is_error_response(ctx.response_json):
parsed = parser.parse_response(ctx.response_json, status_code)
raise EmbeddedErrorException(
provider_name=str(provider.name),
error_code=parsed.embedded_status_code,
error_message=parsed.error_message,
error_status=parsed.error_type,
)
if ctx.needs_conversion_for_error and isinstance(ctx.response_json, dict):
ctx.provider_response_json = ctx.response_json.copy()
registry = get_format_converter_registry()
ctx.response_json = registry.convert_response(
ctx.response_json,
ctx.provider_api_format_for_error or "",
ctx.client_api_format_for_error or "",
requested_model=model,
)
return ctx.response_json if isinstance(ctx.response_json, dict) else {}
async def _finalize_rust_stream_sync_result(
self,
*,
prepared_plan: PreparedExecutionPlan,
provider: Provider,
model: str,
response_body_bytes: bytes,
) -> dict[str, Any]:
async def _iter_body() -> Any:
yield response_body_bytes
return await self._aggregate_upstream_stream_response(
byte_iter=_iter_body(),
prepared_plan=prepared_plan,
provider=provider,
model=model,
)
async def _aggregate_upstream_stream_response(
self,
*,
byte_iter: Any,
prepared_plan: PreparedExecutionPlan,
provider: Provider,
model: str,
) -> dict[str, Any]:
ctx = self._ctx
provider_parser = (
get_parser_for_format(ctx.provider_api_format_for_error)
if ctx.provider_api_format_for_error
else None
)
if (
prepared_plan.provider_type == "kiro"
and prepared_plan.envelope
and prepared_plan.envelope.force_stream_rewrite()
):
from src.services.provider.adapters.kiro.eventstream_rewriter import (
apply_kiro_stream_rewrite,
)
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
from src.api.handlers.base.upstream_stream_bridge import (
aggregate_upstream_stream_to_internal_response,
)
internal_resp = await aggregate_upstream_stream_to_internal_response(
byte_iter,
provider_api_format=ctx.provider_api_format_for_error or "",
provider_name=str(provider.name),
model=str(model or ""),
request_id=str(self._handler.request_id or ""),
envelope=prepared_plan.envelope,
provider_parser=provider_parser,
)
registry = get_format_converter_registry()
tgt_norm = (
registry.get_normalizer(ctx.client_api_format_for_error)
if ctx.client_api_format_for_error
else None
)
if tgt_norm is None:
raise RuntimeError(f"未注册 Normalizer: {ctx.client_api_format_for_error}")
ctx.response_json = tgt_norm.response_from_internal(
internal_resp,
requested_model=model,
)
ctx.response_json = ctx.response_json if isinstance(ctx.response_json, dict) else {}
return ctx.response_json
async def _execute_sync_plan_locally(
self,
*,
prepared_plan: PreparedExecutionPlan,
provider: Provider,
model: str,
) -> dict[str, Any]:
handler = self._handler
ctx = self._ctx
from src.clients.http_client import HTTPClientPool
from src.services.proxy_node.resolver import (
build_post_kwargs_async,
build_stream_kwargs_async,
)
http_client = await HTTPClientPool.get_upstream_client(
prepared_plan.delegate_config,
proxy_config=prepared_plan.proxy_config,
tls_profile=prepared_plan.contract.tls_profile,
)
# 注意:不使用 async with,因为复用的客户端不应该被关闭
# 超时通过 timeout 参数控制
resp: httpx.Response | None = None
if not upstream_is_stream:
if not prepared_plan.upstream_is_stream:
try:
_pkw = await build_post_kwargs_async(
delegate_cfg,
url=url,
headers=provider_hdrs,
payload=provider_payload,
timeout=request_timeout,
client_content_encoding=client_content_encoding,
prepared_plan.delegate_config,
url=prepared_plan.contract.url,
headers=prepared_plan.headers,
payload=prepared_plan.payload,
timeout=prepared_plan.request_timeout,
client_content_encoding=prepared_plan.client_content_encoding,
)
resp = await http_client.post(**_pkw)
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.TimeoutException) as e:
if envelope:
envelope.on_connection_error(base_url=selected_base_url_cached, exc=e)
if selected_base_url_cached:
if prepared_plan.envelope:
prepared_plan.envelope.on_connection_error(
base_url=prepared_plan.selected_base_url,
exc=e,
)
if prepared_plan.selected_base_url:
logger.warning(
f"[{envelope.name}] Connection error: "
f"{selected_base_url_cached} ({e})"
f"[{prepared_plan.envelope.name}] Connection error: "
f"{prepared_plan.selected_base_url} ({e})"
)
raise
else:
# Forced upstream streaming: aggregate SSE to a sync JSON response.
provider_parser = (
get_parser_for_format(provider_api_format) if provider_api_format else None
)
try:
_stream_args = await build_stream_kwargs_async(
delegate_cfg,
url=url,
headers=provider_hdrs,
payload=provider_payload,
timeout=request_timeout,
client_content_encoding=client_content_encoding,
prepared_plan.delegate_config,
url=prepared_plan.contract.url,
headers=prepared_plan.headers,
payload=prepared_plan.payload,
timeout=prepared_plan.request_timeout,
client_content_encoding=prepared_plan.client_content_encoding,
)
async with http_client.stream(**_stream_args) as stream_resp:
resp = stream_resp
@@ -602,58 +949,31 @@ class ChatSyncExecutor:
ctx.response_headers = dict(stream_resp.headers)
extract_proxy_timing(ctx.sync_proxy_info, ctx.response_headers)
if envelope:
envelope.on_http_status(
base_url=selected_base_url_cached,
if prepared_plan.envelope:
prepared_plan.envelope.on_http_status(
base_url=prepared_plan.selected_base_url,
status_code=ctx.status_code,
)
stream_resp.raise_for_status()
byte_iter = stream_resp.aiter_bytes()
if provider_type == "kiro" and envelope and envelope.force_stream_rewrite():
from src.services.provider.adapters.kiro.eventstream_rewriter import (
apply_kiro_stream_rewrite,
)
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
from src.api.handlers.base.upstream_stream_bridge import (
aggregate_upstream_stream_to_internal_response,
)
internal_resp = await aggregate_upstream_stream_to_internal_response(
byte_iter,
provider_api_format=provider_api_format,
provider_name=str(provider.name),
model=str(model or ""),
request_id=str(handler.request_id or ""),
envelope=envelope,
provider_parser=provider_parser,
)
registry = get_format_converter_registry()
tgt_norm = (
registry.get_normalizer(client_api_format) if client_api_format else None
)
if tgt_norm is None:
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
ctx.response_json = tgt_norm.response_from_internal(
internal_resp,
requested_model=model,
)
ctx.response_json = (
ctx.response_json if isinstance(ctx.response_json, dict) else {}
ctx.response_json = await self._aggregate_upstream_stream_response(
byte_iter=stream_resp.aiter_bytes(),
prepared_plan=prepared_plan,
provider=provider,
model=model,
)
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.TimeoutException) as e:
if envelope:
envelope.on_connection_error(base_url=selected_base_url_cached, exc=e)
if selected_base_url_cached:
if prepared_plan.envelope:
prepared_plan.envelope.on_connection_error(
base_url=prepared_plan.selected_base_url,
exc=e,
)
if prepared_plan.selected_base_url:
logger.warning(
f"[{envelope.name}] Connection error: "
f"{selected_base_url_cached} ({e})"
f"[{prepared_plan.envelope.name}] Connection error: "
f"{prepared_plan.selected_base_url} ({e})"
)
raise
@@ -661,35 +981,32 @@ class ChatSyncExecutor:
ctx.response_headers = dict(resp.headers)
extract_proxy_timing(ctx.sync_proxy_info, ctx.response_headers)
if envelope:
envelope.on_http_status(base_url=selected_base_url_cached, status_code=ctx.status_code)
if prepared_plan.envelope:
prepared_plan.envelope.on_http_status(
base_url=prepared_plan.selected_base_url,
status_code=ctx.status_code,
)
# Forced upstream streaming already built response_json via aggregator.
if upstream_is_stream:
if prepared_plan.upstream_is_stream:
return ctx.response_json if isinstance(ctx.response_json, dict) else {}
# 统一使用 HTTPStatusError,让 TaskService/error_classifier 负责分类
# (客户端错误/兼容性错误/限流等)
try:
resp.raise_for_status()
except httpx.HTTPStatusError as e:
error_body = ""
try:
if envelope and hasattr(envelope, "extract_error_text"):
error_body = await envelope.extract_error_text(resp)
if prepared_plan.envelope and hasattr(prepared_plan.envelope, "extract_error_text"):
error_body = await prepared_plan.envelope.extract_error_text(resp)
else:
error_body = resp.text[:4000] if resp.text else ""
except Exception:
error_body = ""
# 供 ErrorClassifier 优先读取
e.upstream_response = error_body # type: ignore[attr-defined]
raise
# 安全解析 JSON 响应,处理可能的编码错误
try:
ctx.response_json = resp.json()
except (UnicodeDecodeError, json.JSONDecodeError) as e:
# 获取原始响应内容用于调试(存入 upstream_response
raw_content = ""
try:
raw_content = resp.text[:500] if resp.text else "(empty)"
@@ -699,7 +1016,6 @@ class ChatSyncExecutor:
except Exception:
raw_content = "(unable to read)"
logger.error(f"[{handler.request_id}] 无法解析响应 JSON: {e}, 原始内容: {raw_content}")
# 判断错误类型,生成友好的客户端错误消息(不暴露提供商信息)
if raw_content == "(empty)" or not raw_content.strip():
client_message = "上游服务返回了空响应"
elif raw_content.strip().startswith(("<", "<!doctype", "<!DOCTYPE")):
@@ -713,13 +1029,15 @@ class ChatSyncExecutor:
upstream_response=raw_content,
)
if envelope:
ctx.response_json = envelope.unwrap_response(ctx.response_json)
envelope.postprocess_unwrapped_response(model=model, data=ctx.response_json)
if prepared_plan.envelope:
ctx.response_json = prepared_plan.envelope.unwrap_response(ctx.response_json)
prepared_plan.envelope.postprocess_unwrapped_response(
model=model,
data=ctx.response_json,
)
# 检查响应体中的嵌套错误(HTTP 200 但响应体包含错误)
if isinstance(ctx.response_json, dict):
parser = get_parser_for_format(provider_api_format)
parser = get_parser_for_format(ctx.provider_api_format_for_error or "")
if parser.is_error_response(ctx.response_json):
parsed = parser.parse_response(ctx.response_json, 200)
logger.warning(
@@ -736,15 +1054,14 @@ class ChatSyncExecutor:
error_status=parsed.error_type,
)
# 跨格式:响应转换回 client_format(失败触发 failover
if needs_conversion and isinstance(ctx.response_json, dict):
if ctx.needs_conversion_for_error and isinstance(ctx.response_json, dict):
ctx.provider_response_json = ctx.response_json.copy()
registry = get_format_converter_registry()
ctx.response_json = registry.convert_response(
ctx.response_json,
provider_api_format,
client_api_format,
requested_model=model, # 使用用户请求的原始模型名
ctx.provider_api_format_for_error or "",
ctx.client_api_format_for_error or "",
requested_model=model,
)
return ctx.response_json if isinstance(ctx.response_json, dict) else {}
+312 -62
View File
@@ -38,6 +38,17 @@ from src.core.exceptions import (
)
from src.core.logger import logger
from src.services.provider.behavior import get_provider_behavior
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
ExecutionProxySnapshot,
build_execution_plan_body,
is_remote_contract_eligible,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
from src.services.scheduling.aware_scheduler import ProviderCandidate
from src.services.system.config import SystemConfigService
from src.services.task.request_state import MutableRequestBodyState
@@ -53,6 +64,67 @@ if TYPE_CHECKING:
class CliStreamMixin:
"""流式处理核心方法的 Mixin"""
def _streamify_sync_response(
self: CliHandlerProtocol,
*,
ctx: StreamContext,
response_json: dict[str, Any],
client_api_format: str,
provider_api_format: str,
) -> AsyncGenerator[bytes]:
registry = get_format_converter_registry()
src_norm = registry.get_normalizer(provider_api_format) if provider_api_format else None
if src_norm is None:
raise RuntimeError(f"未注册 Normalizer: {provider_api_format}")
internal_resp = src_norm.response_to_internal(
response_json if isinstance(response_json, dict) else {}
)
internal_resp.model = str(ctx.model or internal_resp.model or "")
if internal_resp.id:
ctx.response_id = internal_resp.id
if internal_resp.usage:
ctx.input_tokens = int(internal_resp.usage.input_tokens or 0)
ctx.output_tokens = int(internal_resp.usage.output_tokens or 0)
ctx.cached_tokens = int(internal_resp.usage.cache_read_tokens or 0)
ctx.cache_creation_tokens = int(internal_resp.usage.cache_write_tokens or 0)
from src.core.api_format.conversion.stream_state import StreamState
tgt_norm = registry.get_normalizer(client_api_format) if client_api_format else None
if tgt_norm is None:
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
state = StreamState(
model=str(ctx.model or ""),
message_id=str(ctx.response_id or ctx.request_id or self.request_id or ""),
)
output_state = {"first_yield": True, "streaming_updated": False}
async def _streamified() -> AsyncGenerator[bytes]:
for ev in iter_internal_response_as_stream_events(internal_resp):
converted_events = tgt_norm.stream_event_from_internal(ev, state)
if not converted_events:
continue
self._record_converted_chunks(ctx, converted_events)
for sse_line in _format_converted_events_to_sse(
converted_events, client_api_format
):
if not sse_line:
continue
ctx.chunk_count += 1
self._mark_first_output(ctx, output_state)
yield (sse_line + "\n").encode("utf-8")
if str(client_api_format or "").strip().lower() == "openai:chat":
ctx.chunk_count += 1
self._mark_first_output(ctx, output_state)
yield b"data: [DONE]\n\n"
ctx.has_completion = True
return _streamified()
async def process_stream(
self: CliHandlerProtocol,
original_request_body: dict[str, Any],
@@ -341,24 +413,156 @@ class CliStreamMixin:
ctx.selected_base_url = upstream_request.selected_base_url
# 解析有效代理(Key 级别优先于 Provider 级别)
from src.services.proxy_node.resolver import build_proxy_url_async as _bpua
from src.services.proxy_node.resolver import get_proxy_label as _gpl
from src.services.proxy_node.resolver import resolve_delegate_config_async as _rda
from src.services.proxy_node.resolver import resolve_effective_proxy as _rep
from src.services.proxy_node.resolver import resolve_proxy_info_async as _rpi_async
effective_proxy = _rep(provider.proxy, getattr(key, "proxy", None))
ctx.proxy_info = await _rpi_async(effective_proxy)
delegate_cfg = await _rda(effective_proxy)
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
proxy_url: str | None = None
if effective_proxy and not is_tunnel_delegate:
proxy_url = await _bpua(effective_proxy)
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
ctx.proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if is_tunnel_delegate else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if is_tunnel_delegate
else None
),
)
# If upstream is forced to non-stream mode, we execute a sync request and then
# simulate streaming to the client (sync -> stream bridge).
if not upstream_is_stream:
from src.clients.http_client import HTTPClientPool
from src.services.proxy_node.resolver import (
build_post_kwargs_async,
resolve_delegate_config_async,
)
from src.services.proxy_node.resolver import build_post_kwargs_async
request_timeout_sync = provider.request_timeout or config.http_request_timeout
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
rust_plan = ExecutionPlan(
request_id=str(self.request_id or ""),
candidate_id=str(
getattr(candidate, "request_candidate_id", "")
or getattr(candidate, "id", "")
or ""
)
or None,
provider_name=str(provider.name),
provider_id=str(provider.id),
endpoint_id=str(endpoint.id),
key_id=str(key.id),
method="POST",
url=url,
headers=dict(provider_headers),
body=build_execution_plan_body(
provider_payload,
content_type=str(provider_headers.get("content-type") or "").strip() or None,
),
stream=False,
provider_api_format=provider_api_format,
client_api_format=client_api_format,
model_name=str(ctx.model or ""),
content_type=str(provider_headers.get("content-type") or "").strip() or None,
content_encoding=client_content_encoding,
proxy=proxy_snapshot,
tls_profile=envelope_tls_profile,
timeouts=ExecutionPlanTimeouts(
connect_ms=int(config.http_connect_timeout * 1000),
read_ms=int(config.http_read_timeout * 1000),
write_ms=int(config.http_write_timeout * 1000),
pool_ms=int(config.http_pool_timeout * 1000),
total_ms=int(request_timeout_sync * 1000),
),
)
if config.executor_backend == "rust" and is_remote_contract_eligible(rust_plan):
try:
rust_result = await RustExecutorClient().execute_sync_json(rust_plan)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning(
"[{}] CLI Rust executor(sync->stream) 不可用,回退 Python 执行: {}",
self.request_id,
exc,
)
else:
ctx.status_code = rust_result.status_code
ctx.response_headers = dict(rust_result.headers)
ctx.set_proxy_timing(ctx.response_headers)
if envelope:
envelope.on_http_status(
base_url=ctx.selected_base_url,
status_code=ctx.status_code,
)
request = httpx.Request("POST", url, headers=provider_headers)
synthetic_content = rust_result.response_body_bytes
if synthetic_content is None:
synthetic_content = json.dumps(
rust_result.response_json or {},
ensure_ascii=False,
).encode("utf-8")
synthetic_response = httpx.Response(
ctx.status_code,
request=request,
headers=ctx.response_headers,
content=synthetic_content,
)
if ctx.status_code >= 400:
error = httpx.HTTPStatusError(
f"Upstream status error: {ctx.status_code}",
request=request,
response=synthetic_response,
)
error_body = ""
try:
if envelope and hasattr(envelope, "extract_error_text"):
error_body = await envelope.extract_error_text(synthetic_response)
else:
error_body = (
synthetic_response.text[:4000]
if synthetic_response.text
else ""
)
except Exception:
error_body = (
synthetic_response.text[:4000] if synthetic_response.text else ""
)
error.upstream_response = error_body # type: ignore[attr-defined]
raise error
response_json = rust_result.response_json or {}
if envelope:
response_json = envelope.unwrap_response(response_json)
envelope.postprocess_unwrapped_response(model=ctx.model, data=response_json)
if isinstance(response_json, dict) and provider_api_format:
parser = get_parser_for_format(provider_api_format)
if parser.is_error_response(response_json):
parsed = parser.parse_response(response_json, 200)
raise EmbeddedErrorException(
provider_name=str(provider.name),
error_code=parsed.embedded_status_code,
error_message=parsed.error_message,
error_status=parsed.error_type,
)
if isinstance(response_json, dict):
ctx.response_metadata = self._extract_response_metadata(response_json)
return self._streamify_sync_response(
ctx=ctx,
response_json=response_json if isinstance(response_json, dict) else {},
client_api_format=client_api_format,
provider_api_format=provider_api_format,
)
http_client = await HTTPClientPool.get_upstream_client(
delegate_cfg,
proxy_config=effective_proxy,
@@ -488,60 +692,12 @@ class CliStreamMixin:
if isinstance(response_json, dict):
ctx.response_metadata = self._extract_response_metadata(response_json)
# Convert sync JSON -> InternalResponse, then InternalResponse -> client stream events.
registry = get_format_converter_registry()
src_norm = registry.get_normalizer(provider_api_format) if provider_api_format else None
if src_norm is None:
raise RuntimeError(f"未注册 Normalizer: {provider_api_format}")
internal_resp = src_norm.response_to_internal(
response_json if isinstance(response_json, dict) else {}
return self._streamify_sync_response(
ctx=ctx,
response_json=response_json if isinstance(response_json, dict) else {},
client_api_format=client_api_format,
provider_api_format=provider_api_format,
)
internal_resp.model = str(ctx.model or internal_resp.model or "")
if internal_resp.id:
ctx.response_id = internal_resp.id
if internal_resp.usage:
ctx.input_tokens = int(internal_resp.usage.input_tokens or 0)
ctx.output_tokens = int(internal_resp.usage.output_tokens or 0)
ctx.cached_tokens = int(internal_resp.usage.cache_read_tokens or 0)
ctx.cache_creation_tokens = int(internal_resp.usage.cache_write_tokens or 0)
from src.core.api_format.conversion.stream_state import StreamState
tgt_norm = registry.get_normalizer(client_api_format) if client_api_format else None
if tgt_norm is None:
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
state = StreamState(
model=str(ctx.model or ""),
message_id=str(ctx.response_id or ctx.request_id or self.request_id or ""),
)
output_state = {"first_yield": True, "streaming_updated": False}
async def _streamified() -> AsyncGenerator[bytes]:
for ev in iter_internal_response_as_stream_events(internal_resp):
converted_events = tgt_norm.stream_event_from_internal(ev, state)
if not converted_events:
continue
self._record_converted_chunks(ctx, converted_events)
for sse_line in _format_converted_events_to_sse(
converted_events, client_api_format
):
if not sse_line:
continue
ctx.chunk_count += 1
self._mark_first_output(ctx, output_state)
yield (sse_line + "\n").encode("utf-8")
# OpenAI chat clients expect a final [DONE] marker.
if str(client_api_format or "").strip().lower() == "openai:chat":
ctx.chunk_count += 1
self._mark_first_output(ctx, output_state)
yield b"data: [DONE]\n\n"
ctx.has_completion = True
return _streamified()
# 流式请求使用 stream_first_byte_timeout 作为首字节超时
# 优先使用 Provider 配置,否则使用全局配置
@@ -560,12 +716,106 @@ class CliStreamMixin:
# 获取 HTTP 客户端(支持代理配置,Key 级别优先于 Provider 级别)
# 使用连接池复用客户端,避免每次流式请求都新建 TCP/TLS 连接
from src.clients.http_client import HTTPClientPool
from src.services.proxy_node.resolver import (
build_stream_kwargs_async,
resolve_delegate_config_async,
from src.services.proxy_node.resolver import build_stream_kwargs_async
rust_plan = ExecutionPlan(
request_id=str(self.request_id or ""),
candidate_id=str(
getattr(candidate, "request_candidate_id", "") or getattr(candidate, "id", "") or ""
)
or None,
provider_name=str(provider.name),
provider_id=str(provider.id),
endpoint_id=str(endpoint.id),
key_id=str(key.id),
method="POST",
url=url,
headers=dict(provider_headers),
body=build_execution_plan_body(
provider_payload,
content_type=str(provider_headers.get("content-type") or "").strip() or None,
),
stream=True,
provider_api_format=provider_api_format,
client_api_format=client_api_format,
model_name=str(ctx.model or ""),
content_type=str(provider_headers.get("content-type") or "").strip() or None,
content_encoding=client_content_encoding,
proxy=proxy_snapshot,
tls_profile=envelope_tls_profile,
timeouts=ExecutionPlanTimeouts(
connect_ms=int(config.http_connect_timeout * 1000),
read_ms=int(config.http_read_timeout * 1000),
write_ms=int(config.http_write_timeout * 1000),
pool_ms=int(config.http_pool_timeout * 1000),
total_ms=None,
),
)
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
if config.executor_backend == "rust" and is_remote_contract_eligible(rust_plan):
try:
rust_stream = await RustExecutorClient().execute_stream(rust_plan)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning(
"[{}] CLI Rust executor stream 不可用,回退 Python 执行: {}",
self.request_id,
exc,
)
else:
ctx.status_code = rust_stream.status_code
ctx.response_headers = dict(rust_stream.headers)
ctx.set_proxy_timing(ctx.response_headers)
if envelope:
envelope.on_http_status(
base_url=ctx.selected_base_url,
status_code=ctx.status_code,
)
try:
if ctx.status_code >= 400:
error_chunks: list[bytes] = []
async for chunk in rust_stream.byte_iterator:
if chunk:
error_chunks.append(chunk)
if sum(len(item) for item in error_chunks) >= 4000:
break
error_body = b"".join(error_chunks)[:4000].decode(
"utf-8",
errors="replace",
)
request = httpx.Request("POST", url, headers=provider_headers)
response = httpx.Response(
ctx.status_code,
request=request,
headers=rust_stream.headers,
content=b"".join(error_chunks),
)
error = httpx.HTTPStatusError(
f"Upstream status error: {ctx.status_code}",
request=request,
response=response,
)
error.upstream_response = error_body # type: ignore[attr-defined]
raise error
prefetched_chunks = await self._prefetch_and_check_embedded_error(
rust_stream.byte_iterator,
provider,
endpoint,
ctx,
)
except Exception:
await rust_stream.response_ctx.__aexit__(None, None, None)
raise
return self._create_response_stream_with_prefetch(
ctx,
rust_stream.byte_iterator,
rust_stream.response_ctx,
prefetched_chunks,
)
http_client = await HTTPClientPool.get_upstream_client(
delegate_cfg,
proxy_config=effective_proxy,
+200 -35
View File
@@ -31,6 +31,17 @@ from src.core.exceptions import (
ThinkingSignatureException,
)
from src.core.logger import logger
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
ExecutionProxySnapshot,
build_execution_plan_body,
is_remote_contract_eligible,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
from src.services.scheduling.aware_scheduler import ProviderCandidate
from src.services.task.request_state import MutableRequestBodyState
@@ -42,6 +53,54 @@ if TYPE_CHECKING:
class CliSyncMixin:
"""同步处理相关方法的 Mixin"""
async def _aggregate_upstream_stream_sync_response(
self: CliHandlerProtocol,
*,
body_bytes: bytes,
provider_api_format: str,
client_api_format: str,
provider_name: str,
provider_type: str,
model: str,
request_id: str,
envelope: Any,
) -> dict[str, Any]:
registry = get_format_converter_registry()
provider_parser = (
get_parser_for_format(provider_api_format) if provider_api_format else None
)
async def _byte_iter() -> Any:
yield body_bytes
byte_iter = _byte_iter()
if provider_type == "kiro" and envelope and envelope.force_stream_rewrite():
from src.services.provider.adapters.kiro.eventstream_rewriter import (
apply_kiro_stream_rewrite,
)
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
internal_resp = await aggregate_upstream_stream_to_internal_response(
byte_iter,
provider_api_format=provider_api_format,
provider_name=provider_name,
model=model,
request_id=request_id,
envelope=envelope,
provider_parser=provider_parser,
)
tgt_norm = registry.get_normalizer(client_api_format) if client_api_format else None
if tgt_norm is None:
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
response_json = tgt_norm.response_from_internal(
internal_resp,
requested_model=model,
)
return response_json if isinstance(response_json, dict) else {}
async def process_sync(
self: CliHandlerProtocol,
original_request_body: dict[str, Any],
@@ -179,6 +238,7 @@ class CliSyncMixin:
from src.clients.http_client import HTTPClientPool
from src.services.proxy_node.resolver import (
build_post_kwargs_async,
build_proxy_url_async,
build_stream_kwargs_async,
resolve_delegate_config_async,
)
@@ -188,6 +248,141 @@ class CliSyncMixin:
request_timeout = provider.request_timeout or config.http_request_timeout
delegate_cfg = await resolve_delegate_config_async(_effective_proxy)
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
proxy_url: str | None = None
if _effective_proxy and not is_tunnel_delegate:
proxy_url = await build_proxy_url_async(_effective_proxy)
rust_plan = ExecutionPlan(
request_id=str(self.request_id or ""),
candidate_id=str(
getattr(candidate, "request_candidate_id", "")
or getattr(candidate, "id", "")
or ""
)
or None,
provider_name=str(provider.name),
provider_id=str(provider.id),
endpoint_id=str(endpoint.id),
key_id=str(key.id),
method="POST",
url=url,
headers=dict(provider_headers),
body=build_execution_plan_body(
provider_payload,
content_type=str(provider_headers.get("content-type") or "").strip() or None,
),
stream=upstream_is_stream,
provider_api_format=provider_api_format,
client_api_format=client_api_format,
model_name=str(model or ""),
content_type=str(provider_headers.get("content-type") or "").strip() or None,
content_encoding=effective_client_content_encoding,
proxy=ExecutionProxySnapshot.from_proxy_info(
sync_proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if is_tunnel_delegate else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if is_tunnel_delegate
else None
),
),
tls_profile=envelope_tls_profile,
timeouts=ExecutionPlanTimeouts(
connect_ms=int(config.http_connect_timeout * 1000),
read_ms=int(config.http_read_timeout * 1000),
write_ms=int(config.http_write_timeout * 1000),
pool_ms=int(config.http_pool_timeout * 1000),
total_ms=int(request_timeout * 1000),
),
)
if config.executor_backend == "rust" and is_remote_contract_eligible(rust_plan):
try:
rust_result = await RustExecutorClient().execute_sync_json(rust_plan)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning(
"[{}] CLI Rust executor 不可用,回退 Python 执行: {}",
self.request_id,
exc,
)
else:
status_code = rust_result.status_code
response_headers = dict(rust_result.headers)
extract_proxy_timing(sync_proxy_info, response_headers)
if envelope:
envelope.on_http_status(
base_url=selected_base_url_cached,
status_code=status_code,
)
request = httpx.Request("POST", url, headers=provider_headers)
synthetic_content = rust_result.response_body_bytes
if synthetic_content is None:
synthetic_content = json.dumps(
rust_result.response_json or {},
ensure_ascii=False,
).encode("utf-8")
synthetic_response = httpx.Response(
status_code,
request=request,
headers=response_headers,
content=synthetic_content,
)
if status_code >= 400:
error = httpx.HTTPStatusError(
f"Upstream status error: {status_code}",
request=request,
response=synthetic_response,
)
error_body = ""
try:
if envelope and hasattr(envelope, "extract_error_text"):
error_body = await envelope.extract_error_text(synthetic_response)
else:
error_body = (
synthetic_response.text[:4000]
if synthetic_response.text
else ""
)
except Exception:
error_body = (
synthetic_response.text[:4000] if synthetic_response.text else ""
)
error.upstream_response = error_body[:4000] # type: ignore[attr-defined]
raise error
if upstream_is_stream:
if rust_result.response_body_bytes is None:
raise RustExecutorClientError(
"Rust executor stream sync result must contain body bytes"
)
response_json = await self._aggregate_upstream_stream_sync_response(
body_bytes=rust_result.response_body_bytes,
provider_api_format=provider_api_format,
client_api_format=client_api_format,
provider_name=str(provider.name),
provider_type=str(getattr(provider, "provider_type", "") or "").lower(),
model=str(model or ""),
request_id=str(self.request_id or ""),
envelope=envelope,
)
response_metadata_result = self._extract_response_metadata(
response_json or {}
)
return response_json if isinstance(response_json, dict) else {}
response_json = rust_result.response_json or {}
if envelope:
response_json = envelope.unwrap_response(response_json)
envelope.postprocess_unwrapped_response(model=model, data=response_json)
response_metadata_result = self._extract_response_metadata(response_json)
return response_json if isinstance(response_json, dict) else {}
http_client = await HTTPClientPool.get_upstream_client(
delegate_cfg,
proxy_config=_effective_proxy,
@@ -218,11 +413,6 @@ class CliSyncMixin:
raise
else:
# Forced upstream streaming: aggregate SSE to a sync JSON response.
registry = get_format_converter_registry()
provider_parser = (
get_parser_for_format(provider_api_format) if provider_api_format else None
)
try:
_stream_args = await build_stream_kwargs_async(
delegate_cfg,
@@ -246,41 +436,16 @@ class CliSyncMixin:
)
stream_resp.raise_for_status()
byte_iter = stream_resp.aiter_bytes()
_provider_type = str(getattr(provider, "provider_type", "") or "").lower()
if (
_provider_type == "kiro"
and envelope
and envelope.force_stream_rewrite()
):
from src.services.provider.adapters.kiro.eventstream_rewriter import (
apply_kiro_stream_rewrite,
)
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
internal_resp = await aggregate_upstream_stream_to_internal_response(
byte_iter,
response_body = await stream_resp.aread()
response_json = await self._aggregate_upstream_stream_sync_response(
body_bytes=response_body,
provider_api_format=provider_api_format,
client_api_format=client_api_format,
provider_name=str(provider.name),
provider_type=str(getattr(provider, "provider_type", "") or "").lower(),
model=str(model or ""),
request_id=str(self.request_id or ""),
envelope=envelope,
provider_parser=provider_parser,
)
tgt_norm = (
registry.get_normalizer(client_api_format)
if client_api_format
else None
)
if tgt_norm is None:
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
response_json = tgt_norm.response_from_internal(
internal_resp,
requested_model=model,
)
response_json = response_json if isinstance(response_json, dict) else {}
+350 -100
View File
@@ -17,6 +17,7 @@
from __future__ import annotations
import asyncio
import codecs
import json
import time
import uuid
@@ -26,6 +27,7 @@ from typing import Any
import httpx
from src.config.settings import config
from src.core.api_format import (
CORE_REDACT_HEADERS,
merge_headers_with_protection,
@@ -655,6 +657,16 @@ class HttpRequestExecutor:
is_stream = True
try:
rust_result = await self._execute_via_rust(
request=request,
request_id=request_id,
is_stream=is_stream,
start_time=start_time,
effective_timeout=effective_timeout,
)
if rust_result is not None:
return rust_result
from src.services.proxy_node.resolver import build_proxy_client_kwargs
# 统一通过 build_proxy_client_kwargs 构建(支持 tunnel 模式 + 普通代理 + 系统默认回退)
@@ -758,112 +770,350 @@ class HttpRequestExecutor:
async with client.stream(
"POST", request.url, json=request.json_body, headers=request.headers
) as response:
headers = dict(response.headers)
if response.status_code != 200:
# 读取上限 16KB 用于 debug response_body
# error 字段仅取前 500 字符,避免日志/展示过长
error_body = ""
async for chunk in response.aiter_text():
error_body += chunk
if len(error_body) > 16384:
break
logger.debug(
"[{}] check_endpoint | stream error | {}",
request.api_format,
error_body[:500],
)
return {
"error": f"HTTP {response.status_code}: {error_body[:500]}",
"status_code": response.status_code,
"headers": headers,
"response_body": error_body,
}
# 收集 SSE 事件(兼容多种 API 格式)
final_response: dict[str, Any] = {}
collected_text = ""
async for line in response.aiter_lines():
if not line or not line.startswith("data:"):
continue
data_str = line[5:].strip()
if data_str == "[DONE]":
break
try:
event = json.loads(data_str)
if "/v1internal:" in (request.url or ""):
event = self._unwrap_gemini_cli_response_wrapper(event)
event_type = event.get("type", "")
# OpenAI Responses API 事件
if event_type == "response.output_text.delta":
delta = event.get("delta", "")
if isinstance(delta, str):
collected_text += delta
elif event_type == "response.completed":
final_response = event.get("response", {})
break
# OpenAI Chat Completions 格式
elif "choices" in event:
for choice in event.get("choices", []):
delta = choice.get("delta", {})
content = delta.get("content")
if content:
collected_text += content
if choice.get("finish_reason"):
final_response = event
break
# Claude Messages API 格式
elif event_type == "content_block_delta":
delta = event.get("delta", {})
text = delta.get("text", "")
if text:
collected_text += text
elif event_type == "message_stop":
break
# Gemini SSE 格式
elif "candidates" in event:
for candidate in event.get("candidates", []):
content = candidate.get("content", {})
for part in content.get("parts", []):
text = part.get("text", "")
if text:
collected_text += text
except json.JSONDecodeError:
continue
# 如果没有收到最终响应事件,构建一个基本响应
if not final_response:
final_response = {
"status": "completed",
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": collected_text}],
}
],
}
logger.debug(
"[{}] check_endpoint | stream completed | text_length={}",
request.api_format,
len(collected_text),
return await self._consume_stream_lines(
line_iter=response.aiter_lines(),
request=request,
status_code=response.status_code,
headers=dict(response.headers),
error_text_iter=response.aiter_text(),
)
return {"final_response": final_response, "headers": headers}
except Exception as e:
logger.warning("[{}] check_endpoint | stream error | {}", request.api_format, e)
return {"error": str(e), "status_code": 500, "headers": {}, "response_body": None}
async def _execute_via_rust(
self,
*,
request: EndpointCheckRequest,
request_id: str,
is_stream: bool,
start_time: float,
effective_timeout: float,
) -> EndpointCheckResult | None:
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
build_execution_plan_body,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
proxy_snapshot = await self._build_rust_proxy_snapshot(request.proxy_config)
plan = ExecutionPlan(
request_id=f"endpoint-check-{request_id}",
candidate_id=None,
provider_name=str(request.provider_name or request.api_format or ""),
provider_id=str(request.provider_id or ""),
endpoint_id="",
key_id=str(request.api_key_id or ""),
method="POST",
url=request.url,
headers=dict(request.headers),
body=build_execution_plan_body(request.json_body, content_type="application/json"),
stream=is_stream,
provider_api_format=str(request.api_format or ""),
client_api_format=str(request.api_format or ""),
model_name=str(request.model_name or ""),
content_type="application/json",
proxy=proxy_snapshot,
timeouts=ExecutionPlanTimeouts(
connect_ms=min(int(effective_timeout * 1000), 30_000),
read_ms=(None if is_stream else int(effective_timeout * 1000)),
write_ms=int(effective_timeout * 1000),
pool_ms=min(int(effective_timeout * 1000), 30_000),
total_ms=(None if is_stream else int(effective_timeout * 1000)),
),
)
try:
if is_stream:
rust_stream = await RustExecutorClient().execute_stream(plan)
try:
if rust_stream.status_code >= 400:
error_bytes = await self._read_limited_bytes(rust_stream.byte_iterator)
return await ErrorHandler.handle_error(
httpx.HTTPStatusError(
message=f"HTTP {rust_stream.status_code}",
request=httpx.Request("POST", request.url, headers=request.headers),
response=httpx.Response(
rust_stream.status_code,
request=httpx.Request(
"POST", request.url, headers=request.headers
),
headers=rust_stream.headers,
content=error_bytes,
),
),
request,
)
stream_result = await self._consume_stream_lines(
line_iter=self._aiter_lines_from_bytes(rust_stream.byte_iterator),
request=request,
status_code=rust_stream.status_code,
headers=dict(rust_stream.headers),
)
finally:
await rust_stream.response_ctx.__aexit__(None, None, None)
response_time_ms = int((time.time() - start_time) * 1000)
if stream_result.get("error"):
return EndpointCheckResult(
status_code=stream_result.get("status_code", 500),
headers=stream_result.get("headers", {}),
response_time_ms=response_time_ms,
request_id=request_id,
response_data=None,
error_message=stream_result.get("error"),
raw_response_body=stream_result.get("response_body"),
)
return EndpointCheckResult(
status_code=rust_stream.status_code,
headers=stream_result.get("headers", {}),
response_time_ms=response_time_ms,
request_id=request_id,
response_data=stream_result.get("final_response"),
raw_response_body=stream_result.get("final_response"),
)
result = await RustExecutorClient().execute_sync_json(plan)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError, ValueError) as exc:
logger.warning(
"[{}] endpoint check rust fallback | provider={} model={} error={}",
request.api_format,
request.provider_name,
request.model_name,
exc,
)
return None
response = self._build_httpx_response_from_rust(
method="POST",
url=request.url,
request_headers=request.headers,
status_code=result.status_code,
headers=result.headers,
response_json=result.response_json,
response_body_bytes=result.response_body_bytes,
)
if result.status_code >= 400:
return await ErrorHandler.handle_error(
httpx.HTTPStatusError(
message=f"HTTP {result.status_code}",
request=response.request,
response=response,
),
request,
)
response_time_ms = int((time.time() - start_time) * 1000)
response_data: dict[str, Any] | None
if isinstance(result.response_json, dict):
response_data = result.response_json
else:
try:
parsed_body = response.json()
response_data = parsed_body if isinstance(parsed_body, dict) else None
except Exception:
response_data = None
return EndpointCheckResult(
status_code=result.status_code,
headers=dict(result.headers),
response_time_ms=response_time_ms,
request_id=request_id,
response_data=response_data,
raw_response_body=response_data if response_data is not None else response.text,
)
async def _build_rust_proxy_snapshot(
self,
proxy_config: dict[str, Any] | None,
) -> Any:
from src.services.proxy_node.resolver import (
build_proxy_url_async,
get_system_proxy_config_async,
resolve_delegate_config_async,
resolve_proxy_info_async,
)
from src.services.request.executor_plan import ExecutionProxySnapshot
effective_proxy = proxy_config
if not effective_proxy or not effective_proxy.get("enabled", True):
effective_proxy = await get_system_proxy_config_async()
if not effective_proxy:
return None
try:
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
proxy_url: str | None = None
if effective_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
proxy_url = await build_proxy_url_async(effective_proxy)
proxy_info = await resolve_proxy_info_async(effective_proxy)
return ExecutionProxySnapshot.from_proxy_info(
proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if delegate_cfg and delegate_cfg.get("tunnel")
else None
),
)
except Exception as exc:
logger.warning("endpoint check proxy snapshot build failed: {}", exc)
return None
async def _consume_stream_lines(
self,
*,
line_iter: Any,
request: EndpointCheckRequest,
status_code: int,
headers: dict[str, str],
error_text_iter: Any | None = None,
) -> dict[str, Any]:
if status_code != 200:
error_body = ""
if error_text_iter is not None:
async for chunk in error_text_iter:
error_body += chunk
if len(error_body) > 16384:
break
logger.debug(
"[{}] check_endpoint | stream error | {}",
request.api_format,
error_body[:500],
)
return {
"error": f"HTTP {status_code}: {error_body[:500]}",
"status_code": status_code,
"headers": headers,
"response_body": error_body,
}
final_response: dict[str, Any] = {}
collected_text = ""
async for line in line_iter:
if not line or not line.startswith("data:"):
continue
data_str = line[5:].strip()
if data_str == "[DONE]":
break
try:
event = json.loads(data_str)
if "/v1internal:" in (request.url or ""):
event = self._unwrap_gemini_cli_response_wrapper(event)
event_type = event.get("type", "")
if event_type == "response.output_text.delta":
delta = event.get("delta", "")
if isinstance(delta, str):
collected_text += delta
elif event_type == "response.completed":
final_response = event.get("response", {})
break
elif "choices" in event:
for choice in event.get("choices", []):
delta = choice.get("delta", {})
content = delta.get("content")
if content:
collected_text += content
if choice.get("finish_reason"):
final_response = event
break
elif event_type == "content_block_delta":
delta = event.get("delta", {})
text = delta.get("text", "")
if text:
collected_text += text
elif event_type == "message_stop":
break
elif "candidates" in event:
for candidate in event.get("candidates", []):
content = candidate.get("content", {})
for part in content.get("parts", []):
text = part.get("text", "")
if text:
collected_text += text
except json.JSONDecodeError:
continue
if not final_response:
final_response = {
"status": "completed",
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": collected_text}],
}
],
}
logger.debug(
"[{}] check_endpoint | stream completed | text_length={}",
request.api_format,
len(collected_text),
)
return {"final_response": final_response, "headers": headers}
async def _aiter_lines_from_bytes(self, byte_iter: Any) -> Any:
decoder = codecs.getincrementaldecoder("utf-8")()
buffer = ""
async for chunk in byte_iter:
buffer += decoder.decode(chunk)
while "\n" in buffer:
line, buffer = buffer.split("\n", 1)
yield line.rstrip("\r")
buffer += decoder.decode(b"", final=True)
if buffer:
yield buffer.rstrip("\r")
async def _read_limited_bytes(self, byte_iter: Any, limit: int = 16_384) -> bytes:
chunks: list[bytes] = []
total = 0
async for chunk in byte_iter:
if not chunk:
continue
chunks.append(chunk)
total += len(chunk)
if total >= limit:
break
return b"".join(chunks)
@staticmethod
def _build_httpx_response_from_rust(
*,
method: str,
url: str,
request_headers: dict[str, str],
status_code: int,
headers: dict[str, str],
response_json: Any,
response_body_bytes: bytes | None,
) -> httpx.Response:
if response_json is not None:
content = json.dumps(response_json, ensure_ascii=False).encode("utf-8")
else:
content = response_body_bytes or b""
return httpx.Response(
status_code=status_code,
request=httpx.Request(method, url, headers=request_headers),
headers=headers,
content=content,
)
class UsageCalculator:
"""用量计算器 - 专门负责Token计数和费用计算"""
+105 -2
View File
@@ -6,10 +6,12 @@ Video Handler 基类
from __future__ import annotations
import json
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any
import httpx
from fastapi import HTTPException, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse
from sqlalchemy.orm import Session
@@ -28,8 +30,6 @@ from src.services.billing.rule_service import BillingRuleLookupResult
from src.services.scheduling.aware_scheduler import ProviderCandidate
if TYPE_CHECKING:
import httpx
from src.services.candidate.submit import SubmitOutcome
@@ -167,6 +167,109 @@ class VideoHandlerBase(ABC):
content={"error": fallback_payload},
)
async def _try_rust_sync_http_response(
self,
*,
method: str,
url: str,
headers: dict[str, str],
body: Any = None,
provider_name: str | None = None,
provider_id: str | None = None,
endpoint_id: str | None = None,
key_id: str | None = None,
provider_api_format: str | None = None,
client_api_format: str | None = None,
model_name: str | None = None,
content_type: str | None = None,
content_encoding: str | None = None,
proxy: Any = None,
tls_profile: str | None = None,
request_timeout_ms: int = 300_000,
connect_timeout_ms: int = 30_000,
pool_timeout_ms: int = 30_000,
log_label: str = "VideoRequest",
) -> httpx.Response | None:
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
build_execution_plan_body,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
request_headers = dict(headers)
if (
body is not None
and content_type
and not any(str(key).lower() == "content-type" for key in request_headers)
):
request_headers["content-type"] = content_type
plan = ExecutionPlan(
request_id=str(self.request_id or ""),
candidate_id=None,
provider_name=str(provider_name or ""),
provider_id=str(provider_id or ""),
endpoint_id=str(endpoint_id or ""),
key_id=str(key_id or ""),
method=str(method or "POST").upper(),
url=url,
headers=request_headers,
body=build_execution_plan_body(body, content_type=content_type),
stream=False,
provider_api_format=str(provider_api_format or self.FORMAT_ID),
client_api_format=str(client_api_format or self.FORMAT_ID),
model_name=str(model_name or ""),
content_type=content_type,
content_encoding=content_encoding,
proxy=proxy,
tls_profile=tls_profile,
timeouts=ExecutionPlanTimeouts(
connect_ms=connect_timeout_ms,
read_ms=request_timeout_ms,
write_ms=request_timeout_ms,
pool_ms=pool_timeout_ms,
total_ms=request_timeout_ms,
),
)
try:
rust_result = await RustExecutorClient().execute_sync_json(plan)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning(
"[{}] Rust executor unavailable request_id={} method={} url={}: {}",
log_label,
self.request_id,
method,
url,
sanitize_error_message(str(exc)),
)
return None
response_headers = dict(rust_result.headers)
if rust_result.response_json is not None:
response_headers.setdefault("content-type", "application/json")
response_body = json.dumps(rust_result.response_json, ensure_ascii=False).encode(
"utf-8"
)
elif rust_result.response_body_bytes is not None:
response_body = rust_result.response_body_bytes
else:
response_body = b""
return httpx.Response(
status_code=rust_result.status_code,
request=httpx.Request(str(method or "POST").upper(), url, headers=request_headers),
headers=response_headers,
content=response_body,
)
def _format_error_payload(self, error: dict[str, Any], status_code: int) -> dict[str, Any]:
"""
格式化错误负载,子类可重写以匹配特定 API 格式
+168
View File
@@ -183,6 +183,25 @@ class GeminiVeoHandler(VideoHandlerBase):
original_body=original_request_body,
)
rust_response = await self._try_rust_sync_http_response(
method="POST",
url=upstream_url,
headers=headers,
body=converted_body,
provider_name=str(candidate.provider.name),
provider_id=str(candidate.provider.id),
endpoint_id=str(endpoint.id),
key_id=str(_key.id),
provider_api_format=provider_format,
client_api_format=self.FORMAT_ID,
model_name=internal_request.model,
content_type=str(headers.get("content-type") or "").strip()
or "application/json",
log_label="GeminiVideoCreate",
)
if rust_response is not None:
return rust_response
client = await HTTPClientPool.get_default_client_async()
return await client.post(upstream_url, headers=headers, json=converted_body)
else:
@@ -206,6 +225,25 @@ class GeminiVeoHandler(VideoHandlerBase):
body=request_body,
original_body=original_request_body,
)
rust_response = await self._try_rust_sync_http_response(
method="POST",
url=upstream_url,
headers=headers,
body=request_body,
provider_name=str(candidate.provider.name),
provider_id=str(candidate.provider.id),
endpoint_id=str(endpoint.id),
key_id=str(_key.id),
provider_api_format=provider_format,
client_api_format=self.FORMAT_ID,
model_name=internal_request.model,
content_type=str(headers.get("content-type") or "").strip()
or "application/json",
log_label="GeminiVideoCreate",
)
if rust_response is not None:
return rust_response
client = await HTTPClientPool.get_default_client_async()
return await client.post(upstream_url, headers=headers, json=request_body)
@@ -470,6 +508,17 @@ class GeminiVeoHandler(VideoHandlerBase):
# 代理下载而非直接重定向,避免暴露上游存储 URL
# 使用 httpx 支持重定向(Gemini 视频 URL 会重定向到实际存储位置)
rust_response = await self._try_rust_download_stream(
url=task.video_url,
headers=download_headers,
task_id=str(task.id),
endpoint=endpoint,
key=key,
model_name=str(getattr(task, "model", "") or "") or None,
)
if rust_response is not None:
return rust_response
import httpx
try:
@@ -516,6 +565,125 @@ class GeminiVeoHandler(VideoHandlerBase):
media_type=response.headers.get("content-type", "video/mp4"),
)
async def _try_rust_download_stream(
self,
*,
url: str,
headers: dict[str, str],
task_id: str,
endpoint: ProviderEndpoint,
key: ProviderAPIKey,
model_name: str | None = None,
) -> Response | StreamingResponse | None:
import httpx
from src.services.proxy_node.resolver import (
build_proxy_url_async,
get_system_proxy_config_async,
resolve_delegate_config_async,
resolve_effective_proxy,
resolve_proxy_info_async,
)
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
ExecutionPlanTimeouts,
ExecutionProxySnapshot,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
effective_proxy = resolve_effective_proxy(
resolve_provider_proxy(endpoint=endpoint, key=key),
getattr(key, "proxy", None),
)
if not effective_proxy or not effective_proxy.get("enabled", True):
effective_proxy = await get_system_proxy_config_async()
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
proxy_url: str | None = None
if effective_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
proxy_url = await build_proxy_url_async(effective_proxy)
proxy_info = await resolve_proxy_info_async(effective_proxy)
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if delegate_cfg and delegate_cfg.get("tunnel")
else None
),
)
try:
rust_stream = await RustExecutorClient().execute_stream(
ExecutionPlan(
request_id=str(self.request_id or ""),
candidate_id=None,
provider_name="gemini",
provider_id=str(getattr(endpoint, "provider_id", "") or ""),
endpoint_id=str(getattr(endpoint, "id", "") or ""),
key_id=str(getattr(key, "id", "") or ""),
method="GET",
url=url,
headers=dict(headers),
body=ExecutionPlanBody(),
stream=True,
provider_api_format=self.FORMAT_ID,
client_api_format=self.FORMAT_ID,
model_name=str(model_name or ""),
proxy=proxy_snapshot,
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=300_000,
write_ms=300_000,
pool_ms=30_000,
total_ms=None,
),
)
)
except (RustExecutorClientError, httpx.HTTPError, ValueError) as exc:
logger.warning(
"[VideoDownload] Rust executor unavailable task={} url={}: {}",
task_id,
url,
sanitize_error_message(str(exc)),
)
return None
safe_headers = {
k: v for k, v in rust_stream.headers.items() if k.lower() not in HOP_BY_HOP_HEADERS
}
if rust_stream.status_code >= 400:
try:
async for _ in rust_stream.byte_iterator:
pass
finally:
await rust_stream.response_ctx.__aexit__(None, None, None)
raise HTTPException(status_code=rust_stream.status_code, detail="Upstream error")
async def _iter_bytes() -> AsyncIterator[bytes]:
try:
async for chunk in rust_stream.byte_iterator:
yield chunk
finally:
await rust_stream.response_ctx.__aexit__(None, None, None)
return StreamingResponse(
_iter_bytes(),
status_code=rust_stream.status_code,
headers=safe_headers,
media_type=safe_headers.get("content-type", "video/mp4"),
)
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
+211 -4
View File
@@ -43,6 +43,15 @@ from src.core.crypto import crypto_service
from src.core.logger import logger
from src.models.database import ApiKey, ProviderAPIKey, ProviderEndpoint, User, VideoTask
from src.services.billing.rule_service import BillingRuleLookupResult, BillingRuleService
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
ExecutionPlanTimeouts,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
from src.services.scheduling.aware_scheduler import ProviderCandidate
from src.services.usage.service import UsageService
@@ -180,6 +189,25 @@ class OpenAIVideoHandler(VideoHandlerBase):
original_body=original_request_body,
)
rust_response = await self._try_rust_sync_http_response(
method="POST",
url=upstream_url,
headers=headers,
body=converted_body,
provider_name=str(candidate.provider.name),
provider_id=str(candidate.provider.id),
endpoint_id=str(endpoint.id),
key_id=str(_provider_key.id),
provider_api_format=provider_format,
client_api_format=self.FORMAT_ID,
model_name=internal_request.model,
content_type=str(headers.get("content-type") or "").strip()
or "application/json",
log_label="OpenAIVideoCreate",
)
if rust_response is not None:
return rust_response
client = await HTTPClientPool.get_default_client_async()
return await client.post(upstream_url, headers=headers, json=converted_body)
else:
@@ -199,6 +227,25 @@ class OpenAIVideoHandler(VideoHandlerBase):
body=request_body,
original_body=original_request_body,
)
rust_response = await self._try_rust_sync_http_response(
method="POST",
url=upstream_url,
headers=headers,
body=request_body,
provider_name=str(candidate.provider.name),
provider_id=str(candidate.provider.id),
endpoint_id=str(endpoint.id),
key_id=str(_provider_key.id),
provider_api_format=provider_format,
client_api_format=self.FORMAT_ID,
model_name=internal_request.model,
content_type=str(headers.get("content-type") or "").strip()
or "application/json",
log_label="OpenAIVideoCreate",
)
if rust_response is not None:
return rust_response
client = await HTTPClientPool.get_default_client_async()
return await client.post(upstream_url, headers=headers, json=request_body)
@@ -478,8 +525,23 @@ class OpenAIVideoHandler(VideoHandlerBase):
)
headers = self._build_upstream_headers(original_headers, upstream_key, endpoint)
client = await HTTPClientPool.get_default_client_async()
response = await client.delete(upstream_url, headers=headers)
response = await self._try_rust_sync_http_response(
method="DELETE",
url=upstream_url,
headers=headers,
body=None,
provider_name="openai",
provider_id=str(getattr(endpoint, "provider_id", "") or ""),
endpoint_id=str(getattr(endpoint, "id", "") or ""),
key_id=str(getattr(key, "id", "") or ""),
provider_api_format=self.FORMAT_ID,
client_api_format=self.FORMAT_ID,
model_name=str(getattr(task, "model", "") or ""),
log_label="OpenAIVideoDelete",
)
if response is None:
client = await HTTPClientPool.get_default_client_async()
response = await client.delete(upstream_url, headers=headers)
if response.status_code >= 400 and response.status_code != 404:
# 404 表示上游已删除,不算错误
return self._build_error_response(response)
@@ -549,8 +611,24 @@ class OpenAIVideoHandler(VideoHandlerBase):
original_body=original_request_body,
)
client = await HTTPClientPool.get_default_client_async()
response = await client.post(upstream_url, headers=headers, json=request_body)
response = await self._try_rust_sync_http_response(
method="POST",
url=upstream_url,
headers=headers,
body=request_body,
provider_name="openai",
provider_id=str(getattr(endpoint, "provider_id", "") or ""),
endpoint_id=str(getattr(endpoint, "id", "") or ""),
key_id=str(getattr(key, "id", "") or ""),
provider_api_format=self.FORMAT_ID,
client_api_format=self.FORMAT_ID,
model_name=str(getattr(original_task, "model", "") or ""),
content_type=str(headers.get("content-type") or "").strip() or "application/json",
log_label="OpenAIVideoRemix",
)
if response is None:
client = await HTTPClientPool.get_default_client_async()
response = await client.post(upstream_url, headers=headers, json=request_body)
if response.status_code >= 400:
return self._build_error_response(response)
@@ -669,6 +747,16 @@ class OpenAIVideoHandler(VideoHandlerBase):
task_id,
task.video_url,
)
rust_response = await self._try_rust_download_stream(
url=task.video_url,
headers={},
task_id=task_id,
model_name=str(getattr(task, "model", "") or "") or None,
default_media_type="video/mp4",
default_error_message="Video not available",
)
if rust_response is not None:
return rust_response
return await self._proxy_direct_url(task.video_url, task_id)
if not task.external_task_id:
@@ -697,6 +785,20 @@ class OpenAIVideoHandler(VideoHandlerBase):
task_id,
task.external_task_id,
)
rust_response = await self._try_rust_download_stream(
url=upstream_url,
headers=headers,
task_id=task_id,
provider_id=str(getattr(endpoint, "provider_id", "") or "") or None,
endpoint_id=str(getattr(endpoint, "id", "") or "") or None,
key_id=str(getattr(key, "id", "") or "") or None,
model_name=str(getattr(task, "model", "") or "") or None,
default_media_type="application/octet-stream",
default_error_message="Upstream connection failed",
)
if rust_response is not None:
return rust_response
try:
# 使用 httpx 的 stream 方法并正确管理上下文
# 视频下载可能较大,设置 5 分钟超时
@@ -811,6 +913,111 @@ class OpenAIVideoHandler(VideoHandlerBase):
media_type=response.headers.get("content-type", "video/mp4"),
)
async def _try_rust_download_stream(
self,
*,
url: str,
headers: dict[str, str],
task_id: str,
provider_id: str | None = None,
endpoint_id: str | None = None,
key_id: str | None = None,
model_name: str | None = None,
default_media_type: str,
default_error_message: str,
) -> Response | StreamingResponse | None:
if config.executor_backend != "rust":
return None
plan = ExecutionPlan(
request_id=str(self.request_id or ""),
candidate_id=None,
provider_name="openai",
provider_id=str(provider_id or ""),
endpoint_id=str(endpoint_id or ""),
key_id=str(key_id or ""),
method="GET",
url=url,
headers=dict(headers),
body=ExecutionPlanBody(),
stream=True,
provider_api_format=self.FORMAT_ID,
client_api_format=self.FORMAT_ID,
model_name=str(model_name or "") or "",
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=300_000,
write_ms=300_000,
pool_ms=30_000,
total_ms=None,
),
)
try:
rust_stream = await RustExecutorClient().execute_stream(plan)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning(
"[VideoDownload] Rust executor unavailable task={} url={}: {}",
task_id,
url,
sanitize_error_message(str(exc)),
)
return None
safe_headers = {
k: v for k, v in rust_stream.headers.items() if k.lower() not in HOP_BY_HOP_HEADERS
}
if rust_stream.status_code >= 400:
error_chunks: list[bytes] = []
try:
async for chunk in rust_stream.byte_iterator:
if chunk:
error_chunks.append(chunk)
if sum(len(item) for item in error_chunks) >= 16_384:
break
finally:
await rust_stream.response_ctx.__aexit__(None, None, None)
error_body = b"".join(error_chunks)[:16_384]
content_type = str(safe_headers.get("content-type") or "").lower()
if "application/json" in content_type:
try:
data = json.loads(error_body)
if isinstance(data, dict) and isinstance(data.get("error"), dict):
if "message" in data["error"]:
data["error"]["message"] = sanitize_error_message(
str(data["error"]["message"])
)
return JSONResponse(status_code=rust_stream.status_code, content=data)
except json.JSONDecodeError:
pass
message = sanitize_error_message(error_body.decode(errors="ignore"))
return JSONResponse(
status_code=rust_stream.status_code,
content={
"error": {
"type": "upstream_error",
"message": message or default_error_message,
}
},
)
async def _iter_bytes() -> AsyncIterator[bytes]:
try:
async for chunk in rust_stream.byte_iterator:
yield chunk
finally:
await rust_stream.response_ctx.__aexit__(None, None, None)
return StreamingResponse(
_iter_bytes(),
status_code=rust_stream.status_code,
headers=safe_headers,
media_type=safe_headers.get("content-type", default_media_type),
)
def _build_upstream_url(self, base_url: str | None, suffix: str | None = None) -> str:
base = (base_url or self.DEFAULT_BASE_URL).rstrip("/")
if base.endswith("/v1"):
+8 -1
View File
@@ -1,3 +1,10 @@
from .hub import router
from fastapi import APIRouter
from .gateway import router as gateway_router
from .hub import router as hub_router
router = APIRouter()
router.include_router(hub_router)
router.include_router(gateway_router)
__all__ = ["router"]
+14
View File
@@ -0,0 +1,14 @@
from __future__ import annotations
import ipaddress
from fastapi import HTTPException, Request
def ensure_loopback(request: Request) -> None:
host = request.client.host if request.client else ""
try:
if not ipaddress.ip_address(host).is_loopback:
raise ValueError(host)
except ValueError as exc:
raise HTTPException(status_code=403, detail="loopback access only") from exc
File diff suppressed because it is too large Load Diff
+4 -12
View File
@@ -1,7 +1,6 @@
from __future__ import annotations
import asyncio
import ipaddress
from typing import Any
from fastapi import APIRouter, HTTPException, Request
@@ -9,6 +8,8 @@ from pydantic import BaseModel, Field
from src.services.proxy_node.service import ProxyNodeService, build_heartbeat_ack
from .common import ensure_loopback
router = APIRouter(prefix="/api/internal/hub", tags=["Internal - Hub"], include_in_schema=False)
@@ -31,18 +32,9 @@ class HubNodeStatusRequest(BaseModel):
conn_count: int = Field(0, ge=0)
def _ensure_loopback(request: Request) -> None:
host = request.client.host if request.client else ""
try:
if not ipaddress.ip_address(host).is_loopback:
raise ValueError(host)
except ValueError as exc:
raise HTTPException(status_code=403, detail="loopback access only") from exc
@router.post("/heartbeat")
async def hub_heartbeat(request: Request, payload: HubHeartbeatRequest) -> dict[str, Any]:
_ensure_loopback(request)
ensure_loopback(request)
def _sync_apply() -> dict[str, Any]:
from src.database import create_session
@@ -74,7 +66,7 @@ async def hub_heartbeat(request: Request, payload: HubHeartbeatRequest) -> dict[
@router.post("/node-status")
async def hub_node_status(request: Request, payload: HubNodeStatusRequest) -> dict[str, Any]:
_ensure_loopback(request)
ensure_loopback(request)
def _sync_apply() -> dict[str, Any]:
from src.database import create_session
+442 -105
View File
@@ -21,12 +21,15 @@ https://ai.google.dev/api/files
from __future__ import annotations
import json
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple
from typing import Any
from urllib.parse import urlencode
from uuid import uuid4
from fastapi import APIRouter, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from fastapi.responses import JSONResponse, StreamingResponse
from sqlalchemy.orm import Session
from src.clients.http_client import HTTPClientPool
@@ -37,6 +40,7 @@ from src.database import create_session
from src.models.database import ApiKey, GlobalModel, Model, Provider, ProviderEndpoint, User
from src.services.auth.service import AuthService
from src.services.gemini_files_mapping import delete_file_key_mapping, store_file_key_mapping
from src.services.provider.provider_context import resolve_provider_proxy
from src.services.provider.transport import redact_url_for_log
from src.services.scheduling.aware_scheduler import CacheAwareScheduler, ProviderCandidate
from src.services.usage.service import UsageService
@@ -50,6 +54,13 @@ class UpstreamContext:
base_url: str
file_key_id: str
user_id: str
provider_id: str
endpoint_id: str
provider_proxy: dict[str, Any] | None = None
key_proxy: dict[str, Any] | None = None
proxy_config: dict[str, Any] | None = None
delegate_config: dict[str, Any] | None = None
proxy_snapshot: Any = None
router = APIRouter(tags=["Gemini Files API"])
@@ -252,7 +263,7 @@ async def _select_provider_candidate(
async def _resolve_upstream_context(
request: Request,
db: Session,
) -> tuple[str, str, str, str]:
) -> UpstreamContext:
"""
解析上游 Key 与 Base URL(需要外部提供 db session
@@ -263,7 +274,7 @@ async def _resolve_upstream_context(
db: 数据库会话
Returns:
(upstream_key, base_url, key_id, user_id)
UpstreamContext
"""
client_key = _extract_gemini_api_key(request)
@@ -341,7 +352,53 @@ async def _resolve_upstream_context(
)
base_url = candidate.endpoint.base_url or GEMINI_FILES_BASE_URL
return upstream_key, base_url, str(candidate.key.id), str(user.id)
return UpstreamContext(
upstream_key,
base_url,
str(candidate.key.id),
str(user.id),
str(candidate.provider.id),
str(candidate.endpoint.id),
provider_proxy=resolve_provider_proxy(endpoint=candidate.endpoint, key=candidate.key),
key_proxy=(
candidate.key.proxy if isinstance(getattr(candidate.key, "proxy", None), dict) else None
),
)
async def _enrich_upstream_context_proxy(ctx: UpstreamContext) -> UpstreamContext:
from src.services.proxy_node.resolver import (
build_proxy_url_async,
get_system_proxy_config_async,
resolve_delegate_config_async,
resolve_effective_proxy,
resolve_proxy_info_async,
)
from src.services.request.executor_plan import ExecutionProxySnapshot
effective_proxy = resolve_effective_proxy(ctx.provider_proxy, ctx.key_proxy)
if not effective_proxy or not effective_proxy.get("enabled", True):
effective_proxy = await get_system_proxy_config_async()
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
proxy_url: str | None = None
if effective_proxy and not is_tunnel_delegate:
proxy_url = await build_proxy_url_async(effective_proxy)
proxy_info = await resolve_proxy_info_async(effective_proxy)
ctx.proxy_config = effective_proxy
ctx.delegate_config = delegate_cfg
ctx.proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if is_tunnel_delegate else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None if is_tunnel_delegate else None
),
)
return ctx
async def _resolve_upstream_context_standalone(request: Request) -> UpstreamContext:
@@ -357,13 +414,170 @@ async def _resolve_upstream_context_standalone(request: Request) -> UpstreamCont
UpstreamContext: 包含所有必要信息的上下文对象
"""
with create_session() as db:
upstream_key, base_url, file_key_id, user_id = await _resolve_upstream_context(request, db)
return UpstreamContext(
upstream_key=upstream_key,
base_url=base_url,
file_key_id=file_key_id,
user_id=user_id,
ctx = await _resolve_upstream_context(request, db)
return await _enrich_upstream_context_proxy(ctx)
def _build_response_headers(headers: dict[str, str]) -> dict[str, str]:
response_headers = {}
hop_by_hop = {"connection", "keep-alive", "transfer-encoding", "upgrade"}
for name, value in headers.items():
if name.lower() not in hop_by_hop:
response_headers[name] = value
return response_headers
async def _maybe_store_file_mapping_from_payload(
*,
status_code: int,
headers: dict[str, str],
content_bytes: bytes,
file_key_id: str | None = None,
user_id: str | None = None,
) -> None:
if (
not file_key_id
or status_code >= 300
or not headers.get("content-type", "").startswith("application/json")
):
return
try:
payload = json.loads(content_bytes)
file_name = None
file_obj = None
if isinstance(payload, dict):
file_name = payload.get("name")
file_obj = payload
if not file_name and isinstance(payload.get("file"), dict):
file_name = payload["file"].get("name")
file_obj = payload["file"]
if file_name and file_obj:
display_name = file_obj.get("displayName") or file_obj.get("display_name")
mime_type = file_obj.get("mimeType") or file_obj.get("mime_type")
await store_file_key_mapping(
file_name,
file_key_id,
user_id=user_id,
display_name=display_name,
mime_type=mime_type,
)
logger.debug(f"Gemini file→key 映射已存储: {file_name} → key_id={file_key_id}")
files_list = payload.get("files")
if isinstance(files_list, list):
mapped_count = 0
for item in files_list:
if isinstance(item, dict) and item.get("name"):
item_display_name = item.get("displayName") or item.get("display_name")
item_mime_type = item.get("mimeType") or item.get("mime_type")
await store_file_key_mapping(
item["name"],
file_key_id,
user_id=user_id,
display_name=item_display_name,
mime_type=item_mime_type,
)
mapped_count += 1
if mapped_count > 0:
logger.debug(
"Gemini list_files 批量映射已存储: {} 个文件 → key_id={}",
mapped_count,
file_key_id,
)
except (ValueError, KeyError) as e:
logger.debug("Failed to store Gemini file mapping: {}", e)
async def _try_rust_sync_proxy_request(
method: str,
upstream_url: str,
headers: dict[str, str],
*,
content: bytes | None = None,
json_body: dict[str, Any] | None = None,
file_key_id: str | None = None,
user_id: str | None = None,
provider_id: str = "",
endpoint_id: str = "",
proxy: Any = None,
) -> Response | None:
from src.config.settings import config
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
build_execution_plan_body,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
request_headers = dict(headers)
content_type = str(request_headers.get("content-type") or "").strip() or None
request_body = json_body if json_body is not None else content
try:
result = await RustExecutorClient().execute_sync_json(
ExecutionPlan(
request_id=f"gemini-files-{uuid4().hex}",
candidate_id=None,
provider_name="gemini",
provider_id=provider_id,
endpoint_id=endpoint_id,
key_id=str(file_key_id or ""),
method=method.upper(),
url=upstream_url,
headers=request_headers,
body=build_execution_plan_body(request_body, content_type=content_type),
stream=False,
provider_api_format="gemini:files",
client_api_format="gemini:files",
model_name="gemini-files",
proxy=proxy,
content_type=content_type,
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=300_000,
write_ms=300_000,
pool_ms=30_000,
total_ms=300_000,
),
)
)
except (RustExecutorClientError, HTTPException, json.JSONDecodeError, ValueError) as exc:
logger.warning("Gemini Files Rust proxy unavailable: {}", redact_url_for_log(str(exc)))
return None
response_headers = _build_response_headers(dict(result.headers))
if result.response_json is not None:
response_headers.setdefault("content-type", "application/json")
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
elif result.response_body_bytes is not None:
response_body = result.response_body_bytes
else:
response_body = b""
await _maybe_store_file_mapping_from_payload(
status_code=result.status_code,
headers=response_headers,
content_bytes=response_body,
file_key_id=file_key_id,
user_id=user_id,
)
return Response(
content=response_body,
status_code=result.status_code,
headers=response_headers,
media_type=response_headers.get("content-type", "application/json"),
)
async def _proxy_request(
@@ -374,6 +588,11 @@ async def _proxy_request(
json_body: dict[str, Any] | None = None,
file_key_id: str | None = None,
user_id: str | None = None,
provider_id: str = "",
endpoint_id: str = "",
proxy: Any = None,
proxy_config: dict[str, Any] | None = None,
delegate_config: dict[str, Any] | None = None,
) -> Response:
"""
代理请求到上游 Gemini API
@@ -390,7 +609,25 @@ async def _proxy_request(
Returns:
FastAPI Response 对象
"""
client = await HTTPClientPool.get_default_client_async()
rust_response = await _try_rust_sync_proxy_request(
method,
upstream_url,
headers,
content=content,
json_body=json_body,
file_key_id=file_key_id,
user_id=user_id,
provider_id=provider_id,
endpoint_id=endpoint_id,
proxy=proxy,
)
if rust_response is not None:
return rust_response
client = await HTTPClientPool.get_upstream_client(
delegate_config,
proxy_config=proxy_config,
)
try:
if method.upper() == "GET":
@@ -408,73 +645,14 @@ async def _proxy_request(
raise HTTPException(status_code=405, detail="Method not allowed")
# 构建响应头(排除 hop-by-hop 头部)
response_headers = {}
hop_by_hop = {"connection", "keep-alive", "transfer-encoding", "upgrade"}
for name, value in response.headers.items():
if name.lower() not in hop_by_hop:
response_headers[name] = value
if (
file_key_id
and response.status_code < 300
and response.headers.get("content-type", "").startswith("application/json")
):
try:
payload = response.json()
file_name = None
file_obj = None
if isinstance(payload, dict):
# 单文件上传响应
file_name = payload.get("name")
file_obj = payload
# 嵌套格式:{"file": {...}}
if not file_name and isinstance(payload.get("file"), dict):
file_name = payload["file"].get("name")
file_obj = payload["file"]
if file_name and file_obj:
display_name = file_obj.get("displayName") or file_obj.get("display_name")
mime_type = file_obj.get("mimeType") or file_obj.get("mime_type")
await store_file_key_mapping(
file_name,
file_key_id,
user_id=user_id,
display_name=display_name,
mime_type=mime_type,
)
logger.debug(
f"Gemini file→key 映射已存储: {file_name} → key_id={file_key_id}"
)
# 为 list_files 响应中的所有文件建立映射
# 这是正确的:Gemini API 按 Key 隔离文件,返回的文件必然属于当前 Key
files_list = payload.get("files")
if isinstance(files_list, list):
mapped_count = 0
for item in files_list:
if isinstance(item, dict) and item.get("name"):
item_display_name = item.get("displayName") or item.get(
"display_name"
)
item_mime_type = item.get("mimeType") or item.get("mime_type")
await store_file_key_mapping(
item["name"],
file_key_id,
user_id=user_id,
display_name=item_display_name,
mime_type=item_mime_type,
)
mapped_count += 1
if mapped_count > 0:
logger.debug(
"Gemini list_files 批量映射已存储: {} 个文件 → key_id={}",
mapped_count,
file_key_id,
)
except (ValueError, KeyError) as e:
logger.debug("Failed to store Gemini file mapping: {}", e)
response_headers = _build_response_headers(dict(response.headers))
await _maybe_store_file_mapping_from_payload(
status_code=response.status_code,
headers=response_headers,
content_bytes=response.content,
file_key_id=file_key_id,
user_id=user_id,
)
return Response(
content=response.content,
@@ -560,6 +738,11 @@ async def upload_file(
content=body,
file_key_id=ctx.file_key_id,
user_id=ctx.user_id,
provider_id=ctx.provider_id,
endpoint_id=ctx.endpoint_id,
proxy=ctx.proxy_snapshot,
proxy_config=ctx.proxy_config,
delegate_config=ctx.delegate_config,
)
@@ -624,7 +807,16 @@ async def list_files(
logger.debug("Gemini Files list proxy: GET {}", redact_url_for_log(upstream_url))
return await _proxy_request(
"GET", upstream_url, headers, file_key_id=ctx.file_key_id, user_id=ctx.user_id
"GET",
upstream_url,
headers,
file_key_id=ctx.file_key_id,
user_id=ctx.user_id,
provider_id=ctx.provider_id,
endpoint_id=ctx.endpoint_id,
proxy=ctx.proxy_snapshot,
proxy_config=ctx.proxy_config,
delegate_config=ctx.delegate_config,
)
@@ -716,9 +908,19 @@ async def download_file(
优化:HTTP 下载期间不持有数据库连接
"""
import httpx
from fastapi import HTTPException
from fastapi.responses import JSONResponse, Response
from src.config.settings import config
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
ExecutionPlanTimeouts,
ExecutionProxySnapshot,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
# ========== 阶段 1:数据库操作(短暂持有连接)==========
client_key = _extract_gemini_api_key(request)
@@ -730,6 +932,8 @@ async def download_file(
},
)
regular_file_ctx: UpstreamContext | None = None
# 在数据库会话内完成所有查询
with create_session() as db:
auth_result = AuthService.authenticate_api_key(db, client_key)
@@ -766,12 +970,13 @@ async def download_file(
},
)
upstream_url = video_url
file_key_id = ""
provider_id = ""
endpoint_id = ""
else:
# 普通文件下载:透传到 Gemini
try:
upstream_key, base_url, _file_key_id, _user_id = await _resolve_upstream_context(
request, db
)
regular_file_ctx = await _resolve_upstream_context(request, db)
except HTTPException:
raise HTTPException(
status_code=404,
@@ -783,26 +988,145 @@ async def download_file(
}
},
)
file_name = f"files/{file_id}" if not file_id.startswith("files/") else file_id
upstream_url = _build_upstream_url(
base_url,
f"/v1beta/{file_name}:download",
dict(request.query_params),
)
if regular_file_ctx is not None:
ctx = await _enrich_upstream_context_proxy(regular_file_ctx)
upstream_key = ctx.upstream_key
file_key_id = ctx.file_key_id
provider_id = ctx.provider_id
endpoint_id = ctx.endpoint_id
file_name = f"files/{file_id}" if not file_id.startswith("files/") else file_id
upstream_url = _build_upstream_url(
ctx.base_url,
f"/v1beta/{file_name}:download",
dict(request.query_params),
)
# ========== 阶段 2HTTP 下载(不持有数据库连接)==========
headers = _build_upstream_headers(dict(request.headers), upstream_key)
logger.debug("Gemini Files download proxy: GET {}", redact_url_for_log(upstream_url))
proxy_snapshot = None
proxy_config = None
delegate_config = None
if regular_file_ctx is not None:
proxy_snapshot = ctx.proxy_snapshot
proxy_config = ctx.proxy_config
delegate_config = ctx.delegate_config
if config.executor_backend == "rust":
try:
if proxy_snapshot is None and file_id.startswith("aev_"):
from src.services.proxy_node.resolver import (
build_proxy_url_async,
get_system_proxy_config_async,
resolve_delegate_config_async,
resolve_proxy_info_async,
)
system_proxy = await get_system_proxy_config_async()
delegate_cfg = await resolve_delegate_config_async(system_proxy)
proxy_config = system_proxy
delegate_config = delegate_cfg
proxy_url: str | None = None
if system_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
proxy_url = await build_proxy_url_async(system_proxy)
proxy_info = await resolve_proxy_info_async(system_proxy)
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if delegate_cfg and delegate_cfg.get("tunnel")
else None
),
)
except Exception as exc:
logger.warning("Gemini Files download proxy snapshot build failed: {}", exc)
try:
rust_stream = await RustExecutorClient().execute_stream(
ExecutionPlan(
request_id=f"gemini-files-download-{uuid4().hex}",
candidate_id=None,
provider_name="gemini",
provider_id=provider_id,
endpoint_id=endpoint_id,
key_id=str(file_key_id or ""),
method="GET",
url=upstream_url,
headers=headers,
body=ExecutionPlanBody(),
stream=True,
provider_api_format="gemini:files",
client_api_format="gemini:files",
model_name="gemini-files",
proxy=proxy_snapshot,
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=300_000,
write_ms=300_000,
pool_ms=30_000,
total_ms=None,
),
)
)
except (RustExecutorClientError, ValueError) as exc:
logger.warning("Gemini Files Rust download unavailable: {}", exc)
else:
safe_headers = _build_response_headers(dict(rust_stream.headers))
if rust_stream.status_code >= 400:
error_chunks: list[bytes] = []
try:
async for chunk in rust_stream.byte_iterator:
if chunk:
error_chunks.append(chunk)
if sum(len(item) for item in error_chunks) >= 16_384:
break
finally:
await rust_stream.response_ctx.__aexit__(None, None, None)
raw_error = b"".join(error_chunks)
if safe_headers.get("content-type", "").startswith("application/json"):
try:
return JSONResponse(
content=json.loads(raw_error.decode("utf-8")),
status_code=rust_stream.status_code,
)
except Exception:
pass
return JSONResponse(
content={"error": raw_error.decode("utf-8", errors="replace")},
status_code=rust_stream.status_code,
)
async def _iter_bytes() -> AsyncIterator[bytes]:
try:
async for chunk in rust_stream.byte_iterator:
yield chunk
finally:
await rust_stream.response_ctx.__aexit__(None, None, None)
return StreamingResponse(
_iter_bytes(),
status_code=rust_stream.status_code,
headers=safe_headers,
media_type=safe_headers.get("content-type", "application/octet-stream"),
)
# 使用 follow_redirects=True 跟随重定向(Gemini 文件下载会重定向)
try:
from src.services.proxy_node.resolver import build_proxy_client_kwargs
async with httpx.AsyncClient(
**build_proxy_client_kwargs(timeout=httpx.Timeout(300.0), follow_redirects=True)
) as client:
response = await client.get(upstream_url, headers=headers)
client = await HTTPClientPool.get_upstream_client(
delegate_config,
proxy_config=proxy_config,
)
response = await client.get(
upstream_url,
headers=headers,
timeout=300.0,
)
except Exception as exc:
logger.error("Gemini Files download failed: {}", exc)
raise HTTPException(status_code=502, detail="Failed to download file")
@@ -818,15 +1142,10 @@ async def download_file(
content = {"error": response.text}
return JSONResponse(content=content, status_code=response.status_code)
# 返回文件内容
return Response(
content=response.content,
status_code=response.status_code,
headers={
k: v
for k, v in response.headers.items()
if k.lower() not in {"transfer-encoding", "connection", "keep-alive"}
},
headers=_build_response_headers(dict(response.headers)),
media_type=response.headers.get("content-type", "application/octet-stream"),
)
@@ -887,7 +1206,16 @@ async def get_file(
logger.debug("Gemini Files get proxy: GET {}", redact_url_for_log(upstream_url))
return await _proxy_request(
"GET", upstream_url, headers, file_key_id=ctx.file_key_id, user_id=ctx.user_id
"GET",
upstream_url,
headers,
file_key_id=ctx.file_key_id,
user_id=ctx.user_id,
provider_id=ctx.provider_id,
endpoint_id=ctx.endpoint_id,
proxy=ctx.proxy_snapshot,
proxy_config=ctx.proxy_config,
delegate_config=ctx.delegate_config,
)
@@ -933,7 +1261,16 @@ async def delete_file(
logger.debug("Gemini Files delete proxy: DELETE {}", redact_url_for_log(upstream_url))
response = await _proxy_request("DELETE", upstream_url, headers)
response = await _proxy_request(
"DELETE",
upstream_url,
headers,
provider_id=ctx.provider_id,
endpoint_id=ctx.endpoint_id,
proxy=ctx.proxy_snapshot,
proxy_config=ctx.proxy_config,
delegate_config=ctx.delegate_config,
)
if response.status_code < 300:
await delete_file_key_mapping(file_name)
else:
+154 -4
View File
@@ -19,13 +19,15 @@ from src.api.handlers.base.request_builder import (
build_test_request_body,
get_provider_auth,
)
from src.clients.http_client import HTTPClientPool
from src.clients.redis_client import get_redis_client
from src.config.settings import config
from src.core.logger import logger
from src.database import get_db
from src.database.database import get_pool_status
from src.models.database import GlobalModel, Model, Provider, ProviderAPIKey, ProviderEndpoint
from src.services.provider.provider_context import resolve_provider_proxy
from src.services.provider.transport import build_provider_url
from src.utils.ssl_utils import get_ssl_context
router = APIRouter(tags=["System Catalog"])
@@ -98,6 +100,132 @@ def _select_provider(db: Session, provider_name: str | None) -> Provider | None:
return query.order_by(Provider.provider_priority.asc()).first()
async def _build_test_connection_transport_context(
endpoint: ProviderEndpoint,
key: ProviderAPIKey,
) -> tuple[dict[str, Any] | None, dict[str, Any] | None, Any]:
from src.services.proxy_node.resolver import (
build_proxy_url_async,
get_system_proxy_config_async,
resolve_delegate_config_async,
resolve_effective_proxy,
resolve_proxy_info_async,
)
from src.services.request.executor_plan import ExecutionProxySnapshot
try:
effective_proxy = resolve_effective_proxy(
resolve_provider_proxy(endpoint=endpoint, key=key),
getattr(key, "proxy", None),
)
if not effective_proxy or not effective_proxy.get("enabled", True):
effective_proxy = await get_system_proxy_config_async()
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
proxy_url: str | None = None
if effective_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
proxy_url = await build_proxy_url_async(effective_proxy)
proxy_info = await resolve_proxy_info_async(effective_proxy)
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if delegate_cfg and delegate_cfg.get("tunnel")
else None
),
)
return effective_proxy, delegate_cfg, proxy_snapshot
except Exception as exc:
logger.warning(
"Failed to build test-connection transport context endpoint={} key={}: {}",
getattr(endpoint, "id", None),
getattr(key, "id", None),
exc,
)
return None, None, None
async def _try_rust_test_connection_response(
*,
request_id: str,
url: str,
headers: dict[str, str],
body: dict[str, Any],
provider_name: str,
provider_id: str | None,
endpoint_id: str | None,
key_id: str | None,
api_format: str,
model_name: str,
proxy_snapshot: Any,
) -> httpx.Response | None:
import json
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
build_execution_plan_body,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
try:
result = await RustExecutorClient().execute_sync_json(
ExecutionPlan(
request_id=request_id,
candidate_id=None,
provider_name=provider_name,
provider_id=str(provider_id or ""),
endpoint_id=str(endpoint_id or ""),
key_id=str(key_id or ""),
method="POST",
url=url,
headers=dict(headers),
body=build_execution_plan_body(body, content_type="application/json"),
stream=False,
provider_api_format=api_format,
client_api_format=api_format,
model_name=model_name,
content_type="application/json",
proxy=proxy_snapshot,
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=30_000,
write_ms=30_000,
pool_ms=30_000,
total_ms=30_000,
),
)
)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning("Rust test-connection fallback url={}: {}", url, exc)
return None
response_headers = dict(result.headers)
if result.response_json is not None:
response_headers.setdefault("content-type", "application/json")
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
elif result.response_body_bytes is not None:
response_body = result.response_body_bytes
else:
response_body = b""
return httpx.Response(
status_code=result.status_code,
request=httpx.Request("POST", url, headers=headers),
headers=response_headers,
content=response_body,
)
# ============== 端点 ==============
@@ -375,11 +503,33 @@ async def test_connection(
key=key,
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
)
proxy_config, delegate_cfg, proxy_snapshot = await _build_test_connection_transport_context(
endpoint,
key,
)
async with httpx.AsyncClient(timeout=30.0, verify=get_ssl_context()) as client:
resp = await _try_rust_test_connection_response(
request_id=f"test-connection:{selected_provider.id}:{model}",
url=url,
headers=provider_headers,
body=provider_payload,
provider_name=selected_provider.name,
provider_id=getattr(selected_provider, "id", None),
endpoint_id=getattr(endpoint, "id", None),
key_id=getattr(key, "id", None),
api_format=format_value,
model_name=model,
proxy_snapshot=proxy_snapshot,
)
if resp is None:
client = await HTTPClientPool.get_upstream_client(
delegate_cfg,
proxy_config=proxy_config,
)
resp = await client.post(url, json=provider_payload, headers=provider_headers)
resp.raise_for_status()
response = resp.json()
resp.raise_for_status()
response = resp.json()
return {
"status": "success",
+14
View File
@@ -192,6 +192,20 @@ class Config:
# HTTP_REQUEST_TIMEOUT: 非流式请求整体超时(秒),默认 300 秒
self.http_request_timeout = float(os.getenv("HTTP_REQUEST_TIMEOUT", "300.0"))
# 内部 executor 配置
# EXECUTOR_BACKEND:
# - rust: 当前 pioneer 分支默认,优先将可序列化的执行计划转发给 Rust executor
# - python: 显式回退到现有 httpx 路径
self.executor_backend = os.getenv("EXECUTOR_BACKEND", "rust").strip().lower()
self.executor_transport = os.getenv("EXECUTOR_TRANSPORT", "unix_socket").strip().lower()
self.executor_socket_path = os.getenv(
"EXECUTOR_SOCKET_PATH", "/tmp/aether-executor.sock"
).strip()
self.executor_base_url = os.getenv("EXECUTOR_BASE_URL", "http://127.0.0.1:5219").strip()
self.executor_request_timeout = float(
os.getenv("EXECUTOR_REQUEST_TIMEOUT", str(self.http_request_timeout))
)
# HTTP 连接池配置
# HTTP_MAX_CONNECTIONS: 最大连接数,影响并发能力
# - 每个连接占用一个 socket,过多会耗尽系统资源
@@ -13,6 +13,7 @@ import httpx
from sqlalchemy.orm import Session
from src.core.crypto import crypto_service
from src.core.logger import logger
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
from src.services.provider.auth import get_provider_auth
from src.services.provider.oauth_token import looks_like_token_invalidated
@@ -70,6 +71,138 @@ def _extract_error_message_from_response(response: httpx.Response) -> str:
return text[:300] if text else ""
async def _build_codex_proxy_snapshot(
effective_proxy: dict[str, Any] | None,
) -> Any:
from importlib import import_module
from src.services.request.executor_plan import ExecutionProxySnapshot
resolver_module = import_module("src.services.proxy_node.resolver")
build_proxy_url_async = getattr(resolver_module, "build_proxy_url_async", None)
get_system_proxy_config_async = getattr(resolver_module, "get_system_proxy_config_async", None)
resolve_delegate_config_async = getattr(resolver_module, "resolve_delegate_config_async", None)
resolve_proxy_info_async = getattr(resolver_module, "resolve_proxy_info_async", None)
proxy_config = effective_proxy
if (not proxy_config or not proxy_config.get("enabled", True)) and callable(
get_system_proxy_config_async
):
proxy_config = await get_system_proxy_config_async()
if not proxy_config:
return None
try:
delegate_cfg = (
await resolve_delegate_config_async(proxy_config)
if callable(resolve_delegate_config_async)
else None
)
proxy_url: str | None = None
if (
proxy_config
and not (delegate_cfg and delegate_cfg.get("tunnel"))
and callable(build_proxy_url_async)
):
proxy_url = await build_proxy_url_async(proxy_config)
proxy_info = (
await resolve_proxy_info_async(proxy_config)
if callable(resolve_proxy_info_async)
else {"url": proxy_url}
)
return ExecutionProxySnapshot.from_proxy_info(
proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if delegate_cfg and delegate_cfg.get("tunnel")
else None
),
)
except Exception as exc:
logger.warning("Codex quota proxy snapshot build failed: {}", exc)
return None
async def _try_rust_codex_quota_response(
*,
key: ProviderAPIKey,
provider: Provider,
endpoint: ProviderEndpoint,
url: str,
headers: dict[str, str],
proxy_snapshot: Any,
) -> httpx.Response | None:
from src.config import config
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
ExecutionPlanTimeouts,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
try:
result = await RustExecutorClient().execute_sync_json(
ExecutionPlan(
request_id=f"codex-quota:{key.id}",
candidate_id=None,
provider_name="codex",
provider_id=str(getattr(provider, "id", "") or ""),
endpoint_id=str(getattr(endpoint, "id", "") or ""),
key_id=str(getattr(key, "id", "") or ""),
method="GET",
url=url,
headers=dict(headers),
body=ExecutionPlanBody(),
stream=False,
provider_api_format="openai:cli",
client_api_format="openai:cli",
model_name="codex-wham-usage",
proxy=proxy_snapshot,
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=30_000,
write_ms=30_000,
pool_ms=30_000,
total_ms=30_000,
),
)
)
except (
RustExecutorClientError,
httpx.HTTPError,
json.JSONDecodeError,
ValueError,
AttributeError,
TypeError,
) as exc:
logger.warning("Codex quota Rust fallback key_id={} url={}: {}", key.id, url, exc)
return None
response_headers = dict(result.headers)
if result.response_json is not None:
response_headers.setdefault("content-type", "application/json")
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
elif result.response_body_bytes is not None:
response_body = result.response_body_bytes
else:
response_body = b""
return httpx.Response(
status_code=result.status_code,
request=httpx.Request("GET", url, headers=headers),
headers=response_headers,
content=response_body,
)
def _looks_like_account_deactivated(message: str | None) -> bool:
lowered = str(message or "").strip().lower()
return "account has been deactivated" in lowered or "account deactivated" in lowered
@@ -212,12 +345,22 @@ async def refresh_codex_key_quota(
getattr(provider, "proxy", None),
getattr(key, "proxy", None),
)
proxy_snapshot = await _build_codex_proxy_snapshot(effective_proxy)
# 使用 wham/usage API 获取限额信息
async with httpx.AsyncClient(
**build_proxy_client_kwargs(effective_proxy, timeout=30.0)
) as client:
response = await client.get(codex_wham_usage_url, headers=headers)
response = await _try_rust_codex_quota_response(
key=key,
provider=provider,
endpoint=endpoint,
url=codex_wham_usage_url,
headers=headers,
proxy_snapshot=proxy_snapshot,
)
if response is None:
async with httpx.AsyncClient(
**build_proxy_client_kwargs(effective_proxy, timeout=30.0)
) as client:
response = await client.get(codex_wham_usage_url, headers=headers)
if response.status_code != 200:
status_code = int(response.status_code)
+182 -55
View File
@@ -145,6 +145,123 @@ class ProviderOpsService:
finally:
self.db.expire_on_commit = original_expire_on_commit
@staticmethod
def _build_ops_proxy_snapshot(
proxy: Any,
tunnel_node_id: str | None,
) -> Any:
from src.services.request.executor_plan import ExecutionProxySnapshot
if tunnel_node_id:
return ExecutionProxySnapshot(
enabled=True,
mode="tunnel",
node_id=str(tunnel_node_id).strip() or None,
)
if proxy is None:
return None
proxy_url = getattr(proxy, "url", None)
if proxy_url is not None:
proxy_url = str(proxy_url).strip()
else:
proxy_url = str(proxy).strip()
if not proxy_url:
return None
mode = proxy_url.split("://", 1)[0].strip().lower() or None
return ExecutionProxySnapshot(
enabled=True,
mode=mode,
url=proxy_url,
)
async def _try_rust_verify_response(
self,
*,
architecture_id: str,
verify_endpoint: str,
headers: dict[str, str],
proxy: Any,
tunnel_node_id: str | None,
) -> Any:
import json
import httpx
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
ExecutionPlanTimeouts,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
try:
result = await RustExecutorClient().execute_sync_json(
ExecutionPlan(
request_id=f"provider-ops-verify:{architecture_id}",
candidate_id=None,
provider_name=architecture_id,
provider_id="",
endpoint_id="",
key_id="",
method="GET",
url=verify_endpoint,
headers=dict(headers),
body=ExecutionPlanBody(),
stream=False,
provider_api_format="provider_ops:verify",
client_api_format="provider_ops:verify",
model_name="verify-auth",
proxy=self._build_ops_proxy_snapshot(proxy, tunnel_node_id),
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=30_000,
write_ms=30_000,
pool_ms=30_000,
total_ms=30_000,
),
)
)
except (
RustExecutorClientError,
httpx.HTTPError,
json.JSONDecodeError,
ValueError,
AttributeError,
TypeError,
) as exc:
logger.warning(
"Provider ops Rust verify fallback architecture={} endpoint={}: {}",
architecture_id,
verify_endpoint,
exc,
)
return None
response_headers = dict(result.headers)
if result.response_json is not None:
response_headers.setdefault("content-type", "application/json")
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
elif result.response_body_bytes is not None:
response_body = result.response_body_bytes
else:
response_body = b""
return httpx.Response(
status_code=result.status_code,
request=httpx.Request("GET", verify_endpoint, headers=headers),
headers=response_headers,
content=response_body,
)
# ==================== 配置管理 ====================
def get_config(self, provider_id: str) -> ProviderOpsConfig | None:
@@ -1046,64 +1163,74 @@ class ProviderOpsService:
proxy, tunnel_node_id = await resolve_ops_proxy_config_async(config)
# 构建 httpx client 参数
client_kwargs: dict[str, Any] = {
"timeout": 30.0,
"verify": get_ssl_context(),
}
if tunnel_node_id:
from src.services.proxy_node.tunnel_transport import create_tunnel_transport
response = await self._try_rust_verify_response(
architecture_id=architecture_id,
verify_endpoint=verify_endpoint,
headers=headers,
proxy=proxy,
tunnel_node_id=tunnel_node_id,
)
if response is None:
# 构建 httpx client 参数
client_kwargs: dict[str, Any] = {
"timeout": 30.0,
"verify": get_ssl_context(),
}
if tunnel_node_id:
from src.services.proxy_node.tunnel_transport import create_tunnel_transport
client_kwargs["transport"] = create_tunnel_transport(tunnel_node_id, timeout=30.0)
logger.debug("使用 tunnel 代理: node_id={}", tunnel_node_id)
elif proxy:
client_kwargs["proxy"] = proxy
logger.debug("使用代理: {}", proxy)
async with httpx.AsyncClient(**client_kwargs) as client:
response = await client.get(verify_endpoint, headers=headers)
logger.debug(
"验证响应: status={}, content_type={}",
response.status_code,
response.headers.get("content-type"),
)
# 尝试解析 JSON
try:
data = response.json()
except Exception:
data = {}
# 将预处理获取的额外数据合并到响应中
if "_combined_data" in merged_config:
data["_combined_data"] = merged_config["_combined_data"]
elif "_balance_data" in merged_config:
data["_balance_data"] = merged_config["_balance_data"]
# 使用架构的方法解析响应
result = architecture.parse_verify_response(response.status_code, data)
result_dict = result.to_dict()
# 将凭据更新信息附加到响应,供前端同步更新表单
# 过滤掉内部缓存字段(以 _ 开头),前端不需要这些
frontend_creds = {k: v for k, v in updated_creds.items() if not k.startswith("_")}
if frontend_creds:
result_dict["updated_credentials"] = frontend_creds
# 验证成功且有 provider_id 时,缓存余额
if result.success and provider_id and result.quota is not None:
# 从架构的默认配置获取 quota_divisor
balance_config = architecture.default_action_configs.get(
ProviderActionType.QUERY_BALANCE, {}
client_kwargs["transport"] = create_tunnel_transport(
tunnel_node_id, timeout=30.0
)
quota_divisor = balance_config.get("quota_divisor", 1)
# 转换为美元值后缓存
quota_usd = result.quota / quota_divisor
# 传入 extra 信息(如窗口限额)
await self._cache_balance_from_verify(provider_id, quota_usd, result.extra)
logger.debug("使用 tunnel 代理: node_id={}", tunnel_node_id)
elif proxy:
client_kwargs["proxy"] = proxy
logger.debug("使用代理: {}", proxy)
return result_dict
async with httpx.AsyncClient(**client_kwargs) as client:
response = await client.get(verify_endpoint, headers=headers)
logger.debug(
"验证响应: status={}, content_type={}",
response.status_code,
response.headers.get("content-type"),
)
# 尝试解析 JSON
try:
data = response.json()
except Exception:
data = {}
# 将预处理获取的额外数据合并到响应中
if "_combined_data" in merged_config:
data["_combined_data"] = merged_config["_combined_data"]
elif "_balance_data" in merged_config:
data["_balance_data"] = merged_config["_balance_data"]
# 使用架构的方法解析响应
result = architecture.parse_verify_response(response.status_code, data)
result_dict = result.to_dict()
# 将凭据更新信息附加到响应,供前端同步更新表单
# 过滤掉内部缓存字段(以 _ 开头),前端不需要这些
frontend_creds = {k: v for k, v in updated_creds.items() if not k.startswith("_")}
if frontend_creds:
result_dict["updated_credentials"] = frontend_creds
# 验证成功且有 provider_id 时,缓存余额
if result.success and provider_id and result.quota is not None:
# 从架构的默认配置获取 quota_divisor
balance_config = architecture.default_action_configs.get(
ProviderActionType.QUERY_BALANCE, {}
)
quota_divisor = balance_config.get("quota_divisor", 1)
# 转换为美元值后缓存
quota_usd = result.quota / quota_divisor
# 传入 extra 信息(如窗口限额)
await self._cache_balance_from_verify(provider_id, quota_usd, result.extra)
return result_dict
except ValueError as e:
return {"success": False, "message": str(e)}
+183
View File
@@ -0,0 +1,183 @@
"""
执行计划契约
用于在 Python 控制面和未来的 Rust executor 之间传递稳定的请求执行信息
当前阶段先服务于非流式 chat 路径的计划构建与本地执行拆分
"""
from __future__ import annotations
import base64
import json
from dataclasses import asdict, dataclass
from typing import Any
def _drop_none(value: Any) -> Any:
"""递归移除 None 字段,便于序列化为紧凑 payload。"""
if isinstance(value, dict):
return {
key: _drop_none(item)
for key, item in value.items()
if item is not None
}
if isinstance(value, list):
return [_drop_none(item) for item in value]
return value
@dataclass(slots=True)
class ExecutionPlanTimeouts:
connect_ms: int | None = None
read_ms: int | None = None
write_ms: int | None = None
pool_ms: int | None = None
total_ms: int | None = None
@dataclass(slots=True)
class ExecutionPlanBody:
json_body: Any = None
body_bytes_b64: str | None = None
body_ref: str | None = None
@dataclass(slots=True)
class ExecutionProxySnapshot:
enabled: bool
mode: str | None = None
node_id: str | None = None
label: str | None = None
url: str | None = None
extra: dict[str, Any] | None = None
@classmethod
def from_proxy_info(
cls,
proxy_info: dict[str, Any] | None,
*,
proxy_url: str | None = None,
mode_override: str | None = None,
node_id_override: str | None = None,
extra: dict[str, Any] | None = None,
) -> ExecutionProxySnapshot | None:
if not proxy_info and not proxy_url:
return None
mode = mode_override or str(proxy_info.get("type") or proxy_info.get("mode") or "").strip() or None
if not mode and proxy_url:
mode = proxy_url.split("://", 1)[0].strip().lower() or None
return cls(
enabled=True,
mode=mode,
node_id=node_id_override
or str((proxy_info or {}).get("node_id") or "").strip()
or None,
label=str((proxy_info or {}).get("label") or "").strip() or None,
url=str(proxy_url or (proxy_info or {}).get("url") or "").strip() or None,
extra=extra or None,
)
@dataclass(slots=True)
class ExecutionPlan:
request_id: str
candidate_id: str | None
provider_name: str
provider_id: str
endpoint_id: str
key_id: str
method: str
url: str
headers: dict[str, str]
body: ExecutionPlanBody
stream: bool
provider_api_format: str
client_api_format: str
model_name: str
content_type: str | None = None
content_encoding: str | None = None
proxy: ExecutionProxySnapshot | None = None
tls_profile: str | None = None
timeouts: ExecutionPlanTimeouts | None = None
def to_payload(self) -> dict[str, Any]:
return _drop_none(asdict(self))
def is_remote_proxy_supported(proxy: ExecutionProxySnapshot | None) -> bool:
proxy_mode = str((proxy.mode if proxy else "") or "").strip().lower()
return proxy is None or (
str(proxy.url or "").strip() != ""
and proxy_mode not in {"tunnel"}
) or (
proxy_mode == "tunnel"
and str(proxy.node_id or "").strip() != ""
)
def is_remote_contract_eligible(contract: ExecutionPlan) -> bool:
content_encoding = str(contract.content_encoding or "").strip().lower()
has_json_body = contract.body.json_body is not None
has_raw_body = bool(str(contract.body.body_bytes_b64 or "").strip())
has_body = has_json_body or has_raw_body
return (
((not has_body and content_encoding == "") or has_body)
and (not has_json_body or content_encoding in {"", "gzip"} or has_raw_body)
and is_remote_proxy_supported(contract.proxy)
)
def build_execution_plan_body(
payload: Any,
*,
content_type: str | None = None,
) -> ExecutionPlanBody:
normalized_content_type = str(content_type or "").strip().lower()
if isinstance(payload, dict):
return ExecutionPlanBody(json_body=payload)
if isinstance(payload, list) and "json" in normalized_content_type:
return ExecutionPlanBody(json_body=payload)
if isinstance(payload, (bytes, bytearray, memoryview)):
return ExecutionPlanBody(
body_bytes_b64=base64.b64encode(bytes(payload)).decode("ascii")
)
if isinstance(payload, str):
return ExecutionPlanBody(
body_bytes_b64=base64.b64encode(payload.encode("utf-8")).decode("ascii")
)
if payload is None:
return ExecutionPlanBody()
return ExecutionPlanBody(
body_bytes_b64=base64.b64encode(
json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
).decode("ascii")
)
@dataclass(slots=True)
class PreparedExecutionPlan:
"""本地执行所需的运行时上下文;`contract` 是可序列化的稳定边界。"""
contract: ExecutionPlan
payload: dict[str, Any]
headers: dict[str, str]
upstream_is_stream: bool
needs_conversion: bool
provider_type: str
request_timeout: float
delegate_config: dict[str, Any] | None = None
proxy_config: dict[str, Any] | None = None
envelope: Any = None
selected_base_url: str | None = None
client_content_encoding: str | None = None
proxy_info: dict[str, Any] | None = None
@property
def remote_eligible(self) -> bool:
return is_remote_contract_eligible(self.contract)
@@ -0,0 +1,245 @@
"""
Rust executor 客户端
当前实现支持两类同步请求
- 普通 JSON 响应
- 上游流式但客户端同步聚合的请求返回原始流字节给 Python 后处理
整体仍然是 best-effort
- backend 未就绪或连接失败时由调用方决定是否回退到 Python 路径
- 传输协议优先支持 Unix Socket + HTTP其次是 TCP + HTTP
"""
from __future__ import annotations
import base64
import json
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from typing import Any
import httpx
from src.config.settings import config
from src.services.request.executor_plan import ExecutionPlan
class RustExecutorClientError(RuntimeError):
"""Rust executor 客户端错误。"""
@dataclass(slots=True)
class RustExecutorSyncResult:
status_code: int
response_json: Any = None
headers: dict[str, str] = field(default_factory=dict)
provider_response_json: Any = None
response_body_bytes: bytes | None = None
@dataclass(slots=True)
class RustExecutorStreamResult:
status_code: int
headers: dict[str, str]
byte_iterator: AsyncIterator[bytes]
response_ctx: Any
class _RustExecutorManagedStreamContext:
def __init__(self, client: httpx.AsyncClient, response_ctx: Any) -> None:
self._client = client
self._response_ctx = response_ctx
self._closed = False
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
if self._closed:
return
self._closed = True
try:
await self._response_ctx.__aexit__(exc_type, exc, tb)
finally:
await self._client.aclose()
class RustExecutorClient:
"""Python 控制面访问 Rust executor 的轻量客户端。"""
def __init__(
self,
*,
transport: str | None = None,
base_url: str | None = None,
socket_path: str | None = None,
request_timeout: float | None = None,
) -> None:
self.transport = (transport or config.executor_transport).strip().lower()
self.base_url = (base_url or config.executor_base_url).strip()
self.socket_path = (socket_path or config.executor_socket_path).strip()
self.request_timeout = (
request_timeout if request_timeout is not None else config.executor_request_timeout
)
def _build_client(self, *, streaming: bool = False) -> httpx.AsyncClient:
if streaming:
timeout = httpx.Timeout(
connect=min(self.request_timeout, 30.0),
read=None,
write=self.request_timeout,
pool=self.request_timeout,
)
else:
timeout = httpx.Timeout(self.request_timeout)
if self.transport == "unix_socket":
if not self.socket_path:
raise RustExecutorClientError("EXECUTOR_SOCKET_PATH is required for unix_socket")
transport = httpx.AsyncHTTPTransport(uds=self.socket_path, retries=0)
return httpx.AsyncClient(
transport=transport,
base_url=self.base_url,
timeout=timeout,
)
if self.transport != "tcp":
raise RustExecutorClientError(f"Unsupported executor transport: {self.transport}")
return httpx.AsyncClient(
base_url=self.base_url,
timeout=timeout,
transport=httpx.AsyncHTTPTransport(retries=0),
)
async def execute_sync_json(self, plan: ExecutionPlan) -> RustExecutorSyncResult:
async with self._build_client() as client:
response = await client.post(
"/v1/execute/sync",
json=plan.to_payload(),
)
response.raise_for_status()
payload = response.json()
status_code = int(payload.get("status_code") or 200)
headers = payload.get("headers") or {}
if not isinstance(headers, dict):
raise RustExecutorClientError("Executor response headers must be an object")
response_json = payload.get("response_json")
provider_response_json = payload.get("provider_response_json")
body_payload = payload.get("body")
body_bytes_b64 = None
if response_json is None and isinstance(body_payload, dict):
response_json = body_payload.get("json_body")
body_bytes_b64 = body_payload.get("body_bytes_b64")
response_body_bytes: bytes | None = None
if body_bytes_b64 is not None:
if not isinstance(body_bytes_b64, str):
raise RustExecutorClientError("Executor body_bytes_b64 must be a string")
try:
response_body_bytes = base64.b64decode(body_bytes_b64)
except Exception as exc: # noqa: BLE001
raise RustExecutorClientError(
"Executor body_bytes_b64 must be valid base64"
) from exc
return RustExecutorSyncResult(
status_code=status_code,
response_json=response_json,
headers={str(k): str(v) for k, v in headers.items()},
provider_response_json=provider_response_json,
response_body_bytes=response_body_bytes,
)
async def execute_stream(self, plan: ExecutionPlan) -> RustExecutorStreamResult:
client = self._build_client(streaming=True)
response_ctx = client.stream(
"POST",
"/v1/execute/stream",
json=plan.to_payload(),
)
try:
response = await response_ctx.__aenter__()
response.raise_for_status()
line_iter = response.aiter_lines()
headers_frame = await self._read_first_stream_frame(line_iter)
payload = headers_frame.get("payload")
if not isinstance(payload, dict) or payload.get("kind") != "headers":
raise RustExecutorClientError("Executor stream must start with headers frame")
status_code = int(payload.get("status_code") or 200)
headers = payload.get("headers") or {}
if not isinstance(headers, dict):
raise RustExecutorClientError("Executor stream headers must be an object")
async def _byte_iter() -> AsyncIterator[bytes]:
async for line in line_iter:
if not line:
continue
frame = self._decode_stream_frame(line)
frame_payload = frame["payload"]
kind = str(frame_payload.get("kind") or "").strip().lower()
if kind == "data":
chunk_b64 = frame_payload.get("chunk_b64")
if isinstance(chunk_b64, str):
if chunk_b64:
try:
yield base64.b64decode(chunk_b64)
except Exception as exc: # noqa: BLE001
raise RustExecutorClientError(
"Executor stream chunk_b64 must be valid base64"
) from exc
continue
text = frame_payload.get("text")
if isinstance(text, str):
if text:
yield text.encode("utf-8")
continue
if kind == "error":
error = frame_payload.get("error") or {}
message = str(error.get("message") or "executor stream error")
raise httpx.ReadError(message)
if kind == "telemetry":
continue
if kind == "eof":
break
raise RustExecutorClientError(f"Unexpected executor stream frame kind: {kind}")
return RustExecutorStreamResult(
status_code=status_code,
headers={str(k): str(v) for k, v in headers.items()},
byte_iterator=_byte_iter(),
response_ctx=_RustExecutorManagedStreamContext(client, response_ctx),
)
except Exception:
try:
await response_ctx.__aexit__(None, None, None)
finally:
await client.aclose()
raise
@staticmethod
def _decode_stream_frame(line: str) -> dict[str, Any]:
try:
frame = json.loads(line)
except json.JSONDecodeError as exc:
raise RustExecutorClientError("Executor stream frame must be valid JSON") from exc
if not isinstance(frame, dict):
raise RustExecutorClientError("Executor stream frame must be an object")
payload = frame.get("payload")
if not isinstance(payload, dict):
raise RustExecutorClientError("Executor stream frame payload must be an object")
return frame
async def _read_first_stream_frame(
self,
line_iter: AsyncIterator[str],
) -> dict[str, Any]:
async for line in line_iter:
if not line:
continue
return self._decode_stream_frame(line)
raise RustExecutorClientError("Executor stream ended before headers frame")
+103 -3
View File
@@ -4,6 +4,7 @@ from typing import Any
from sqlalchemy.orm import Session
from src.config.settings import config
from src.core.logger import logger
from src.models.database import ProviderAPIKey, ProviderEndpoint
from src.services.usage.service import UsageService
@@ -29,8 +30,10 @@ class VideoTaskCancelService:
- None on success
- upstream httpx.Response when upstream returns an error (status >= 400)
"""
import json
from datetime import datetime, timezone
import httpx
from fastapi import HTTPException
from src.clients.http_client import HTTPClientPool
@@ -94,12 +97,100 @@ class VideoTaskCancelService:
header_rules=getattr(endpoint, "header_rules", None),
)
client = await HTTPClientPool.get_default_client_async()
async def _try_rust_cancel_response(
*,
method: str,
url: str,
request_headers: dict[str, str],
body: Any,
content_type: str | None = None,
) -> httpx.Response | None:
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanTimeouts,
build_execution_plan_body,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
final_headers = dict(request_headers)
if (
body is not None
and content_type
and not any(str(key).lower() == "content-type" for key in final_headers)
):
final_headers["content-type"] = content_type
try:
result = await RustExecutorClient().execute_sync_json(
ExecutionPlan(
request_id=str(getattr(task, "request_id", "") or task_id),
candidate_id=None,
provider_name=provider_format_norm.split(":", 1)[0],
provider_id=str(getattr(endpoint, "provider_id", "") or ""),
endpoint_id=str(getattr(endpoint, "id", "") or ""),
key_id=str(getattr(key, "id", "") or ""),
method=method,
url=url,
headers=final_headers,
body=build_execution_plan_body(body, content_type=content_type),
stream=False,
provider_api_format=provider_format,
client_api_format=provider_format,
model_name=str(getattr(task, "model", "") or ""),
content_type=content_type,
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=300_000,
write_ms=300_000,
pool_ms=30_000,
total_ms=300_000,
),
)
)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
logger.warning(
"[VideoCancel] Rust executor unavailable task={} method={} url={}: {}",
getattr(task, "id", task_id),
method,
url,
str(exc),
)
return None
response_headers = dict(result.headers)
if result.response_json is not None:
response_headers.setdefault("content-type", "application/json")
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
elif result.response_body_bytes is not None:
response_body = result.response_body_bytes
else:
response_body = b""
return httpx.Response(
status_code=result.status_code,
request=httpx.Request(method, url, headers=final_headers),
headers=response_headers,
content=response_body,
)
if provider_format_norm.startswith("openai:"):
upstream_url = build_provider_url(endpoint, is_stream=False, key=key)
upstream_url = f"{upstream_url.rstrip('/')}/{str(external_task_id).lstrip('/')}"
response = await client.delete(upstream_url, headers=headers)
response = await _try_rust_cancel_response(
method="DELETE",
url=upstream_url,
request_headers=headers,
body=None,
)
if response is None:
client = await HTTPClientPool.get_default_client_async()
response = await client.delete(upstream_url, headers=headers)
if response.status_code >= 400:
return response
@@ -125,7 +216,16 @@ class VideoTaskCancelService:
headers.pop("x-goog-api-key", None)
headers[auth_info.auth_header] = auth_info.auth_value
response = await client.post(upstream_url, headers=headers, json={})
response = await _try_rust_cancel_response(
method="POST",
url=upstream_url,
request_headers=headers,
body={},
content_type="application/json",
)
if response is None:
client = await HTTPClientPool.get_default_client_async()
response = await client.post(upstream_url, headers=headers, json={})
if response.status_code >= 400:
return response
+172 -17
View File
@@ -37,6 +37,7 @@ from src.core.video_utils import (
from src.database import create_session
from src.models.database import ProviderAPIKey, ProviderEndpoint, VideoTask
from src.services.provider.auth import get_provider_auth
from src.services.provider.provider_context import resolve_provider_proxy
@dataclass(slots=True)
@@ -55,6 +56,9 @@ class VideoPollContext:
poll_interval_seconds: int
max_poll_count: int
current_status: str
proxy_config: dict[str, Any] | None = None
delegate_config: dict[str, Any] | None = None
proxy_snapshot: Any = None
# 永久性错误指示词(用于降级判断,不应重试)
@@ -200,6 +204,10 @@ class VideoTaskPollerAdapter:
else:
auth_info = None
headers = self._build_headers(provider_format, upstream_key, endpoint, auth_info)
proxy_config, delegate_config, proxy_snapshot = await self._build_transport_context(
endpoint=endpoint,
key=key,
)
return VideoPollContext(
task_id=task.id,
@@ -213,6 +221,9 @@ class VideoTaskPollerAdapter:
poll_interval_seconds=task.poll_interval_seconds,
max_poll_count=task.max_poll_count,
current_status=task.status,
proxy_config=proxy_config,
delegate_config=delegate_config,
proxy_snapshot=proxy_snapshot,
)
async def poll_task_http(self, ctx: VideoPollContext) -> InternalVideoPollResult:
@@ -351,13 +362,18 @@ class VideoTaskPollerAdapter:
"""使用上下文进行 OpenAI 轮询(不需要数据库)"""
url = self._build_openai_url(ctx.base_url, ctx.external_task_id)
client = await HTTPClientPool.get_default_client_async()
response = await client.get(url, headers=ctx.headers)
if response.status_code >= 400:
error_message = self._extract_error_message(response.text, response.status_code)
raise PollHTTPError(response.status_code, error_message)
payload = await self._try_rust_poll_payload(ctx=ctx, url=url)
if payload is None:
client = await HTTPClientPool.get_upstream_client(
ctx.delegate_config,
proxy_config=ctx.proxy_config,
)
response = await client.get(url, headers=ctx.headers)
if response.status_code >= 400:
error_message = self._extract_error_message(response.text, response.status_code)
raise PollHTTPError(response.status_code, error_message)
payload = response.json()
payload = response.json()
return self._openai_normalizer.video_poll_to_internal(payload)
async def _poll_gemini_with_context(self, ctx: VideoPollContext) -> InternalVideoPollResult:
@@ -372,21 +388,111 @@ class VideoTaskPollerAdapter:
url,
)
client = await HTTPClientPool.get_default_client_async()
response = await client.get(url, headers=ctx.headers)
if response.status_code >= 400:
logger.warning(
"[VideoPoller] Gemini poll failed: task={} status={} response={}",
ctx.task_id,
response.status_code,
response.text[:500] if response.text else "(empty)",
payload = await self._try_rust_poll_payload(ctx=ctx, url=url)
if payload is None:
client = await HTTPClientPool.get_upstream_client(
ctx.delegate_config,
proxy_config=ctx.proxy_config,
)
error_message = self._extract_error_message(response.text, response.status_code)
raise PollHTTPError(response.status_code, error_message)
response = await client.get(url, headers=ctx.headers)
if response.status_code >= 400:
logger.warning(
"[VideoPoller] Gemini poll failed: task={} status={} response={}",
ctx.task_id,
response.status_code,
response.text[:500] if response.text else "(empty)",
)
error_message = self._extract_error_message(response.text, response.status_code)
raise PollHTTPError(response.status_code, error_message)
payload = response.json()
payload = response.json()
return self._gemini_normalizer.video_poll_to_internal(payload)
async def _try_rust_poll_payload(
self,
*,
ctx: VideoPollContext,
url: str,
) -> dict[str, Any] | None:
import httpx
from src.services.request.executor_plan import (
ExecutionPlan,
ExecutionPlanBody,
ExecutionPlanTimeouts,
)
from src.services.request.rust_executor_client import (
RustExecutorClient,
RustExecutorClientError,
)
if config.executor_backend != "rust":
return None
try:
result = await RustExecutorClient().execute_sync_json(
ExecutionPlan(
request_id=f"video-poll-{ctx.task_id}",
candidate_id=None,
provider_name=ctx.provider_api_format.split(":", 1)[0],
provider_id="",
endpoint_id="",
key_id="",
method="GET",
url=url,
headers=dict(ctx.headers),
body=ExecutionPlanBody(),
stream=False,
provider_api_format=ctx.provider_api_format,
client_api_format=ctx.provider_api_format,
model_name="video-poll",
proxy=ctx.proxy_snapshot,
timeouts=ExecutionPlanTimeouts(
connect_ms=30_000,
read_ms=300_000,
write_ms=300_000,
pool_ms=30_000,
total_ms=300_000,
),
)
)
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError, ValueError) as exc:
logger.warning(
"[VideoPoller] Rust poll fallback task={} url={} error={}",
ctx.task_id,
url,
sanitize_error_message(str(exc)),
)
return None
if result.status_code >= 400:
response_text = ""
if result.response_json is not None:
response_text = json.dumps(result.response_json, ensure_ascii=False)
elif result.response_body_bytes is not None:
response_text = result.response_body_bytes.decode("utf-8", errors="replace")
raise PollHTTPError(
result.status_code,
self._extract_error_message(response_text, result.status_code),
)
if isinstance(result.response_json, dict):
return result.response_json
if result.response_body_bytes is not None:
try:
payload = json.loads(result.response_body_bytes.decode("utf-8"))
if isinstance(payload, dict):
return payload
except Exception:
logger.warning(
"[VideoPoller] Rust poll returned non-json body task={} url={}",
ctx.task_id,
url,
)
return None
# ==================== 旧版方法(保留兼容性)====================
async def poll_single_task(
@@ -591,6 +697,55 @@ class VideoTaskPollerAdapter:
headers[auth_info.auth_header] = auth_info.auth_value
return headers
async def _build_transport_context(
self,
*,
endpoint: ProviderEndpoint,
key: ProviderAPIKey,
) -> tuple[dict[str, Any] | None, dict[str, Any] | None, Any]:
from src.services.proxy_node.resolver import (
build_proxy_url_async,
get_system_proxy_config_async,
resolve_delegate_config_async,
resolve_effective_proxy,
resolve_proxy_info_async,
)
from src.services.request.executor_plan import ExecutionProxySnapshot
try:
effective_proxy = resolve_effective_proxy(
resolve_provider_proxy(endpoint=endpoint, key=key),
getattr(key, "proxy", None),
)
if not effective_proxy or not effective_proxy.get("enabled", True):
effective_proxy = await get_system_proxy_config_async()
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
proxy_url: str | None = None
if effective_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
proxy_url = await build_proxy_url_async(effective_proxy)
proxy_info = await resolve_proxy_info_async(effective_proxy)
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
proxy_info,
proxy_url=proxy_url,
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
node_id_override=(
str(delegate_cfg.get("node_id") or "").strip() or None
if delegate_cfg and delegate_cfg.get("tunnel")
else None
),
)
return effective_proxy, delegate_cfg, proxy_snapshot
except Exception as exc:
logger.warning(
"[VideoPoller] Failed to build transport context endpoint={} key={}: {}",
getattr(endpoint, "id", None),
getattr(key, "id", None),
sanitize_error_message(str(exc)),
)
return None, None, None
def _get_endpoint(self, db: Session, endpoint_id: str) -> ProviderEndpoint:
endpoint = db.query(ProviderEndpoint).filter(ProviderEndpoint.id == endpoint_id).first()
if not endpoint: