Merge pull request #44 from keyboardstaff/ws-rework

Ws rework

Author: keyboardstaff

AGENTS.md

@@ -71,12 +71,10 @@ When running in Docker, Agent Zero uses two distinct Python runtimes to isolate
 ├── initialize.py         # Framework initialization logic
 ├── models.py             # LLM provider configurations
 ├── run_ui.py             # WebUI server entry point
-├── python/
-│   ├── api/              # API Handlers (ApiHandler subclasses)
-│   ├── extensions/       # Backend lifecycle extensions
-│   ├── helpers/          # Shared Python utilities (plugins, files, etc.)
-│   ├── tools/            # Agent tools (Tool subclasses)
-│   └── websocket_handlers/# WebSocket event handlers
+├── api/                  # API Handlers (ApiHandler subclasses) + WsHandler subclasses (ws_*.py)
+├── extensions/           # Backend lifecycle extensions
+├── helpers/              # Shared Python utilities (plugins, files, etc.)
+├── tools/                # Agent tools (Tool subclasses)
 ├── webui/
 │   ├── components/       # Alpine.js components
 │   ├── js/               # Core frontend logic (modals, stores, etc.)

api/image_get.py

@@ -1,5 +1,6 @@
 import base64
 import os
+from urllib.parse import quote
 from helpers.api import ApiHandler, Request, Response, send_file
 from helpers import files, runtime
 import io
@@ -47,24 +48,31 @@ async def process(self, input: dict, request: Request) -> dict | Response:
 
             # in development environment, try to serve the image from local file system if exists, otherwise from docker
             if runtime.is_development():
-                if files.exists(path):
-                    response = send_file(path)
-                elif await runtime.call_development_function(files.exists, path):
-                    b64_content = await runtime.call_development_function(
-                        files.read_file_base64, path
-                    )
-                    file_content = base64.b64decode(b64_content)
-                    mime_type, _ = guess_type(filename)
-                    if not mime_type:
-                        mime_type = "application/octet-stream"
-                    response = send_file(
-                        io.BytesIO(file_content),
-                        mimetype=mime_type,
-                        as_attachment=False,
-                        download_name=filename,
-                    )
+                # Convert /a0/... Docker paths to local absolute paths
+                local_path = files.fix_dev_path(path)
+                if files.exists(local_path):
+                    response = send_file(local_path)
                 else:
-                    response = _send_fallback_icon("image")
+                    # Try fetching from Docker via RFC as fallback
+                    try:
+                        if await runtime.call_development_function(files.exists, path):
+                            b64_content = await runtime.call_development_function(
+                                files.read_file_base64, path
+                            )
+                            file_content = base64.b64decode(b64_content)
+                            mime_type, _ = guess_type(filename)
+                            if not mime_type:
+                                mime_type = "application/octet-stream"
+                            response = send_file(
+                                io.BytesIO(file_content),
+                                mimetype=mime_type,
+                                as_attachment=False,
+                                download_name=filename,
+                            )
+                        else:
+                            response = _send_fallback_icon("image")
+                    except Exception:
+                        response = _send_fallback_icon("image")
             else:
                 if files.exists(path):
                     response = send_file(path)
@@ -74,7 +82,7 @@ async def process(self, input: dict, request: Request) -> dict | Response:
             # Add cache headers for better device sync performance
             response.headers["Cache-Control"] = "public, max-age=3600"
             response.headers["X-File-Type"] = "image"
-            response.headers["X-File-Name"] = filename
+            response.headers["X-File-Name"] = quote(filename)
             return response
         else:
             # Handle non-image files with fallback icons
@@ -135,7 +143,7 @@ def _send_file_type_icon(file_ext, filename=None):
         response.headers["X-File-Type"] = "icon"
         response.headers["X-Icon-Type"] = icon_name
         if filename:
-            response.headers["X-File-Name"] = filename
+            response.headers["X-File-Name"] = quote(filename)
 
     return response
 

api/ws_dev_test.py

@@ -0,0 +1,77 @@
+import asyncio
+from typing import Any
+
+from helpers.ws import WsHandler
+from helpers.ws_manager import WsResult
+from helpers.print_style import PrintStyle
+from helpers import runtime
+
+
+class WsDevTest(WsHandler):
+    """Developer-only WebSocket test harness handler."""
+
+    async def process(self, event: str, data: dict, sid: str) -> dict[str, Any] | WsResult | None:
+        if event == "ws_event_console_subscribe":
+            if not runtime.is_development():
+                return WsResult.error(
+                    code="NOT_AVAILABLE",
+                    message="Event console is available only in development mode",
+                )
+            registered = self.manager.register_diagnostic_watcher(self.namespace, sid)
+            if not registered:
+                return WsResult.error(
+                    code="SUBSCRIBE_FAILED",
+                    message="Unable to subscribe to diagnostics",
+                )
+            return {"status": "subscribed", "timestamp": data.get("requestedAt")}
+
+        if event == "ws_event_console_unsubscribe":
+            self.manager.unregister_diagnostic_watcher(self.namespace, sid)
+            return {"status": "unsubscribed"}
+
+        if event == "ws_tester_emit":
+            message = data.get("message", "emit")
+            payload = {"message": message, "echo": True, "timestamp": data.get("timestamp")}
+            await self.broadcast("ws_tester_broadcast", payload)
+            PrintStyle.info(f"Harness emit broadcasted message='{message}'")
+            return None
+
+        if event == "ws_tester_request":
+            value = data.get("value")
+            PrintStyle.debug("Harness request responded with echo %s", value)
+            return {"echo": value, "handler": self.identifier, "status": "ok"}
+
+        if event == "ws_tester_request_delayed":
+            delay_ms = int(data.get("delay_ms", 0))
+            await asyncio.sleep(delay_ms / 1000)
+            PrintStyle.warning("Harness delayed request finished after %s ms", delay_ms)
+            return {"status": "delayed", "delay_ms": delay_ms, "handler": self.identifier}
+
+        if event == "ws_tester_trigger_persistence":
+            phase = data.get("phase", "unknown")
+            payload = {"phase": phase, "handler": self.identifier}
+            await self.emit_to(sid, "ws_tester_persistence", payload)
+            PrintStyle.info(f"Harness persistence event phase='{phase}' -> {sid}")
+            return None
+
+        if event == "ws_tester_broadcast_demo_trigger":
+            payload = {"demo": True, "requested_at": data.get("requested_at")}
+            await self.broadcast("ws_tester_broadcast_demo", payload)
+            PrintStyle.info("Harness broadcast demo event dispatched")
+            return None
+
+        if event == "ws_tester_request_all":
+            correlation_id = data.get("correlationId")
+            aggregated = await self.dispatch_to_all_sids(
+                "ws_tester_request",
+                {"value": data.get("marker", "aggregate")},
+                correlation_id=correlation_id,
+            )
+            return {"results": aggregated}
+
+        # Ignore events not targeted at this handler (other activated handlers
+        # may process them).  Only warn for events that look like dev-harness
+        # traffic so we don't spam logs with unrelated events.
+        if event.startswith("ws_tester_"):
+            PrintStyle.warning(f"Harness received unknown event '{event}'")
+        return None

api/ws_hello.py

@@ -0,0 +1,13 @@
+from helpers.ws import WsHandler
+from helpers.print_style import PrintStyle
+
+
+class WsHello(WsHandler):
+    """Simple echo handler used for foundational testing."""
+
+    async def process(self, event: str, data: dict, sid: str) -> dict | None:
+        if event != "hello_request":
+            return None
+        name = data.get("name") or "stranger"
+        PrintStyle.info(f"hello_request from {sid} ({name})")
+        return {"message": f"Hello, {name}!", "handler": self.identifier}

…thon/websocket_handlers/webui_handler.py api/ws_webui.pypython/websocket_handlers/webui_handler.py renamed to api/ws_webui.py python/websocket_handlers/webui_handler.py renamed to api/ws_webui.py

@@ -1,8 +1,10 @@
-from helpers.websocket import WebSocketHandler, WebSocketResult
+from helpers.ws import WsHandler
 from helpers import extension
 
 
-class WebuiHandler(WebSocketHandler):
+class WsWebui(WsHandler):
+    """State synchronisation handler — the primary WebSocket endpoint for the UI."""
+
     async def on_connect(self, sid: str) -> None:
         await extension.call_extensions_async(
             "webui_ws_connect", agent=None, instance=self, sid=sid
@@ -13,22 +15,18 @@ async def on_disconnect(self, sid: str) -> None:
             "webui_ws_disconnect", agent=None, instance=self, sid=sid
         )
 
-    async def process_event(
-        self, event_type: str, data: dict, sid: str
-    ) -> dict | WebSocketResult | None:
+    async def process(self, event: str, data: dict, sid: str) -> dict | None:
         response_data: dict = {}
 
         await extension.call_extensions_async(
             "webui_ws_event",
             agent=None,
             instance=self,
             sid=sid,
-            event_type=event_type,
+            event_type=event,
             data=data,
             response_data=response_data,
         )
 
-        return self.result_ok(
-            response_data,
-            correlation_id=data.get("correlationId"),
-        )
+        # Return None (fire-and-forget) when no extension populated the response.
+        return response_data if response_data else None