Support selecting agent context for screenshots

2025-10-27 09:21:57 -06:00 · 2025-10-15 04:40:21 -06:00
parent a86231c8f5
commit 41c52590f4
6 changed files with 291 additions and 56 deletions
--- a/Roles/Node_Agent_Role_Screenshot.jsx
+++ b/Roles/Node_Agent_Role_Screenshot.jsx
@@ -1,5 +1,5 @@
 ////////// PROJECT FILE SEPARATION LINE ////////// CODE AFTER THIS LINE ARE FROM: <ProjectRoot>/Data/WebUI/src/nodes/Agent/Node_Agent_Role_Screenshot.jsx
-import React, { useEffect, useRef, useState } from "react";
+import React, { useCallback, useEffect, useRef, useState } from "react";
 import { Handle, Position, useReactFlow, useStore } from "reactflow";
 import ShareIcon from "@mui/icons-material/Share";
 import IconButton from "@mui/material/IconButton";
@@ -21,6 +21,17 @@ const AgentScreenshotNode = ({ id, data }) => {
  const { setNodes, getNodes } = useReactFlow();
  const edges = useStore(state => state.edges);

+  const resolveAgentData = useCallback(() => {
+    try {
+      const agentEdge = edges.find(e => e.target === id && e.sourceHandle === "provisioner");
+      const agentNode = getNodes().find(n => n.id === agentEdge?.source);
+      return agentNode?.data || null;
+    } catch (err) {
+      return null;
+    }
+  }, [edges, getNodes, id]);
+
+
  // Core config values pulled from sidebar config (with defaults)
  const interval = parseInt(data?.interval || 1000, 10) || 1000;
  const region = {
@@ -32,6 +43,11 @@ const AgentScreenshotNode = ({ id, data }) => {
  const visible = (data?.visible ?? "true") === "true";
  const alias = data?.alias || "";
  const [imageBase64, setImageBase64] = useState(data?.value || "");
+  const agentData = resolveAgentData();
+  const targetModeLabel = ((agentData?.agent_mode || "").toString().toLowerCase() === "system")
+    ? "SYSTEM Agent"
+    : "CURRENTUSER Agent";
+  const targetHostLabel = (agentData?.agent_host || "").toString();

  // Always push current imageBase64 into BorealisValueBus at the global update rate
  useEffect(() => {
@@ -56,14 +72,9 @@ 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
-      }
+      const agentData = resolveAgentData();
+      const selectedAgentId = agentData?.agent_id;
+      if (!selectedAgentId || payload?.agent_id !== selectedAgentId) return;

      if (payload.image_base64) {
        setImageBase64(payload.image_base64);
@@ -86,24 +97,30 @@ const AgentScreenshotNode = ({ id, data }) => {

    socket.on("agent_screenshot_task", handleScreenshot);
    return () => socket.off("agent_screenshot_task", handleScreenshot);
-  }, [id, setNodes, edges, getNodes]);
+  }, [id, setNodes, resolveAgentData]);

  // Register this node for the agent provisioning sync
  window.__BorealisInstructionNodes = window.__BorealisInstructionNodes || {};
-  window.__BorealisInstructionNodes[id] = () => ({
-    node_id: id,
-    role: "screenshot",
-    interval,
-    visible,
-    alias,
-    ...region
-  });
+  window.__BorealisInstructionNodes[id] = () => {
+    const agentData = resolveAgentData() || {};
+    const modeRaw = (agentData.agent_mode || "").toString().toLowerCase();
+    const targetMode = modeRaw === "system" ? "system" : "currentuser";
+    return {
+      node_id: id,
+      role: "screenshot",
+      interval,
+      visible,
+      alias,
+      target_agent_mode: targetMode,
+      target_agent_host: agentData.agent_host || "",
+      ...region
+    };
+  };

  // Manual live view copy button
  const handleCopyLiveViewLink = () => {
-    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;
+    const agentData = resolveAgentData();
+    const selectedAgentId = agentData?.agent_id;

    if (!selectedAgentId) {
      alert("No valid agent connection found.");
@@ -132,6 +149,17 @@ const AgentScreenshotNode = ({ id, data }) => {
        <div>
          <b>Interval:</b> {interval} ms
        </div>
+        <div>
+          <b>Agent Context:</b> {targetModeLabel}
+        </div>
+        <div>
+          <b>Target Host:</b>{" "}
+          {targetHostLabel ? (
+            targetHostLabel
+          ) : (
+            <span style={{ color: "#666" }}>unknown</span>
+          )}
+        </div>
        <div>
          <b>Overlay:</b> {visible ? "Yes" : "No"}
        </div>
@@ -165,6 +193,7 @@ Capture a live screenshot of a defined region from a remote Borealis Agent.
 - Optionally show a visual overlay with a label
 - Pushes base64 PNG stream to downstream nodes
 - Use copy button to share live view URL
+- Targets the CURRENTUSER or SYSTEM agent context selected upstream
 `.trim(),
  content: "Capture screenshot region via agent",
  component: AgentScreenshotNode,
--- a/Data/Server/WebUI/src/nodes/Agent/Node_Agent.jsx
+++ b/Data/Server/WebUI/src/nodes/Agent/Node_Agent.jsx
@@ -8,22 +8,61 @@ const BorealisAgentNode = ({ id, data }) => {
  const edges = useStore((state) => state.edges);
  const [agents, setAgents] = useState({});
  const [selectedAgent, setSelectedAgent] = useState(data.agent_id || "");
+  const [selectedHost, setSelectedHost] = useState(data.agent_host || "");
+  const initialMode = (data.agent_mode || "currentuser").toLowerCase();
+  const [selectedMode, setSelectedMode] = useState(
+    initialMode === "system" ? "system" : "currentuser"
+  );
  const [isConnected, setIsConnected] = useState(false);
  const prevRolesRef = useRef([]);

-  // Agent List Sorted (Online First)
-  const agentList = useMemo(() => {
-    if (!agents || typeof agents !== "object") return [];
-    return Object.entries(agents)
-      .map(([aid, info]) => ({
-        id: aid,
-        status: info?.status || "offline",
-        last_seen: info?.last_seen || 0
-      }))
-      .filter(({ status }) => status !== "offline")
-      .sort((a, b) => b.last_seen - a.last_seen);
+  // Group agents by hostname and execution context
+  const agentsByHostname = useMemo(() => {
+    if (!agents || typeof agents !== "object") return {};
+    const grouped = {};
+    Object.entries(agents).forEach(([aid, info]) => {
+      if (!info || typeof info !== "object") return;
+      const status = (info.status || "").toString().toLowerCase();
+      if (status === "offline") return;
+      const host = (info.hostname || info.agent_hostname || "").trim() || "unknown";
+      const modeRaw = (info.service_mode || "").toString().toLowerCase();
+      const mode = modeRaw === "system" ? "system" : "currentuser";
+      if (!grouped[host]) {
+        grouped[host] = { currentuser: null, system: null };
+      }
+      grouped[host][mode] = {
+        agent_id: aid,
+        status: info.status || "offline",
+        last_seen: info.last_seen || 0,
+        info,
+      };
+    });
+    return grouped;
  }, [agents]);

+  const hostOptions = useMemo(() => {
+    const entries = Object.entries(agentsByHostname)
+      .map(([host, contexts]) => {
+        const candidates = [contexts.currentuser, contexts.system].filter(Boolean);
+        if (!candidates.length) return null;
+        const badge = (record) => {
+          if (!record) return "✕";
+          const st = (record.status || "").toString().toLowerCase();
+          if (st === "provisioned") return "✓";
+          return "•";
+        };
+        const label = `${host} (CURRENTUSER ${badge(contexts.currentuser)}, SYSTEM ${badge(contexts.system)})`;
+        const latest = Math.max(...candidates.map((r) => r.last_seen || 0));
+        return { host, label, contexts, latest };
+      })
+      .filter(Boolean)
+      .sort((a, b) => {
+        if (b.latest !== a.latest) return b.latest - a.latest;
+        return a.host.localeCompare(b.host);
+      });
+    return entries;
+  }, [agentsByHostname]);
+
  // Fetch Agents Periodically
  useEffect(() => {
    const fetchAgents = () => {
@@ -37,15 +76,79 @@ const BorealisAgentNode = ({ id, data }) => {
    return () => clearInterval(interval);
  }, []);

+  // Ensure host selection stays aligned with available agents
+  useEffect(() => {
+    const hostExists = hostOptions.some((opt) => opt.host === selectedHost);
+    if (hostExists) return;
+
+    if (selectedAgent && agents[selectedAgent]) {
+      const info = agents[selectedAgent];
+      const inferredHost = (info?.hostname || info?.agent_hostname || "").trim() || "unknown";
+      if (inferredHost && inferredHost !== selectedHost) {
+        setSelectedHost(inferredHost);
+        return;
+      }
+    }
+
+    const fallbackHost = hostOptions[0]?.host || "";
+    if (fallbackHost !== selectedHost) {
+      setSelectedHost(fallbackHost);
+    }
+    if (!fallbackHost && selectedAgent) {
+      setSelectedAgent("");
+    }
+  }, [hostOptions, selectedHost, selectedAgent, agents]);
+
+  // Align agent selection with host/mode choice
+  useEffect(() => {
+    if (!selectedHost) {
+      if (selectedMode !== "currentuser") setSelectedMode("currentuser");
+      if (selectedAgent) setSelectedAgent("");
+      return;
+    }
+    const contexts = agentsByHostname[selectedHost];
+    if (!contexts) {
+      if (selectedMode !== "currentuser") setSelectedMode("currentuser");
+      if (selectedAgent) setSelectedAgent("");
+      return;
+    }
+    if (!contexts[selectedMode]) {
+      const fallbackMode = contexts.currentuser
+        ? "currentuser"
+        : contexts.system
+        ? "system"
+        : selectedMode;
+      if (fallbackMode !== selectedMode) {
+        setSelectedMode(fallbackMode);
+        return;
+      }
+    }
+    const activeContext = contexts[selectedMode];
+    const targetAgentId = activeContext?.agent_id || "";
+    if (targetAgentId !== selectedAgent) {
+      setSelectedAgent(targetAgentId);
+    }
+  }, [selectedHost, selectedMode, agentsByHostname, selectedAgent]);
+
  // Sync node data with sidebar changes
  useEffect(() => {
    setNodes((nds) =>
      nds.map((n) =>
-        n.id === id ? { ...n, data: { ...n.data, agent_id: selectedAgent } } : n
+        n.id === id
+          ? {
+              ...n,
+              data: {
+                ...n.data,
+                agent_id: selectedAgent,
+                agent_host: selectedHost,
+                agent_mode: selectedMode,
+              },
+            }
+          : n
      )
    );
    setIsConnected(false);
-  }, [selectedAgent, setNodes, id]);
+  }, [selectedAgent, selectedHost, selectedMode, setNodes, id]);

  // Attached Roles logic
  const attachedRoleIds = useMemo(
@@ -109,11 +212,19 @@ const BorealisAgentNode = ({ id, data }) => {

  // Status Label
  const selectedAgentStatus = useMemo(() => {
-    if (!selectedAgent) return "Unassigned";
-    const agent = agents[selectedAgent];
-    if (!agent) return "Reconnecting...";
-    return agent.status === "provisioned" ? "Connected" : "Available";
-  }, [agents, selectedAgent]);
+    if (!selectedHost) return "Unassigned";
+    const contexts = agentsByHostname[selectedHost];
+    if (!contexts) return "Offline";
+    const activeContext = contexts[selectedMode];
+    if (!selectedAgent || !activeContext) return "Unavailable";
+    const status = (activeContext.status || "").toString().toLowerCase();
+    if (status === "provisioned") return "Connected";
+    if (status === "orphaned") return "Available";
+    if (!status) return "Available";
+    return status.charAt(0).toUpperCase() + status.slice(1);
+  }, [agentsByHostname, selectedHost, selectedMode, selectedAgent]);
+
+  const activeHostContexts = selectedHost ? agentsByHostname[selectedHost] : null;

  // Render (Sidebar handles config)
  return (
@@ -128,20 +239,44 @@ const BorealisAgentNode = ({ id, data }) => {

      <div className="borealis-node-header">Borealis Agent</div>
      <div className="borealis-node-content" style={{ fontSize: "9px" }}>
-        <label>Agent:</label>
+        <label>Device:</label>
        <select
-          value={selectedAgent}
-          onChange={(e) => setSelectedAgent(e.target.value)}
+          value={selectedHost}
+          onChange={(e) => setSelectedHost(e.target.value)}
          style={{ width: "100%", marginBottom: "6px", fontSize: "9px" }}
        >
          <option value="">-- Select --</option>
-          {agentList.map(({ id: aid, status }) => (
-            <option key={aid} value={aid}>
-              {aid} ({status})
+          {hostOptions.map(({ host, label }) => (
+            <option key={host} value={host}>
+              {label}
            </option>
          ))}
        </select>

+        <label>Agent Context:</label>
+        <select
+          value={selectedMode}
+          onChange={(e) => setSelectedMode(e.target.value)}
+          style={{ width: "100%", marginBottom: "6px", fontSize: "9px" }}
+          disabled={!selectedHost}
+        >
+          <option value="currentuser" disabled={!activeHostContexts?.currentuser}>
+            CURRENTUSER Agent
+          </option>
+          <option value="system" disabled={!activeHostContexts?.system}>
+            SYSTEM Agent
+          </option>
+        </select>
+
+        <div style={{ fontSize: "8px", color: "#aaa", marginBottom: "4px" }}>
+          Target Agent ID:{" "}
+          {selectedAgent ? (
+            <span style={{ color: "#eee" }}>{selectedAgent}</span>
+          ) : (
+            <span style={{ color: "#666" }}>none</span>
+          )}
+        </div>
+
        {isConnected ? (
          <button
            onClick={handleDisconnect}
@@ -180,6 +315,7 @@ Select and connect to a remote Borealis Agent.
 - Assign roles to agent dynamically by connecting "Agent Role" nodes.
 - Auto-provisions agent as role assignments change.
 - See live agent status and re-connect/disconnect easily.
+- Choose between CURRENTUSER and SYSTEM contexts for each device.
 `.trim(),
  content: "Select and manage an Agent with dynamic roles",
  component: BorealisAgentNode,
@@ -197,7 +333,7 @@ Select and connect to a remote Borealis Agent.
 This node represents an available Borealis Agent (Python client) you can control from your workflow.

 #### Features
- **Select** an agent from the list of online agents.
+- **Select** a device and agent context (CURRENTUSER vs SYSTEM).
 - **Connect/Disconnect** from the agent at any time.
 - **Attach roles** (by connecting "Agent Role" nodes to this node's output handle) to assign behaviors dynamically.
 - **Live status** shows if the agent is available, connected, or offline.