Janady13/AI_Brain / sub_brains/money_agents/worker.py

1#!/usr/bin/env python3 2import json 3import logging 4import os 5import re 6import signal 7import subprocess 8import sys 9import time 10from datetime import datetime 11from typing import Optional 12 13import psycopg 14from psycopg.rows import dict_row 15 16DB_URL = os.environ.get("DATABASE_URL") 17if not DB_URL: 18 print("ERR: DATABASE_URL not set", file=sys.stderr) 19 sys.exit(1) 20 21WORKER_ID = int(os.environ.get("WORKER_ID", str(os.getpid()))) 22LOGS_DIR = os.environ.get("LOGS_DIR", "/Volumes/Omnicron/AI_Brain/logs") 23os.makedirs(LOGS_DIR, exist_ok=True) 24FILE_LOG = os.path.join(LOGS_DIR, f"money_agents_worker_{WORKER_ID}.out") 25 26logger = logging.getLogger("money_agents.worker") 27logger.setLevel(logging.INFO) 28fh = logging.FileHandler(FILE_LOG) 29fh.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s")) 30logger.addHandler(fh) 31sh = logging.StreamHandler(sys.stdout) 32sh.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s")) 33logger.addHandler(sh) 34 35 36def db_connect(): 37 return psycopg.connect(DB_URL) 38 39 40def db_log(conn, level: str, msg: str, job_id: Optional[int] = None): 41 with conn.cursor() as cur: 42 cur.execute( 43 "INSERT INTO logs(worker_id, job_id, message, level) VALUES (%s,%s,%s,%s)", 44 (WORKER_ID, job_id, msg, level), 45 ) 46 conn.commit() 47 logger.info("[log:%s] job=%s %s", level, str(job_id), msg) 48 49 50ALLOWED_EXEC = { 51 # domain allowlist enforced via regex match below 52 "curl": re.compile( 53 r"^curl\s+https://(api\.(stripe|paypal)\.com|cash\.app|api\.cashapp\.com|[^\s]+)" 54 ), 55} 56 57 58def safe_execute(payload: str) -> int: 59 payload = payload.strip() 60 if not payload: 61 return 0 62 parts = payload.split() 63 cmd = parts[0] 64 if cmd not in ALLOWED_EXEC: 65 logger.warning("Blocked command: %s", cmd) 66 return 126 67 if not ALLOWED_EXEC[cmd].match(payload): 68 logger.warning("Blocked payload by allowlist policy") 69 return 126 70 try: 71 # Avoid printing secrets: no shell expansion, no env dump 72 res = subprocess.run(parts, capture_output=True, timeout=60) 73 if res.returncode != 0: 74 logger.error("payload stderr: %s", res.stderr.decode(errors="ignore")[:500]) 75 else: 76 logger.info("payload stdout: %s", res.stdout.decode(errors="ignore")[:200]) 77 return res.returncode 78 except Exception as e: 79 logger.exception("payload execution error: %s", e) 80 return 1 81 82 83def extract_service(payload: str) -> Optional[str]: 84 m = re.search(r"https://([^/\s]+)", payload or "") 85 return m.group(1) if m else None 86 87 88STOP = False 89 90 91def handle_sigterm(signum, frame): 92 global STOP 93 STOP = True 94 95 96signal.signal(signal.SIGTERM, handle_sigterm) 97signal.signal(signal.SIGINT, handle_sigterm) 98 99 100def main(): 101 poll_sleep = float(os.environ.get("POLL_INTERVAL", "1.5")) 102 heartbeat_every = 10 # seconds 103 last_heartbeat = 0.0 104 with db_connect() as conn: 105 conn.autocommit = False 106 # initial heartbeat 107 db_log(conn, "info", "worker starting", None) 108 while not STOP: 109 now = time.time() 110 if now - last_heartbeat > heartbeat_every: 111 db_log(conn, "debug", "heartbeat", None) 112 last_heartbeat = now 113 114 job = None 115 try: 116 with conn.cursor(row_factory=dict_row) as cur: 117 cur.execute("BEGIN") 118 cur.execute( 119 "SELECT id, payload, amount FROM jobs WHERE status='pending' FOR UPDATE SKIP LOCKED LIMIT 1" 120 ) 121 row = cur.fetchone() 122 if row: 123 job = row 124 cur.execute( 125 "UPDATE jobs SET status='claimed', worker_id=%s, started_at=NOW() WHERE id=%s", 126 (WORKER_ID, job["id"]), 127 ) 128 cur.execute("COMMIT") 129 db_log(conn, "info", "claimed job", job["id"]) 130 else: 131 cur.execute("ROLLBACK") 132 except Exception as e: 133 conn.rollback() 134 logger.exception("claim failure: %s", e) 135 time.sleep(poll_sleep) 136 continue 137 138 if not job: 139 time.sleep(poll_sleep) 140 continue 141 142 job_id = job["id"] 143 payload = (job.get("payload") or "").strip() 144 amount = job.get("amount") or 0 145 service = extract_service(payload) or "unknown" 146 147 try: 148 db_log(conn, "info", "executed payload", job_id) 149 rc = safe_execute(payload) 150 status = "success" if rc == 0 else "error" 151 152 if amount and float(amount) != 0.0: 153 with conn.cursor() as cur: 154 cur.execute( 155 "INSERT INTO transactions(job_id, amount, service, status) VALUES (%s,%s,%s,%s)", 156 (job_id, amount, service, status), 157 ) 158 conn.commit() 159 db_log( 160 conn, 161 "info", 162 f"transaction status={status} amount={amount} service={service}", 163 job_id, 164 ) 165 166 with conn.cursor() as cur: 167 cur.execute( 168 "UPDATE jobs SET status=%s, finished_at=NOW() WHERE id=%s", 169 ("completed" if rc == 0 else "error", job_id), 170 ) 171 conn.commit() 172 db_log(conn, "info", "completed" if rc == 0 else "error", job_id) 173 except Exception as e: 174 conn.rollback() 175 logger.exception("job failure: %s", e) 176 try: 177 with db_connect() as c2: 178 with c2, c2.cursor() as cur: 179 cur.execute( 180 "UPDATE jobs SET status='error', finished_at=NOW() WHERE id=%s", 181 (job_id,), 182 ) 183 db_log(c2, "error", f"exception: {e}", job_id) 184 except Exception: 185 pass 186 finally: 187 time.sleep(0.2) 188 189 190if __name__ == "__main__": 191 main() 192