from __future__ import annotations import asyncio import json import time from dataclasses import dataclass, field from pathlib import Path from typing import Any import websockets from app.codex_runner import CodexResult class AppServerProtocolError(RuntimeError): pass @dataclass class _TurnCollector: agent_deltas: list[str] = field(default_factory=list) final_messages: list[str] = field(default_factory=list) commands: list[str] = field(default_factory=list) errors: list[str] = field(default_factory=list) done: bool = False status: str | None = None def handle(self, message: dict[str, Any]) -> None: method = message.get("method") params = message.get("params") or {} if not isinstance(params, dict): return if method == "item/agentMessage/delta": delta = params.get("delta") if isinstance(delta, str): self.agent_deltas.append(delta) return if method == "item/completed": item = params.get("item") if isinstance(item, dict): self._handle_completed_item(item) return if method == "turn/completed": self.done = True turn = params.get("turn") if isinstance(turn, dict): self.status = turn.get("status") error = turn.get("error") if isinstance(error, dict) and error.get("message"): self.errors.append(str(error["message"])) return def _handle_completed_item(self, item: dict[str, Any]) -> None: item_type = item.get("type") if item_type in {"agentMessage", "agent_message"} and item.get("text"): self.final_messages.append(str(item["text"])) return if item_type in {"commandExecution", "command_execution"}: command = item.get("command") if isinstance(command, list): self.commands.append(" ".join(str(part) for part in command)) elif command: self.commands.append(str(command)) def final_text(self) -> str: if self.final_messages: return self.final_messages[-1] if self.agent_deltas: return "".join(self.agent_deltas).strip() if self.errors: return "\n".join(self.errors) return "(Codex finished without text output.)" class AppServerRunner: def __init__( self, url: str, timeout_seconds: int, model: str | None = None, effort: str | None = None, service_tier: str | None = None, ) -> None: self.url = url self.timeout_seconds = timeout_seconds self.model = model self.effort = effort self.service_tier = service_tier self._next_id = 1 async def run( self, prompt: str, cwd: Path, sandbox: str, session_id: str | None, skip_git_repo_check: bool = False, ) -> CodexResult: del skip_git_repo_check try: return await asyncio.wait_for( self._run_once(prompt, cwd, sandbox, session_id), timeout=self.timeout_seconds, ) except asyncio.TimeoutError: return CodexResult( ok=False, final_message=f"Codex app-server timed out after {self.timeout_seconds} seconds.", session_id=session_id, ) except Exception as exc: return CodexResult( ok=False, final_message=f"Codex app-server failed: {exc}", session_id=session_id, ) async def call(self, method: str, params: dict[str, Any] | None = None) -> Any: collector = _TurnCollector() async with websockets.connect( self.url, open_timeout=15, ping_interval=20, max_size=16 * 1024 * 1024, ) as ws: await self._initialize(ws, collector) return await self._request(ws, method, params or {}, collector) async def list_threads( self, limit: int = 10, archived: bool = False, cursor: str | None = None, ) -> dict[str, Any]: params: dict[str, Any] = { "limit": limit, "archived": archived, } if cursor: params["cursor"] = cursor result = await self.call("thread/list", params) return result if isinstance(result, dict) else {} async def loaded_threads(self) -> dict[str, Any]: result = await self.call("thread/loaded/list", {}) return result if isinstance(result, dict) else {} async def archive_thread(self, thread_id: str) -> None: await self.call("thread/archive", {"threadId": thread_id}) async def unarchive_thread(self, thread_id: str) -> None: await self.call("thread/unarchive", {"threadId": thread_id}) async def resume_thread(self, thread_id: str) -> str: result = await self.call("thread/resume", {"threadId": thread_id}) return _extract_thread_id(result) or thread_id async def _run_once( self, prompt: str, cwd: Path, sandbox: str, session_id: str | None, ) -> CodexResult: collector = _TurnCollector() async with websockets.connect( self.url, open_timeout=15, ping_interval=20, max_size=16 * 1024 * 1024, ) as ws: await self._initialize(ws, collector) thread_id = await self._ensure_thread(ws, cwd, sandbox, session_id, collector) turn_params: dict[str, Any] = { "threadId": thread_id, "input": [{"type": "text", "text": prompt}], } await self._request_with_settings_fallback( ws, "turn/start", turn_params, collector, ) while not collector.done: message = await self._recv(ws) await self._handle_non_response(ws, message, collector) ok = collector.status not in {"failed", "interrupted"} return CodexResult( ok=ok, final_message=collector.final_text(), session_id=thread_id, commands=collector.commands, stderr="\n".join(collector.errors), returncode=0 if ok else 1, ) async def _ensure_thread( self, ws: Any, cwd: Path, sandbox: str, session_id: str | None, collector: _TurnCollector, ) -> str: if session_id: try: result = await self._request_with_settings_fallback( ws, "thread/resume", {"threadId": session_id}, collector, ) return _extract_thread_id(result) or session_id except AppServerProtocolError: pass result = await self._request_with_settings_fallback( ws, "thread/start", {}, collector, ) thread_id = _extract_thread_id(result) if not thread_id: raise AppServerProtocolError("thread/start did not return a thread id") return thread_id async def _initialize(self, ws: Any, collector: _TurnCollector) -> None: await self._request( ws, "initialize", { "clientInfo": { "name": "telegram_codex_bot", "title": "Telegram Codex Bot", "version": "0.1.0", }, "capabilities": {"experimentalApi": True}, }, collector, ) await self._notify(ws, "initialized", {}) def _settings_params(self) -> dict[str, Any]: params: dict[str, Any] = {} if self.model: params["model"] = self.model if self.effort: params["effort"] = self.effort if self.service_tier: params["serviceTier"] = self.service_tier return params def _settings_variants(self) -> list[tuple[str, dict[str, Any]]]: full = self._settings_params() variants: list[tuple[str, dict[str, Any]]] = [] seen: set[str] = set() def add(label: str, params: dict[str, Any]) -> None: key = json.dumps(params, sort_keys=True, ensure_ascii=True) if key in seen: return seen.add(key) variants.append((label, params)) add("model/effort/speed", full) if "serviceTier" in full: add("model/effort", {k: v for k, v in full.items() if k != "serviceTier"}) if "effort" in full: add("model/speed", {k: v for k, v in full.items() if k != "effort"}) if "model" in full: add("model", {"model": full["model"]}) add("default settings", {}) return variants async def _request_with_settings_fallback( self, ws: Any, method: str, base_params: dict[str, Any], collector: _TurnCollector, ) -> Any: errors: list[str] = [] last_error: AppServerProtocolError | None = None for label, settings in self._settings_variants(): params = dict(base_params) params.update(settings) try: result = await self._request(ws, method, params, collector) if errors: collector.errors.append( "Codex app-server rejected one or more runtime setting " f"overrides; continued with {label}. Last error: {errors[-1]}" ) return result except AppServerProtocolError as exc: last_error = exc errors.append(f"{label}: {exc}") if not settings: break raise last_error or AppServerProtocolError(f"{method} failed") async def _request( self, ws: Any, method: str, params: dict[str, Any], collector: _TurnCollector, ) -> Any: request_id = self._next_request_id() await ws.send(json.dumps({"method": method, "id": request_id, "params": params})) while True: message = await self._recv(ws) if message.get("id") == request_id: if message.get("error"): error = message["error"] if isinstance(error, dict): raise AppServerProtocolError(error.get("message") or str(error)) raise AppServerProtocolError(str(error)) return message.get("result") await self._handle_non_response(ws, message, collector) async def _notify(self, ws: Any, method: str, params: dict[str, Any]) -> None: await ws.send(json.dumps({"method": method, "params": params})) async def _recv(self, ws: Any) -> dict[str, Any]: raw = await ws.recv() if isinstance(raw, bytes): raw = raw.decode("utf-8", errors="replace") try: message = json.loads(raw) except json.JSONDecodeError as exc: raise AppServerProtocolError(f"invalid JSON from app-server: {exc}") from exc if not isinstance(message, dict): raise AppServerProtocolError("app-server returned a non-object message") return message async def _handle_non_response( self, ws: Any, message: dict[str, Any], collector: _TurnCollector, ) -> None: if "id" in message and "method" in message: await ws.send( json.dumps( { "id": message["id"], "error": { "code": -32601, "message": "Telegram bot client cannot handle this app-server request.", }, } ) ) return collector.handle(message) def _next_request_id(self) -> int: request_id = self._next_id self._next_id += 1 return request_id def _extract_thread_id(result: Any) -> str | None: if not isinstance(result, dict): return None thread = result.get("thread") if isinstance(thread, dict) and thread.get("id"): return str(thread["id"]) if result.get("threadId"): return str(result["threadId"]) return None class PersistentAppServerRunner(AppServerRunner): def __init__( self, url: str, timeout_seconds: int, model: str | None = None, effort: str | None = None, service_tier: str | None = None, ) -> None: super().__init__( url=url, timeout_seconds=timeout_seconds, model=model, effort=effort, service_tier=service_tier, ) self._ws: Any | None = None self._lock = asyncio.Lock() self.last_used_monotonic = time.monotonic() @property def is_connected(self) -> bool: return self._ws is not None def idle_seconds(self) -> float: return time.monotonic() - self.last_used_monotonic async def close(self) -> None: ws = self._ws self._ws = None if ws is not None: await ws.close() async def close_if_idle(self, idle_seconds: int) -> bool: if self._ws is None: return False if self.idle_seconds() < idle_seconds: return False async with self._lock: if self._ws is None or self.idle_seconds() < idle_seconds: return False await self.close() return True async def run( self, prompt: str, cwd: Path, sandbox: str, session_id: str | None, skip_git_repo_check: bool = False, ) -> CodexResult: del skip_git_repo_check async with self._lock: try: return await asyncio.wait_for( self._run_persistent(prompt, cwd, sandbox, session_id), timeout=self.timeout_seconds, ) except asyncio.TimeoutError: await self.close() return CodexResult( ok=False, final_message=f"Codex app-server timed out after {self.timeout_seconds} seconds.", session_id=session_id, ) except Exception as exc: await self.close() return CodexResult( ok=False, final_message=f"Codex app-server failed: {exc}", session_id=session_id, ) async def _run_persistent( self, prompt: str, cwd: Path, sandbox: str, session_id: str | None, ) -> CodexResult: try: return await self._run_on_connection(prompt, cwd, sandbox, session_id) except Exception: await self.close() return await self._run_on_connection(prompt, cwd, sandbox, session_id) async def _run_on_connection( self, prompt: str, cwd: Path, sandbox: str, session_id: str | None, ) -> CodexResult: self.last_used_monotonic = time.monotonic() collector = _TurnCollector() ws = await self._ensure_connection(collector) thread_id = await self._ensure_thread(ws, cwd, sandbox, session_id, collector) turn_params: dict[str, Any] = { "threadId": thread_id, "input": [{"type": "text", "text": prompt}], } await self._request_with_settings_fallback( ws, "turn/start", turn_params, collector, ) while not collector.done: message = await self._recv(ws) await self._handle_non_response(ws, message, collector) ok = collector.status not in {"failed", "interrupted"} self.last_used_monotonic = time.monotonic() return CodexResult( ok=ok, final_message=collector.final_text(), session_id=thread_id, commands=collector.commands, stderr="\n".join(collector.errors), returncode=0 if ok else 1, ) async def _ensure_connection(self, collector: _TurnCollector) -> Any: if self._ws is not None: return self._ws self._ws = await websockets.connect( self.url, open_timeout=15, ping_interval=20, max_size=16 * 1024 * 1024, ) await self._initialize(self._ws, collector) return self._ws def _thread_sandbox(sandbox: str) -> str: return { "read-only": "readOnly", "workspace-write": "workspaceWrite", "danger-full-access": "dangerFullAccess", }.get(sandbox, "readOnly")