import grpc
import os
from concurrent import futures
import time
import agent_pb2
import agent_pb2_grpc
import threading
import queue
import jwt
import hmac
import hashlib
import json
SECRET_KEY = "cortex-secret-shared-key"
class TaskJournal:
"""State machine for tracking tasks through their lifecycle."""
def __init__(self):
self.lock = threading.Lock()
self.tasks = {} # task_id -> { "event": Event, "result": None, "node_id": str }
def register(self, task_id, node_id=None):
event = threading.Event()
with self.lock:
self.tasks[task_id] = {"event": event, "result": None, "node_id": node_id}
return event
def fulfill(self, task_id, result):
with self.lock:
if task_id in self.tasks:
self.tasks[task_id]["result"] = result
self.tasks[task_id]["event"].set()
return True
return False
def get_result(self, task_id):
with self.lock:
data = self.tasks.get(task_id)
return data["result"] if data else None
def pop(self, task_id):
with self.lock:
return self.tasks.pop(task_id, None)
class AbstractNodeRegistry:
"""Interface for finding and tracking Managers."""
def register(self, node_id, data): raise NotImplementedError
def update_stats(self, node_id, stats): raise NotImplementedError
def get_best(self): raise NotImplementedError
def get_node(self, node_id): raise NotImplementedError
class MemoryNodeRegistry(AbstractNodeRegistry):
def __init__(self):
self.lock = threading.Lock()
self.nodes = {} # node_id -> { stats: {}, queue: queue, metadata: {} }
def register(self, node_id, q, metadata):
with self.lock:
self.nodes[node_id] = {"stats": {}, "queue": q, "metadata": metadata}
print(f"[📋] Registered: {node_id}")
def update_stats(self, node_id, stats):
with self.lock:
if node_id in self.nodes: self.nodes[node_id]["stats"].update(stats)
def get_best(self):
with self.lock:
if not self.nodes: return None
# Pick based on active worker count
return sorted(self.nodes.items(), key=lambda x: x[1]["stats"].get("active_worker_count", 999))[0][0]
def get_node(self, node_id):
with self.lock: return self.nodes.get(node_id)
class GlobalWorkPool:
def __init__(self):
self.lock = threading.Lock()
self.available = {"shared-001": "uname -a", "shared-002": "uptime"}
def claim(self, task_id, node_id):
with self.lock:
if task_id in self.available:
return True, self.available.pop(task_id)
return False, None
class TaskAssistant:
"""The High-Level AI API."""
def __init__(self, registry, journal, pool):
self.registry = registry
self.journal = journal
self.pool = pool
def dispatch_single(self, node_id, cmd, timeout=30):
# Implementation of retry logic and signing
node = self.registry.get_node(node_id)
if not node: return {"error": "Offline"}
tid = f"task-{int(time.time()*1000)}"
event = self.journal.register(tid, node_id)
# Don't wrap in JSON, use raw cmd for shells
msg_json = json.dumps({"command": cmd}) # Simulation legacy check
sig = hmac.new(SECRET_KEY.encode(), cmd.encode(), hashlib.sha256).hexdigest()
req = agent_pb2.ServerTaskMessage(task_request=agent_pb2.TaskRequest(
task_id=tid, payload_json=cmd, signature=sig))
print(f"[📤] Dispatching {tid} to {node_id}")
node["queue"].put(req)
if event.wait(timeout):
res = self.journal.get_result(tid)
self.journal.pop(tid)
return res
self.journal.pop(tid)
return {"error": "Timeout"}
def dispatch_browser(self, node_id, action, timeout=60):
node = self.registry.get_node(node_id)
if not node: return {"error": "Offline"}
tid = f"br-{int(time.time()*1000)}"
event = self.journal.register(tid, node_id)
# Basic signature for POC: Sign the action enum name + URL
sign_base = f"{action.action}:{action.url}:{action.session_id}"
sig = hmac.new(SECRET_KEY.encode(), sign_base.encode(), hashlib.sha256).hexdigest()
req = agent_pb2.ServerTaskMessage(task_request=agent_pb2.TaskRequest(
task_id=tid, browser_action=action, signature=sig))
node["queue"].put(req)
if event.wait(timeout):
res = self.journal.get_result(tid)
self.journal.pop(tid)
return res
self.journal.pop(tid)
return {"error": "Timeout"}
class AgentOrchestrator(agent_pb2_grpc.AgentOrchestratorServicer):
def __init__(self):
self.registry = MemoryNodeRegistry()
self.journal = TaskJournal()
self.pool = GlobalWorkPool()
self.assistant = TaskAssistant(self.registry, self.journal, self.pool)
def SyncConfiguration(self, request, context):
# Pre-registration for metadata search
self.registry.register(request.node_id, queue.Queue(), {"desc": request.node_description, "caps": dict(request.capabilities)})
return agent_pb2.RegistrationResponse(success=True,
policy=agent_pb2.SandboxPolicy(mode=agent_pb2.SandboxPolicy.STRICT,
allowed_commands=["ls", "uname", "echo", "sleep"]))
def TaskStream(self, request_iterator, context):
try:
# 1. Blocking wait for identity
first_msg = next(request_iterator)
if first_msg.WhichOneof('payload') != 'announce':
print("[!] Stream rejected: No NodeAnnounce")
return
node_id = first_msg.announce.node_id
node = self.registry.get_node(node_id)
if not node:
print(f"[!] Stream rejected: Node {node_id} not registered via SyncConfiguration")
return
print(f"[📶] Stream established for {node_id}")
# 2. Results Listener
def _read_results():
for msg in request_iterator:
kind = msg.WhichOneof('payload')
if kind == 'task_claim':
success, payload = self.pool.claim(msg.task_claim.task_id, node_id)
if success:
sig = hmac.new(SECRET_KEY.encode(), payload.encode(), hashlib.sha256).hexdigest()
node["queue"].put(agent_pb2.ServerTaskMessage(
task_request=agent_pb2.TaskRequest(task_id=msg.task_claim.task_id, payload_json=payload, signature=sig)))
elif kind == 'task_response':
res_obj = {"stdout": msg.task_response.stdout, "status": msg.task_response.status}
if msg.task_response.HasField("browser_result"):
br = msg.task_response.browser_result
res_obj["browser"] = {
"url": br.url, "title": br.title, "has_snapshot": len(br.snapshot) > 0,
"a11y": br.a11y_tree[:100] + "..." if br.a11y_tree else None,
"eval": br.eval_result
}
self.journal.fulfill(msg.task_response.task_id, res_obj)
elif kind == 'browser_event':
e = msg.browser_event
if e.HasField("console_msg"):
print(f" [🖥️] Live Browser Console: {e.console_msg.text}", flush=True)
elif e.HasField("network_req"):
print(f" [🌐] Live Network Request: {e.network_req.method} {e.network_req.url}", flush=True)
threading.Thread(target=_read_results, daemon=True).start()
# 3. Work Dispatcher (Main Stream)
while context.is_active():
try:
msg = node["queue"].get(timeout=1.0)
yield msg
print(f"[🚀] Pushed message from queue to stream: {node_id}")
except queue.Empty:
continue
except StopIteration:
pass
except Exception as e:
print(f"[!] TaskStream Error: {e}")
# Broadcast pool
if self.pool.available:
yield agent_pb2.ServerTaskMessage(work_pool_update=agent_pb2.WorkPoolUpdate(available_task_ids=list(self.pool.available.keys())))
# Send direct tasks
if node_id and self.registry.get_node(node_id):
try:
msg = self.registry.get_node(node_id)["queue"].get(timeout=2)
yield msg
except queue.Empty: pass
else: time.sleep(1)
def ReportHealth(self, request_iterator, context):
for hb in request_iterator:
self.registry.update_stats(hb.node_id, {"active_worker_count": hb.active_worker_count, "running": list(hb.running_task_ids)})
yield agent_pb2.HealthCheckResponse(server_time_ms=int(time.time()*1000))
def serve():
with open('certs/server.key', 'rb') as f: pkey = f.read()
with open('certs/server.crt', 'rb') as f: cert = f.read()
with open('certs/ca.crt', 'rb') as f: ca = f.read()
creds = grpc.ssl_server_credentials([(pkey, cert)], ca, True)
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
orch = AgentOrchestrator()
agent_pb2_grpc.add_AgentOrchestratorServicer_to_server(orch, server)
server.add_secure_port('[::]:50051', creds)
print("[🛡️] Boss Plane Refactored & Online.")
server.start()
# Simple AI Simulation loop
time.sleep(10)
print("\n[🧠] AI Simulation Start...", flush=True)
res_single = orch.assistant.dispatch_single('agent-node-007', 'uname -a')
print(f" Uname: {res_single}", flush=True)
# NEW: Browser Phase
print("\n[🧠] AI Phase 4: Navigating Browser (Antigravity Bridge)...")
nav_action = agent_pb2.BrowserAction(
action=agent_pb2.BrowserAction.NAVIGATE,
url="https://example.com",
session_id="antigravity-session-1"
)
res_nav = orch.assistant.dispatch_browser("agent-node-007", nav_action)
print(f" Nav Result: {res_nav}")
print("\n[🧠] AI Phase 5: Multi-Action Persistence (Screenshot)...")
snap_action = agent_pb2.BrowserAction(
action=agent_pb2.BrowserAction.SCREENSHOT,
session_id="antigravity-session-1"
)
res_snap = orch.assistant.dispatch_browser("agent-node-007", snap_action)
print(f" Snap Result: {res_snap.get('browser', {}).get('title')} | Snapshot captured: {res_snap.get('browser', {}).get('has_snapshot')}")
# NEW: Phase 4 Pro Features
print("\n[🧠] AI Phase 4 Pro: Perception & Advanced Logic...")
a11y_action = agent_pb2.BrowserAction(
action=agent_pb2.BrowserAction.GET_A11Y,
session_id="antigravity-session-1"
)
res_a11y = orch.assistant.dispatch_browser("agent-node-007", a11y_action)
print(f" A11y Result: {res_a11y.get('browser', {}).get('a11y')}")
eval_action = agent_pb2.BrowserAction(
action=agent_pb2.BrowserAction.EVAL,
text="window.performance.now()",
session_id="antigravity-session-1"
)
res_eval = orch.assistant.dispatch_browser("agent-node-007", eval_action)
print(f" Eval Result (Timestamp): {res_eval.get('browser', {}).get('eval')}")
# NEW: Phase 4 Pro Features - Real-time Events
print("\n[🧠] AI Phase 4 Pro: Triggering Real-time Events (Tunneling)...")
trigger_action = agent_pb2.BrowserAction(
action=agent_pb2.BrowserAction.EVAL,
text="console.log('Hello from Antigravity Bridge!'); fetch('https://example.com/api/ping');",
session_id="antigravity-session-1"
)
orch.assistant.dispatch_browser("agent-node-007", trigger_action)
server.wait_for_termination()
if __name__ == '__main__':
serve()
