Spaces:
Running
Running
File size: 12,457 Bytes
19491c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 | """
Warm Codex app-server — ONE persistent process, reused across requests, to remove
per-request cold start (spawn + initialize + cold thread setup). A lock serializes
turns; this also avoids two processes refreshing (and rotating) auth.json at once.
Opt-in via CODEX_ENGINE=pool. Same event contract as codex_engine.run_turn:
{"type": "delta"|"reasoning", "text": ...}
{"type": "final", "text", "thread_id", "usage", "images"}
Falls back safely: on any process/protocol failure the process is killed and the
next request respawns it (and re-reads auth.json, e.g. after a re-upload).
"""
import asyncio
import json
import os
from pathlib import Path
from typing import AsyncIterator, Optional
from codex_engine import CodexError
_AUTH_DEAD = ("session has ended", "log in again", "failed to refresh token")
_STREAM_LIMIT = 16 * 1024 * 1024
class WarmCodex:
def __init__(self, codex_bin: str, codex_home: str, read_timeout: float):
self.bin = codex_bin
self.home = codex_home
self.read_timeout = read_timeout
self.proc: Optional[asyncio.subprocess.Process] = None
self.lock = asyncio.Lock() # one turn at a time
self._pending: dict[int, asyncio.Future] = {}
self._queue: Optional[asyncio.Queue] = None # current turn's notifications
self._reader: Optional[asyncio.Task] = None
self._next_id = 100
self._threads: dict[str, str] = {} # session_id -> thread_id (warm)
# -- process lifecycle ---------------------------------------------------
async def _ensure(self) -> None:
if self.proc and self.proc.returncode is None:
return
self.proc = await asyncio.create_subprocess_exec(
self.bin, "app-server",
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.DEVNULL,
cwd="/tmp",
env={**os.environ, "CODEX_HOME": self.home},
limit=_STREAM_LIMIT,
)
self._pending = {}
self._threads = {} # in-memory threads are gone after a respawn
self._reader = asyncio.create_task(self._read_loop())
await self._request("initialize", {
"clientInfo": {"name": "codex-as-api", "title": "Codex as API",
"version": "1.0.0"},
"capabilities": {"experimentalApi": True, "requestAttestation": False},
})
await self._send({"method": "initialized"})
async def _kill(self) -> None:
if self._reader and not self._reader.done():
self._reader.cancel()
if self.proc:
try:
self.proc.kill()
except Exception:
pass
self.proc = None
for f in self._pending.values():
if not f.done():
f.set_exception(CodexError("app-server restarted"))
self._pending = {}
async def _read_loop(self) -> None:
try:
while self.proc and self.proc.stdout:
line = await self.proc.stdout.readline()
if not line:
break
line = line.strip()
if not line:
continue
try:
msg = json.loads(line)
except json.JSONDecodeError:
continue
mid = msg.get("id")
if mid is not None and ("result" in msg or "error" in msg):
fut = self._pending.pop(mid, None)
if fut and not fut.done():
fut.set_result(msg)
elif msg.get("method"):
if self._queue is not None:
self._queue.put_nowait(msg)
except asyncio.CancelledError:
pass
except Exception:
pass
finally:
for f in self._pending.values():
if not f.done():
f.set_exception(CodexError("app-server closed"))
self._pending = {}
# -- io ------------------------------------------------------------------
async def _send(self, obj: dict) -> None:
if not self.proc or not self.proc.stdin:
raise CodexError("app-server not running")
self.proc.stdin.write((json.dumps(obj) + "\n").encode("utf-8"))
await self.proc.stdin.drain()
async def _request(self, method: str, params: dict) -> dict:
self._next_id += 1
rid = self._next_id
fut: asyncio.Future = asyncio.get_event_loop().create_future()
self._pending[rid] = fut
await self._send({"method": method, "id": rid, "params": params})
try:
msg = await asyncio.wait_for(fut, timeout=self.read_timeout)
except asyncio.TimeoutError:
raise CodexError(f"app-server timed out on {method}")
if "error" in msg:
raise CodexError(f"app-server error on {method}: {msg['error']}")
return msg.get("result", {})
# -- one turn ------------------------------------------------------------
async def run(self, *, prompt, workspace, thread_id, sandbox, model,
effort, input_items, output_schema, session_id,
developer_instructions=None) -> AsyncIterator[dict]:
async with self.lock:
try:
await self._ensure()
except (FileNotFoundError, OSError) as e:
raise CodexError(f"could not start codex app-server ('{self.bin}'): {e}")
self._queue = asyncio.Queue()
try:
tid = (self._threads.get(session_id) if session_id else None) or thread_id
resolved = None
if tid:
try:
rp = {"threadId": tid, "cwd": str(workspace),
"approvalPolicy": "never", "sandbox": sandbox,
"excludeTurns": True}
if developer_instructions:
rp["developerInstructions"] = developer_instructions
res = await self._request("thread/resume", rp)
resolved = (res.get("thread") or {}).get("id")
except CodexError:
resolved = None
if not resolved:
params = {"cwd": str(workspace), "approvalPolicy": "never",
"sandbox": sandbox}
if model:
params["model"] = model
if developer_instructions:
params["developerInstructions"] = developer_instructions
res = await self._request("thread/start", params)
resolved = (res.get("thread") or {}).get("id")
if not resolved:
raise CodexError("app-server did not return a thread id")
if session_id:
self._threads[session_id] = resolved
turn_input = input_items or [
{"type": "text", "text": prompt, "text_elements": []}]
tp = {"threadId": resolved, "input": turn_input}
if model:
tp["model"] = model
if effort:
tp["effort"] = effort
if output_schema:
tp["outputSchema"] = output_schema
await self._request("turn/start", tp) # returns quickly
delta_parts: list[str] = []
final_text: Optional[str] = None
images: list[str] = []
usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0,
"prompt_tokens_details": {"cached_tokens": 0}}
while True:
try:
msg = await asyncio.wait_for(self._queue.get(),
timeout=self.read_timeout)
except asyncio.TimeoutError:
raise CodexError("app-server timed out during the turn")
method = msg.get("method")
if method == "item/agentMessage/delta":
d = (msg.get("params") or {}).get("delta", "")
if d:
delta_parts.append(d)
yield {"type": "delta", "text": d}
elif method in ("item/reasoning/textDelta",
"item/reasoning/summaryTextDelta"):
rd = (msg.get("params") or {}).get("delta") \
or (msg.get("params") or {}).get("text", "")
if rd:
yield {"type": "reasoning", "text": rd}
elif method == "item/completed":
item = (msg.get("params") or {}).get("item", {})
it = item.get("type")
if it == "agentMessage" and item.get("text") is not None:
final_text = item["text"]
elif it == "imageGeneration":
p = item.get("savedPath") or item.get("result")
if p:
images.append(p)
elif method == "thread/tokenUsage/updated":
last = ((msg.get("params") or {}).get("tokenUsage") or {}).get("last", {})
usage = {
"prompt_tokens": last.get("inputTokens", 0) or 0,
"completion_tokens": last.get("outputTokens", 0) or 0,
"total_tokens": last.get("totalTokens", 0) or 0,
"prompt_tokens_details": {
"cached_tokens": last.get("cachedInputTokens", 0) or 0},
}
elif method == "error":
err = (msg.get("params") or {}).get("error", {}) or {}
blob = f"{err.get('message','')} {err.get('additionalDetails') or ''}".lower()
if any(s in blob for s in _AUTH_DEAD):
await self._kill() # respawn next time (re-read auth)
raise CodexError(
"Codex login expired (session ended). Re-upload a "
"fresh auth.json to /data/.codex/auth.json.")
elif method == "turn/completed":
break
elif msg.get("id") is not None and method is not None:
# server->client request (approval); decline to avoid hang
await self._send({"id": msg["id"],
"error": {"code": -32601,
"message": "approvals disabled"}})
text = final_text if final_text is not None else "".join(delta_parts)
yield {"type": "final", "text": text, "thread_id": resolved,
"usage": usage, "images": images}
except CodexError:
raise
except Exception as e:
await self._kill()
raise CodexError(f"warm pool error: {e}")
finally:
self._queue = None
_POOL: Optional[WarmCodex] = None
def _get_pool(codex_bin, codex_home, read_timeout) -> WarmCodex:
global _POOL
if _POOL is None:
_POOL = WarmCodex(codex_bin, codex_home, read_timeout)
return _POOL
async def run_turn_pool(*, codex_bin, codex_home, prompt, workspace, thread_id,
sandbox, model, read_timeout, effort=None,
input_items=None, output_schema=None,
session_id=None, developer_instructions=None) -> AsyncIterator[dict]:
pool = _get_pool(codex_bin, codex_home, read_timeout)
async for evt in pool.run(prompt=prompt, workspace=workspace, thread_id=thread_id,
sandbox=sandbox, model=model, effort=effort,
input_items=input_items, output_schema=output_schema,
session_id=session_id,
developer_instructions=developer_instructions):
yield evt
|