mirror of
https://github.com/hansjone/oclaw.git
synced 2026-10-09 05:50:44 +08:00
oliver
7253f6795d
本次将握手鉴权、Origin 校验、限流、重连补偿、发送背压与观测字段打通，同时修复 gateway 异常路径与备份清理兼容问题，确保工具历史压缩语义一致并恢复全量测试通过。 Made-with: Cursor
102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from starlette.websockets import WebSocketState
|
|
|
|
|
|
def _ws_is_disconnected(conn: Any) -> bool:
|
|
ws = getattr(conn, "ws", None)
|
|
if ws is None:
|
|
return True
|
|
state = getattr(ws, "application_state", None)
|
|
if state == WebSocketState.DISCONNECTED:
|
|
return True
|
|
state = getattr(ws, "client_state", None)
|
|
if state == WebSocketState.DISCONNECTED:
|
|
return True
|
|
return False
|
|
|
|
|
|
async def _safe_send_text(conn: Any, text: str, *, use_queue: bool = True) -> None:
|
|
if _ws_is_disconnected(conn):
|
|
return
|
|
queue_send = getattr(conn, "_queue_send_text", None)
|
|
if use_queue and callable(queue_send):
|
|
ok = await queue_send(text)
|
|
if ok:
|
|
return
|
|
try:
|
|
await conn.ws.close(code=1013, reason="backpressure")
|
|
except Exception:
|
|
pass
|
|
return
|
|
try:
|
|
await conn.ws.send_text(text)
|
|
except Exception as e:
|
|
msg = str(e or "")
|
|
if "Unexpected ASGI message 'websocket.send'" in msg or "response already completed" in msg:
|
|
return
|
|
raise
|
|
|
|
|
|
async def send_res(conn: Any, req_id: str, *, ok: bool, payload: Any | None = None, error: Any | None = None) -> None:
|
|
frame: dict[str, Any] = {"type": "res", "id": str(req_id or "invalid"), "ok": bool(ok)}
|
|
if payload is not None:
|
|
frame["payload"] = payload
|
|
if error is not None:
|
|
frame["error"] = error
|
|
await _safe_send_text(conn, json.dumps(frame, ensure_ascii=False), use_queue=False)
|
|
|
|
|
|
async def send_event(conn: Any, event: str, payload: Any | None = None) -> None:
|
|
if _ws_is_disconnected(conn):
|
|
return
|
|
conn.seq += 1
|
|
frame: dict[str, Any] = {"type": "event", "event": str(event or "event"), "seq": int(conn.seq)}
|
|
if payload is not None:
|
|
frame["payload"] = payload
|
|
raw = json.dumps(frame, ensure_ascii=False)
|
|
frame["_raw"] = raw
|
|
remember = getattr(conn, "remember_event", None)
|
|
if callable(remember):
|
|
remember(frame)
|
|
await _safe_send_text(conn, raw)
|
|
|
|
|
|
async def emit_agent_event(conn: Any, *, run_id: str, stream: str, data: dict[str, Any], now_ms: int) -> None:
|
|
payload = {"runId": str(run_id), "seq": int(conn.seq + 1), "stream": str(stream), "ts": int(now_ms), "data": dict(data or {})}
|
|
await send_event(conn, "agent.event", payload)
|
|
|
|
|
|
async def emit_chat_event(
|
|
conn: Any,
|
|
*,
|
|
run_id: str,
|
|
state: str,
|
|
delta: str = "",
|
|
reply: str = "",
|
|
error: str = "",
|
|
message: dict[str, Any] | None = None,
|
|
session_key: str | None = None,
|
|
seq: int | None = None,
|
|
) -> None:
|
|
payload: dict[str, Any] = {"runId": str(run_id), "state": str(state or "")}
|
|
if session_key:
|
|
payload["sessionKey"] = str(session_key)
|
|
if seq is not None:
|
|
payload["seq"] = int(seq)
|
|
if delta:
|
|
payload["message"] = {"role": "assistant", "content": [{"type": "text", "text": str(delta)}]}
|
|
elif reply and message is None:
|
|
payload["message"] = {"role": "assistant", "content": [{"type": "text", "text": str(reply)}]}
|
|
if error:
|
|
payload["errorMessage"] = str(error)
|
|
if message is not None:
|
|
payload["message"] = message
|
|
await send_event(conn, "chat", payload)
|
|
|
|
|
|
__all__ = ["send_res", "send_event", "emit_agent_event", "emit_chat_event"]
|
|
|