Initial commit
This commit is contained in:
525
app/app_server_runner.py
Normal file
525
app/app_server_runner.py
Normal file
@@ -0,0 +1,525 @@
|
||||
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")
|
||||
Reference in New Issue
Block a user