Agent Multi-Role Milestone 3
This commit is contained in:
parent
4c42f53cae
commit
5f80ed59ca
@ -10,6 +10,7 @@ import concurrent.futures
|
||||
from functools import partial
|
||||
from io import BytesIO
|
||||
import base64
|
||||
import traceback
|
||||
|
||||
import socketio
|
||||
from qasync import QEventLoop
|
||||
@ -17,7 +18,7 @@ from PyQt5 import QtCore, QtGui, QtWidgets
|
||||
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")
|
||||
DEFAULT_CONFIG = {
|
||||
@ -36,7 +37,9 @@ class ConfigManager:
|
||||
self.load()
|
||||
|
||||
def load(self):
|
||||
print("[DEBUG] Loading config from disk.")
|
||||
if not os.path.exists(self.path):
|
||||
print("[DEBUG] Config file not found. Creating default.")
|
||||
self.data = DEFAULT_CONFIG.copy()
|
||||
self._write()
|
||||
else:
|
||||
@ -44,6 +47,7 @@ class ConfigManager:
|
||||
with open(self.path, 'r') as f:
|
||||
loaded = json.load(f)
|
||||
self.data = {**DEFAULT_CONFIG, **loaded}
|
||||
print("[DEBUG] Config loaded:", self.data)
|
||||
except Exception as e:
|
||||
print(f"[WARN] Failed to parse config: {e}")
|
||||
self.data = DEFAULT_CONFIG.copy()
|
||||
@ -56,6 +60,7 @@ class ConfigManager:
|
||||
try:
|
||||
with open(self.path, 'w') as f:
|
||||
json.dump(self.data, f, indent=2)
|
||||
print("[DEBUG] Config written to disk.")
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Could not write config: {e}")
|
||||
|
||||
@ -71,83 +76,92 @@ class ConfigManager:
|
||||
return False
|
||||
|
||||
CONFIG = ConfigManager(CONFIG_PATH)
|
||||
CONFIG.data['regions'] = {}
|
||||
CONFIG._write()
|
||||
# //////////////////////////////////////////////////////////////////////////
|
||||
# END CORE SECTION: CONFIG MANAGER
|
||||
# //////////////////////////////////////////////////////////////////////////
|
||||
CONFIG.load()
|
||||
|
||||
host = socket.gethostname().lower()
|
||||
stored_id = CONFIG.data.get('agent_id')
|
||||
if stored_id:
|
||||
AGENT_ID = stored_id
|
||||
else:
|
||||
AGENT_ID = f"{host}-agent-{uuid.uuid4().hex[:8]}"
|
||||
CONFIG.data['agent_id'] = AGENT_ID
|
||||
def init_agent_id():
|
||||
if not CONFIG.data.get('agent_id'):
|
||||
CONFIG.data['agent_id'] = f"{socket.gethostname().lower()}-agent-{uuid.uuid4().hex[:8]}"
|
||||
CONFIG._write()
|
||||
return CONFIG.data['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()
|
||||
|
||||
# //////////////////////////////////////////////////////////////////////////
|
||||
# CORE SECTION: WEBSOCKET SETUP & HANDLERS (do not modify unless absolutely necessary)
|
||||
clear_regions_only()
|
||||
|
||||
# //////////////////////////////////////////////////////////////////////////
|
||||
|
||||
sio = socketio.AsyncClient(reconnection=True, reconnection_attempts=0, reconnection_delay=5)
|
||||
role_tasks = {}
|
||||
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
|
||||
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('request_config', {"agent_id": AGENT_ID})
|
||||
|
||||
@sio.event
|
||||
async def disconnect():
|
||||
print("[WebSocket] Disconnected from Borealis server.")
|
||||
for task in list(role_tasks.values()):
|
||||
task.cancel()
|
||||
role_tasks.clear()
|
||||
for widget in list(overlay_widgets.values()):
|
||||
try: widget.close()
|
||||
except: pass
|
||||
overlay_widgets.clear()
|
||||
await stop_all_roles()
|
||||
CONFIG.data['regions'].clear()
|
||||
CONFIG._write()
|
||||
CONFIG.load()
|
||||
|
||||
@sio.on('agent_config')
|
||||
async def on_agent_config(cfg):
|
||||
print(f"[CONNECTED] Received config with {len(cfg.get('roles',[]))} roles.")
|
||||
new_ids = {r.get('node_id') for r in cfg.get('roles', []) if r.get('node_id')}
|
||||
print("[DEBUG] agent_config event received.")
|
||||
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())
|
||||
removed = old_ids - new_ids
|
||||
|
||||
# Cancel removed roles
|
||||
for rid in removed:
|
||||
if rid in CONFIG.data['regions']:
|
||||
CONFIG.data['regions'].pop(rid, None)
|
||||
print(f"[DEBUG] Removing node {rid} from regions/overlays.")
|
||||
CONFIG.data['regions'].pop(rid, None)
|
||||
w = overlay_widgets.pop(rid, None)
|
||||
if w:
|
||||
try: w.close()
|
||||
except: pass
|
||||
|
||||
if removed:
|
||||
CONFIG._write()
|
||||
|
||||
# Cancel all existing to ensure clean state
|
||||
for task in list(role_tasks.values()):
|
||||
task.cancel()
|
||||
role_tasks.clear()
|
||||
|
||||
# Restart everything to ensure roles are re-applied
|
||||
for role_cfg in cfg.get('roles', []):
|
||||
for role_cfg in roles:
|
||||
nid = role_cfg.get('node_id')
|
||||
if role_cfg.get('role') == 'screenshot':
|
||||
print(f"[DEBUG] Starting screenshot task for {nid}")
|
||||
task = asyncio.create_task(screenshot_task(role_cfg))
|
||||
role_tasks[nid] = task
|
||||
# //////////////////////////////////////////////////////////////////////////
|
||||
# END CORE SECTION: WEBSOCKET SETUP & HANDLERS
|
||||
# //////////////////////////////////////////////////////////////////////////
|
||||
|
||||
# ---------------- Overlay Widget ----------------
|
||||
class ScreenshotRegion(QtWidgets.QWidget):
|
||||
def __init__(self, node_id, x=100, y=100, w=300, h=200):
|
||||
super().__init__()
|
||||
@ -169,14 +183,14 @@ class ScreenshotRegion(QtWidgets.QWidget):
|
||||
p = QtGui.QPainter(self)
|
||||
p.setRenderHint(QtGui.QPainter.Antialiasing)
|
||||
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())
|
||||
hr = self.resize_handle_size
|
||||
hrect = QtCore.QRect(self.width()-hr, self.height()-hr, hr, hr)
|
||||
p.fillRect(hrect, QtGui.QColor(0,255,0))
|
||||
hrect = QtCore.QRect(self.width() - hr, self.height() - hr, hr, hr)
|
||||
p.fillRect(hrect, QtGui.QColor(0, 255, 0))
|
||||
|
||||
def mousePressEvent(self, e):
|
||||
if e.button()==QtCore.Qt.LeftButton:
|
||||
if e.button() == QtCore.Qt.LeftButton:
|
||||
x, y = e.pos().x(), e.pos().y()
|
||||
if x > self.width() - self.resize_handle_size and y > self.height() - self.resize_handle_size:
|
||||
self.resizing = True
|
||||
@ -202,15 +216,13 @@ class ScreenshotRegion(QtWidgets.QWidget):
|
||||
# ---------------- Screenshot Task ----------------
|
||||
async def screenshot_task(cfg):
|
||||
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)
|
||||
if r:
|
||||
region = (r['x'], r['y'], r['w'], r['h'])
|
||||
else:
|
||||
region = (cfg.get('x', 100), cfg.get('y', 100), cfg.get('w', 300), cfg.get('h', 200))
|
||||
CONFIG.data['regions'][nid] = {
|
||||
'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3]
|
||||
}
|
||||
CONFIG.data['regions'][nid] = {'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3]}
|
||||
CONFIG._write()
|
||||
|
||||
if nid not in overlay_widgets:
|
||||
@ -218,21 +230,25 @@ async def screenshot_task(cfg):
|
||||
overlay_widgets[nid] = widget
|
||||
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
|
||||
loop = asyncio.get_event_loop()
|
||||
executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=CONFIG.data.get('max_task_workers', DEFAULT_CONFIG['max_task_workers'])
|
||||
)
|
||||
executor = concurrent.futures.ThreadPoolExecutor(max_workers=CONFIG.data.get('max_task_workers', 8))
|
||||
|
||||
try:
|
||||
while True:
|
||||
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}
|
||||
if prev != new_geom:
|
||||
if CONFIG.data['regions'].get(nid) != new_geom:
|
||||
CONFIG.data['regions'][nid] = new_geom
|
||||
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)
|
||||
buf = BytesIO()
|
||||
img.save(buf, format='PNG')
|
||||
@ -241,43 +257,79 @@ async def screenshot_task(cfg):
|
||||
'agent_id': AGENT_ID,
|
||||
'node_id': nid,
|
||||
'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)
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
print(f"[TASK] Screenshot role {nid} cancelled.")
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Screenshot task {nid} failed: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
# ---------------- Config Watcher ----------------
|
||||
async def config_watcher():
|
||||
print("[DEBUG] Starting config watcher")
|
||||
while True:
|
||||
if CONFIG.watch(): pass
|
||||
await asyncio.sleep(CONFIG.data.get('config_file_watcher_interval', DEFAULT_CONFIG['config_file_watcher_interval']))
|
||||
CONFIG.watch()
|
||||
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():
|
||||
retry = 5
|
||||
while True:
|
||||
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}...")
|
||||
await sio.connect(url, transports=['websocket'])
|
||||
break
|
||||
except:
|
||||
print(f"[WebSocket] Server unavailable, retrying in {retry}s...")
|
||||
except Exception as e:
|
||||
print(f"[WebSocket] Server unavailable: {e}. Retrying in {retry}s...")
|
||||
await asyncio.sleep(retry)
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("[DEBUG] Starting QApplication and QEventLoop")
|
||||
app = QtWidgets.QApplication(sys.argv)
|
||||
loop = QEventLoop(app)
|
||||
asyncio.set_event_loop(loop)
|
||||
with loop:
|
||||
loop.create_task(config_watcher())
|
||||
loop.create_task(connect_loop())
|
||||
loop.run_forever()
|
||||
# //////////////////////////////////////////////////////////////////////////
|
||||
# END CORE SECTION: MAIN & EVENT LOOP
|
||||
# //////////////////////////////////////////////////////////////////////////
|
||||
|
||||
# Dummy window to keep PyQt event loop alive
|
||||
dummy_window = PersistentWindow()
|
||||
dummy_window.show()
|
||||
print("[DEBUG] Dummy window shown to prevent Qt exit")
|
||||
|
||||
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.")
|
||||
|
||||
@ -74,7 +74,7 @@ export default function StatusBar() {
|
||||
</Box>
|
||||
|
||||
<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
|
||||
href="http://localhost:5000/health"
|
||||
target="_blank"
|
||||
|
||||
@ -11,6 +11,7 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
const [isConnected, setIsConnected] = useState(false);
|
||||
const prevRolesRef = useRef([]);
|
||||
|
||||
// ---------------- Agent List & Sorting ----------------
|
||||
const agentList = useMemo(() => {
|
||||
if (!agents || typeof agents !== "object") return [];
|
||||
return Object.entries(agents)
|
||||
@ -23,6 +24,7 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
.sort((a, b) => b.last_seen - a.last_seen);
|
||||
}, [agents]);
|
||||
|
||||
// ---------------- Periodic Agent Fetching ----------------
|
||||
useEffect(() => {
|
||||
const fetchAgents = () => {
|
||||
fetch("/api/agents")
|
||||
@ -35,6 +37,7 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
return () => clearInterval(interval);
|
||||
}, []);
|
||||
|
||||
// ---------------- Node Data Sync ----------------
|
||||
useEffect(() => {
|
||||
setNodes((nds) =>
|
||||
nds.map((n) =>
|
||||
@ -44,6 +47,7 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
setIsConnected(false);
|
||||
}, [selectedAgent]);
|
||||
|
||||
// ---------------- Attached Role Collection ----------------
|
||||
const attachedRoleIds = useMemo(
|
||||
() =>
|
||||
edges
|
||||
@ -62,8 +66,9 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
.filter((r) => r);
|
||||
}, [attachedRoleIds, getNodes]);
|
||||
|
||||
// ---------------- Provision Role Logic ----------------
|
||||
const provisionRoles = useCallback((roles) => {
|
||||
if (!selectedAgent || roles.length === 0) return;
|
||||
if (!selectedAgent) return; // Allow empty roles but require agent
|
||||
fetch("/api/agent/provision", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
@ -78,7 +83,7 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
|
||||
const handleConnect = useCallback(() => {
|
||||
const roles = getAttachedRoles();
|
||||
provisionRoles(roles);
|
||||
provisionRoles(roles); // Always call even with empty roles
|
||||
}, [getAttachedRoles, provisionRoles]);
|
||||
|
||||
const handleDisconnect = useCallback(() => {
|
||||
@ -95,6 +100,7 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
.catch(() => {});
|
||||
}, [selectedAgent]);
|
||||
|
||||
// ---------------- Auto-Provision When Roles Change ----------------
|
||||
useEffect(() => {
|
||||
const newRoles = getAttachedRoles();
|
||||
const prevSerialized = JSON.stringify(prevRolesRef.current || []);
|
||||
@ -102,8 +108,9 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
if (isConnected && newSerialized !== prevSerialized) {
|
||||
provisionRoles(newRoles);
|
||||
}
|
||||
}, [attachedRoleIds, isConnected]);
|
||||
}, [attachedRoleIds, isConnected, getAttachedRoles, provisionRoles]);
|
||||
|
||||
// ---------------- Status Label ----------------
|
||||
const selectedAgentStatus = useMemo(() => {
|
||||
if (!selectedAgent) return "Unassigned";
|
||||
const agent = agents[selectedAgent];
|
||||
@ -111,6 +118,7 @@ const BorealisAgentNode = ({ id, data }) => {
|
||||
return agent.status === "provisioned" ? "Connected" : "Available";
|
||||
}, [agents, selectedAgent]);
|
||||
|
||||
// ---------------- Render ----------------
|
||||
return (
|
||||
<div className="borealis-node">
|
||||
<Handle
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user