mirror of
https://github.com/hansjone/oclaw.git
synced 2026-10-09 01:50:44 +08:00
完成 WebSocket 与 workspace 架构优化收尾,并修复全量回归阻塞。
本次将握手鉴权、Origin 校验、限流、重连补偿、发送背压与观测字段打通,同时修复 gateway 异常路径与备份清理兼容问题,确保工具历史压缩语义一致并恢复全量测试通过。 Made-with: Cursor
This commit is contained in:
parent
09836dddfa
commit
7253f6795d
12 changed files with 536 additions and 24 deletions
|
|
@ -2,7 +2,10 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
from collections import defaultdict, deque
|
||||
from typing import Any
|
||||
import threading
|
||||
|
||||
|
|
@ -16,9 +19,17 @@ from oclaw.interfaces.ws.common import (
|
|||
PREAUTH_HANDSHAKE_TIMEOUT_MS,
|
||||
PROTOCOL_VERSION,
|
||||
TICK_INTERVAL_MS,
|
||||
WS_EVENT_REPLAY_MAX,
|
||||
WS_RATE_LIMIT_CONN_PER_WINDOW,
|
||||
WS_RATE_LIMIT_IP_PER_WINDOW,
|
||||
WS_RATE_LIMIT_USER_PER_WINDOW,
|
||||
WS_RATE_LIMIT_WINDOW_MS,
|
||||
WS_SEND_QUEUE_MAX_BYTES,
|
||||
WS_SEND_QUEUE_MAX_MESSAGES,
|
||||
error_shape as _error_shape,
|
||||
normalize_ws_attachments as _normalize_ws_attachments,
|
||||
now_ms as _now_ms,
|
||||
origin_is_allowed,
|
||||
)
|
||||
from oclaw.interfaces.ws.events import (
|
||||
emit_agent_event as emit_agent_event_impl,
|
||||
|
|
@ -34,8 +45,18 @@ from oclaw.interfaces.ws.turn_runner import run_agent_turn_via_bridge
|
|||
from oclaw.interfaces.ws.ws_schema import format_validation_errors, get_ws_schemas, validate_or_errors
|
||||
from oclaw.runtime.relay_pointer import validate_relay_share_envelope
|
||||
|
||||
_LOG = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OclawWsGatewayConnection:
|
||||
_rate_lock = threading.Lock()
|
||||
_rate_by_ip: dict[str, deque[int]] = defaultdict(deque)
|
||||
_rate_by_user: dict[str, deque[int]] = defaultdict(deque)
|
||||
_stats: dict[str, int] = defaultdict(int)
|
||||
_event_buffer_by_user: dict[str, deque[dict[str, Any]]] = defaultdict(
|
||||
lambda: deque(maxlen=max(1, int(WS_EVENT_REPLAY_MAX)))
|
||||
)
|
||||
|
||||
def __init__(self, ws: WebSocket):
|
||||
self.ws = ws
|
||||
self.schemas = get_ws_schemas()
|
||||
|
|
@ -55,12 +76,30 @@ class OclawWsGatewayConnection:
|
|||
self._abort_lock = threading.Lock()
|
||||
self._aborted_run_ids: set[str] = set()
|
||||
self._active_run_session: dict[str, str] = {}
|
||||
self._send_queue: asyncio.Queue[tuple[str, int]] = asyncio.Queue(maxsize=WS_SEND_QUEUE_MAX_MESSAGES)
|
||||
self._send_pending_bytes = 0
|
||||
self._send_pending_lock = asyncio.Lock()
|
||||
self._sender_task: asyncio.Task[None] | None = None
|
||||
self._event_buffer: deque[dict[str, Any]] = deque(maxlen=max(1, int(WS_EVENT_REPLAY_MAX)))
|
||||
self._rate_local: deque[int] = deque()
|
||||
self._gateway_handlers = build_gateway_method_handlers()
|
||||
self._now_ms = _now_ms
|
||||
self._error_shape = _error_shape
|
||||
|
||||
async def run(self) -> None:
|
||||
await run_connection_loop(self)
|
||||
self._inc_stat("ws_connections_opened")
|
||||
self._sender_task = asyncio.create_task(self._sender_loop())
|
||||
try:
|
||||
await run_connection_loop(self)
|
||||
finally:
|
||||
await self._drain_sender()
|
||||
self._inc_stat("ws_connections_closed")
|
||||
_LOG.info(
|
||||
"ws connection closed conn_id=%s total_opened=%s total_closed=%s",
|
||||
self.conn_id,
|
||||
self._stats.get("ws_connections_opened", 0),
|
||||
self._stats.get("ws_connections_closed", 0),
|
||||
)
|
||||
|
||||
async def _recv_frame(self, *, preauth: bool = False) -> dict[str, Any] | None:
|
||||
return await recv_frame_impl(
|
||||
|
|
@ -88,8 +127,31 @@ class OclawWsGatewayConnection:
|
|||
await close_ws_impl(self, code=code, reason=reason)
|
||||
|
||||
async def _dispatch_connected(self, req_id: str, method: str, params: Any) -> None:
|
||||
limited, bucket = self._rate_limited()
|
||||
if limited:
|
||||
self._inc_stat("ws_rate_limited")
|
||||
await self.send_res(
|
||||
req_id,
|
||||
ok=False,
|
||||
error=_error_shape("RATE_LIMITED", "too many requests", details={"bucket": bucket}),
|
||||
)
|
||||
return
|
||||
await dispatch_connected_impl(self, req_id=req_id, method=method, params=params)
|
||||
|
||||
@classmethod
|
||||
def _inc_stat(cls, key: str, delta: int = 1) -> None:
|
||||
with cls._rate_lock:
|
||||
cls._stats[str(key)] = int(cls._stats.get(str(key), 0)) + int(delta)
|
||||
|
||||
def mark_handshake(self, *, ok: bool) -> None:
|
||||
self._inc_stat("ws_handshake_ok" if ok else "ws_handshake_failed")
|
||||
_LOG.info(
|
||||
"ws handshake conn_id=%s ok=%s user_id=%s",
|
||||
self.conn_id,
|
||||
int(bool(ok)),
|
||||
str((self.auth_ctx or {}).get("user_id") or ""),
|
||||
)
|
||||
|
||||
async def _dispatch_via_server_methods(self, *, req_id: str, method: str, params: Any) -> bool:
|
||||
return await dispatch_via_server_methods(
|
||||
req_id=req_id,
|
||||
|
|
@ -117,6 +179,122 @@ class OclawWsGatewayConnection:
|
|||
now_ms=_now_ms,
|
||||
)
|
||||
|
||||
def validate_origin(self) -> bool:
|
||||
headers = getattr(self.ws, "headers", None)
|
||||
origin = str(headers.get("origin") or "").strip() if headers is not None else ""
|
||||
host = str(headers.get("host") or "").strip() if headers is not None else ""
|
||||
allowed = origin_is_allowed(origin, host)
|
||||
if not allowed:
|
||||
_LOG.warning("ws origin blocked conn_id=%s origin=%s host=%s", self.conn_id, origin, host)
|
||||
return allowed
|
||||
|
||||
def _prune_window(self, dq: deque[int], now: int) -> None:
|
||||
cutoff = now - int(WS_RATE_LIMIT_WINDOW_MS)
|
||||
while dq and dq[0] < cutoff:
|
||||
dq.popleft()
|
||||
|
||||
def _rate_limited(self) -> tuple[bool, str]:
|
||||
now = _now_ms()
|
||||
self._prune_window(self._rate_local, now)
|
||||
if len(self._rate_local) >= int(WS_RATE_LIMIT_CONN_PER_WINDOW):
|
||||
return True, "connection"
|
||||
self._rate_local.append(now)
|
||||
|
||||
headers = getattr(self.ws, "headers", None)
|
||||
ip = str(headers.get("x-forwarded-for") or headers.get("x-real-ip") or "").split(",")[0].strip() if headers is not None else ""
|
||||
user_id = str((self.auth_ctx or {}).get("user_id") or "").strip()
|
||||
with self._rate_lock:
|
||||
if ip:
|
||||
ip_bucket = self._rate_by_ip[ip]
|
||||
self._prune_window(ip_bucket, now)
|
||||
if len(ip_bucket) >= int(WS_RATE_LIMIT_IP_PER_WINDOW):
|
||||
return True, "ip"
|
||||
ip_bucket.append(now)
|
||||
if user_id:
|
||||
user_bucket = self._rate_by_user[user_id]
|
||||
self._prune_window(user_bucket, now)
|
||||
if len(user_bucket) >= int(WS_RATE_LIMIT_USER_PER_WINDOW):
|
||||
return True, "user"
|
||||
user_bucket.append(now)
|
||||
return False, ""
|
||||
|
||||
async def _queue_send_text(self, text: str) -> bool:
|
||||
payload = str(text or "")
|
||||
payload_size = len(payload.encode("utf-8", errors="ignore"))
|
||||
async with self._send_pending_lock:
|
||||
if self._send_pending_bytes + payload_size > int(WS_SEND_QUEUE_MAX_BYTES):
|
||||
_LOG.warning("ws send queue bytes exceeded conn_id=%s", self.conn_id)
|
||||
return False
|
||||
self._send_pending_bytes += payload_size
|
||||
try:
|
||||
self._send_queue.put_nowait((payload, payload_size))
|
||||
return True
|
||||
except asyncio.QueueFull:
|
||||
async with self._send_pending_lock:
|
||||
self._send_pending_bytes = max(0, self._send_pending_bytes - payload_size)
|
||||
_LOG.warning("ws send queue full conn_id=%s", self.conn_id)
|
||||
return False
|
||||
|
||||
async def _sender_loop(self) -> None:
|
||||
while True:
|
||||
item = await self._send_queue.get()
|
||||
if item[0] == "__STOP__":
|
||||
self._send_queue.task_done()
|
||||
return
|
||||
payload, payload_size = item
|
||||
try:
|
||||
await self.ws.send_text(payload)
|
||||
except Exception:
|
||||
return
|
||||
finally:
|
||||
async with self._send_pending_lock:
|
||||
self._send_pending_bytes = max(0, self._send_pending_bytes - payload_size)
|
||||
self._send_queue.task_done()
|
||||
|
||||
async def _drain_sender(self) -> None:
|
||||
if self._sender_task is None:
|
||||
return
|
||||
try:
|
||||
self._send_queue.put_nowait(("__STOP__", 0))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
await self._sender_task
|
||||
except Exception:
|
||||
pass
|
||||
self._sender_task = None
|
||||
|
||||
def remember_event(self, frame: dict[str, Any]) -> None:
|
||||
snap = dict(frame or {})
|
||||
self._event_buffer.append(snap)
|
||||
user_id = str((self.auth_ctx or {}).get("user_id") or "").strip()
|
||||
if not user_id:
|
||||
return
|
||||
with self._rate_lock:
|
||||
bucket = self._event_buffer_by_user.get(user_id)
|
||||
if bucket is None or bucket.maxlen != max(1, int(WS_EVENT_REPLAY_MAX)):
|
||||
bucket = deque(maxlen=max(1, int(WS_EVENT_REPLAY_MAX)))
|
||||
self._event_buffer_by_user[user_id] = bucket
|
||||
bucket.append(dict(snap))
|
||||
|
||||
async def replay_events_since(self, seq: int) -> None:
|
||||
after = int(seq or 0)
|
||||
frames = list(self._event_buffer)
|
||||
user_id = str((self.auth_ctx or {}).get("user_id") or "").strip()
|
||||
if user_id:
|
||||
with self._rate_lock:
|
||||
shared = list(self._event_buffer_by_user.get(user_id) or [])
|
||||
if shared:
|
||||
frames = shared
|
||||
for frame in frames:
|
||||
fseq = int(frame.get("seq") or 0)
|
||||
if fseq <= after:
|
||||
continue
|
||||
text = frame.get("_raw")
|
||||
if not isinstance(text, str) or not text:
|
||||
continue
|
||||
await self._queue_send_text(text)
|
||||
|
||||
def build_hello_ok(self, _connect_params: dict[str, Any] | None) -> dict[str, Any]:
|
||||
methods = ["connect", *method_names()]
|
||||
return build_hello_ok_payload(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue