From ecd3c90e80ce83d8b8f53f69aa69d71515dd2e09 Mon Sep 17 00:00:00 2001 From: Kris <2893855659@qq.com> Date: Thu, 21 May 2026 10:26:33 +0800 Subject: [PATCH] =?UTF-8?q?feat(channel/gateway):=20=E6=96=B0=E5=A2=9E?= =?UTF-8?q?=E5=AE=8C=E6=95=B4=E7=BD=91=E5=85=B3=E9=80=9A=E9=81=93=E6=A8=A1?= =?UTF-8?q?=E5=9D=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增设备身份管理、认证限流、并发通道、Webhook路由、RBAC权限控制、SSE/轮询降级等全套网关通道功能，包含： 1. 设备身份生成与签名验证 2. 设备令牌认证与速率限制 3. 内存+数据库双重设备注册表 4. 并发通道限流管理 5. Webhook安全处理与路由 6. RBAC权限校验系统 7. OpenAI API兼容适配层 8. Tailscale认证支持 9. HTTP轮询降级机制 --- .../package/yuxi/channel/gateway/__init__.py | 177 ++ backend/package/yuxi/channel/gateway/auth.py | 387 +++++ .../yuxi/channel/gateway/auth_rate_limiter.py | 111 ++ .../yuxi/channel/gateway/broadcaster.py | 359 ++++ .../yuxi/channel/gateway/channel_plugin.py | 192 +++ .../package/yuxi/channel/gateway/client.py | 557 +++++++ .../yuxi/channel/gateway/device_auth.py | 108 ++ .../yuxi/channel/gateway/device_identity.py | 68 + .../yuxi/channel/gateway/device_registry.py | 61 + .../channel/gateway/device_registry_db.py | 31 + backend/package/yuxi/channel/gateway/lanes.py | 114 ++ .../package/yuxi/channel/gateway/net_utils.py | 373 +++++ .../yuxi/channel/gateway/openai_adapter.py | 161 ++ .../package/yuxi/channel/gateway/polling.py | 102 ++ backend/package/yuxi/channel/gateway/probe.py | 289 ++++ .../package/yuxi/channel/gateway/protocol.py | 207 +++ backend/package/yuxi/channel/gateway/rbac.py | 149 ++ .../package/yuxi/channel/gateway/routes.py | 161 ++ .../yuxi/channel/gateway/rpc_dispatcher.py | 193 +++ .../yuxi/channel/gateway/rpc_handlers.py | 1444 +++++++++++++++++ .../package/yuxi/channel/gateway/server.py | 382 +++++ backend/package/yuxi/channel/gateway/sse.py | 277 ++++ .../yuxi/channel/gateway/tailscale_auth.py | 173 ++ .../yuxi/channel/gateway/validation.py | 110 ++ .../yuxi/channel/gateway/webhook_security.py | 304 ++++ 25 files changed, 6490 insertions(+) create mode 100644 backend/package/yuxi/channel/gateway/__init__.py create mode 100644 backend/package/yuxi/channel/gateway/auth.py create mode 100644 backend/package/yuxi/channel/gateway/auth_rate_limiter.py create mode 100644 backend/package/yuxi/channel/gateway/broadcaster.py create mode 100644 backend/package/yuxi/channel/gateway/channel_plugin.py create mode 100644 backend/package/yuxi/channel/gateway/client.py create mode 100644 backend/package/yuxi/channel/gateway/device_auth.py create mode 100644 backend/package/yuxi/channel/gateway/device_identity.py create mode 100644 backend/package/yuxi/channel/gateway/device_registry.py create mode 100644 backend/package/yuxi/channel/gateway/device_registry_db.py create mode 100644 backend/package/yuxi/channel/gateway/lanes.py create mode 100644 backend/package/yuxi/channel/gateway/net_utils.py create mode 100644 backend/package/yuxi/channel/gateway/openai_adapter.py create mode 100644 backend/package/yuxi/channel/gateway/polling.py create mode 100644 backend/package/yuxi/channel/gateway/probe.py create mode 100644 backend/package/yuxi/channel/gateway/protocol.py create mode 100644 backend/package/yuxi/channel/gateway/rbac.py create mode 100644 backend/package/yuxi/channel/gateway/routes.py create mode 100644 backend/package/yuxi/channel/gateway/rpc_dispatcher.py create mode 100644 backend/package/yuxi/channel/gateway/rpc_handlers.py create mode 100644 backend/package/yuxi/channel/gateway/server.py create mode 100644 backend/package/yuxi/channel/gateway/sse.py create mode 100644 backend/package/yuxi/channel/gateway/tailscale_auth.py create mode 100644 backend/package/yuxi/channel/gateway/validation.py create mode 100644 backend/package/yuxi/channel/gateway/webhook_security.py diff --git a/backend/package/yuxi/channel/gateway/__init__.py b/backend/package/yuxi/channel/gateway/__init__.py new file mode 100644 index 00000000..43815899 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/__init__.py @@ -0,0 +1,177 @@ +from yuxi.channel.gateway.auth import ( + GatewayAuthMode, + GatewayAuthResult, + TrustedProxyConfig, + authenticate_gateway_connect, + authenticate_tailscale, + authenticate_trusted_proxy, + clear_bootstrap_token, + generate_bootstrap_token, + get_bootstrap_token, + set_bootstrap_token, +) +from yuxi.channel.gateway.auth_rate_limiter import AuthRateLimiter, AuthRateScope, auth_rate_limiter +from yuxi.channel.gateway.device_auth import build_device_token, parse_device_token +from yuxi.channel.gateway.device_identity import ( + derive_device_id, + generate_device_identity, + sign_challenge, + verify_signature, +) +from yuxi.channel.gateway.device_registry import ( + list_devices, + lookup_public_key, + register_device, + unregister_device, +) +from yuxi.channel.gateway.lanes import ChannelLane, LaneManager, lane_manager +from yuxi.channel.gateway.net_utils import ( + GatewayBindMode, + build_ws_security_error, + can_bind_to_host, + can_bind_to_host_cached, + has_forwarded_request_headers, + has_tailscale_proxy_headers, + is_local_direct_request, + is_localish_host, + is_loopback_address, + is_loopback_host, + is_private_host, + is_private_or_loopback_host, + is_secure_ws_url, + is_tailscale_proxy_request, + is_trusted_proxy_address, + resolve_client_ip, + resolve_forwarded_client_ip, + resolve_gateway_bind_host, + resolve_tailscale_client_ip, +) +from yuxi.channel.gateway.protocol import ( + DeliveryMode, + FrameType, + GatewayErrorCode, + GatewayRpcMethod, + HelloOk, + RpcEvent, + RpcFrame, + RpcRequest, + RpcResponse, + marshal_frame, + unmarshal_frame, +) +from yuxi.channel.gateway.routes import ( + WebhookRegistry, + build_health_response, + build_webhook_path, + webhook_registry, +) +from yuxi.channel.gateway.openai_adapter import ( + OPENAI_CHAT_COMPLETIONS_PATH, + OPENAI_MODELS_PATH, + OpenAIMessageConverter, + format_sse_stream, +) +from yuxi.channel.gateway.rpc_dispatcher import RpcDispatcher, rpc_dispatcher +from yuxi.channel.gateway.server import GatewayWsServer, gateway_ws_server +from yuxi.channel.gateway.sse import GatewaySseEndpoint, gateway_sse_endpoint +from yuxi.channel.gateway.polling import PollingFallback, polling_fallback, start_polling_cleanup +from yuxi.channel.gateway.tailscale_auth import ( + TailscaleWhoisIdentity, + get_tailscale_user_from_headers, + read_tailscale_whois_identity, +) +from yuxi.channel.gateway.webhook_security import ( + WebhookAnomalyTracker, + WebhookConcurrencyGuard, + WebhookGuard, + WebhookGuardConfig, + WebhookGuardResult, + WebhookSigner, + webhook_guard, +) + +__all__ = [ + "AuthRateLimiter", + "AuthRateScope", + "ChannelLane", + "DeliveryMode", + "FrameType", + "GatewayAuthMode", + "GatewayAuthResult", + "GatewayBindMode", + "GatewayErrorCode", + "GatewayRpcMethod", + "HelloOk", + "GatewayWsServer", + "GatewaySseEndpoint", + "LaneManager", + "OPENAI_CHAT_COMPLETIONS_PATH", + "OPENAI_MODELS_PATH", + "OpenAIMessageConverter", + "PollingFallback", + "RpcDispatcher", + "RpcEvent", + "RpcFrame", + "RpcRequest", + "RpcResponse", + "TailscaleWhoisIdentity", + "TrustedProxyConfig", + "WebhookAnomalyTracker", + "WebhookConcurrencyGuard", + "WebhookGuard", + "WebhookGuardConfig", + "WebhookGuardResult", + "WebhookRegistry", + "WebhookSigner", + "authenticate_gateway_connect", + "authenticate_tailscale", + "authenticate_trusted_proxy", + "auth_rate_limiter", + "build_device_token", + "build_health_response", + "build_webhook_path", + "build_ws_security_error", + "can_bind_to_host", + "can_bind_to_host_cached", + "clear_bootstrap_token", + "derive_device_id", + "gateway_ws_server", + "gateway_sse_endpoint", + "format_sse_stream", + "polling_fallback", + "generate_bootstrap_token", + "generate_device_identity", + "get_bootstrap_token", + "get_tailscale_user_from_headers", + "has_forwarded_request_headers", + "has_tailscale_proxy_headers", + "is_local_direct_request", + "is_localish_host", + "is_loopback_address", + "is_loopback_host", + "is_private_host", + "is_private_or_loopback_host", + "is_secure_ws_url", + "is_tailscale_proxy_request", + "is_trusted_proxy_address", + "lane_manager", + "list_devices", + "lookup_public_key", + "marshal_frame", + "parse_device_token", + "read_tailscale_whois_identity", + "register_device", + "resolve_client_ip", + "resolve_forwarded_client_ip", + "resolve_gateway_bind_host", + "resolve_tailscale_client_ip", + "rpc_dispatcher", + "set_bootstrap_token", + "sign_challenge", + "start_polling_cleanup", + "unmarshal_frame", + "unregister_device", + "verify_signature", + "webhook_guard", + "webhook_registry", +] diff --git a/backend/package/yuxi/channel/gateway/auth.py b/backend/package/yuxi/channel/gateway/auth.py new file mode 100644 index 00000000..f34caa88 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/auth.py @@ -0,0 +1,387 @@ +import hmac +import logging +import secrets +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from enum import StrEnum + +from server.utils.auth_utils import AuthUtils +from yuxi.channel.gateway.auth_rate_limiter import ( + AUTH_LOCKOUT_SECONDS, + AuthRateScope, + auth_rate_limiter, +) +from yuxi.channel.gateway.rbac import GatewayRole, map_user_role + +logger = logging.getLogger(__name__) + +AUTH_RATE_LOCKOUT_MINUTES = int(AUTH_LOCKOUT_SECONDS / 60) + +LookupPublicKey = Callable[[str], Awaitable[str | None]] + + +class GatewayAuthMode(StrEnum): + TOKEN = "token" + PASSWORD = "password" + DEVICE_TOKEN = "device_token" + BOOTSTRAP = "bootstrap" + TAILSCALE = "tailscale" + TRUSTED_PROXY = "trusted_proxy" + + +@dataclass +class TrustedProxyConfig: + user_header: str + required_headers: list[str] = field(default_factory=list) + allow_users: list[str] = field(default_factory=list) + allow_loopback: bool = False + + +@dataclass +class GatewayAuthResult: + authenticated: bool + user_id: str | None = None + mode: GatewayAuthMode | None = None + error: str | None = None + metadata: dict = field(default_factory=dict) + roles: list[GatewayRole] = field(default_factory=list) + + +import threading + +_bootstrap_token: str | None = None +_bootstrap_token_lock = threading.Lock() + + +def set_bootstrap_token(token: str | None) -> None: + global _bootstrap_token + with _bootstrap_token_lock: + _bootstrap_token = token + + +def get_bootstrap_token() -> str | None: + with _bootstrap_token_lock: + return _bootstrap_token + + +def clear_bootstrap_token() -> None: + global _bootstrap_token + with _bootstrap_token_lock: + _bootstrap_token = None + + +def generate_bootstrap_token() -> str: + return secrets.token_hex(24) + + +def _timing_safe_compare(a: str, b: str) -> bool: + if len(a) != len(b): + return False + return hmac.compare_digest(a.encode(), b.encode()) + + +def _authenticate_bootstrap_token(token: str, client_ip: str | None) -> GatewayAuthResult: + scope_key = client_ip or "unknown" + stored = get_bootstrap_token() + + is_locked = auth_rate_limiter.is_locked(scope_key, AuthRateScope.BOOTSTRAP) + if is_locked: + return GatewayAuthResult( + authenticated=False, + error=f"Bootstrap 认证尝试次数过多，请 {AUTH_RATE_LOCKOUT_MINUTES} 分钟后重试。", + metadata={"rate_limited": True, "retry_after_seconds": AUTH_LOCKOUT_SECONDS}, + ) + + if not stored: + auth_rate_limiter.record_failure(scope_key, AuthRateScope.BOOTSTRAP) + return GatewayAuthResult( + authenticated=False, + error="Bootstrap token 未配置或已过期。", + ) + + if not _timing_safe_compare(token, stored): + auth_rate_limiter.record_failure(scope_key, AuthRateScope.BOOTSTRAP) + return GatewayAuthResult( + authenticated=False, + error="Bootstrap token 验证失败。", + ) + + auth_rate_limiter.reset(scope_key, AuthRateScope.BOOTSTRAP) + return GatewayAuthResult( + authenticated=True, + user_id="bootstrap", + mode=GatewayAuthMode.BOOTSTRAP, + roles=[GatewayRole.ADMIN], + metadata={"bootstrap": True}, + ) + + +async def authenticate_tailscale( + headers: dict | None, + client_ip: str | None = None, + remote_addr: str | None = None, + tailscale_whois: Callable[[str], Awaitable] | None = None, +) -> GatewayAuthResult: + from yuxi.channel.gateway.net_utils import ( + is_tailscale_proxy_request, + resolve_tailscale_client_ip, + ) + from yuxi.channel.gateway.tailscale_auth import ( + get_tailscale_user_from_headers, + read_tailscale_whois_identity, + ) + + tailscale_user = get_tailscale_user_from_headers(headers) + if not tailscale_user: + return GatewayAuthResult( + authenticated=False, + error="缺少 Tailscale 用户身份头信息。", + ) + + login, name = tailscale_user + + if not is_tailscale_proxy_request(remote_addr, headers): + return GatewayAuthResult( + authenticated=False, + error="请求未经过 Tailscale Serve 代理。", + ) + + real_client_ip = resolve_tailscale_client_ip(remote_addr, headers) + if not real_client_ip: + return GatewayAuthResult( + authenticated=False, + error="无法解析 Tailscale 客户端 IP。", + ) + + whois_func = tailscale_whois or read_tailscale_whois_identity + whois = await whois_func(real_client_ip) + if not whois or not whois.login: + return GatewayAuthResult( + authenticated=False, + error="Tailscale whois 查询失败。", + ) + + if whois.login != login: + return GatewayAuthResult( + authenticated=False, + error="Tailscale 用户身份不匹配。", + ) + + return GatewayAuthResult( + authenticated=True, + user_id=f"tailscale:{whois.login}", + mode=GatewayAuthMode.TAILSCALE, + metadata={"tailscale_login": whois.login, "tailscale_name": whois.name}, + roles=[GatewayRole.OPERATOR], + ) + + +def authenticate_trusted_proxy( + headers: dict | None, + client_ip: str | None, + remote_addr: str | None, + config: TrustedProxyConfig, + trusted_proxies: list[str], +) -> GatewayAuthResult: + from yuxi.channel.gateway.net_utils import ( + is_loopback_address, + is_trusted_proxy_address, + ) + + if not headers: + return GatewayAuthResult( + authenticated=False, + error="缺少请求头信息，无法进行受信代理认证。", + ) + + if not remote_addr or not is_trusted_proxy_address(remote_addr, trusted_proxies): + return GatewayAuthResult( + authenticated=False, + error="请求来源 IP 不在受信代理列表中。", + ) + + remote_is_loopback = is_loopback_address(remote_addr) + if remote_is_loopback and not config.allow_loopback: + return GatewayAuthResult( + authenticated=False, + error="受信代理不允许回环地址来源，请启用 allow_loopback 配置。", + ) + + for header_name in config.required_headers: + value = headers.get(header_name.lower()) + if not value or not value.strip(): + return GatewayAuthResult( + authenticated=False, + error=f"受信代理缺少必要的请求头: {header_name}。", + ) + + user_header_name = config.user_header.lower() + user_header_value = headers.get(user_header_name) + if not user_header_value or not user_header_value.strip(): + return GatewayAuthResult( + authenticated=False, + error=f"受信代理未提供用户身份请求头: {config.user_header}。", + ) + + user = user_header_value.strip() + + if config.allow_users and user not in config.allow_users: + return GatewayAuthResult( + authenticated=False, + error=f"用户 {user} 不在受信代理的允许列表中。", + ) + + return GatewayAuthResult( + authenticated=True, + user_id=f"proxy:{user}", + mode=GatewayAuthMode.TRUSTED_PROXY, + metadata={"proxy_user": user, "proxy_remote": remote_addr}, + roles=[GatewayRole.ADMIN], + ) + + +async def authenticate_gateway_connect( + auth_header: str | None, + query_token: str | None = None, + client_ip: str | None = None, + shared_password: str | None = None, + lookup_public_key: LookupPublicKey | None = None, + bootstrap_token_param: str | None = None, + headers: dict | None = None, + remote_addr: str | None = None, + trusted_proxy_config: TrustedProxyConfig | None = None, + trusted_proxies: list[str] | None = None, + allow_tailscale: bool = False, +) -> GatewayAuthResult: + scope_key = client_ip or "unknown" + + if trusted_proxy_config and trusted_proxies: + result = authenticate_trusted_proxy( + headers=headers, + client_ip=client_ip, + remote_addr=remote_addr, + config=trusted_proxy_config, + trusted_proxies=trusted_proxies, + ) + if result.authenticated: + return result + + if bootstrap_token_param: + return _authenticate_bootstrap_token(bootstrap_token_param, client_ip) + + token = None + if auth_header and auth_header.startswith("Bearer "): + token = auth_header[7:] + elif query_token: + token = query_token + + if allow_tailscale: + from yuxi.channel.gateway.net_utils import is_local_direct_request + from yuxi.channel.gateway.tailscale_auth import get_tailscale_user_from_headers + + if not is_local_direct_request(remote_addr, headers): + tailscale_user = get_tailscale_user_from_headers(headers) + if tailscale_user and not (shared_password and token): + result = await authenticate_tailscale( + headers=headers, + client_ip=client_ip, + remote_addr=remote_addr, + ) + if result.authenticated: + return result + + if token and token.startswith("dv."): + return await _authenticate_device_token(token, client_ip, lookup_public_key) + + if shared_password and token: + if not _timing_safe_compare(token, shared_password): + auth_rate_limiter.record_failure(scope_key, AuthRateScope.PASSWORD) + return GatewayAuthResult( + authenticated=False, + error="密码认证失败。", + ) + auth_rate_limiter.reset(scope_key, AuthRateScope.PASSWORD) + return GatewayAuthResult( + authenticated=True, + user_id="shared", + mode=GatewayAuthMode.PASSWORD, + roles=[GatewayRole.ADMIN], + ) + + if shared_password and not token: + return GatewayAuthResult( + authenticated=False, + error="缺少认证 token。", + ) + + if not token: + return GatewayAuthResult( + authenticated=False, + error="缺少认证 token。WebSocket 头部须携带 Bearer token，或连接 URL 附加 ?token=xxx 参数。", + ) + + is_locked = auth_rate_limiter.is_locked(scope_key, AuthRateScope.TOKEN) + if is_locked: + remaining = auth_rate_limiter.remaining_attempts(scope_key, AuthRateScope.TOKEN) + return GatewayAuthResult( + authenticated=False, + error=f"认证尝试次数过多，请 {AUTH_RATE_LOCKOUT_MINUTES} 分钟后重试。", + metadata={"rate_limited": True, "retry_after_seconds": AUTH_LOCKOUT_SECONDS}, + ) + + try: + payload = AuthUtils.verify_access_token(token) + user_id = payload.get("sub") + if not user_id: + raise ValueError("Token 缺少 subject 声明") + user_role = payload.get("role") + roles = [map_user_role(user_role)] + auth_rate_limiter.reset(scope_key, AuthRateScope.TOKEN) + return GatewayAuthResult( + authenticated=True, + user_id=user_id, + mode=GatewayAuthMode.TOKEN, + metadata={"token_payload": payload}, + roles=roles, + ) + except Exception as e: + auth_rate_limiter.record_failure(scope_key, AuthRateScope.TOKEN) + logger.debug("Token verification failed for scope=%s: %s", scope_key, e) + remaining = auth_rate_limiter.remaining_attempts(scope_key, AuthRateScope.TOKEN) + return GatewayAuthResult( + authenticated=False, + error="Token 验证失败，请检查凭据是否有效。", + metadata={"remaining_attempts": remaining}, + ) + + +async def _authenticate_device_token( + token: str, + client_ip: str | None, + lookup_public_key: LookupPublicKey | None, +) -> GatewayAuthResult: + from yuxi.channel.gateway.device_auth import authenticate_device + from yuxi.channel.gateway.device_registry import lookup_public_key as default_lookup + + scope_key = client_ip or "unknown" + + is_locked = auth_rate_limiter.is_locked(scope_key, AuthRateScope.DEVICE_TOKEN) + if is_locked: + return GatewayAuthResult( + authenticated=False, + error=f"设备认证尝试次数过多，请 {AUTH_RATE_LOCKOUT_MINUTES} 分钟后重试。", + metadata={"rate_limited": True, "retry_after_seconds": AUTH_LOCKOUT_SECONDS}, + ) + + result = await authenticate_device( + token=token, + lookup_public_key=lookup_public_key or default_lookup, + client_ip=client_ip, + ) + + if not result.authenticated: + auth_rate_limiter.record_failure(scope_key, AuthRateScope.DEVICE_TOKEN) + else: + auth_rate_limiter.reset(scope_key, AuthRateScope.DEVICE_TOKEN) + + return result diff --git a/backend/package/yuxi/channel/gateway/auth_rate_limiter.py b/backend/package/yuxi/channel/gateway/auth_rate_limiter.py new file mode 100644 index 00000000..8dab1ae8 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/auth_rate_limiter.py @@ -0,0 +1,111 @@ +import logging +import threading +import time +from enum import StrEnum +from typing import NamedTuple + +logger = logging.getLogger(__name__) + +AUTH_FAIL_LIMIT = 10 +AUTH_LOCKOUT_SECONDS = 300 +AUTH_RATE_SCOPE_BUCKETS = 10_000 + + +class AuthRateScope(StrEnum): + TOKEN = "token" + PASSWORD = "password" + DEVICE_TOKEN = "device_token" + BOOTSTRAP = "bootstrap" + + +class _ScopeState(NamedTuple): + failures: int + locked_until: float + + +class AuthRateLimiter: + def __init__( + self, + fail_limit: int = AUTH_FAIL_LIMIT, + lockout_seconds: float = AUTH_LOCKOUT_SECONDS, + max_buckets: int = AUTH_RATE_SCOPE_BUCKETS, + ): + self._fail_limit = fail_limit + self._lockout_seconds = lockout_seconds + self._max_buckets = max_buckets + self._buckets: dict[tuple[str, AuthRateScope], _ScopeState] = {} + self._lock = threading.Lock() + + def _cleanup_expired(self) -> None: + if len(self._buckets) <= self._max_buckets: + return + now = time.monotonic() + expired = [ + key + for key, state in self._buckets.items() + if now >= state.locked_until and state.failures < self._fail_limit + ] + for key in expired: + del self._buckets[key] + + def record_failure(self, identifier: str, scope: AuthRateScope) -> bool: + with self._lock: + self._cleanup_expired() + key = (identifier, scope) + now = time.monotonic() + state = self._buckets.get(key) + + if state is not None and now < state.locked_until: + return False + + new_failures = (state.failures + 1) if state else 1 + locked_until = now + self._lockout_seconds if new_failures >= self._fail_limit else 0.0 + self._buckets[key] = _ScopeState(new_failures, locked_until) + + if new_failures >= self._fail_limit: + logger.warning( + "Auth rate limit locked: identifier=%s scope=%s", + identifier, + scope.value, + ) + + if len(self._buckets) > self._max_buckets * 2: + self._force_evict(now) + + return True + + def reset(self, identifier: str, scope: AuthRateScope) -> None: + key = (identifier, scope) + with self._lock: + self._buckets.pop(key, None) + + def is_locked(self, identifier: str, scope: AuthRateScope) -> bool: + key = (identifier, scope) + with self._lock: + state = self._buckets.get(key) + if state is None: + return False + return time.monotonic() < state.locked_until + + def remaining_attempts(self, identifier: str, scope: AuthRateScope) -> int: + key = (identifier, scope) + with self._lock: + state = self._buckets.get(key) + if state is None: + return self._fail_limit + now = time.monotonic() + if state.locked_until > 0.0 and now >= state.locked_until: + return self._fail_limit + return max(0, self._fail_limit - state.failures) + + def _force_evict(self, now: float) -> None: + by_expiry = sorted( + self._buckets.items(), + key=lambda item: item[1].locked_until, + ) + to_remove = len(self._buckets) - self._max_buckets + for i in range(to_remove): + del self._buckets[by_expiry[i][0]] + + +auth_rate_limiter = AuthRateLimiter() diff --git a/backend/package/yuxi/channel/gateway/broadcaster.py b/backend/package/yuxi/channel/gateway/broadcaster.py new file mode 100644 index 00000000..4fdb7bd2 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/broadcaster.py @@ -0,0 +1,359 @@ +import logging +import time as _time +from collections.abc import Collection +from dataclasses import dataclass, field +from typing import Any + +from yuxi.channel.gateway.auth import GatewayAuthResult +from yuxi.channel.gateway.protocol import RpcEvent, marshal_frame +from yuxi.channel.gateway.rbac import ROLE_HIERARCHY, GatewayRole + +logger = logging.getLogger(__name__) + +DEFAULT_BROADCAST_RATE = 100 +DEFAULT_BROADCAST_BURST = 200 +DEFAULT_MAX_BUFFERED_FRAMES = 256 +DEFAULT_SLOW_CONSUMER_THRESHOLD = 50 + +PUBLIC_EVENTS: set[str] = { + "heartbeat", + "ping", + "server.shutdown", + "server.tick", +} + +READ_EVENTS: set[str] = PUBLIC_EVENTS | { + "channel.created", + "channel.updated", + "channel.deleted", + "channel.state_change", + "channel.message", + "session.created", + "session.updated", + "session.deleted", + "session.message", + "session.tool", + "agent.status", + "agent.event", + "cron.status", +} + +WRITE_EVENTS: set[str] = READ_EVENTS | { + "config.changed", +} + +ADMIN_EVENTS: set[str] = WRITE_EVENTS | { + "plugin.installed", + "plugin.removed", + "device.pair.requested", + "device.pair.resolved", +} + +EVENT_MIN_ROLE: dict[str, GatewayRole] = {} + +for _event in PUBLIC_EVENTS: + EVENT_MIN_ROLE[_event] = GatewayRole.VIEWER +for _event in READ_EVENTS - PUBLIC_EVENTS: + EVENT_MIN_ROLE[_event] = GatewayRole.VIEWER +for _event in WRITE_EVENTS - READ_EVENTS: + EVENT_MIN_ROLE[_event] = GatewayRole.OPERATOR +for _event in ADMIN_EVENTS - WRITE_EVENTS: + EVENT_MIN_ROLE[_event] = GatewayRole.ADMIN + + +@dataclass +class _ConnInfo: + conn_id: str + user_id: str | None = None + roles: list[GatewayRole] = field(default_factory=list) + subscribed_events: set[str] = field(default_factory=set) + pending_frames: int = 0 + + +@dataclass +class BroadcastFilter: + conn_ids: Collection[str] | None = None + user_ids: Collection[str] | None = None + roles: Collection[GatewayRole] | None = None + drop_if_slow: bool = True + + +@dataclass +class BroadcastResult: + sent: int = 0 + dropped: int = 0 + filtered: int = 0 + + +class TokenBucket: + def __init__(self, rate: float, burst: float): + self.rate = rate + self.burst = burst + self.tokens = burst + self.last_refill = _time.monotonic() + + def consume(self, n: float = 1.0) -> bool: + self._refill() + if self.tokens >= n: + self.tokens -= n + return True + return False + + def _refill(self): + now = _time.monotonic() + elapsed = now - self.last_refill + self.tokens = min(self.burst, self.tokens + elapsed * self.rate) + self.last_refill = now + + +class GatewayBroadcaster: + def __init__( + self, + broadcast_rate: float = DEFAULT_BROADCAST_RATE, + broadcast_burst: float = DEFAULT_BROADCAST_BURST, + max_buffered_frames: int = DEFAULT_MAX_BUFFERED_FRAMES, + slow_consumer_threshold: int = DEFAULT_SLOW_CONSUMER_THRESHOLD, + ): + self._connections: dict[str, _ConnInfo] = {} + self._rate_limiter = TokenBucket(broadcast_rate, broadcast_burst) + self._max_buffered_frames = max_buffered_frames + self._slow_consumer_threshold = slow_consumer_threshold + self._send_func: dict[str, Any] = {} + self._global_seq = 0 + + @property + def global_seq(self) -> int: + return self._global_seq + + @property + def connection_count(self) -> int: + return len(self._connections) + + def register_connection( + self, + conn_id: str, + auth: GatewayAuthResult, + send_fn: Any, + ) -> None: + info = _ConnInfo( + conn_id=conn_id, + user_id=auth.user_id, + roles=list(auth.roles) if auth.roles else [], + ) + self._connections[conn_id] = info + self._send_func[conn_id] = send_fn + + def unregister_connection(self, conn_id: str) -> None: + self._connections.pop(conn_id, None) + self._send_func.pop(conn_id, None) + + def subscribe(self, conn_id: str, events: Collection[str]) -> None: + info = self._connections.get(conn_id) + if info is None: + return + info.subscribed_events.update(events) + + def unsubscribe(self, conn_id: str, events: Collection[str]) -> None: + info = self._connections.get(conn_id) + if info is None: + return + info.subscribed_events.difference_update(events) + if not events or "*" in events: + info.subscribed_events.clear() + + def clear_subscriptions(self, conn_id: str) -> None: + info = self._connections.get(conn_id) + if info is None: + return + info.subscribed_events.clear() + + def get_subscriptions(self, conn_id: str) -> set[str]: + info = self._connections.get(conn_id) + if info is None: + return set() + return set(info.subscribed_events) + + async def broadcast( + self, + event: str, + data: dict | None = None, + filter_: BroadcastFilter | None = None, + ) -> BroadcastResult: + return await self._broadcast_internal(event, data, filter_) + + async def broadcast_to_conn_ids( + self, + event: str, + data: dict | None, + conn_ids: Collection[str], + ) -> BroadcastResult: + return await self._broadcast_internal(event, data, BroadcastFilter(conn_ids=conn_ids)) + + async def broadcast_to_user( + self, + event: str, + data: dict | None, + user_id: str, + ) -> BroadcastResult: + return await self._broadcast_internal(event, data, BroadcastFilter(user_ids={user_id})) + + async def broadcast_to_role( + self, + event: str, + data: dict | None, + role: GatewayRole, + ) -> BroadcastResult: + return await self._broadcast_internal(event, data, BroadcastFilter(roles={role})) + + async def send_event( + self, + conn_id: str, + event: str, + data: dict | None = None, + ) -> bool: + send_fn = self._send_func.get(conn_id) + if send_fn is None: + return False + self._global_seq += 1 + rpc_event = RpcEvent( + event=event, + data=data, + seq=self._global_seq, + state_version=self._resolve_state_version(event), + ) + try: + await send_fn(marshal_frame(rpc_event)) + return True + except Exception: + return False + + async def _broadcast_internal( + self, + event: str, + data: dict | None, + filter_: BroadcastFilter | None, + ) -> BroadcastResult: + + result = BroadcastResult() + + if not self._connections: + return result + + candidates = [ + conn_id + for conn_id, info in self._connections.items() + if self._passes_filter(info, event, filter_) and self._has_event_scope(info, event) + ] + candidate_count = len(candidates) + result.filtered = len(self._connections) - candidate_count + + if not candidates: + return result + + if not self._rate_limiter.consume(candidate_count): + logger.warning("Broadcast rate limited: event=%s candidates=%d", event, candidate_count) + return result + + self._global_seq += 1 + rpc_event = RpcEvent( + event=event, + data=data, + seq=self._global_seq, + state_version=self._resolve_state_version(event), + ) + frame = marshal_frame(rpc_event) + + for conn_id in candidates: + info = self._connections.get(conn_id) + if info is None: + continue + + info.pending_frames += 1 + is_slow = info.pending_frames > self._slow_consumer_threshold + drop = filter_ is not None and filter_.drop_if_slow + + if is_slow and drop: + info.pending_frames = max(0, info.pending_frames - 1) + result.dropped += 1 + continue + + if info.pending_frames > self._max_buffered_frames: + logger.warning( + "Slow consumer exceeded max buffer: %s (event=%s pending=%d)", + conn_id, + event, + info.pending_frames, + ) + info.pending_frames = max(0, info.pending_frames - 1) + result.dropped += 1 + continue + + send_fn = self._send_func.get(conn_id) + if send_fn is None: + continue + + try: + await send_fn(frame) + info.pending_frames = max(0, info.pending_frames - 1) + result.sent += 1 + except Exception: + logger.debug("Failed to send event %s to %s", event, conn_id) + + return result + + def _passes_filter( + self, + info: _ConnInfo, + event: str, + filter_: BroadcastFilter | None, + ) -> bool: + if info.subscribed_events: + if event not in info.subscribed_events and "*" not in info.subscribed_events: + return False + + if filter_ is None: + return True + + if filter_.conn_ids is not None and info.conn_id not in filter_.conn_ids: + return False + if filter_.user_ids is not None and info.user_id not in filter_.user_ids: + return False + if filter_.roles is not None: + user_has_role = any(r in filter_.roles for r in info.roles) + if not user_has_role: + return False + + return True + + @staticmethod + def _has_event_scope(info: _ConnInfo, event: str) -> bool: + min_role = EVENT_MIN_ROLE.get(event) + if min_role is None: + return True + + if not info.roles: + return min_role == GatewayRole.VIEWER + + max_level = max(ROLE_HIERARCHY.get(r, 0) for r in info.roles) + required_level = ROLE_HIERARCHY.get(min_role, 0) + return max_level >= required_level + + def pending_frame_count(self, conn_id: str) -> int: + info = self._connections.get(conn_id) + if info is None: + return 0 + return info.pending_frames + + def _resolve_state_version(self, event: str) -> int | None: + versioned_events = { + "channel.state_change", + "config.changed", + "plugin.installed", + "plugin.removed", + } + if event in versioned_events: + return self._global_seq + return None + + +gateway_broadcaster = GatewayBroadcaster() diff --git a/backend/package/yuxi/channel/gateway/channel_plugin.py b/backend/package/yuxi/channel/gateway/channel_plugin.py new file mode 100644 index 00000000..ec74fd8c --- /dev/null +++ b/backend/package/yuxi/channel/gateway/channel_plugin.py @@ -0,0 +1,192 @@ +import logging + +from yuxi.channel.runtime.backoff import BackoffConfig, ErrorBackoff +from yuxi.channel.extensions.base import BaseChannelPlugin +from yuxi.channel.gateway.client import GatewayClient + +logger = logging.getLogger(__name__) + +DEFAULT_RECONNECT_LEVELS = [1.0, 2.0, 4.0, 8.0, 15.0, 30.0] + + +class GatewayChannelPlugin(BaseChannelPlugin): + """使用 Gateway Client SDK 的渠道插件基类。 + + 封装常用连接模式： + - 自动重连（基于 ErrorBackoff 的状态管理） + - 请求重试（基于 ErrorBackoff.execute） + - 设备身份管理（自动生成/注册/刷新） + - 生命周期回调（on_connect / on_disconnect） + + 用法:: + + class MyChannelPlugin(GatewayChannelPlugin): + id = "my_channel" + name = "MyChannel" + + async def on_connect(self, client: GatewayClient) -> None: + await client.request("channel.start") + + async def on_event(self, client: GatewayClient, event) -> None: + ... + """ + + id: str = "" + name: str = "" + + def __init__( + self, + *, + gateway_url: str = "ws://127.0.0.1:9001", + token: str | None = None, + device_token: str | None = None, + request_timeout: float = 30.0, + reconnect_base_ms: float = 1000.0, + reconnect_max_ms: float = 30000.0, + reconnect_factor: float = 2.0, + reconnect_jitter: float = 0.1, + ): + self._gateway_url = gateway_url + self._token = token + self._device_token = device_token + self._request_timeout = request_timeout + self._reconnect_base_ms = reconnect_base_ms + self._reconnect_max_ms = reconnect_max_ms + self._reconnect_factor = reconnect_factor + self._reconnect_jitter = reconnect_jitter + + self._client: GatewayClient | None = None + self._backoff = ErrorBackoff( + config=BackoffConfig( + base_delay=reconnect_base_ms / 1000.0, + max_delay=reconnect_max_ms / 1000.0, + exponent=reconnect_factor, + jitter=True, + jitter_factor=reconnect_jitter, + max_retries=0, + ), + ) + + @property + def client(self) -> GatewayClient | None: + return self._client + + @property + def connected(self) -> bool: + return self._client is not None and self._client.connected + + async def start(self, ctx) -> object: + self._client = GatewayClient( + url=self._gateway_url, + token=self._token, + device_token=self._device_token, + request_timeout=self._request_timeout, + reconnect_base_ms=self._reconnect_base_ms, + reconnect_max_ms=self._reconnect_max_ms, + reconnect_factor=self._reconnect_factor, + reconnect_jitter=self._reconnect_jitter, + on_connect=self._on_gw_connect, + on_disconnect=self._on_gw_disconnect, + on_event=self._on_gw_event, + on_error=self._on_gw_error, + on_close=self._on_gw_close, + ) + + try: + await self._client.start() + except ValueError as e: + logger.error("[%s] Gateway client start failed: %s", self.id, e) + raise + + logger.info("[%s] Gateway channel plugin started", self.id) + return self._client + + async def stop(self, ctx) -> None: + if self._client: + await self._client.stop() + self._client = None + self._backoff.reset(f"ch-{self.id}") + logger.info("[%s] Gateway channel plugin stopped", self.id) + + async def send_request( + self, + method: str, + params: dict | None = None, + *, + timeout: float | None = None, + max_attempts: int = 3, + ) -> dict: + """发送 RPC 请求，带自动重试。""" + if self._client is None: + raise RuntimeError(f"[{self.id}] gateway client not started") + + return await self._client.request_with_retry( + method, + params=params, + timeout=timeout, + max_attempts=max_attempts, + ) + + async def init_device_identity(self) -> str: + """初始化设备身份，返回 device_token。""" + if self._client is None: + self._client = GatewayClient(url=self._gateway_url, token=self._token) + return await self._client.init_device_identity() + + async def refresh_device_token(self) -> str | None: + """刷新设备令牌。""" + if self._client is None: + return None + return await self._client.refresh_device_token() + + def create_backoff(self, name: str, levels: list[float] | None = None) -> ErrorBackoff: + """创建 ErrorBackoff 实例，供子类管理独立退避状态。""" + return ErrorBackoff( + config=BackoffConfig( + base_delay=levels[0] if levels else 1.0, + max_delay=levels[-1] if levels else 30.0, + exponent=2.0, + jitter=True, + jitter_factor=0.1, + max_retries=0, + ), + ) + + async def on_connect(self, client: GatewayClient) -> None: + """子类可重写：Gateway 连接成功时回调。""" + + async def on_disconnect(self, client: GatewayClient, code: int, reason: str) -> None: + """子类可重写：Gateway 断开连接时回调。""" + + async def on_event(self, client: GatewayClient, event) -> None: + """子类可重写：收到 Gateway 事件时回调。""" + + async def on_error(self, client: GatewayClient, error: Exception) -> None: + """子类可重写：Gateway 错误时回调。""" + + async def _on_gw_connect(self, client: GatewayClient) -> None: + try: + await self.on_connect(client) + except Exception: + logger.exception("[%s] on_connect callback failed", self.id) + + async def _on_gw_disconnect(self, client: GatewayClient, code: int, reason: str) -> None: + try: + await self.on_disconnect(client, code, reason) + except Exception: + logger.exception("[%s] on_disconnect callback failed", self.id) + + async def _on_gw_event(self, client: GatewayClient, event) -> None: + try: + await self.on_event(client, event) + except Exception: + logger.exception("[%s] on_event callback failed", self.id) + + async def _on_gw_error(self, client: GatewayClient, error: Exception) -> None: + try: + await self.on_error(client, error) + except Exception: + logger.exception("[%s] on_error callback failed", self.id) + + async def _on_gw_close(self, client: GatewayClient, code: int, reason: str) -> None: + logger.info("[%s] Gateway connection closed: code=%d reason=%s", self.id, code, reason) diff --git a/backend/package/yuxi/channel/gateway/client.py b/backend/package/yuxi/channel/gateway/client.py new file mode 100644 index 00000000..5f9deb3c --- /dev/null +++ b/backend/package/yuxi/channel/gateway/client.py @@ -0,0 +1,557 @@ +import asyncio +import logging +import time +from collections.abc import Awaitable, Callable + +import aiohttp + +from yuxi.channel.config.defaults import TIMEOUT +from yuxi.channel.runtime.backoff import ( + BackoffConfig, + ErrorBackoff, +) +from yuxi.channel.gateway.net_utils import is_secure_ws_url +from yuxi.channel.gateway.protocol import ( + GatewayErrorCode, + RpcEvent, + RpcRequest, + RpcResponse, + marshal_frame, + unmarshal_frame, +) + +logger = logging.getLogger(__name__) + +DEFAULT_GATEWAY_URL = "ws://127.0.0.1:9001" +DEFAULT_REQUEST_TIMEOUT = TIMEOUT.http.normal +DEFAULT_RECONNECT_BASE_MS = 1000.0 +DEFAULT_RECONNECT_MAX_MS = 30000.0 +DEFAULT_RECONNECT_FACTOR = 2.0 +DEFAULT_RECONNECT_JITTER = 0.1 + +OnConnectCallback = Callable[["GatewayClient"], Awaitable[None]] +OnDisconnectCallback = Callable[["GatewayClient", int, str], Awaitable[None]] +OnEventCallback = Callable[["GatewayClient", RpcEvent], Awaitable[None]] +OnErrorCallback = Callable[["GatewayClient", Exception], Awaitable[None]] +OnCloseCallback = Callable[["GatewayClient", int, str], Awaitable[None]] + + +class GatewayClientError(Exception): + def __init__(self, code: GatewayErrorCode | str, message: str, details: dict | None = None): + self.code = code + self.details = details or {} + super().__init__(message) + + +class _PendingRequest: + __slots__ = ("future", "timeout_handle") + future: asyncio.Future + timeout_handle: asyncio.TimerHandle | None + + def __init__(self): + self.future = asyncio.get_running_loop().create_future() + self.timeout_handle = None + + +class GatewayClient: + """统一 Gateway WebSocket 客户端 SDK。 + + 连接 → 认证 → 收发 RPC 请求/事件 → 自动重连。 + 对齐 OpenClaw GatewayClient 语义，适配项目现有 Gateway 协议。 + + 用法:: + + client = GatewayClient( + url="ws://127.0.0.1:9001", + token="my-token", + on_event=my_event_handler, + ) + await client.start() + + resp = await client.request("system.health") + ... + + await client.stop() + """ + + def __init__( + self, + *, + url: str = DEFAULT_GATEWAY_URL, + token: str | None = None, + password: str | None = None, + device_token: str | None = None, + request_timeout: float = DEFAULT_REQUEST_TIMEOUT, + reconnect_base_ms: float = DEFAULT_RECONNECT_BASE_MS, + reconnect_max_ms: float = DEFAULT_RECONNECT_MAX_MS, + reconnect_factor: float = DEFAULT_RECONNECT_FACTOR, + reconnect_jitter: float = DEFAULT_RECONNECT_JITTER, + on_connect: OnConnectCallback | None = None, + on_disconnect: OnDisconnectCallback | None = None, + on_event: OnEventCallback | None = None, + on_error: OnErrorCallback | None = None, + on_close: OnCloseCallback | None = None, + ): + self._url = url + self._token = token + self._password = password + self._device_token = device_token + self._request_timeout = max(1.0, request_timeout) + self._reconnect_base_ms = max(100.0, reconnect_base_ms) + self._reconnect_max_ms = max(reconnect_base_ms, reconnect_max_ms) + self._reconnect_factor = max(1.0, reconnect_factor) + self._reconnect_jitter = float(reconnect_jitter) + + self._on_connect = on_connect + self._on_disconnect = on_disconnect + self._on_event = on_event + self._on_error = on_error + self._on_close = on_close + + self._ws: aiohttp.ClientWebSocketResponse | None = None + self._session: aiohttp.ClientSession | None = None + self._pending: dict[str, _PendingRequest] = {} + self._running = False + self._close_event = asyncio.Event() + self._reconnect_task: asyncio.Task | None = None + + self._device_id: str | None = None + self._device_private_key_pem: str | None = None + + self._last_tick: float | None = None + self._tick_interval_ms = 30_000 + self._tick_watch_task: asyncio.Task | None = None + + self._reconnect_backoff = ErrorBackoff( + config=BackoffConfig( + base_delay=reconnect_base_ms / 1000.0, + max_delay=reconnect_max_ms / 1000.0, + exponent=reconnect_factor, + jitter=True, + jitter_factor=reconnect_jitter, + max_retries=0, + ), + ) + + @property + def connected(self) -> bool: + return self._ws is not None and not self._ws.closed + + @property + def url(self) -> str: + return self._url + + async def start(self) -> None: + """启动客户端，开始连接和自动重连循环。""" + if self._running: + return + self._running = True + self._close_event.clear() + self._validate_url() + self._reconnect_task = asyncio.create_task(self._reconnect_loop()) + + async def stop(self) -> None: + """停止客户端，断开连接并等待清理完成。""" + self._running = False + self._close_event.set() + self._stop_tick_watch() + self._cancel_all_pending(RuntimeError("gateway client stopped")) + await self._disconnect() + if self._session and not self._session.closed: + await self._session.close() + self._session = None + if self._reconnect_task: + self._reconnect_task.cancel() + try: + await self._reconnect_task + except asyncio.CancelledError: + pass + self._reconnect_task = None + + async def request(self, method: str, params: dict | None = None, timeout: float | None = None) -> dict: + """发送 RPC 请求，返回响应结果。 + + Raises: + GatewayClientError: 服务端返回错误 + RuntimeError: 未连接 + TimeoutError: 请求超时 + """ + if not self.connected: + raise RuntimeError("gateway not connected") + + req = RpcRequest(method=method, params=params) + frame = marshal_frame(req) + request_id = req.id + effective_timeout = timeout if timeout is not None else self._request_timeout + + pending = _PendingRequest() + self._pending[request_id] = pending + + timeout_coro: asyncio.Task | None = None + try: + timeout_coro = asyncio.create_task(asyncio.sleep(effective_timeout)) + send_task = asyncio.create_task(self._ws.send_str(frame)) + + done, _pending_set = await asyncio.wait( + [pending.future, send_task, timeout_coro], + return_when=asyncio.FIRST_COMPLETED, + ) + for task in _pending_set: + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + + if timeout_coro in done and not pending.future.done(): + raise TimeoutError(f"gateway request timeout for {method}") + + if send_task in done: + send_exc = send_task.exception() + if send_exc: + raise send_exc + + if pending.future in done: + exc = pending.future.exception() + if exc: + raise exc + return pending.future.result() + + raise RuntimeError(f"gateway request failed for {method}: unexpected state") + + finally: + self._pending.pop(request_id, None) + if timeout_coro and not timeout_coro.done(): + timeout_coro.cancel() + try: + await timeout_coro + except asyncio.CancelledError: + pass + + async def request_with_retry( + self, + method: str, + params: dict | None = None, + *, + timeout: float | None = None, + max_attempts: int = 3, + ) -> dict: + """发送 RPC 请求，失败时自动重试。 + + Raises: + GatewayClientError: 服务端返回错误 + RuntimeError: 未连接 + TimeoutError: 请求超时 + """ + backoff = ErrorBackoff( + config=BackoffConfig( + base_delay=0.4, + max_delay=5.0, + exponent=2.0, + jitter=True, + jitter_factor=0.1, + max_retries=max_attempts - 1, + ), + ) + + async def _do_request() -> dict: + return await self.request(method, params=params, timeout=timeout) + + return await backoff.execute(f"gw-rpc:{method}", _do_request) + + async def init_device_identity(self) -> str: + """生成并注册设备身份，返回 device_token。 + + 首次调用生成 Ed25519 密钥对，注册到设备注册表， + 并构建可用于认证的 device_token。 + """ + from yuxi.channel.gateway.device_auth import build_device_token as _build_token + from yuxi.channel.gateway.device_identity import generate_device_identity + from yuxi.channel.gateway.device_registry import register_device + + device_id, public_key_pem, private_key_pem = generate_device_identity() + self._device_id = device_id + self._device_private_key_pem = private_key_pem + + register_device(device_id, public_key_pem) + + token = self._build_device_token_internal() + self._device_token = token + + logger.info("GatewayClient device identity initialized: device_id=%s", device_id) + return token + + async def refresh_device_token(self) -> str | None: + """刷新设备令牌（重新签名时间戳），返回新的 device_token。""" + if not self._device_id or not self._device_private_key_pem: + return None + token = self._build_device_token_internal() + self._device_token = token + return token + + @property + def device_id(self) -> str | None: + return self._device_id + + # ---- internal ---- + + def _validate_url(self) -> None: + if not is_secure_ws_url(self._url): + raise ValueError( + f"不安全连接: 不允许通过明文 ws:// 连接到非回环地址 " + f"(CWE-319)。请使用 wss:// 或将 gateway.bind 设为 loopback。" + f" 当前 URL: {self._url}" + ) + + async def _ensure_session(self) -> None: + if self._session is None or self._session.closed: + self._session = aiohttp.ClientSession() + + async def _connect(self) -> bool: + await self._ensure_session() + headers = {} + auth_value = self._resolve_auth_header() + if auth_value: + headers["Authorization"] = auth_value + + try: + self._ws = await self._session.ws_connect(self._url, headers=headers, heartbeat=30.0) + except aiohttp.ClientError as e: + logger.warning("Gateway client connect failed: %s", e) + await self._fire_error(e) + return False + except Exception as e: + logger.warning("Gateway client connect error: %s", e) + await self._fire_error(e) + return False + + try: + request_id = f"auth-check-{id(self):x}" + await asyncio.wait_for( + self._ws.send_str(marshal_frame(RpcRequest(id=request_id, method="system.health", params={}))), + timeout=TIMEOUT.gateway.connect, + ) + auth_ok = False + async for msg in self._ws: + if msg.type == aiohttp.WSMsgType.TEXT: + frame = unmarshal_frame(msg.data) + if isinstance(frame, RpcResponse) and frame.id == request_id: + auth_ok = frame.ok + break + elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR): + break + if not auth_ok: + logger.warning("Gateway client auth check failed") + await self._ws.close(code=4001, message=b"auth failed") + self._ws = None + return False + return True + except Exception as e: + logger.warning("Gateway client auth check error: %s", e) + ws = self._ws + self._ws = None + if ws and not ws.closed: + try: + await ws.close() + except Exception: + pass + return False + + def _resolve_auth_header(self) -> str | None: + if self._device_token: + return f"Bearer {self._device_token}" + if self._password: + return f"Bearer {self._password}" + if self._token: + return f"Bearer {self._token}" + return None + + async def _disconnect(self) -> None: + self._stop_tick_watch() + self._cancel_all_pending(RuntimeError("gateway disconnected")) + ws = self._ws + self._ws = None + if ws and not ws.closed: + try: + await ws.close(code=1000, message=b"client shutdown") + except Exception: + pass + + def _start_tick_watch(self) -> None: + self._last_tick = None + + async def _watch(): + try: + while True: + await asyncio.sleep(self._tick_interval_ms / 1000.0 * 2) + if self._last_tick is None: + continue + gap_ms = (asyncio.get_event_loop().time() - self._last_tick) * 1000 + if gap_ms > self._tick_interval_ms * 3: + logger.warning("Gateway tick timeout: gap=%.0fms", gap_ms) + await self._disconnect() + break + except asyncio.CancelledError: + pass + + self._tick_watch_task = asyncio.create_task(_watch()) + + def _stop_tick_watch(self) -> None: + if self._tick_watch_task and not self._tick_watch_task.done(): + self._tick_watch_task.cancel() + self._tick_watch_task = None + self._last_tick = None + + async def _reconnect_loop(self) -> None: + backoff_ms = self._reconnect_base_ms + + while self._running: + connected = await self._connect() + if not connected: + delay = self._reconnect_backoff.config.compute_delay(0) + try: + await asyncio.wait_for(self._close_event.wait(), timeout=delay) + return + except TimeoutError: + backoff_ms = min(backoff_ms * self._reconnect_factor, self._reconnect_max_ms) + continue + + backoff_ms = self._reconnect_base_ms + self._start_tick_watch() + await self._fire_connect() + + try: + await self._recv_loop() + except aiohttp.ClientError as e: + logger.warning("Gateway client recv error: %s", e) + await self._fire_error(e) + except asyncio.CancelledError: + return + except Exception as e: + logger.exception("Gateway client recv unexpected error") + await self._fire_error(e) + + close_code = 1006 + close_reason = "connection lost" + if self._ws is not None: + close_code = self._ws.close_code or close_code + close_reason = self._ws._close_reason or close_reason + + await self._disconnect() + await self._fire_close(close_code, close_reason) + + if not self._running: + return + + delay = self._reconnect_backoff.config.compute_delay(0) + try: + await asyncio.wait_for(self._close_event.wait(), timeout=delay) + return + except TimeoutError: + backoff_ms = min(backoff_ms * self._reconnect_factor, self._reconnect_max_ms) + + async def _recv_loop(self) -> None: + ws = self._ws + if ws is None: + return + + async for msg in ws: + if msg.type == aiohttp.WSMsgType.TEXT: + try: + frame = unmarshal_frame(msg.data) + except Exception: + logger.warning("Gateway client: invalid frame received") + continue + + if isinstance(frame, RpcResponse): + self._handle_response(frame) + elif isinstance(frame, RpcEvent): + self._handle_event(frame) + + elif msg.type == aiohttp.WSMsgType.CLOSED: + break + elif msg.type == aiohttp.WSMsgType.ERROR: + break + + def _handle_response(self, resp: RpcResponse) -> None: + pending = self._pending.get(resp.id) + if pending is None: + return + self._pending.pop(resp.id, None) + + if resp.ok: + pending.future.set_result(resp.result or {}) + else: + code = resp.error_code or GatewayErrorCode.INTERNAL_ERROR + msg = resp.error_message or "unknown error" + err = GatewayClientError(code=code, message=msg) + pending.future.set_exception(err) + + def _handle_event(self, evt: RpcEvent) -> None: + if evt.event == "tick": + self._last_tick = asyncio.get_event_loop().time() + if self._on_event: + asyncio.create_task(self._fire_event(evt)) + + def _cancel_all_pending(self, exc: Exception) -> None: + for pending in list(self._pending.values()): + if not pending.future.done(): + pending.future.set_exception(exc) + self._pending.clear() + + def _compute_reconnect_levels(self) -> list[float]: + levels = [] + current = self._reconnect_base_ms / 1000.0 + for _ in range(10): + if current >= self._reconnect_max_ms / 1000.0: + break + levels.append(current) + current *= self._reconnect_factor + if not levels or levels[-1] < self._reconnect_max_ms / 1000.0: + levels.append(self._reconnect_max_ms / 1000.0) + return levels + + def _build_device_token_internal(self) -> str: + from yuxi.channel.gateway.device_auth import build_device_token as _build_token, build_challenge + from yuxi.channel.gateway.device_identity import sign_challenge + + timestamp_ms = int(time.time() * 1000) + challenge = build_challenge(self._device_id, timestamp_ms) + signature = sign_challenge(self._device_private_key_pem, challenge) + return _build_token(self._device_id, timestamp_ms, signature) + + async def _fire_connect(self) -> None: + if self._on_connect: + try: + await self._on_connect(self) + except Exception: + logger.exception("on_connect callback failed") + + async def _fire_close(self, code: int, reason: str) -> None: + await self._fire_disconnect(code, reason) + if self._on_close: + try: + await self._on_close(self, code, reason) + except Exception: + logger.exception("on_close callback failed") + + async def _fire_disconnect(self, code: int, reason: str) -> None: + if self._on_disconnect: + try: + await self._on_disconnect(self, code, reason) + except Exception: + logger.exception("on_disconnect callback failed") + + async def _fire_event(self, evt: RpcEvent) -> None: + if self._on_event: + try: + await self._on_event(self, evt) + except Exception: + logger.exception("on_event callback failed") + + async def _fire_error(self, err: Exception) -> None: + if self._on_error: + try: + await self._on_error(self, err) + except Exception: + logger.exception("on_error callback failed") diff --git a/backend/package/yuxi/channel/gateway/device_auth.py b/backend/package/yuxi/channel/gateway/device_auth.py new file mode 100644 index 00000000..8c19efd7 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/device_auth.py @@ -0,0 +1,108 @@ +import base64 +import logging +import time + +from yuxi.channel.gateway.auth import GatewayAuthMode, GatewayAuthResult +from yuxi.channel.gateway.device_identity import verify_signature +from yuxi.channel.gateway.rbac import GatewayRole + +logger = logging.getLogger(__name__) + +DEVICE_TOKEN_PREFIX = "dv." +DEVICE_TOKEN_TIMESTAMP_TOLERANCE_SECONDS = 300 + + +def parse_device_token(token: str) -> tuple[str, int, bytes] | None: + """解析设备令牌，格式: dv.{base64(device_id)}.{base64(timestamp_ms)}.{base64(signature)}""" + if not token.startswith(DEVICE_TOKEN_PREFIX): + return None + parts = token[len(DEVICE_TOKEN_PREFIX) :].split(".") + if len(parts) != 3: + return None + try: + def _add_padding(s: str) -> str: + missing = (4 - len(s) % 4) % 4 + return s + "=" * missing if missing else s + + device_id = base64.urlsafe_b64decode(_add_padding(parts[0])).decode("utf-8") + timestamp_ms = int(base64.urlsafe_b64decode(_add_padding(parts[1])).decode("utf-8")) + signature = base64.urlsafe_b64decode(_add_padding(parts[2])) + except (ValueError, UnicodeDecodeError, base64.binascii.Error): + return None + return device_id, timestamp_ms, signature + + +def build_device_token(device_id: str, timestamp_ms: int, signature: bytes) -> str: + """构建设备令牌""" + device_id_b64 = base64.urlsafe_b64encode(device_id.encode("utf-8")).rstrip(b"=").decode("ascii") + ts_b64 = base64.urlsafe_b64encode(str(timestamp_ms).encode("utf-8")).rstrip(b"=").decode("ascii") + sig_b64 = base64.urlsafe_b64encode(signature).rstrip(b"=").decode("ascii") + return f"{DEVICE_TOKEN_PREFIX}{device_id_b64}.{ts_b64}.{sig_b64}" + + +def build_challenge(device_id: str, timestamp_ms: int) -> bytes: + return f"{device_id}:{timestamp_ms}".encode("utf-8") + + +def validate_timestamp(timestamp_ms: int) -> bool: + now_ms = int(time.time() * 1000) + if timestamp_ms > now_ms: + return False + diff_ms = now_ms - timestamp_ms + return diff_ms <= DEVICE_TOKEN_TIMESTAMP_TOLERANCE_SECONDS * 1000 + + +async def authenticate_device( + token: str, + lookup_public_key, + client_ip: str | None = None, +) -> GatewayAuthResult: + parsed = parse_device_token(token) + if parsed is None: + return GatewayAuthResult( + authenticated=False, + mode=GatewayAuthMode.DEVICE_TOKEN, + error="设备令牌格式无效。预期格式: dv.{deviceId}.{ts}.{sig}", + ) + + device_id, timestamp_ms, signature = parsed + + if not validate_timestamp(timestamp_ms): + return GatewayAuthResult( + authenticated=False, + mode=GatewayAuthMode.DEVICE_TOKEN, + error="设备令牌时间戳超出容忍范围（±5 分钟）。", + metadata={"device_id": device_id}, + ) + + public_key_pem = await lookup_public_key(device_id) + if public_key_pem is None: + return GatewayAuthResult( + authenticated=False, + mode=GatewayAuthMode.DEVICE_TOKEN, + error=f"未知设备: {device_id}", + metadata={"device_id": device_id}, + ) + + challenge = build_challenge(device_id, timestamp_ms) + if not verify_signature(public_key_pem, challenge, signature): + return GatewayAuthResult( + authenticated=False, + mode=GatewayAuthMode.DEVICE_TOKEN, + error="设备签名验证失败。", + metadata={"device_id": device_id}, + ) + + logger.info( + "Device authenticated: device_id=%s ip=%s", + device_id, + client_ip or "unknown", + ) + + return GatewayAuthResult( + authenticated=True, + user_id=f"device:{device_id}", + mode=GatewayAuthMode.DEVICE_TOKEN, + metadata={"device_id": device_id}, + roles=[GatewayRole.OPERATOR], + ) diff --git a/backend/package/yuxi/channel/gateway/device_identity.py b/backend/package/yuxi/channel/gateway/device_identity.py new file mode 100644 index 00000000..c2ea8c95 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/device_identity.py @@ -0,0 +1,68 @@ +import hashlib +import logging + +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey + +logger = logging.getLogger(__name__) + +DEVICE_ID_HEX_LENGTH = 16 + + +def generate_device_identity() -> tuple[str, str, str]: + """生成 Ed25519 设备身份。 + + Returns: + (device_id, public_key_pem, private_key_pem) + """ + private_key = Ed25519PrivateKey.generate() + public_key = private_key.public_key() + + public_key_pem = public_key.public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo, + ).decode("utf-8") + + private_key_pem = private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.BestAvailableEncryption(b"yuxi-device-key"), + ).decode("utf-8") + + device_id = _derive_device_id_from_key(public_key) + + logger.info("Generated device identity: device_id=%s", device_id) + return device_id, public_key_pem, private_key_pem + + +def derive_device_id(public_key_pem: str) -> str: + public_key = serialization.load_pem_public_key(public_key_pem.encode("utf-8")) + return _derive_device_id_from_key(public_key) + + +_PRIVATE_KEY_PASSWORD = b"yuxi-device-key" + + +def sign_challenge(private_key_pem: str, challenge: bytes) -> bytes: + private_key = serialization.load_pem_private_key( + private_key_pem.encode("utf-8"), + password=_PRIVATE_KEY_PASSWORD, + ) + return private_key.sign(challenge) + + +def verify_signature(public_key_pem: str, challenge: bytes, signature: bytes) -> bool: + try: + public_key = serialization.load_pem_public_key(public_key_pem.encode("utf-8")) + public_key.verify(signature, challenge) + return True + except Exception: + return False + + +def _derive_device_id_from_key(public_key) -> str: + public_key_raw = public_key.public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw, + ) + return hashlib.sha256(public_key_raw).hexdigest()[:DEVICE_ID_HEX_LENGTH] diff --git a/backend/package/yuxi/channel/gateway/device_registry.py b/backend/package/yuxi/channel/gateway/device_registry.py new file mode 100644 index 00000000..f21cfb75 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/device_registry.py @@ -0,0 +1,61 @@ +import logging +import threading + +logger = logging.getLogger(__name__) + +_registry: dict[str, str] = {} +_registry_lock = threading.Lock() + + +def register_device(device_id: str, public_key_pem: str) -> None: + with _registry_lock: + _registry[device_id] = public_key_pem + logger.info("Device registered: device_id=%s", device_id) + try: + from yuxi.channel.gateway.device_registry_db import db_save_device + + import asyncio + + try: + loop = asyncio.get_running_loop() + loop.create_task(db_save_device(device_id, public_key_pem)) + except RuntimeError: + pass + except Exception: + logger.debug("DB device save skipped: device_id=%s", device_id) + + +def unregister_device(device_id: str) -> bool: + with _registry_lock: + if device_id in _registry: + del _registry[device_id] + logger.info("Device unregistered: device_id=%s", device_id) + return True + return False + + +async def lookup_public_key(device_id: str) -> str | None: + with _registry_lock: + key = _registry.get(device_id) + if key is not None: + return key + key = await _db_lookup(device_id) + if key is not None: + with _registry_lock: + _registry[device_id] = key + return key + + +async def _db_lookup(device_id: str) -> str | None: + try: + from yuxi.channel.gateway.device_registry_db import db_lookup_public_key + + return await db_lookup_public_key(device_id) + except Exception: + logger.debug("DB device lookup skipped: device_id=%s", device_id) + return None + + +def list_devices() -> list[str]: + with _registry_lock: + return list(_registry.keys()) diff --git a/backend/package/yuxi/channel/gateway/device_registry_db.py b/backend/package/yuxi/channel/gateway/device_registry_db.py new file mode 100644 index 00000000..008f0b49 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/device_registry_db.py @@ -0,0 +1,31 @@ +import logging + +from yuxi.repositories.channel_device_identity_repo import DeviceIdentityRepository + +logger = logging.getLogger(__name__) + +_repo = DeviceIdentityRepository() + + +async def db_lookup_public_key(device_id: str) -> str | None: + record = await _repo.get_by_device_id(device_id) + if record is None or record.status != "active": + return None + + await _repo.mark_used(device_id) + return record.public_key_pem + + +async def db_save_device(device_id: str, public_key_pem: str) -> None: + existing = await _repo.get_by_device_id(device_id) + if existing is not None: + await _repo.update(device_id, {"public_key_pem": public_key_pem, "status": "active"}) + else: + await _repo.create( + { + "device_id": device_id, + "public_key_pem": public_key_pem, + "status": "active", + } + ) + logger.debug("Device saved to DB: device_id=%s", device_id) diff --git a/backend/package/yuxi/channel/gateway/lanes.py b/backend/package/yuxi/channel/gateway/lanes.py new file mode 100644 index 00000000..0ca3d662 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/lanes.py @@ -0,0 +1,114 @@ +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +logger = logging.getLogger(__name__) + +DEFAULT_MAX_LANES = 5 +DEFAULT_GLOBAL_MAX_LANES = 50 + + +class ChannelLane: + """单个渠道账户的并发通道。 + + 每个账户一个独立 lane，控制该账户下同时运行的 Agent 数量。 + """ + + def __init__(self, max_lanes: int = DEFAULT_MAX_LANES): + self._semaphore = asyncio.Semaphore(max_lanes) + self._active_runs: int = 0 + self._busy: bool = False + self._max_lanes = max_lanes + + @asynccontextmanager + async def acquire(self) -> AsyncIterator[None]: + async with self._semaphore: + self._active_runs += 1 + if self._active_runs >= self._max_lanes: + self._busy = True + try: + yield + finally: + self._active_runs = max(0, self._active_runs - 1) + if self._active_runs < self._max_lanes: + self._busy = False + + @property + def active_runs(self) -> int: + return self._active_runs + + @property + def is_busy(self) -> bool: + return self._busy + + @property + def available_slots(self) -> int: + return max(0, self._max_lanes - self._active_runs) + + +class LaneManager: + """并发通道管理器。 + + 按账户维度管理并发通道，同时提供全局上限保护。 + """ + + def __init__( + self, + per_account_max: int = DEFAULT_MAX_LANES, + global_max: int = DEFAULT_GLOBAL_MAX_LANES, + ): + self._per_account_max = per_account_max + self._lanes: dict[str, ChannelLane] = {} + self._global_semaphore = asyncio.Semaphore(global_max) + + def _key(self, channel_type: str, account_id: str) -> str: + return f"{channel_type}:{account_id}" + + def get_lane(self, channel_type: str, account_id: str) -> ChannelLane: + key = self._key(channel_type, account_id) + if key not in self._lanes: + self._lanes[key] = ChannelLane(max_lanes=self._per_account_max) + return self._lanes[key] + + @asynccontextmanager + async def run_with_lane(self, channel_type: str, account_id: str) -> AsyncIterator[ChannelLane]: + lane = self.get_lane(channel_type, account_id) + if lane.is_busy: + logger.warning( + "Channel lane busy: %s/%s, active=%d", + channel_type, + account_id, + lane.active_runs, + ) + async with lane.acquire(): + async with self._global_semaphore: + yield lane + + def remove_lane(self, channel_type: str, account_id: str) -> None: + key = self._key(channel_type, account_id) + self._lanes.pop(key, None) + + def get_stats(self) -> dict[str, dict]: + return { + key: { + "active_runs": lane.active_runs, + "busy": lane.is_busy, + "available_slots": lane.available_slots, + } + for key, lane in self._lanes.items() + } + + def get_total_active(self) -> int: + return sum(lane.active_runs for lane in self._lanes.values()) + + def cleanup_idle(self) -> int: + idle_keys = [key for key, lane in self._lanes.items() if lane.active_runs == 0] + for key in idle_keys: + del self._lanes[key] + if idle_keys: + logger.debug("LaneManager cleaned up %d idle lanes", len(idle_keys)) + return len(idle_keys) + + +lane_manager = LaneManager() diff --git a/backend/package/yuxi/channel/gateway/net_utils.py b/backend/package/yuxi/channel/gateway/net_utils.py new file mode 100644 index 00000000..d10af6bb --- /dev/null +++ b/backend/package/yuxi/channel/gateway/net_utils.py @@ -0,0 +1,373 @@ +import asyncio +import ipaddress +import logging +import os +import socket +from enum import StrEnum +from urllib.parse import urlparse + +logger = logging.getLogger(__name__) + +_LOOPBACK_HOSTS = frozenset(["localhost", "127.0.0.1", "::1"]) + +_PRIVATE_IP_NETS = [ + ipaddress.ip_network("10.0.0.0/8"), + ipaddress.ip_network("172.16.0.0/12"), + ipaddress.ip_network("192.168.0.0/16"), + ipaddress.ip_network("127.0.0.0/8"), + ipaddress.ip_network("::1/128"), + ipaddress.ip_network("fc00::/7"), + ipaddress.ip_network("fe80::/10"), +] + +_TAILNET_IPV4_NET = ipaddress.ip_network("100.64.0.0/10") + +_ENV_ALLOW_INSECURE_PRIVATE_WS = "YUXI_ALLOW_INSECURE_PRIVATE_WS" + + +class GatewayBindMode(StrEnum): + LOOPBACK = "loopback" + LAN = "lan" + TAILNET = "tailnet" + AUTO = "auto" + CUSTOM = "custom" + + +def is_loopback_address(ip: str | None) -> bool: + if not ip: + return False + try: + return ipaddress.ip_address(ip.strip()).is_loopback + except ValueError: + return False + + +def is_trusted_proxy_address(ip: str | None, trusted_proxies: list[str] | None) -> bool: + if not ip or not trusted_proxies: + return False + try: + addr = ipaddress.ip_address(ip.strip()) + except ValueError: + return False + for proxy in trusted_proxies: + candidate = proxy.strip() + if not candidate: + continue + try: + net = ipaddress.ip_network(candidate, strict=False) + except ValueError: + continue + if addr in net: + return True + return False + + +TAILSCALE_TRUSTED_PROXIES = ["127.0.0.1", "::1"] + + +def _parse_ip_literal(raw: str | None) -> str | None: + if not raw: + return None + trimmed = raw.strip() + if not trimmed: + return None + if trimmed.startswith("[") and "]" in trimmed: + trimmed = trimmed[1 : trimmed.index("]")] + if ":" in trimmed and "." in trimmed: + last_colon = trimmed.rfind(":") + candidate = trimmed[:last_colon] + try: + ipaddress.IPv4Address(candidate) + trimmed = candidate + except ValueError: + pass + try: + ipaddress.ip_address(trimmed) + return trimmed + except ValueError: + return None + + +def resolve_forwarded_client_ip( + forwarded_for: str | None, + trusted_proxies: list[str] | None, +) -> str | None: + if not trusted_proxies: + return None + chain: list[str] = [] + for entry in (forwarded_for or "").split(","): + normalized = _parse_ip_literal(entry) + if normalized: + chain.append(normalized) + if not chain: + return None + for hop in reversed(chain): + if is_loopback_address(hop): + continue + if not is_trusted_proxy_address(hop, trusted_proxies): + return hop + return None + + +def resolve_client_ip( + remote_addr: str | None, + forwarded_for: str | None = None, + real_ip: str | None = None, + trusted_proxies: list[str] | None = None, + allow_real_ip_fallback: bool = False, +) -> str | None: + remote = _parse_ip_literal(remote_addr) + if not remote: + return None + if not is_trusted_proxy_address(remote, trusted_proxies): + return remote + forwarded = resolve_forwarded_client_ip(forwarded_for, trusted_proxies) + if forwarded: + return forwarded + if allow_real_ip_fallback: + return _parse_ip_literal(real_ip) + return None + + +def has_forwarded_request_headers(headers: dict) -> bool: + return bool( + headers.get("forwarded") + or headers.get("x-forwarded-for") + or headers.get("x-forwarded-proto") + or headers.get("x-real-ip") + or headers.get("x-forwarded-host") + ) + + +def is_local_direct_request( + remote_addr: str | None, + headers: dict | None = None, +) -> bool: + if not remote_addr: + return False + if headers and has_forwarded_request_headers(headers): + return False + return is_loopback_address(remote_addr) + + +def has_tailscale_proxy_headers(headers: dict | None) -> bool: + if not headers: + return False + return bool(headers.get("x-forwarded-for") and headers.get("x-forwarded-proto") and headers.get("x-forwarded-host")) + + +def is_tailscale_proxy_request( + remote_addr: str | None, + headers: dict | None = None, +) -> bool: + if not remote_addr: + return False + return is_loopback_address(remote_addr) and has_tailscale_proxy_headers(headers) + + +def resolve_tailscale_client_ip( + remote_addr: str | None, + headers: dict | None = None, +) -> str | None: + return resolve_client_ip( + remote_addr=remote_addr, + forwarded_for=headers.get("x-forwarded-for") if headers else None, + trusted_proxies=list(TAILSCALE_TRUSTED_PROXIES), + ) + + +def is_loopback_host(host: str) -> bool: + host = host.strip().lower().rstrip(".") + if not host: + return False + if host in _LOOPBACK_HOSTS: + return True + try: + addr = ipaddress.ip_address(host) + except ValueError: + return False + return addr.is_loopback + + +def is_private_host(host: str) -> bool: + host = host.strip().lower().rstrip(".") + if not host: + return False + try: + addr = ipaddress.ip_address(host) + except ValueError: + return False + return addr.is_private + + +def is_private_or_loopback_host(host: str) -> bool: + host = host.strip().lower().rstrip(".") + if not host: + return False + try: + addr = ipaddress.ip_address(host) + except ValueError: + return False + return addr.is_private or addr.is_loopback or addr.is_link_local + + +def is_localish_host(host: str | None) -> bool: + if not host: + return False + host = host.strip().lower().rstrip(".") + return is_loopback_host(host) or host.endswith(".ts.net") + + +def is_secure_ws_url(url: str, allow_private_ws: bool = False) -> bool: + try: + parsed = urlparse(url) + except ValueError: + return False + + protocol = parsed.scheme.lower() + if protocol == "wss": + return True + if protocol not in ("ws", "http"): + return False + + hostname = _extract_ws_hostname(parsed) + if not hostname: + return False + + if is_loopback_host(hostname): + return True + + if allow_private_ws and is_private_or_loopback_host(hostname): + return True + + return False + + +def _extract_ws_hostname(parsed) -> str: + netloc = parsed.netloc or parsed.hostname or "" + if "@" in netloc: + netloc = netloc.rsplit("@", 1)[-1] + if "[" in netloc and "]" in netloc: + start = netloc.index("[") + 1 + end = netloc.index("]") + return netloc[start:end] + if ":" in netloc: + return netloc.rsplit(":", 1)[0] + return netloc + + +async def can_bind_to_host(host: str) -> bool: + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, _sync_can_bind_to_host, host) + + +def _sync_can_bind_to_host(host: str) -> bool: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind((host, 0)) + return True + except OSError: + return False + finally: + sock.close() + + +_can_bind_cache: dict[str, bool] = {} + + +async def can_bind_to_host_cached(host: str) -> bool: + cached = _can_bind_cache.get(host) + if cached is not None: + return cached + result = await can_bind_to_host(host) + _can_bind_cache[host] = result + return result + + +def _is_container_environment() -> bool: + if os.path.exists("/.dockerenv"): + return True + try: + with open("/proc/1/cgroup") as f: + content = f.read() + if "docker" in content or "kubepods" in content: + return True + except OSError: + pass + return False + + +def _pick_primary_tailnet_ipv4() -> str | None: + try: + for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET): + ip = info[4][0] + try: + addr = ipaddress.ip_address(ip) + except ValueError: + continue + if addr in _TAILNET_IPV4_NET: + return ip + except OSError: + pass + return None + + +async def resolve_gateway_bind_host( + mode: GatewayBindMode | None = None, + custom_host: str | None = None, +) -> str: + mode = mode or GatewayBindMode.LOOPBACK + + if mode == GatewayBindMode.LOOPBACK: + if await can_bind_to_host("127.0.0.1"): + return "127.0.0.1" + return "0.0.0.0" + + if mode == GatewayBindMode.TAILNET: + tailnet_ip = _pick_primary_tailnet_ipv4() + if tailnet_ip and await can_bind_to_host(tailnet_ip): + return tailnet_ip + if await can_bind_to_host("127.0.0.1"): + return "127.0.0.1" + return "0.0.0.0" + + if mode == GatewayBindMode.LAN: + return "0.0.0.0" + + if mode == GatewayBindMode.CUSTOM: + host = (custom_host or "").strip() + if not host: + return "0.0.0.0" + try: + ipaddress.ip_address(host) + except ValueError: + logger.warning("gateway bind=custom: invalid IP '%s', falling back to 0.0.0.0", host) + return "0.0.0.0" + if await can_bind_to_host(host): + return host + logger.warning("gateway bind=custom: cannot bind '%s', falling back to 0.0.0.0", host) + return "0.0.0.0" + + if mode == GatewayBindMode.AUTO: + if _is_container_environment(): + return "0.0.0.0" + if await can_bind_to_host("127.0.0.1"): + return "127.0.0.1" + return "0.0.0.0" + + return "0.0.0.0" + + +def build_ws_security_error(display_host: str) -> str: + allow_private = os.environ.get(_ENV_ALLOW_INSECURE_PRIVATE_WS) == "1" + msg = ( + f'SECURITY ERROR: Cannot connect to "{display_host}" over plaintext ws://. ' + "Both credentials and chat data would be exposed to network interception. " + "Use wss:// for remote URLs. Safe defaults: keep gateway.bind=loopback and " + "connect via SSH tunnel " + "(ssh -N -L 18789:127.0.0.1:18789 user@gateway-host), or use Tailscale Serve/Funnel." + ) + if not allow_private: + msg += f" Break-glass (trusted private networks only): set {_ENV_ALLOW_INSECURE_PRIVATE_WS}=1." + return msg diff --git a/backend/package/yuxi/channel/gateway/openai_adapter.py b/backend/package/yuxi/channel/gateway/openai_adapter.py new file mode 100644 index 00000000..f4526338 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/openai_adapter.py @@ -0,0 +1,161 @@ +"""OpenAI API 兼容层 — 请求格式转换与 SS 流式响应""" + +from __future__ import annotations + +import json +import time +from collections.abc import AsyncIterator + +OPENAI_CHAT_COMPLETIONS_PATH = "/v1/chat/completions" +OPENAI_MODELS_PATH = "/v1/models" + + +class OpenAIMessageConverter: + @staticmethod + def extract_system_prompt(messages: list[dict]) -> str | None: + for msg in messages: + if msg.get("role") == "system": + return msg.get("content", "").strip() or None + return None + + @staticmethod + def extract_user_query(messages: list[dict]) -> tuple[str, str | None]: + user_texts: list[str] = [] + image_base64: str | None = None + + for msg in messages: + if msg.get("role") != "user": + continue + content = msg.get("content", "") + if isinstance(content, str): + user_texts.append(content.strip()) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + if part.get("type") == "text": + user_texts.append(part.get("text", "").strip()) + elif part.get("type") == "image_url": + url = part.get("image_url", {}).get("url", "") + if url.startswith("data:"): + image_base64 = url.split(",", 1)[1] if "," in url else url + + return "\n".join(t for t in user_texts if t), image_base64 + + @staticmethod + def extract_conversation_history(messages: list[dict]) -> list[dict]: + history: list[dict] = [] + for msg in messages: + role = msg.get("role", "") + content = msg.get("content", "") + if role in ("user", "assistant") and content: + if isinstance(content, str) and content.strip(): + history.append({"role": role, "content": content.strip()}) + return history + + @staticmethod + def extract_model_name(request_data: dict) -> str: + return request_data.get("model", "default") + + @staticmethod + def extract_stream_flag(request_data: dict) -> bool: + return bool(request_data.get("stream", False)) + + @staticmethod + def extract_max_tokens(request_data: dict) -> int | None: + return request_data.get("max_tokens") + + @staticmethod + def extract_temperature(request_data: dict) -> float | None: + return request_data.get("temperature") + + +def _format_sse_chunk( + content: str | None = None, + status: str = "streaming", + *, + finish_reason: str | None = None, + model: str = "default", + index: int = 0, +) -> bytes: + delta: dict = {} + if content is not None: + delta["content"] = content + if status is not None: + payload: dict = { + "id": f"chatcmpl-{int(time.time() * 1000)}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [ + { + "index": index, + "delta": delta, + "finish_reason": finish_reason, + } + ], + } + return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n".encode() + + return b"data: [DONE]\n\n" + + +def _format_openai_non_stream_response( + content: str, + *, + model: str = "default", + finish_reason: str = "stop", +) -> dict: + return { + "id": f"chatcmpl-{int(time.time() * 1000)}", + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": content, + }, + "finish_reason": finish_reason, + } + ], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + }, + } + + +async def format_sse_stream( + raw_stream: AsyncIterator[bytes], + *, + model: str = "default", +) -> AsyncIterator[bytes]: + accumulated: list[str] = [] + async for raw in raw_stream: + try: + data = json.loads(raw.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + + status = data.get("status", "") + content = data.get("response", "") + + if status == "streaming" and content: + accumulated.append(content) + yield _format_sse_chunk(content=content, status="streaming", model=model) + elif status == "finished": + pass + elif status == "error": + error_msg = data.get("error_message", content or "未知错误") + yield _format_sse_chunk( + content=error_msg, + status="error", + finish_reason="error", + model=model, + ) + return + + yield _format_sse_chunk(content="", status=None, finish_reason="stop", model=model) diff --git a/backend/package/yuxi/channel/gateway/polling.py b/backend/package/yuxi/channel/gateway/polling.py new file mode 100644 index 00000000..8c64219c --- /dev/null +++ b/backend/package/yuxi/channel/gateway/polling.py @@ -0,0 +1,102 @@ +"""HTTP Polling 降级机制 — SSE/WS 不可用时的备选路径。 + +按 session_id 维护响应队列，前端通过 HTTP GET 定时轮询获取回复。 +与 SSE 端点协同使用，形成三级降级链：SSE → Polling → 单次 fetch。 + +Usage: + polling = PollingFallback() + await polling.push("session_abc", {"type": "delta", "data": {"content": "hello"}}) + events = await polling.poll("session_abc") +""" + +from __future__ import annotations + +import asyncio +import logging +import time as _time + +logger = logging.getLogger(__name__) + +DEFAULT_POLL_TTL = 600 + + +class PollingFallback: + def __init__(self, max_queue_size: int = 100, ttl_seconds: int = DEFAULT_POLL_TTL): + self._queues: dict[str, asyncio.Queue[dict]] = {} + self._ttl = ttl_seconds + self._last_active: dict[str, float] = {} + self._max_queue_size = max_queue_size + self._lock = asyncio.Lock() + + async def ensure(self, session_id: str) -> None: + async with self._lock: + if session_id not in self._queues: + self._queues[session_id] = asyncio.Queue(maxsize=self._max_queue_size) + self._last_active[session_id] = _time.monotonic() + + async def push(self, session_id: str, event: dict) -> None: + async with self._lock: + self._last_active[session_id] = _time.monotonic() + q = self._queues.get(session_id) + if q is None: + q = asyncio.Queue(maxsize=self._max_queue_size) + self._queues[session_id] = q + try: + q.put_nowait(event) + except asyncio.QueueFull: + logger.warning("Polling queue full for session %s, dropping event", session_id) + + async def poll(self, session_id: str) -> list[dict]: + async with self._lock: + self._last_active[session_id] = _time.monotonic() + q = self._queues.get(session_id) + if q is None: + return [] + + events: list[dict] = [] + while not q.empty(): + try: + events.append(q.get_nowait()) + except asyncio.QueueEmpty: + break + return events + + async def cleanup_stale(self) -> int: + async with self._lock: + now = _time.monotonic() + stale = [sid for sid, ts in self._last_active.items() if now - ts > self._ttl] + for sid in stale: + self._queues.pop(sid, None) + self._last_active.pop(sid, None) + if stale: + logger.info("PollingFallback cleaned up %d stale sessions", len(stale)) + return len(stale) + + def active_sessions(self) -> int: + return len(self._queues) + + +polling_fallback = PollingFallback() + + +async def _polling_cleanup_loop(interval: int = 300) -> None: + while True: + await asyncio.sleep(interval) + await polling_fallback.cleanup_stale() + + +_cleanup_task: asyncio.Task | None = None + + +def start_polling_cleanup(interval: int = 300) -> None: + global _cleanup_task + if _cleanup_task is None or _cleanup_task.done(): + _cleanup_task = asyncio.ensure_future(_polling_cleanup_loop(interval)) + logger.info("PollingFallback cleanup loop started (interval=%ds)", interval) + + +def stop_polling_cleanup() -> None: + global _cleanup_task + if _cleanup_task and not _cleanup_task.done(): + _cleanup_task.cancel() + _cleanup_task = None diff --git a/backend/package/yuxi/channel/gateway/probe.py b/backend/package/yuxi/channel/gateway/probe.py new file mode 100644 index 00000000..61894bb3 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/probe.py @@ -0,0 +1,289 @@ +import asyncio +import logging +import time +from dataclasses import dataclass, field +from enum import StrEnum + +import aiohttp + +from yuxi.channel.gateway.protocol import ( + RpcRequest, + RpcResponse, + marshal_frame, + unmarshal_frame, +) + +logger = logging.getLogger(__name__) + +DEFAULT_PROBE_TIMEOUT_MS = 10_000 +MIN_PROBE_TIMEOUT_MS = 250 +MAX_PROBE_TIMEOUT_MS = 30_000 + + +class ProbeCapability(StrEnum): + UNKNOWN = "unknown" + PAIRING_PENDING = "pairing_pending" + CONNECTED_NO_OPERATOR_SCOPE = "connected_no_operator_scope" + READ_ONLY = "read_only" + WRITE_CAPABLE = "write_capable" + ADMIN_CAPABLE = "admin_capable" + + +@dataclass +class ProbeAuthSummary: + role: str | None = None + scopes: list[str] = field(default_factory=list) + capability: ProbeCapability = ProbeCapability.UNKNOWN + + +@dataclass +class ProbeServerSummary: + version: str | None = None + conn_id: str | None = None + + +@dataclass +class ProbeClose: + code: int + reason: str + hint: str | None = None + + +@dataclass +class GatewayProbeResult: + ok: bool + url: str + connect_latency_ms: float | None = None + error: str | None = None + close: ProbeClose | None = None + auth: ProbeAuthSummary = field(default_factory=ProbeAuthSummary) + server: ProbeServerSummary = field(default_factory=ProbeServerSummary) + health: dict | None = None + status: dict | None = None + presence: list | None = None + config_snapshot: dict | None = None + + +def _clamp_timeout_ms(timeout_ms: float | int) -> float: + return max(MIN_PROBE_TIMEOUT_MS, min(float(timeout_ms), MAX_PROBE_TIMEOUT_MS)) + + +def _resolve_capability( + scopes: list[str], + connect_latency_ms: float | None, + auth_metadata_present: bool, +) -> ProbeCapability: + if "operator.admin" in scopes: + return ProbeCapability.ADMIN_CAPABLE + if "operator.write" in scopes: + return ProbeCapability.WRITE_CAPABLE + if "operator.read" in scopes or "admin" in scopes: + return ProbeCapability.READ_ONLY + if connect_latency_ms is not None and auth_metadata_present: + return ProbeCapability.CONNECTED_NO_OPERATOR_SCOPE + return ProbeCapability.UNKNOWN + + +async def probe_gateway( + url: str, + token: str | None = None, + password: str | None = None, + timeout_ms: float = DEFAULT_PROBE_TIMEOUT_MS, + detail_level: str = "full", +) -> GatewayProbeResult: + connect_latency_ms: float | None = None + connect_error: str | None = None + close_info: ProbeClose | None = None + auth_summary = ProbeAuthSummary() + server_summary = ProbeServerSummary() + auth_metadata_present = False + + effective_timeout = _clamp_timeout_ms(timeout_ms) + timeout_sec = effective_timeout / 1000.0 + + headers: dict[str, str] = {} + if token: + headers["Authorization"] = f"Bearer {token}" + elif password: + headers["Authorization"] = f"Bearer {password}" + + conn_start = time.monotonic() + + try: + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=timeout_sec), + ) as session: + async with session.ws_connect( + url, + headers=headers, + heartbeat=15.0, + ) as ws: + connect_latency_ms = (time.monotonic() - conn_start) * 1000 + + auth_metadata_present = True + + if detail_level == "none": + ws_close = ws.close_code + if ws_close is not None: + close_info = ProbeClose( + code=ws_close, + reason="connection closed by server", + ) + return GatewayProbeResult( + ok=ws_close is None, + url=url, + connect_latency_ms=connect_latency_ms, + error=None, + close=close_info, + auth=auth_summary, + server=server_summary, + ) + + rpc_request_id = "probe" + rpc_request = RpcRequest( + id=rpc_request_id, + method="system.health", + params={}, + ) + await ws.send_str(marshal_frame(rpc_request)) + + try: + raw = await asyncio.wait_for( + ws.receive_str(), + timeout=effective_timeout / 1000.0, + ) + except asyncio.TimeoutError: + return GatewayProbeResult( + ok=False, + url=url, + connect_latency_ms=connect_latency_ms, + error="timeout waiting for health response", + auth=auth_summary, + server=server_summary, + ) + + try: + frame = unmarshal_frame(raw) + except Exception: + return GatewayProbeResult( + ok=False, + url=url, + connect_latency_ms=connect_latency_ms, + error="invalid response frame from gateway", + auth=auth_summary, + server=server_summary, + ) + + if not isinstance(frame, RpcResponse) or not frame.ok: + error_msg = frame.error_message if isinstance(frame, RpcResponse) else "unexpected frame type" + return GatewayProbeResult( + ok=False, + url=url, + connect_latency_ms=connect_latency_ms, + error=error_msg or "gateway returned error", + auth=auth_summary, + server=server_summary, + ) + + health_data = frame.result + + auth_summary = ProbeAuthSummary( + role="viewer", + scopes=[], + capability=ProbeCapability.READ_ONLY, + ) + + if detail_level == "presence": + return GatewayProbeResult( + ok=True, + url=url, + connect_latency_ms=connect_latency_ms, + health=health_data, + auth=auth_summary, + server=server_summary, + presence=health_data.get("channels") if health_data else None, + ) + + status_data: dict | None = None + try: + status_request = RpcRequest( + id="probe-status", + method="channels.status", + params={}, + ) + await ws.send_str(marshal_frame(status_request)) + raw_status = await asyncio.wait_for( + ws.receive_str(), + timeout=effective_timeout / 1000.0, + ) + status_frame = unmarshal_frame(raw_status) + if isinstance(status_frame, RpcResponse) and status_frame.ok: + status_data = status_frame.result + except Exception: + logger.debug("Failed to get channel status during probe", exc_info=True) + + return GatewayProbeResult( + ok=True, + url=url, + connect_latency_ms=connect_latency_ms, + health=health_data, + status=status_data, + auth=auth_summary, + server=server_summary, + ) + + except aiohttp.ClientConnectorError as e: + connect_error = f"连接失败: {e}" + except aiohttp.WSServerHandshakeError as e: + connect_error = f"WebSocket 握手失败: {e.status} {e.message}" + except aiohttp.ClientError as e: + connect_error = f"客户端错误: {e}" + except asyncio.TimeoutError: + connect_error = "连接超时" + except Exception as e: + connect_error = f"未知错误: {e}" + + return GatewayProbeResult( + ok=False, + url=url, + connect_latency_ms=connect_latency_ms, + error=connect_error, + close=close_info, + auth=auth_summary, + server=server_summary, + ) + + +def format_probe_result(result: GatewayProbeResult, verbose: bool = False) -> str: + lines: list[str] = [] + + if result.ok: + lines.append(f"✓ 网关探测成功: {result.url}") + if result.connect_latency_ms is not None: + lines.append(f" 连接延迟: {result.connect_latency_ms:.1f}ms") + lines.append(f" 能力级别: {result.auth.capability.value}") + else: + lines.append(f"✗ 网关探测失败: {result.url}") + if result.connect_latency_ms is not None: + lines.append(f" 连接延迟: {result.connect_latency_ms:.1f}ms") + if result.error: + lines.append(f" 错误: {result.error}") + if result.close: + lines.append(f" 关闭信息: code={result.close.code} reason={result.close.reason}") + + if verbose and result.health: + lines.append("") + lines.append("健康报告:") + health = result.health + if isinstance(health, dict): + lines.append(f" 状态: {health.get('status', 'unknown')}") + summary = health.get("summary", {}) + if summary: + lines.append( + f" 频道: 总计={summary.get('total', 0)}, " + f"运行={summary.get('running', 0)}, " + f"停止={summary.get('stopped', 0)}, " + f"异常={summary.get('unhealthy', 0)}" + ) + + return "\n".join(lines) diff --git a/backend/package/yuxi/channel/gateway/protocol.py b/backend/package/yuxi/channel/gateway/protocol.py new file mode 100644 index 00000000..bfe85d05 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/protocol.py @@ -0,0 +1,207 @@ +import json as _json +import time as _time +import uuid as _uuid +from dataclasses import dataclass, field +from enum import StrEnum + + +class FrameType(StrEnum): + REQUEST = "request" + RESPONSE = "response" + EVENT = "event" + + +class GatewayErrorCode(StrEnum): + INVALID_REQUEST = "INVALID_REQUEST" + METHOD_NOT_FOUND = "METHOD_NOT_FOUND" + INVALID_PARAMS = "INVALID_PARAMS" + INTERNAL_ERROR = "INTERNAL_ERROR" + AUTH_ERROR = "AUTH_ERROR" + PERMISSION_DENIED = "PERMISSION_DENIED" + RATE_LIMITED = "RATE_LIMITED" + TIMEOUT = "TIMEOUT" + NOT_CONNECTED = "NOT_CONNECTED" + UNAVAILABLE = "UNAVAILABLE" + DEVICE_IDENTITY_REQUIRED = "DEVICE_IDENTITY_REQUIRED" + + +class GatewayRpcMethod(StrEnum): + CONNECT = "connect" + CHAT_SEND = "chat.send" + CHAT_STREAM = "chat.stream" + CHANNELS_LIST = "channels.list" + CHANNELS_STATUS = "channels.status" + CHANNELS_CONFIGURE = "channels.configure" + SESSIONS_LIST = "sessions.list" + SESSIONS_HISTORY = "sessions.history" + PLUGINS_LIST = "plugins.list" + PLUGINS_INSTALL = "plugins.install" + SYSTEM_HEALTH = "system.health" + SYSTEM_CONFIG = "system.config" + SYSTEM_VERSION = "system.version" + CHANNELS_START = "channels.start" + CHANNELS_STOP = "channels.stop" + CHANNELS_RESTART = "channels.restart" + PAIRING_LIST = "pairing.list" + PAIRING_APPROVE = "pairing.approve" + PAIRING_REJECT = "pairing.reject" + ALLOWLIST_GET = "allowlist.get" + ALLOWLIST_ADD = "allowlist.add" + ALLOWLIST_REMOVE = "allowlist.remove" + CHAT_HISTORY = "chat.history" + CHAT_CANCEL = "chat.cancel" + CRON_LIST = "cron.list" + CRON_CREATE = "cron.create" + CRON_UPDATE = "cron.update" + CRON_DELETE = "cron.delete" + CRON_FORCE_RUN = "cron.force_run" + CRON_PAUSE = "cron.pause" + CRON_RESUME = "cron.resume" + CRON_DIAGNOSTICS = "cron.diagnostics" + CRON_RUN_LOG = "cron.run_log" + AGENT_TOOLS_LIST = "agentTools.list" + AGENT_TOOLS_EXECUTE = "agentTools.execute" + SYSTEM_LOG_TAIL = "system.log.tail" + SESSIONS_CREATE = "sessions.create" + SESSIONS_DELETE = "sessions.delete" + CONFIG_GET = "config.get" + CONFIG_SET = "config.set" + MESSAGE_SEND = "message.send" + MESSAGE_ACTION = "message.action" + CHANNELS_PROBE = "channels.probe" + IDENTITY_LINKS_LIST = "identity_links.list" + IDENTITY_LINKS_ADD = "identity_links.add" + IDENTITY_LINKS_REMOVE = "identity_links.remove" + + +class DeliveryMode(StrEnum): + DIRECT = "direct" + GATEWAY = "gateway" + HYBRID = "hybrid" + + +@dataclass +class RpcFrame: + frame_type: FrameType + + +@dataclass +class RpcRequest(RpcFrame): + frame_type: FrameType = FrameType.REQUEST + id: str = field(default_factory=lambda: _uuid.uuid4().hex[:12]) + method: str = "" + params: dict | None = None + session_id: str | None = None + caller_user_id: str | None = None + + +@dataclass +class RpcResponse(RpcFrame): + frame_type: FrameType = FrameType.RESPONSE + id: str = "" + ok: bool = True + result: dict | None = None + error_code: GatewayErrorCode | None = None + error_message: str | None = None + error_data: dict | None = None + retryable: bool = False + retry_after_ms: int | None = None + + +@dataclass +class HelloOk(RpcResponse): + frame_type: FrameType = FrameType.RESPONSE + id: str = "" + ok: bool = True + result: dict | None = field(default_factory=lambda: { + "protocolVersion": 1, + "server": {"version": "1.0.0", "connId": ""}, + "features": {"methods": [], "events": []}, + "auth": {"role": "", "scopes": []}, + "policy": { + "maxPayload": 25 * 1024 * 1024, + "maxBufferedBytes": 256 * 1024, + "tickIntervalMs": 30_000, + }, + }) + + +@dataclass +class RpcEvent(RpcFrame): + frame_type: FrameType = FrameType.EVENT + event: str = "" + data: dict | None = None + timestamp: float = field(default_factory=_time.time) + seq: int | None = None + state_version: int | None = None + + +def marshal_frame(frame: RpcFrame) -> str: + if isinstance(frame, RpcRequest): + payload = {"type": frame.frame_type.value, "id": frame.id, "method": frame.method} + if frame.params is not None: + payload["params"] = frame.params + if frame.session_id is not None: + payload["sessionId"] = frame.session_id + elif isinstance(frame, RpcResponse): + payload = {"type": frame.frame_type.value, "id": frame.id, "ok": frame.ok} + if frame.result is not None: + payload["result"] = frame.result + if frame.error_code is not None: + error_payload = { + "code": frame.error_code.value, + "message": frame.error_message or "", + } + if frame.error_data is not None: + error_payload["data"] = frame.error_data + if frame.retryable: + error_payload["retryable"] = True + if frame.retry_after_ms is not None: + error_payload["retryAfterMs"] = frame.retry_after_ms + payload["error"] = error_payload + elif isinstance(frame, RpcEvent): + payload = {"type": frame.frame_type.value, "event": frame.event, "timestamp": frame.timestamp} + if frame.data is not None: + payload["data"] = frame.data + if frame.seq is not None: + payload["seq"] = frame.seq + if frame.state_version is not None: + payload["stateVersion"] = frame.state_version + else: + raise TypeError(f"Unknown frame type: {type(frame)}") + return _json.dumps(payload, ensure_ascii=False, default=str) + + +def unmarshal_frame(data: str | bytes) -> RpcFrame: + if isinstance(data, bytes): + data = data.decode("utf-8") + raw = _json.loads(data) + frame_type = raw.get("type", FrameType.REQUEST.value) + if frame_type == FrameType.REQUEST.value: + return RpcRequest( + id=raw.get("id", ""), + method=raw.get("method", ""), + params=raw.get("params"), + session_id=raw.get("sessionId"), + ) + elif frame_type == FrameType.RESPONSE.value: + error = raw.get("error") + return RpcResponse( + id=raw.get("id", ""), + ok=raw.get("ok", True), + result=raw.get("result"), + error_code=GatewayErrorCode(error["code"]) if error else None, + error_message=error.get("message") if error else None, + error_data=error.get("data") if error else None, + retryable=error.get("retryable", False) if error else False, + retry_after_ms=error.get("retryAfterMs") if error else None, + ) + elif frame_type == FrameType.EVENT.value: + return RpcEvent( + event=raw.get("event", ""), + data=raw.get("data"), + timestamp=raw.get("timestamp", _time.time()), + seq=raw.get("seq"), + state_version=raw.get("stateVersion"), + ) + raise ValueError(f"Unknown frame type: {frame_type}") diff --git a/backend/package/yuxi/channel/gateway/rbac.py b/backend/package/yuxi/channel/gateway/rbac.py new file mode 100644 index 00000000..2523a15a --- /dev/null +++ b/backend/package/yuxi/channel/gateway/rbac.py @@ -0,0 +1,149 @@ +import logging +from collections.abc import Callable +from enum import StrEnum +from functools import wraps +from typing import Any + +logger = logging.getLogger(__name__) + + +class GatewayRole(StrEnum): + SUPERADMIN = "superadmin" + ADMIN = "admin" + OPERATOR = "operator" + VIEWER = "viewer" + + +ROLE_HIERARCHY: dict[GatewayRole, int] = { + GatewayRole.SUPERADMIN: 4, + GatewayRole.ADMIN: 3, + GatewayRole.OPERATOR: 2, + GatewayRole.VIEWER: 1, +} + + +_USER_ROLE_MAP: dict[str, GatewayRole] = { + "superadmin": GatewayRole.SUPERADMIN, + "admin": GatewayRole.ADMIN, + "user": GatewayRole.OPERATOR, +} + + +def map_user_role(user_role: str | None) -> GatewayRole: + if not user_role: + return GatewayRole.VIEWER + return _USER_ROLE_MAP.get(user_role.lower(), GatewayRole.VIEWER) + + +READ_ONLY_METHODS: set[str] = { + "system.health", + "system.version", + "channels.list", + "channels.status", + "channels.probe", + "chat.history", + "sessions.list", + "sessions.history", + "plugins.list", + "agentTools.list", + "config.get", +} + +OPERATOR_METHODS: set[str] = READ_ONLY_METHODS | { + "chat.send", + "chat.stream", + "chat.cancel", + "sessions.create", + "sessions.delete", + "allowlist.get", + "cron.list", + "agentTools.execute", + "message.send", + "message.action", +} + +ADMIN_METHODS: set[str] = OPERATOR_METHODS | { + "system.log.tail", + "system.config", + "channels.start", + "channels.stop", + "channels.restart", + "channels.configure", + "pairing.list", + "pairing.approve", + "pairing.reject", + "allowlist.add", + "allowlist.remove", + "identity_links.list", + "identity_links.add", + "identity_links.remove", + "plugins.install", + "config.set", +} + +SUPERADMIN_METHODS: set[str] = ADMIN_METHODS | { + "cron.create", + "cron.delete", + "cron.force_run", + "cron.pause", + "cron.resume", +} + +METHOD_MIN_ROLE: dict[str, GatewayRole] = {} + +for _method in READ_ONLY_METHODS: + METHOD_MIN_ROLE[_method] = GatewayRole.VIEWER + +for _method in OPERATOR_METHODS - READ_ONLY_METHODS: + METHOD_MIN_ROLE[_method] = GatewayRole.OPERATOR + +for _method in ADMIN_METHODS - OPERATOR_METHODS: + METHOD_MIN_ROLE[_method] = GatewayRole.ADMIN + +for _method in {"cron.create", "cron.delete", "cron.force_run", "cron.pause", "cron.resume"}: + METHOD_MIN_ROLE[_method] = GatewayRole.SUPERADMIN + + +def check_permission(role: GatewayRole | None, method: str) -> bool: + effective_role = role or GatewayRole.VIEWER + min_role = METHOD_MIN_ROLE.get(method) + if min_role is None: + return True + return ROLE_HIERARCHY[effective_role] >= ROLE_HIERARCHY[min_role] + + +def require_role(min_role: GatewayRole): + def decorator( + func: Callable[..., Any], + ) -> Callable[..., Any]: + @wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + caller_role: GatewayRole | None = kwargs.pop("_caller_role", None) + if caller_role is None: + logger.warning( + "RBAC: _caller_role not found in handler kwargs for %s, allowing by default", + func.__name__, + ) + return await func(*args, **kwargs) + + if ROLE_HIERARCHY[caller_role] < ROLE_HIERARCHY[min_role]: + from yuxi.channel.gateway.protocol import GatewayErrorCode, RpcResponse + + logger.warning( + "RBAC: %s requires %s, caller has %s — denied", + func.__name__, + min_role.value, + caller_role.value, + ) + return RpcResponse( + id=kwargs.get("request_id", ""), + ok=False, + error_code=GatewayErrorCode.PERMISSION_DENIED, + error_message=f"需要 {min_role.value} 权限，当前角色为 {caller_role.value}", + ) + + return await func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/backend/package/yuxi/channel/gateway/routes.py b/backend/package/yuxi/channel/gateway/routes.py new file mode 100644 index 00000000..3c9eca8a --- /dev/null +++ b/backend/package/yuxi/channel/gateway/routes.py @@ -0,0 +1,161 @@ +import asyncio +import logging +from collections.abc import Callable +from typing import Any + +from yuxi.channel.gateway.protocol import DeliveryMode +from yuxi.channel.gateway.webhook_security import ( + WebhookGuard, + WebhookGuardConfig, + WebhookGuardResult, +) +from yuxi.channel.gateway.webhook_security import ( + webhook_guard as _default_webhook_guard, +) + +logger = logging.getLogger(__name__) + +WebhookHandler = Callable[[dict[str, Any]], asyncio.Future[dict[str, Any]]] + + +class WebhookRegistry: + """渠道 Webhook 处理器注册表。 + + 每个渠道插件注册自己的 webhook 处理函数， + HTTP 层统一通过 FastAPI router 将请求分发到此处。 + + 集成了多层 Webhook 安全 Guard： + - Method 检查（仅允许注册的 HTTP 方法） + - Content-Type 验证 + - Body 大小限制 + - HMAC-SHA256 签名校验 + - 并发限制 + - 异常追踪告警 + """ + + def __init__(self, webhook_guard: WebhookGuard | None = None): + self._handlers: dict[str, WebhookHandler] = {} + self._delivery_modes: dict[str, DeliveryMode] = {} + self._guard = webhook_guard or _default_webhook_guard + + def register( + self, + channel_type: str, + handler: WebhookHandler, + delivery_mode: DeliveryMode = DeliveryMode.DIRECT, + guard_config: WebhookGuardConfig | None = None, + ) -> None: + if channel_type in self._handlers: + logger.warning("Overwriting webhook handler for channel: %s", channel_type) + self._handlers[channel_type] = handler + self._delivery_modes[channel_type] = delivery_mode + self._guard.register_channel(channel_type, guard_config) + + def get_handler(self, channel_type: str) -> WebhookHandler | None: + return self._handlers.get(channel_type) + + def get_delivery_mode(self, channel_type: str) -> DeliveryMode: + return self._delivery_modes.get(channel_type, DeliveryMode.DIRECT) + + def list_channels(self) -> list[dict[str, str]]: + return [ + { + "channel_type": ct, + "delivery_mode": self._delivery_modes.get(ct, DeliveryMode.DIRECT).value, + } + for ct in self._handlers + ] + + def remove(self, channel_type: str) -> None: + self._handlers.pop(channel_type, None) + self._delivery_modes.pop(channel_type, None) + self._guard.unregister_channel(channel_type) + + def clear(self) -> None: + for ct in list(self._handlers): + self._guard.unregister_channel(ct) + self._handlers.clear() + self._delivery_modes.clear() + + async def dispatch_guarded( + self, + channel_type: str, + method: str, + content_type: str | None, + body: bytes, + signature: str | None = None, + extra_headers: dict[str, str] | None = None, + ) -> tuple[WebhookGuardResult | None, dict[str, Any] | None]: + guard_result = self._guard.run_pipeline( + channel_type, + method, + content_type, + body, + signature, + ) + if not guard_result.allowed: + return guard_result, None + + handler = self._handlers.get(channel_type) + if handler is None: + self._guard.release_concurrency(channel_type) + return WebhookGuardResult( + allowed=False, + error_code=None, + error_message=f"No handler registered for channel: {channel_type}", + http_status=404, + ), None + + try: + import json + + payload = json.loads(body) + except Exception: + self._guard.release_concurrency(channel_type) + return WebhookGuardResult( + allowed=False, + error_code=None, + error_message="Invalid JSON payload", + http_status=400, + ), None + + if extra_headers: + payload["_headers"] = extra_headers + payload["_raw_body"] = body + + try: + result = await handler(payload) + except Exception: + logger.exception("Webhook handler failed for channel: %s", channel_type) + result = None + finally: + self._guard.release_concurrency(channel_type) + + return None, result + + @property + def guard(self) -> WebhookGuard: + return self._guard + + +webhook_registry = WebhookRegistry() + + +import re + +_WEBHOOK_PATH_RE = re.compile(r"^[a-zA-Z0-9_-]+$") + + +def build_webhook_path(channel_type: str) -> str: + if not channel_type or not _WEBHOOK_PATH_RE.match(channel_type): + raise ValueError(f"Invalid channel_type: {channel_type!r}") + return f"/webhook/{channel_type}" + + +def build_health_response(status: str, channels: dict, summary: dict, ts: str | None = None) -> dict: + return { + "status": status, + "timestamp": ts, + "channels": channels, + "summary": summary, + } diff --git a/backend/package/yuxi/channel/gateway/rpc_dispatcher.py b/backend/package/yuxi/channel/gateway/rpc_dispatcher.py new file mode 100644 index 00000000..3c789e61 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/rpc_dispatcher.py @@ -0,0 +1,193 @@ +import logging +from collections.abc import AsyncGenerator, Callable, Coroutine +from typing import Any + +from yuxi.channel.gateway.protocol import ( + GatewayErrorCode, + GatewayRpcMethod, + RpcEvent, + RpcRequest, + RpcResponse, +) +from yuxi.channel.gateway.rbac import GatewayRole, check_permission + +logger = logging.getLogger(__name__) + +RpcHandler = Callable[[RpcRequest], Coroutine[Any, Any, RpcResponse]] +RpcStreamHandler = Callable[[RpcRequest], AsyncGenerator[RpcEvent | RpcResponse, None]] + + +class RpcDispatcher: + def __init__(self): + self._handlers: dict[str, RpcHandler] = {} + self._stream_handlers: dict[str, RpcStreamHandler] = {} + self._lazy_handlers: dict[str, Callable[[], RpcHandler]] = {} + self._lazy_stream_handlers: dict[str, Callable[[], RpcStreamHandler]] = {} + self._aliases: dict[str, str] = {} + self._enforce_rbac: bool = True + + @property + def enforce_rbac(self) -> bool: + return self._enforce_rbac + + @enforce_rbac.setter + def enforce_rbac(self, value: bool) -> None: + self._enforce_rbac = value + + def register(self, method: GatewayRpcMethod | str, handler: RpcHandler) -> None: + key = method.value if isinstance(method, GatewayRpcMethod) else method + if key in self._handlers or key in self._stream_handlers: + logger.warning("RpcDispatcher: 覆盖已注册方法 %s", key) + self._handlers[key] = handler + + def register_stream(self, method: GatewayRpcMethod | str, handler: RpcStreamHandler) -> None: + key = method.value if isinstance(method, GatewayRpcMethod) else method + if key in self._handlers or key in self._stream_handlers: + logger.warning("RpcDispatcher: 覆盖已注册方法 %s", key) + self._stream_handlers[key] = handler + + def register_lazy(self, method: GatewayRpcMethod | str, factory: Callable[[], RpcHandler]) -> None: + key = method.value if isinstance(method, GatewayRpcMethod) else method + if key in self._handlers or key in self._lazy_handlers: + logger.warning("RpcDispatcher: 覆盖已有懒加载方法 %s", key) + self._lazy_handlers[key] = factory + + def register_stream_lazy( + self, method: GatewayRpcMethod | str, factory: Callable[[], RpcStreamHandler] + ) -> None: + key = method.value if isinstance(method, GatewayRpcMethod) else method + if key in self._stream_handlers or key in self._lazy_stream_handlers: + logger.warning("RpcDispatcher: 覆盖已有懒加载流方法 %s", key) + self._lazy_stream_handlers[key] = factory + + def add_alias(self, alias: str, method: str) -> None: + self._aliases[alias] = method + + def _resolve_method(self, method: str) -> str: + return self._aliases.get(method, method) + + def _resolve_handler(self, method: str) -> RpcHandler | None: + handler = self._handlers.get(method) + if handler is not None: + return handler + + factory = self._lazy_handlers.get(method) + if factory is not None: + handler = factory() + self._handlers[method] = handler + del self._lazy_handlers[method] + return handler + + return None + + def _resolve_stream_handler(self, method: str) -> RpcStreamHandler | None: + handler = self._stream_handlers.get(method) + if handler is not None: + return handler + + factory = self._lazy_stream_handlers.get(method) + if factory is not None: + handler = factory() + self._stream_handlers[method] = handler + del self._lazy_stream_handlers[method] + return handler + + return None + + def is_stream(self, method: str) -> bool: + resolved = self._resolve_method(method) + return resolved in self._stream_handlers or resolved in self._lazy_stream_handlers + + async def dispatch(self, request: RpcRequest, caller_roles: list[GatewayRole] | None = None) -> RpcResponse: + resolved_method = self._resolve_method(request.method) + + if self._enforce_rbac and caller_roles: + primary_role = caller_roles[0] if caller_roles else None + if not check_permission(primary_role, resolved_method): + logger.warning( + "RpcDispatcher: 权限拒绝 method=%s role=%s", + resolved_method, + primary_role, + ) + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.PERMISSION_DENIED, + error_message=f"方法 '{resolved_method}' 需要更高权限，当前角色为 {primary_role}", + ) + + handler = self._resolve_handler(resolved_method) + if handler is None: + logger.warning("RpcDispatcher: 未注册方法 %s", resolved_method) + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.METHOD_NOT_FOUND, + error_message=f"方法 '{resolved_method}' 未注册", + ) + try: + return await handler(request) + except Exception: + logger.exception("RpcDispatcher: 方法 %s 执行异常", resolved_method) + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message=f"方法 '{resolved_method}' 执行异常", + ) + + async def dispatch_stream( + self, request: RpcRequest, caller_roles: list[GatewayRole] | None = None + ) -> AsyncGenerator[RpcEvent | RpcResponse, None]: + resolved_method = self._resolve_method(request.method) + + if self._enforce_rbac and caller_roles: + primary_role = caller_roles[0] if caller_roles else None + if not check_permission(primary_role, resolved_method): + logger.warning( + "RpcDispatcher: 流式权限拒绝 method=%s role=%s", + resolved_method, + primary_role, + ) + yield RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.PERMISSION_DENIED, + error_message=f"方法 '{resolved_method}' 需要更高权限，当前角色为 {primary_role}", + ) + return + + handler = self._resolve_stream_handler(resolved_method) + if handler is None: + logger.warning("RpcDispatcher: 未注册流式方法 %s", resolved_method) + yield RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.METHOD_NOT_FOUND, + error_message=f"方法 '{resolved_method}' 未注册", + ) + return + try: + async for item in handler(request): + yield item + except Exception: + logger.exception("RpcDispatcher: 流式方法 %s 执行异常", resolved_method) + yield RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message=f"方法 '{resolved_method}' 执行异常", + ) + + def list_methods(self) -> list[str]: + methods = set(self._handlers.keys()) + methods.update(self._lazy_handlers.keys()) + methods.update(self._stream_handlers.keys()) + methods.update(self._lazy_stream_handlers.keys()) + return list(methods) + + def list_aliases(self) -> dict[str, str]: + return dict(self._aliases) + + +rpc_dispatcher = RpcDispatcher() diff --git a/backend/package/yuxi/channel/gateway/rpc_handlers.py b/backend/package/yuxi/channel/gateway/rpc_handlers.py new file mode 100644 index 00000000..1069e866 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/rpc_handlers.py @@ -0,0 +1,1444 @@ +import json +import logging +import uuid +from collections.abc import AsyncGenerator + +from sqlalchemy import select + +from yuxi.channel.gateway.protocol import ( + GatewayErrorCode, + GatewayRpcMethod, + RpcEvent, + RpcRequest, + RpcResponse, +) +from yuxi.channel.gateway.rpc_dispatcher import rpc_dispatcher +from yuxi.channel.gateway.validation import ( + ProbeParams, + SendMessageParams, + StartAccountParams, + StopAccountParams, + validate_params, +) +from yuxi.channel.runtime.manager import gateway +from yuxi.storage.postgres.manager import pg_manager + +logger = logging.getLogger(__name__) + + +def _is_placeholder_user(user_id: str | None) -> bool: + if not user_id: + return True + return user_id in ("shared", "bootstrap") or user_id.startswith(("tailscale:", "device:", "proxy:")) + + +async def _resolve_caller_user(request: RpcRequest, db): + from yuxi.storage.postgres.models_business import User + + user_id = request.caller_user_id + if user_id and not _is_placeholder_user(user_id): + try: + int_id = int(user_id) + result = await db.execute(select(User).where(User.id == int_id)) + user = result.scalar_one_or_none() + if user: + return user + except (ValueError, TypeError): + pass + try: + from sqlalchemy import func as sa_func + + result = await db.execute(select(User).where(sa_func.cast(User.id, str) == user_id)) + user = result.scalar_one_or_none() + if user: + return user + except Exception: + pass + + result = await db.execute(select(User).limit(1)) + user = result.scalar_one_or_none() + if user: + logger.warning( + "RPC handler using fallback user %s for caller_user_id=%s method=%s", + user.id, + user_id, + request.method, + ) + return user + + +async def _handle_health(request: RpcRequest) -> RpcResponse: + report = gateway.get_health() + + event_loop = None + try: + from yuxi.channel.monitoring.event_loop_monitor import event_loop_monitor + + loop_health = event_loop_monitor.health() + loop_stats = event_loop_monitor.stats() + event_loop = { + "degraded": loop_health.degraded, + "reasons": loop_health.reasons, + "interval_ms": loop_health.interval_ms, + "delay": { + "p50_ms": loop_stats.p50_ms, + "p90_ms": loop_stats.p90_ms, + "p99_ms": loop_stats.p99_ms, + "p999_ms": loop_stats.p999_ms, + "max_ms": loop_stats.max_ms, + "min_ms": loop_stats.min_ms, + "avg_ms": loop_stats.avg_ms, + }, + "utilization": loop_stats.utilization, + "cpu_core_ratio": loop_health.cpu_core_ratio, + "samples": loop_stats.samples, + "warnings": loop_stats.warnings, + } + except Exception: + logger.exception("Failed to collect event loop health") + + return RpcResponse( + id=request.id, + ok=True, + result={ + "status": report.status, + "summary": { + "total": report.summary.total, + "running": report.summary.running, + "stopped": report.summary.stopped, + "degraded": report.summary.degraded, + "unhealthy": report.summary.unhealthy, + "unknown": report.summary.unknown, + "connected": report.summary.connected, + "total_active_runs": report.summary.total_active_runs, + }, + "channels": report.channels, + "event_loop": event_loop, + }, + ) + + +async def _handle_channels_list(request: RpcRequest) -> RpcResponse: + snapshots = gateway.get_all_snapshots() + return RpcResponse( + id=request.id, + ok=True, + result={"channels": [s.to_dict() for s in snapshots.values()]}, + ) + + +async def _handle_channels_status(request: RpcRequest) -> RpcResponse: + channel_type = request.params.get("channel_type") if request.params else None + account_id = request.params.get("account_id", "default") if request.params else "default" + + if channel_type: + snapshot = gateway.get_snapshot(channel_type, account_id) + if snapshot is None: + return RpcResponse( + id=request.id, + ok=False, + error_code=None, + error_message=f"渠道 {channel_type}:{account_id} 未找到", + ) + return RpcResponse(id=request.id, ok=True, result=snapshot.to_dict()) + + snapshots = gateway.get_all_snapshots() + return RpcResponse( + id=request.id, + ok=True, + result={"channels": [s.to_dict() for s in snapshots.values()]}, + ) + + +@validate_params(ProbeParams) +async def _handle_channels_probe(request: RpcRequest) -> RpcResponse: + validated = request._validated_params + + snapshot = gateway.get_snapshot(validated.channel_type, validated.account_id) + if snapshot is None: + return RpcResponse( + id=request.id, + ok=False, + error_code=None, + error_message=f"渠道 {validated.channel_type}:{validated.account_id} 未找到", + ) + + return RpcResponse( + id=request.id, + ok=True, + result={ + "channel_type": snapshot.channel_type, + "account_id": snapshot.account_id, + "state": snapshot.state.value, + "connected": snapshot.connected, + "configured": snapshot.configured, + "enabled": snapshot.enabled, + "healthy": snapshot.health_state.value if snapshot.health_state else "unknown", + "last_connected_at": snapshot.last_connected_at, + "last_disconnect": snapshot.last_disconnect.to_dict() if snapshot.last_disconnect else None, + "last_error": snapshot.last_error, + }, + ) + + +@validate_params(StartAccountParams) +async def _handle_channels_start(request: RpcRequest) -> RpcResponse: + validated = request._validated_params + channel_type = validated.channel_type + account_id = validated.account_id + + if channel_type: + snapshot = await gateway.start_channel(channel_type, account_id) + return RpcResponse(id=request.id, ok=True, result=snapshot.to_dict()) + + await gateway.start_all() + snapshots = gateway.get_all_snapshots() + return RpcResponse( + id=request.id, + ok=True, + result={"channels": [s.to_dict() for s in snapshots.values()]}, + ) + + +@validate_params(StopAccountParams) +async def _handle_channels_stop(request: RpcRequest) -> RpcResponse: + validated = request._validated_params + channel_type = validated.channel_type + account_id = validated.account_id + + if channel_type: + snapshot = await gateway.stop_channel(channel_type, account_id) + return RpcResponse(id=request.id, ok=True, result=snapshot.to_dict()) + + await gateway.stop_all() + return RpcResponse(id=request.id, ok=True, result={"stopped": True}) + + +async def _handle_channels_restart(request: RpcRequest) -> RpcResponse: + channel_type = request.params.get("channel_type") if request.params else None + account_id = request.params.get("account_id", "default") if request.params else "default" + + if not channel_type: + return RpcResponse( + id=request.id, + ok=False, + error_code=None, + error_message="restart 需要指定 channel_type", + ) + + await gateway.stop_channel(channel_type, account_id) + snapshot = await gateway.start_channel(channel_type, account_id) + return RpcResponse(id=request.id, ok=True, result=snapshot.to_dict()) + + +async def _handle_channels_configure(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + account_id = params.get("account_id", "default") + + if not channel_type: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 参数") + + try: + snapshot = gateway.get_snapshot(channel_type, account_id) + if snapshot is None: + return RpcResponse( + id=request.id, + ok=False, + error_message=f"渠道 {channel_type}:{account_id} 未找到", + ) + + config_update = params.get("config", {}) + if config_update: + from yuxi.channel.config.reloader import reload_channel_config + + await reload_channel_config(channel_type, account_id, config_update) + + return RpcResponse( + id=request.id, + ok=True, + result={"channel_type": channel_type, "account_id": account_id, "configured": True}, + ) + except Exception: + logger.exception("channels.configure failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="channels.configure 执行异常", + ) + + +async def _handle_chat_send(request: RpcRequest) -> RpcResponse: + params = request.params or {} + query = params.get("query", "") + agent_config_id = params.get("agent_config_id") + thread_id = params.get("thread_id") + request_id = str(uuid.uuid4()) + + if not query: + return RpcResponse(id=request.id, ok=False, error_message="缺少 query 参数") + if not agent_config_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 agent_config_id 参数") + + try: + from yuxi.services.chat_service import agent_chat + + async with pg_manager.get_async_session_context() as db: + current_user = await _resolve_caller_user(request, db) + + meta = { + "source": "gateway_rpc", + "channel_type": params.get("channel_type", "rpc"), + "account_id": params.get("account_id", "default"), + "request_id": request_id, + } + + reply = await agent_chat( + query=query, + agent_config_id=int(agent_config_id), + thread_id=thread_id, + meta=meta, + image_content=params.get("image_content"), + current_user=current_user, + db=db, + ) + + return RpcResponse( + id=request.id, + ok=True, + result={ + "response": reply.get("response", ""), + "thread_id": reply.get("thread_id", thread_id), + "status": reply.get("status", "finished"), + "request_id": request_id, + }, + ) + except Exception: + logger.exception("chat.send failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="chat.send 执行异常", + ) + + +async def _handle_chat_stream(request: RpcRequest) -> AsyncGenerator[RpcEvent, None]: + params = request.params or {} + query = params.get("query", "") + agent_config_id = params.get("agent_config_id") + thread_id = params.get("thread_id") + request_id = str(uuid.uuid4()) + + if not query or not agent_config_id: + yield RpcResponse(id=request.id, ok=False, error_message="缺少 query 或 agent_config_id 参数") + return + + try: + from yuxi.services.chat_service import stream_agent_chat + + async with pg_manager.get_async_session_context() as db: + current_user = await _resolve_caller_user(request, db) + + meta = { + "source": "gateway_rpc", + "channel_type": params.get("channel_type", "rpc"), + "account_id": params.get("account_id", "default"), + "peer_id": params.get("peer_id", ""), + "request_id": request_id, + } + + async for chunk in stream_agent_chat( + query=query, + agent_config_id=int(agent_config_id), + thread_id=thread_id, + meta=meta, + image_content=params.get("image_content"), + current_user=current_user, + db=db, + ): + try: + data = json.loads(chunk.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + + yield RpcEvent( + event="chat.chunk", + data={ + "id": request.id, + "status": data.get("status", ""), + "response": data.get("response", ""), + "thread_id": data.get("thread_id", thread_id), + "request_id": request_id, + }, + ) + + yield RpcResponse(id=request.id, ok=True, result={"status": "finished", "request_id": request_id}) + except Exception: + logger.exception("chat.stream failed") + yield RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="chat.stream 执行异常", + ) + + +async def _handle_chat_history(request: RpcRequest) -> RpcResponse: + params = request.params or {} + thread_id = params.get("thread_id") + limit = params.get("limit", 50) + offset = params.get("offset", 0) + + if not thread_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 thread_id 参数") + + try: + from yuxi.repositories.conversation_repository import ConversationRepository + + async with pg_manager.get_async_session_context() as db: + conv_repo = ConversationRepository(db) + messages = await conv_repo.get_messages_by_thread_id(thread_id, limit=limit, offset=offset) + + return RpcResponse( + id=request.id, + ok=True, + result={ + "thread_id": thread_id, + "messages": [ + { + "id": m.id, + "role": m.role, + "content": m.content, + "message_type": m.message_type, + "created_at": m.created_at.isoformat() if m.created_at else None, + } + for m in messages + ], + }, + ) + except Exception: + logger.exception("chat.history failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="chat.history 执行异常", + ) + + +async def _handle_chat_cancel(request: RpcRequest) -> RpcResponse: + params = request.params or {} + run_id = params.get("run_id") + thread_id = params.get("thread_id") + + if not run_id and not thread_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 run_id 或 thread_id 参数") + + try: + async with pg_manager.get_async_session_context() as db: + current_user = await _resolve_caller_user(request, db) + current_user_id = str(current_user.id) if current_user else None + + if run_id and current_user_id: + from yuxi.services.agent_run_service import cancel_agent_run_view + + async with pg_manager.get_async_session_context() as db: + await cancel_agent_run_view( + run_id=run_id, + current_user_id=current_user_id, + db=db, + ) + elif thread_id and current_user_id: + from yuxi.services.agent_run_service import get_active_run_by_thread + + async with pg_manager.get_async_session_context() as db: + active = await get_active_run_by_thread( + thread_id=thread_id, + current_user_id=current_user_id, + db=db, + ) + if active and active.get("run_id"): + from yuxi.services.agent_run_service import cancel_agent_run_view + + await cancel_agent_run_view( + run_id=active["run_id"], + current_user_id=current_user_id, + db=db, + ) + + return RpcResponse(id=request.id, ok=True, result={"cancelled": True}) + except Exception: + logger.exception("chat.cancel failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="chat.cancel 执行异常", + ) + + +async def _handle_sessions_list(request: RpcRequest) -> RpcResponse: + params = request.params or {} + limit = params.get("limit", 50) + offset = params.get("offset", 0) + + try: + from yuxi.repositories.conversation_repository import ConversationRepository + + async with pg_manager.get_async_session_context() as db: + conv_repo = ConversationRepository(db) + conversations = await conv_repo.list_conversations(limit=limit, offset=offset) + + return RpcResponse( + id=request.id, + ok=True, + result={ + "sessions": [ + { + "id": conv.thread_id, + "title": conv.title, + "created_at": conv.created_at.isoformat() if conv.created_at else None, + "updated_at": conv.updated_at.isoformat() if conv.updated_at else None, + } + for conv in conversations + ], + }, + ) + except Exception: + logger.exception("sessions.list failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="sessions.list 执行异常", + ) + + +async def _handle_sessions_history(request: RpcRequest) -> RpcResponse: + params = request.params or {} + thread_id = params.get("thread_id") + limit = params.get("limit", 100) + offset = params.get("offset", 0) + + if not thread_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 thread_id 参数") + + try: + from yuxi.repositories.conversation_repository import ConversationRepository + + async with pg_manager.get_async_session_context() as db: + conv_repo = ConversationRepository(db) + messages = await conv_repo.get_messages_by_thread_id(thread_id, limit=limit, offset=offset) + + return RpcResponse( + id=request.id, + ok=True, + result={ + "thread_id": thread_id, + "messages": [ + { + "id": m.id, + "role": m.role, + "content": m.content, + "message_type": m.message_type, + "created_at": m.created_at.isoformat() if m.created_at else None, + } + for m in messages + ], + }, + ) + except Exception: + logger.exception("sessions.history failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="sessions.history 执行异常", + ) + + +async def _handle_plugins_list(request: RpcRequest) -> RpcResponse: + from yuxi.channel.plugins.registry import ChannelPluginRegistry + + plugins = ChannelPluginRegistry.list_all() + return RpcResponse( + id=request.id, + ok=True, + result={ + "plugins": [ + { + "id": p.id, + "name": p.name, + "order": p.order, + } + for p in plugins + ], + }, + ) + + +async def _handle_plugins_install(request: RpcRequest) -> RpcResponse: + return RpcResponse( + id=request.id, + ok=True, + result={"message": "plugins.install: 插件动态安装需要通过 Plugin SDK 完成，当前仅支持内置渠道"}, + ) + + +async def _handle_system_config(request: RpcRequest) -> RpcResponse: + snapshots = gateway.get_all_snapshots() + channel_configs = { + channel_type: { + "account_count": len(snapshot.accounts) if hasattr(snapshot, "accounts") else 1, + "status": snapshot.status, + } + for channel_type, snapshot in snapshots.items() + } + return RpcResponse(id=request.id, ok=True, result={"channels": channel_configs}) + + +async def _handle_system_version(request: RpcRequest) -> RpcResponse: + try: + from yuxi import __version__ as yuxi_version + except ImportError: + yuxi_version = "unknown" + return RpcResponse( + id=request.id, + ok=True, + result={"version": yuxi_version, "platform": "yuxi"}, + ) + + +async def _handle_system_log_tail(request: RpcRequest) -> RpcResponse: + params = request.params or {} + lines = params.get("lines", 50) + return RpcResponse( + id=request.id, + ok=True, + result={ + "message": f"system.log.tail: 请求尾部 {lines} 行日志，请通过 HTTP API /api/logs 获取", + "hint": "use GET /api/logs?tail={lines} instead", + }, + ) + + +async def _handle_pairing_list(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + account_id = params.get("account_id", "default") + + if not channel_type: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 参数") + + try: + from yuxi.channel.security.pairing import PairingManager + + manager = PairingManager() + pending = await manager.list_pending(channel_type, account_id) + return RpcResponse( + id=request.id, + ok=True, + result={ + "pending": [ + { + "channel_type": r.channel_type, + "peer_id": r.peer_id, + "account_id": r.account_id, + "code": r.code, + "created_at": r.created_at, + } + for r in pending + ], + }, + ) + except Exception: + logger.exception("pairing.list failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="pairing.list 执行异常", + ) + + +async def _handle_pairing_approve(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + peer_id = params.get("peer_id") + account_id = params.get("account_id", "default") + + if not channel_type or not peer_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 或 peer_id 参数") + + try: + from yuxi.channel.security.allowlist import get_allowlist_checker + from yuxi.channel.security.pairing import PairingManager + + allowlist = await get_allowlist_checker() + manager = PairingManager(allowlist=allowlist) + await manager.approve(channel_type, peer_id, account_id) + return RpcResponse( + id=request.id, + ok=True, + result={"approved": True, "peer_id": peer_id, "channel_type": channel_type}, + ) + except Exception: + logger.exception("pairing.approve failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="pairing.approve 执行异常", + ) + + +async def _handle_pairing_reject(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + peer_id = params.get("peer_id") + account_id = params.get("account_id", "default") + + if not channel_type or not peer_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 或 peer_id 参数") + + try: + from yuxi.channel.security.allowlist import get_allowlist_checker + from yuxi.channel.security.pairing import PairingManager + + allowlist = await get_allowlist_checker() + manager = PairingManager(allowlist=allowlist) + await manager.reject(channel_type, peer_id, account_id) + return RpcResponse( + id=request.id, + ok=True, + result={"rejected": True, "peer_id": peer_id, "channel_type": channel_type}, + ) + except Exception: + logger.exception("pairing.reject failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="pairing.reject 执行异常", + ) + + +async def _handle_allowlist_get(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + + if not channel_type: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 参数") + + try: + from yuxi.channel.security.allowlist import get_allowlist_checker + + checker = await get_allowlist_checker() + await checker.ensure_loaded(channel_type) + result = { + "channel_type": channel_type, + "dm_policy": checker.dm_policy.value, + "group_policy": checker.group_policy.value, + "dm_entries": checker.list_entries(channel_type, "dm"), + "group_entries": checker.list_entries(channel_type, "group"), + "allow_from_entries": checker.list_entries(channel_type, "allow_from"), + } + return RpcResponse(id=request.id, ok=True, result=result) + except Exception: + logger.exception("allowlist.get failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="allowlist.get 执行异常", + ) + + +async def _handle_allowlist_add(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + list_type = params.get("list_type", "dm") + entry = params.get("entry") + + if not channel_type or not entry: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 或 entry 参数") + + try: + from yuxi.channel.security.allowlist import get_allowlist_checker + + checker = await get_allowlist_checker() + await checker.ensure_loaded(channel_type) + await checker.add_entry(channel_type, list_type, entry) + return RpcResponse( + id=request.id, + ok=True, + result={"message": f"将 {entry} 添加到 {channel_type} 的 {list_type} 白名单", "added": True}, + ) + except Exception: + logger.exception("allowlist.add failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="allowlist.add 执行异常", + ) + + +async def _handle_allowlist_remove(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + list_type = params.get("list_type", "dm") + entry = params.get("entry") + + if not channel_type or not entry: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 或 entry 参数") + + try: + from yuxi.channel.security.allowlist import get_allowlist_checker + + checker = await get_allowlist_checker() + await checker.ensure_loaded(channel_type) + await checker.remove_entry(channel_type, list_type, entry) + return RpcResponse( + id=request.id, + ok=True, + result={"message": f"从 {channel_type} 的 {list_type} 白名单移除 {entry}", "removed": True}, + ) + except Exception: + logger.exception("allowlist.remove failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="allowlist.remove 执行异常", + ) + + +async def _handle_cron_list(request: RpcRequest) -> RpcResponse: + from yuxi.channel.cron.engine import get_cron_engine + + engine = get_cron_engine() + jobs = engine.list_jobs() + return RpcResponse( + id=request.id, + ok=True, + result={"jobs": [j.to_dict() for j in jobs]}, + ) + + +async def _handle_cron_create(request: RpcRequest) -> RpcResponse: + params = request.params or {} + job_id = params.get("id") + name = params.get("name", "") + cron_expr = params.get("cron_expr", "*/5 * * * *") + handler_ref = params.get("handler", "") + + if not job_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 id 参数") + + from yuxi.channel.cron.engine import CronJob, get_cron_engine + from yuxi.channel.cron.handler_registry import register_handler + from yuxi.channel.cron.types import CronDeliveryTarget, CronFailureAlert, ScheduleKind + + engine = get_cron_engine() + + schedule_kind_raw = params.get("schedule_kind", "cron") + try: + schedule_kind = ScheduleKind(schedule_kind_raw) + except ValueError: + schedule_kind = ScheduleKind.CRON + schedule_value = params.get("schedule_value", cron_expr) + schedule_tz = params.get("schedule_tz") + stagger_ms = params.get("stagger_ms", 0) + every_ms = params.get("every_ms") + anchor_ms = params.get("anchor_ms") + max_runs = params.get("max_runs", 0) + delete_after_run = params.get("delete_after_run", False) + + delivery_raw = params.get("delivery", []) + delivery = [CronDeliveryTarget.from_dict(t) for t in delivery_raw if isinstance(t, dict)] + failure_alert = CronFailureAlert.from_dict(params.get("failure_alert")) + + handler_name = params.get("handler_name", handler_ref or job_id) + handler = _build_cron_handler(handler_ref, params.get("handler_args", {})) + + try: + job = CronJob( + id=job_id, + name=name, + handler_name=handler_name, + handler=handler, + handler_args=params.get("handler_args", {}), + schedule_kind=schedule_kind, + schedule_value=schedule_value, + schedule_tz=schedule_tz, + stagger_ms=stagger_ms, + every_ms=every_ms, + anchor_ms=anchor_ms, + max_runs=max_runs, + delete_after_run=delete_after_run, + delivery=delivery, + failure_alert=failure_alert, + ) + engine.add_job(job) + if handler_ref: + register_handler(handler_name, handler) + return RpcResponse( + id=request.id, + ok=True, + result={"job": job.to_dict()}, + ) + except Exception as e: + return RpcResponse(id=request.id, ok=False, error_message=str(e)) + + +async def _handle_cron_delete(request: RpcRequest) -> RpcResponse: + params = request.params or {} + job_id = params.get("id") + if not job_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 id 参数") + + from yuxi.channel.cron.engine import get_cron_engine + + engine = get_cron_engine() + engine.remove_job(job_id) + return RpcResponse(id=request.id, ok=True, result={"deleted": True}) + + +async def _handle_cron_update(request: RpcRequest) -> RpcResponse: + params = dict(request.params or {}) + job_id = params.pop("id", None) + if not job_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 id 参数") + + from yuxi.channel.cron.engine import get_cron_engine + + engine = get_cron_engine() + job = engine.update_job(job_id, **params) + if job is None: + return RpcResponse(id=request.id, ok=False, error_message=f"任务 {job_id} 不存在") + return RpcResponse(id=request.id, ok=True, result={"job": job.to_dict()}) + + +async def _handle_cron_force_run(request: RpcRequest) -> RpcResponse: + params = request.params or {} + job_id = params.get("id") + if not job_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 id 参数") + + from yuxi.channel.cron.engine import get_cron_engine + + engine = get_cron_engine() + result = await engine.force_run(job_id) + return RpcResponse(id=request.id, ok=True, result=result) + + +async def _handle_cron_pause(request: RpcRequest) -> RpcResponse: + params = request.params or {} + job_id = params.get("id") + if not job_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 id 参数") + + from yuxi.channel.cron.engine import get_cron_engine + + engine = get_cron_engine() + engine.pause_job(job_id) + return RpcResponse(id=request.id, ok=True, result={"paused": True}) + + +async def _handle_cron_resume(request: RpcRequest) -> RpcResponse: + params = request.params or {} + job_id = params.get("id") + if not job_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 id 参数") + + from yuxi.channel.cron.engine import get_cron_engine + + engine = get_cron_engine() + engine.resume_job(job_id) + return RpcResponse(id=request.id, ok=True, result={"resumed": True}) + + +async def _handle_cron_diagnostics(request: RpcRequest) -> RpcResponse: + params = request.params or {} + job_id = params.get("id") + limit = int(params.get("limit", 50)) + offset = int(params.get("offset", 0)) + + from yuxi.channel.cron.engine import get_cron_engine + + engine = get_cron_engine() + result = engine.read_diagnostics(job_id=job_id, limit=limit, offset=offset) + return RpcResponse(id=request.id, ok=True, result=result) + + +async def _handle_cron_run_log(request: RpcRequest) -> RpcResponse: + params = request.params or {} + job_id = params.get("id") + if not job_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 id 参数") + limit = int(params.get("limit", 50)) + offset = int(params.get("offset", 0)) + + from yuxi.channel.cron.engine import get_cron_engine + + engine = get_cron_engine() + result = engine.read_run_log(job_id=job_id, limit=limit, offset=offset) + return RpcResponse(id=request.id, ok=True, result=result) + + +def _build_cron_handler(handler_ref: str, args: dict): + if handler_ref == "health_check": + + async def _health_check(**kwargs): + from yuxi.channel.runtime.manager import gateway + + report = gateway.get_health() + return { + "status": report.status, + "summary": { + "total": report.summary.total, + "running": report.summary.running, + }, + } + + return _health_check + + if handler_ref == "channel_heartbeat": + + async def _heartbeat(**kwargs): + channel_type = kwargs.get("channel_type", "") + from yuxi.channel.protocols import HeartbeatProtocol + from yuxi.channel.plugins.registry import ChannelPluginRegistry + + plugin = ChannelPluginRegistry.get(channel_type) + if plugin is not None and isinstance(plugin, HeartbeatProtocol): + ok = await plugin.on_ping({}) + return {"ok": ok, "channel_type": channel_type} + return {"ok": False, "error": f"Channel {channel_type} not found"} + + return _heartbeat + + async def _noop(**kwargs): + return {"message": f"noop handler: {handler_ref}"} + + return _noop + + +async def _handle_agent_tools_list(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + if not channel_type: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 参数") + + from yuxi.channel.message.bridge import AgentBridge + + bridge = AgentBridge(lambda **kw: _empty_stream(), None) + tools = bridge.collect_channel_tools(channel_type) + return RpcResponse(id=request.id, ok=True, result={"tools": tools}) + + +async def _handle_agent_tools_execute(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + tool_name = params.get("tool_name") + tool_params = params.get("params", {}) + context = params.get("context", {}) + + if not channel_type or not tool_name: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 或 tool_name 参数") + + from yuxi.channel.message.bridge import AgentBridge + + bridge = AgentBridge(lambda **kw: _empty_stream(), None) + result = await bridge.execute_channel_tool(channel_type, tool_name, tool_params, context) + return RpcResponse(id=request.id, ok=True, result=result) + + +async def _empty_stream(): + if False: + yield b"" + + +async def _handle_sessions_create(request: RpcRequest) -> RpcResponse: + params = request.params or {} + agent_id = params.get("agent_id", "default") + title = params.get("title", "新对话") + thread_id = params.get("thread_id") + user_id = params.get("user_id") + + try: + from yuxi.repositories.conversation_repository import ConversationRepository + + async with pg_manager.get_async_session_context() as db: + conv_repo = ConversationRepository(db) + + if not user_id: + current_user = await _resolve_caller_user(request, db) + user_id = str(current_user.id) if current_user else "system" + + conversation = await conv_repo.create_conversation( + user_id=user_id, + agent_id=agent_id, + title=title, + thread_id=thread_id, + metadata=params.get("metadata"), + ) + + return RpcResponse( + id=request.id, + ok=True, + result={ + "key": conversation.thread_id, + "session_id": conversation.thread_id, + "title": conversation.title, + "agent_id": agent_id, + "created_at": conversation.created_at.isoformat() if conversation.created_at else None, + }, + ) + except Exception: + logger.exception("sessions.create failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="sessions.create 执行异常", + ) + + +async def _handle_sessions_delete(request: RpcRequest) -> RpcResponse: + params = request.params or {} + thread_id = params.get("thread_id") + + if not thread_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 thread_id 参数") + + try: + from yuxi.repositories.conversation_repository import ConversationRepository + + async with pg_manager.get_async_session_context() as db: + conv_repo = ConversationRepository(db) + success = await conv_repo.delete_conversation(thread_id, soft_delete=True) + + return RpcResponse( + id=request.id, + ok=True, + result={"key": thread_id, "deleted": success}, + ) + except Exception: + logger.exception("sessions.delete failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="sessions.delete 执行异常", + ) + + +async def _handle_config_get(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + + try: + from yuxi.repositories.channel_config_repo import ChannelConfigRepository + + repo = ChannelConfigRepository() + if channel_type: + configs = await repo.list_by_type(channel_type) + else: + configs = await repo.list_all() + + return RpcResponse( + id=request.id, + ok=True, + result={ + "configs": [ + c.to_dict() + if hasattr(c, "to_dict") + else {"id": str(c.id), "channel_type": c.channel_type, "enabled": c.enabled} + for c in configs + ], + }, + ) + except Exception: + logger.exception("config.get failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="config.get 执行异常", + ) + + +async def _handle_config_set(request: RpcRequest) -> RpcResponse: + params = request.params or {} + config_id = params.get("config_id") + channel_type = params.get("channel_type") + config_data = params.get("config", {}) + + if not config_id and not channel_type: + return RpcResponse(id=request.id, ok=False, error_message="缺少 config_id 或 channel_type 参数") + + try: + from yuxi.repositories.channel_config_repo import ChannelConfigRepository + + repo = ChannelConfigRepository() + + if config_id: + updated = await repo.update(config_id, config_data) + affected_channel_type = updated.channel_type if updated else channel_type + affected_account_id = config_id + else: + configs = await repo.list_by_type(channel_type) + if not configs: + return RpcResponse( + id=request.id, + ok=False, + error_message=f"渠道 {channel_type} 未找到配置", + ) + updated = await repo.update(str(configs[0].id), config_data) + affected_channel_type = channel_type + affected_account_id = str(configs[0].id) if configs else None + + result = RpcResponse( + id=request.id, + ok=True, + result={ + "updated": True, + "config": updated.to_dict() if updated and hasattr(updated, "to_dict") else None, + }, + ) + + if updated and affected_channel_type and affected_account_id: + try: + from yuxi.channel.config.reloader import reload_channel_config + + await reload_channel_config(affected_channel_type, affected_account_id, config_data) + except Exception: + logger.exception( + "config.set: reload_channel_config failed for %s:%s", + affected_channel_type, + affected_account_id, + ) + + return result + except Exception: + logger.exception("config.set failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="config.set 执行异常", + ) + + +@validate_params(SendMessageParams) +async def _handle_message_send(request: RpcRequest) -> RpcResponse: + validated = request._validated_params + channel_type = validated.channel_type + target_id = validated.target_id + content = validated.text + account_id = validated.account_id + reply_to_id = request.params.get("reply_to_id") if request.params else None + thread_id = request.params.get("thread_id") if request.params else None + + try: + from yuxi.channel.protocols import OutboundProtocol + from yuxi.channel.plugins.registry import ChannelPluginRegistry + + plugin = ChannelPluginRegistry.get(channel_type) + if plugin is None: + return RpcResponse( + id=request.id, + ok=False, + error_message=f"渠道 {channel_type} 未注册", + ) + + if not isinstance(plugin, OutboundProtocol): + return RpcResponse( + id=request.id, + ok=False, + error_message=f"渠道 {channel_type} 不支持出站消息发送", + ) + + await plugin.send_text( + target_id=target_id, + content=content, + reply_to_id=reply_to_id, + thread_id=thread_id, + account_id=account_id, + ) + + return RpcResponse( + id=request.id, + ok=True, + result={ + "channel": channel_type, + "to": target_id, + "sent": True, + }, + ) + except Exception: + logger.exception("message.send failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="message.send 执行异常", + ) + + +async def _handle_message_action(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + action = params.get("action", "") + action_params = params.get("params", {}) + context = params.get("context", {}) + + if not channel_type: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 参数") + if not action: + return RpcResponse(id=request.id, ok=False, error_message="缺少 action 参数") + + try: + from yuxi.channel.sdk.actions.dispatch import ( + DispatchContext, + dispatch_message_action, + ) + + ctx = DispatchContext( + channel_type=channel_type, + action=action, + params=action_params, + account_id=params.get("account_id", "default"), + requester_sender_id=context.get("sender_id"), + session_key=context.get("session_key"), + agent_id=context.get("agent_id"), + ) + result = await dispatch_message_action(ctx) + + return RpcResponse( + id=request.id, + ok=result["success"], + result={ + "channel": channel_type, + "action": action, + "handled": result["success"], + "message": result.get("message", ""), + "data": result.get("data", {}), + } if result["success"] else None, + error_message=result.get("message") if not result["success"] else None, + ) + except Exception: + logger.exception("message.action failed") + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INTERNAL_ERROR, + error_message="message.action 执行异常", + ) + + +async def _handle_identity_links_list(request: RpcRequest) -> RpcResponse: + from yuxi.channel.security.identity_link import IdentityLinkResolver + + resolver = IdentityLinkResolver() + await resolver.ensure_loaded() + return RpcResponse(id=request.id, ok=True, result={ + "links": resolver.list_links(), + }) + + +async def _handle_identity_links_add(request: RpcRequest) -> RpcResponse: + params = request.params or {} + identity = params.get("identity") + channel_type = params.get("channel_type") + peer_id = params.get("peer_id") + + if not all([identity, channel_type, peer_id]): + return RpcResponse(id=request.id, ok=False, error_message="缺少 identity、channel_type 或 peer_id 参数") + + from yuxi.channel.security.identity_link import IdentityLinkResolver + + resolver = IdentityLinkResolver() + await resolver.ensure_loaded() + await resolver.add_link(identity, channel_type, peer_id) + return RpcResponse(id=request.id, ok=True, result={"added": True}) + + +async def _handle_identity_links_remove(request: RpcRequest) -> RpcResponse: + params = request.params or {} + channel_type = params.get("channel_type") + peer_id = params.get("peer_id") + + if not channel_type or not peer_id: + return RpcResponse(id=request.id, ok=False, error_message="缺少 channel_type 或 peer_id 参数") + + from yuxi.channel.security.identity_link import IdentityLinkResolver + + resolver = IdentityLinkResolver() + await resolver.ensure_loaded() + await resolver.remove_link(channel_type, peer_id) + return RpcResponse(id=request.id, ok=True, result={"removed": True}) + + +def register_all_handlers() -> None: + rpc_dispatcher.register(GatewayRpcMethod.SYSTEM_HEALTH, _handle_health) + rpc_dispatcher.register(GatewayRpcMethod.SYSTEM_CONFIG, _handle_system_config) + rpc_dispatcher.register(GatewayRpcMethod.SYSTEM_VERSION, _handle_system_version) + rpc_dispatcher.register(GatewayRpcMethod.SYSTEM_LOG_TAIL, _handle_system_log_tail) + rpc_dispatcher.register(GatewayRpcMethod.CHANNELS_LIST, _handle_channels_list) + rpc_dispatcher.register(GatewayRpcMethod.CHANNELS_STATUS, _handle_channels_status) + rpc_dispatcher.register(GatewayRpcMethod.CHANNELS_PROBE, _handle_channels_probe) + rpc_dispatcher.register(GatewayRpcMethod.CHANNELS_START, _handle_channels_start) + rpc_dispatcher.register(GatewayRpcMethod.CHANNELS_STOP, _handle_channels_stop) + rpc_dispatcher.register(GatewayRpcMethod.CHANNELS_RESTART, _handle_channels_restart) + rpc_dispatcher.register(GatewayRpcMethod.CHANNELS_CONFIGURE, _handle_channels_configure) + rpc_dispatcher.register(GatewayRpcMethod.CHAT_SEND, _handle_chat_send) + rpc_dispatcher.register_stream(GatewayRpcMethod.CHAT_STREAM, _handle_chat_stream) + rpc_dispatcher.register(GatewayRpcMethod.CHAT_HISTORY, _handle_chat_history) + rpc_dispatcher.register(GatewayRpcMethod.CHAT_CANCEL, _handle_chat_cancel) + rpc_dispatcher.register(GatewayRpcMethod.SESSIONS_LIST, _handle_sessions_list) + rpc_dispatcher.register(GatewayRpcMethod.SESSIONS_HISTORY, _handle_sessions_history) + rpc_dispatcher.register(GatewayRpcMethod.PLUGINS_LIST, _handle_plugins_list) + rpc_dispatcher.register(GatewayRpcMethod.PLUGINS_INSTALL, _handle_plugins_install) + rpc_dispatcher.register(GatewayRpcMethod.PAIRING_LIST, _handle_pairing_list) + rpc_dispatcher.register(GatewayRpcMethod.PAIRING_APPROVE, _handle_pairing_approve) + rpc_dispatcher.register(GatewayRpcMethod.PAIRING_REJECT, _handle_pairing_reject) + rpc_dispatcher.register(GatewayRpcMethod.ALLOWLIST_GET, _handle_allowlist_get) + rpc_dispatcher.register(GatewayRpcMethod.ALLOWLIST_ADD, _handle_allowlist_add) + rpc_dispatcher.register(GatewayRpcMethod.ALLOWLIST_REMOVE, _handle_allowlist_remove) + rpc_dispatcher.register(GatewayRpcMethod.CRON_LIST, _handle_cron_list) + rpc_dispatcher.register(GatewayRpcMethod.CRON_CREATE, _handle_cron_create) + rpc_dispatcher.register(GatewayRpcMethod.CRON_UPDATE, _handle_cron_update) + rpc_dispatcher.register(GatewayRpcMethod.CRON_DELETE, _handle_cron_delete) + rpc_dispatcher.register(GatewayRpcMethod.CRON_FORCE_RUN, _handle_cron_force_run) + rpc_dispatcher.register(GatewayRpcMethod.CRON_PAUSE, _handle_cron_pause) + rpc_dispatcher.register(GatewayRpcMethod.CRON_RESUME, _handle_cron_resume) + rpc_dispatcher.register(GatewayRpcMethod.CRON_DIAGNOSTICS, _handle_cron_diagnostics) + rpc_dispatcher.register(GatewayRpcMethod.CRON_RUN_LOG, _handle_cron_run_log) + rpc_dispatcher.register(GatewayRpcMethod.AGENT_TOOLS_LIST, _handle_agent_tools_list) + rpc_dispatcher.register(GatewayRpcMethod.AGENT_TOOLS_EXECUTE, _handle_agent_tools_execute) + rpc_dispatcher.register(GatewayRpcMethod.SESSIONS_CREATE, _handle_sessions_create) + rpc_dispatcher.register(GatewayRpcMethod.SESSIONS_DELETE, _handle_sessions_delete) + rpc_dispatcher.register(GatewayRpcMethod.CONFIG_GET, _handle_config_get) + rpc_dispatcher.register(GatewayRpcMethod.CONFIG_SET, _handle_config_set) + rpc_dispatcher.register(GatewayRpcMethod.MESSAGE_SEND, _handle_message_send) + rpc_dispatcher.register(GatewayRpcMethod.MESSAGE_ACTION, _handle_message_action) + rpc_dispatcher.register(GatewayRpcMethod.IDENTITY_LINKS_LIST, _handle_identity_links_list) + rpc_dispatcher.register(GatewayRpcMethod.IDENTITY_LINKS_ADD, _handle_identity_links_add) + rpc_dispatcher.register(GatewayRpcMethod.IDENTITY_LINKS_REMOVE, _handle_identity_links_remove) + + rpc_dispatcher.add_alias("health", GatewayRpcMethod.SYSTEM_HEALTH) + rpc_dispatcher.add_alias("config", GatewayRpcMethod.SYSTEM_CONFIG) + rpc_dispatcher.add_alias("version", GatewayRpcMethod.SYSTEM_VERSION) + rpc_dispatcher.add_alias("chat.message", GatewayRpcMethod.CHAT_SEND) + rpc_dispatcher.add_alias("channels", GatewayRpcMethod.CHANNELS_LIST) + rpc_dispatcher.add_alias("channels.start", GatewayRpcMethod.CHANNELS_START) + rpc_dispatcher.add_alias("channels.stop", GatewayRpcMethod.CHANNELS_STOP) + rpc_dispatcher.add_alias("channels.restart", GatewayRpcMethod.CHANNELS_RESTART) + rpc_dispatcher.add_alias("sessions", GatewayRpcMethod.SESSIONS_LIST) + + logger.info("Gateway RPC handlers registered: %s", rpc_dispatcher.list_methods()) diff --git a/backend/package/yuxi/channel/gateway/server.py b/backend/package/yuxi/channel/gateway/server.py new file mode 100644 index 00000000..b28a8e72 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/server.py @@ -0,0 +1,382 @@ +import asyncio +import logging +import os +import time as _time +from collections.abc import Callable, Collection + +from fastapi import WebSocket, WebSocketDisconnect + +from yuxi.channel.gateway.auth import GatewayAuthResult, authenticate_gateway_connect +from yuxi.channel.gateway.broadcaster import BroadcastFilter, BroadcastResult, gateway_broadcaster +from yuxi.channel.gateway.net_utils import ( + _ENV_ALLOW_INSECURE_PRIVATE_WS, + is_loopback_address, + is_secure_ws_url, +) +from yuxi.channel.gateway.protocol import ( + GatewayErrorCode, + GatewayRpcMethod, + HelloOk, + RpcEvent, + RpcRequest, + RpcResponse, + marshal_frame, + unmarshal_frame, +) +from yuxi.channel.config.defaults import TIMEOUT +from yuxi.channel.gateway.rbac import GatewayRole +from yuxi.channel.gateway.rpc_dispatcher import rpc_dispatcher + +logger = logging.getLogger(__name__) + +OnConnectCallback = Callable[[str, GatewayAuthResult], None] +OnDisconnectCallback = Callable[[str, GatewayAuthResult], None] + +HEARTBEAT_INTERVAL = TIMEOUT.gateway.heartbeat_interval +HEARTBEAT_TIMEOUT = TIMEOUT.gateway.heartbeat +SEND_TIMEOUT = TIMEOUT.gateway.send +_SEND_SLOW_WARN_THRESHOLD = 5.0 +CONNECT_NEGOTIATION_TIMEOUT = TIMEOUT.gateway.negotiation +TICK_INTERVAL_MS = 30_000 + + +async def _send_with_timeout(ws: WebSocket, text: str, timeout: float = SEND_TIMEOUT) -> None: + try: + await asyncio.wait_for(ws.send_text(text), timeout=timeout) + except TimeoutError: + logger.warning( + "Gateway WS send timed out after %.0fs, closing connection", + timeout, + ) + raise + + +class GatewayWsServer: + def __init__(self, allow_private_ws: bool | None = None): + self._active_connections: dict[str, WebSocket] = {} + self._connection_auth: dict[str, GatewayAuthResult] = {} + self._heartbeat_tasks: dict[str, asyncio.Task] = {} + self._tick_tasks: dict[str, asyncio.Task] = {} + self._on_connect: list[OnConnectCallback] = [] + self._on_disconnect: list[OnDisconnectCallback] = [] + self._shutting_down = False + self._allow_private_ws = ( + allow_private_ws + if allow_private_ws is not None + else os.environ.get(_ENV_ALLOW_INSECURE_PRIVATE_WS) == "1" + ) + + @property + def active_count(self) -> int: + return len(self._active_connections) + + @property + def shutting_down(self) -> bool: + return self._shutting_down + + @property + def allow_private_ws(self) -> bool: + return self._allow_private_ws + + def set_allow_private_ws(self, value: bool) -> None: + self._allow_private_ws = value + + async def handle_connection(self, ws: WebSocket, token: str | None = None): + client_ip = _resolve_ws_client_ip(ws) + is_loopback_client = is_loopback_address(client_ip) + + if not is_loopback_client and not is_secure_ws_url(str(ws.url), allow_private_ws=self._allow_private_ws): + display_host = _format_ws_display_host(str(ws.url)) + logger.warning( + "Gateway WS: rejected insecure ws:// connection from %s to %s", + client_ip or "unknown", + display_host, + ) + await ws.close(code=4400, reason="insecure_ws_connection") + return + + await ws.accept() + + auth_header = ws.headers.get("authorization") + auth_result = await authenticate_gateway_connect( + auth_header, token, client_ip=client_ip, headers=dict(ws.headers), remote_addr=client_ip + ) + if not auth_result.authenticated: + await _send_with_timeout( + ws, + marshal_frame( + RpcResponse( + id="auth", + ok=False, + error_code=GatewayErrorCode.AUTH_ERROR, + error_message=auth_result.error or "认证失败", + ) + ), + ) + await ws.close(code=4001, reason="unauthorized") + return + + conn_id = f"{auth_result.user_id}-{id(ws):x}" + self._active_connections[conn_id] = ws + self._connection_auth[conn_id] = auth_result + + async def _send_to_ws(text: str) -> None: + await ws.send_text(text) + + gateway_broadcaster.register_connection(conn_id, auth_result, _send_to_ws) + + heartbeat_task = asyncio.create_task( + self._heartbeat_loop(conn_id, ws), + name=f"gateway_heartbeat:{conn_id}", + ) + self._heartbeat_tasks[conn_id] = heartbeat_task + + tick_task = asyncio.create_task( + self._tick_loop(conn_id, ws), + name=f"gateway_tick:{conn_id}", + ) + self._tick_tasks[conn_id] = tick_task + + try: + first_frame = await self._negotiate_connect(conn_id, ws, auth_result) + except Exception: + logger.warning("Gateway connect negotiation failed: %s", conn_id) + first_frame = None + + try: + for cb in self._on_connect: + try: + cb(conn_id, auth_result) + except Exception: + logger.exception("on_connect callback failed: %s", conn_id) + + logger.info( + "Gateway WS connected: %s (user=%s, roles=%s)", + conn_id, + auth_result.user_id, + [r.value for r in auth_result.roles], + ) + + if first_frame is not None: + await self._dispatch_frame(conn_id, ws, first_frame) + + while True: + raw = await ws.receive_text() + try: + frame = unmarshal_frame(raw) + except Exception: + logger.warning("Gateway WS: 无效帧 from %s", conn_id) + await _send_with_timeout( + ws, + marshal_frame( + RpcResponse( + id="", + ok=False, + error_code=GatewayErrorCode.INVALID_REQUEST, + error_message="无效的 RPC 帧格式", + ) + ), + ) + continue + + await self._dispatch_frame(conn_id, ws, frame) + + except WebSocketDisconnect: + logger.info("Gateway WS disconnected: %s", conn_id) + except Exception: + logger.exception("Gateway WS error: %s", conn_id) + finally: + gateway_broadcaster.unregister_connection(conn_id) + for task_key, task_dict in [("heartbeat", self._heartbeat_tasks), ("tick", self._tick_tasks)]: + task = task_dict.get(conn_id) + if task and not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + task_dict.pop(conn_id, None) + self._active_connections.pop(conn_id, None) + auth = self._connection_auth.pop(conn_id, None) + for cb in self._on_disconnect: + try: + cb(conn_id, auth or GatewayAuthResult(authenticated=False)) + except Exception: + logger.exception("on_disconnect callback failed: %s", conn_id) + + async def _heartbeat_loop(self, conn_id: str, ws: WebSocket) -> None: + try: + while True: + await asyncio.sleep(HEARTBEAT_INTERVAL) + try: + await ws.send_text(marshal_frame(RpcEvent(event="ping"))) + except Exception: + logger.warning("Gateway WS heartbeat failed: %s", conn_id) + break + except asyncio.CancelledError: + pass + + async def _negotiate_connect(self, conn_id: str, ws: WebSocket, auth_result: GatewayAuthResult): + try: + raw = await asyncio.wait_for(ws.receive_text(), timeout=CONNECT_NEGOTIATION_TIMEOUT) + frame = unmarshal_frame(raw) + if isinstance(frame, RpcRequest) and frame.method == GatewayRpcMethod.CONNECT: + hello_ok = self._build_hello_ok(conn_id, auth_result) + await _send_with_timeout(ws, marshal_frame(hello_ok)) + return None + return frame + except TimeoutError: + return None + + def _build_hello_ok(self, conn_id: str, auth_result: GatewayAuthResult) -> HelloOk: + hello = HelloOk() + hello.result["server"]["connId"] = conn_id + hello.result["auth"] = { + "role": auth_result.roles[0].value if auth_result.roles else "viewer", + "scopes": [r.value for r in auth_result.roles], + } + return hello + + async def _dispatch_frame(self, conn_id: str, ws: WebSocket, frame): + auth_result = self._connection_auth.get(conn_id) + if auth_result is None: + return + if isinstance(frame, RpcRequest): + frame.caller_user_id = auth_result.user_id + if frame.method == "event.subscribe": + resp = await self._handle_event_subscribe(conn_id, frame) + await _send_with_timeout(ws, marshal_frame(resp)) + elif frame.method == "event.unsubscribe": + resp = await self._handle_event_unsubscribe(conn_id, frame) + await _send_with_timeout(ws, marshal_frame(resp)) + elif rpc_dispatcher.is_stream(frame.method): + async for item in rpc_dispatcher.dispatch_stream(frame, auth_result.roles): + await _send_with_timeout(ws, marshal_frame(item), timeout=TIMEOUT.gateway.stream_rpc) + else: + resp = await rpc_dispatcher.dispatch(frame, auth_result.roles) + await _send_with_timeout(ws, marshal_frame(resp)) + + async def _tick_loop(self, conn_id: str, ws: WebSocket) -> None: + try: + while True: + await asyncio.sleep(TICK_INTERVAL_MS / 1000.0) + tick_event = RpcEvent( + event="tick", + data={"ts": int(_time.time() * 1000)}, + ) + await _send_with_timeout(ws, marshal_frame(tick_event)) + except asyncio.CancelledError: + pass + except Exception: + logger.warning("Gateway tick failed for %s", conn_id) + + async def _handle_event_subscribe(self, conn_id: str, req: RpcRequest) -> RpcResponse: + events = req.params.get("events", []) if req.params else [] + if not isinstance(events, list) or not all(isinstance(e, str) for e in events): + return RpcResponse( + id=req.id, + ok=False, + error_code=GatewayErrorCode.INVALID_PARAMS, + error_message="events 必须是字符串数组", + ) + gateway_broadcaster.subscribe(conn_id, events) + subscribed = gateway_broadcaster.get_subscriptions(conn_id) + return RpcResponse(id=req.id, ok=True, result={"subscribed": list(subscribed)}) + + async def _handle_event_unsubscribe(self, conn_id: str, req: RpcRequest) -> RpcResponse: + events = req.params.get("events", []) if req.params else [] + if not isinstance(events, list) or not all(isinstance(e, str) for e in events): + return RpcResponse( + id=req.id, + ok=False, + error_code=GatewayErrorCode.INVALID_PARAMS, + error_message="events 必须是字符串数组", + ) + gateway_broadcaster.unsubscribe(conn_id, events) + subscribed = gateway_broadcaster.get_subscriptions(conn_id) + return RpcResponse(id=req.id, ok=True, result={"subscribed": list(subscribed)}) + + async def send_event(self, conn_id: str, event: RpcEvent) -> bool: + return await gateway_broadcaster.send_event(conn_id, event.event, event.data) + + async def broadcast_event(self, event: RpcEvent) -> int: + result = await gateway_broadcaster.broadcast(event.event, event.data) + return result.sent + + async def broadcast_event_filtered( + self, + event: str, + data: dict | None = None, + *, + conn_ids: Collection[str] | None = None, + user_ids: Collection[str] | None = None, + roles: Collection[GatewayRole] | None = None, + drop_if_slow: bool = True, + ) -> BroadcastResult: + filter_ = BroadcastFilter( + conn_ids=conn_ids, + user_ids=user_ids, + roles=roles, + drop_if_slow=drop_if_slow, + ) + return await gateway_broadcaster.broadcast(event, data, filter_) + + def register_on_connect(self, cb: OnConnectCallback) -> None: + self._on_connect.append(cb) + + def register_on_disconnect(self, cb: OnDisconnectCallback) -> None: + self._on_disconnect.append(cb) + + async def shutdown(self, timeout: float = 10.0) -> None: + self._shutting_down = True + logger.info("Gateway WS server shutting down (%d connections)", len(self._active_connections)) + + await gateway_broadcaster.broadcast("server.shutdown", {"reason": "server_shutdown"}) + + for task_dict in [self._heartbeat_tasks, self._tick_tasks]: + for conn_id, task in list(task_dict.items()): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + gateway_broadcaster.unregister_connection(conn_id) + + async def _close_conn(conn_id: str, ws: WebSocket) -> None: + try: + await asyncio.wait_for(ws.close(code=1001, reason="server_shutdown"), timeout=timeout) + except Exception: + pass + + tasks = [_close_conn(cid, ws) for cid, ws in self._active_connections.items()] + await asyncio.gather(*tasks, return_exceptions=True) + + self._active_connections.clear() + self._connection_auth.clear() + self._heartbeat_tasks.clear() + self._tick_tasks.clear() + logger.info("Gateway WS server shut down complete") + + +gateway_ws_server = GatewayWsServer() + + +def _resolve_ws_client_ip(ws: WebSocket) -> str | None: + client = getattr(ws, "client", None) + if client and hasattr(client, "host"): + return client.host + return None + + +def _format_ws_display_host(url_str: str) -> str: + from urllib.parse import urlparse + + try: + parsed = urlparse(url_str) + host = parsed.hostname or parsed.netloc or "unknown" + port = parsed.port + if port and port not in (80, 443): + return f"{host}:{port}" + return host + except Exception: + return url_str diff --git a/backend/package/yuxi/channel/gateway/sse.py b/backend/package/yuxi/channel/gateway/sse.py new file mode 100644 index 00000000..afa6c05c --- /dev/null +++ b/backend/package/yuxi/channel/gateway/sse.py @@ -0,0 +1,277 @@ +"""Gateway HTTP SSE 端点 — 为不支持 WebSocket 的客户端提供备选流式通道。 + +复用 stream_agent_chat 的流式输出能力，通过 HTTP SSE (text/event-stream) 推送给浏览器客户端。 + +SSE 事件类型对齐 CowAgent 标准： + reasoning / delta / tool_start / tool_end / message_end / + phase / image / file / video / done / error + +路由： + POST /api/sse/chat — 发送聊天消息，返回 request_id + GET /api/sse/stream — 订阅 SSE 事件流 + GET /api/poll — Polling 降级轮询 +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import time +import uuid + +from fastapi import APIRouter, Header, Query, Request +from fastapi.responses import StreamingResponse + +from yuxi.channel.protocols import SseEventType +from yuxi.storage.postgres.manager import pg_manager + +logger = logging.getLogger(__name__) + +SSE_KEEP_ALIVE_SEC = 15.0 +SSE_QUEUE_TTL_SEC = 600 +MAX_QUEUE_SIZE = 256 + + +def _sse_frame(event_type: str, data: dict | str) -> str: + content = json.dumps(data, ensure_ascii=False) if isinstance(data, dict) else data + return f"event: {event_type}\ndata: {content}\n\n" + + +def _sse_comment(comment: str) -> str: + return f": {comment}\n\n" + + +class GatewaySseEndpoint: + def __init__(self): + self._queues: dict[str, asyncio.Queue[dict]] = {} + self._tasks: dict[str, asyncio.Task] = {} + self._last_active: dict[str, float] = {} + + def create(self, rid: str, q: asyncio.Queue[dict], task: asyncio.Task) -> None: + self._queues[rid] = q + self._tasks[rid] = task + self._last_active[rid] = time.monotonic() + + def get(self, rid: str) -> asyncio.Queue[dict] | None: + return self._queues.get(rid) + + def touch(self, rid: str) -> None: + self._last_active[rid] = time.monotonic() + + def remove(self, rid: str) -> None: + self._queues.pop(rid, None) + task = self._tasks.pop(rid, None) + if task and not task.done(): + task.cancel() + self._last_active.pop(rid, None) + + def cleanup_stale(self) -> int: + now = time.monotonic() + stale = [rid for rid, ts in self._last_active.items() if now - ts > SSE_QUEUE_TTL_SEC] + for rid in stale: + self.remove(rid) + return len(stale) + + +gateway_sse_endpoint = GatewaySseEndpoint() + + +async def _run_chat_feed( + query: str, + agent_config_id: int, + thread_id: str | None, + image_content: str | None, + current_user, + db, + q: asyncio.Queue[dict], + rid: str, +) -> None: + from yuxi.services.chat_service import stream_agent_chat + + meta = { + "source": "gateway_sse", + "channel_type": "sse", + "account_id": "default", + "request_id": rid, + } + + try: + async for chunk in stream_agent_chat( + query=query, + agent_config_id=agent_config_id, + thread_id=thread_id, + meta=meta, + image_content=image_content, + current_user=current_user, + db=db, + ): + try: + data = json.loads(chunk.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + + status = data.get("status", "") + content = data.get("response", "") + + if status == "error": + await q.put({"type": SseEventType.ERROR, "data": {"message": data.get("error_message", content)}}) + await q.put({"type": SseEventType.DONE, "data": {}}) + return + + if status == "init": + continue + + if status in ("streaming", "reasoning"): + event_type = SseEventType.REASONING if data.get("reasoning") else SseEventType.DELTA + if content: + await q.put({"type": event_type, "data": {"content": content}}) + + if status == "thinking": + await q.put({"type": SseEventType.PHASE, "data": {"content": content or "thinking"}}) + + if status == "finished": + final_content = data.get("final_response", content) + await q.put( + { + "type": SseEventType.MESSAGE_END, + "data": {"content": final_content, "thread_id": data.get("thread_id", thread_id)}, + } + ) + + await q.put({"type": SseEventType.DONE, "data": {}}) + except asyncio.CancelledError: + pass + except Exception: + logger.exception("SSE chat stream failed for %s", rid) + await q.put({"type": SseEventType.ERROR, "data": {"message": "SSE stream error"}}) + await q.put({"type": SseEventType.DONE, "data": {}}) + + +async def _resolve_sse_user(db, authorization: str | None): + from sqlalchemy import select + + from server.utils.auth_utils import AuthUtils + from yuxi.storage.postgres.models_business import User + + if authorization and authorization.startswith("Bearer "): + token = authorization[7:] + try: + payload = AuthUtils.verify_access_token(token) + user_id = payload.get("sub") + if user_id: + result = await db.execute(select(User).where(User.id == int(user_id))) + user = result.scalar_one_or_none() + if user: + return user + except Exception: + pass + + result = await db.execute( + select(User).where(User.is_deleted == 0).order_by(User.id).limit(1) + ) + user = result.scalar_one_or_none() + if user: + logger.warning("SSE endpoint using fallback user id=%s (no valid auth provided)", user.id) + return user + + +router = APIRouter(prefix="/api/sse", tags=["sse"]) + + +@router.post("/chat") +async def sse_chat( + query: str = Query(...), + agent_config_id: int = Query(...), + thread_id: str | None = Query(None), + image_content: str | None = Query(None), + authorization: str | None = Header(None), +): + rid = str(uuid.uuid4()) + + async with pg_manager.get_async_session_context() as db: + current_user = await _resolve_sse_user(db, authorization) + + q: asyncio.Queue[dict] = asyncio.Queue(maxsize=MAX_QUEUE_SIZE) + + task = asyncio.create_task( + _run_chat_feed( + query=query, + agent_config_id=agent_config_id, + thread_id=thread_id, + image_content=image_content, + current_user=current_user, + db=db, + q=q, + rid=rid, + ), + name=f"sse_chat:{rid}", + ) + gateway_sse_endpoint.create(rid, q, task) + + return {"request_id": rid, "stream_url": f"/api/sse/stream?request_id={rid}"} + + +@router.get("/stream") +async def sse_stream(request: Request, request_id: str = Query(...)): + q = gateway_sse_endpoint.get(request_id) + if q is None: + return StreamingResponse( + iter([_sse_frame(SseEventType.ERROR, {"message": "request_id invalid or expired"})]), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + status_code=404, + ) + + gateway_sse_endpoint.touch(request_id) + + async def generate(): + try: + while True: + if await request.is_disconnected(): + logger.info("SSE client disconnected: %s", request_id) + break + try: + event = await asyncio.wait_for(q.get(), timeout=SSE_KEEP_ALIVE_SEC) + yield _sse_frame(str(event["type"]), event["data"]) + if event["type"] == SseEventType.DONE: + break + except TimeoutError: + yield _sse_comment("keepalive") + finally: + gateway_sse_endpoint.remove(request_id) + + return StreamingResponse( + generate(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +# ── OPT-7: 实时日志 SSE ────────────────────────────────── + + +@router.get("/logs") +async def sse_logs( + request: Request, + levels: str | None = Query(None, description="comma-separated: DEBUG,INFO,WARNING,ERROR"), +): + logger.info("SSE log stream requested (not yet implemented), levels=%s", levels) + + async def generate(): + yield _sse_frame(SseEventType.ERROR, {"message": "SSE log streaming is not yet implemented"}) + yield _sse_frame(SseEventType.DONE, {}) + + return StreamingResponse( + generate(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) diff --git a/backend/package/yuxi/channel/gateway/tailscale_auth.py b/backend/package/yuxi/channel/gateway/tailscale_auth.py new file mode 100644 index 00000000..a063bae3 --- /dev/null +++ b/backend/package/yuxi/channel/gateway/tailscale_auth.py @@ -0,0 +1,173 @@ +import asyncio +import json +import logging +import time +from dataclasses import dataclass + +logger = logging.getLogger(__name__) + +WHOIS_CACHE_TTL_MS = 60_000 +WHOIS_ERROR_TTL_MS = 5_000 +WHOIS_TIMEOUT_S = 5.0 + + +@dataclass +class TailscaleWhoisIdentity: + login: str + name: str | None = None + + +@dataclass +class _CacheEntry: + value: TailscaleWhoisIdentity | None + expires_at: float + + +_whois_cache: dict[str, _CacheEntry] = {} + + +def _read_cached_whois(ip: str, now_ms: float) -> TailscaleWhoisIdentity | None | bool: + """Returns None=cache miss, False=cached null, or the identity.""" + entry = _whois_cache.get(ip) + if entry is None: + return None + if entry.expires_at <= now_ms: + del _whois_cache[ip] + return None + return entry.value if entry.value is not None else False # type: ignore[return-value] + + +def _write_cached_whois(ip: str, value: TailscaleWhoisIdentity | None, ttl_ms: float) -> None: + _whois_cache[ip] = _CacheEntry(value=value, expires_at=time.time() * 1000 + ttl_ms) + + +def _parse_whois_identity(payload: dict) -> TailscaleWhoisIdentity | None: + user_profile = payload.get("UserProfile") or payload.get("userProfile") or payload.get("User") + if not isinstance(user_profile, dict): + return None + login = ( + user_profile.get("LoginName") + or user_profile.get("Login") + or user_profile.get("loginName") + or user_profile.get("login") + ) + if not login or not isinstance(login, str): + return None + name = ( + user_profile.get("DisplayName") + or user_profile.get("displayName") + or user_profile.get("Name") + or user_profile.get("name") + ) + return TailscaleWhoisIdentity( + login=login.lower().strip(), + name=name.strip() if isinstance(name, str) and name.strip() else None, + ) + + +def _parse_possibly_noisy_json(stdout: str) -> dict: + trimmed = stdout.strip() + start = trimmed.find("{") + end = trimmed.rfind("}") + if start >= 0 and end > start: + return json.loads(trimmed[start : end + 1]) + return json.loads(trimmed) + + +async def _get_tailscale_binary() -> str | None: + try: + proc = await asyncio.create_subprocess_exec( + "which", + "tailscale", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, + ) + stdout, _ = await proc.communicate() + if proc.returncode == 0 and stdout: + return stdout.decode().strip() + except FileNotFoundError: + logger.debug("tailscale binary not found via 'which'") + except Exception: + logger.debug("tailscale binary lookup failed via 'which'", exc_info=True) + + import os + + mac_path = "/Applications/Tailscale.app/Contents/MacOS/Tailscale" + if os.path.exists(mac_path) and os.access(mac_path, os.X_OK): + return mac_path + return None + + +async def read_tailscale_whois_identity( + ip: str, + timeout_s: float = WHOIS_TIMEOUT_S, + cache_ttl_ms: float = WHOIS_CACHE_TTL_MS, + error_ttl_ms: float = WHOIS_ERROR_TTL_MS, +) -> TailscaleWhoisIdentity | None: + normalized = ip.strip() + if not normalized: + return None + + now_ms = time.time() * 1000 + cached = _read_cached_whois(normalized, now_ms) + if cached is not None: + return cached if cached is not False else None + + tailscale_bin = await _get_tailscale_binary() + if not tailscale_bin: + logger.debug("tailscale binary not found, whois unavailable") + _write_cached_whois(normalized, None, error_ttl_ms) + return None + + proc: asyncio.subprocess.Process | None = None + try: + proc = await asyncio.create_subprocess_exec( + tailscale_bin, + "whois", + "--json", + normalized, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await asyncio.wait_for( + proc.communicate(), + timeout=timeout_s, + ) + if proc.returncode != 0: + logger.debug("tailscale whois failed for %s: %s", normalized, stderr.decode().strip()) + _write_cached_whois(normalized, None, error_ttl_ms) + return None + payload = _parse_possibly_noisy_json(stdout.decode()) + identity = _parse_whois_identity(payload) + _write_cached_whois(normalized, identity, cache_ttl_ms) + return identity + except TimeoutError: + logger.debug("tailscale whois timeout for %s", normalized) + if proc is not None: + try: + proc.kill() + await proc.wait() + except Exception: + pass + _write_cached_whois(normalized, None, error_ttl_ms) + return None + except Exception: + logger.debug("tailscale whois error for %s", normalized, exc_info=True) + if proc is not None: + try: + proc.kill() + await proc.wait() + except Exception: + pass + _write_cached_whois(normalized, None, error_ttl_ms) + return None + + +def get_tailscale_user_from_headers(headers: dict | None) -> tuple[str, str] | None: + if not headers: + return None + login = headers.get("tailscale-user-login") + if not login: + return None + name = headers.get("tailscale-user-name") or login + return (login.lower().strip(), name.strip() if isinstance(name, str) else login) diff --git a/backend/package/yuxi/channel/gateway/validation.py b/backend/package/yuxi/channel/gateway/validation.py new file mode 100644 index 00000000..d7f008dc --- /dev/null +++ b/backend/package/yuxi/channel/gateway/validation.py @@ -0,0 +1,110 @@ +from functools import wraps +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from yuxi.channel.gateway.protocol import GatewayErrorCode, RpcRequest, RpcResponse + + +class ErrorShape(BaseModel): + code: str + message: str = "" + data: Any | None = None + retryable: bool = False + retryAfterMs: int | None = None + + +class RequestFrameSchema(BaseModel): + type: str = "request" + id: str + method: str + params: dict | None = None + sessionId: str | None = None + + +class ResponseFrameSchema(BaseModel): + type: str = "response" + id: str + ok: bool = True + result: dict | None = None + error: ErrorShape | None = None + + +class EventFrameSchema(BaseModel): + type: str = "event" + event: str + data: dict | None = None + timestamp: float = 0.0 + seq: int | None = None + stateVersion: int | None = None + + +class StartAccountParams(BaseModel): + model_config = ConfigDict(extra="allow") + + channel_type: str = Field(..., min_length=1, max_length=50) + account_id: str = Field(default="default", min_length=1, max_length=100) + config: dict = Field(default_factory=dict) + + +class StopAccountParams(BaseModel): + model_config = ConfigDict(extra="allow") + + channel_type: str = Field(..., min_length=1, max_length=50) + account_id: str = Field(default="default", min_length=1, max_length=100) + force: bool = False + + +class SendMessageParams(BaseModel): + model_config = ConfigDict(extra="allow") + + channel_type: str = Field(..., min_length=1, max_length=50) + account_id: str = Field(default="default", min_length=1, max_length=100) + target_id: str = Field(..., min_length=1, max_length=200) + text: str = Field(..., min_length=1, max_length=4096) + msg_type: str = Field(default="text") + media_url: str | None = None + extra: dict = Field(default_factory=dict) + + +class ProbeParams(BaseModel): + model_config = ConfigDict(extra="allow") + + channel_type: str = Field(..., min_length=1, max_length=50) + account_id: str = Field(default="default", min_length=1, max_length=100) + + +class DiagnoseParams(BaseModel): + model_config = ConfigDict(extra="allow") + + channel_type: str = Field(..., min_length=1, max_length=50) + account_id: str = Field(default="default", min_length=1, max_length=100) + + +class RepairParams(BaseModel): + model_config = ConfigDict(extra="allow") + + channel_type: str = Field(..., min_length=1, max_length=50) + step_id: str = Field(..., min_length=1, max_length=100) + account_id: str = Field(default="default", min_length=1, max_length=100) + + +def validate_params(schema_cls: type[BaseModel]): + def decorator(func): + @wraps(func) + async def wrapper(request: RpcRequest, *args, **kwargs): + try: + validated = schema_cls(**(request.params or {})) + except ValidationError as e: + return RpcResponse( + id=request.id, + ok=False, + error_code=GatewayErrorCode.INVALID_PARAMS, + error_message=f"参数校验失败: {e}", + ) + request._validated_params = validated + return await func(request, *args, **kwargs) + + return wrapper + + return decorator diff --git a/backend/package/yuxi/channel/gateway/webhook_security.py b/backend/package/yuxi/channel/gateway/webhook_security.py new file mode 100644 index 00000000..c81a100f --- /dev/null +++ b/backend/package/yuxi/channel/gateway/webhook_security.py @@ -0,0 +1,304 @@ +import hashlib +import hmac +import logging +import time +from dataclasses import dataclass +from enum import StrEnum + +logger = logging.getLogger(__name__) + +PRE_AUTH_BODY_SIZE_LIMIT = 64 * 1024 +POST_AUTH_BODY_SIZE_LIMIT = 1 * 1024 * 1024 +PRE_AUTH_READ_TIMEOUT = 5.0 +POST_AUTH_READ_TIMEOUT = 30.0 +MAX_CONCURRENT_PER_KEY = 8 +MAX_TRACKING_KEYS = 4096 + +_VALID_WEBHOOK_METHODS = frozenset({"POST", "PUT", "PATCH"}) + + +class WebhookErrorCode(StrEnum): + METHOD_NOT_ALLOWED = "METHOD_NOT_ALLOWED" + CONTENT_TYPE_UNSUPPORTED = "CONTENT_TYPE_UNSUPPORTED" + BODY_TOO_LARGE = "BODY_TOO_LARGE" + TOO_MANY_REQUESTS = "TOO_MANY_REQUESTS" + SIGNATURE_MISSING = "SIGNATURE_MISSING" + SIGNATURE_INVALID = "SIGNATURE_INVALID" + PAYLOAD_INVALID = "PAYLOAD_INVALID" + + +@dataclass +class WebhookGuardConfig: + channel_type: str + allowed_methods: tuple[str, ...] = ("POST",) + allowed_content_types: tuple[str, ...] = ("application/json",) + pre_auth: bool = True + verify_signature: bool = True + body_size_limit: int = PRE_AUTH_BODY_SIZE_LIMIT + + +import threading + +class WebhookAnomalyTracker: + def __init__(self, sample_window: float = 60.0, alert_threshold: int = 50): + self._sample_window = sample_window + self._alert_threshold = alert_threshold + self._counters: dict[str, dict[str, int]] = {} + self._last_reset: float = time.monotonic() + self._lock = threading.Lock() + + def record(self, channel_type: str, error_code: str) -> None: + now = time.monotonic() + with self._lock: + if now - self._last_reset > self._sample_window: + self._counters.clear() + self._last_reset = now + + channel_counters = self._counters.setdefault(channel_type, {}) + channel_counters[error_code] = channel_counters.get(error_code, 0) + 1 + + total = sum(channel_counters.values()) + if total >= self._alert_threshold and total % self._alert_threshold == 0: + self._emit_alert(channel_type, channel_counters) + + def _emit_alert(self, channel_type: str, counters: dict[str, int]) -> None: + details = ", ".join(f"{k}={v}" for k, v in counters.items()) + logger.warning( + "Webhook anomaly alert [%s]: %d anomalies in window — %s", + channel_type, + sum(counters.values()), + details, + ) + + def get_stats(self, channel_type: str | None = None) -> dict[str, dict[str, int]]: + with self._lock: + if channel_type: + return {channel_type: self._counters.get(channel_type, {})} + return dict(self._counters) + + +class WebhookSigner: + def __init__(self): + self._signing_keys: dict[str, str] = {} + + def set_channel_key(self, channel_type: str, signing_key: str) -> None: + self._signing_keys[channel_type] = signing_key + + def remove_channel_key(self, channel_type: str) -> None: + self._signing_keys.pop(channel_type, None) + + def compute_signature(self, channel_type: str, body: bytes) -> str | None: + key = self._signing_keys.get(channel_type) + if not key: + return None + return hmac.new(key.encode(), body, hashlib.sha256).hexdigest() + + def verify( + self, + channel_type: str, + body: bytes, + signature: str, + ) -> bool: + expected = self.compute_signature(channel_type, body) + if expected is None: + return False + return hmac.compare_digest(expected, signature) + + @staticmethod + def compute_sha256_signature(secret: str, body: bytes) -> str: + return hmac.new(secret.encode(), body, hashlib.sha256).hexdigest() + + +class WebhookConcurrencyGuard: + def __init__(self, max_per_key: int = MAX_CONCURRENT_PER_KEY, max_keys: int = MAX_TRACKING_KEYS): + self._max_per_key = max_per_key + self._max_keys = max_keys + self._inflight: dict[str, int] = {} + self._lock = threading.Lock() + + def try_acquire(self, channel_type: str) -> bool: + with self._lock: + current = self._inflight.get(channel_type, 0) + if current >= self._max_per_key: + return False + if len(self._inflight) >= self._max_keys and channel_type not in self._inflight: + return False + self._inflight[channel_type] = current + 1 + return True + + def release(self, channel_type: str) -> None: + with self._lock: + current = self._inflight.get(channel_type, 0) + if current <= 1: + self._inflight.pop(channel_type, None) + else: + self._inflight[channel_type] = current - 1 + + @property + def inflight_count(self) -> dict[str, int]: + with self._lock: + return dict(self._inflight) + + +@dataclass +class WebhookGuardResult: + allowed: bool + error_code: WebhookErrorCode | None = None + error_message: str | None = None + http_status: int = 200 + + +class WebhookGuard: + def __init__( + self, + signer: WebhookSigner | None = None, + concurrency_guard: WebhookConcurrencyGuard | None = None, + anomaly_tracker: WebhookAnomalyTracker | None = None, + ): + self._signer = signer or WebhookSigner() + self._concurrency = concurrency_guard or WebhookConcurrencyGuard() + self._anomaly = anomaly_tracker or WebhookAnomalyTracker() + self._channel_configs: dict[str, WebhookGuardConfig] = {} + + def register_channel(self, channel_type: str, config: WebhookGuardConfig | None = None) -> None: + self._channel_configs[channel_type] = config or WebhookGuardConfig(channel_type=channel_type) + + def unregister_channel(self, channel_type: str) -> None: + self._channel_configs.pop(channel_type, None) + + def check_method(self, channel_type: str, method: str) -> WebhookGuardResult: + cfg = self._channel_configs.get(channel_type) + allowed = list(cfg.allowed_methods) if cfg else ["POST"] + if method not in _VALID_WEBHOOK_METHODS or method not in allowed: + self._anomaly.record(channel_type, "405") + return WebhookGuardResult( + allowed=False, + error_code=WebhookErrorCode.METHOD_NOT_ALLOWED, + error_message=f"Method {method} not allowed", + http_status=405, + ) + return WebhookGuardResult(allowed=True) + + def check_content_type(self, channel_type: str, content_type: str | None) -> WebhookGuardResult: + cfg = self._channel_configs.get(channel_type) + allowed = list(cfg.allowed_content_types) if cfg else ["application/json"] + if not content_type: + self._anomaly.record(channel_type, "415") + return WebhookGuardResult( + allowed=False, + error_code=WebhookErrorCode.CONTENT_TYPE_UNSUPPORTED, + error_message="Content-Type header missing", + http_status=415, + ) + ctype_main = content_type.split(";")[0].strip().lower() + if ctype_main not in allowed: + self._anomaly.record(channel_type, "415") + return WebhookGuardResult( + allowed=False, + error_code=WebhookErrorCode.CONTENT_TYPE_UNSUPPORTED, + error_message=f"Content-Type {content_type} not supported", + http_status=415, + ) + return WebhookGuardResult(allowed=True) + + def check_body_size(self, channel_type: str, body: bytes) -> WebhookGuardResult: + cfg = self._channel_configs.get(channel_type) + limit = cfg.body_size_limit if cfg else PRE_AUTH_BODY_SIZE_LIMIT + if len(body) > limit: + self._anomaly.record(channel_type, "413") + return WebhookGuardResult( + allowed=False, + error_code=WebhookErrorCode.BODY_TOO_LARGE, + error_message=f"Body size {len(body)} exceeds limit {limit}", + http_status=413, + ) + return WebhookGuardResult(allowed=True) + + def check_concurrency(self, channel_type: str) -> WebhookGuardResult: + if not self._concurrency.try_acquire(channel_type): + self._anomaly.record(channel_type, "429") + return WebhookGuardResult( + allowed=False, + error_code=WebhookErrorCode.TOO_MANY_REQUESTS, + error_message="Too many concurrent requests", + http_status=429, + ) + return WebhookGuardResult(allowed=True) + + def release_concurrency(self, channel_type: str) -> None: + self._concurrency.release(channel_type) + + def check_signature( + self, + channel_type: str, + body: bytes, + signature: str | None, + ) -> WebhookGuardResult: + cfg = self._channel_configs.get(channel_type) + if cfg is None or not cfg.verify_signature: + return WebhookGuardResult(allowed=True) + + if not signature: + self._anomaly.record(channel_type, "401") + return WebhookGuardResult( + allowed=False, + error_code=WebhookErrorCode.SIGNATURE_MISSING, + error_message="Webhook signature missing", + http_status=401, + ) + + if not self._signer.verify(channel_type, body, signature): + self._anomaly.record(channel_type, "401") + return WebhookGuardResult( + allowed=False, + error_code=WebhookErrorCode.SIGNATURE_INVALID, + error_message="Webhook signature invalid", + http_status=401, + ) + + return WebhookGuardResult(allowed=True) + + def run_pipeline( + self, + channel_type: str, + method: str, + content_type: str | None, + body: bytes, + signature: str | None = None, + ) -> WebhookGuardResult: + result = self.check_method(channel_type, method) + if not result.allowed: + return result + + result = self.check_content_type(channel_type, content_type) + if not result.allowed: + return result + + result = self.check_body_size(channel_type, body) + if not result.allowed: + return result + + result = self.check_signature(channel_type, body, signature) + if not result.allowed: + return result + + result = self.check_concurrency(channel_type) + if not result.allowed: + return result + + return WebhookGuardResult(allowed=True) + + @property + def signer(self) -> WebhookSigner: + return self._signer + + @property + def anomaly_tracker(self) -> WebhookAnomalyTracker: + return self._anomaly + + @property + def concurrency_guard(self) -> WebhookConcurrencyGuard: + return self._concurrency + + +webhook_guard = WebhookGuard()