feat(gateway): outbound agent with task dispatch and event bridge
- OutboundAgent dials /v1/agent with bearer access key, blocking recv - heartbeat thread reports version; backoff reconnect on failure - task frames dispatch through shared route_gateway_request with duplicate-id idempotent rejection (409) and explicit error mapping - event bridge pushes SubscriptionManager pending deliveries, honors event_ack/event_nack, repushes unconfirmed batches after reconnect
This commit is contained in:
@@ -0,0 +1,377 @@
|
||||
"""Gateway outbound agent: dial the platform, dispatch tasks, push events.
|
||||
|
||||
方向反转后网关只有一条自己发起的 WS 连接（/v1/agent，连接层 Bearer key 认证）：
|
||||
- 心跳线程周期上报 {type: heartbeat, version}，平台据此维护 online/last_seen_at；
|
||||
- 读循环接收 task（线程池执行，复用 HTTP 路由语义）与 event_ack/event_nack；
|
||||
- result 帧回传路由结果；重复 task id 幂等拒绝，不重复执行有副作用的操作；
|
||||
- 事件桥把 SubscriptionManager 的 pending deliveries 推为 event 帧，
|
||||
平台事务保存成功后回 event_ack；未确认批次在 nack 或断线重连后重推。
|
||||
协议与平台侧 internal/controlplane/api/gateway_channel.go 一一对应。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable
|
||||
|
||||
import websocket
|
||||
|
||||
LOG = logging.getLogger("creatorhub.gateway")
|
||||
|
||||
GW_VERSION = "1.0.0-outbound"
|
||||
HEARTBEAT_INTERVAL = 20.0
|
||||
RECONNECT_BASE = 1.0
|
||||
RECONNECT_MAX = 30.0
|
||||
WORKER_COUNT = 4
|
||||
IDEMPOTENT_STATUS = 409
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
platform_url: str
|
||||
access_key: str
|
||||
version: str = GW_VERSION
|
||||
heartbeat_interval: float = HEARTBEAT_INTERVAL
|
||||
reconnect_base: float = RECONNECT_BASE
|
||||
reconnect_max: float = RECONNECT_MAX
|
||||
worker_count: int = WORKER_COUNT
|
||||
|
||||
|
||||
class OutboundAgent:
|
||||
"""单网关出站连接的完整生命周期；run() 阻塞运行，stop() 置位退出。"""
|
||||
|
||||
def __init__(self, gateway: Any, config: AgentConfig) -> None:
|
||||
self.gateway = gateway
|
||||
self.config = config
|
||||
self.connected = threading.Event()
|
||||
self.last_error = ""
|
||||
self.reconnect_attempts = 0
|
||||
self._stop_requested = threading.Event()
|
||||
self._write_lock = threading.Lock()
|
||||
self._conn: websocket.WebSocket | None = None
|
||||
self._sent_events: dict[str, dict[str, float]] = {}
|
||||
self._sent_lock = threading.Lock()
|
||||
self._seen_tasks: set[int] = set()
|
||||
self._seen_tasks_lock = threading.Lock()
|
||||
self._push_wakeup = threading.Event()
|
||||
|
||||
# ---- 生命周期 -------------------------------------------------------
|
||||
|
||||
def run(self) -> None:
|
||||
backoff = self.config.reconnect_base
|
||||
while not self._stop_requested.is_set():
|
||||
try:
|
||||
self.reconnect_attempts += 1
|
||||
conn = self._connect()
|
||||
except Exception as exc:
|
||||
self.last_error = str(exc) or exc.__class__.__name__
|
||||
LOG.warning("platform connection failed: %s", self.last_error)
|
||||
if self._wait(backoff):
|
||||
break
|
||||
backoff = min(backoff * 2, self.config.reconnect_max)
|
||||
continue
|
||||
self.reconnect_attempts = 0
|
||||
self.last_error = ""
|
||||
self._serve(conn)
|
||||
backoff = self.config.reconnect_base
|
||||
if not self._stop_requested.is_set():
|
||||
LOG.info("platform connection lost; reconnecting")
|
||||
|
||||
def stop(self) -> None:
|
||||
self._stop_requested.set()
|
||||
with self._write_lock:
|
||||
conn = self._conn
|
||||
if conn is not None:
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _wait(self, seconds: float) -> bool:
|
||||
return self._stop_requested.wait(seconds)
|
||||
|
||||
def _connect(self) -> websocket.WebSocket:
|
||||
url = self.config.platform_url
|
||||
if url.startswith("https://"):
|
||||
url = "wss://" + url[len("https://"):]
|
||||
elif url.startswith("http://"):
|
||||
url = "ws://" + url[len("http://"):]
|
||||
conn = websocket.create_connection(
|
||||
url,
|
||||
header=[f"Authorization: Bearer {self.config.access_key}"],
|
||||
timeout=10,
|
||||
)
|
||||
# timeout 仅用于握手；读循环必须阻塞，否则空闲连接会在 recv 超时后假性断线。
|
||||
conn.settimeout(None)
|
||||
with self._write_lock:
|
||||
self._conn = conn
|
||||
return conn
|
||||
|
||||
# ---- 会话线程 -------------------------------------------------------
|
||||
|
||||
def _serve(self, conn: websocket.WebSocket) -> None:
|
||||
self.connected.set()
|
||||
with self._sent_lock:
|
||||
self._sent_events.clear()
|
||||
executor = ThreadPoolExecutor(max_workers=self.config.worker_count, thread_name_prefix="gw-task")
|
||||
heartbeat = threading.Thread(target=self._heartbeat_loop, args=(conn,), daemon=True)
|
||||
sender = threading.Thread(target=self._event_sender_loop, args=(conn,), daemon=True)
|
||||
heartbeat.start()
|
||||
sender.start()
|
||||
failure = ""
|
||||
try:
|
||||
while not self._stop_requested.is_set():
|
||||
raw = conn.recv()
|
||||
if not isinstance(raw, str):
|
||||
continue
|
||||
try:
|
||||
frame = json.loads(raw)
|
||||
except ValueError:
|
||||
failure = "gateway channel message is not JSON"
|
||||
break
|
||||
if not isinstance(frame, dict):
|
||||
failure = "gateway channel message invalid"
|
||||
break
|
||||
kind = frame.get("type")
|
||||
if kind == "task":
|
||||
future = executor.submit(self._handle_task, conn, frame)
|
||||
future.add_done_callback(self._log_task_failure)
|
||||
elif kind == "event_ack":
|
||||
self._handle_event_ack(frame)
|
||||
elif kind == "event_nack":
|
||||
self._handle_event_nack(frame)
|
||||
else:
|
||||
failure = f"gateway channel message type invalid: {kind!r}"
|
||||
break
|
||||
except Exception as exc:
|
||||
if not self._stop_requested.is_set():
|
||||
failure = str(exc) or exc.__class__.__name__
|
||||
LOG.exception("gateway channel receive loop failed: %s", failure)
|
||||
finally:
|
||||
self.connected.clear()
|
||||
executor.shutdown(wait=False, cancel_futures=True)
|
||||
self._stop_threads((heartbeat, sender))
|
||||
with self._write_lock:
|
||||
if self._conn is conn:
|
||||
self._conn = None
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
if failure:
|
||||
self.last_error = failure
|
||||
|
||||
@staticmethod
|
||||
def _stop_threads(threads: tuple[threading.Thread, ...]) -> None:
|
||||
for thread in threads:
|
||||
thread.join(timeout=2)
|
||||
|
||||
def _heartbeat_loop(self, conn: websocket.WebSocket) -> None:
|
||||
while not self._stop_requested.wait(self.config.heartbeat_interval):
|
||||
try:
|
||||
self._send(conn, {"type": "heartbeat", "version": self.config.version})
|
||||
except Exception:
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
# ---- 帧读写 ---------------------------------------------------------
|
||||
|
||||
def _send(self, conn: websocket.WebSocket, value: dict) -> None:
|
||||
payload = json.dumps(value, ensure_ascii=False, separators=(",", ":"))
|
||||
with self._write_lock:
|
||||
conn.send(payload)
|
||||
|
||||
# ---- 任务分发 -------------------------------------------------------
|
||||
|
||||
def _handle_task(self, conn: websocket.WebSocket, frame: dict) -> None:
|
||||
task_id = frame.get("id")
|
||||
method = frame.get("method")
|
||||
path = frame.get("path")
|
||||
if not isinstance(task_id, int) or not isinstance(method, str) or not isinstance(path, str):
|
||||
self.last_error = "gateway task frame invalid"
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
with self._seen_tasks_lock:
|
||||
duplicate = task_id in self._seen_tasks
|
||||
self._seen_tasks.add(task_id)
|
||||
if duplicate:
|
||||
self._send(conn, {
|
||||
"type": "result",
|
||||
"id": task_id,
|
||||
"status": IDEMPOTENT_STATUS,
|
||||
"body": "",
|
||||
"error": "duplicate task id; previous execution retained",
|
||||
})
|
||||
return
|
||||
status, body, error = self._execute_task(method, path, frame.get("payload"))
|
||||
self._send(conn, {
|
||||
"type": "result",
|
||||
"id": task_id,
|
||||
"status": status,
|
||||
"body": base64.b64encode(body).decode("ascii"),
|
||||
"error": error,
|
||||
})
|
||||
|
||||
def _execute_task(self, method: str, path: str, payload: Any) -> tuple[int, bytes, str]:
|
||||
from .server.http import RequestError, route_gateway_request
|
||||
|
||||
body: dict = {}
|
||||
if payload:
|
||||
try:
|
||||
body = json.loads(payload)
|
||||
except ValueError as exc:
|
||||
return 400, b"", f"task payload invalid: {exc}"
|
||||
if not isinstance(body, dict):
|
||||
return 400, b"", "task payload must be a JSON object"
|
||||
try:
|
||||
result = route_gateway_request(self.gateway, method, path, {}, body)
|
||||
except RequestError as exc:
|
||||
return exc.status, b"", str(exc)
|
||||
except FileNotFoundError as exc:
|
||||
return 404, b"", str(exc)
|
||||
except ValueError as exc:
|
||||
return 400, b"", str(exc)
|
||||
except Exception as exc: # 路由内部故障必须回传给平台，不能吞掉
|
||||
LOG.exception("gateway task %s %s failed", method, path)
|
||||
return 500, b"", "gateway operation failed"
|
||||
if result is None:
|
||||
return 204, b"", ""
|
||||
if isinstance(result, tuple):
|
||||
status, value = result
|
||||
return int(status), json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8"), ""
|
||||
return 200, json.dumps(result, ensure_ascii=False, separators=(",", ":")).encode("utf-8"), ""
|
||||
|
||||
@staticmethod
|
||||
def _log_task_failure(future: Any) -> None:
|
||||
exc = future.exception()
|
||||
if exc is not None:
|
||||
LOG.exception("gateway task dispatch failed", exc_info=exc)
|
||||
|
||||
# ---- 事件桥 ---------------------------------------------------------
|
||||
|
||||
def _event_sender_loop(self, conn: websocket.WebSocket) -> None:
|
||||
manager = self.gateway.subscriptions
|
||||
while not self._stop_requested.is_set():
|
||||
try:
|
||||
with manager.changed:
|
||||
version = manager.version
|
||||
self._push_wakeup.clear()
|
||||
self._push_pending(conn, version)
|
||||
with manager.changed:
|
||||
manager.changed.wait_for(
|
||||
lambda: self._stop_requested.is_set()
|
||||
or manager.version != version
|
||||
or self._push_wakeup.is_set(),
|
||||
timeout=5,
|
||||
)
|
||||
except Exception:
|
||||
if self._stop_requested.is_set():
|
||||
return
|
||||
LOG.exception("event sender loop failed; closing connection")
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
def _subscription_items(manager: Any) -> list[tuple[str, Any]]:
|
||||
# 从注册表快照枚举活动订阅；_items 缺失（测试替身）时视为空。
|
||||
items = getattr(manager, "_items", {})
|
||||
if not isinstance(items, dict):
|
||||
return []
|
||||
lock = getattr(manager, "_lock", None)
|
||||
if not isinstance(lock, type(threading.RLock())):
|
||||
return list(items.items())
|
||||
with lock:
|
||||
return list(items.items())
|
||||
|
||||
def _push_pending(self, conn: websocket.WebSocket, version: int) -> None:
|
||||
# 每轮全量对账 pending 与已推记录的差集：新事件推送一次；nack/重连清空
|
||||
# 已推记录后，同一批次会自然重推；无新差集时为空操作。
|
||||
from .platform.douyin import DouyinError
|
||||
|
||||
manager = self.gateway.subscriptions
|
||||
with self._sent_lock:
|
||||
sent = self._sent_events
|
||||
for alias, item in self._subscription_items(manager):
|
||||
try:
|
||||
current = manager._get(alias)
|
||||
error = item.error or (
|
||||
"event subscription replaced or stopped"
|
||||
if current is not item or item.stopped
|
||||
else ""
|
||||
)
|
||||
except (DouyinError, Exception) as exc:
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
if error:
|
||||
self._send(conn, {"type": "error", "subscription": "", "alias": alias, "error": error})
|
||||
sent.pop(alias, None)
|
||||
continue
|
||||
try:
|
||||
already = sent.setdefault(alias, {})
|
||||
pending = [d for d in item.pending() if d["delivery_id"] not in already]
|
||||
except Exception as exc:
|
||||
self._send(conn, {"type": "error", "subscription": "", "alias": alias, "error": str(exc)})
|
||||
sent.pop(alias, None)
|
||||
continue
|
||||
if not pending:
|
||||
continue
|
||||
subscription = f"{alias}:{item.session_id}"
|
||||
self._send(conn, {
|
||||
"type": "event",
|
||||
"event": {"subscription": subscription, "alias": alias, "deliveries": pending},
|
||||
})
|
||||
stamp = time.monotonic()
|
||||
already.update({d["delivery_id"]: stamp for d in pending})
|
||||
LOG.info("event pushed alias=%s subscription=%s count=%s", alias, subscription, len(pending))
|
||||
|
||||
def _handle_event_ack(self, frame: dict) -> None:
|
||||
alias = frame.get("alias")
|
||||
ids = frame.get("delivery_ids")
|
||||
if not isinstance(alias, str) or not isinstance(ids, list):
|
||||
self.last_error = "event acknowledgement invalid"
|
||||
return
|
||||
manager = self.gateway.subscriptions
|
||||
with self._sent_lock:
|
||||
already = self._sent_events.get(alias, {})
|
||||
unknown = [i for i in ids if i not in already]
|
||||
if unknown:
|
||||
# 平台确认了网关没有记录的投递（重连后旧 ack 等），仅清理已知项。
|
||||
LOG.warning("event ack references unknown deliveries alias=%s count=%s", alias, len(unknown))
|
||||
known = [i for i in ids if i in already]
|
||||
if not known:
|
||||
return
|
||||
manager.ack(alias, known)
|
||||
with self._sent_lock:
|
||||
already = self._sent_events.get(alias, {})
|
||||
for delivery_id in known:
|
||||
already.pop(delivery_id, None)
|
||||
|
||||
def _handle_event_nack(self, frame: dict) -> None:
|
||||
alias = frame.get("alias")
|
||||
if not isinstance(alias, str):
|
||||
return
|
||||
with self._sent_lock:
|
||||
self._sent_events.pop(alias, None)
|
||||
self._wake_sender()
|
||||
|
||||
def _wake_sender(self) -> None:
|
||||
# Condition.wait_for 只在 notify/超时时求值谓词；清空已推记录后必须唤醒。
|
||||
self._push_wakeup.set()
|
||||
manager = self.gateway.subscriptions
|
||||
changed = getattr(manager, "changed", None)
|
||||
if isinstance(changed, threading.Condition):
|
||||
with changed:
|
||||
changed.notify_all()
|
||||
@@ -585,6 +585,62 @@ class GatewayHTTPServer(ThreadingHTTPServer):
|
||||
connection.close()
|
||||
|
||||
|
||||
def route_gateway_request(gateway: "Gateway", method: str, path: str, query: dict, body: dict):
|
||||
"""统一任务路由：HTTP handler 与出站通道 task 分发共用同一套 method/path 语义。"""
|
||||
if method == "GET" and path == "/v1/info":
|
||||
return gateway.info()
|
||||
if method == "GET" and path == "/v1/browsers":
|
||||
return gateway.list_browsers()
|
||||
if method == "POST" and path == "/v1/browsers":
|
||||
return 201, gateway.create(body)
|
||||
match = re.fullmatch(r"/v1/browsers/([a-z0-9][a-z0-9-]{0,31})", path)
|
||||
if match and method == "DELETE":
|
||||
gateway.remove(match.group(1), body)
|
||||
return None
|
||||
match = re.fullmatch(
|
||||
r"/v1/browsers/([a-z0-9][a-z0-9-]{0,31})/(start|stop|cancel|proxy)", path
|
||||
)
|
||||
if match:
|
||||
alias, action = match.groups()
|
||||
if method == "POST" and action in {"start", "stop", "cancel"}:
|
||||
gateway.change_state(alias, action, body)
|
||||
return None
|
||||
if method == "POST" and action == "proxy":
|
||||
gateway.restore_proxy(alias, body)
|
||||
return None
|
||||
match = re.fullmatch(
|
||||
r"/v1/browsers/([a-z0-9][a-z0-9-]{0,31})/douyin/(get|media|image|identity|resolve|login-qr|action|messages|inbox|events)",
|
||||
path,
|
||||
)
|
||||
if match:
|
||||
alias, action = match.groups()
|
||||
if action == "get" and method == "POST":
|
||||
return gateway.get_douyin(alias, body)
|
||||
if action == "resolve" and method == "POST":
|
||||
return gateway.resolve_douyin(alias, body)
|
||||
if action == "media" and method == "POST":
|
||||
return gateway.get_douyin_media(alias, body)
|
||||
if action == "image" and method == "POST":
|
||||
return gateway.get_douyin_image(alias, body)
|
||||
if action == "identity" and method == "POST":
|
||||
return gateway.douyin_identity(alias, body)
|
||||
if action == "login-qr" and method == "POST":
|
||||
return gateway.douyin_login_qr(alias, body)
|
||||
if action == "action" and method == "POST":
|
||||
return gateway.douyin_action(alias, body)
|
||||
if action == "inbox" and method == "POST":
|
||||
return gateway.douyin_inbox(alias, body)
|
||||
if action == "messages" and method == "POST":
|
||||
return gateway.douyin_message_history(alias, body)
|
||||
if action == "events":
|
||||
if method == "POST":
|
||||
return gateway.start_douyin_events(alias, body)
|
||||
if method == "DELETE":
|
||||
gateway.stop_douyin_events(alias, body)
|
||||
return None
|
||||
raise RequestError("not found", 404)
|
||||
|
||||
|
||||
class GatewayHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
@@ -659,59 +715,7 @@ class GatewayHandler(BaseHTTPRequestHandler):
|
||||
self._respond(500, json_bytes({"error": "gateway operation failed"}))
|
||||
|
||||
def _route(self, method: str, path: str, query: dict, body: dict):
|
||||
gateway = self.server_as_gateway().gateway
|
||||
if method == "GET" and path == "/v1/info":
|
||||
return gateway.info()
|
||||
if method == "GET" and path == "/v1/browsers":
|
||||
return gateway.list_browsers()
|
||||
if method == "POST" and path == "/v1/browsers":
|
||||
return 201, gateway.create(body)
|
||||
match = re.fullmatch(r"/v1/browsers/([a-z0-9][a-z0-9-]{0,31})", path)
|
||||
if match and method == "DELETE":
|
||||
gateway.remove(match.group(1), body)
|
||||
return None
|
||||
match = re.fullmatch(
|
||||
r"/v1/browsers/([a-z0-9][a-z0-9-]{0,31})/(start|stop|cancel|proxy)", path
|
||||
)
|
||||
if match:
|
||||
alias, action = match.groups()
|
||||
if method == "POST" and action in {"start", "stop", "cancel"}:
|
||||
gateway.change_state(alias, action, body)
|
||||
return None
|
||||
if method == "POST" and action == "proxy":
|
||||
gateway.restore_proxy(alias, body)
|
||||
return None
|
||||
match = re.fullmatch(
|
||||
r"/v1/browsers/([a-z0-9][a-z0-9-]{0,31})/douyin/(get|media|image|identity|resolve|login-qr|action|messages|inbox|events)",
|
||||
path,
|
||||
)
|
||||
if match:
|
||||
alias, action = match.groups()
|
||||
if action == "get" and method == "POST":
|
||||
return gateway.get_douyin(alias, body)
|
||||
if action == "resolve" and method == "POST":
|
||||
return gateway.resolve_douyin(alias, body)
|
||||
if action == "media" and method == "POST":
|
||||
return gateway.get_douyin_media(alias, body)
|
||||
if action == "image" and method == "POST":
|
||||
return gateway.get_douyin_image(alias, body)
|
||||
if action == "identity" and method == "POST":
|
||||
return gateway.douyin_identity(alias, body)
|
||||
if action == "login-qr" and method == "POST":
|
||||
return gateway.douyin_login_qr(alias, body)
|
||||
if action == "action" and method == "POST":
|
||||
return gateway.douyin_action(alias, body)
|
||||
if action == "inbox" and method == "POST":
|
||||
return gateway.douyin_inbox(alias, body)
|
||||
if action == "messages" and method == "POST":
|
||||
return gateway.douyin_message_history(alias, body)
|
||||
if action == "events":
|
||||
if method == "POST":
|
||||
return gateway.start_douyin_events(alias, body)
|
||||
if method == "DELETE":
|
||||
gateway.stop_douyin_events(alias, body)
|
||||
return None
|
||||
raise RequestError("not found", 404)
|
||||
return route_gateway_request(self.server_as_gateway().gateway, method, path, query, body)
|
||||
|
||||
def server_as_gateway(self) -> GatewayHTTPServer:
|
||||
if not isinstance(self.server, GatewayHTTPServer):
|
||||
|
||||
@@ -0,0 +1,384 @@
|
||||
"""Outbound agent contract tests: connection, heartbeat, task dispatch, event push."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from importlib import import_module
|
||||
from unittest.mock import Mock
|
||||
|
||||
import websocket
|
||||
from websockets.sync.server import serve
|
||||
|
||||
agent_module = import_module(f"{__package__}.agent")
|
||||
gateway_module = import_module(f"{__package__}.server.http")
|
||||
|
||||
AgentConfig = agent_module.AgentConfig
|
||||
OutboundAgent = agent_module.OutboundAgent
|
||||
RequestError = gateway_module.RequestError
|
||||
|
||||
|
||||
class FakePlatform:
|
||||
"""Minimal platform-side WS peer capturing frames and replaying scripted ones."""
|
||||
|
||||
def __init__(self, *, reject: bool = False) -> None:
|
||||
self.reject = reject
|
||||
self.server = serve(self._handler, "127.0.0.1", 0, process_request=self._process_request)
|
||||
self.port = self.server.socket.getsockname()[1]
|
||||
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
||||
self.thread.start()
|
||||
self.received: list[dict] = []
|
||||
self.lock = threading.Lock()
|
||||
self.connections: list = []
|
||||
self.on_task: dict[int, dict] = {}
|
||||
self.reject_task_ids: set[int] = set()
|
||||
self.keep_open = threading.Event()
|
||||
|
||||
def url(self) -> str:
|
||||
return f"ws://127.0.0.1:{self.port}/v1/agent"
|
||||
|
||||
def _process_request(self, connection, request):
|
||||
if self.reject:
|
||||
from websockets.http11 import Response
|
||||
return Response(403, "Forbidden", b"invalid access key\r\n")
|
||||
return None
|
||||
|
||||
def _handler(self, ws) -> None:
|
||||
with self.lock:
|
||||
self.connections.append(ws)
|
||||
try:
|
||||
while True:
|
||||
raw = ws.recv(timeout=5)
|
||||
if not isinstance(raw, str):
|
||||
continue
|
||||
frame = json.loads(raw)
|
||||
with self.lock:
|
||||
self.received.append(frame)
|
||||
if frame.get("type") == "task":
|
||||
task_id = frame["id"]
|
||||
if task_id in self.reject_task_ids:
|
||||
ws.send(json.dumps({"type": "result", "id": task_id, "status": 409, "error": "duplicate task id"}))
|
||||
continue
|
||||
scripted = self.on_task.get(task_id)
|
||||
if scripted is not None:
|
||||
ws.send(json.dumps({"type": "result", "id": task_id, **scripted}))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def wait_for(self, predicate, timeout: float = 3.0) -> dict | None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
with self.lock:
|
||||
for frame in self.received:
|
||||
if predicate(frame):
|
||||
return frame
|
||||
time.sleep(0.01)
|
||||
return None
|
||||
|
||||
def latest_connection(self):
|
||||
with self.lock:
|
||||
return self.connections[-1] if self.connections else None
|
||||
|
||||
def close(self) -> None:
|
||||
self.keep_open.set()
|
||||
self.server.shutdown()
|
||||
for ws in list(self.connections):
|
||||
try:
|
||||
ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def make_gateway() -> Mock:
|
||||
gateway = Mock()
|
||||
gateway.token = "unused"
|
||||
manager = Mock()
|
||||
# 事件桥依赖真实 Condition/版本语义；替身必须与 SubscriptionManager 同构。
|
||||
manager.changed = threading.Condition()
|
||||
manager.version = 0
|
||||
manager._items = {}
|
||||
gateway.subscriptions = manager
|
||||
return gateway
|
||||
|
||||
|
||||
def make_agent(gateway: Mock, platform: FakePlatform) -> OutboundAgent:
|
||||
return OutboundAgent(gateway, AgentConfig(
|
||||
platform_url=platform.url(),
|
||||
access_key="agent-test-key-0001",
|
||||
heartbeat_interval=0.1,
|
||||
reconnect_base=0.05,
|
||||
reconnect_max=0.2,
|
||||
))
|
||||
|
||||
|
||||
class OutboundAgentConnectionTests(unittest.TestCase):
|
||||
def test_heartbeats_flow_with_version_until_stopped(self) -> None:
|
||||
platform = FakePlatform()
|
||||
agent = make_agent(make_gateway(), platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
beat = platform.wait_for(lambda f: f.get("type") == "heartbeat")
|
||||
self.assertIsNotNone(beat, "first heartbeat missing")
|
||||
self.assertEqual(beat.get("version"), agent.config.version)
|
||||
second = platform.wait_for(
|
||||
lambda f: f.get("type") == "heartbeat" and f is not beat, timeout=3.0
|
||||
)
|
||||
self.assertIsNotNone(second, "periodic heartbeat missing")
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
self.assertFalse(runner.is_alive())
|
||||
|
||||
def test_auth_failure_is_reported_and_retried_with_backoff(self) -> None:
|
||||
platform = FakePlatform(reject=True)
|
||||
agent = make_agent(make_gateway(), platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
deadline = time.monotonic() + 3.0
|
||||
while time.monotonic() < deadline and agent.reconnect_attempts < 2:
|
||||
time.sleep(0.01)
|
||||
self.assertGreaterEqual(agent.reconnect_attempts, 2, "auth failure must trigger retry")
|
||||
self.assertTrue(agent.last_error, "auth failure must be reported")
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
|
||||
class OutboundAgentTaskTests(unittest.TestCase):
|
||||
def test_task_dispatches_route_and_returns_result(self) -> None:
|
||||
platform = FakePlatform()
|
||||
gateway = make_gateway()
|
||||
gateway.info.return_value = {"service": "browser-gateway", "node_id": "node-a"}
|
||||
agent = make_agent(gateway, platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
self.wait_until_connected(agent)
|
||||
platform.latest_connection().send(json.dumps({
|
||||
"type": "task", "id": 7, "method": "GET", "path": "/v1/info", "payload": None,
|
||||
}))
|
||||
result = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 7)
|
||||
self.assertIsNotNone(result, "task result missing")
|
||||
self.assertEqual(result.get("status"), 200)
|
||||
import base64
|
||||
self.assertEqual(
|
||||
json.loads(base64.b64decode(result.get("body"))),
|
||||
{"service": "browser-gateway", "node_id": "node-a"},
|
||||
)
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
def test_unknown_route_returns_explicit_404_result(self) -> None:
|
||||
platform = FakePlatform()
|
||||
gateway = make_gateway()
|
||||
agent = make_agent(gateway, platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
self.wait_until_connected(agent)
|
||||
platform.latest_connection().send(json.dumps({
|
||||
"type": "task", "id": 8, "method": "GET", "path": "/v1/nothing", "payload": None,
|
||||
}))
|
||||
result = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 8)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.get("status"), 404)
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
def test_duplicate_task_id_is_rejected_idempotently(self) -> None:
|
||||
platform = FakePlatform()
|
||||
gateway = make_gateway()
|
||||
gateway.list_browsers.return_value = []
|
||||
agent = make_agent(gateway, platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
self.wait_until_connected(agent)
|
||||
frame = {"type": "task", "id": 9, "method": "GET", "path": "/v1/browsers", "payload": None}
|
||||
platform.latest_connection().send(json.dumps(frame))
|
||||
first = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 9)
|
||||
self.assertIsNotNone(first)
|
||||
platform.latest_connection().send(json.dumps(frame))
|
||||
second = platform.wait_for(
|
||||
lambda f: f.get("type") == "result" and f.get("id") == 9 and f is not first,
|
||||
timeout=3.0,
|
||||
)
|
||||
self.assertIsNotNone(second)
|
||||
self.assertEqual(second.get("status"), 409)
|
||||
self.assertIn("duplicate", second.get("error", ""))
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
def test_route_exception_maps_to_error_result(self) -> None:
|
||||
platform = FakePlatform()
|
||||
gateway = make_gateway()
|
||||
gateway.list_browsers.side_effect = ValueError("invalid alias")
|
||||
agent = make_agent(gateway, platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
self.wait_until_connected(agent)
|
||||
platform.latest_connection().send(json.dumps({
|
||||
"type": "task", "id": 10, "method": "GET", "path": "/v1/browsers", "payload": None,
|
||||
}))
|
||||
result = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 10)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.get("status"), 400)
|
||||
self.assertIn("invalid alias", result.get("error", ""))
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
@staticmethod
|
||||
def wait_until_connected(agent: OutboundAgent, timeout: float = 5.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline and not agent.connected.is_set():
|
||||
time.sleep(0.01)
|
||||
if not agent.connected.is_set():
|
||||
raise AssertionError("agent did not connect")
|
||||
|
||||
|
||||
class OutboundAgentEventTests(unittest.TestCase):
|
||||
def test_pending_events_are_pushed_and_ack_removes_them(self) -> None:
|
||||
platform = FakePlatform()
|
||||
gateway = make_gateway()
|
||||
manager = gateway.subscriptions
|
||||
manager.version = 1
|
||||
item = Mock()
|
||||
item.error = ""
|
||||
item.stopped = False
|
||||
item.pending.return_value = [
|
||||
{"delivery_id": "d-1", "kind": "notice"},
|
||||
{"delivery_id": "d-2", "kind": "checkpoint"},
|
||||
]
|
||||
manager._items["alias-a"] = item
|
||||
manager._get.return_value = item
|
||||
agent = make_agent(gateway, platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
OutboundAgentTaskTests.wait_until_connected(agent)
|
||||
batch = platform.wait_for(lambda f: f.get("type") == "event")
|
||||
self.assertIsNotNone(batch, "event batch missing")
|
||||
self.assertEqual(batch["event"]["alias"], "alias-a")
|
||||
self.assertEqual(len(batch["event"]["deliveries"]), 2)
|
||||
platform.latest_connection().send(json.dumps({
|
||||
"type": "event_ack", "alias": "alias-a", "delivery_ids": ["d-1", "d-2"],
|
||||
}))
|
||||
deadline = time.monotonic() + 3.0
|
||||
while time.monotonic() < deadline and manager.ack.call_count == 0:
|
||||
time.sleep(0.01)
|
||||
manager.ack.assert_called_once()
|
||||
args = manager.ack.call_args
|
||||
self.assertEqual(args.args[0], "alias-a")
|
||||
self.assertEqual(sorted(args.args[1]), ["d-1", "d-2"])
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
def test_nack_resets_unconfirmed_batch_for_repush(self) -> None:
|
||||
platform = FakePlatform()
|
||||
gateway = make_gateway()
|
||||
manager = gateway.subscriptions
|
||||
manager.changed = threading.Condition()
|
||||
item = Mock()
|
||||
item.error = ""
|
||||
item.stopped = False
|
||||
item.pending.return_value = [{"delivery_id": "d-9", "kind": "notice"}]
|
||||
manager._items["alias-a"] = item
|
||||
manager._get.return_value = item
|
||||
agent = make_agent(gateway, platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
OutboundAgentTaskTests.wait_until_connected(agent)
|
||||
first = platform.wait_for(lambda f: f.get("type") == "event")
|
||||
self.assertIsNotNone(first)
|
||||
platform.latest_connection().send(json.dumps({
|
||||
"type": "event_nack", "subscription": first["event"]["subscription"], "alias": "alias-a",
|
||||
}))
|
||||
repush = platform.wait_for(
|
||||
lambda f: f.get("type") == "event" and f is not first, timeout=3.0
|
||||
)
|
||||
self.assertIsNotNone(repush, "nacked batch was not re-pushed")
|
||||
self.assertEqual([d["delivery_id"] for d in repush["event"]["deliveries"]], ["d-9"])
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
def test_reconnect_clears_unconfirmed_state_and_repushes(self) -> None:
|
||||
platform = FakePlatform()
|
||||
gateway = make_gateway()
|
||||
manager = gateway.subscriptions
|
||||
manager.changed = threading.Condition()
|
||||
item = Mock()
|
||||
item.error = ""
|
||||
item.stopped = False
|
||||
item.pending.return_value = [{"delivery_id": "d-7", "kind": "notice"}]
|
||||
manager._items["alias-a"] = item
|
||||
manager._get.return_value = item
|
||||
agent = make_agent(gateway, platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
OutboundAgentTaskTests.wait_until_connected(agent)
|
||||
first = platform.wait_for(lambda f: f.get("type") == "event")
|
||||
self.assertIsNotNone(first)
|
||||
# 断开连接：agent 应退避重连，并重推未确认批次
|
||||
first_connection = platform.latest_connection()
|
||||
first_connection.close()
|
||||
deadline = time.monotonic() + 5.0
|
||||
repush = None
|
||||
while time.monotonic() < deadline:
|
||||
repush = platform.wait_for(
|
||||
lambda f: f.get("type") == "event" and f is not first, timeout=0.2
|
||||
)
|
||||
if repush is not None:
|
||||
break
|
||||
self.assertIsNotNone(repush, "unconfirmed batch not re-pushed after reconnect")
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
def test_subscription_error_is_surfaced_as_error_frame(self) -> None:
|
||||
platform = FakePlatform()
|
||||
gateway = make_gateway()
|
||||
manager = gateway.subscriptions
|
||||
manager.changed = threading.Condition()
|
||||
item = Mock()
|
||||
item.error = "browser identity changed"
|
||||
item.stopped = False
|
||||
item.pending.return_value = []
|
||||
manager._items["alias-a"] = item
|
||||
manager._get.return_value = item
|
||||
agent = make_agent(gateway, platform)
|
||||
runner = threading.Thread(target=agent.run, daemon=True)
|
||||
runner.start()
|
||||
try:
|
||||
OutboundAgentTaskTests.wait_until_connected(agent)
|
||||
error = platform.wait_for(lambda f: f.get("type") == "error" and f.get("alias"))
|
||||
self.assertIsNotNone(error, "subscription error frame missing")
|
||||
self.assertIn("identity changed", error.get("error", ""))
|
||||
finally:
|
||||
agent.stop()
|
||||
runner.join(timeout=5)
|
||||
platform.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user