Janady13/AI_Brain / skynet/engine.py

1#!/usr/bin/env python3 2from __future__ import annotations 3 4import os 5import time 6import socket 7import threading 8from typing import Any, Callable, Dict, List, Optional 9 10try: 11 # Prefer real Omnicron interface if available 12 from omnicron import iface as omni # type: ignore 13except Exception: # pragma: no cover 14 # Fallback to local shim 15 from infra import omnicron_iface as omni # type: ignore 16from skynet.agents_api_client import AgentsApiClient, AgentsApiError 17from skynet import handlers 18 19 20Handler = Callable[[], Dict[str, Any]] 21 22 23DEFAULT_HANDLERS: Dict[str, Handler] = { 24 "planner.run": handlers.handle_planner_run, 25 "alignment.evaluate": handlers.handle_alignment_evaluate, 26 "metabrain.consolidate": handlers.handle_metabrain_consolidate, 27 "executor.run": handlers.handle_executor_run, 28 "github.sync": handlers.handle_github_sync, 29 # Revenue-oriented kinds 30 "marketing.publish": handlers.handle_marketing_publish, 31 "affiliate.publish": handlers.handle_affiliate_publish, 32 "commerce.cashapp.promote": handlers.handle_commerce_cashapp_promote, 33 "commerce.venmo.promote": handlers.handle_commerce_venmo_promote, 34 "revenue.stripe.refresh_keys": handlers.handle_revenue_stripe_refresh_keys, 35 "revenue.stripe.create_checkout": handlers.handle_revenue_stripe_checkout, 36 "revenue.stripe.consume": handlers.handle_revenue_stripe_consume, 37} 38 39 40class SkynetEngine: 41 def __init__(self, kinds: Optional[List[str]] = None, agent_id: Optional[str] = None, client: Optional[AgentsApiClient] = None, registry: Optional[Dict[str, Handler]] = None) -> None: 42 self.kinds = kinds or list(DEFAULT_HANDLERS.keys()) 43 self.agent_id = agent_id or os.getenv("AGENT_ID") or os.getenv("HOSTNAME") or "skynet-engine" 44 self.client = client or AgentsApiClient() 45 self.registry: Dict[str, Handler] = registry or DEFAULT_HANDLERS 46 self._hb_stop = threading.Event() 47 self._hb_thread: Optional[threading.Thread] = None 48 49 def _start_presence_heartbeat(self, interval_s: float = 30.0) -> None: 50 hostname = socket.gethostname() 51 ip = os.getenv("TS_IP") or None 52 kinds = self.kinds 53 tags_env = os.getenv("AGENT_TAGS") or os.getenv("SKYNET_TAGS") or "skynet" 54 tags = [t.strip() for t in tags_env.split(",") if t.strip()] 55 def loop() -> None: 56 while not self._hb_stop.is_set(): 57 try: 58 self.client.agents_heartbeat(agent=self.agent_id, hostname=hostname, ip=ip, kinds=kinds, meta={"tags": tags}) 59 except Exception as exc: 60 omni.emit_log("skynet_engine_presence_error", level="debug", payload={"error": str(exc)}) 61 self._hb_stop.wait(interval_s) 62 t = threading.Thread(target=loop, name="skynet-presence", daemon=True) 63 t.start() 64 self._hb_thread = t 65 66 def run_forever(self, idle_sleep_s: float = 10.0) -> None: 67 omni.emit_log("skynet_engine_start", payload={"kinds": self.kinds, "agent": self.agent_id}) 68 self._start_presence_heartbeat() 69 while True: 70 try: 71 claim = self.client.claim(agent=self.agent_id, kinds=self.kinds) 72 except AgentsApiError as e: 73 # No tasks or temporary failure 74 omni.emit_log("skynet_engine_idle", level="debug", payload={"error": str(e)}) 75 time.sleep(idle_sleep_s) 76 continue 77 78 task = claim.get("task") or claim 79 tid = task.get("id") 80 kind = task.get("kind") or "" 81 handler = self.registry.get(kind) 82 if not handler: 83 omni.emit_log("skynet_engine_unknown_kind", level="warn", payload={"kind": kind, "task_id": tid}) 84 # Mark failed to avoid deadlocks 85 try: 86 self.client.fail(tid, self.agent_id, f"unknown kind: {kind}") 87 except Exception: 88 pass 89 continue 90 91 try: 92 self.client.heartbeat(tid, self.agent_id, status="running") 93 result = handler() or {} 94 self.client.complete(tid, self.agent_id, result=result, message="ok", return_code=0) 95 omni.emit_log("skynet_engine_task_done", payload={"task_id": tid, "kind": kind}) 96 except Exception as exc: 97 omni.emit_log("skynet_engine_task_error", level="warn", payload={"task_id": tid, "kind": kind, "error": str(exc)}) 98 try: 99 self.client.fail(tid, self.agent_id, str(exc)) 100 except Exception: 101 pass 102 103 104def run() -> None: # convenient entrypoint 105 kinds_env = os.getenv("SKYNET_KINDS", "").strip() 106 kinds = [k.strip() for k in kinds_env.split(",") if k.strip()] if kinds_env else None 107 SkynetEngine(kinds=kinds).run_forever() 108