import grpc import time import os import agent_pb2 import agent_pb2_grpc import threading import subprocess import json import jwt import datetime import hmac import hashlib import queue import sys import platform from concurrent import futures from playwright.sync_api import sync_playwright SECRET_KEY = "cortex-secret-shared-key" class BaseSkill: """Interface for pluggable node capabilities.""" def execute(self, task, sandbox, on_complete, on_event=None): raise NotImplementedError def cancel(self, task_id): return False class ShellSkill(BaseSkill): """Default Skill: Executing shell commands.""" def __init__(self): self.processes = {} # task_id -> Popen self.lock = threading.Lock() def execute(self, task, sandbox, on_complete, on_event=None): try: cmd = task.payload_json allowed, status_msg = sandbox.verify(cmd) if not allowed: return on_complete(task.task_id, {"stderr": f"SANDBOX_VIOLATION: {status_msg}", "status": 2}, task.trace_id) print(f" [🐚] Executing Shell: {cmd}", flush=True) p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) with self.lock: self.processes[task.task_id] = p timeout = task.timeout_ms / 1000.0 if task.timeout_ms > 0 else None stdout, stderr = p.communicate(timeout=timeout) print(f" [🐚] Shell Done: {cmd} | Stdout Size: {len(stdout)}", flush=True) on_complete(task.task_id, {"stdout": stdout, "stderr": stderr, "status": 1 if p.returncode == 0 else 2}, task.trace_id) except subprocess.TimeoutExpired: self.cancel(task.task_id) on_complete(task.task_id, {"stderr": "TIMEOUT", "status": 2}, task.trace_id) except Exception as e: on_complete(task.task_id, {"stderr": str(e), "status": 2}, task.trace_id) finally: with self.lock: self.processes.pop(task.task_id, None) def cancel(self, task_id): with self.lock: p = self.processes.get(task_id) if p: print(f"[🛑] Killing Shell Process: {task_id}") p.kill() return True return False class BrowserSkill(BaseSkill): """The 'Antigravity Bridge': Persistent Browser Skill using a dedicated Actor thread.""" def __init__(self): self.task_queue = queue.Queue() self.sessions = {} # session_id -> { "context": Context, "page": Page } self.lock = threading.Lock() threading.Thread(target=self._browser_actor, daemon=True, name="BrowserActor").start() def _setup_listeners(self, sid, page, on_event): if not on_event: return page.on("console", lambda msg: on_event(agent_pb2.BrowserEvent( session_id=sid, console_msg=agent_pb2.ConsoleMessage( level=msg.type, text=msg.text, timestamp_ms=int(time.time()*1000) ) ))) page.on("requestfinished", lambda req: on_event(agent_pb2.BrowserEvent( session_id=sid, network_req=agent_pb2.NetworkRequest( method=req.method, url=req.url, status=req.response().status if req.response() else 0, resource_type=req.resource_type, latency_ms=0 ) ))) def _browser_actor(self): print("[🌐] Browser Actor Starting...", flush=True) try: pw = sync_playwright().start() browser = pw.chromium.launch(headless=True, args=[ '--no-sandbox', '--disable-setuid-sandbox', '--disable-dev-shm-usage', '--disable-gpu' ]) print("[🌐] Browser Engine Online.", flush=True) except Exception as e: print(f"[!] Browser Actor Startup Fail: {e}", flush=True) return while True: try: task, sandbox, on_complete, on_event = self.task_queue.get() action = task.browser_action sid = action.session_id or "default" with self.lock: if sid not in self.sessions: context = browser.new_context() page = context.new_page() self._setup_listeners(sid, page, on_event) self.sessions[sid] = {"context": context, "page": page} page = self.sessions[sid]["page"] print(f" [🌐] Browser Actor Processing: {agent_pb2.BrowserAction.ActionType.Name(action.action)} | Session: {sid}", flush=True) res_data = {} if action.action == agent_pb2.BrowserAction.NAVIGATE: page.goto(action.url, wait_until="commit") elif action.action == agent_pb2.BrowserAction.CLICK: page.click(action.selector) elif action.action == agent_pb2.BrowserAction.TYPE: page.fill(action.selector, action.text) elif action.action == agent_pb2.BrowserAction.SCREENSHOT: res_data["snapshot"] = page.screenshot() elif action.action == agent_pb2.BrowserAction.GET_DOM: res_data["dom_content"] = page.content() elif action.action == agent_pb2.BrowserAction.HOVER: page.hover(action.selector) elif action.action == agent_pb2.BrowserAction.SCROLL: page.mouse.wheel(x=0, y=action.y) elif action.action == agent_pb2.BrowserAction.EVAL: res_data["eval_result"] = str(page.evaluate(action.text)) elif action.action == agent_pb2.BrowserAction.GET_A11Y: res_data["a11y_tree"] = json.dumps(page.accessibility.snapshot()) elif action.action == agent_pb2.BrowserAction.CLOSE: with self.lock: sess = self.sessions.pop(sid, None) if sess: sess["context"].close() # Refresh metadata After br_res = agent_pb2.BrowserResponse( url=page.url, title=page.title(), snapshot=res_data.get("snapshot", b""), dom_content=res_data.get("dom_content", ""), a11y_tree=res_data.get("a11y_tree", ""), eval_result=res_data.get("eval_result", "") ) on_complete(task.task_id, {"status": 1, "browser_result": br_res}, task.trace_id) except Exception as e: print(f" [!] Browser Actor Error: {e}", flush=True) on_complete(task.task_id, {"stderr": str(e), "status": 2}, task.trace_id) def execute(self, task, sandbox, on_complete, on_event=None): self.task_queue.put((task, sandbox, on_complete, on_event)) def cancel(self, task_id): return False class SkillManager: """Orchestrates multiple skills and manages the worker thread pool.""" def __init__(self, max_workers=5): self.executor = futures.ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="skill-worker") self.active_tasks = {} # task_id -> future self.skills = { "shell": ShellSkill(), "browser": BrowserSkill() } self.max_workers = max_workers self.lock = threading.Lock() def submit(self, task, sandbox, on_complete, on_event=None): with self.lock: if len(self.active_tasks) >= self.max_workers: return False, "Node Capacity Reached" # Decide Skill if task.HasField("browser_action") or task.task_type == "browser": skill = self.skills["browser"] else: skill = self.skills["shell"] future = self.executor.submit(skill.execute, task, sandbox, on_complete, on_event) self.active_tasks[task.task_id] = future # Cleanup hook future.add_done_callback(lambda f: self._cleanup(task.task_id)) return True, "Accepted" def cancel(self, task_id): with self.lock: # Tell all skills to try and cancel this ID cancelled = any(s.cancel(task_id) for s in self.skills.values()) return cancelled def get_active_ids(self): with self.lock: return list(self.active_tasks.keys()) def _cleanup(self, task_id): with self.lock: self.active_tasks.pop(task_id, None) class SandboxEngine: def __init__(self): self.policy = None def sync(self, p): self.policy = {"MODE": "STRICT" if p.mode == agent_pb2.SandboxPolicy.STRICT else "PERMISSIVE", "ALLOWED": list(p.allowed_commands), "DENIED": list(p.denied_commands), "SENSITIVE": list(p.sensitive_commands)} def verify(self, command_str): if not self.policy: return False, "No Policy" parts = (command_str or "").strip().split() if not parts: return False, "Empty" base_cmd = parts[0] if base_cmd in self.policy["DENIED"]: return False, "Forbidden" if self.policy["MODE"] == "STRICT" and base_cmd not in self.policy["ALLOWED"]: return False, "Not Whitelisted" return True, "OK" class AgentNode: def __init__(self, node_id="agent-node-007"): self.node_id = node_id self.skills = SkillManager() self.sandbox = SandboxEngine() self.task_queue = queue.Queue() # gRPC Setup with open('certs/client.key', 'rb') as f: pkey = f.read() with open('certs/client.crt', 'rb') as f: cert = f.read() with open('certs/ca.crt', 'rb') as f: ca = f.read() creds = grpc.ssl_channel_credentials(ca, pkey, cert) self.channel = grpc.secure_channel('localhost:50051', creds) self.stub = agent_pb2_grpc.AgentOrchestratorStub(self.channel) def _create_token(self): return jwt.encode({"sub": self.node_id, "iat": datetime.datetime.utcnow(), "exp": datetime.datetime.utcnow() + datetime.timedelta(minutes=10)}, SECRET_KEY, algorithm="HS256") def sync_configuration(self): print(f"[*] Handshake: {self.node_id}") reg = agent_pb2.RegistrationRequest(node_id=self.node_id, auth_token=self._create_token(), node_description="Refactored Stateful Node with Browser Skill", capabilities={"shell": "v1", "browser": "playwright-1.42"}) res = self.stub.SyncConfiguration(reg) if res.success: self.sandbox.sync(res.policy); print("[OK] Policy Synced.") else: print(f"[!] Rejected: {res.error_message}"); sys.exit(1) def start_health_reporting(self): def _gen(): while True: ids = self.skills.get_active_ids() yield agent_pb2.Heartbeat(node_id=self.node_id, cpu_usage_percent=1.0, active_worker_count=len(ids), max_worker_capacity=5, running_task_ids=ids) time.sleep(10) threading.Thread(target=lambda: list(self.stub.ReportHealth(_gen())), daemon=True).start() def run_task_stream(self): def _gen(): yield agent_pb2.ClientTaskMessage(announce=agent_pb2.NodeAnnounce(node_id=self.node_id)) while True: yield self.task_queue.get() responses = self.stub.TaskStream(_gen()) print(f"[*] Stream processing started: {self.node_id}", flush=True) try: for msg in responses: kind = msg.WhichOneof('payload') print(f"[*] Received message from server: {kind}", flush=True) if kind == 'task_request': self._handle_task(msg.task_request) elif kind == 'task_cancel': if self.skills.cancel(msg.task_cancel.task_id): self.task_queue.put(agent_pb2.ClientTaskMessage(task_response=agent_pb2.TaskResponse( task_id=msg.task_cancel.task_id, status=agent_pb2.TaskResponse.CANCELLED))) elif kind == 'work_pool_update': for tid in msg.work_pool_update.available_task_ids: if len(self.skills.get_active_ids()) < self.skills.max_workers: self.task_queue.put(agent_pb2.ClientTaskMessage(task_claim=agent_pb2.TaskClaimRequest(task_id=tid, node_id=self.node_id))) except Exception as e: print(f"[!] Stream Error: {e}", flush=True) def _handle_task(self, task): print(f"[*] Handling Task: {task.task_id}", flush=True) # Sig Verify logic based on payload type if task.HasField("browser_action"): a = task.browser_action sign_base = f"{a.action}:{a.url}:{a.session_id}".encode() else: sign_base = task.payload_json.encode() expected_sig = hmac.new(SECRET_KEY.encode(), sign_base, hashlib.sha256).hexdigest() if not hmac.compare_digest(task.signature, expected_sig): return print(f"[!] Sig Fail for {task.task_id} | Raw: {sign_base}", flush=True) print(f"[✅] Signature Verified for {task.task_id}", flush=True) self.skills.submit(task, self.sandbox, self._on_finish, self._on_event) def _on_event(self, event): self.task_queue.put(agent_pb2.ClientTaskMessage(browser_event=event)) def _on_finish(self, tid, res, trace): print(f"[*] Task {tid} finished.", flush=True) status = agent_pb2.TaskResponse.SUCCESS if res['status'] == 1 else agent_pb2.TaskResponse.ERROR tr = agent_pb2.TaskResponse( task_id=tid, status=status, stdout=res.get('stdout',''), stderr=res.get('stderr',''), trace_id=trace, browser_result=res.get("browser_result") ) self.task_queue.put(agent_pb2.ClientTaskMessage(task_response=tr)) if __name__ == '__main__': node = AgentNode() node.sync_configuration() node.start_health_reporting() node.run_task_stream()