bunny-lab/Borealis
bunny-lab/Borealis
2
0
Fork 0
Code Issues 10 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

Implemented Websocket Support for Agents

Browse Source
This commit is contained in:
Nicole Rappe 2025-04-09 01:20:46 -06:00
parent 5c35447a15
commit 4e48d119f7
5 changed files with 131 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
Data/Agent/borealis-agent.py
View File

@ -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()

9
Data/Agent/requirements.txt
View File

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

66
Data/WebUI/src/nodes/Agent/Node_Borealis_Agent.jsx
View File

@ -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 || {};
window.BorealisValueBus[id] = data.image_base64;
imageRef.current = 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,18 +76,11 @@ const APINode = ({ id, data }) => {
visible: overlayVisible,
task: selectedType
})
}).then(() => {
socket.emit('request_config', { agent_id: selectedAgent });
setNodes(nds =>
nds.map(n => n.id === id
? {
...n,
data: { agent_id: selectedAgent, data_type: selectedType, interval: intervalMs }
}
: n
)
);
});
})
.then(res => res.json())
.then(() => {
console.log("[DEBUG] Agent provisioned");
});
};
const toggleOverlay = () => {
@ -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",
label: "API Data Collector",
description: "Real-time provisioning and image collection via WebSocket.",
content: "Publishes real-time agent-collected data into the workflow.",
component: APINode
type: "Borealis_Agent",
label: "Borealis Agent",
description: "Connects to and controls a Borealis Agent via WebSocket in real-time.",
content: "Provisions a Borealis Agent and streams collected data into the workflow graph.",
component: BorealisAgentNode
};

125
Data/server.py
View File

@ -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/agent/checkin", methods=["POST"])
def agent_checkin():
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/agents")
def get_agents():
return jsonify(registered_agents)
@app.route("/api/agent/provision", methods=["POST"])
def provision_agent():
@ -65,59 +55,92 @@ 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):
agent_id = data.get('agent_id')
hostname = data.get('hostname', 'unknown')
def connect_agent(data):
agent_id = data.get("agent_id")
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('request_config')
def send_agent_config(data):
agent_id = data.get("agent_id")
config = agent_configurations.get(agent_id)
if config:
emit('agent_config', config)
@socketio.on('screenshot')
def handle_screenshot(data):
agent_id = data.get('agent_id')
image_base64 = data.get('image_base64')
def receive_screenshot(data):
agent_id = data.get("agent_id")
image = data.get("image_base64")
if agent_id and image_base64:
if agent_id and image:
latest_images[agent_id] = {
'image_base64': image_base64,
'timestamp': time.time()
"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)
# 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):
agent_id = data.get('agent_id')
config = agent_configurations.get(agent_id, {})
emit('agent_config', config)
@socketio.on('disconnect')
def on_disconnect():
print("[WS] Agent disconnected")
# ---------------------------------------------
# Server Start
# ---------------------------------------------
if __name__ == "__main__":
socketio.run(app, host="0.0.0.0", port=5000, debug=False)
socketio.run(app, host="0.0.0.0", port=5000, debug=False)

4
Launch-Borealis.ps1
View File

@ -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
}