diff --git a/Data/Server/WebUI/src/nodes/Agent Roles/Node_Agent_Role_Screenshot.jsx b/Data/Server/WebUI/src/nodes/Agent Roles/Node_Agent_Role_Screenshot.jsx index 54f9266..50a87f7 100644 --- a/Data/Server/WebUI/src/nodes/Agent Roles/Node_Agent_Role_Screenshot.jsx +++ b/Data/Server/WebUI/src/nodes/Agent Roles/Node_Agent_Role_Screenshot.jsx @@ -55,6 +55,15 @@ const AgentScreenshotNode = ({ id, data }) => { const handleScreenshot = (payload) => { if (payload?.node_id !== id) return; + // Additionally ensure payload is from the agent connected upstream of this node + try { + const agentEdge = edges.find(e => e.target === id && e.sourceHandle === "provisioner"); + const agentNode = getNodes().find(n => n.id === agentEdge?.source); + const selectedAgentId = agentNode?.data?.agent_id; + if (!selectedAgentId || payload?.agent_id !== selectedAgentId) return; + } catch (err) { + return; // fail-closed if we cannot resolve upstream agent + } if (payload.image_base64) { setImageBase64(payload.image_base64); @@ -77,7 +86,7 @@ const AgentScreenshotNode = ({ id, data }) => { socket.on("agent_screenshot_task", handleScreenshot); return () => socket.off("agent_screenshot_task", handleScreenshot); - }, [id, setNodes]); + }, [id, setNodes, edges, getNodes]); // Register this node for the agent provisioning sync window.__BorealisInstructionNodes = window.__BorealisInstructionNodes || {}; diff --git a/Data/Server/server.py b/Data/Server/server.py index f9db76f..cd3d49d 100644 --- a/Data/Server/server.py +++ b/Data/Server/server.py @@ -6,7 +6,7 @@ eventlet.monkey_patch() import requests from flask import Flask, request, jsonify, Response, send_from_directory, make_response -from flask_socketio import SocketIO, emit +from flask_socketio import SocketIO, emit, join_room from flask_cors import CORS import time @@ -1335,7 +1335,12 @@ def provision_agent(): if agent_id in registered_agents: registered_agents[agent_id]["status"] = "provisioned" - socketio.emit("agent_config", config) + # Target only the intended agent by emitting to its room + try: + socketio.emit("agent_config", {**config, "agent_id": agent_id}, room=agent_id) + except TypeError: + # Compatibility with older flask-socketio versions that use 'to' + socketio.emit("agent_config", {**config, "agent_id": agent_id}, to=agent_id) return jsonify({"status": "provisioned", "roles": roles}) # --------------------------------------------- @@ -1454,6 +1459,12 @@ def connect_agent(data): return print(f"Agent connected: {agent_id}") + # Join per-agent room so we can address this connection specifically + try: + join_room(agent_id) + except Exception: + pass + rec = registered_agents.setdefault(agent_id, {}) rec["agent_id"] = agent_id rec["hostname"] = rec.get("hostname", "unknown") diff --git a/Workflows/Agent Screenshot Testing.json b/Workflows/Agent Screenshot Testing.json new file mode 100644 index 0000000..f3e8508 --- /dev/null +++ b/Workflows/Agent Screenshot Testing.json @@ -0,0 +1,5 @@ +{ + "tab_name": "Agent Screenshot Testing", + "nodes": [], + "edges": [] +} \ No newline at end of file