526 lines
17 KiB
Python
526 lines
17 KiB
Python
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")
|