multi-agent-mux/.agents/skills/lib_py/atomic_yaml.py

# 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
from lib_py import registry_format

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")
        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'] = []

        # Issue #71: an older registry is migrated here, under the lock and
        # before the caller's mutation. It changes row fields only, so the
        # YAML is rewritten whether or not the mutation changes anything.
        from_version = registry_format.stored_version(d)
        migrated_rows = 0
        force_yaml = False
        if from_version < registry_format.CURRENT:
            d, migrated_rows = registry_format.migrate(d)
            force_yaml = True

        old_sessions = get_all_sessions_status(d)
        # a mutation that replaces a row the monitor registered for a pane it
        # created itself (create_session.sh) drops that name from old_roles
        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'], '<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', 'grok_session_id_own', 'opencode_session_id_own', 'codex_session_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 force_yaml or 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
    return from_version, registry_format.stored_version(d), migrated_rows

if __name__ == '__main__':
    atomic_dump_yaml_main()