springbootAI 1.8.0__py3-none-any.whl

spring/websocket/__init__.py

@@ -0,0 +1,86 @@
+"""SpringBootAI WebSocket 实时通信模块 —— 注解驱动的 WebSocket 端点与消息路由（对齐
+JSR-356 ``@ServerEndpoint`` + Spring WebSocket ``@MessageMapping`` / ``@SendTo`` 体系）。
+
+模块组成：
+- **session**:   ``WebSocketSession`` 会话抽象 + ``WebSocketSessionRegistry`` 全局会话注册表
+                  （对齐 Spring ``WebSocketSession`` / ``WebSocketHandlerRegistry``）
+- **handler**:   ``WebSocketHandler`` 接口 + ``TextWebSocketHandler`` / ``BinaryWebSocketHandler``
+                  便捷基类 + ``@ServerEndpoint`` 类级注解（JSR-356 风格生命周期钩子）
+- **annotations**: ``@MessageMapping`` / ``@SendTo`` / ``@SendToUser`` / ``@SubscribeMapping``
+                  方法级注解（Spring STOMP 风格消息路由）
+- **broker**:    ``InMemoryBroker`` 主题式发布订阅代理（对齐 Spring ``SimpleBrokerMessageHandler``）
+                  + ``SimpMessageSendingOperations`` 发送操作 API
+- **router**:    ``WebSocketRouter`` 把注解端点转换为 Starlette/FastAPI WebSocket 路由
+- **exceptions**: ``WebSocketException`` 异常族
+
+设计原则：**复用项目既有范式，不重复造轮子**。
+- ``@ServerEndpoint`` / ``@MessageMapping`` 等注解复用 ``SpringAnnotation`` 元数据描述符范式
+  （与 ``spring.excel`` / ``spring.csv`` / ``spring.validation`` 一致）。
+- 不依赖第三方 WebSocket 库；直接基于 Starlette ``WebSocket`` 原生能力。
+- 不实现完整 STOMP 协议（复杂度过高）；采用简化 JSON 消息帧，对齐 Spring ``SimpleBroker``
+  的核心语义（destination 路由 + topic 订阅）。
+
+与 Java 的差异：
+- JSR-356 在 Java EE 容器内注册端点；本实现通过 ``WebSocketRouter.install(app)`` 挂载到
+  Starlette/FastAPI 应用。
+- Spring STOMP 用 ``@EnableWebSocketMessageBroker`` + ``WebSocketMessageBrokerConfigurer``；
+  本实现用 ``MessageBrokerConfigurer`` + ``WebSocketRouter`` 组合，更轻量。
+- 不支持 STOMP ACK/NACK/事务（可按需扩展 ``InMemoryBroker``）。
+"""
+from .exceptions import (
+    WebSocketException,
+    WebSocketConnectionException,
+    WebSocketHandlerException,
+    MessageBrokerException,
+)
+from .session import (
+    WebSocketSession,
+    WebSocketSessionRegistry,
+    global_session_registry,
+)
+from .handler import (
+    WebSocketHandler,
+    TextWebSocketHandler,
+    BinaryWebSocketHandler,
+    ServerEndpoint,
+    AnnotatedEndpointHandler,
+    discover_server_endpoints,
+)
+from .annotations import (
+    MessageMapping,
+    SendTo,
+    SendToUser,
+    SubscribeMapping,
+    MessageEndpoint,
+    collect_message_mappings,
+    MessageMappingModel,
+)
+from .broker import (
+    InMemoryBroker,
+    SimpMessageSendingOperations,
+    MessageBrokerConfigurer,
+    broker_registry,
+)
+from .router import WebSocketRouter, MessageEndpointDispatcher, install_websocket_routes
+
+__version__ = "1.0.0"
+
+__all__ = [
+    # 异常
+    "WebSocketException", "WebSocketConnectionException",
+    "WebSocketHandlerException", "MessageBrokerException",
+    # 会话
+    "WebSocketSession", "WebSocketSessionRegistry", "global_session_registry",
+    # Handler / @ServerEndpoint
+    "WebSocketHandler", "TextWebSocketHandler", "BinaryWebSocketHandler",
+    "ServerEndpoint", "AnnotatedEndpointHandler", "discover_server_endpoints",
+    # 注解
+    "MessageMapping", "SendTo", "SendToUser", "SubscribeMapping", "MessageEndpoint",
+    "MessageMappingModel", "collect_message_mappings",
+    # Broker
+    "InMemoryBroker", "SimpMessageSendingOperations",
+    "MessageBrokerConfigurer", "broker_registry",
+    # Router
+    "WebSocketRouter", "MessageEndpointDispatcher", "install_websocket_routes",
+    "__version__",
+]

spring/websocket/annotations.py

@@ -0,0 +1,169 @@
+"""WebSocket 消息映射注解（对齐 Spring ``@MessageMapping`` / ``@SendTo`` / ``@SendToUser``
+/ ``@SubscribeMapping``）。
+
+注解语义（与 Spring STOMP 一致）：
+- ``@MessageMapping("/chat")``       方法处理发往 ``/chat`` 的消息（``/app`` 前缀由配置剥离）。
+- ``@SendTo("/topic/greetings")``    方法返回值自动广播到 ``/topic/greetings``。
+                                      不标注时默认回发给发送者（对齐 Spring ``@SendToUser`` 默认）。
+- ``@SendToUser``                    方法返回值定向发给发送者（对齐 Spring ``@SendToUser``）。
+- ``@SubscribeMapping("/topic/x")``  客户端订阅时触发，方法返回值作为初始数据回发给订阅者。
+- ``@MessageEndpoint``               类级注解，标记一个类包含 ``@MessageMapping`` 方法
+                                      （对齐 Spring ``@MessageMapping``+``@Controller`` 组合）。
+
+注解本身只注册元数据；实际方法包装在 ``MessageEndpointDispatcher``（见 router.py）中完成。
+"""
+from __future__ import annotations
+
+from typing import Any, Callable, List, Optional, Type, Union
+
+from spring.annotations.core import SpringAnnotation
+
+
+# ==================== 类级注解 ====================
+
+class MessageEndpoint(SpringAnnotation):
+    """``@MessageEndpoint`` 类级注解，标记一个类为消息端点（含 ``@MessageMapping`` 方法）。
+
+    等价于 Spring ``@Controller`` + ``@MessageMapping`` 的组合约定。
+    """
+
+    _annotation_type = "message_endpoint"
+
+
+# ==================== 方法级注解 ====================
+
+class MessageMapping(SpringAnnotation):
+    """``@MessageMapping("/chat")`` 方法级注解，声明消息处理方法。
+
+    用法::
+
+        @MessageEndpoint
+        class GreetingController:
+            @MessageMapping("/greet")
+            @SendTo("/topic/greetings")
+            def greet(self, message):
+                return {"content": "Hello, " + message["name"]}
+
+    客户端发送 ``{"destination": "/app/greet", "payload": {"name": "Tom"}}`` 时触发。
+    """
+
+    _annotation_type = "message_mapping"
+
+    def __init__(self, value: str = ""):
+        super().__init__(value=value)
+
+
+class SendTo(SpringAnnotation):
+    """``@SendTo("/topic/greetings")`` 方法级注解，把返回值广播到指定 destination。
+
+    不标注时，方法返回值默认回发给发送者（等价于 ``@SendToUser``）。
+    """
+
+    _annotation_type = "send_to"
+
+    def __init__(self, value: str = ""):
+        # 支持单 destination 或多 destination（逗号分隔）
+        dests = [v.strip() for v in value.split(",") if v.strip()] if value else []
+        super().__init__(value=value, destinations=dests)
+
+
+class SendToUser(SpringAnnotation):
+    """``@SendToUser`` 方法级注解，把返回值定向发给发送者（不广播）。
+
+    可指定 ``broadcast``（False 时只发给当前会话；True 时发给该用户所有会话）。
+    """
+
+    _annotation_type = "send_to_user"
+
+    def __init__(self, broadcast: bool = False):
+        super().__init__(broadcast=broadcast)
+
+
+class SubscribeMapping(SpringAnnotation):
+    """``@SubscribeMapping("/topic/init")`` 方法级注解，订阅时触发并返回初始数据。
+
+    与 ``@MessageMapping`` 区别：``@SubscribeMapping`` 在客户端订阅 destination 时触发，
+    返回值直接回发给订阅者（不经过 broker 广播）。
+    """
+
+    _annotation_type = "subscribe_mapping"
+
+    def __init__(self, value: str = ""):
+        super().__init__(value=value)
+
+
+# ==================== 元数据模型 ====================
+
+class MessageMappingModel:
+    """解析后的消息映射元数据，供 ``router.MessageEndpointDispatcher`` 消费。"""
+
+    __slots__ = (
+        "method_name", "destination", "send_to", "send_to_user",
+        "send_to_user_broadcast", "subscribe_destination", "is_subscribe",
+    )
+
+    def __init__(
+        self,
+        method_name: str,
+        destination: str,
+        send_to: Optional[List[str]] = None,
+        send_to_user: bool = False,
+        send_to_user_broadcast: bool = False,
+        subscribe_destination: str = "",
+        is_subscribe: bool = False,
+    ):
+        self.method_name = method_name
+        self.destination = destination
+        self.send_to = send_to or []
+        self.send_to_user = send_to_user
+        self.send_to_user_broadcast = send_to_user_broadcast
+        self.subscribe_destination = subscribe_destination
+        self.is_subscribe = is_subscribe
+
+
+# ==================== 元数据收集 ====================
+
+def collect_message_mappings(cls: Type) -> List[MessageMappingModel]:
+    """收集类上所有 ``@MessageMapping`` / ``@SubscribeMapping`` 方法的元数据。
+
+    遍历 ``cls.__dict__``（仅本类，不含继承），为每个标注方法构造 ``MessageMappingModel``。
+    同时读取方法上的 ``@SendTo`` / ``@SendToUser`` 决定返回值去向。
+    """
+    models: List[MessageMappingModel] = []
+    for name, method in vars(cls).items():
+        if not callable(method):
+            continue
+        annotations = getattr(method, "__spring_annotations__", []) or []
+        # 优先 @MessageMapping
+        msg_mapping = next((a for a in annotations if isinstance(a, MessageMapping)), None)
+        sub_mapping = next((a for a in annotations if isinstance(a, SubscribeMapping)), None)
+        if msg_mapping is None and sub_mapping is None:
+            continue
+        send_to = next((a for a in annotations if isinstance(a, SendTo)), None)
+        send_to_user = next((a for a in annotations if isinstance(a, SendToUser)), None)
+
+        send_to_dests = send_to.destinations if send_to else []
+        is_subscribe = sub_mapping is not None
+        destination = (sub_mapping.value if is_subscribe else msg_mapping.value) or ""
+
+        models.append(MessageMappingModel(
+            method_name=name,
+            destination=destination,
+            send_to=send_to_dests,
+            send_to_user=bool(send_to_user) or (not send_to_dests and not is_subscribe),
+            send_to_user_broadcast=getattr(send_to_user, "broadcast", False) if send_to_user else False,
+            subscribe_destination=destination if is_subscribe else "",
+            is_subscribe=is_subscribe,
+        ))
+    return models
+
+
+__all__ = [
+    "MessageEndpoint",
+    "MessageMapping",
+    "SendTo",
+    "SendToUser",
+    "SubscribeMapping",
+    "MessageMappingModel",
+    "collect_message_mappings",
+]

spring/websocket/broker.py

@@ -0,0 +1,238 @@
+"""内存消息代理 + 发送操作 API（对齐 Spring ``SimpleBrokerMessageHandler`` +
+``SimpMessageSendingOperations``）。
+
+``InMemoryBroker`` 主题式发布订阅代理：
+- ``subscribe(destination, session)``        会话订阅 destination
+- ``unsubscribe(destination, session_id)``   会话取消订阅（按 session_id）
+- ``unsubscribe_all(session_id)``            会话退出时取消其所有订阅
+- ``publish(destination, message, exclude)`` 向 destination 的所有订阅者推送 JSON 消息
+- ``subscribers(destination)``               返回 destination 的订阅会话列表
+
+destination 命名约定（对齐 Spring STOMP）：
+- ``/topic/*``   主题广播（多客户端可订阅）
+- ``/queue/*``   队列（点对点，通常与 ``@SendToUser`` 配合）
+- ``/app/*``     应用消息（``@MessageMapping`` 入口，前缀由 ``MessageBrokerConfigurer`` 剥离）
+
+``SimpMessageSendingOperations`` 提供高阶 API：
+- ``convert_and_send(destination, payload)``  转换 payload 为 JSON 并发布到 destination
+- ``convert_and_send_to_user(user, payload)`` 定向推送给用户
+
+``MessageBrokerConfigurer`` 配置入口：
+- ``application_destination_prefixes``：``@MessageMapping`` 入口前缀（默认 ``["/app"]``）
+- ``broker_prefixes``：broker 处理的前缀（默认 ``["/topic", "/queue"]``）
+- ``user_destination_prefix``：用户私有目的地前缀（默认 ``/user``）
+
+``broker_registry`` 全局单例（``InMemoryBroker`` + ``SimpMessageSendingOperations``）。
+"""
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any, Dict, Iterable, List, Optional, Set
+
+from .exceptions import MessageBrokerException
+from .session import WebSocketSession, global_session_registry
+
+logger = logging.getLogger("Spring.WebSocket.Broker")
+
+
+class InMemoryBroker:
+    """内存主题式消息代理（对齐 Spring ``SimpleBrokerMessageHandler``）。
+
+    线程安全：所有订阅/发布操作加锁；推送是 ``async``，需事件循环驱动。
+    """
+
+    def __init__(self):
+        # destination -> {session_id: WebSocketSession}
+        self._subscriptions: Dict[str, Dict[str, WebSocketSession]] = {}
+        self._lock = threading.RLock()
+
+    def subscribe(self, destination: str, session: WebSocketSession) -> int:
+        """会话订阅 destination；返回该 destination 当前订阅数。"""
+        if not destination:
+            raise MessageBrokerException("destination 不能为空")
+        with self._lock:
+            bucket = self._subscriptions.setdefault(destination, {})
+            bucket[session.id] = session
+            return len(bucket)
+
+    def unsubscribe(self, destination: str, session_id: str) -> bool:
+        """会话取消订阅指定 destination；返回是否成功取消。"""
+        with self._lock:
+            bucket = self._subscriptions.get(destination)
+            if bucket is None:
+                return False
+            return bucket.pop(session_id, None) is not None
+
+    def unsubscribe_all(self, session_id: str) -> int:
+        """会话退出时取消其所有订阅；返回取消的订阅数。"""
+        n = 0
+        with self._lock:
+            for destination, bucket in list(self._subscriptions.items()):
+                if bucket.pop(session_id, None) is not None:
+                    n += 1
+                if not bucket:
+                    self._subscriptions.pop(destination, None)
+        return n
+
+    def subscribers(self, destination: str) -> List[WebSocketSession]:
+        """返回 destination 的订阅会话列表（拷贝）。"""
+        with self._lock:
+            bucket = self._subscriptions.get(destination, {})
+            return list(bucket.values())
+
+    def subscriber_count(self, destination: str) -> int:
+        with self._lock:
+            return len(self._subscriptions.get(destination, {}))
+
+    def destinations(self) -> List[str]:
+        with self._lock:
+            return list(self._subscriptions.keys())
+
+    def clear(self) -> None:
+        with self._lock:
+            self._subscriptions.clear()
+
+    async def publish(self, destination: str, message: Any,
+                      exclude: Optional[Iterable[str]] = None) -> int:
+        """向 destination 的所有订阅者推送 JSON 消息；返回成功推送数。
+
+        - ``exclude``：要排除的 session_id 列表。
+        - 已关闭的会话自动跳过，并在推送后清理其订阅。
+        """
+        if not destination:
+            raise MessageBrokerException("destination 不能为空")
+        excluded: Set[str] = set(exclude or [])
+        sent = 0
+        stale: List[str] = []
+        with self._lock:
+            bucket = self._subscriptions.get(destination, {})
+            targets = [s for sid, s in bucket.items() if sid not in excluded]
+        for session in targets:
+            if not session.is_open:
+                stale.append(session.id)
+                continue
+            try:
+                await session.send_json(_wrap_message(destination, message))
+                sent += 1
+            except Exception as exc:
+                logger.warning("publish to session %s failed: %s", session.id, exc)
+                stale.append(session.id)
+        # 清理失效会话的订阅
+        if stale:
+            with self._lock:
+                bucket = self._subscriptions.get(destination)
+                if bucket is not None:
+                    for sid in stale:
+                        bucket.pop(sid, None)
+                    if not bucket:
+                        self._subscriptions.pop(destination, None)
+        return sent
+
+
+def _wrap_message(destination: str, payload: Any) -> Dict[str, Any]:
+    """构造标准消息帧：``{"destination": ..., "payload": ...}``。"""
+    return {"destination": destination, "payload": payload}
+
+
+# ==================== SimpMessageSendingOperations ====================
+
+class SimpMessageSendingOperations:
+    """高阶消息发送 API（对齐 Spring ``SimpMessageSendingOperations``）。"""
+
+    def __init__(self, broker: InMemoryBroker,
+                 session_registry=global_session_registry):
+        self._broker = broker
+        self._session_registry = session_registry
+
+    @property
+    def broker(self) -> InMemoryBroker:
+        return self._broker
+
+    async def convert_and_send(self, destination: str, payload: Any,
+                               exclude: Optional[Iterable[str]] = None) -> int:
+        """转换 payload 为 JSON 并发布到 destination。"""
+        return await self._broker.publish(destination, payload, exclude=exclude)
+
+    async def convert_and_send_to_user(self, user: str, destination: str, payload: Any) -> int:
+        """定向推送给用户：消息发到该用户的所有会话，destination 作为元数据。
+
+        实现：直接通过 ``session_registry.send_to_user`` 推送，不走 broker 订阅。
+        """
+        message = _wrap_message(destination, payload)
+        return await self._session_registry.send_to_user(user, message, as_json=True)
+
+
+# ==================== MessageBrokerConfigurer ====================
+
+class MessageBrokerConfigurer:
+    """消息代理配置器（对齐 Spring ``@EnableWebSocketMessageBroker`` +
+    ``WebSocketMessageBrokerConfigurer``）。
+
+    配置项：
+    - ``application_destination_prefixes``：``@MessageMapping`` 入口前缀（默认 ``["/app"]``）。
+      客户端发往 ``/app/greet`` 的消息被路由到 ``@MessageMapping("/greet")``。
+    - ``broker_prefixes``：broker 处理的前缀（默认 ``["/topic", "/queue"]``）。
+      ``@SendTo("/topic/x")`` 的消息直接由 broker 广播。
+    - ``user_destination_prefix``：用户私有目的地前缀（默认 ``/user``）。
+    """
+
+    def __init__(
+        self,
+        application_destination_prefixes: Optional[List[str]] = None,
+        broker_prefixes: Optional[List[str]] = None,
+        user_destination_prefix: str = "/user",
+    ):
+        self._app_prefixes: List[str] = list(application_destination_prefixes or ["/app"])
+        self._broker_prefixes: List[str] = list(broker_prefixes or ["/topic", "/queue"])
+        self._user_prefix: str = user_destination_prefix
+        self._broker = InMemoryBroker()
+        self._sending_ops = SimpMessageSendingOperations(self._broker)
+
+    @property
+    def broker(self) -> InMemoryBroker:
+        return self._broker
+
+    @property
+    def sending_operations(self) -> SimpMessageSendingOperations:
+        return self._sending_ops
+
+    @property
+    def application_destination_prefixes(self) -> List[str]:
+        return list(self._app_prefixes)
+
+    @property
+    def broker_prefixes(self) -> List[str]:
+        return list(self._broker_prefixes)
+
+    @property
+    def user_destination_prefix(self) -> str:
+        return self._user_prefix
+
+    def strip_app_prefix(self, destination: str) -> Optional[str]:
+        """剥离 ``/app`` 前缀，返回 ``@MessageMapping`` 匹配路径；非入口返回 None。"""
+        for prefix in self._app_prefixes:
+            if destination == prefix:
+                return ""
+            if destination.startswith(prefix + "/"):
+                return destination[len(prefix):]
+        return None
+
+    def is_broker_destination(self, destination: str) -> bool:
+        """destination 是否由 broker 直接处理（``/topic`` / ``/queue`` 等）。"""
+        for prefix in self._broker_prefixes:
+            if destination == prefix or destination.startswith(prefix + "/"):
+                return True
+        return False
+
+
+# 全局单例（默认配置）
+broker_registry = MessageBrokerConfigurer()
+
+
+__all__ = [
+    "InMemoryBroker",
+    "SimpMessageSendingOperations",
+    "MessageBrokerConfigurer",
+    "broker_registry",
+]

spring/websocket/exceptions.py

@@ -0,0 +1,26 @@
+"""WebSocket 模块异常族（对齐 Spring ``WebSocketException`` 体系）。"""
+from __future__ import annotations
+
+
+class WebSocketException(Exception):
+    """WebSocket 模块根异常（对齐 Spring ``WebSocketException``）。"""
+
+
+class WebSocketConnectionException(WebSocketException):
+    """连接级异常（握手失败、连接断开等）。"""
+
+
+class WebSocketHandlerException(WebSocketException):
+    """处理器执行异常（``on_message`` 抛错、消息路由失败等）。"""
+
+
+class MessageBrokerException(WebSocketException):
+    """消息代理异常（订阅失败、destination 解析失败等）。"""
+
+
+__all__ = [
+    "WebSocketException",
+    "WebSocketConnectionException",
+    "WebSocketHandlerException",
+    "MessageBrokerException",
+]