Codex

#!/usr/bin/env python3
"""Run Codex normally, or connect its IDE stdio to one durable local app server."""
import fcntl
import hashlib
import json
import os
from pathlib import Path
import re
import socket
import subprocess
import sys
import threading
import time
import signal
import tempfile

sys.path.insert(0, str(Path(__file__).resolve().parent / 'vendor'))

from websockets.sync.client import unix_connect
from websockets.exceptions import ConnectionClosed


def native_binary():
    explicit = os.environ.get('ADOM_CODEX_NATIVE')
    if explicit:
        path = Path(explicit).resolve()
        if not path.is_file() or path == Path(__file__).resolve():
            raise RuntimeError('ADOM_CODEX_NATIVE must identify the native Codex executable')
        return path
    # The extension appends its own native bin directory to this child's PATH.
    # Prefer that exact version, even if a newer VSIX is already on disk.
    for entry in reversed(os.environ.get('PATH', '').split(os.pathsep)):
        directory = Path(entry)
        if directory.name not in ('linux-x86_64', 'linux-aarch64'):
            continue
        package = directory.parent.parent
        if package.name.startswith('openai.chatgpt-') and (directory / 'codex').is_file():
            return directory / 'codex'
    extensions = Path.home() / '.local/share/code-server/extensions'
    candidates = []
    for package in extensions.glob('openai.chatgpt-*-linux-*'):
        version = re.search(r'openai\.chatgpt-([\d.]+)', package.name)
        if not version:
            continue
        arch = 'linux-aarch64' if os.uname().machine == 'aarch64' else 'linux-x86_64'
        binary = package / 'bin' / arch / 'codex'
        if binary.is_file():
            candidates.append((tuple(map(int, version[1].split('.'))), binary))
    if not candidates:
        raise RuntimeError('No installed Linux Codex extension binary was found')
    if len(candidates) != 1:
        raise RuntimeError('Multiple Codex versions installed; launch through the extension '
                           'so its exact native binary can be selected')
    return candidates[0][1]


def server_args(argv):
    """Only replace the plain IDE server invocation; leave other CLI verbs alone."""
    index = 0
    while index < len(argv):
        value = argv[index]
        if value in ('-c', '--config', '--enable', '--disable', '--code-mode-host'):
            index += 2
        elif value.startswith(('--config=', '--enable=', '--disable=')):
            index += 1
        elif value == 'app-server':
            tail = argv[index + 1:]
            # Explicit transports, help, schema generation and daemon management
            # are native CLI operations, not IDE session launches.
            if any(x in tail for x in ('--listen', '--stdio', '--help', '-h',
                                       'daemon', 'proxy', 'generate-ts', 'generate-json-schema')):
                return None
            if any(x.startswith('--listen=') for x in tail):
                return None
            return list(argv)
        else:
            return None
    return None


def state_dir():
    home = Path(os.environ.get('CODEX_HOME', str(Path.home() / '.codex'))).resolve()
    identity = hashlib.sha256(os.fsencode(home)).hexdigest()[:16]
    base = Path(os.environ.get('ADOM_CODEX_RUNTIME_STATE',
                str(Path.home() / '.local/state/adom/codex-runtime')))
    directory = base / identity
    directory.mkdir(parents=True, exist_ok=True, mode=0o700)
    directory.chmod(0o700)
    if len(os.fsencode(directory / 'server.sock')) >= 104:
        raise RuntimeError('Runtime state path is too long for a Unix domain socket')
    return directory


def process_matches(metadata):
    try:
        pid = int(metadata['pid'])
        stat = Path(f'/proc/{pid}/stat').read_text().rsplit(')', 1)[1].split()
        return stat[19] == metadata['start_ticks'] and stat[0] != 'Z'
    except (OSError, KeyError, ValueError):
        return False


def socket_alive(path):
    try:
        with socket.socket(socket.AF_UNIX) as connection:
            connection.settimeout(0.2)
            connection.connect(str(path))
        return True
    except OSError:
        return False


def identity(binary):
    stat = binary.stat()
    return [str(binary.resolve()), stat.st_size, stat.st_mtime_ns]


def extension_version(binary):
    match = re.search(r'openai\.chatgpt-([\d.]+)-', str(binary))
    return tuple(map(int, match[1].split('.'))) if match else None


def protocol_fingerprint(binary, directory):
    """Cache the complete native schema by executable identity, not VSIX version."""
    before = identity(binary)
    key = hashlib.sha256(json.dumps(before).encode()).hexdigest()
    path = directory / 'protocols.json'
    try:
        cache = json.loads(path.read_text())
        if not isinstance(cache, dict): cache = {}
    except (OSError, ValueError):
        cache = {}
    if key in cache: return cache[key]
    with tempfile.TemporaryDirectory(prefix='codex-protocol-') as temporary:
        subprocess.run([str(binary), 'app-server', 'generate-json-schema',
                        '--experimental', '--out', temporary], check=True,
                       stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
                       stderr=subprocess.PIPE, timeout=15)
        files = sorted(Path(temporary).rglob('*.json'))
        if not files: raise RuntimeError('Native app-server schema generation returned no schemas')
        digest = hashlib.sha256()
        for file in files:
            schema = json.loads(file.read_text())
            digest.update(json.dumps([str(file.relative_to(temporary)), schema],
                                     sort_keys=True, separators=(',', ':')).encode())
        if identity(binary) != before:
            raise RuntimeError('Codex binary changed while checking its protocol')
    cache[key] = digest.hexdigest()
    cache = dict(list(cache.items())[-8:])
    temporary = path.with_suffix('.tmp')
    temporary.write_text(json.dumps(cache) + '\n')
    temporary.replace(path)
    return cache[key]


def compatible_backend(metadata, binary, directory):
    running = Path(metadata.get('binary', ''))
    if not extension_version(binary) or not extension_version(running): return False
    try:
        running_protocol = metadata.get('protocolFingerprint')
        if not running_protocol:
            # Legacy servers have no saved attestation. A replaced/deleted path
            # cannot prove which protocol their running image implements.
            if metadata.get('identity') != identity(running): return False
            running_protocol = protocol_fingerprint(running, directory)
        return running_protocol == protocol_fingerprint(binary, directory)
    except (OSError, ValueError, RuntimeError, subprocess.SubprocessError):
        return False


def defer_upgrade(directory, metadata, binary, reason):
    record = {'state': 'deferred', 'reason': reason, 'protocolCompatible': True,
              'runningBinary': metadata['binary'], 'requestedBinary': str(binary),
              'pid': metadata['pid'], 'checkedAt': time.time(),
              'message': 'Connected to the existing compatible backend. Active work is preserved; '
                         'the next connection can upgrade after all views close and work is idle.'}
    path = directory / 'upgrade.json'
    temporary = path.with_suffix('.tmp')
    temporary.write_text(json.dumps(record, indent=2) + '\n')
    temporary.replace(path)


def backend_idle(sock):
    """Ask the backend; disk history can lag an active turn."""
    with unix_connect(str(sock), compression=None, max_size=None, open_timeout=5,
                      close_timeout=1) as connection:
        request_id = 0
        def rpc(method, params):
            nonlocal request_id
            request_id += 1
            connection.send(json.dumps({'id': request_id, 'method': method, 'params': params}))
            deadline = time.monotonic() + 5
            while time.monotonic() < deadline:
                reply = json.loads(connection.recv(timeout=max(.01, deadline-time.monotonic())))
                if reply.get('id') == request_id:
                    if 'error' in reply: raise RuntimeError(str(reply['error']))
                    return reply['result']
            raise RuntimeError('Could not verify backend idle state')
        rpc('initialize', {'clientInfo': {'name': 'adom-codex-updater', 'version': '1'},
                           'capabilities': {'experimentalApi': True}})
        connection.send(json.dumps({'method': 'initialized'}))
        cursor = None
        while True:
            result = rpc('thread/loaded/list', {'cursor': cursor, 'limit': 100})
            for thread_id in result['data']:
                thread = rpc('thread/read', {'threadId': thread_id, 'includeTurns': False})['thread']
                if thread['status']['type'] != 'idle': return False
            cursor = result.get('nextCursor')
            if not cursor: return True


def ensure_server(binary, args, client_guard=None):
    directory = state_dir()
    sock = directory / 'server.sock'
    metadata_path = directory / 'server.json'
    # One owner per CODEX_HOME, even during overlapping iframe reloads.
    with (directory / 'startup.lock').open('a') as guard:
        fcntl.flock(guard, fcntl.LOCK_EX)
        try:
            metadata = json.loads(metadata_path.read_text())
        except (OSError, ValueError):
            metadata = {}
        if process_matches(metadata):
            if metadata.get('args') != args:
                raise RuntimeError('The running Codex backend has different startup options; '
                                   'finish its work before changing runtime options')
            same_binary = (metadata.get('identity') == identity(binary) if 'identity' in metadata
                           else metadata.get('binary') == str(binary))
            if not socket_alive(sock):
                raise RuntimeError('The existing Codex backend is alive but unavailable; '
                                   'refusing to start a competing thread writer')
            if same_binary: return sock
            requested_version = extension_version(binary)
            running_version = extension_version(metadata.get('binary', ''))
            if requested_version and running_version and requested_version < running_version:
                raise RuntimeError('This editor uses an older Codex extension than the running backend. '
                                   'Reload the editor to use the updated extension; backend downgrade refused.')
            # Matching generated schemas permit reconnecting without replacing
            # the writer. Different/unknown protocols retain the conservative gate.
            compatible = compatible_backend(metadata, binary, directory)
            if client_guard is None:
                raise RuntimeError('Backend version differs; upgrade requires exclusive client ownership')
            try: fcntl.flock(client_guard, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except BlockingIOError:
                if compatible:
                    # Failed lock conversion may drop our shared lock on Linux.
                    fcntl.flock(client_guard, fcntl.LOCK_SH)
                    defer_upgrade(directory, metadata, binary, 'connected_views')
                    return sock
                raise RuntimeError('Codex updated; the previous version still has connected views. '
                                   'Its work is preserved. Close those views after work finishes, then retry.')
            try:
                idle = backend_idle(sock)
            except Exception:
                if not compatible: raise
                idle = False
            if not idle:
                if compatible:
                    fcntl.flock(client_guard, fcntl.LOCK_SH)
                    defer_upgrade(directory, metadata, binary, 'backend_busy_or_unverified')
                    return sock
                raise RuntimeError('Codex updated while a turn is running. Its work is preserved; '
                                   'retry when it finishes to activate the new version.')
            if process_matches(metadata): os.kill(metadata['pid'], signal.SIGTERM)
            deadline = time.monotonic() + 10
            while process_matches(metadata) and time.monotonic() < deadline: time.sleep(.05)
            if process_matches(metadata):
                raise RuntimeError('Old backend has not exited; no forced termination performed')
        if socket_alive(sock):
            raise RuntimeError('An unrecognized process owns the Codex runtime socket')
        sock.unlink(missing_ok=True)
        protocol = None
        if extension_version(binary):
            try: protocol = protocol_fingerprint(binary, directory)
            except (OSError, ValueError, RuntimeError, subprocess.SubprocessError): pass
        log = directory / 'server.log'
        with log.open('ab') as output:
            process = subprocess.Popen(
                [str(binary), *args, '--listen', 'unix://' + str(sock)],
                stdin=subprocess.DEVNULL, stdout=output, stderr=output,
                start_new_session=True, close_fds=True,
            )
        metadata = {'pid': process.pid, 'binary': str(binary), 'identity': identity(binary), 'args': args,
                    'protocolFingerprint': protocol,
                    'socket': str(sock), 'started_at': time.time(),
                    'start_ticks': Path(f'/proc/{process.pid}/stat').read_text()
                        .rsplit(')', 1)[1].split()[19]}
        temporary = metadata_path.with_suffix('.tmp')
        temporary.write_text(json.dumps(metadata, indent=2) + '\n')
        temporary.replace(metadata_path)
        deadline = time.monotonic() + 20
        while time.monotonic() < deadline:
            if process.poll() is not None:
                raise RuntimeError(f'Codex backend exited ({process.returncode}); inspect {log}')
            if socket_alive(sock):
                (directory / 'upgrade.json').unlink(missing_ok=True)
                return sock
            time.sleep(0.05)
        raise RuntimeError(f'Codex backend is still starting; inspect {log}')


def proxy(sock):
    # The native Unix listener speaks WebSocket; the extension speaks JSONL.
    # No message rewriting: IDs, approvals, notifications and usage pass through.
    with unix_connect(str(sock), compression=None, max_size=None,
                      open_timeout=10, close_timeout=1) as connection:
        def forward_input():
            try:
                for line in sys.stdin:
                    if line.strip():
                        connection.send(line.rstrip('\r\n'))
            except (ConnectionClosed, OSError):
                pass
            finally:
                connection.close()

        reader = threading.Thread(target=forward_input, daemon=True)
        reader.start()
        try:
            for message in connection:
                if isinstance(message, bytes):
                    message = message.decode('utf-8')
                sys.stdout.write(message + '\n')
                sys.stdout.flush()
        except BrokenPipeError:
            return


def main():
    os.umask(0o077)
    argv = sys.argv[1:]
    if argv == ['--runtime-status']:
        path = state_dir() / 'server.json'
        metadata = json.loads(path.read_text()) if path.exists() else {}
        pending = path.with_name('upgrade.json')
        upgrade = json.loads(pending.read_text()) if pending.exists() else None
        if upgrade and upgrade.get('pid') != metadata.get('pid'): upgrade = None
        print(json.dumps({**metadata, 'alive': process_matches(metadata), 'upgrade': upgrade}, indent=2))
        return
    binary = native_binary()
    args = server_args(argv)
    if args is None:
        os.execv(str(binary), [str(binary), *argv])
    else:
        with (state_dir() / 'clients.lock').open('a') as client_guard:
            fcntl.flock(client_guard, fcntl.LOCK_SH)
            sock = ensure_server(binary, args, client_guard)
            fcntl.flock(client_guard, fcntl.LOCK_SH)
            proxy(sock)


if __name__ == '__main__':
    try:
        main()
    except Exception as error:
        print(f'Adom Codex runtime: {error}', file=sys.stderr)
        sys.exit(1)