pulse-framework 0.1.40__py3-none-any.whl → 0.1.42__py3-none-any.whl

pulse/middleware.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
-from collections.abc import Callable, Sequence
+import asyncio
+from collections.abc import Awaitable, Callable, Sequence
 from typing import Any, Generic, TypeVar, overload, override
 
+from pulse.env import env
 from pulse.messages import (
 	ClientMessage,
+	Prerender,
 	PrerenderPayload,
-	PrerenderResult,
 	ServerInitMessage,
 )
 from pulse.request import PulseRequest
@@ -26,20 +28,21 @@ class NotFound: ...
 
 
 class Ok(Generic[T]):
-	payload: T | None
+	payload: T
 
 	@overload
 	def __init__(self, payload: T) -> None: ...
 	@overload
-	def __init__(self, payload: T | None = None) -> None: ...
+	def __init__(self, payload: None = None) -> None: ...
 	def __init__(self, payload: T | None = None) -> None:
-		self.payload = payload
+		self.payload = payload  # pyright: ignore[reportAttributeAccessIssue]
 
 
 class Deny: ...
 
 
-PrerenderResponse = Ok[ServerInitMessage] | Redirect | NotFound
+RoutePrerenderResponse = Ok[ServerInitMessage] | Redirect | NotFound
+PrerenderResponse = Ok[Prerender] | Redirect | NotFound
 ConnectResponse = Ok[None] | Deny
 
 
@@ -50,56 +53,65 @@ class PulseMiddleware:
 	for later use. Return a decision to allow or short-circuit the flow.
 	"""
 
-	def prerender(
+	dev: bool
+
+	def __init__(self, dev: bool = False) -> None:
+		"""Initialize middleware.
+
+		Args:
+			dev: If True, this middleware is only active in dev environments.
+		"""
+		self.dev = dev
+
+	async def prerender(
 		self,
 		*,
 		payload: "PrerenderPayload",
-		result: "PrerenderResult",
 		request: PulseRequest,
 		session: dict[str, Any],
-		next: Callable[[], "PrerenderResult"],
-	) -> "PrerenderResult":
+		next: Callable[[], Awaitable[PrerenderResponse]],
+	) -> PrerenderResponse:
 		"""Handle batch prerender at the top level.
 
-		Receives the full PrerenderPayload and can modify the PrerenderResult
-		(views and directives) before it's returned to the client.
+		Receives the full PrerenderPayload. Call next() to get the PrerenderResult
+		and can modify it (views and directives) before returning to the client.
 		"""
-		return next()
+		return await next()
 
-	def prerender_route(
+	async def prerender_route(
 		self,
 		*,
 		path: str,
 		request: PulseRequest,
 		route_info: RouteInfo,
 		session: dict[str, Any],
-		next: Callable[[], PrerenderResponse],
-	) -> PrerenderResponse:
-		return next()
+		next: Callable[[], Awaitable[RoutePrerenderResponse]],
+	) -> RoutePrerenderResponse:
+		return await next()
 
-	def connect(
+	async def connect(
 		self,
 		*,
 		request: PulseRequest,
 		session: dict[str, Any],
-		next: Callable[[], ConnectResponse],
+		next: Callable[[], Awaitable[ConnectResponse]],
 	) -> ConnectResponse:
-		return next()
+		return await next()
 
-	def message(
+	async def message(
 		self,
 		*,
 		data: ClientMessage,
 		session: dict[str, Any],
-		next: Callable[[], Ok[None]],
+		next: Callable[[], Awaitable[Ok[None]]],
 	) -> Ok[None] | Deny:
 		"""Handle per-message authorization.
 
 		Return Deny() to block, Ok(None) to allow.
 		"""
-		return next()
+		return await next()
 
-	def channel(
+	async def channel(
 		self,
 		*,
 		channel_id: str,
@@ -107,9 +119,9 @@ class PulseMiddleware:
 		payload: Any,
 		request_id: str | None,
 		session: dict[str, Any],
-		next: Callable[[], Ok[None]],
+		next: Callable[[], Awaitable[Ok[None]]],
 	) -> Ok[None] | Deny:
-		return next()
+		return await next()
 
 
 class MiddlewareStack(PulseMiddleware):
@@ -119,56 +131,58 @@ class MiddlewareStack(PulseMiddleware):
 	middleware returns without calling `next`, the chain short-circuits.
 	"""
 
-	def __init__(self, middlewares: Sequence[PulseMiddleware]):
+	def __init__(self, middlewares: Sequence[PulseMiddleware]) -> None:
+		super().__init__(dev=False)
+		# Filter out dev middlewares when not in dev environment
+		if env.pulse_env != "dev":
+			middlewares = [mw for mw in middlewares if not mw.dev]
 		self._middlewares: list[PulseMiddleware] = list(middlewares)
 
 	@override
-	def prerender(
+	async def prerender(
 		self,
 		*,
 		payload: "PrerenderPayload",
-		result: "PrerenderResult",
 		request: PulseRequest,
 		session: dict[str, Any],
-		next: Callable[[], "PrerenderResult"],
-	) -> "PrerenderResult":
-		def dispatch(index: int) -> "PrerenderResult":
+		next: Callable[[], Awaitable[PrerenderResponse]],
+	) -> PrerenderResponse:
+		async def dispatch(index: int) -> PrerenderResponse:
 			if index >= len(self._middlewares):
-				return next()
+				return await next()
 			mw = self._middlewares[index]
 
-			def _next() -> "PrerenderResult":
-				return dispatch(index + 1)
+			async def _next() -> PrerenderResponse:
+				return await dispatch(index + 1)
 
-			return mw.prerender(
+			return await mw.prerender(
 				payload=payload,
-				result=result,
 				request=request,
 				session=session,
 				next=_next,
 			)
 
-		return dispatch(0)
+		return await dispatch(0)
 
 	@override
-	def prerender_route(
+	async def prerender_route(
 		self,
 		*,
 		path: str,
 		request: PulseRequest,
 		route_info: RouteInfo,
 		session: dict[str, Any],
-		next: Callable[[], PrerenderResponse],
-	) -> PrerenderResponse:
-		def dispatch(index: int) -> PrerenderResponse:
+		next: Callable[[], Awaitable[RoutePrerenderResponse]],
+	) -> RoutePrerenderResponse:
+		async def dispatch(index: int) -> RoutePrerenderResponse:
 			if index >= len(self._middlewares):
-				return next()
+				return await next()
 			mw = self._middlewares[index]
 
-			def _next() -> PrerenderResponse:
-				return dispatch(index + 1)
+			async def _next() -> RoutePrerenderResponse:
+				return await dispatch(index + 1)
 
-			return mw.prerender_route(
+			return await mw.prerender_route(
 				path=path,
 				route_info=route_info,
 				request=request,
@@ -176,50 +190,56 @@ class MiddlewareStack(PulseMiddleware):
 				next=_next,
 			)
 
-		return dispatch(0)
+		return await dispatch(0)
 
 	@override
-	def connect(
+	async def connect(
 		self,
 		*,
 		request: PulseRequest,
 		session: dict[str, Any],
-		next: Callable[[], ConnectResponse],
+		next: Callable[[], Awaitable[ConnectResponse]],
 	) -> ConnectResponse:
-		def dispatch(index: int) -> ConnectResponse:
+		async def dispatch(index: int) -> ConnectResponse:
 			if index >= len(self._middlewares):
-				return next()
+				return await next()
 			mw = self._middlewares[index]
 
-			def _next() -> ConnectResponse:
-				return dispatch(index + 1)
+			async def _next() -> ConnectResponse:
+				return await dispatch(index + 1)
 
-			return mw.connect(request=request, session=session, next=_next)
+			return await mw.connect(request=request, session=session, next=_next)
 
-		return dispatch(0)
+		return await dispatch(0)
 
 	@override
-	def message(
+	async def message(
 		self,
 		*,
 		data: ClientMessage,
 		session: dict[str, Any],
-		next: Callable[[], Ok[None]],
+		next: Callable[[], Awaitable[Ok[None]]],
 	) -> Ok[None] | Deny:
-		def dispatch(index: int) -> Ok[None] | Deny:
+		async def dispatch(index: int) -> Ok[None] | Deny:
 			if index >= len(self._middlewares):
-				return next()
+				return await next()
 			mw = self._middlewares[index]
 
-			def _next() -> Ok[None]:
-				return dispatch(index + 1)  # pyright: ignore[reportReturnType]
+			async def _next() -> Ok[None]:
+				result = await dispatch(index + 1)
+				# If dispatch returns Deny, the middleware should have short-circuited
+				# This should only be called when continuing the chain
+				if isinstance(result, Deny):
+					# This shouldn't happen, but handle it gracefully
+					return Ok(None)
+				return result
 
-			return mw.message(session=session, data=data, next=_next)
+			return await mw.message(session=session, data=data, next=_next)
 
-		return dispatch(0)
+		return await dispatch(0)
 
 	@override
-	def channel(
+	async def channel(
 		self,
 		*,
 		channel_id: str,
@@ -227,17 +247,23 @@ class MiddlewareStack(PulseMiddleware):
 		payload: Any,
 		request_id: str | None,
 		session: dict[str, Any],
-		next: Callable[[], Ok[None]],
+		next: Callable[[], Awaitable[Ok[None]]],
 	) -> Ok[None] | Deny:
-		def dispatch(index: int) -> Ok[None] | Deny:
+		async def dispatch(index: int) -> Ok[None] | Deny:
 			if index >= len(self._middlewares):
-				return next()
+				return await next()
 			mw = self._middlewares[index]
 
-			def _next() -> Ok[None]:
-				return dispatch(index + 1)  # pyright: ignore[reportReturnType]
+			async def _next() -> Ok[None]:
+				result = await dispatch(index + 1)
+				# If dispatch returns Deny, the middleware should have short-circuited
+				# This should only be called when continuing the chain
+				if isinstance(result, Deny):
+					# This shouldn't happen, but handle it gracefully
+					return Ok(None)
+				return result
 
-			return mw.channel(
+			return await mw.channel(
 				channel_id=channel_id,
 				event=event,
 				payload=payload,
@@ -246,7 +272,7 @@ class MiddlewareStack(PulseMiddleware):
 				next=_next,
 			)
 
-		return dispatch(0)
+		return await dispatch(0)
 
 
 def stack(*middlewares: PulseMiddleware) -> PulseMiddleware:
@@ -258,84 +284,108 @@ def stack(*middlewares: PulseMiddleware) -> PulseMiddleware:
 	return MiddlewareStack(list(middlewares))
 
 
-class PulseCoreMiddleware(PulseMiddleware):
-	"""Core middleware that ensures a PulseContext is mounted around the chain.
+class LatencyMiddleware(PulseMiddleware):
+	"""Middleware that adds artificial latency to simulate network conditions.
 
-	It executes first to set up the context, then lets subsequent middlewares
-	run, and finally returns their response unchanged.
+	Useful for testing and development to simulate real-world network delays.
+	Defaults are realistic for typical web applications.
+
+	Example:
+	    ```python
+	    app = ps.App(
+	        middleware=ps.LatencyMiddleware(
+	            prerender_ms=100,
+	            connect_ms=50,
+	        )
+	    )
+	    ```
 	"""
 
+	prerender_ms: float
+	prerender_route_ms: float
+	connect_ms: float
+	message_ms: float
+	channel_ms: float
+
+	def __init__(
+		self,
+		*,
+		prerender_ms: float = 80.0,
+		prerender_route_ms: float = 60.0,
+		connect_ms: float = 40.0,
+		message_ms: float = 25.0,
+		channel_ms: float = 20.0,
+	) -> None:
+		"""Initialize latency middleware.
+
+		Args:
+			prerender_ms: Latency for batch prerender requests (HTTP). Default: 80ms
+			prerender_route_ms: Latency for individual route prerenders. Default: 60ms
+			connect_ms: Latency for WebSocket connections. Default: 40ms
+			message_ms: Latency for WebSocket messages (including API calls). Default: 25ms
+			channel_ms: Latency for channel messages. Default: 20ms
+			dev: If True, only active in dev environments. Default: True
+		"""
+		super().__init__(dev=True)
+		self.prerender_ms = prerender_ms
+		self.prerender_route_ms = prerender_route_ms
+		self.connect_ms = connect_ms
+		self.message_ms = message_ms
+		self.channel_ms = channel_ms
+
 	@override
-	def prerender(
+	async def prerender(
 		self,
 		*,
 		payload: "PrerenderPayload",
-		result: "PrerenderResult",
 		request: PulseRequest,
 		session: dict[str, Any],
-		next: Callable[[], "PrerenderResult"],
-	) -> "PrerenderResult":
-		res = next()
-		# Return the result as-is (no normalization needed)
-		return res
-
-	# --- Normalization helpers -------------------------------------------------
-	def _normalize_prerender_response(self, res: Any) -> PrerenderResponse:
-		if isinstance(res, (Ok, Redirect, NotFound)):
-			return res  # type: ignore[return-value]
-		# Treat any other value as a VDOM payload
-		return Ok(res)
-
-	def _normalize_connect_response(self, res: Any) -> ConnectResponse:
-		if isinstance(res, (Ok, Deny)):
-			return res  # type: ignore[return-value]
-		# Treat any other value as allow
-		return Ok(None)
-
-	def _normalize_message_response(self, res: Any) -> Ok[None] | Deny:
-		if isinstance(res, (Ok, Deny)):
-			return res  # type: ignore[return-value]
-		# Treat any other value as allow
-		return Ok(None)
+		next: Callable[[], Awaitable[PrerenderResponse]],
+	) -> PrerenderResponse:
+		if self.prerender_ms > 0:
+			await asyncio.sleep(self.prerender_ms / 1000.0)
+		return await next()
 
 	@override
-	def prerender_route(
+	async def prerender_route(
 		self,
 		*,
 		path: str,
 		request: PulseRequest,
 		route_info: RouteInfo,
 		session: dict[str, Any],
-		next: Callable[[], PrerenderResponse],
-	) -> PrerenderResponse:
-		# No render object is available during prerender middleware
-		res = next()
-		return self._normalize_prerender_response(res)
+		next: Callable[[], Awaitable[RoutePrerenderResponse]],
+	) -> RoutePrerenderResponse:
+		if self.prerender_route_ms > 0:
+			await asyncio.sleep(self.prerender_route_ms / 1000.0)
+		return await next()
 
 	@override
-	def connect(
+	async def connect(
 		self,
 		*,
 		request: PulseRequest,
 		session: dict[str, Any],
-		next: Callable[[], ConnectResponse],
+		next: Callable[[], Awaitable[ConnectResponse]],
 	) -> ConnectResponse:
-		res = next()
-		return self._normalize_connect_response(res)
+		if self.connect_ms > 0:
+			await asyncio.sleep(self.connect_ms / 1000.0)
+		return await next()
 
 	@override
-	def message(
+	async def message(
 		self,
 		*,
 		data: ClientMessage,
 		session: dict[str, Any],
-		next: Callable[[], Ok[None]],
+		next: Callable[[], Awaitable[Ok[None]]],
 	) -> Ok[None] | Deny:
-		res = next()
-		return self._normalize_message_response(res)
+		if self.message_ms > 0:
+			await asyncio.sleep(self.message_ms / 1000.0)
+		return await next()
 
 	@override
-	def channel(
+	async def channel(
 		self,
 		*,
 		channel_id: str,
@@ -343,7 +393,8 @@ class PulseCoreMiddleware(PulseMiddleware):
 		payload: Any,
 		request_id: str | None,
 		session: dict[str, Any],
-		next: Callable[[], Ok[None]],
+		next: Callable[[], Awaitable[Ok[None]]],
 	) -> Ok[None] | Deny:
-		res = next()
-		return self._normalize_message_response(res)
+		if self.channel_ms > 0:
+			await asyncio.sleep(self.channel_ms / 1000.0)
+		return await next()

pulse/plugin.py

@@ -16,9 +16,6 @@ class Plugin:
 	def routes(self) -> list[Route | Layout]:
 		return []
 
-	def dev_routes(self) -> list[Route | Layout]:
-		return []
-
 	def middleware(self) -> list[PulseMiddleware]:
 		return []
 

pulse/proxy.py

@@ -2,32 +2,36 @@
 Proxy handler for forwarding requests to React Router server in single-server mode.
 """
 
+import asyncio
 import logging
-from typing import Callable
+from typing import cast
 
 import httpx
+import websockets
 from fastapi.responses import StreamingResponse
 from starlette.background import BackgroundTask
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse, Response
+from starlette.websockets import WebSocket, WebSocketDisconnect
+from websockets.typing import Subprotocol
 
 logger = logging.getLogger(__name__)
 
 
-class ReactProxyHandler:
+class ReactProxy:
 	"""
-	Handles proxying HTTP requests to React Router server.
+	Handles proxying HTTP requests and WebSocket connections to React Router server.
 	"""
 
-	get_react_server_address: Callable[[], str | None]
+	react_server_address: str
 	_client: httpx.AsyncClient | None
 
-	def __init__(self, get_react_server_address: Callable[[], str | None]):
+	def __init__(self, react_server_address: str):
 		"""
 		Args:
-		    get_react_server_address: Callable that returns the React Router server full URL (or None if not started)
+		    react_server_address: React Router server full URL (required in single-server mode)
 		"""
-		self.get_react_server_address = get_react_server_address
+		self.react_server_address = react_server_address
 		self._client = None
 
 	@property
@@ -40,20 +44,134 @@ class ReactProxyHandler:
 			)
 		return self._client
 
-	async def __call__(self, request: Request) -> Response:
+	def _is_websocket_upgrade(self, request: Request) -> bool:
+		"""Check if request is a WebSocket upgrade."""
+		upgrade = request.headers.get("upgrade", "").lower()
+		connection = request.headers.get("connection", "").lower()
+		return upgrade == "websocket" and "upgrade" in connection
+
+	def _http_to_ws_url(self, http_url: str) -> str:
+		"""Convert HTTP URL to WebSocket URL."""
+		if http_url.startswith("https://"):
+			return http_url.replace("https://", "wss://", 1)
+		elif http_url.startswith("http://"):
+			return http_url.replace("http://", "ws://", 1)
+		return http_url
+
+	async def proxy_websocket(self, websocket: WebSocket) -> None:
 		"""
-		Forward HTTP request to React Router server and stream response back.
+		Proxy WebSocket connection to React Router server.
+		Only allowed in dev mode and on root path "/".
 		"""
-		# Get the React server address
-		react_server_address = self.get_react_server_address()
-		if react_server_address is None:
-			# React server not started yet, return error
-			return PlainTextResponse(
-				"Service Unavailable: React server not ready", status_code=503
+
+		# Build target WebSocket URL
+		ws_url = self._http_to_ws_url(self.react_server_address)
+		target_url = ws_url.rstrip("/") + websocket.url.path
+		if websocket.url.query:
+			target_url += "?" + websocket.url.query
+
+		# Extract subprotocols from client request
+		subprotocol_header = websocket.headers.get("sec-websocket-protocol")
+		subprotocols: list[Subprotocol] | None = None
+		if subprotocol_header:
+			# Parse comma-separated list of subprotocols
+			# Subprotocol is a NewType (just a type annotation), so cast strings to it
+			subprotocols = cast(
+				list[Subprotocol], [p.strip() for p in subprotocol_header.split(",")]
 			)
 
+		# Extract headers for WebSocket connection (excluding WebSocket-specific headers)
+		headers = {
+			k: v
+			for k, v in websocket.headers.items()
+			if k.lower()
+			not in (
+				"host",
+				"upgrade",
+				"connection",
+				"sec-websocket-key",
+				"sec-websocket-version",
+				"sec-websocket-protocol",
+			)
+		}
+
+		# Accept the client WebSocket connection first
+		# We'll accept without subprotocol initially, then update if target accepts one
+		await websocket.accept()
+
+		# Connect to target WebSocket server
+		try:
+			async with websockets.connect(
+				target_url,
+				additional_headers=headers,
+				subprotocols=subprotocols,
+				ping_interval=None,  # Let the target server handle ping/pong
+			) as target_ws:
+				# Forward messages bidirectionally
+				async def forward_client_to_target():
+					try:
+						async for message in websocket.iter_text():
+							await target_ws.send(message)
+					except (WebSocketDisconnect, websockets.ConnectionClosed):
+						# Client disconnected, close target connection
+						logger.debug("Client disconnected, closing target connection")
+						try:
+							await target_ws.close()
+						except Exception:
+							pass
+					except Exception as e:
+						logger.error(f"Error forwarding client message: {e}")
+						raise
+
+				async def forward_target_to_client():
+					try:
+						async for message in target_ws:
+							if isinstance(message, str):
+								await websocket.send_text(message)
+							else:
+								await websocket.send_bytes(message)
+					except (WebSocketDisconnect, websockets.ConnectionClosed) as e:
+						# Client or target disconnected, stop forwarding
+						logger.debug(
+							"Connection closed, stopping forward_target_to_client"
+						)
+						# If target disconnected, close client connection
+						if isinstance(e, websockets.ConnectionClosed):
+							try:
+								await websocket.close()
+							except Exception:
+								pass
+					except Exception as e:
+						logger.error(f"Error forwarding target message: {e}")
+						raise
+
+				# Run both forwarding tasks concurrently
+				# If one side closes, the other will detect it and stop gracefully
+				await asyncio.gather(
+					forward_client_to_target(),
+					forward_target_to_client(),
+					return_exceptions=True,
+				)
+
+		except (websockets.WebSocketException, websockets.ConnectionClosedError) as e:
+			logger.error(f"WebSocket proxy connection failed: {e}")
+			await websocket.close(
+				code=1014,  # Bad Gateway
+				reason="Bad Gateway: Could not connect to React Router server",
+			)
+		except Exception as e:
+			logger.error(f"WebSocket proxy error: {e}")
+			await websocket.close(
+				code=1011,  # Internal Server Error
+				reason="Bad Gateway: Proxy error",
+			)
+
+	async def __call__(self, request: Request) -> Response:
+		"""
+		Forward HTTP request to React Router server and stream response back.
+		"""
 		# Build target URL
-		url = react_server_address.rstrip("/") + request.url.path
+		url = self.react_server_address.rstrip("/") + request.url.path
 		if request.url.query:
 			url += "?" + request.url.query