codex-api-proxy 0.1.0__py3-none-any.whl

codex_api_proxy/__init__.py

@@ -0,0 +1,3 @@
+"""OpenAI-compatible HTTP proxy backed by the local Codex CLI."""
+
+__version__ = "0.1.0"

codex_api_proxy/app_server_runner.py

@@ -0,0 +1,554 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import signal
+import time
+import tomllib
+from collections.abc import AsyncIterator, Callable
+from pathlib import Path
+from typing import Any, Protocol
+
+from . import __version__
+
+
+class AppServerRunError(RuntimeError):
+    """Raised when the app-server engine cannot complete a turn."""
+
+
+class AppServerBusy(RuntimeError):
+    """Raised when no app-server worker is available before queue limits are hit."""
+
+
+DEFAULT_DISABLED_SKILL_NAMES = {
+    "imagegen",
+    "openai-docs",
+    "plugin-creator",
+    "skill-creator",
+    "skill-installer",
+}
+
+
+class AppServerClient(Protocol):
+    async def start(self) -> None: ...
+
+    async def close(self) -> None: ...
+
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: ...
+
+    async def read_message(self) -> dict[str, Any]: ...
+
+
+def _parse_config_value(raw: str) -> Any:
+    try:
+        return tomllib.loads(f"value = {raw}")["value"]
+    except tomllib.TOMLDecodeError:
+        return raw
+
+
+def codex_configs_to_object(codex_configs: list[str]) -> dict[str, Any]:
+    config: dict[str, Any] = {}
+    for item in codex_configs:
+        key, separator, raw_value = item.partition("=")
+        if not separator:
+            continue
+        target = config
+        parts = [part for part in key.strip().split(".") if part]
+        if not parts:
+            continue
+        for part in parts[:-1]:
+            next_target = target.setdefault(part, {})
+            if not isinstance(next_target, dict):
+                next_target = {}
+                target[part] = next_target
+            target = next_target
+        target[parts[-1]] = _parse_config_value(raw_value.strip())
+    return config
+
+
+def default_source_codex_home() -> Path:
+    return Path(os.environ.get("CODEX_HOME", Path.home() / ".codex")).expanduser()
+
+
+def default_app_server_codex_home() -> Path:
+    return Path.home() / ".codex-api-proxy" / "codex-home"
+
+
+def default_skill_roots(*, source_codex_home: Path | None = None) -> list[Path]:
+    source_home = source_codex_home or default_source_codex_home()
+    return [source_home / "skills", Path.home() / ".agents" / "skills"]
+
+
+def _parse_skill_name(skill_file: Path) -> str | None:
+    try:
+        with skill_file.open(encoding="utf-8") as handle:
+            for index, line in enumerate(handle):
+                if index > 80:
+                    return None
+                key, separator, value = line.partition(":")
+                if separator and key.strip() == "name":
+                    return value.strip().strip("\"'")
+    except OSError:
+        return None
+    return None
+
+
+def discover_skill_names(skill_roots: list[Path]) -> set[str]:
+    names = set(DEFAULT_DISABLED_SKILL_NAMES)
+    for root in skill_roots:
+        expanded_root = root.expanduser()
+        if not expanded_root.exists():
+            continue
+        for skill_file in [*expanded_root.glob("*/SKILL.md"), *expanded_root.glob(".system/*/SKILL.md")]:
+            names.add(_parse_skill_name(skill_file) or skill_file.parent.name)
+    return names
+
+
+def build_disabled_skills_config(*, skill_roots: list[Path] | None = None) -> str:
+    names = discover_skill_names(skill_roots or default_skill_roots())
+    entries = ",".join(f"{{name={json.dumps(name)},enabled=false}}" for name in sorted(names))
+    return f"skills.config=[{entries}]"
+
+
+def build_app_server_command(
+    *,
+    codex_bin: str,
+    codex_configs: list[str],
+    disabled_skills_config: str,
+) -> list[str]:
+    command = [
+        codex_bin,
+        "app-server",
+        "--stdio",
+        "--disable",
+        "apps",
+        "--disable",
+        "plugins",
+        "--disable",
+        "skill_mcp_dependency_install",
+    ]
+    for config in codex_configs:
+        command.extend(["-c", config])
+    command.extend(["-c", "mcp_servers={}", "-c", disabled_skills_config])
+    return command
+
+
+def prepare_isolated_codex_home(codex_home: Path, *, source_codex_home: Path | None = None) -> Path:
+    source_home = (source_codex_home or default_source_codex_home()).expanduser()
+    target_home = codex_home.expanduser()
+    if target_home.resolve(strict=False) == source_home.resolve(strict=False):
+        raise AppServerRunError("app-server CODEX_HOME must be isolated from the current Codex user home")
+
+    target_home.mkdir(parents=True, exist_ok=True)
+    source_auth = source_home / "auth.json"
+    target_auth = target_home / "auth.json"
+    if not source_auth.exists():
+        return target_home
+
+    if target_auth.is_symlink():
+        try:
+            if target_auth.resolve() == source_auth.resolve():
+                return target_home
+        except FileNotFoundError:
+            pass
+        target_auth.unlink()
+
+    if not target_auth.exists():
+        target_auth.symlink_to(source_auth)
+    return target_home
+
+
+def build_app_server_env(*, proxy: str | None, codex_home: Path) -> dict[str, str]:
+    env = os.environ.copy()
+    env["CODEX_HOME"] = str(codex_home)
+    if proxy:
+        env["http_proxy"] = proxy
+        env["https_proxy"] = proxy
+        env["HTTP_PROXY"] = proxy
+        env["HTTPS_PROXY"] = proxy
+    return env
+
+
+class StdioJsonRpcClient:
+    def __init__(
+        self,
+        *,
+        codex_bin: str,
+        codex_configs: list[str],
+        proxy: str | None,
+        codex_home: Path | None,
+        timeout_seconds: float,
+    ) -> None:
+        self.codex_bin = codex_bin
+        self.codex_configs = codex_configs
+        self.proxy = proxy
+        self.codex_home = codex_home
+        self.timeout_seconds = timeout_seconds
+        self.process: asyncio.subprocess.Process | None = None
+        self._next_id = 1
+        self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {}
+        self._notifications: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        self._reader_task: asyncio.Task[None] | None = None
+        self._stderr_task: asyncio.Task[None] | None = None
+        self._stderr_tail: list[str] = []
+        self._write_lock = asyncio.Lock()
+
+    async def start(self) -> None:
+        if self.process:
+            return
+        source_codex_home = default_source_codex_home()
+        codex_home = prepare_isolated_codex_home(
+            self.codex_home or default_app_server_codex_home(),
+            source_codex_home=source_codex_home,
+        )
+        command = build_app_server_command(
+            codex_bin=self.codex_bin,
+            codex_configs=self.codex_configs,
+            disabled_skills_config=build_disabled_skills_config(
+                skill_roots=default_skill_roots(source_codex_home=source_codex_home)
+            ),
+        )
+        env = build_app_server_env(proxy=self.proxy, codex_home=codex_home)
+        self.process = await asyncio.create_subprocess_exec(
+            *command,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=env,
+            start_new_session=True,
+        )
+        self._reader_task = asyncio.create_task(self._read_stdout())
+        self._stderr_task = asyncio.create_task(self._read_stderr())
+        await self.request(
+            "initialize",
+            {
+                "clientInfo": {"name": "codex-api-proxy", "version": __version__},
+                "capabilities": {"experimentalApi": True},
+            },
+        )
+
+    async def close(self) -> None:
+        process = self.process
+        if not process:
+            return
+        if process.returncode is None:
+            try:
+                process.send_signal(signal.SIGTERM)
+            except ProcessLookupError:
+                pass
+            try:
+                await asyncio.wait_for(process.wait(), timeout=2)
+            except TimeoutError:
+                process.kill()
+                await process.wait()
+        for task in (self._reader_task, self._stderr_task):
+            if task:
+                task.cancel()
+        self.process = None
+
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
+        await self.start()
+        request_id = self._next_id
+        self._next_id += 1
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future[dict[str, Any]] = loop.create_future()
+        self._pending[request_id] = future
+        await self._send({"jsonrpc": "2.0", "id": request_id, "method": method, "params": params})
+        try:
+            return await asyncio.wait_for(future, timeout=self.timeout_seconds)
+        finally:
+            self._pending.pop(request_id, None)
+
+    async def read_message(self) -> dict[str, Any]:
+        await self.start()
+        return await asyncio.wait_for(self._notifications.get(), timeout=self.timeout_seconds)
+
+    async def _send(self, payload: dict[str, Any]) -> None:
+        process = self.process
+        if not process or not process.stdin:
+            raise AppServerRunError("codex app-server process is not running")
+        data = json.dumps(payload, separators=(",", ":")).encode("utf-8") + b"\n"
+        async with self._write_lock:
+            process.stdin.write(data)
+            await process.stdin.drain()
+
+    async def _read_stdout(self) -> None:
+        process = self.process
+        if not process or not process.stdout:
+            return
+        try:
+            while line := await process.stdout.readline():
+                try:
+                    message = json.loads(line.decode("utf-8"))
+                except json.JSONDecodeError:
+                    continue
+                await self._dispatch_message(message)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            self._fail_pending(AppServerRunError(f"app-server stdout reader failed: {exc}"))
+        finally:
+            if process.returncode is None:
+                return
+            self._fail_pending(AppServerRunError(self._process_error_detail()))
+
+    async def _read_stderr(self) -> None:
+        process = self.process
+        if not process or not process.stderr:
+            return
+        try:
+            while line := await process.stderr.readline():
+                text = line.decode("utf-8", errors="replace").strip()
+                if text:
+                    self._stderr_tail.append(text)
+                    self._stderr_tail = self._stderr_tail[-20:]
+        except asyncio.CancelledError:
+            raise
+
+    async def _dispatch_message(self, message: Any) -> None:
+        if not isinstance(message, dict):
+            return
+        request_id = message.get("id")
+        method = message.get("method")
+        if request_id is not None and method is None:
+            future = self._pending.get(request_id)
+            if future and not future.done():
+                if "error" in message:
+                    future.set_exception(AppServerRunError(str(message["error"])))
+                else:
+                    result = message.get("result") or {}
+                    future.set_result(result if isinstance(result, dict) else {"value": result})
+            return
+        if request_id is not None and method is not None:
+            await self._send(
+                {
+                    "jsonrpc": "2.0",
+                    "id": request_id,
+                    "error": {"code": -32601, "message": f"unsupported server request: {method}"},
+                }
+            )
+            return
+        await self._notifications.put(message)
+
+    def _fail_pending(self, exc: BaseException) -> None:
+        for future in self._pending.values():
+            if not future.done():
+                future.set_exception(exc)
+
+    def _process_error_detail(self) -> str:
+        if self._stderr_tail:
+            return "codex app-server exited: " + " ".join(self._stderr_tail[-5:])
+        return "codex app-server exited"
+
+
+LatencyCallback = Callable[[str, float], None]
+
+
+def _elapsed_ms(started_at: float) -> float:
+    return round((time.perf_counter() - started_at) * 1000, 2)
+
+
+def _record_latency(callback: LatencyCallback | None, name: str, started_at: float) -> None:
+    if callback:
+        callback(name, _elapsed_ms(started_at))
+
+
+async def stream_app_server_turn(
+    *,
+    client: AppServerClient,
+    cwd: Path,
+    prompt: str,
+    model: str | None,
+    codex_configs: list[str],
+    ephemeral: bool,
+    timeout_seconds: float,
+    latency_callback: LatencyCallback | None = None,
+) -> AsyncIterator[str]:
+    await client.start()
+    thread_id: str | None = None
+    config = codex_configs_to_object(codex_configs)
+    thread_params: dict[str, Any] = {
+        "approvalPolicy": "never",
+        "config": config,
+        "cwd": str(cwd),
+        "dynamicTools": [],
+        "environments": [],
+        "ephemeral": ephemeral,
+        "model": model,
+        "sandbox": "read-only",
+    }
+    phase_started_at = time.perf_counter()
+    start_response = await asyncio.wait_for(client.request("thread/start", thread_params), timeout=timeout_seconds)
+    _record_latency(latency_callback, "app_server_thread_start", phase_started_at)
+    thread = start_response.get("thread")
+    if not isinstance(thread, dict) or not isinstance(thread.get("id"), str):
+        raise AppServerRunError("app-server thread/start returned no thread id")
+    thread_id = thread["id"]
+    phase_started_at = time.perf_counter()
+    turn_response = await asyncio.wait_for(
+        client.request("turn/start", {"input": [{"type": "text", "text": prompt}], "threadId": thread_id}),
+        timeout=timeout_seconds,
+    )
+    _record_latency(latency_callback, "app_server_turn_start", phase_started_at)
+    turn = turn_response.get("turn")
+    if not isinstance(turn, dict) or not isinstance(turn.get("id"), str):
+        raise AppServerRunError("app-server turn/start returned no turn id")
+    turn_id = turn["id"]
+    emitted = False
+    first_delta_recorded = False
+    final_text: str | None = None
+    event_wait_started_at = time.perf_counter()
+    try:
+        while True:
+            message = await asyncio.wait_for(client.read_message(), timeout=timeout_seconds)
+            method = message.get("method")
+            params = message.get("params")
+            if not isinstance(params, dict):
+                continue
+            if params.get("threadId") != thread_id:
+                continue
+            if method == "item/agentMessage/delta" and params.get("turnId") == turn_id:
+                delta = params.get("delta")
+                if isinstance(delta, str) and delta:
+                    if not first_delta_recorded:
+                        first_delta_recorded = True
+                        _record_latency(latency_callback, "app_server_first_delta", event_wait_started_at)
+                    emitted = True
+                    yield delta
+            elif method == "item/completed" and params.get("turnId") == turn_id:
+                item = params.get("item")
+                if isinstance(item, dict) and item.get("type") == "agentMessage" and isinstance(item.get("text"), str):
+                    final_text = item["text"]
+            elif method == "turn/completed":
+                turn = params.get("turn")
+                if isinstance(turn, dict) and turn.get("id") == turn_id:
+                    _record_latency(latency_callback, "app_server_turn_complete", event_wait_started_at)
+                    if final_text and not emitted:
+                        yield final_text
+                    return
+    finally:
+        if thread_id:
+            try:
+                phase_started_at = time.perf_counter()
+                await client.request("thread/archive", {"threadId": thread_id})
+                _record_latency(latency_callback, "app_server_thread_archive", phase_started_at)
+            except Exception:
+                pass
+
+
+ClientFactory = Callable[[], AppServerClient]
+
+
+class AppServerWorkerPool:
+    def __init__(
+        self,
+        *,
+        client_factory: ClientFactory,
+        workers: int,
+        max_queue_size: int,
+        queue_timeout_seconds: float,
+    ) -> None:
+        self.client_factory = client_factory
+        self.worker_count = workers
+        self.max_queue_size = max_queue_size
+        self.queue_timeout_seconds = queue_timeout_seconds
+        self._idle: asyncio.Queue[AppServerClient] = asyncio.Queue()
+        self._started = False
+        self._start_lock = asyncio.Lock()
+        self._waiters = 0
+        self._waiters_lock = asyncio.Lock()
+
+    async def start(self) -> None:
+        async with self._start_lock:
+            if self._started:
+                return
+            for _ in range(self.worker_count):
+                client = self.client_factory()
+                await client.start()
+                await self._idle.put(client)
+            self._started = True
+
+    async def close(self) -> None:
+        while not self._idle.empty():
+            client = await self._idle.get()
+            await client.close()
+
+    async def stream_completion(
+        self,
+        *,
+        cwd: Path,
+        prompt: str,
+        model: str | None,
+        codex_configs: list[str],
+        ephemeral: bool,
+        timeout_seconds: float,
+        latency_callback: LatencyCallback | None = None,
+    ) -> AsyncIterator[str]:
+        client = await self._acquire()
+        healthy = True
+        try:
+            async for chunk in stream_app_server_turn(
+                client=client,
+                cwd=cwd,
+                prompt=prompt,
+                model=model,
+                codex_configs=codex_configs,
+                ephemeral=ephemeral,
+                timeout_seconds=timeout_seconds,
+                latency_callback=latency_callback,
+            ):
+                yield chunk
+        except Exception:
+            healthy = False
+            await client.close()
+            raise
+        finally:
+            if healthy:
+                await self._idle.put(client)
+            else:
+                replacement = self.client_factory()
+                try:
+                    await replacement.start()
+                    await self._idle.put(replacement)
+                except Exception:
+                    await replacement.close()
+
+    async def _acquire(self) -> AppServerClient:
+        await self.start()
+        async with self._waiters_lock:
+            if self._waiters >= self.max_queue_size:
+                raise AppServerBusy("app-server worker queue is full")
+            self._waiters += 1
+        try:
+            return await asyncio.wait_for(self._idle.get(), timeout=self.queue_timeout_seconds)
+        except TimeoutError as exc:
+            raise AppServerBusy("timed out waiting for an app-server worker") from exc
+        finally:
+            async with self._waiters_lock:
+                self._waiters -= 1
+
+
+def create_stdio_app_server_pool(
+    *,
+    codex_bin: str,
+    proxy: str | None,
+    codex_home: Path | None,
+    codex_configs: list[str],
+    workers: int,
+    max_queue_size: int,
+    queue_timeout_seconds: float,
+    timeout_seconds: float,
+) -> AppServerWorkerPool:
+    return AppServerWorkerPool(
+        client_factory=lambda: StdioJsonRpcClient(
+            codex_bin=codex_bin,
+            codex_configs=codex_configs,
+            proxy=proxy,
+            codex_home=codex_home,
+            timeout_seconds=timeout_seconds,
+        ),
+        workers=workers,
+        max_queue_size=max_queue_size,
+        queue_timeout_seconds=queue_timeout_seconds,
+    )