Add broadcast

Author: apiad

violetear/app.py

@@ -107,13 +107,27 @@ def get_service_worker(scope_hash: str):
                 headers=headers,
             )
 
+        # In App.__init__
+        self.socket_manager = SocketManager()
+
+        @self.api.websocket("/_violetear/ws")
+        async def websocket_endpoint(websocket: WebSocket):
+            await self.socket_manager.connect(websocket)
+            try:
+                while True:
+                    # Keep the connection alive.
+                    # We can also listen for client-to-server messages here if needed later.
+                    await websocket.receive()
+            except WebSocketDisconnect:
+                self.socket_manager.disconnect(websocket)
+
     def client(self, func: Callable):
         """Decorator to mark a function to be compiled to the client."""
         if not inspect.iscoroutinefunction(func):
             raise ValueError("func must be async")
 
         self.client_functions[func.__name__] = func
-        return func
+        return ClientFunctionWrapper(self, func)
 
     def startup(self, func: Callable):
         """
@@ -304,6 +318,25 @@ def _generate_bundle(self) -> str:
             code = [c for c in code if not c.startswith("@")]  # remove decorators
             user_code.append("\n".join(code))
 
+        # --- SAFETY INJECTION START ---
+        # We attach a dummy .broadcast() method to client functions running in the browser.
+        # This prevents confusion if a user tries to call await my_func.broadcast() in client code.
+        safety_checks = []
+        safety_checks.append(
+            dedent(
+                """
+                def _server_only_broadcast(*args, **kwargs):
+                    raise RuntimeError("❌ .broadcast() cannot be called from the Client (Browser).\\nIt must be called from the Server to trigger client updates.")
+                """
+            )
+        )
+
+        for name in self.client_functions.keys():
+            safety_checks.append(f"{name}.broadcast = _server_only_broadcast")
+
+        safety_code = "\n".join(safety_checks)
+        # --- SAFETY INJECTION END ---
+
         # 5. Generate Server Stubs
         server_stubs = self._generate_server_stubs()
 
@@ -321,6 +354,7 @@ def _generate_bundle(self) -> str:
                 storage_injection,
                 runtime_code,
                 "\n\n".join(user_code),
+                safety_code,
                 server_stubs,
                 init_code,
             ]
@@ -537,3 +571,85 @@ def mount_static(self, directory: str, path: str = "/static"):
     def run(self, host="0.0.0.0", port=8000, **kwargs):
         """Helper to run via uvicorn programmatically."""
         uvicorn.run(self.api, host=host, port=port, **kwargs)
+
+
+class ClientFunctionWrapper:
+    """
+    Wraps a client-side function to provide Server-Side RPC capabilities.
+
+    When you define:
+        @app.client
+        async def my_func(...): ...
+
+    This wrapper ensures that:
+    1. Calling `await my_func(...)` on the server runs the function (or warns).
+    2. Calling `await my_func.broadcast(...)` triggers the WebSocket dispatcher.
+    """
+
+    def __init__(self, app: "App", func: Callable):
+        self.app = app
+        self.func = func
+        # Mimic the original function's identity
+        self.__name__ = func.__name__
+        self.__doc__ = func.__doc__
+
+    def __call__(self, *args, **kwargs):
+        """
+        Standard call: await my_func(...)
+        Executes the function logic locally (useful for testing or shared logic).
+        """
+        raise RuntimeError(
+            f"Cannot call client-side functions in the server! Did you meant {self.func.__name__}.broadcast(...)?"
+        )
+
+    async def broadcast(self, *args, **kwargs):
+        """
+        RPC Call: await my_func.broadcast(...)
+
+        Tells the server to instructing ALL connected clients to run this function.
+        """
+        # FUTURE: This will hook into the WebSocket manager
+        if hasattr(self.app, "socket_manager"):
+            # print(f"Broadcasting {self.__name__} to all clients...")
+            await self.app.socket_manager.broadcast(
+                func_name=self.__name__, args=args, kwargs=kwargs
+            )
+        else:
+            print(
+                f"[Violetear] Warning: Broadcast called on '{self.__name__}' but no SocketManager is active."
+            )
+
+
+# Add this class to violetear/app.py
+
+from fastapi import WebSocket, WebSocketDisconnect
+
+
+class SocketManager:
+    def __init__(self):
+        # Keep track of active connections
+        self.active_connections: List[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket):
+        self.active_connections.remove(websocket)
+
+    async def broadcast(self, func_name: str, args: tuple, kwargs: dict):
+        """
+        Sends a command to all connected clients to run a specific function.
+        """
+        payload = json.dumps(
+            {"type": "rpc", "func": func_name, "args": args, "kwargs": kwargs}
+        )
+
+        # Iterate over all connections and send the message
+        # We use a copy of the list to avoid modification errors during iteration
+        for connection in self.active_connections[:]:
+            try:
+                await connection.send_text(payload)
+            except Exception:
+                # If sending fails (e.g. client disconnected), remove it
+                self.disconnect(connection)

violetear/client.py

@@ -4,6 +4,8 @@
 """
 
 import sys
+import json
+import asyncio
 
 # We define IS_BROWSER to avoid import errors if this is imported on the server
 IS_BROWSER = "pyodide" in sys.modules or "emscripten" in sys.platform
@@ -72,3 +74,43 @@ def hydrate(namespace: dict):
                 )
 
     print(f"[Violetear] Hydrated {bound_count} interactive elements.")
+    setup_socket_listener(namespace)
+
+
+def setup_socket_listener(namespace):
+    """
+    Connects to the server and listens for RPC commands.
+    """
+    from js import WebSocket, window
+
+    # Calculate the WebSocket URL (ws:// or wss://)
+    protocol = "wss" if window.location.protocol == "https:" else "ws"
+    ws_url = f"{protocol}://{window.location.host}/_violetear/ws"
+
+    socket = WebSocket.new(ws_url)
+
+    def on_message(event):
+        data = json.loads(event.data)
+
+        if data.get("type") == "rpc":
+            func_name = data["func"]
+            args = data["args"]
+            kwargs = data["kwargs"]
+
+            # 1. Look up the function in the global scope
+            if func_name in namespace:
+                func = namespace[func_name]
+
+                # 2. Schedule the async function to run on the event loop
+                # We use asyncio.create_task because we are inside a sync callback
+                asyncio.create_task(func(*args, **kwargs))
+            else:
+                print(f"[Violetear] Received RPC for unknown function: {func_name}")
+
+    # Attach the callback (converting Python function to JS proxy not strictly needed for simple events in recent Pyodide, but good practice)
+    socket.onmessage = on_message
+
+    # Keep a reference so it doesn't get garbage collected
+    window.violetear_socket = socket
+
+    print("[Violetear] Attached socket connection")