rain34572/responses-adapter-gateway

1from __future__ import annotations2 3import asyncio4import json5import os6from dataclasses import dataclass7from pathlib import Path8from typing import Any9 10 11class CodexRunnerError(RuntimeError):12    pass13 14 15@dataclass16class CodexRunnerConfig:17    codex_bin: str18    workdir: str19    codex_home: str20    model: str21    model_provider: str22    provider_name: str23    api_base_url: str24    api_key: str25    timeout_seconds: float26 27    @classmethod28    def from_env(29        cls,30        default_model: str,31        default_api_base_url: str,32        default_api_key: str,33        default_timeout_seconds: float,34    ) -> "CodexRunnerConfig":35        return cls(36            codex_bin=os.getenv("CODEX_BIN", "codex").strip() or "codex",37            workdir=os.getenv("CODEX_WORKDIR", "/tmp").strip() or "/tmp",38            codex_home=os.getenv("CODEX_HOME", str(Path.home() / ".codex")).strip() or str(Path.home() / ".codex"),39            model=os.getenv("CODEX_MODEL", default_model).strip() or default_model,40            model_provider=os.getenv("CODEX_MODEL_PROVIDER", "relaygw").strip() or "relaygw",41            provider_name=os.getenv("CODEX_PROVIDER_NAME", "Relay Gateway").strip() or "Relay Gateway",42            api_base_url=os.getenv("CODEX_API_BASE_URL", default_api_base_url).strip().rstrip("/"),43            api_key=os.getenv("CODEX_API_KEY", default_api_key).strip() or default_api_key,44            timeout_seconds=float(os.getenv("CODEX_TIMEOUT_SECONDS", str(default_timeout_seconds))),45        )46 47    @classmethod48    def from_runtime(49        cls,50        model: str,51        api_base_url: str,52        api_key: str,53        timeout_seconds: float,54    ) -> "CodexRunnerConfig":55        return cls(56            codex_bin=os.getenv("CODEX_BIN", "codex").strip() or "codex",57            workdir=os.getenv("CODEX_WORKDIR", "/tmp").strip() or "/tmp",58            codex_home=os.getenv("CODEX_HOME", str(Path.home() / ".codex")).strip() or str(Path.home() / ".codex"),59            model=model,60            model_provider=os.getenv("CODEX_MODEL_PROVIDER", "relaygw").strip() or "relaygw",61            provider_name=os.getenv("CODEX_PROVIDER_NAME", "Relay Gateway").strip() or "Relay Gateway",62            api_base_url=api_base_url.strip().rstrip("/"),63            api_key=api_key.strip(),64            timeout_seconds=timeout_seconds,65        )66 67 68def _escape_toml(value: str) -> str:69    return value.replace("\\", "\\\\").replace('"', '\\"')70 71 72class CodexRunner:73    def __init__(self, config: CodexRunnerConfig):74        self.config = config75        self._session_threads: dict[str, str] = {}76 77    async def run(78        self,79        prompt: str,80        system_prompt: str,81        session_id: str | None,82        previous_response_id: str | None,83        metadata: dict[str, Any],84    ) -> dict[str, Any]:85        self._ensure_provider_files()86        full_prompt = self._build_prompt(87            prompt=prompt,88            system_prompt=system_prompt,89            session_id=session_id,90            previous_response_id=previous_response_id,91            metadata=metadata,92        )93        thread_id = self._resolve_thread_id(session_id, previous_response_id)94        cmd = self._build_command(thread_id)95 96        env = os.environ.copy()97        env["CODEX_HOME"] = self.config.codex_home98        if self.config.api_key:99            env["OPENAI_API_KEY"] = self.config.api_key100 101        process = None102        stdout = b""103        stderr = b""104        try:105            process = await asyncio.create_subprocess_exec(106                *cmd,107                stdin=asyncio.subprocess.PIPE,108                stdout=asyncio.subprocess.PIPE,109                stderr=asyncio.subprocess.PIPE,110                env=env,111                cwd=self.config.workdir,112            )113            stdout, stderr = await asyncio.wait_for(114                process.communicate(input=full_prompt.encode("utf-8")),115                timeout=self.config.timeout_seconds,116            )117        except FileNotFoundError as exc:118            raise CodexRunnerError(f"codex binary not found: {self.config.codex_bin}") from exc119        except asyncio.TimeoutError as exc:120            raise CodexRunnerError("codex exec timed out") from exc121 122        if process is None or process.returncode != 0:123            stderr_text = stderr.decode("utf-8", errors="ignore").strip()124            stdout_text = stdout.decode("utf-8", errors="ignore").strip()125            code = process.returncode if process is not None else "unknown"126            raise CodexRunnerError(stderr_text or stdout_text or f"codex exited with {code}")127 128        parsed = self._parse_exec_jsonl(129            stdout.decode("utf-8", errors="ignore"),130            fallback_thread_id=thread_id,131        )132        reply = parsed["reply"]133        if not reply:134            raise CodexRunnerError("codex returned empty output")135        response_id = parsed["thread_id"]136        if session_id and response_id:137            self._session_threads[session_id] = response_id138 139        return {140            "reply": reply,141            "model": self.config.model,142            "session_id": session_id,143            "response_id": response_id,144            "usage": None,145            "raw": {146                "runtime": "codex_cli",147                "events": parsed["events"],148            },149        }150 151    @staticmethod152    def _build_prompt(153        prompt: str,154        system_prompt: str,155        session_id: str | None,156        previous_response_id: str | None,157        metadata: dict[str, Any],158    ) -> str:159        lines = [160            "You are running as a relay-connected coding agent.",161            "",162            "System instructions:",163            system_prompt,164            "",165        ]166        if session_id:167            lines.append(f"Session ID: {session_id}")168        if previous_response_id:169            lines.append(f"Previous response ID: {previous_response_id}")170        if metadata:171            lines.append(f"Metadata: {metadata}")172        lines.extend(["", "User prompt:", prompt.strip()])173        return "\n".join(lines).strip()174 175    def _resolve_thread_id(176        self,177        session_id: str | None,178        previous_response_id: str | None,179    ) -> str | None:180        if session_id:181            cached = self._session_threads.get(session_id)182            if cached:183                return cached184        if isinstance(previous_response_id, str) and previous_response_id.strip():185            return previous_response_id.strip()186        return None187 188    def _build_command(self, thread_id: str | None) -> list[str]:189        base = [190            self.config.codex_bin,191            "exec",192        ]193        if thread_id:194            base.append("resume")195        base.append("--skip-git-repo-check")196        if not thread_id:197            base.extend([198                "--sandbox",199                "danger-full-access",200            ])201        base.extend([202            "--model",203            self.config.model,204            "-c",205            f'model_provider="{_escape_toml(self.config.model_provider)}"',206        ])207        if thread_id:208            base.extend([thread_id, "--json", "-"])209        else:210            base.extend(["--json", "--cd", self.config.workdir, "-"])211        return base212 213    def _ensure_provider_files(self) -> None:214        home = Path(self.config.codex_home)215        home.mkdir(parents=True, exist_ok=True)216        self._write_provider_config(home)217        self._write_auth_config(home)218 219    def _write_provider_config(self, home: Path) -> None:220        cfg_path = home / "config.toml"221        existing = cfg_path.read_text(encoding="utf-8") if cfg_path.exists() else ""222        section = self._build_provider_section()223        updated = self._upsert_provider_section(existing, section)224        cfg_path.write_text(updated, encoding="utf-8")225 226    def _write_auth_config(self, home: Path) -> None:227        if not self.config.api_key:228            return229        auth_path = home / "auth.json"230        payload = {231            "OPENAI_API_KEY": self.config.api_key,232            "auth_mode": "apikey",233        }234        auth_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")235        os.chmod(auth_path, 0o600)236 237    def _build_provider_section(self) -> str:238        lines = [239            f"[model_providers.{self.config.model_provider}]",240            f'name = "{_escape_toml(self.config.provider_name)}"',241        ]242        if self.config.api_base_url:243            lines.append(f'base_url = "{_escape_toml(self.config.api_base_url)}"')244        lines.extend([245            'env_key = "OPENAI_API_KEY"',246            'wire_api = "responses"',247            "",248        ])249        return "\n".join(lines)250 251    def _upsert_provider_section(self, content: str, section: str) -> str:252        header = f"[model_providers.{self.config.model_provider}]"253        subsection_prefix = f"[model_providers.{self.config.model_provider}."254        if header not in content:255            trimmed = content.rstrip()256            return f"{trimmed}\n\n{section}".strip() + "\n"257 258        lines = content.splitlines()259        kept: list[str] = []260        skipping = False261        inserted = False262        for line in lines:263            stripped = line.strip()264            if stripped == header:265                if not inserted:266                    kept.extend(section.rstrip().splitlines())267                    inserted = True268                skipping = True269                continue270            if skipping and stripped.startswith("[") and stripped != header and not stripped.startswith(subsection_prefix):271                skipping = False272            if not skipping:273                kept.append(line)274        if not inserted:275            kept.extend(["", *section.rstrip().splitlines()])276        return "\n".join(kept).strip() + "\n"277 278    def _parse_exec_jsonl(self, raw_output: str, fallback_thread_id: str | None) -> dict[str, Any]:279        thread_id = fallback_thread_id280        reply_parts: list[str] = []281        events: list[dict[str, Any]] = []282        for line in raw_output.splitlines():283            line = line.strip()284            if not line:285                continue286            try:287                event = json.loads(line)288            except json.JSONDecodeError:289                continue290            if not isinstance(event, dict):291                continue292            events.append(event)293            event_type = event.get("type")294            if event_type == "thread.started" and isinstance(event.get("thread_id"), str):295                thread_id = event["thread_id"].strip() or thread_id296                continue297            if event_type == "turn.failed":298                error = event.get("error")299                if isinstance(error, dict) and isinstance(error.get("message"), str):300                    raise CodexRunnerError(error["message"])301                raise CodexRunnerError("codex turn failed")302            if event_type != "item.completed":303                continue304            item = event.get("item")305            if not isinstance(item, dict):306                continue307            if item.get("type") not in {"agent_message", "message"}:308                continue309            text = self._extract_item_text(item)310            if text:311                reply_parts.append(text)312        return {313            "thread_id": thread_id,314            "reply": "\n".join(reply_parts).strip(),315            "events": events,316        }317 318    def _extract_item_text(self, item: dict[str, Any]) -> str:319        for key in ("output_text", "content", "text"):320            text = self._flatten_text(item.get(key))321            if text:322                return text323        return ""324 325    def _flatten_text(self, value: Any) -> str:326        if isinstance(value, str):327            return value.strip()328        if isinstance(value, list):329            parts = [self._flatten_text(item) for item in value]330            return "\n".join([part for part in parts if part]).strip()331        if isinstance(value, dict):332            for key in ("text", "content", "output_text"):333                text = self._flatten_text(value.get(key))334                if text:335                    return text336        return ""337