multi-agent-mux/.agents/skills/lib_py/agents/adapters/codex.py

import os, json, glob, sqlite3, subprocess
from typing import Optional, Any, List
from lib_py.agents.base import BaseAgentAdapter, DiscoveryContext
from lib_py.verify_session import workspace_key

class CodexAgentAdapter(BaseAgentAdapter):
    @property
    def name(self) -> str:
        return 'codex'

    @property
    def own_key(self) -> str:
        return 'codex_session_id_own'

    @property
    def ready_tokens(self) -> str:
        return 'Codex|Ask Codex to do anything'

    @property
    def exit_key(self) -> str:
        return 'Ctrl+C'

    @property
    def delegate_agent_key(self) -> str:
        return 'codex-cli'

    @property
    def identity_cache_fields(self) -> tuple:
        return ('session_id',)

    @property
    def input_placeholder(self) -> Optional[str]:
        return 'Ask Codex to do anything'

    def _codex_home(self, ctx: DiscoveryContext) -> str:
        return os.environ.get('CODEX_HOME') or f"{ctx.home_dir}/.codex"

    def _state_db(self, ctx: DiscoveryContext) -> Optional[str]:
        matches = glob.glob(f"{self._codex_home(ctx)}/state_*.sqlite")
        if not matches:
            return None
        return max(matches, key=os.path.getmtime)

    def _find_rollout_via_db(self, uuid: str, ctx: DiscoveryContext):
        """Returns (rollout_path, cwd) from the state DB, or (None, None)."""
        db = self._state_db(ctx)
        if not db or not os.path.exists(db):
            return None, None
        try:
            conn = sqlite3.connect(db)
            row = conn.execute("SELECT rollout_path, cwd FROM threads WHERE id=?", (uuid,)).fetchone()
            conn.close()
            if row and row[0] and os.path.exists(row[0]):
                return row[0], row[1]
        except Exception:
            pass
        return None, None

    def _find_rollout(self, uuid: str, ctx: DiscoveryContext) -> Optional[str]:
        path, _ = self._find_rollout_via_db(uuid, ctx)
        if path:
            return path
        matches = glob.glob(f"{self._codex_home(ctx)}/sessions/*/*/*/rollout-*-{uuid}.jsonl")
        return matches[0] if matches else None

    def artifact_path(self, uuid: str, ctx: DiscoveryContext) -> str:
        return self._find_rollout(uuid, ctx) or f"{self._codex_home(ctx)}/sessions/rollout-{uuid}.jsonl"

    def verify_artifact(self, uuid: str, ctx: DiscoveryContext) -> bool:
        db_path, db_cwd = self._find_rollout_via_db(uuid, ctx)
        path = db_path or self._find_rollout(uuid, ctx)
        if not path or not os.path.exists(path):
            return False
        if ctx.epoch and os.path.getmtime(path) < ctx.epoch:
            return False
        found_cwd = db_cwd
        if not found_cwd:
            try:
                with open(path) as f:
                    first = f.readline()
                found_cwd = json.loads(first).get('payload', {}).get('cwd')
            except Exception:
                return False
        if found_cwd and workspace_key(found_cwd) != workspace_key(ctx.cwd):
            return False
        return True

    def purge_artifacts(self, uuid: str, ctx: DiscoveryContext) -> List[str]:
        purged = []
        path = self._find_rollout(uuid, ctx)
        try:
            res = subprocess.run(['codex', 'delete', uuid, '--force'],
                                  capture_output=True, text=True, timeout=5)
            if res.returncode == 0:
                if path:
                    purged.append(path)
                return purged
        except Exception:
            pass
        if path and os.path.exists(path):
            os.remove(path)
            purged.append(path)
        db = self._state_db(ctx)
        if db and os.path.exists(db):
            try:
                conn = sqlite3.connect(db)
                conn.execute("DELETE FROM threads WHERE id=?", (uuid,))
                conn.commit()
                conn.close()
                purged.append(f"state db row for thread: {uuid}")
            except Exception:
                pass
        return purged

    def spawn_spec(self, binary: str, session_uuid: str = "", use_wrapper: bool = False) -> str:
        return f"{binary} --dangerously-bypass-approvals-and-sandbox"

    def resume_spec(self, binary: str, session_uuid: str, materialized: bool = False) -> str:
        if materialized and session_uuid:
            return f"{binary} resume --dangerously-bypass-approvals-and-sandbox {session_uuid}"
        return f"{binary} --dangerously-bypass-approvals-and-sandbox"

    def auth_ok(self, run_cmd: Optional[Any] = None) -> bool:
        home = os.environ.get('HOME_DIR') or os.environ.get('HOME') or os.path.expanduser('~')
        codex_home = os.environ.get('CODEX_HOME') or f"{home}/.codex"
        return os.path.exists(f"{codex_home}/auth.json") or bool(os.environ.get('OPENAI_API_KEY'))

    def discover(self, ctx: DiscoveryContext) -> List[str]:
        db = self._state_db(ctx)
        if db and os.path.exists(db):
            try:
                conn = sqlite3.connect(db)
                cur = conn.cursor()
                cur.execute("PRAGMA table_info(threads)")
                cols = {r[1] for r in cur.fetchall()}
                order_col = "updated_at" if "updated_at" in cols else ("created_at" if "created_at" in cols else None)
                order_clause = f"ORDER BY {order_col} DESC" if order_col else ""
                rows = conn.execute(
                    f"SELECT id, cwd FROM threads {order_clause} LIMIT 50"
                ).fetchall()
                conn.close()
                candidates = []
                for uuid, cwd in rows:
                    if ctx.workspace and cwd and workspace_key(cwd) != workspace_key(ctx.workspace):
                        continue
                    if uuid and self.verify_artifact(uuid, ctx):
                        candidates.append(uuid)
                return candidates
            except Exception:
                pass
        root = self._codex_home(ctx)
        candidates = []
        paths = sorted(glob.glob(f"{root}/sessions/*/*/*/rollout-*.jsonl"), key=os.path.getmtime, reverse=True)[:50]
        for path in paths:
            uuid = os.path.basename(path)[:-6][-36:]
            if uuid and self.verify_artifact(uuid, ctx):
                candidates.append(uuid)
        return candidates