Implemented Websocket Support for Agents

Data/Agent/borealis-agent.py

@@ -11,7 +11,8 @@ from PyQt5 import QtCore, QtGui, QtWidgets
 from PIL import ImageGrab
 
 # ---------------- Configuration ----------------
-SERVER_URL = "http://localhost:5000"  # WebSocket-enabled server URL
+#SERVER_URL = "http://localhost:5000"  # WebSocket-enabled Internal URL
+SERVER_URL = "https://borealis.bunny-lab.io"  # WebSocket-enabled Public URL"
 
 HOSTNAME = socket.gethostname().lower()
 RANDOM_SUFFIX = uuid.uuid4().hex[:8]
@@ -32,7 +33,7 @@ sio = socketio.Client()
 # ---------------- WebSocket Handlers ----------------
 @sio.event
 def connect():
-    print(f"[WS CONNECTED] Agent ID: {AGENT_ID} connected to Borealis.")
+    print(f"[WebSocket] Agent ID: {AGENT_ID} connected to Borealis.")
     sio.emit('connect_agent', {"agent_id": AGENT_ID, "hostname": HOSTNAME})
     sio.emit('request_config', {"agent_id": AGENT_ID})
 
@@ -123,7 +124,6 @@ class ScreenshotRegion(QtWidgets.QWidget):
 
 # ---------------- Screenshot Capture ----------------
 def capture_loop():
-    print("[INFO] Screenshot capture loop started")
     config_ready.wait()
 
     while region_widget is None:
@@ -172,7 +172,7 @@ if __name__ == "__main__":
     app_instance = QtWidgets.QApplication(sys.argv)
     region_launcher = RegionLauncher()
 
-    sio.connect(SERVER_URL)
+    sio.connect(SERVER_URL, transports=['websocket'])
 
     threading.Thread(target=capture_loop, daemon=True).start()
 

Data/Agent/requirements.txt

@@ -1,3 +1,10 @@
+# API / WebSocket Handling
 requests
+python-socketio
+websocket-client
+
+# GUI-related dependencies (Qt for GUI components)
 PyQt5
+
+# Computer Vision & OCR Dependencies
 Pillow

Data/WebUI/src/nodes/Agent/Node_Borealis_Agent.jsx

@@ -2,9 +2,11 @@ import React, { useEffect, useState, useRef } from "react";
 import { Handle, Position, useReactFlow } from "reactflow";
 import { io } from "socket.io-client";
 
-const socket = io();
+const socket = io(window.location.origin, {
+    transports: ["websocket"]
+});
 
-const APINode = ({ id, data }) => {
+const BorealisAgentNode = ({ id, data }) => {
     const { setNodes } = useReactFlow();
     const [agents, setAgents] = useState([]);
     const [selectedAgent, setSelectedAgent] = useState(data.agent_id || "");
@@ -13,6 +15,7 @@ const APINode = ({ id, data }) => {
     const [paused, setPaused] = useState(false);
     const [overlayVisible, setOverlayVisible] = useState(true);
     const [imageData, setImageData] = useState("");
+    const imageRef = useRef("");
 
     useEffect(() => {
         fetch("/api/agents").then(res => res.json()).then(setAgents);
@@ -24,15 +27,39 @@ const APINode = ({ id, data }) => {
 
     useEffect(() => {
         socket.on('new_screenshot', (data) => {
+            console.log("[DEBUG] Screenshot received", data);
             if (data.agent_id === selectedAgent) {
                 setImageData(data.image_base64);
-                window.BorealisValueBus = window.BorealisValueBus || {};
+                imageRef.current = data.image_base64;
-                window.BorealisValueBus[id] = data.image_base64;
             }
         });
 
         return () => socket.off('new_screenshot');
-    }, [selectedAgent, id]);
+    }, [selectedAgent]);
+
+    useEffect(() => {
+        const interval = setInterval(() => {
+            if (!paused && imageRef.current) {
+                window.BorealisValueBus = window.BorealisValueBus || {};
+                window.BorealisValueBus[id] = imageRef.current;
+
+                setNodes(nds => {
+                    const updated = [...nds];
+                    const node = updated.find(n => n.id === id);
+                    if (node) {
+                        node.data = {
+                            ...node.data,
+                            value: imageRef.current
+                        };
+                    }
+                    return updated;
+                });
+
+            }
+        }, window.BorealisUpdateRate || 100);
+
+        return () => clearInterval(interval);
+    }, [id, paused, setNodes]);
 
     const provisionAgent = () => {
         if (!selectedAgent) return;
@@ -49,17 +76,10 @@ const APINode = ({ id, data }) => {
                 visible: overlayVisible,
                 task: selectedType
             })
-        }).then(() => {
+        })
-            socket.emit('request_config', { agent_id: selectedAgent });
+            .then(res => res.json())
-            setNodes(nds =>
+            .then(() => {
-                nds.map(n => n.id === id
+                console.log("[DEBUG] Agent provisioned");
-                    ? {
-                        ...n,
-                        data: { agent_id: selectedAgent, data_type: selectedType, interval: intervalMs }
-                    }
-                    : n
-                )
-            );
             });
     };
 
@@ -72,7 +92,7 @@ const APINode = ({ id, data }) => {
     return (
         <div className="borealis-node">
             <Handle type="source" position={Position.Right} className="borealis-handle" />
-            <div className="borealis-node-header">API Data Collector</div>
+            <div className="borealis-node-header">Borealis Agent</div>
             <div className="borealis-node-content">
                 <label style={{ fontSize: "10px" }}>Agent:</label>
                 <select
@@ -139,9 +159,9 @@ const APINode = ({ id, data }) => {
 };
 
 export default {
-    type: "API_Data_Collector",
+    type: "Borealis_Agent",
-    label: "API Data Collector",
+    label: "Borealis Agent",
-    description: "Real-time provisioning and image collection via WebSocket.",
+    description: "Connects to and controls a Borealis Agent via WebSocket in real-time.",
-    content: "Publishes real-time agent-collected data into the workflow.",
+    content: "Provisions a Borealis Agent and streams collected data into the workflow graph.",
-    component: APINode
+    component: BorealisAgentNode
 };

Data/server.py

@@ -1,7 +1,8 @@
-from flask import Flask, request, jsonify, send_from_directory
+from flask import Flask, request, jsonify, send_from_directory, Response
 from flask_socketio import SocketIO, emit
 import time
 import os
+import base64
 
 # ---------------------------------------------
 # React Frontend Hosting Configuration
@@ -11,7 +12,7 @@ if not os.path.exists(build_folder):
     print("WARNING: web-interface build folder not found. Please build your React app.")
 
 app = Flask(__name__, static_folder=build_folder, static_url_path="/")
-socketio = SocketIO(app, cors_allowed_origins="*", transports=['websocket'])
+socketio = SocketIO(app, cors_allowed_origins="*", async_mode="threading")
 
 @app.route("/")
 def serve_index():
@@ -28,26 +29,15 @@ def serve_react_app(path):
     return send_from_directory(build_folder, "index.html")
 
 # ---------------------------------------------
-# Borealis Agent Management (Hybrid: API + WebSockets)
+# Borealis Agent API Endpoints
 # ---------------------------------------------
 registered_agents = {}
 agent_configurations = {}
 latest_images = {}
 
-# API Endpoints (kept for provisioning and status)
+@app.route("/api/agents")
-@app.route("/api/agent/checkin", methods=["POST"])
+def get_agents():
-def agent_checkin():
+    return jsonify(registered_agents)
-    data = request.json
-    agent_id = data.get("agent_id")
-    hostname = data.get("hostname", "unknown")
-
-    registered_agents[agent_id] = {
-        "agent_id": agent_id,
-        "hostname": hostname,
-        "last_seen": time.time(),
-        "status": "orphaned" if agent_id not in agent_configurations else "provisioned"
-    }
-    return jsonify({"status": "ok"})
 
 @app.route("/api/agent/provision", methods=["POST"])
 def provision_agent():
@@ -65,57 +55,90 @@ def provision_agent():
     agent_configurations[agent_id] = config
     if agent_id in registered_agents:
         registered_agents[agent_id]["status"] = "provisioned"
+
+    # NEW: Emit config update back to the agent via WebSocket
+    socketio.emit("agent_config", config)
+
     return jsonify({"status": "provisioned"})
 
-@app.route("/api/agents")
-def get_agents():
-    return jsonify(registered_agents)
 
-# WebSocket Handlers
+# ---------------------------------------------
+# Raw Image Feed Viewer for Screenshot Agents
+# ---------------------------------------------
+@app.route("/api/agent/<agent_id>/screenshot")
+def screenshot_viewer(agent_id):
+    if agent_configurations.get(agent_id, {}).get("task") != "screenshot":
+        return "<h1>Agent not provisioned as Screenshot Collector</h1>", 400
+
+    html = f"""
+    <html>
+    <head>
+        <title>Borealis - {agent_id} Screenshot</title>
+        <script>
+            setInterval(function() {{
+                var img = document.getElementById('feed');
+                img.src = '/api/agent/{agent_id}/screenshot/raw?rnd=' + Math.random();
+            }}, 1000);
+        </script>
+    </head>
+    <body style='background-color: black;'>
+        <img id='feed' src='/api/agent/{agent_id}/screenshot/raw' style='max-width:100%; height:auto;' />
+    </body>
+    </html>
+    """
+    return html
+
+@app.route("/api/agent/<agent_id>/screenshot/raw")
+def screenshot_raw(agent_id):
+    entry = latest_images.get(agent_id)
+    if not entry:
+        return "", 204
+    try:
+        raw_img = base64.b64decode(entry["image_base64"])
+        return Response(raw_img, mimetype="image/png")
+    except Exception:
+        return "", 204
+
+# ---------------------------------------------
+# WebSocket Events
+# ---------------------------------------------
 @socketio.on('connect_agent')
-def handle_agent_connect(data):
+def connect_agent(data):
-    agent_id = data.get('agent_id')
+    agent_id = data.get("agent_id")
-    hostname = data.get('hostname', 'unknown')
+    hostname = data.get("hostname", "unknown")
+    print(f"Agent connected: {agent_id}")
 
     registered_agents[agent_id] = {
         "agent_id": agent_id,
         "hostname": hostname,
         "last_seen": time.time(),
-        "status": "connected"
+        "status": "orphaned" if agent_id not in agent_configurations else "provisioned"
     }
 
-    print(f"Agent connected: {agent_id}")
-    emit('agent_connected', {'status': 'connected'})
-
-@socketio.on('screenshot')
-def handle_screenshot(data):
-    agent_id = data.get('agent_id')
-    image_base64 = data.get('image_base64')
-
-    if agent_id and image_base64:
-        latest_images[agent_id] = {
-            'image_base64': image_base64,
-            'timestamp': time.time()
-        }
-
-        # Real-time broadcast to connected dashboards
-        emit('new_screenshot', {
-            'agent_id': agent_id,
-            'image_base64': image_base64
-        }, broadcast=True)
-
-        emit('screenshot_received', {'status': 'ok'})
-        print(f"Screenshot received from agent: {agent_id}")
-    else:
-        emit('error', {'message': 'Invalid screenshot data'})
-
 @socketio.on('request_config')
-def handle_request_config(data):
+def send_agent_config(data):
-    agent_id = data.get('agent_id')
+    agent_id = data.get("agent_id")
-    config = agent_configurations.get(agent_id, {})
+    config = agent_configurations.get(agent_id)
-
+    if config:
         emit('agent_config', config)
 
+@socketio.on('screenshot')
+def receive_screenshot(data):
+    agent_id = data.get("agent_id")
+    image = data.get("image_base64")
+
+    if agent_id and image:
+        latest_images[agent_id] = {
+            "image_base64": image,
+            "timestamp": time.time()
+        }
+        print(f"[DEBUG] Screenshot received from agent {agent_id}")
+        emit("new_screenshot", {"agent_id": agent_id, "image_base64": image}, broadcast=True)
+
+@socketio.on('disconnect')
+def on_disconnect():
+    print("[WS] Agent disconnected")
+
 # ---------------------------------------------
 # Server Start
 # ---------------------------------------------

Launch-Borealis.ps1

@@ -137,7 +137,7 @@ switch ($choice) {
                 npm install --silent react-resizable --no-fund --audit=false | Out-Null
                 npm install --silent reactflow --no-fund --audit=false | Out-Null
                 npm install --silent @mui/material @mui/icons-material @emotion/react @emotion/styled --no-fund --audit=false 2>&1 | Out-Null
-
+                npm install  --silent socket.io-client --no-fund --audit=false | Out-Null
                 Pop-Location
             }
         }
@@ -153,7 +153,7 @@ switch ($choice) {
         Push-Location $venvFolder
         Write-Host "`nLaunching Borealis..." -ForegroundColor Green
         Write-Host "===================================================================================="
-        Write-Host "$($symbols.Running) Starting Python Flask Server..." -NoNewline
+        Write-Host "$($symbols.Running) Python Flask Server Started..."
         python "Borealis\server.py"
         Pop-Location
     }