# atomic_yaml.py — atomic YAML/DB update logic # Extracted from lib.sh atomic_dump_yaml PYEOF block import os, sys, tempfile, shutil, glob, subprocess, json, sqlite3, fcntl from datetime import datetime, timezone import yaml def atomic_dump_yaml_main(): yaml_path = os.environ['YAML_PATH'] db_path = os.path.splitext(yaml_path)[0] + '.db' def _validate(d): if not isinstance(d, dict): raise SystemExit("VALIDATE: top-level is not a mapping") sessions = d.get('herdr_sessions', []) if not isinstance(sessions, list): raise SystemExit("VALIDATE: herdr_sessions is not a list") valid = {'running', 'terminated', 'archived', 'stopped'} for i, s in enumerate(sessions): if not isinstance(s, dict): raise SystemExit(f"VALIDATE: herdr_sessions[{i}] not a mapping") if not s.get('name') or not s.get('status'): raise SystemExit(f"VALIDATE: herdr_sessions[{i}] missing name/status") if s.get('role') is not None and (not isinstance(s['role'], str) or not s['role'].strip()): raise SystemExit(f"VALIDATE: herdr_sessions[{i}] {s.get('name')!r} role must be a non-empty string") if s['status'] not in valid: raise SystemExit(f"VALIDATE: herdr_sessions[{i}] {s.get('name')!r} bad status {s['status']!r}") if not isinstance(s.get('pane'), dict): raise SystemExit(f"VALIDATE: herdr_sessions[{i}] {s.get('name')!r} missing pane") iso = s.get('isolation') if iso is not None: if not isinstance(iso, dict) or not iso.get('uuid') or not iso.get('root'): raise SystemExit(f"VALIDATE: herdr_sessions[{i}] {s.get('name')!r} isolation block requires uuid/root") orc_uuids = d.get('orchestrator_uuids') if orc_uuids is not None: if not isinstance(orc_uuids, list): raise SystemExit("VALIDATE: orchestrator_uuids is not a list") seen_orc = set() for u in orc_uuids: if not isinstance(u, str) or not u.strip(): raise SystemExit("VALIDATE: orchestrator_uuids items must be non-empty strings") if u in seen_orc: raise SystemExit("VALIDATE: orchestrator_uuids contains duplicate item") seen_orc.add(u) def get_terminal_set(d): return {s.get('name'): s.get('status') for s in d.get('herdr_sessions', []) if s.get('status') in ('stopped', 'terminated', 'archived')} def get_all_sessions_status(d): import hashlib res = {} for s in d.get('herdr_sessions', []): name = s.get('name') ser = json.dumps(s, sort_keys=True) res[name] = hashlib.sha256(ser.encode('utf-8')).hexdigest() orc_uuids = d.get('orchestrator_uuids') if orc_uuids is not None: ser_orc = json.dumps(orc_uuids, sort_keys=True) res['__orchestrator_uuids__'] = hashlib.sha256(ser_orc.encode('utf-8')).hexdigest() return res os.makedirs(os.path.dirname(db_path) or '.', exist_ok=True) _flock_f = open(db_path + '.lock', 'w') fcntl.flock(_flock_f, fcntl.LOCK_EX) conn = sqlite3.connect(db_path, timeout=60.0) conn.execute('PRAGMA busy_timeout = 60000') for f in [db_path, db_path + '-wal', db_path + '-shm']: if os.path.exists(f): try: os.chmod(f, 0o600) except Exception: pass is_nfs = os.environ.get('MAM_IS_NFS') == 'true' if is_nfs: conn.execute('PRAGMA journal_mode=DELETE') else: conn.execute('PRAGMA journal_mode=WAL') try: for _beg_attempt in range(300): try: conn.execute('BEGIN IMMEDIATE') break except sqlite3.OperationalError: import time time.sleep(0.1) else: conn.execute('BEGIN IMMEDIATE') conn.execute('CREATE TABLE IF NOT EXISTS state (id INTEGER PRIMARY KEY, data TEXT)') conn.execute('CREATE TABLE IF NOT EXISTS sessions (name TEXT PRIMARY KEY, status TEXT, pane_cwd TEXT, data JSON)') conn.execute('CREATE INDEX IF NOT EXISTS idx_sessions_pane_cwd ON sessions(pane_cwd)') row = conn.execute('SELECT data FROM state WHERE id=1').fetchone() if row: d = json.loads(row[0]) else: if os.path.exists(yaml_path): with open(yaml_path) as f: d = yaml.safe_load(f) or {} else: d = {} db_sessions = [] cursor = conn.execute('SELECT name, status, pane_cwd, data FROM sessions') for s_row in cursor.fetchall(): s_data = json.loads(s_row[3]) s_data['name'] = s_row[0] s_data['status'] = s_row[1] if 'pane' not in s_data: s_data['pane'] = {} s_data['pane']['cwd'] = s_row[2] db_sessions.append(s_data) if db_sessions: d['herdr_sessions'] = db_sessions elif 'herdr_sessions' not in d: d['herdr_sessions'] = [] old_sessions = get_all_sessions_status(d) old_roles = {s.get('name'): s.get('role') for s in db_sessions if s.get('role')} # --- caller mutation (module scope: sees d, yaml, os, glob, subprocess) --- mutation_ns = dict(globals()) mutation_ns.update(locals()) exec(compile(os.environ['AGENT_SESSIONS_MUTATION'], '', 'exec'), mutation_ns) if 'd' in mutation_ns: d = mutation_ns['d'] for s in d.get('herdr_sessions', []): name = s.get('name') if name in old_roles and s.get('role') != old_roles[name]: raise SystemExit(f"VALIDATE: role of session {name!r} cannot be modified from {old_roles[name]!r} to {s.get('role')!r}") running_keys = ['claude_session_id_own', 'agy_conversation_id_own', 'hermes_conversation_id_own', 'cline_conversation_id_own'] id_to_session = {} for s in d.get('herdr_sessions', []): if s.get('status') == 'running': s_name = s.get('name') for k in running_keys: v = s.get(k) if v: if v in id_to_session and id_to_session[v] != s_name: raise SystemExit(f"VALIDATE: Duplicate running conversation ID {v!r} detected between {id_to_session[v]!r} and {s_name!r}") id_to_session[v] = s_name _validate(d) d_state = {k: v for k, v in d.items() if k != 'herdr_sessions'} conn.execute('REPLACE INTO state (id, data) VALUES (1, ?)', (json.dumps(d_state),)) current_names = [] for s in d.get('herdr_sessions', []): name = s.get('name') status = s.get('status') pane_cwd = (s.get('pane') or {}).get('cwd', '') conn.execute('REPLACE INTO sessions (name, status, pane_cwd, data) VALUES (?, ?, ?, ?)', (name, status, pane_cwd, json.dumps(s))) current_names.append(name) if current_names: placeholders = ','.join('?' for _ in current_names) conn.execute(f'DELETE FROM sessions WHERE name NOT IN ({placeholders})', current_names) else: conn.execute('DELETE FROM sessions') new_sessions = get_all_sessions_status(d) conn.commit() if new_sessions != old_sessions: if os.path.exists(yaml_path): try: shutil.copy2(yaml_path, yaml_path + '.bak') except Exception: pass dir_ = os.path.dirname(yaml_path) or '.' fd, tmp = tempfile.mkstemp(dir=dir_, prefix='.agent-sessions.', suffix='.tmp') try: with os.fdopen(fd, 'w') as f: yaml.safe_dump(d, f, default_flow_style=False, sort_keys=False, allow_unicode=True, width=4096) os.replace(tmp, yaml_path) except Exception: if os.path.exists(tmp): os.remove(tmp) raise try: conn.execute('PRAGMA wal_checkpoint(TRUNCATE)') except Exception: pass except Exception: conn.rollback() raise finally: conn.close() try: os.chmod(db_path, 0o600) wal = db_path + '-wal' if os.path.exists(wal): os.chmod(wal, 0o600) shm = db_path + '-shm' if os.path.exists(shm): os.chmod(shm, 0o600) except Exception: pass try: fcntl.flock(_flock_f, fcntl.LOCK_UN) _flock_f.close() except Exception: pass if __name__ == '__main__': atomic_dump_yaml_main()