bunny-lab/Borealis-Github-Replica
1
0
Fork 0
mirror of https://github.com/bunny-lab-io/Borealis.git synced 2025-07-27 03:48:29 -06:00
Code Issues Packages Projects Releases Wiki Activity

Agent Multi-Role Milestone 3

Browse Source
This commit is contained in:
Nicole Rappe
2025-05-03 04:58:21 -06:00
parent 4c42f53cae
commit 5f80ed59ca
3 changed files with 129 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -10,6 +10,7 @@ import concurrent.futures
from functools import partial from functools import partial
from io import BytesIO from io import BytesIO
import base64 import base64
import traceback
import socketio import socketio
from qasync import QEventLoop from qasync import QEventLoop
@ -17,7 +18,7 @@ from PyQt5 import QtCore, QtGui, QtWidgets
from PIL import ImageGrab from PIL import ImageGrab
# ////////////////////////////////////////////////////////////////////////// # //////////////////////////////////////////////////////////////////////////
# CORE SECTION: CONFIG MANAGER (do not modify unless you know what you’re doing) # CORE SECTION: CONFIG MANAGER
# ////////////////////////////////////////////////////////////////////////// # //////////////////////////////////////////////////////////////////////////
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "agent_settings.json") CONFIG_PATH = os.path.join(os.path.dirname(__file__), "agent_settings.json")
DEFAULT_CONFIG = { DEFAULT_CONFIG = {
@ -36,7 +37,9 @@ class ConfigManager:
self.load() self.load()
def load(self): def load(self):
print("[DEBUG] Loading config from disk.")
if not os.path.exists(self.path): if not os.path.exists(self.path):
print("[DEBUG] Config file not found. Creating default.")
self.data = DEFAULT_CONFIG.copy() self.data = DEFAULT_CONFIG.copy()
self._write() self._write()
else: else:
@ -44,6 +47,7 @@ class ConfigManager:
with open(self.path, 'r') as f: with open(self.path, 'r') as f:
loaded = json.load(f) loaded = json.load(f)
self.data = {**DEFAULT_CONFIG, **loaded} self.data = {**DEFAULT_CONFIG, **loaded}
print("[DEBUG] Config loaded:", self.data)
except Exception as e: except Exception as e:
print(f"[WARN] Failed to parse config: {e}") print(f"[WARN] Failed to parse config: {e}")
self.data = DEFAULT_CONFIG.copy() self.data = DEFAULT_CONFIG.copy()
@ -56,6 +60,7 @@ class ConfigManager:
try: try:
with open(self.path, 'w') as f: with open(self.path, 'w') as f:
json.dump(self.data, f, indent=2) json.dump(self.data, f, indent=2)
print("[DEBUG] Config written to disk.")
except Exception as e: except Exception as e:
print(f"[ERROR] Could not write config: {e}") print(f"[ERROR] Could not write config: {e}")
@ -71,83 +76,92 @@ class ConfigManager:
return False return False
CONFIG = ConfigManager(CONFIG_PATH) CONFIG = ConfigManager(CONFIG_PATH)
CONFIG.data['regions'] = {} CONFIG.load()
CONFIG._write()
# //////////////////////////////////////////////////////////////////////////
# END CORE SECTION: CONFIG MANAGER
# //////////////////////////////////////////////////////////////////////////
host = socket.gethostname().lower() def init_agent_id():
stored_id = CONFIG.data.get('agent_id') if not CONFIG.data.get('agent_id'):
if stored_id: CONFIG.data['agent_id'] = f"{socket.gethostname().lower()}-agent-{uuid.uuid4().hex[:8]}"
AGENT_ID = stored_id CONFIG._write()
else: return CONFIG.data['agent_id']
AGENT_ID = f"{host}-agent-{uuid.uuid4().hex[:8]}"
CONFIG.data['agent_id'] = AGENT_ID AGENT_ID = init_agent_id()
print(f"[DEBUG] Using AGENT_ID: {AGENT_ID}")
def clear_regions_only():
CONFIG.data['regions'] = CONFIG.data.get('regions', {})
CONFIG._write() CONFIG._write()
# ////////////////////////////////////////////////////////////////////////// clear_regions_only()
# CORE SECTION: WEBSOCKET SETUP & HANDLERS (do not modify unless absolutely necessary)
# ////////////////////////////////////////////////////////////////////////// # //////////////////////////////////////////////////////////////////////////
sio = socketio.AsyncClient(reconnection=True, reconnection_attempts=0, reconnection_delay=5) sio = socketio.AsyncClient(reconnection=True, reconnection_attempts=0, reconnection_delay=5)
role_tasks = {} role_tasks = {}
overlay_widgets = {} overlay_widgets = {}
background_tasks = []
async def stop_all_roles():
print("[DEBUG] Stopping all roles.")
for task in list(role_tasks.values()):
print(f"[DEBUG] Cancelling task for node: {task}")
task.cancel()
role_tasks.clear()
for node_id, widget in overlay_widgets.items():
print(f"[DEBUG] Closing overlay widget: {node_id}")
try: widget.close()
except Exception as e: print(f"[WARN] Error closing widget: {e}")
overlay_widgets.clear()
@sio.event @sio.event
async def connect(): async def connect():
print(f"[WebSocket] Connected to Agent ID: {AGENT_ID}.") print(f"[WebSocket] Connected to Borealis Server with Agent ID: {AGENT_ID}")
await sio.emit('connect_agent', {"agent_id": AGENT_ID}) await sio.emit('connect_agent', {"agent_id": AGENT_ID})
await sio.emit('request_config', {"agent_id": AGENT_ID}) await sio.emit('request_config', {"agent_id": AGENT_ID})
@sio.event @sio.event
async def disconnect(): async def disconnect():
print("[WebSocket] Disconnected from Borealis server.") print("[WebSocket] Disconnected from Borealis server.")
for task in list(role_tasks.values()): await stop_all_roles()
task.cancel()
role_tasks.clear()
for widget in list(overlay_widgets.values()):
try: widget.close()
except: pass
overlay_widgets.clear()
CONFIG.data['regions'].clear() CONFIG.data['regions'].clear()
CONFIG._write() CONFIG._write()
CONFIG.load()
@sio.on('agent_config') @sio.on('agent_config')
async def on_agent_config(cfg): async def on_agent_config(cfg):
print(f"[CONNECTED] Received config with {len(cfg.get('roles',[]))} roles.") print("[DEBUG] agent_config event received.")
new_ids = {r.get('node_id') for r in cfg.get('roles', []) if r.get('node_id')} roles = cfg.get('roles', [])
if not roles:
print("[CONFIG] Config Reset by Borealis Server Operator - Awaiting New Config...")
await stop_all_roles()
return
print(f"[CONFIG] Received New Agent Config with {len(roles)} Role(s).")
new_ids = {r.get('node_id') for r in roles if r.get('node_id')}
old_ids = set(role_tasks.keys()) old_ids = set(role_tasks.keys())
removed = old_ids - new_ids removed = old_ids - new_ids
# Cancel removed roles
for rid in removed: for rid in removed:
if rid in CONFIG.data['regions']: print(f"[DEBUG] Removing node {rid} from regions/overlays.")
CONFIG.data['regions'].pop(rid, None) CONFIG.data['regions'].pop(rid, None)
w = overlay_widgets.pop(rid, None) w = overlay_widgets.pop(rid, None)
if w: if w:
try: w.close() try: w.close()
except: pass except: pass
if removed: if removed:
CONFIG._write() CONFIG._write()
# Cancel all existing to ensure clean state
for task in list(role_tasks.values()): for task in list(role_tasks.values()):
task.cancel() task.cancel()
role_tasks.clear() role_tasks.clear()
# Restart everything to ensure roles are re-applied for role_cfg in roles:
for role_cfg in cfg.get('roles', []):
nid = role_cfg.get('node_id') nid = role_cfg.get('node_id')
if role_cfg.get('role') == 'screenshot': if role_cfg.get('role') == 'screenshot':
print(f"[DEBUG] Starting screenshot task for {nid}")
task = asyncio.create_task(screenshot_task(role_cfg)) task = asyncio.create_task(screenshot_task(role_cfg))
role_tasks[nid] = task role_tasks[nid] = task
# //////////////////////////////////////////////////////////////////////////
# END CORE SECTION: WEBSOCKET SETUP & HANDLERS
# //////////////////////////////////////////////////////////////////////////
# ---------------- Overlay Widget ----------------
class ScreenshotRegion(QtWidgets.QWidget): class ScreenshotRegion(QtWidgets.QWidget):
def __init__(self, node_id, x=100, y=100, w=300, h=200): def __init__(self, node_id, x=100, y=100, w=300, h=200):
super().__init__() super().__init__()
@ -169,14 +183,14 @@ class ScreenshotRegion(QtWidgets.QWidget):
p = QtGui.QPainter(self) p = QtGui.QPainter(self)
p.setRenderHint(QtGui.QPainter.Antialiasing) p.setRenderHint(QtGui.QPainter.Antialiasing)
p.setBrush(QtCore.Qt.transparent) p.setBrush(QtCore.Qt.transparent)
p.setPen(QtGui.QPen(QtGui.QColor(0,255,0),2)) p.setPen(QtGui.QPen(QtGui.QColor(0, 255, 0), 2))
p.drawRect(self.rect()) p.drawRect(self.rect())
hr = self.resize_handle_size hr = self.resize_handle_size
hrect = QtCore.QRect(self.width()-hr, self.height()-hr, hr, hr) hrect = QtCore.QRect(self.width() - hr, self.height() - hr, hr, hr)
p.fillRect(hrect, QtGui.QColor(0,255,0)) p.fillRect(hrect, QtGui.QColor(0, 255, 0))
def mousePressEvent(self, e): def mousePressEvent(self, e):
if e.button()==QtCore.Qt.LeftButton: if e.button() == QtCore.Qt.LeftButton:
x, y = e.pos().x(), e.pos().y() x, y = e.pos().x(), e.pos().y()
if x > self.width() - self.resize_handle_size and y > self.height() - self.resize_handle_size: if x > self.width() - self.resize_handle_size and y > self.height() - self.resize_handle_size:
self.resizing = True self.resizing = True
@ -202,15 +216,13 @@ class ScreenshotRegion(QtWidgets.QWidget):
# ---------------- Screenshot Task ---------------- # ---------------- Screenshot Task ----------------
async def screenshot_task(cfg): async def screenshot_task(cfg):
nid = cfg.get('node_id') nid = cfg.get('node_id')
# If existing region in config, honor that print(f"[DEBUG] Running screenshot_task for {nid}")
r = CONFIG.data['regions'].get(nid) r = CONFIG.data['regions'].get(nid)
if r: if r:
region = (r['x'], r['y'], r['w'], r['h']) region = (r['x'], r['y'], r['w'], r['h'])
else: else:
region = (cfg.get('x', 100), cfg.get('y', 100), cfg.get('w', 300), cfg.get('h', 200)) region = (cfg.get('x', 100), cfg.get('y', 100), cfg.get('w', 300), cfg.get('h', 200))
CONFIG.data['regions'][nid] = { CONFIG.data['regions'][nid] = {'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3]}
'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3]
}
CONFIG._write() CONFIG._write()
if nid not in overlay_widgets: if nid not in overlay_widgets:
@ -218,21 +230,25 @@ async def screenshot_task(cfg):
overlay_widgets[nid] = widget overlay_widgets[nid] = widget
widget.show() widget.show()
await sio.emit('agent_screenshot_task', {
'agent_id': AGENT_ID,
'node_id': nid,
'image_base64': "",
'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3]
})
interval = cfg.get('interval', 1000) / 1000.0 interval = cfg.get('interval', 1000) / 1000.0
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
executor = concurrent.futures.ThreadPoolExecutor( executor = concurrent.futures.ThreadPoolExecutor(max_workers=CONFIG.data.get('max_task_workers', 8))
max_workers=CONFIG.data.get('max_task_workers', DEFAULT_CONFIG['max_task_workers'])
)
try: try:
while True: while True:
x, y, w, h = overlay_widgets[nid].get_geometry() x, y, w, h = overlay_widgets[nid].get_geometry()
prev = CONFIG.data['regions'].get(nid)
new_geom = {'x': x, 'y': y, 'w': w, 'h': h} new_geom = {'x': x, 'y': y, 'w': w, 'h': h}
if prev != new_geom: if CONFIG.data['regions'].get(nid) != new_geom:
CONFIG.data['regions'][nid] = new_geom CONFIG.data['regions'][nid] = new_geom
CONFIG._write() CONFIG._write()
grab = partial(ImageGrab.grab, bbox=(x, y, x+w, y+h)) grab = partial(ImageGrab.grab, bbox=(x, y, x + w, y + h))
img = await loop.run_in_executor(executor, grab) img = await loop.run_in_executor(executor, grab)
buf = BytesIO() buf = BytesIO()
img.save(buf, format='PNG') img.save(buf, format='PNG')
@ -241,43 +257,79 @@ async def screenshot_task(cfg):
'agent_id': AGENT_ID, 'agent_id': AGENT_ID,
'node_id': nid, 'node_id': nid,
'image_base64': encoded, 'image_base64': encoded,
'x': x, 'y': y, 'w': w, 'h': h # Bi-directional live-sync 'x': x, 'y': y, 'w': w, 'h': h
}) })
await asyncio.sleep(interval) await asyncio.sleep(interval)
except asyncio.CancelledError: except asyncio.CancelledError:
return print(f"[TASK] Screenshot role {nid} cancelled.")
except Exception as e: except Exception as e:
print(f"[ERROR] Screenshot task {nid} failed: {e}") print(f"[ERROR] Screenshot task {nid} failed: {e}")
traceback.print_exc()
# ---------------- Config Watcher ---------------- # ---------------- Config Watcher ----------------
async def config_watcher(): async def config_watcher():
print("[DEBUG] Starting config watcher")
while True: while True:
if CONFIG.watch(): pass CONFIG.watch()
await asyncio.sleep(CONFIG.data.get('config_file_watcher_interval', DEFAULT_CONFIG['config_file_watcher_interval'])) await asyncio.sleep(CONFIG.data.get('config_file_watcher_interval', 2))
# ---------------- Persistent Idle Task ----------------
async def idle_task():
print("[Agent] Entering idle state. Awaiting instructions...")
try:
while True:
await asyncio.sleep(60)
print("[DEBUG] Idle task still alive.")
except asyncio.CancelledError:
print("[FATAL] Idle task was cancelled!")
except Exception as e:
print(f"[FATAL] Idle task crashed: {e}")
traceback.print_exc()
# ---------------- Dummy Qt Widget to Prevent Exit ----------------
class PersistentWindow(QtWidgets.QWidget):
def __init__(self):
super().__init__()
self.setWindowTitle("KeepAlive")
self.setGeometry(-1000, -1000, 1, 1)
self.setAttribute(QtCore.Qt.WA_DontShowOnScreen)
self.hide()
# ////////////////////////////////////////////////////////////////////////// # //////////////////////////////////////////////////////////////////////////
# CORE SECTION: MAIN & EVENT LOOP (do not modify unless you know what you’re doing) # MAIN & EVENT LOOP
# ////////////////////////////////////////////////////////////////////////// # //////////////////////////////////////////////////////////////////////////
async def connect_loop(): async def connect_loop():
retry = 5 retry = 5
while True: while True:
try: try:
url = CONFIG.data.get('borealis_server_url', DEFAULT_CONFIG['borealis_server_url']) url = CONFIG.data.get('borealis_server_url', "http://localhost:5000")
print(f"[WebSocket] Connecting to {url}...") print(f"[WebSocket] Connecting to {url}...")
await sio.connect(url, transports=['websocket']) await sio.connect(url, transports=['websocket'])
break break
except: except Exception as e:
print(f"[WebSocket] Server unavailable, retrying in {retry}s...") print(f"[WebSocket] Server unavailable: {e}. Retrying in {retry}s...")
await asyncio.sleep(retry) await asyncio.sleep(retry)
if __name__ == '__main__': if __name__ == '__main__':
print("[DEBUG] Starting QApplication and QEventLoop")
app = QtWidgets.QApplication(sys.argv) app = QtWidgets.QApplication(sys.argv)
loop = QEventLoop(app) loop = QEventLoop(app)
asyncio.set_event_loop(loop) asyncio.set_event_loop(loop)
with loop:
loop.create_task(config_watcher()) # Dummy window to keep PyQt event loop alive
loop.create_task(connect_loop()) dummy_window = PersistentWindow()
loop.run_forever() dummy_window.show()
# ////////////////////////////////////////////////////////////////////////// print("[DEBUG] Dummy window shown to prevent Qt exit")
# END CORE SECTION: MAIN & EVENT LOOP
# ////////////////////////////////////////////////////////////////////////// try:
with loop:
print("[DEBUG] Scheduling tasks into event loop")
background_tasks.append(loop.create_task(config_watcher()))
background_tasks.append(loop.create_task(connect_loop()))
background_tasks.append(loop.create_task(idle_task()))
loop.run_forever()
except Exception as e:
print(f"[FATAL] Event loop crashed: {e}")
traceback.print_exc()
finally:
print("[FATAL] Agent exited unexpectedly.")

2
Data/Server/WebUI/src/Status_Bar.jsx
View File

@ -74,7 +74,7 @@ export default function StatusBar() {
</Box> </Box>
<Box sx={{ fontSize: "1.0rem", display: "flex", alignItems: "center", gap: 1 }}> <Box sx={{ fontSize: "1.0rem", display: "flex", alignItems: "center", gap: 1 }}>
<strong style={{ color: "#58a6ff" }}>API Server</strong>: <strong style={{ color: "#58a6ff" }}>Backend API Server</strong>:
<a <a
href="http://localhost:5000/health" href="http://localhost:5000/health"
target="_blank" target="_blank"

14
Data/Server/WebUI/src/nodes/Agents/Node_Agent.jsx
View File

@ -11,6 +11,7 @@ const BorealisAgentNode = ({ id, data }) => {
const [isConnected, setIsConnected] = useState(false); const [isConnected, setIsConnected] = useState(false);
const prevRolesRef = useRef([]); const prevRolesRef = useRef([]);
// ---------------- Agent List & Sorting ----------------
const agentList = useMemo(() => { const agentList = useMemo(() => {
if (!agents || typeof agents !== "object") return []; if (!agents || typeof agents !== "object") return [];
return Object.entries(agents) return Object.entries(agents)
@ -23,6 +24,7 @@ const BorealisAgentNode = ({ id, data }) => {
.sort((a, b) => b.last_seen - a.last_seen); .sort((a, b) => b.last_seen - a.last_seen);
}, [agents]); }, [agents]);
// ---------------- Periodic Agent Fetching ----------------
useEffect(() => { useEffect(() => {
const fetchAgents = () => { const fetchAgents = () => {
fetch("/api/agents") fetch("/api/agents")
@ -35,6 +37,7 @@ const BorealisAgentNode = ({ id, data }) => {
return () => clearInterval(interval); return () => clearInterval(interval);
}, []); }, []);
// ---------------- Node Data Sync ----------------
useEffect(() => { useEffect(() => {
setNodes((nds) => setNodes((nds) =>
nds.map((n) => nds.map((n) =>
@ -44,6 +47,7 @@ const BorealisAgentNode = ({ id, data }) => {
setIsConnected(false); setIsConnected(false);
}, [selectedAgent]); }, [selectedAgent]);
// ---------------- Attached Role Collection ----------------
const attachedRoleIds = useMemo( const attachedRoleIds = useMemo(
() => () =>
edges edges
@ -62,8 +66,9 @@ const BorealisAgentNode = ({ id, data }) => {
.filter((r) => r); .filter((r) => r);
}, [attachedRoleIds, getNodes]); }, [attachedRoleIds, getNodes]);
// ---------------- Provision Role Logic ----------------
const provisionRoles = useCallback((roles) => { const provisionRoles = useCallback((roles) => {
if (!selectedAgent || roles.length === 0) return; if (!selectedAgent) return; // Allow empty roles but require agent
fetch("/api/agent/provision", { fetch("/api/agent/provision", {
method: "POST", method: "POST",
headers: { "Content-Type": "application/json" }, headers: { "Content-Type": "application/json" },
@ -78,7 +83,7 @@ const BorealisAgentNode = ({ id, data }) => {
const handleConnect = useCallback(() => { const handleConnect = useCallback(() => {
const roles = getAttachedRoles(); const roles = getAttachedRoles();
provisionRoles(roles); provisionRoles(roles); // Always call even with empty roles
}, [getAttachedRoles, provisionRoles]); }, [getAttachedRoles, provisionRoles]);
const handleDisconnect = useCallback(() => { const handleDisconnect = useCallback(() => {
@ -95,6 +100,7 @@ const BorealisAgentNode = ({ id, data }) => {
.catch(() => {}); .catch(() => {});
}, [selectedAgent]); }, [selectedAgent]);
// ---------------- Auto-Provision When Roles Change ----------------
useEffect(() => { useEffect(() => {
const newRoles = getAttachedRoles(); const newRoles = getAttachedRoles();
const prevSerialized = JSON.stringify(prevRolesRef.current || []); const prevSerialized = JSON.stringify(prevRolesRef.current || []);
@ -102,8 +108,9 @@ const BorealisAgentNode = ({ id, data }) => {
if (isConnected && newSerialized !== prevSerialized) { if (isConnected && newSerialized !== prevSerialized) {
provisionRoles(newRoles); provisionRoles(newRoles);
} }
}, [attachedRoleIds, isConnected]); }, [attachedRoleIds, isConnected, getAttachedRoles, provisionRoles]);
// ---------------- Status Label ----------------
const selectedAgentStatus = useMemo(() => { const selectedAgentStatus = useMemo(() => {
if (!selectedAgent) return "Unassigned"; if (!selectedAgent) return "Unassigned";
const agent = agents[selectedAgent]; const agent = agents[selectedAgent];
@ -111,6 +118,7 @@ const BorealisAgentNode = ({ id, data }) => {
return agent.status === "provisioned" ? "Connected" : "Available"; return agent.status === "provisioned" ? "Connected" : "Available";
}, [agents, selectedAgent]); }, [agents, selectedAgent]);
// ---------------- Render ----------------
return ( return (
<div className="borealis-node"> <div className="borealis-node">
<Handle <Handle