from __future__ import annotations import asyncio import contextlib import json import logging import os import signal from websockets.asyncio.client import unix_connect from websockets.exceptions import ConnectionClosed, WebSocketException from websockets.protocol import State as WebSocketState from .models import ControlError class JsonProcess: """A supervised JSON-lines client, never a terminal or stdout scraper.""" def __init__(self, command: list[str], *, env: dict | None = None, on_event=None, on_disconnect=None): self.command, self.env = command, env self.on_event, self.on_disconnect = on_event, on_disconnect self.process = None self.pending = {} self.serial = 0 self.start_lock = asyncio.Lock() self.write_lock = asyncio.Lock() self.reader = None self.stderr_reader = None self.generation = 0 async def start(self) -> None: async with self.start_lock: if self.process and self.process.returncode is None: return try: self.process = await asyncio.create_subprocess_exec( *self.command, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=self.env or os.environ.copy(), limit=8 * 1024 * 1024, start_new_session=True, ) except OSError as exc: raise ControlError("agent transport could not start; check service configuration") from exc self.generation += 1 self.reader = asyncio.create_task(self.read(self.process)) self.stderr_reader = asyncio.create_task(self.drain_stderr(self.process)) async def drain_stderr(self, process): # SDKs may print sensitive provider errors. Drain without exposing them in HTTP/logs. while await process.stderr.read(8192): pass async def write(self, message): async with self.write_lock: self.process.stdin.write((json.dumps(message) + "\n").encode()) await self.process.stdin.drain() async def read(self, process): try: while line := await process.stdout.readline(): message = json.loads(line) if "method" in message: if self.on_event: await self.on_event(message) elif (future := self.pending.pop(message.get("id"), None)) is not None and not future.done(): future.set_result(message) except (ValueError, OSError): logging.getLogger(__name__).warning("agent transport disconnected") finally: for future in self.pending.values(): if not future.done(): future.set_exception(ControlError("agent transport disconnected; input outcome may be unknown")) self.pending.clear() if process.returncode is None: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGTERM) await process.wait() # Reap any tools left in the host's process group after an abnormal exit. with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) if self.on_disconnect: self.on_disconnect() async def call(self, method: str, params: dict, timeout: float = 45): await self.start() self.serial += 1 request_id = self.serial future = asyncio.get_running_loop().create_future() self.pending[request_id] = future try: await self.write({"id": request_id, "method": method, "params": params}) response = await asyncio.wait_for(future, timeout) if "error" in response: error = response["error"] raise ControlError(error.get("message", "agent request rejected"), error.get("status", 502)) return response.get("result") except TimeoutError as exc: raise ControlError("agent request timed out; do not automatically retry input", 504) from exc except (OSError, ConnectionError) as exc: raise ControlError("agent connection failed; input outcome may be unknown") from exc finally: self.pending.pop(request_id, None) async def close(self): if not self.process: return process = self.process if process.returncode is None: process.stdin.close() try: await asyncio.wait_for(process.wait(), 12) except TimeoutError: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGTERM) try: await asyncio.wait_for(process.wait(), 3) except TimeoutError: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) await process.wait() for task in (self.reader, self.stderr_reader): if task: with contextlib.suppress(asyncio.CancelledError): await task self.process = None class JsonUnixSocket(JsonProcess): """Official WebSocket framing over the existing daemon's Unix socket.""" def __init__(self, path: str, **kwargs): super().__init__([], **kwargs) self.path = path self.socket = None async def start(self): async with self.start_lock: # A socket object lingers after the peer goes away. Reusing it would # pin the adapter to a dead daemon forever, so only an open socket # counts as connected and a closed one is replaced. if self.socket is not None: if self.socket.state == WebSocketState.OPEN: return self.socket = None try: # Codex 0.154.0's control socket rejects permessage-deflate offers. self.socket = await unix_connect(self.path, uri="ws://localhost/", compression=None, open_timeout=5, close_timeout=3, max_size=8 * 1024 * 1024) except (OSError, TimeoutError, WebSocketException) as exc: raise ControlError("Codex daemon unavailable; check the configured local socket") from exc self.generation += 1 self.reader = asyncio.create_task(self.read_socket(self.socket)) async def write(self, message): async with self.write_lock: try: await self.socket.send(json.dumps(message)) except ConnectionClosed as exc: raise ControlError("Codex connection lost; input outcome may be unknown") from exc async def reply(self, message, generation): # Never call start() here: server request IDs are connection-bound. async with self.write_lock: if generation != self.generation or self.socket is None: return False socket = self.socket try: await socket.send(json.dumps(message)) return True except ConnectionClosed: return False async def read_socket(self, socket): try: async for payload in socket: message = json.loads(payload) if "method" in message: if self.on_event: await self.on_event(message) elif (future := self.pending.pop(message.get("id"), None)) is not None and not future.done(): future.set_result(message) except (ValueError, ConnectionClosed, ControlError): pass finally: # Only clear the socket this reader owns. A reconnect may already # have installed a replacement, and nulling it would strand the # adapter on a closed connection again. if self.socket is socket: self.socket = None for future in self.pending.values(): if not future.done(): future.set_exception(ControlError("Codex disconnected; input outcome may be unknown")) self.pending.clear() if self.on_disconnect: self.on_disconnect() async def close(self): if self.socket is not None: await self.socket.close() if self.reader: await self.reader