stroblmeandClaude Opus 5.5
391edc8766
The connector now serves /mcp, its OAuth documents and /hooks/ besides /api/v1/, refuses any dot segment (httpx resolves them, which carried /api/v1/../../mcp onto the MCP mount), addresses calls to the hub's host, and bridges served sockets by calling the app with a websocket scope. A hook without a secret answers the tunnel with the 404 (or 1008) of a wrong secret. Requests through the hub get the hub's issuer, resource and consent page in the OAuth metadata and the 401, and the MCP host check admits the hub while enrolled. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WgBfkm4rCmtd1T4vHWjEPx
432 lines
16 KiB
Python
432 lines
16 KiB
Python
"""WebSocket nodes: a socket the engine dials or serves, listening or sending."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import hmac
|
|
import json
|
|
import logging
|
|
from collections.abc import Awaitable, Callable, Iterable
|
|
from enum import Enum
|
|
from typing import TYPE_CHECKING, Any
|
|
from urllib.parse import urlsplit
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from fluksio.flow.messages import MessageSpec
|
|
from fluksio.flow.nodes.base import Node
|
|
from fluksio.flow.nodes.http import HttpNode, _first_catch_all
|
|
|
|
if TYPE_CHECKING:
|
|
from fastapi import FastAPI
|
|
from starlette.websockets import WebSocket
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WebSocketNode(Node):
|
|
"""
|
|
A long-lived WebSocket, as a client or a server.
|
|
|
|
The URL picks the role: ``ws://`` or ``wss://`` dials out and holds the
|
|
connection, a path such as ``/ticker`` is served at
|
|
``/hooks/<flow>/<path>[/<secret>]`` while the flow runs. The ports pick the
|
|
direction, as on the mqtt node: outputs only make a **listener**, inputs a
|
|
**sender**.
|
|
|
|
- client listener: frames the server sends become outputs
|
|
- client sender: inputs go to the server, one frame per message
|
|
- server listener: frames any connected client sends become outputs
|
|
- server sender: inputs go to every connected client
|
|
|
|
Frames are text. An incoming frame is parsed as JSON; an object feeds the
|
|
ports whose names (or ``json_key``) it carries, and a node with one output
|
|
takes anything else whole. Outgoing, a single input is sent as it stands —
|
|
a string raw, anything else as JSON — and several go as one object keyed
|
|
by port.
|
|
"""
|
|
|
|
class Mode(Enum):
|
|
"""Which way the frames go."""
|
|
|
|
LISTENER = "listener"
|
|
SENDER = "sender"
|
|
|
|
# A frame sent twice is two commands to whatever reads it.
|
|
idempotent = False
|
|
|
|
__slots__ = (
|
|
"url",
|
|
"mode",
|
|
"serves",
|
|
"headers",
|
|
"on_connect",
|
|
"secret",
|
|
"timeout",
|
|
"send_queue_size",
|
|
"json_keys",
|
|
"_clients",
|
|
"_route_registered",
|
|
"_task",
|
|
"_queue",
|
|
"_loop",
|
|
)
|
|
|
|
class Params(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
url: str = Field(
|
|
description=(
|
|
"ws:// or wss:// to connect to a server, or a path such as "
|
|
"/ticker to serve one at /hooks/<flow>/<path>."
|
|
)
|
|
)
|
|
headers: dict[str, str] = Field(
|
|
default_factory=dict,
|
|
description=(
|
|
"Handshake headers, client only. A value may be a secret reference."
|
|
),
|
|
)
|
|
on_connect: list[Any] = Field(
|
|
default_factory=list,
|
|
description=(
|
|
"Frames sent after every (re)connect, client only — the "
|
|
"subscribe or auth message many APIs expect. A string goes "
|
|
"raw, anything else as JSON; an item may be a secret reference."
|
|
),
|
|
)
|
|
secret: str = Field(
|
|
default="",
|
|
description=(
|
|
"Shared secret for a served socket, appended to its URL: "
|
|
"/hooks/<flow>/<url>/<secret>. Empty leaves it open to anyone."
|
|
),
|
|
json_schema_extra={"x-secret": True},
|
|
)
|
|
timeout: float = Field(
|
|
default=10.0,
|
|
gt=0,
|
|
description=(
|
|
"Seconds allowed for the handshake, and for a send to one "
|
|
"client of a served socket."
|
|
),
|
|
)
|
|
send_queue_size: int = Field(
|
|
default=256,
|
|
gt=0,
|
|
description=(
|
|
"How many messages may wait to be sent. Past this the oldest "
|
|
"is dropped and the node reports degraded."
|
|
),
|
|
)
|
|
# Same as the mqtt node: a device that wraps its reading in an object.
|
|
json_key: str | dict[str, str] = Field(
|
|
default="",
|
|
description=(
|
|
"Key to lift out of a JSON object frame, for every port or as "
|
|
'{"port": "key"}.'
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def target_key(cls, params: dict[str, Any]) -> str | None:
|
|
# A URL, credentials stripped, exactly as for the http node.
|
|
return HttpNode.target_key(params)
|
|
|
|
def __init__(
|
|
self,
|
|
requires: MessageSpec | Iterable[MessageSpec] = (),
|
|
provides: MessageSpec | Iterable[MessageSpec] = (),
|
|
params: dict[str, Any] | None = None,
|
|
name: str | None = None,
|
|
):
|
|
cfg = self.Params.model_validate(params or {})
|
|
|
|
requires = Node._normalize_ports(requires)
|
|
provides = Node._normalize_ports(provides)
|
|
if not requires and not provides:
|
|
raise ValueError(
|
|
"A WebSocket node needs either inputs (to send) or outputs (to listen)"
|
|
)
|
|
# ponytail: one direction per node, as mqtt and http. Inputs and
|
|
# outputs over one socket is the NOTEPAD follow-up.
|
|
self.mode = (
|
|
WebSocketNode.Mode.SENDER if requires else WebSocketNode.Mode.LISTENER
|
|
)
|
|
|
|
self.serves = cfg.url.startswith("/")
|
|
if not self.serves and urlsplit(cfg.url).scheme not in ("ws", "wss"):
|
|
raise ValueError(
|
|
"A WebSocket url is ws://…, wss://… or a path to serve, "
|
|
f"not '{cfg.url}'"
|
|
)
|
|
|
|
ports = requires or provides
|
|
if isinstance(cfg.json_key, dict):
|
|
self.json_keys: dict[str, str] = dict(cfg.json_key)
|
|
else:
|
|
self.json_keys = (
|
|
{spec.port: cfg.json_key for spec in ports} if cfg.json_key else {}
|
|
)
|
|
|
|
self.url = cfg.url
|
|
self.headers = cfg.headers
|
|
self.on_connect = [_frame(item) for item in cfg.on_connect]
|
|
self.secret = cfg.secret
|
|
self.timeout = cfg.timeout
|
|
self.send_queue_size = cfg.send_queue_size
|
|
|
|
self._clients: set[WebSocket] = set()
|
|
self._route_registered = False
|
|
self._task: asyncio.Task[None] | None = None
|
|
self._queue: asyncio.Queue[dict[str, Any]] | None = None
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
if name is None:
|
|
safe_url = (
|
|
self.url.split("://", 1)[-1].replace("/", "_").replace(":", "")
|
|
).strip("_")
|
|
name = f"websocket_{self.mode.value}_{safe_url}"
|
|
|
|
super().__init__(
|
|
f=(
|
|
self._sender_handler
|
|
if self.mode is WebSocketNode.Mode.SENDER
|
|
else self._noop_listener
|
|
),
|
|
requires=requires,
|
|
provides=provides,
|
|
params=params,
|
|
name=name,
|
|
)
|
|
|
|
@staticmethod
|
|
def _noop_listener(params: dict[str, Any], **kwargs: Any) -> dict[str, Any] | None:
|
|
"""A listener injects what arrives; it is never called with inputs."""
|
|
return None
|
|
|
|
# -- frames ---------------------------------------------------------------
|
|
|
|
def _encode(self, data: dict[str, Any]) -> str:
|
|
"""One message as one frame."""
|
|
if len(data) == 1:
|
|
return _frame(next(iter(data.values())))
|
|
return json.dumps(data, default=str)
|
|
|
|
def _decode(self, text: str) -> dict[str, Any]:
|
|
"""One frame as outputs, keyed by port. Empty when nothing matched."""
|
|
try:
|
|
parsed = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
parsed = text
|
|
outputs: dict[str, Any] = {}
|
|
if isinstance(parsed, dict):
|
|
for spec in self.output_ports:
|
|
key = self.json_keys.get(spec.port, spec.port)
|
|
if key in parsed:
|
|
outputs[spec.port] = spec.coerce(parsed[key])
|
|
if not outputs and len(self.output_ports) == 1:
|
|
spec = self.output_ports[0]
|
|
outputs[spec.port] = spec.coerce(parsed)
|
|
return outputs
|
|
|
|
async def _receive(self, text: str) -> None:
|
|
"""Inject one frame. A bad one is logged, never fatal to the socket."""
|
|
try:
|
|
outputs = self._decode(text)
|
|
if outputs:
|
|
await asyncio.to_thread(self.inject, outputs)
|
|
else:
|
|
logger.debug("[%s] Frame matched no port: %s", self.name, text)
|
|
except Exception as exc:
|
|
logger.error("[%s] Error processing frame: %s", self.name, exc)
|
|
|
|
# -- sending --------------------------------------------------------------
|
|
|
|
def _sender_handler(
|
|
self, params: dict[str, Any], **kwargs: Any
|
|
) -> dict[str, Any] | None:
|
|
"""Hand the inputs to the node's socket task, which owns the send."""
|
|
loop, queue = self._loop, self._queue
|
|
if loop is not None and queue is not None:
|
|
loop.call_soon_threadsafe(self._enqueue, queue, dict(kwargs))
|
|
return None
|
|
|
|
# Not started: a draft preview or a test. A client connects once and
|
|
# sends; a served socket has nobody to send to.
|
|
if not self.serves:
|
|
from websockets.sync.client import connect
|
|
|
|
with connect(
|
|
self.url,
|
|
additional_headers=self.headers,
|
|
open_timeout=self.timeout,
|
|
) as ws:
|
|
for frame in self.on_connect:
|
|
ws.send(frame)
|
|
ws.send(self._encode(dict(kwargs)))
|
|
return None
|
|
|
|
def _enqueue(
|
|
self, queue: asyncio.Queue[dict[str, Any]], data: dict[str, Any]
|
|
) -> None:
|
|
"""Queue a message, dropping the oldest when the socket cannot keep up."""
|
|
if queue.full():
|
|
with contextlib.suppress(asyncio.QueueEmpty):
|
|
queue.get_nowait()
|
|
logger.warning("Send queue full for node '%s', dropped", self.name)
|
|
self.report_health("degraded", "send queue full")
|
|
queue.put_nowait(data)
|
|
|
|
async def _send_queued(self, send: Callable[[str], Awaitable[None]]) -> None:
|
|
"""Send whatever is queued, in order, for as long as the socket lives."""
|
|
queue = self._queue
|
|
if queue is None:
|
|
return
|
|
while True:
|
|
await send(self._encode(await queue.get()))
|
|
|
|
async def _broadcast(self, text: str) -> None:
|
|
"""One frame to every connected client.
|
|
|
|
ponytail: sequential, so a stalled client holds the others up for
|
|
`timeout`. Per-client queues if a served socket ever has many readers.
|
|
"""
|
|
for ws in list(self._clients):
|
|
try:
|
|
await asyncio.wait_for(ws.send_text(text), self.timeout)
|
|
except Exception:
|
|
self._clients.discard(ws)
|
|
with contextlib.suppress(Exception):
|
|
await asyncio.wait_for(ws.close(code=1011), self.timeout)
|
|
|
|
# -- client ---------------------------------------------------------------
|
|
|
|
async def _client_loop(self) -> None:
|
|
"""Hold one connection. A drop raises, and the supervisor reconnects."""
|
|
from websockets.asyncio.client import connect
|
|
|
|
try:
|
|
async with connect(
|
|
self.url,
|
|
additional_headers=self.headers,
|
|
open_timeout=self.timeout,
|
|
) as ws:
|
|
for frame in self.on_connect:
|
|
await ws.send(frame)
|
|
self.report_health("ok")
|
|
|
|
if self.mode is WebSocketNode.Mode.LISTENER:
|
|
async for received in ws:
|
|
if isinstance(received, str):
|
|
await self._receive(received)
|
|
else:
|
|
sender = asyncio.create_task(self._send_queued(ws.send))
|
|
try:
|
|
# Replies are not a sender's to publish, but they are
|
|
# read: an unread buffer stops the pings being answered.
|
|
async for _ in ws:
|
|
pass
|
|
finally:
|
|
await self._cancel_task(sender)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
self.report_health("down", str(exc))
|
|
raise
|
|
self.report_health("down", "connection closed")
|
|
raise ConnectionError(f"WebSocket for '{self.name}' closed")
|
|
|
|
# -- server ---------------------------------------------------------------
|
|
|
|
async def _serve(self, websocket: WebSocket) -> None:
|
|
"""One client of a served socket, for as long as it stays."""
|
|
from fluksio.cloud.config import via_portal
|
|
|
|
if (
|
|
self.secret
|
|
and not hmac.compare_digest(
|
|
str(websocket.path_params.get("secret", "")).encode(),
|
|
self.secret.encode(),
|
|
)
|
|
) or (not self.secret and via_portal(websocket.headers)):
|
|
# Closed before accepting: the same 403 as a path with no route.
|
|
# Through the portal a socket needs a secret, as a webhook does.
|
|
await websocket.close(code=1008)
|
|
return
|
|
await websocket.accept()
|
|
self._clients.add(websocket)
|
|
try:
|
|
while True:
|
|
message = await websocket.receive()
|
|
if message["type"] == "websocket.disconnect":
|
|
break
|
|
text = message.get("text")
|
|
# A sender's clients are read only to notice them leave.
|
|
if text is not None and self.mode is WebSocketNode.Mode.LISTENER:
|
|
await self._receive(text)
|
|
finally:
|
|
self._clients.discard(websocket)
|
|
|
|
def register_route(self, app: FastAPI) -> None:
|
|
"""Mount the served socket, ahead of any catch-all, as webhooks are."""
|
|
if self._route_registered:
|
|
return
|
|
from starlette.routing import WebSocketRoute
|
|
|
|
# The secret is a path parameter, so the route table never holds it.
|
|
path = f"{self.url}/{{secret}}" if self.secret else self.url
|
|
app.routes.insert(
|
|
_first_catch_all(app), WebSocketRoute(path, self._serve, name=self.id)
|
|
)
|
|
self._route_registered = True
|
|
logger.info("Registered websocket route for node '%s': %s", self.name, path)
|
|
|
|
def unregister_route(self, app: FastAPI) -> None:
|
|
if not self._route_registered:
|
|
return
|
|
app.routes[:] = [
|
|
route for route in app.routes if getattr(route, "name", None) != self.id
|
|
]
|
|
self._route_registered = False
|
|
|
|
# -- lifecycle ------------------------------------------------------------
|
|
|
|
async def start(self, app: FastAPI | None = None) -> None:
|
|
if self.serves and app is not None:
|
|
self.register_route(app)
|
|
if self.mode is WebSocketNode.Mode.SENDER and self._queue is None:
|
|
# Outlives a reconnect, so what was queued meanwhile still goes.
|
|
self._queue = asyncio.Queue(maxsize=self.send_queue_size)
|
|
self._loop = asyncio.get_running_loop()
|
|
# Spawned on every start: enabling a flow cancels its supervised loops
|
|
# and starts the nodes again, and the supervisor ignores a duplicate.
|
|
if self._task is not None and not self._task.done():
|
|
return
|
|
if not self.serves:
|
|
self._task = self._run_supervised("websocket", self._client_loop)
|
|
elif self.mode is WebSocketNode.Mode.SENDER:
|
|
self._task = self._run_supervised(
|
|
"websocket-out", lambda: self._send_queued(self._broadcast)
|
|
)
|
|
|
|
async def stop(self, app: FastAPI | None = None) -> None:
|
|
if self.serves and app is not None:
|
|
self.unregister_route(app)
|
|
for ws in list(self._clients):
|
|
with contextlib.suppress(Exception):
|
|
await ws.close(code=1001)
|
|
self._clients.clear()
|
|
if self._task is not None:
|
|
await self._cancel_task(self._task)
|
|
self._task = None
|
|
self._queue = None
|
|
self._loop = None
|
|
|
|
|
|
def _frame(value: Any) -> str:
|
|
"""A value as frame text: a string as it stands, anything else as JSON."""
|
|
# default=str: a value JSON cannot hold must not kill the send loop.
|
|
return value if isinstance(value, str) else json.dumps(value, default=str)
|