Codex

import importlib.util
import json
import os
from pathlib import Path
import queue
import signal
import subprocess
import sys
import tempfile
import threading
import time
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from unittest.mock import patch

ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location('runtime', ROOT / 'runtime/runtime.py')
runtime = importlib.util.module_from_spec(spec)
spec.loader.exec_module(runtime)
sys.path.insert(0, str(ROOT / 'runtime'))


class RoutingTests(unittest.TestCase):
    def test_uses_launching_extension_even_when_newer_vsix_is_downloaded(self):
        with tempfile.TemporaryDirectory() as tmp:
            home = Path(tmp)
            paths = []
            for version in ['1.0.0', '99.0.0']:
                directory = home / f'.local/share/code-server/extensions/openai.chatgpt-{version}-linux-x64/bin/linux-x86_64'
                directory.mkdir(parents=True)
                (directory / 'codex').write_text('fixture')
                paths.append(directory)
            with patch.dict(os.environ, {'PATH': str(paths[0])}, clear=True), patch.object(Path, 'home', return_value=home):
                self.assertEqual(runtime.native_binary(), paths[0] / 'codex')
            with patch.dict(os.environ, {'PATH': '/usr/bin'}, clear=True), patch.object(Path, 'home', return_value=home):
                with self.assertRaisesRegex(RuntimeError, 'Multiple Codex versions'):
                    runtime.native_binary()

    def test_extension_launch_uses_durable_server(self):
        args = ['-c', 'features.code_mode_host=true', 'app-server', '--analytics-default-enabled']
        self.assertEqual(runtime.server_args(args), args)

    def test_other_cli_commands_keep_native_behavior(self):
        for args in [[], ['--version'], ['resume'], ['stdio-to-uds', '/tmp/lsp.sock'],
                     ['app-server', 'proxy'], ['app-server', 'daemon', 'start'],
                     ['app-server', '--help'], ['app-server', '--listen=unix:///tmp/server']]:
            with self.subTest(args=args):
                self.assertIsNone(runtime.server_args(args))

    def test_pid_reuse_is_not_mistaken_for_server(self):
        self.assertFalse(runtime.process_matches({'pid': os.getpid(), 'start_ticks': '-1'}))

    def test_stale_extension_cannot_downgrade_a_new_backend(self):
        with tempfile.TemporaryDirectory() as tmp:
            directory = Path(tmp)
            binary = directory/'openai.chatgpt-1.0.0-linux-x64/bin/codex'
            binary.parent.mkdir(parents=True); binary.write_text('fixture')
            (directory/'server.json').write_text(json.dumps({'binary': '/extensions/openai.chatgpt-2.0.0-linux-x64/bin/codex',
                'args': ['app-server'], 'identity': ['newer']}))
            with patch.object(runtime, 'state_dir', return_value=directory), \
                 patch.object(runtime, 'process_matches', return_value=True), \
                 patch.object(runtime, 'socket_alive', return_value=True), \
                 patch.object(runtime.os, 'kill') as kill:
                with self.assertRaisesRegex(RuntimeError, 'downgrade refused'):
                    runtime.ensure_server(binary, ['app-server'])
                kill.assert_not_called()

    def test_distinct_codex_homes_have_distinct_private_state(self):
        with tempfile.TemporaryDirectory(prefix='cxd-') as directory:
            with patch.dict(os.environ, {'ADOM_CODEX_RUNTIME_STATE': directory, 'CODEX_HOME': '/tmp/a'}):
                first = runtime.state_dir()
                self.assertEqual(first.stat().st_mode & 0o777, 0o700)
            with patch.dict(os.environ, {'ADOM_CODEX_RUNTIME_STATE': directory, 'CODEX_HOME': '/tmp/b'}):
                self.assertNotEqual(runtime.state_dir(), first)

    def test_compatible_upgrade_preserves_busy_backend_and_connected_views(self):
        for connected, idle in [(False, False), (True, True), (False, RuntimeError('probe failed'))]:
            with self.subTest(connected=connected, idle=idle), tempfile.TemporaryDirectory() as tmp:
                directory = Path(tmp)
                old = directory/'openai.chatgpt-1.0.0-linux-x64/codex'
                new = directory/'openai.chatgpt-1.0.1-linux-x64/codex'
                for binary in [old, new]:
                    binary.parent.mkdir(); binary.write_text('fixture')
                metadata = {'pid': 123, 'binary': str(old), 'identity': runtime.identity(old), 'args': ['app-server']}
                (directory/'server.json').write_text(json.dumps(metadata))
                with (directory/'clients.lock').open('a') as guard, (directory/'clients.lock').open('a') as other:
                    runtime.fcntl.flock(guard, runtime.fcntl.LOCK_SH)
                    if connected: runtime.fcntl.flock(other, runtime.fcntl.LOCK_SH)
                    with patch.object(runtime, 'state_dir', return_value=directory), \
                         patch.object(runtime, 'process_matches', return_value=True), \
                         patch.object(runtime, 'socket_alive', return_value=True), \
                         patch.object(runtime, 'protocol_fingerprint', return_value='same-schema'), \
                         patch.object(runtime, 'backend_idle', **({'side_effect': idle} if isinstance(idle, Exception) else {'return_value': idle})), \
                         patch.object(runtime.os, 'kill') as kill, \
                         patch.object(runtime.subprocess, 'Popen') as spawn:
                        self.assertEqual(runtime.ensure_server(new, ['app-server'], guard), directory/'server.sock')
                        kill.assert_not_called(); spawn.assert_not_called()
                    record = json.loads((directory/'upgrade.json').read_text())
                    self.assertEqual(record['pid'], 123)
                    self.assertTrue(record['protocolCompatible'])
                    runtime.fcntl.flock(other, runtime.fcntl.LOCK_UN)
                    with self.assertRaises(BlockingIOError):
                        runtime.fcntl.flock(other, runtime.fcntl.LOCK_EX | runtime.fcntl.LOCK_NB)

    def test_unknown_changed_or_different_protocol_does_not_attest_compatibility(self):
        with tempfile.TemporaryDirectory() as tmp:
            directory = Path(tmp)
            old = directory/'openai.chatgpt-1.0.0-linux-x64/codex'
            new = directory/'openai.chatgpt-1.0.1-linux-x64/codex'
            for binary in [old, new]:
                binary.parent.mkdir(); binary.write_text('fixture')
            metadata = {'binary': str(old), 'identity': runtime.identity(old)}
            with patch.object(runtime, 'protocol_fingerprint', side_effect=['old-schema', 'new-schema']):
                self.assertFalse(runtime.compatible_backend(metadata, new, directory))
            with patch.object(runtime, 'protocol_fingerprint', side_effect=RuntimeError('unsupported')):
                self.assertFalse(runtime.compatible_backend(metadata, new, directory))
            old.write_text('replacement')
            with patch.object(runtime, 'protocol_fingerprint') as fingerprint:
                self.assertFalse(runtime.compatible_backend(metadata, new, directory))
                fingerprint.assert_not_called()
            metadata['protocolFingerprint'] = 'attested-at-start'
            old.unlink()
            with patch.object(runtime, 'protocol_fingerprint', return_value='attested-at-start'):
                self.assertTrue(runtime.compatible_backend(metadata, new, directory))

    def test_schema_fingerprint_is_cached_by_binary_identity(self):
        with tempfile.TemporaryDirectory() as tmp:
            directory = Path(tmp); binary = directory/'codex'; binary.write_text('fixture')
            def generate(args, **kwargs):
                (Path(args[-1])/'schema.json').write_text('{"type":"object","properties":{}}')
            with patch.object(runtime.subprocess, 'run', side_effect=generate) as run:
                first = runtime.protocol_fingerprint(binary, directory)
                self.assertEqual(runtime.protocol_fingerprint(binary, directory), first)
                self.assertEqual(run.call_count, 1)
                binary.write_text('changed fixture')
                self.assertEqual(runtime.protocol_fingerprint(binary, directory), first)
                self.assertEqual(run.call_count, 2)


class Client:
    def __init__(self, env):
        self.process = subprocess.Popen(
            [sys.executable, str(ROOT / 'runtime/runtime.py'), '-c', 'features.code_mode_host=true',
             'app-server', '--analytics-default-enabled'], env=env,
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            text=True, bufsize=1)
        self.messages = queue.Queue()
        def read():
            for line in self.process.stdout:
                self.messages.put(json.loads(line))
        threading.Thread(target=read, daemon=True).start()
        self.rpc(1, 'initialize', {'clientInfo': {'name': 'reload-test', 'version': '1'},
                                 'capabilities': {'experimentalApi': True}})
        self.send({'method': 'initialized'})

    def send(self, message):
        self.process.stdin.write(json.dumps(message) + '\n')
        self.process.stdin.flush()

    def rpc(self, request_id, method, params):
        self.send({'id': request_id, 'method': method, 'params': params})
        deadline = time.monotonic() + 25
        while time.monotonic() < deadline:
            result = self.messages.get(timeout=25)
            if result.get('id') == request_id:
                if 'error' in result:
                    raise AssertionError(result['error'])
                return result['result']
        raise TimeoutError(method)

    def close(self, abrupt=False):
        if self.process.poll() is None:
            if abrupt:
                self.process.kill()
            else:
                self.process.stdin.close()
            self.process.wait(timeout=10)
        for stream in (self.process.stdin, self.process.stdout, self.process.stderr):
            if not stream.closed:
                stream.close()


class ReconnectTest(unittest.TestCase):
    def test_same_active_turn_survives_three_proxy_replacements(self):
        started = threading.Event()
        finish = threading.Event()
        requests = []

        class LocalFixture(BaseHTTPRequestHandler):
            def log_message(self, *args):
                pass

            def do_POST(self):
                self.rfile.read(int(self.headers.get('Content-Length', '0')))
                requests.append(self.path)
                self.send_response(200)
                self.send_header('Content-Type', 'text/event-stream')
                self.end_headers()
                self.event({'type': 'response.created', 'response': {'id': 'reload_test'}})
                started.set()
                finish.wait(30)
                message = {'type': 'message', 'id': 'msg_reload', 'role': 'assistant',
                           'content': [{'type': 'output_text', 'text': 'Reload survived.'}]}
                self.event({'type': 'response.output_item.done', 'output_index': 0, 'item': message})
                self.event({'type': 'response.completed', 'response': {
                    'id': 'reload_test', 'status': 'completed', 'output': [message],
                    'usage': {'input_tokens': 1, 'output_tokens': 1, 'total_tokens': 2}}})

            def event(self, data):
                self.wfile.write(('data: ' + json.dumps(data) + '\n\n').encode())
                self.wfile.flush()

        server = ThreadingHTTPServer(('127.0.0.1', 0), LocalFixture)
        threading.Thread(target=server.serve_forever, daemon=True).start()
        clients = []
        with tempfile.TemporaryDirectory(prefix='cxd-') as directory:
            test_root = Path(directory)
            home = test_root / 'home'
            home.mkdir()
            (home / 'config.toml').write_text(
                '[model_providers.reload_mock]\nname = "Local fixture"\n'
                f'base_url = "http://127.0.0.1:{server.server_port}"\n'
                'wire_api = "responses"\nrequires_openai_auth = false\n')
            env = dict(os.environ, CODEX_HOME=str(home), ADOM_CODEX_RUNTIME_STATE=str(test_root / 'state'))
            env.pop('OPENAI_API_KEY', None)
            if os.environ.get('ADOM_TEST_CODEX_OLD'):
                env['ADOM_CODEX_NATIVE'] = os.environ['ADOM_TEST_CODEX_OLD']
            try:
                # Simultaneous clients must share one server, even before any
                # conversation exists. IDs are scoped to each connection.
                first = Client(env)
                clients.append(first)
                second = Client(env)
                clients.append(second)
                metadata_path = next((test_root / 'state').glob('*/server.json'))
                original_pid = json.loads(metadata_path.read_text())['pid']
                created = first.rpc(2, 'thread/start', {
                    'cwd': directory, 'approvalPolicy': 'never', 'sandbox': 'danger-full-access',
                    'model': 'reload-test', 'modelProvider': 'reload_mock',
                    'baseInstructions': 'Local transport fixture; no real model.',
                    'config': {'model_providers.reload_mock': {'name': 'Local fixture',
                        'base_url': f'http://127.0.0.1:{server.server_port}',
                        'wire_api': 'responses', 'requires_openai_auth': False}},
                })
                thread_id = created['thread']['id']
                turn = first.rpc(3, 'turn/start', {'threadId': thread_id,
                    'input': [{'type': 'text', 'text': 'Local fixture', 'text_elements': []}]})['turn']
                self.assertTrue(started.wait(10), 'Fixture did not receive the model request')
                self.assertFalse(runtime.backend_idle(Path(json.loads(metadata_path.read_text())['socket'])))
                if os.environ.get('ADOM_TEST_CODEX_NEW'):
                    upgrade_env = dict(env, ADOM_CODEX_NATIVE=os.environ['ADOM_TEST_CODEX_NEW'])
                    newer = Client(upgrade_env)
                    clients.append(newer)
                    observed = newer.rpc(2, 'thread/read', {'threadId': thread_id, 'includeTurns': True})['thread']
                    self.assertEqual(observed['turns'][-1]['id'], turn['id'])
                    self.assertEqual(json.loads(metadata_path.read_text())['pid'], original_pid)
                    self.assertEqual(json.loads(metadata_path.with_name('upgrade.json').read_text())['reason'], 'connected_views')
                    newer.close(); first.close(); second.close()
                    env = upgrade_env
                    first = Client(env)
                    clients.append(first)
                    self.assertEqual(json.loads(metadata_path.read_text())['pid'], original_pid)
                    self.assertEqual(json.loads(metadata_path.with_name('upgrade.json').read_text())['reason'], 'backend_busy_or_unverified')
                second.close()
                for index in range(3):
                    first.close(abrupt=index % 2 == 0)
                    first = Client(env)
                    clients.append(first)
                    resumed = first.rpc(2, 'thread/resume', {'threadId': thread_id})['thread']
                    self.assertEqual(resumed['id'], thread_id)
                    self.assertEqual(resumed['status']['type'], 'active')
                    self.assertEqual(resumed['turns'][-1]['id'], turn['id'])
                    self.assertEqual(resumed['turns'][-1]['status'], 'inProgress')
                    self.assertEqual(json.loads(metadata_path.read_text())['pid'], original_pid)
                # An unattested binary change must still fail closed, preserving
                # the active turn even when there are no proxy clients.
                first.close()
                original_metadata = json.loads(metadata_path.read_text())
                changed_metadata = {**original_metadata, 'identity': ['older-extension'], 'protocolFingerprint': None}
                metadata_path.write_text(json.dumps(changed_metadata))
                rejected = subprocess.run([sys.executable, str(ROOT/'runtime/runtime.py'),
                    '-c', 'features.code_mode_host=true', 'app-server', '--analytics-default-enabled'],
                    env=env, input='', capture_output=True, text=True, timeout=10)
                self.assertNotEqual(rejected.returncode, 0)
                self.assertIn('turn is running', rejected.stderr)
                self.assertTrue(runtime.process_matches(original_metadata))
                metadata_path.write_text(json.dumps(original_metadata))
                first = Client(env)
                clients.append(first)
                first.rpc(2, 'thread/resume', {'threadId': thread_id})
                finish.set()
                deadline = time.monotonic() + 10
                while time.monotonic() < deadline:
                    message = first.messages.get(timeout=10)
                    if message.get('method') == 'turn/completed':
                        completed = message['params']['turn']
                        self.assertEqual(completed['id'], turn['id'])
                        self.assertEqual(completed['status'], 'completed')
                        self.assertEqual(completed['items'][0]['text'], 'Reload survived.')
                        break
                else:
                    self.fail('Reconnected client did not receive completion')
                self.assertEqual(len(requests), 1, 'Reconnect restarted model work')
                self.assertTrue(runtime.backend_idle(Path(json.loads(metadata_path.read_text())['socket'])))
                # Completed history must still reopen after the last client leaves.
                first.close()
                # With no clients and completed work, the next connection can
                # retire an older version and reopen its persisted conversation.
                metadata_path.write_text(json.dumps(changed_metadata))
                first = Client(env)
                clients.append(first)
                self.assertNotEqual(json.loads(metadata_path.read_text())['pid'], original_pid)
                resumed = first.rpc(2, 'thread/resume', {'threadId': thread_id})['thread']
                self.assertEqual(resumed['turns'][-1]['status'], 'completed')
                print('Verified: one backend PID, one model fixture request, three reloads, '
                      'same turn completed, saved history reopened.')
            finally:
                finish.set()
                for client in clients:
                    client.close(abrupt=True)
                for path in (test_root / 'state').glob('*/server.json'):
                    metadata = json.loads(path.read_text())
                    if runtime.process_matches(metadata):
                        os.kill(metadata['pid'], signal.SIGTERM)
                server.shutdown()
                server.server_close()


if __name__ == '__main__':
    unittest.main()