feat: add codex telegram bridge scripts
This commit is contained in:
banteg
@@ -0,0 +1,195 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from bridge_common import TelegramClient, RouteStore, parse_allowed_chat_ids
|
||||
|
||||
# -------------------- Codex runner --------------------
|
||||
|
||||
|
||||
class CodexExecRunner:
|
||||
"""
|
||||
Runs Codex in non-interactive mode:
|
||||
- new: codex exec --json ... -
|
||||
- resume: codex exec --json ... resume <SESSION_ID> -
|
||||
"""
|
||||
|
||||
def __init__(self, codex_cmd: str, workspace: Optional[str], extra_args: list[str]) -> None:
|
||||
self.codex_cmd = codex_cmd
|
||||
self.workspace = workspace
|
||||
self.extra_args = extra_args
|
||||
|
||||
# per-session locks to prevent concurrent resumes to same session_id
|
||||
self._locks: dict[str, threading.Lock] = {}
|
||||
self._locks_guard = threading.Lock()
|
||||
|
||||
def _lock_for(self, session_id: str) -> threading.Lock:
|
||||
with self._locks_guard:
|
||||
if session_id not in self._locks:
|
||||
self._locks[session_id] = threading.Lock()
|
||||
return self._locks[session_id]
|
||||
|
||||
def run(self, prompt: str, session_id: Optional[str]) -> Tuple[str, str]:
|
||||
"""
|
||||
Returns (session_id, final_agent_message_text)
|
||||
"""
|
||||
args = [self.codex_cmd, "exec", "--json"]
|
||||
args.extend(self.extra_args)
|
||||
if self.workspace:
|
||||
args.extend(["--cd", self.workspace])
|
||||
|
||||
# Always pipe prompt via stdin ("-") to avoid quoting issues.
|
||||
if session_id:
|
||||
args.extend(["resume", session_id, "-"])
|
||||
else:
|
||||
args.append("-")
|
||||
|
||||
# read both stdout+stderr without deadlock
|
||||
proc = subprocess.Popen(
|
||||
args,
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
assert proc.stdin and proc.stdout and proc.stderr
|
||||
|
||||
# send prompt then close stdin
|
||||
proc.stdin.write(prompt)
|
||||
proc.stdin.close()
|
||||
|
||||
stderr_lines: list[str] = []
|
||||
|
||||
def _drain_stderr() -> None:
|
||||
for line in proc.stderr:
|
||||
stderr_lines.append(line)
|
||||
|
||||
t = threading.Thread(target=_drain_stderr, daemon=True)
|
||||
t.start()
|
||||
|
||||
found_session: Optional[str] = session_id
|
||||
last_agent_text: Optional[str] = None
|
||||
|
||||
for line in proc.stdout:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
evt = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# From Codex JSONL event stream
|
||||
if evt.get("type") == "thread.started":
|
||||
found_session = evt.get("thread_id") or found_session
|
||||
|
||||
if evt.get("type") == "item.completed":
|
||||
item = evt.get("item") or {}
|
||||
if item.get("type") == "agent_message" and isinstance(item.get("text"), str):
|
||||
last_agent_text = item["text"]
|
||||
|
||||
rc = proc.wait()
|
||||
t.join(timeout=2.0)
|
||||
|
||||
if rc != 0:
|
||||
tail = "".join(stderr_lines[-200:])
|
||||
raise RuntimeError(f"codex exec failed (rc={rc}). stderr tail:\n{tail}")
|
||||
|
||||
if not found_session:
|
||||
raise RuntimeError("codex exec finished but no session_id/thread_id was captured")
|
||||
|
||||
return found_session, (last_agent_text or "(No agent_message captured from JSON stream.)")
|
||||
|
||||
def run_serialized(self, prompt: str, session_id: Optional[str]) -> Tuple[str, str]:
|
||||
"""
|
||||
If resuming, serialize per-session.
|
||||
"""
|
||||
if not session_id:
|
||||
return self.run(prompt, session_id=None)
|
||||
lock = self._lock_for(session_id)
|
||||
with lock:
|
||||
return self.run(prompt, session_id=session_id)
|
||||
|
||||
|
||||
# -------------------- Telegram loop --------------------
|
||||
|
||||
|
||||
def main() -> None:
|
||||
token = os.environ.get("TELEGRAM_BOT_TOKEN", "")
|
||||
db_path = os.environ.get("BRIDGE_DB", "./bridge_routes.sqlite3")
|
||||
allowed = parse_allowed_chat_ids(os.environ.get("ALLOWED_CHAT_IDS", ""))
|
||||
|
||||
codex_cmd = os.environ.get("CODEX_CMD", "codex")
|
||||
workspace = os.environ.get("CODEX_WORKSPACE") # optional
|
||||
extra_args = shlex.split(os.environ.get("CODEX_EXEC_ARGS", "")) # e.g. "--full-auto --search"
|
||||
|
||||
bot = TelegramClient(token)
|
||||
store = RouteStore(db_path)
|
||||
runner = CodexExecRunner(codex_cmd=codex_cmd, workspace=workspace, extra_args=extra_args)
|
||||
|
||||
pool = ThreadPoolExecutor(max_workers=int(os.environ.get("MAX_WORKERS", "4")))
|
||||
offset: Optional[int] = None
|
||||
|
||||
print("Option1 bridge running (codex exec). Long-polling Telegram...")
|
||||
|
||||
def handle(chat_id: int, user_msg_id: int, text: str, resume_session: Optional[str]) -> None:
|
||||
try:
|
||||
session_id, answer = runner.run_serialized(text, resume_session)
|
||||
sent_msgs = bot.send_message_chunked(
|
||||
chat_id=chat_id,
|
||||
text=answer,
|
||||
reply_to_message_id=user_msg_id,
|
||||
)
|
||||
for m in sent_msgs:
|
||||
store.link(chat_id, m["message_id"], "exec", session_id, meta={"workspace": workspace})
|
||||
except Exception as e:
|
||||
err = f"❌ Error:\n{e}"
|
||||
sent_msgs = bot.send_message_chunked(chat_id=chat_id, text=err, reply_to_message_id=user_msg_id)
|
||||
for m in sent_msgs:
|
||||
store.link(chat_id, m["message_id"], "exec", resume_session or "unknown", meta={"error": True})
|
||||
|
||||
while True:
|
||||
try:
|
||||
updates = bot.get_updates(offset=offset, timeout_s=50, allowed_updates=["message"])
|
||||
except Exception as e:
|
||||
print(f"[telegram] get_updates error: {e}")
|
||||
time.sleep(2.0)
|
||||
continue
|
||||
|
||||
for upd in updates:
|
||||
offset = upd["update_id"] + 1
|
||||
msg = upd.get("message") or {}
|
||||
if "text" not in msg:
|
||||
continue
|
||||
|
||||
chat_id = msg["chat"]["id"]
|
||||
if allowed is not None and int(chat_id) not in allowed:
|
||||
continue
|
||||
|
||||
if msg.get("from", {}).get("is_bot"):
|
||||
continue
|
||||
|
||||
text = msg["text"]
|
||||
user_msg_id = msg["message_id"]
|
||||
|
||||
# If user replied to a bot message, route to that session
|
||||
resume_session: Optional[str] = None
|
||||
r = msg.get("reply_to_message")
|
||||
if r and "message_id" in r:
|
||||
route = store.resolve(chat_id, r["message_id"])
|
||||
if route and route.route_type == "exec":
|
||||
resume_session = route.route_id
|
||||
|
||||
pool.submit(handle, chat_id, user_msg_id, text, resume_session)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user