From 9f04c61bd3b7e55d0f6455fa360d64d03ed912a5 Mon Sep 17 00:00:00 2001
From: Nicole Rappe <nicole.rappe@bunny-lab.io>
Date: Sun, 7 Sep 2025 22:11:19 -0600
Subject: [PATCH] Fixed Multiple Agent Screenshot Streams Overlapping

---
 .../Agent Roles/Node_Agent_Role_Screenshot.jsx    | 11 ++++++++++-
 Data/Server/server.py                             | 15 +++++++++++++--
 Workflows/Agent Screenshot Testing.json           |  5 +++++
 3 files changed, 28 insertions(+), 3 deletions(-)
 create mode 100644 Workflows/Agent Screenshot Testing.json

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