feat(channel/gateway): 新增完整网关通道模块

新增设备身份管理、认证限流、并发通道、Webhook路由、RBAC权限控制、SSE/轮询降级等全套网关通道功能,包含:
1. 设备身份生成与签名验证
2. 设备令牌认证与速率限制
3. 内存+数据库双重设备注册表
4. 并发通道限流管理
5. Webhook安全处理与路由
6. RBAC权限校验系统
7. OpenAI API兼容适配层
8. Tailscale认证支持
9. HTTP轮询降级机制
This commit is contained in:
Kris 2026-05-21 10:26:33 +08:00
parent c9d0a26f52
commit ecd3c90e80
25 changed files with 6490 additions and 0 deletions

View File

@ -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",
]

View File

@ -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

View File

@ -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()

View File

@ -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()

View File

@ -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)

View File

@ -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")

View File

@ -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],
)

View File

@ -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]

View File

@ -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())

View File

@ -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)

View File

@ -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()

View File

@ -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

View File

@ -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)

View File

@ -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

View File

@ -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)

View File

@ -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}")

View File

@ -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

View File

@ -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,
}

View File

@ -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()

File diff suppressed because it is too large Load Diff

View File

@ -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

View File

@ -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",
},
)

View File

@ -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)

View File

@ -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

View File

@ -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()