team_mail_websocket.py
Rama main · réplica de lectura del repositorio original.
"""Authenticated local passive Mail socket, outside managed Hub ASGI drains.
No global subscription copies, mailbox fetches, shared executor waits or idle
queries. At most sixteen dedicated Mail workers each own one coalescing broker
lease and one serial bounded ASGI writer. Production remains feature-gated.
"""
from __future__ import annotations
import asyncio
from contextlib import suppress
import json
import threading
import uuid
from typing import Any, Callable
from agentsdock_team_hub.mail_hints import MAIL_IDENTIFIER_RE, MailArrival
from agentsdock_team_hub.notification_hints import NotificationCursor
MAIL_WEBSOCKET_PROTOCOL = "agentsdock.team-mail-hints.v1"
NOTIFICATION_WEBSOCKET_PROTOCOL = "agentsdock.team-mail-hints.v2"
_SLOTS = threading.BoundedSemaphore(16)
class _MailSocketWriter:
"""Bridge one dedicated worker to an actual event-loop Task completion.
A cancelled concurrent Future is not proof its asyncio Task settled. Keep
the worker's lease/slot until a Task done callback proves completion, even
if an ASGI implementation resists cancellation. That consumes only this
bounded Mail capacity, never a control request slot or event-loop wait.
"""
def __init__(self, loop, *, timeout=5):
self.loop, self.timeout = loop, timeout
self.guard = threading.Lock()
self.stopped = False
self.task = None
def cancel(self):
with self.guard:
self.stopped = True
def cancel_task():
with self.guard:
task = self.task
if task is not None:
task.cancel()
with suppress(RuntimeError):
self.loop.call_soon_threadsafe(cancel_task)
def send(self, coroutine_factory):
settled = threading.Event()
outcome = {}
def finished(task):
try:
outcome["result"] = task.result()
except BaseException as exc:
outcome["error"] = exc
with self.guard:
if self.task is task:
self.task = None
settled.set()
def start():
try:
task = self.loop.create_task(coroutine_factory())
# This callback also runs when cancelled before its first step.
task.add_done_callback(finished)
with self.guard:
self.task = task
stopped = self.stopped
if stopped:
task.cancel()
except BaseException as exc:
outcome["error"] = exc
settled.set()
self.loop.call_soon_threadsafe(start)
if not settled.wait(self.timeout):
self.cancel()
settled.wait()
if "error" in outcome:
raise outcome["error"]
async def serve_team_mail_hints(ws: Any, runtime: Any, *, server_identity: str,
authorized: Callable[[], bool], protocols: list[str]) -> None:
protocol = (NOTIFICATION_WEBSOCKET_PROTOCOL if NOTIFICATION_WEBSOCKET_PROTOCOL in protocols
else MAIL_WEBSOCKET_PROTOCOL if MAIL_WEBSOCKET_PROTOCOL in protocols else None)
version = 2 if protocol == NOTIFICATION_WEBSOCKET_PROTOCOL else 1
await ws.accept(subprotocol=protocol)
if not authorized() or ws.query_params:
await ws.close(code=4401)
return
capability = (getattr(runtime, "team_notification_hint_capability", lambda: {}) if version == 2
else runtime.team_mail_hint_capability)
if protocol is None or not capability().get("enabled"):
await ws.close(code=4406)
return
if not _SLOTS.acquire(blocking=False):
await ws.close(code=1013)
return
loop = asyncio.get_running_loop()
done = asyncio.Event()
stopped = threading.Event()
guard = threading.Lock()
holder: dict[str, Any] = {}
writer = _MailSocketWriter(loop)
stream_id = uuid.uuid4().hex
worker = None
try:
raw = await asyncio.wait_for(ws.receive_text(), timeout=5)
if len(raw.encode("utf-8")) > 4096:
raise ValueError("Mail subscription frame is too large")
request = json.loads(raw)
if (not isinstance(request, dict) or set(request) != {"version", "team_id", "previous_cursor"}
or type(request["version"]) is not int or request["version"] != version
or not isinstance(request["team_id"], str) or MAIL_IDENTIFIER_RE.fullmatch(request["team_id"]) is None):
raise ValueError("Mail subscription frame is invalid")
if request["previous_cursor"] is not None:
previous = (NotificationCursor.from_dict(request["previous_cursor"]) if version == 2
else MailArrival.from_dict(request["previous_cursor"]))
if previous.mailbox[0] != request["team_id"]:
await ws.close(code=4403)
return
def abort():
stopped.set()
writer.cancel()
def run():
lease = None
try:
subscriber = (runtime.subscribe_team_notification_hints if version == 2
else runtime.subscribe_team_mail_hints)
lease = subscriber(request["team_id"], request["previous_cursor"])
with guard:
holder["lease"] = lease
lease.set_aborter(abort)
if stopped.is_set():
return
def write(kind, cursor):
async def send():
if stopped.is_set() or lease.closed or not authorized():
raise RuntimeError("Mail socket authority changed")
frame = {"type": kind, "server_identity": server_identity,
"hub_id": lease.hub_id, "stream_id": stream_id, "cursor": cursor}
data = json.dumps(frame, separators=(",", ":"))
if len(data.encode("utf-8")) > 4096:
raise ValueError("Mail socket metadata is too large")
await ws.send_text(data)
writer.send(send)
lease.write(lambda cursor: write("snapshot", cursor), lease.snapshot)
while not stopped.is_set():
arrival = lease.take()
if arrival is None:
break
lease.write(lambda cursor: write("hint", cursor), arrival.as_dict(reset=False))
except asyncio.CancelledError:
pass
except Exception as exc:
holder["close_code"] = (4403 if getattr(exc, "status_code", None) in (401, 403)
else 4406 if getattr(exc, "status_code", None) in (404, 501)
else 1008 if isinstance(exc, ValueError) or getattr(exc, "code", None) == "remote_invalid"
else 1012)
finally:
if lease is not None:
holder.setdefault("close_code", lease.close_code)
lease.close()
_SLOTS.release()
with suppress(RuntimeError):
loop.call_soon_threadsafe(done.set)
pending_worker = threading.Thread(target=run, name="agentsdock-mail-local", daemon=True)
pending_worker.start()
worker = pending_worker
receive = asyncio.create_task(ws.receive_text())
finished = asyncio.create_task(done.wait())
try:
completed, _pending = await asyncio.wait((receive, finished), return_when=asyncio.FIRST_COMPLETED)
if receive in completed:
# Subscriptions have no command channel or acknowledgement.
receive.result()
await ws.close(code=1008)
else:
await ws.close(code=holder.get("close_code", 1012))
finally:
for task in (receive, finished):
task.cancel()
with suppress(BaseException):
await task
except (ValueError, asyncio.TimeoutError):
with suppress(Exception):
await ws.close(code=1008)
except Exception:
pass
finally:
stopped.set()
with guard:
lease = holder.get("lease")
if lease is not None:
lease.cancel()
if worker is not None:
# No join/default executor on the event loop. The socket's bounded
# worker signals completion after closing its subscription.
with suppress(asyncio.TimeoutError):
await asyncio.wait_for(done.wait(), timeout=15)
else:
_SLOTS.release()