banteg
532 lines
19 KiB
Python
532 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
# /// script
|
|
# requires-python = ">=3.10"
|
|
# dependencies = ["markdown-it-py", "sulguk", "tomli; python_version < '3.11'"]
|
|
# ///
|
|
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, Callable, Dict, Optional, Tuple
|
|
|
|
from bridge_common import (
|
|
TelegramClient,
|
|
RouteStore,
|
|
TELEGRAM_HARD_LIMIT,
|
|
config_get,
|
|
load_telegram_config,
|
|
render_markdown,
|
|
resolve_chat_ids,
|
|
)
|
|
|
|
# -------------------- Codex runner --------------------
|
|
|
|
|
|
def log(msg: str) -> None:
|
|
ts = time.strftime("%Y-%m-%d %H:%M:%S")
|
|
print(f"[{ts}] {msg}", flush=True)
|
|
|
|
|
|
def _one_line(text: Optional[str]) -> str:
|
|
if text is None:
|
|
return "None"
|
|
return text.replace("\r", "\\r").replace("\n", "\\n")
|
|
|
|
|
|
TELEGRAM_TEXT_LIMIT = TELEGRAM_HARD_LIMIT
|
|
|
|
|
|
def _clamp_tg_text(text: str, limit: int = TELEGRAM_TEXT_LIMIT) -> str:
|
|
if len(text) <= limit:
|
|
return text
|
|
return text[: limit - 20] + "\n...(truncated)"
|
|
|
|
|
|
def _summarize_item(item: Dict[str, Any]) -> str:
|
|
item_type = item.get("type")
|
|
if isinstance(item_type, str):
|
|
if item_type == "agent_message" and isinstance(item.get("text"), str):
|
|
snippet = item["text"].strip().replace("\n", " ")
|
|
if len(snippet) > 140:
|
|
snippet = snippet[:140] + "..."
|
|
return f"agent_message: {snippet}"
|
|
name = item.get("name") or item.get("tool_name") or item.get("id")
|
|
if isinstance(name, str):
|
|
return f"{item_type}: {name}"
|
|
return item_type
|
|
return "item.completed"
|
|
|
|
|
|
class ProgressEditor:
|
|
def __init__(
|
|
self,
|
|
bot: TelegramClient,
|
|
chat_id: int,
|
|
message_id: int,
|
|
edit_every_s: float,
|
|
) -> None:
|
|
self.bot = bot
|
|
self.chat_id = chat_id
|
|
self.message_id = message_id
|
|
self.edit_every_s = edit_every_s
|
|
|
|
self._lock = threading.Lock()
|
|
self._pending: Optional[str] = None
|
|
self._last_sent: Optional[str] = None
|
|
self._last_edit_at = 0.0
|
|
|
|
self._stop = threading.Event()
|
|
self._thread = threading.Thread(target=self._run, daemon=True)
|
|
self._thread.start()
|
|
|
|
def set(self, text: str) -> None:
|
|
text = _clamp_tg_text(text)
|
|
with self._lock:
|
|
self._pending = text
|
|
|
|
def stop(self) -> None:
|
|
self._stop.set()
|
|
self._thread.join(timeout=1.0)
|
|
|
|
def _edit(self, text: str) -> None:
|
|
try:
|
|
self.bot.edit_message_text(
|
|
chat_id=self.chat_id,
|
|
message_id=self.message_id,
|
|
text=text,
|
|
)
|
|
except Exception as e:
|
|
log(
|
|
"[progress] edit failed "
|
|
f"chat_id={self.chat_id} message_id={self.message_id}: {e}"
|
|
)
|
|
|
|
def _run(self) -> None:
|
|
while not self._stop.is_set():
|
|
to_send: Optional[str] = None
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
if self._pending is not None and (now - self._last_edit_at) >= self.edit_every_s:
|
|
if self._pending != self._last_sent:
|
|
to_send = self._pending
|
|
self._last_sent = self._pending
|
|
self._last_edit_at = now
|
|
self._pending = None
|
|
|
|
if to_send is not None:
|
|
self._edit(to_send)
|
|
|
|
self._stop.wait(0.2)
|
|
|
|
|
|
def _typing_loop(bot: TelegramClient, chat_id: int, stop_evt: threading.Event) -> None:
|
|
while not stop_evt.is_set():
|
|
try:
|
|
bot.send_chat_action(chat_id=chat_id, action="typing")
|
|
except Exception as e:
|
|
log(f"[typing] send_chat_action failed chat_id={chat_id}: {e}")
|
|
stop_evt.wait(4.0)
|
|
|
|
|
|
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],
|
|
on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
|
|
) -> Tuple[str, str]:
|
|
"""
|
|
Returns (session_id, final_agent_message_text)
|
|
"""
|
|
log(f"[codex] start run session_id={session_id!r} workspace={self.workspace!r}")
|
|
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:
|
|
log(f"[codex][stderr] {line.rstrip()}")
|
|
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
|
|
log(f"[codex][event] {line}")
|
|
try:
|
|
evt = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if on_event is not None:
|
|
try:
|
|
on_event(evt)
|
|
except Exception as e:
|
|
log(f"[codex][on_event] callback error: {e}")
|
|
|
|
# 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")
|
|
|
|
log(f"[codex] done run session_id={found_session!r}")
|
|
return found_session, (last_agent_text or "(No agent_message captured from JSON stream.)")
|
|
|
|
def run_serialized(
|
|
self,
|
|
prompt: str,
|
|
session_id: Optional[str],
|
|
on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
|
|
) -> Tuple[str, str]:
|
|
"""
|
|
If resuming, serialize per-session.
|
|
"""
|
|
if not session_id:
|
|
return self.run(prompt, session_id=None, on_event=on_event)
|
|
lock = self._lock_for(session_id)
|
|
with lock:
|
|
return self.run(prompt, session_id=session_id, on_event=on_event)
|
|
|
|
|
|
# -------------------- Telegram loop --------------------
|
|
|
|
|
|
def main() -> None:
|
|
config = load_telegram_config()
|
|
token = config_get(config, "bot_token") or ""
|
|
db_path = config_get(config, "bridge_db") or "./bridge_routes.sqlite3"
|
|
chat_ids = resolve_chat_ids(config)
|
|
allowed = chat_ids
|
|
startup_ids = chat_ids
|
|
startup_msg = config_get(config, "startup_message") or "✅ exec_bridge started (codex exec)."
|
|
startup_pwd = os.getcwd()
|
|
startup_msg = f"{startup_msg}\nPWD: {startup_pwd}"
|
|
|
|
codex_cmd = config_get(config, "codex_cmd") or "codex"
|
|
workspace = config_get(config, "codex_workspace")
|
|
raw_exec_args = config_get(config, "codex_exec_args") or ""
|
|
if isinstance(raw_exec_args, list):
|
|
extra_args = [str(v) for v in raw_exec_args]
|
|
else:
|
|
extra_args = shlex.split(str(raw_exec_args)) # e.g. "--full-auto --search"
|
|
|
|
def _has_notify_override(args: list[str]) -> bool:
|
|
for i, arg in enumerate(args):
|
|
if arg in ("-c", "--config"):
|
|
if i + 1 >= len(args):
|
|
continue
|
|
key = args[i + 1].split("=", 1)[0].strip()
|
|
if key == "notify" or key.endswith(".notify"):
|
|
return True
|
|
elif arg.startswith(("--config=", "-c=")):
|
|
key = arg.split("=", 1)[1].split("=", 1)[0].strip()
|
|
if key == "notify" or key.endswith(".notify"):
|
|
return True
|
|
return False
|
|
|
|
# Default: disable notify hook for exec-bridge runs to avoid duplicate messages.
|
|
if not _has_notify_override(extra_args):
|
|
extra_args.extend(["-c", "notify=[]"])
|
|
|
|
bot = TelegramClient(token)
|
|
store = RouteStore(db_path)
|
|
runner = CodexExecRunner(codex_cmd=codex_cmd, workspace=workspace, extra_args=extra_args)
|
|
|
|
max_workers = config_get(config, "max_workers")
|
|
if isinstance(max_workers, str):
|
|
max_workers = int(max_workers) if max_workers.strip() else None
|
|
elif not isinstance(max_workers, int):
|
|
max_workers = None
|
|
pool = ThreadPoolExecutor(max_workers=max_workers or 4)
|
|
offset: Optional[int] = None
|
|
|
|
log(f"[startup] pwd={startup_pwd}")
|
|
log("Option1 bridge running (codex exec). Long-polling Telegram...")
|
|
if startup_ids:
|
|
for chat_id in startup_ids:
|
|
try:
|
|
bot.send_message(chat_id=chat_id, text=startup_msg)
|
|
log(f"[startup] sent startup message to chat_id={chat_id}")
|
|
except Exception as e:
|
|
log(f"[startup] failed to send startup message to chat_id={chat_id}: {e}")
|
|
else:
|
|
log("[startup] no chat_id configured; skipping startup message")
|
|
|
|
def handle(chat_id: int, user_msg_id: int, text: str, resume_session: Optional[str]) -> None:
|
|
log(
|
|
"[handle] start "
|
|
f"chat_id={chat_id} user_msg_id={user_msg_id} resume_session={resume_session!r}"
|
|
)
|
|
try:
|
|
edit_every_s = float(os.environ.get("TG_PROGRESS_EDIT_EVERY_S", "2.5"))
|
|
except ValueError:
|
|
edit_every_s = 2.5
|
|
silent_progress = os.environ.get("TG_PROGRESS_SILENT", "1") == "1"
|
|
loud_final = os.environ.get("TG_FINAL_NOTIFY", "1") == "1"
|
|
|
|
typing_stop = threading.Event()
|
|
typing_thread = threading.Thread(
|
|
target=_typing_loop,
|
|
args=(bot, chat_id, typing_stop),
|
|
daemon=True,
|
|
)
|
|
typing_thread.start()
|
|
|
|
progress_id: Optional[int] = None
|
|
progress: Optional[ProgressEditor] = None
|
|
try:
|
|
progress_msg = bot.send_message(
|
|
chat_id=chat_id,
|
|
text="Working...",
|
|
reply_to_message_id=user_msg_id,
|
|
disable_notification=silent_progress,
|
|
)
|
|
progress_id = int(progress_msg["message_id"])
|
|
except Exception as e:
|
|
log(f"[handle] failed to send progress message chat_id={chat_id}: {e}")
|
|
|
|
if progress_id is not None:
|
|
progress = ProgressEditor(bot, chat_id, progress_id, edit_every_s)
|
|
|
|
started_at = time.monotonic()
|
|
session_box: dict[str, Optional[str]] = {"id": resume_session}
|
|
|
|
def on_event(evt: Dict[str, Any]) -> None:
|
|
event_type = evt.get("type")
|
|
if event_type == "thread.started":
|
|
thread_id = evt.get("thread_id")
|
|
if isinstance(thread_id, str) and thread_id:
|
|
session_box["id"] = thread_id
|
|
if progress_id is not None:
|
|
store.link(
|
|
chat_id,
|
|
progress_id,
|
|
"exec",
|
|
thread_id,
|
|
meta={"workspace": workspace},
|
|
)
|
|
elif event_type == "item.completed":
|
|
item = evt.get("item") or {}
|
|
elapsed = int(time.monotonic() - started_at)
|
|
line = _summarize_item(item) if isinstance(item, dict) else "item.completed"
|
|
session_id = session_box["id"]
|
|
header = f"Working... ({elapsed}s)"
|
|
if session_id:
|
|
header += f"\nSession: {session_id}"
|
|
msg = f"{header}\n\n{line}"
|
|
if progress is not None:
|
|
progress.set(msg)
|
|
|
|
def _stop_background() -> None:
|
|
typing_stop.set()
|
|
typing_thread.join(timeout=1.0)
|
|
if progress is not None:
|
|
progress.stop()
|
|
|
|
try:
|
|
session_id, answer = runner.run_serialized(text, resume_session, on_event=on_event)
|
|
except Exception as e:
|
|
_stop_background()
|
|
err = _clamp_tg_text(f"Error:\n{e}")
|
|
route_id = session_box["id"] or resume_session or "unknown"
|
|
if progress_id is not None and len(err) <= TELEGRAM_TEXT_LIMIT:
|
|
try:
|
|
bot.edit_message_text(chat_id=chat_id, message_id=progress_id, text=err)
|
|
store.link(chat_id, progress_id, "exec", route_id, meta={"error": True})
|
|
log(
|
|
"[handle] error "
|
|
f"chat_id={chat_id} user_msg_id={user_msg_id} "
|
|
f"resume_session={resume_session!r} err={e}"
|
|
)
|
|
return
|
|
except Exception as ee:
|
|
log(f"[handle] failed to edit progress into error: {ee}")
|
|
|
|
sent_msgs = bot.send_message_markdown_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", route_id, meta={"error": True})
|
|
log(
|
|
"[handle] error "
|
|
f"chat_id={chat_id} user_msg_id={user_msg_id} resume_session={resume_session!r} err={e}"
|
|
)
|
|
return
|
|
|
|
_stop_background()
|
|
|
|
answer = answer or "(No agent_message captured from JSON stream.)"
|
|
final_text, final_entities = render_markdown(answer)
|
|
can_edit_final = progress_id is not None and len(final_text) <= TELEGRAM_TEXT_LIMIT
|
|
|
|
if loud_final or not can_edit_final:
|
|
if progress_id is not None:
|
|
try:
|
|
bot.delete_message(chat_id=chat_id, message_id=progress_id)
|
|
except Exception as e:
|
|
log(f"[handle] delete progress failed chat_id={chat_id} message_id={progress_id}: {e}")
|
|
|
|
sent_msgs = bot.send_message_markdown_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})
|
|
else:
|
|
bot.edit_message_text(
|
|
chat_id=chat_id,
|
|
message_id=progress_id,
|
|
text=final_text,
|
|
entities=final_entities or None,
|
|
)
|
|
store.link(chat_id, progress_id, "exec", session_id, meta={"workspace": workspace})
|
|
|
|
log(
|
|
"[handle] done "
|
|
f"chat_id={chat_id} user_msg_id={user_msg_id} session_id={session_id!r}"
|
|
)
|
|
|
|
while True:
|
|
try:
|
|
updates = bot.get_updates(offset=offset, timeout_s=50, allowed_updates=["message"])
|
|
except Exception as e:
|
|
log(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 {}
|
|
chat_id = msg.get("chat", {}).get("id")
|
|
from_bot = msg.get("from", {}).get("is_bot")
|
|
msg_text = msg.get("text")
|
|
reply_to = (msg.get("reply_to_message") or {}).get("message_id")
|
|
log(
|
|
"[telegram] received "
|
|
f"update_id={upd.get('update_id')} chat_id={chat_id} "
|
|
f"from_bot={from_bot} has_text={msg_text is not None} "
|
|
f"reply_to={reply_to} text={_one_line(msg_text)}"
|
|
)
|
|
if "text" not in msg:
|
|
log(
|
|
"[telegram] ignoring non-text message "
|
|
f"chat_id={chat_id} update_id={upd.get('update_id')}"
|
|
)
|
|
continue
|
|
|
|
if allowed is not None and int(chat_id) not in allowed:
|
|
log(
|
|
"[telegram] rejected by ACL "
|
|
f"chat_id={chat_id} allowed={sorted(allowed)}"
|
|
)
|
|
continue
|
|
|
|
if msg.get("from", {}).get("is_bot"):
|
|
log(
|
|
"[telegram] ignoring bot message "
|
|
f"chat_id={chat_id} update_id={upd.get('update_id')}"
|
|
)
|
|
continue
|
|
|
|
text = msg["text"]
|
|
user_msg_id = msg["message_id"]
|
|
log(
|
|
"[telegram] accepted message "
|
|
f"chat_id={chat_id} user_msg_id={user_msg_id} text={_one_line(text)}"
|
|
)
|
|
|
|
# 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
|
|
log(
|
|
"[telegram] resolved reply route "
|
|
f"chat_id={chat_id} bot_message_id={r['message_id']} session_id={resume_session!r}"
|
|
)
|
|
else:
|
|
log(
|
|
"[telegram] reply has no exec route "
|
|
f"chat_id={chat_id} bot_message_id={r['message_id']}"
|
|
)
|
|
|
|
pool.submit(handle, chat_id, user_msg_id, text, resume_session)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|