vention-communication 0.1.0__py3-none-any.whl

communication/connect_router.py

@@ -0,0 +1,325 @@
+from __future__ import annotations
+import asyncio
+import json
+import time
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict
+
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from .entries import ActionEntry, StreamEntry
+from .errors import error_envelope, to_connect_error
+
+logger = logging.getLogger(__name__)
+
+CONTENT_TYPE = "application/connect+json"
+
+
+def _frame(payload: Dict[str, Any], *, trailer: bool = False) -> bytes:
+    """Create a Connect protocol frame for streaming responses.
+
+    Connect uses a binary framing format for streaming JSON over HTTP. Each frame
+    consists of:
+    - 1 byte: Flags (0x00 for data, 0x80 for trailer/end-of-stream)
+    - 4 bytes: Body length in big-endian format
+    - N bytes: JSON-encoded payload
+    """
+    body = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode(
+        "utf-8"
+    )
+    flag = 0x80 if trailer else 0x00
+    header = bytes([flag]) + len(body).to_bytes(4, byteorder="big", signed=False)
+    return header + body
+
+
+# ---------- Subscriber ----------
+
+
+@dataclass(eq=False, unsafe_hash=True)
+class _Subscriber:
+    queue: asyncio.Queue[Any]
+    joined_at: float = field(default_factory=lambda: time.time())
+    last_send_at: float = field(default_factory=lambda: time.time())
+
+
+class StreamManager:
+    """Topic-oriented fan-out with a distributor task per stream.
+
+    Supports configurable delivery policies: "latest" (drops old items when queue is full)
+    or "fifo" (waits for space to ensure all items are delivered).
+    """
+
+    def __init__(self) -> None:
+        """Initialize the StreamManager with an empty topic registry."""
+        self._topics: Dict[str, Dict[str, Any]] = {}
+
+    def ensure_topic(self, entry: StreamEntry) -> None:
+        """Create a topic if it doesn't exist (synchronous, safe before loop)."""
+        if entry.name in self._topics:
+            return
+
+        publish_queue: asyncio.Queue[Any] = asyncio.Queue()
+        subscribers: set[_Subscriber] = set()
+        topic = {
+            "entry": entry,
+            "publish_queue": publish_queue,
+            "subscribers": subscribers,
+            "last_value": None,
+            "task": None,
+        }
+        self._topics[entry.name] = topic
+
+        # If an event loop is running, start distributor immediately
+        try:
+            loop = asyncio.get_running_loop()
+            topic["task"] = loop.create_task(self._distributor(entry.name))
+        except RuntimeError:
+            pass
+
+    def start_distributor_if_needed(self, stream_name: str) -> None:
+        """Start the distributor task if it doesn't exist or has completed.
+
+        Args:
+            stream_name: Name of the stream topic
+        """
+        topic = self._topics[stream_name]
+        if topic["task"] is None or topic["task"].done():
+            try:
+                loop = asyncio.get_running_loop()
+            except RuntimeError:
+                loop = asyncio.get_event_loop()
+            topic["task"] = loop.create_task(self._distributor(stream_name))
+
+    def add_stream(self, entry: StreamEntry) -> None:
+        """Register a stream entry and ensure its topic exists.
+
+        Args:
+            entry: Stream entry to register
+        """
+        self.ensure_topic(entry)
+
+    def publish(self, stream_name: str, item: Any) -> None:
+        """Publish an item to a stream topic.
+
+        Args:
+            stream_name: Name of the stream topic
+            item: Item to publish
+
+        Raises:
+            KeyError: If the stream topic doesn't exist
+        """
+        topic = self._topics.get(stream_name)
+        if topic is None:
+            raise KeyError(f"Unknown stream: {stream_name}")
+        self.start_distributor_if_needed(stream_name)
+        topic["publish_queue"].put_nowait(item)
+
+    def subscribe(self, stream_name: str) -> _Subscriber:
+        """Subscribe to a stream topic and return a subscriber.
+
+        If replay is enabled, the last published value will be enqueued
+        for the new subscriber.
+
+        Args:
+            stream_name: Name of the stream topic
+
+        Returns:
+            A subscriber instance with its own queue
+        """
+        topic = self._topics[stream_name]
+        entry: StreamEntry = topic["entry"]
+        subscriber_queue: asyncio.Queue[Any] = asyncio.Queue(
+            maxsize=entry.queue_maxsize
+        )
+        subscriber = _Subscriber(queue=subscriber_queue)
+        topic["subscribers"].add(subscriber)
+        self.start_distributor_if_needed(stream_name)
+
+        if entry.replay and topic["last_value"] is not None:
+            try:
+                subscriber_queue.put_nowait(topic["last_value"])
+            except asyncio.QueueFull:
+                _ = subscriber_queue.get_nowait()
+                subscriber_queue.put_nowait(topic["last_value"])
+        return subscriber
+
+    def unsubscribe(self, stream_name: str, subscriber: _Subscriber) -> None:
+        """Remove a subscriber from a stream topic.
+
+        Args:
+            stream_name: Name of the stream topic
+            subscriber: Subscriber to remove
+        """
+        topic = self._topics.get(stream_name)
+        if topic:
+            topic["subscribers"].discard(subscriber)
+
+    async def _distributor(self, stream_name: str) -> None:
+        """Distribute published items to all subscribers of a stream.
+
+        For FIFO policy, waits for queue space. For latest-wins policy,
+        drops oldest items when queues are full.
+
+        Args:
+            stream_name: Name of the stream topic to distribute
+        """
+        topic = self._topics[stream_name]
+        publish_queue: asyncio.Queue[Any] = topic["publish_queue"]
+        subscribers: set[_Subscriber] = topic["subscribers"]
+
+        while True:
+            item = await publish_queue.get()
+            topic["last_value"] = item
+            dead_subscribers: list[_Subscriber] = []
+            entry: StreamEntry = topic["entry"]
+
+            for subscriber in list(subscribers):
+                try:
+                    if entry.policy == "fifo":
+                        await subscriber.queue.put(item)
+                    else:
+                        subscriber.queue.put_nowait(item)
+                except asyncio.QueueFull:
+                    try:
+                        _ = subscriber.queue.get_nowait()
+                    except Exception:
+                        pass
+                    try:
+                        subscriber.queue.put_nowait(item)
+                    except Exception:
+                        dead_subscribers.append(subscriber)
+
+            for dead_subscriber in dead_subscribers:
+                subscribers.discard(dead_subscriber)
+
+
+class ConnectRouter:
+    """Router for Connect-style RPC endpoints."""
+
+    def __init__(self) -> None:
+        """Initialize the ConnectRouter with empty registries."""
+        self.router = APIRouter()
+        self._unaries: Dict[str, ActionEntry] = {}
+        self._streams: Dict[str, StreamEntry] = {}
+        self.manager = StreamManager()
+
+    def add_unary(self, entry: ActionEntry, service_fqn: str) -> None:
+        """Register a unary RPC action endpoint.
+
+        Args:
+            entry: Action entry containing the handler function and types
+            service_fqn: Fully qualified service name for the path
+        """
+        self._unaries[entry.name] = entry
+        path = f"/{service_fqn}/{entry.name}"
+
+        @self.router.post(path)
+        async def _handler(request: Request) -> JSONResponse:
+            try:
+                request_body = await request.json()
+            except Exception:
+                request_body = {}
+
+            try:
+                if entry.input_type is None:
+                    result = await _maybe_await(entry.func())
+                else:
+                    input_type = entry.input_type
+                    if hasattr(input_type, "model_validate"):
+                        validated_arg = input_type.model_validate(request_body)
+                    else:
+                        validated_arg = request_body
+                    result = await _maybe_await(entry.func(validated_arg))
+
+                if hasattr(result, "model_dump"):
+                    result = result.model_dump()
+                return JSONResponse(result or {})
+            except Exception as exc:
+                return JSONResponse(error_envelope(exc))
+
+    def add_stream(self, entry: StreamEntry, service_fqn: str) -> None:
+        """Register a streaming RPC endpoint.
+
+        Args:
+            entry: Stream entry containing configuration
+            service_fqn: Fully qualified service name for the path
+        """
+        self._streams[entry.name] = entry
+        self.manager.add_stream(entry)
+
+        path = f"/{service_fqn}/{entry.name}"
+
+        async def _stream_iter(subscriber: _Subscriber) -> Any:
+            """Async generator yielding Connect frames for each stream item."""
+            try:
+                while True:
+                    stream_item = await subscriber.queue.get()
+                    payload = _serialize_stream_item(stream_item)
+                    yield _frame(payload)
+                    await asyncio.sleep(0)
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                error_trailer = error_envelope(to_connect_error(exc))
+                yield _frame(error_trailer, trailer=True)
+            finally:
+                self.manager.unsubscribe(entry.name, subscriber)
+
+        @self.router.post(path)
+        async def _handler(_: Request) -> StreamingResponse:
+            self.manager.start_distributor_if_needed(entry.name)
+            subscriber = self.manager.subscribe(entry.name)
+            logger.debug(f"Stream subscribed: {entry.name}")
+            return StreamingResponse(
+                _stream_iter(subscriber),
+                media_type=CONTENT_TYPE,
+                headers={"Transfer-Encoding": "chunked"},
+            )
+
+    def publish(self, stream_name: str, item: Any) -> None:
+        """Publish an item to a stream.
+
+        Creates the stream topic if it doesn't exist yet.
+
+        Args:
+            stream_name: Name of the stream
+            item: Item to publish
+
+        Raises:
+            KeyError: If the stream doesn't exist and can't be found
+        """
+        if stream_name not in self.manager._topics:
+            entry = self._streams.get(stream_name)
+            if not entry:
+                raise KeyError(f"Unknown stream: {stream_name}")
+            self.manager.add_stream(entry)
+        self.manager.publish(stream_name, item)
+
+
+def _serialize_stream_item(item: Any) -> Dict[str, Any]:
+    if hasattr(item, "model_dump"):
+        dumped = item.model_dump()
+        if isinstance(dumped, dict):
+            return dumped
+        return {"value": dumped}
+    if isinstance(item, dict):
+        return item
+    if isinstance(item, list):
+        return {"value": item}
+    return {"value": item}
+
+
+async def _maybe_await(value: Any) -> Any:
+    """Await a value if it's a coroutine or Future, otherwise return it.
+
+    Args:
+        value: Value that may or may not be awaitable
+
+    Returns:
+        The result of awaiting or the value itself
+    """
+    if asyncio.iscoroutine(value) or isinstance(value, asyncio.Future):
+        return await value
+    return value

communication/decorators.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+from typing import Any, Callable, Literal, Optional, Type, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .app import VentionApp
+    from .entries import RpcBundle
+
+from .entries import ActionEntry, StreamEntry
+from .typing_utils import infer_input_type, infer_output_type, is_pydantic_model
+
+_actions: List[ActionEntry] = []
+_streams: List[StreamEntry] = []
+_GLOBAL_APP: Optional["VentionApp"] = None
+
+
+def set_global_app(app: Any) -> None:
+    """Set the global app instance for use by stream publishers.
+
+    Args:
+        app: VentionApp instance to make globally available
+    """
+    global _GLOBAL_APP
+    _GLOBAL_APP = app
+
+
+def action(
+    name: Optional[str] = None,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Decorator to register a function as an RPC action.
+
+    Args:
+        name: Optional name for the action. If not provided, uses the function name.
+
+    Returns:
+        Decorator function that registers the action
+    """
+
+    def decorator(function: Callable[..., Any]) -> Callable[..., Any]:
+        input_type = infer_input_type(function)
+        output_type = infer_output_type(function)
+        entry = ActionEntry(
+            name or function.__name__, function, input_type, output_type
+        )
+        _actions.append(entry)
+        return function
+
+    return decorator
+
+
+def stream(
+    name: str,
+    *,
+    payload: Type[Any],
+    replay: bool = True,
+    queue_maxsize: int = 1,
+    policy: Literal["latest", "fifo"] = "latest",
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Register a server-broadcast stream.
+
+    The decorated function becomes a publisher that publishes its return value
+    to the stream when called.
+
+    Args:
+        name: Name of the stream
+        payload: Type of the payload (Pydantic model or JSON-serializable type)
+        replay: Whether to replay the last value to new subscribers
+        queue_maxsize: Maximum size of the per-subscriber queue
+        policy: Delivery policy, either "latest" or "fifo"
+
+    Returns:
+        Decorator function that registers the stream
+    """
+    if not (is_pydantic_model(payload) or payload in (int, float, str, bool, dict)):
+        raise ValueError(
+            "payload must be a pydantic BaseModel or a JSON-serializable scalar/dict"
+        )
+
+    def decorator(function: Callable[..., Any]) -> Callable[..., Any]:
+        entry = StreamEntry(
+            name=name,
+            func=None,
+            payload_type=payload,
+            replay=replay,
+            queue_maxsize=queue_maxsize,
+            policy=policy,
+        )
+        _streams.append(entry)
+
+        async def publisher_wrapper(*args: Any, **kwargs: Any) -> Any:
+            if _GLOBAL_APP is None or _GLOBAL_APP.connect_router is None:
+                raise RuntimeError("Stream publish called before app.finalize()")
+            result = await function(*args, **kwargs)
+            _GLOBAL_APP.connect_router.publish(name, result)
+            return None
+
+        entry.func = publisher_wrapper
+        return publisher_wrapper
+
+    return decorator
+
+
+def collect_bundle() -> RpcBundle:
+    """Collect all registered actions and streams into an RpcBundle.
+
+    Returns:
+        RpcBundle containing all actions and streams registered via decorators
+    """
+    from .entries import RpcBundle
+
+    return RpcBundle(actions=list(_actions), streams=list(_streams))

communication/entries.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import Any, Callable, Literal, Optional, Type
+
+
+@dataclass
+class ActionEntry:
+    """Entry for a unary RPC action."""
+
+    name: str
+    func: Callable[..., Any]
+    input_type: Optional[Type[Any]]
+    output_type: Optional[Type[Any]]
+
+
+@dataclass
+class StreamEntry:
+    """Entry for a streaming RPC."""
+
+    name: str
+    func: Optional[Callable[..., Any]]
+    payload_type: Type[Any]
+    replay: bool = True
+    queue_maxsize: int = 1
+    policy: Literal["latest", "fifo"] = "latest"
+
+
+@dataclass
+class RpcBundle:
+    """Bundle of RPC actions and streams."""
+
+    actions: list[ActionEntry] = field(default_factory=list)
+    streams: list[StreamEntry] = field(default_factory=list)
+
+    def extend(self, other: "RpcBundle") -> None:
+        """Extend this bundle with actions and streams from another bundle.
+
+        Args:
+            other: RPC bundle to merge into this one
+        """
+        self.actions.extend(other.actions)
+        self.streams.extend(other.streams)

communication/errors.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+from typing import Any, Dict, List, Optional
+
+
+class ConnectError(Exception):
+    """Application-level error to send over Connect transport."""
+
+    def __init__(
+        self, code: str, message: str, *, details: Optional[List[Any]] = None
+    ) -> None:
+        super().__init__(message)
+        self.code = code
+        self.message = message
+        self.details = details or []
+
+
+def to_connect_error(exception: BaseException) -> ConnectError:
+    """Map arbitrary exceptions to a ConnectError.
+
+    Args:
+        exception: Exception to convert
+
+    Returns:
+        ConnectError with appropriate error code based on exception type
+    """
+    if isinstance(exception, ConnectError):
+        return exception
+
+    import asyncio
+
+    if isinstance(exception, asyncio.TimeoutError):
+        return ConnectError("deadline_exceeded", str(exception) or "Deadline exceeded")
+
+    if isinstance(exception, (KeyError, ValueError)):
+        return ConnectError("invalid_argument", str(exception) or "Invalid argument")
+
+    if isinstance(exception, PermissionError):
+        return ConnectError("permission_denied", str(exception) or "Permission denied")
+
+    return ConnectError("internal", str(exception) or exception.__class__.__name__)
+
+
+def error_envelope(exception: BaseException) -> Dict[str, Any]:
+    """Wrap an exception in a Connect error envelope format.
+
+    Args:
+        exception: Exception to wrap
+
+    Returns:
+        Dictionary with error code, message, and details in Connect format
+    """
+    connect_error = to_connect_error(exception)
+    return {
+        "error": {
+            "code": connect_error.code,
+            "message": connect_error.message,
+            "details": connect_error.details,
+        }
+    }

communication/typing_utils.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+import inspect
+from typing import Any, Optional, Type, get_type_hints, cast
+
+from pydantic import BaseModel
+
+
+class TypingError(Exception):
+    """Raised when type inference fails."""
+
+
+def _strip_self(params: list[inspect.Parameter]) -> list[inspect.Parameter]:
+    if not params:
+        return params
+    first = params[0]
+    if first.name in ("self", "cls"):
+        return params[1:]
+    return params
+
+
+def infer_input_type(function: Any) -> Optional[Type[Any]]:
+    """Infer the input type annotation from a function's first parameter.
+
+    Args:
+        function: Function to inspect
+
+    Returns:
+        Type annotation of the first parameter, or None if no parameters or type is Any
+
+    Raises:
+        TypingError: If the first parameter lacks a type annotation
+    """
+    signature = inspect.signature(function)
+    parameters = _strip_self(list(signature.parameters.values()))
+    if not parameters:
+        return None
+
+    first_param = parameters[0]
+    type_hints = get_type_hints(function)
+    if first_param.name not in type_hints:
+        raise TypingError(f"First parameter '{first_param.name}' must be annotated")
+
+    hint_type = type_hints[first_param.name]
+    if hint_type is Any:
+        return None
+
+    if isinstance(hint_type, type):
+        return cast(Type[Any], hint_type)
+
+    return None
+
+
+def infer_output_type(function: Any) -> Optional[Type[Any]]:
+    """Infer the return type annotation from a function.
+
+    Args:
+        function: Function to inspect
+
+    Returns:
+        Return type annotation, or None if not annotated or type is Any
+    """
+    type_hints = get_type_hints(function)
+    return_type = type_hints.get("return")
+    if return_type is None or return_type is Any:
+        return None
+
+    if isinstance(return_type, type) and return_type in (type(None),):
+        return None
+
+    if isinstance(return_type, type):
+        return cast(Type[Any], return_type)
+
+    return None
+
+
+def is_pydantic_model(type_annotation: Any) -> bool:
+    """Check if a type annotation is a Pydantic BaseModel.
+
+    Args:
+        type_annotation: Type to check
+
+    Returns:
+        True if the type is a Pydantic BaseModel subclass, False otherwise
+    """
+    try:
+        return isinstance(type_annotation, type) and issubclass(
+            type_annotation, BaseModel
+        )
+    except Exception:
+        return False