banteg
576 lines
16 KiB
Python
576 lines
16 KiB
Python
import anyio
|
|
import pytest
|
|
|
|
from takopi.telegram.bridge import (
|
|
TelegramBridgeConfig,
|
|
TelegramTransport,
|
|
_build_bot_commands,
|
|
_handle_cancel,
|
|
_is_cancel_command,
|
|
_send_with_resume,
|
|
_strip_engine_command,
|
|
run_main_loop,
|
|
)
|
|
from takopi.runner_bridge import ExecBridgeConfig, RunningTask
|
|
from takopi.markdown import MarkdownPresenter
|
|
from takopi.model import EngineId, ResumeToken
|
|
from takopi.router import AutoRouter, RunnerEntry
|
|
from takopi.runners.mock import Return, ScriptRunner, Sleep, Wait
|
|
from takopi.transport import MessageRef, RenderedMessage, SendOptions
|
|
|
|
CODEX_ENGINE = EngineId("codex")
|
|
|
|
|
|
def _make_router(runner) -> AutoRouter:
|
|
return AutoRouter(
|
|
entries=[RunnerEntry(engine=runner.engine, runner=runner)],
|
|
default_engine=runner.engine,
|
|
)
|
|
|
|
|
|
class _FakeTransport:
|
|
def __init__(self, progress_ready: anyio.Event | None = None) -> None:
|
|
self._next_id = 1
|
|
self.send_calls: list[dict] = []
|
|
self.edit_calls: list[dict] = []
|
|
self.delete_calls: list[MessageRef] = []
|
|
self.progress_ready = progress_ready
|
|
self.progress_ref: MessageRef | None = None
|
|
|
|
async def send(
|
|
self,
|
|
*,
|
|
channel_id: int | str,
|
|
message: RenderedMessage,
|
|
options: SendOptions | None = None,
|
|
) -> MessageRef:
|
|
ref = MessageRef(channel_id=channel_id, message_id=self._next_id)
|
|
self._next_id += 1
|
|
self.send_calls.append(
|
|
{
|
|
"ref": ref,
|
|
"channel_id": channel_id,
|
|
"message": message,
|
|
"options": options,
|
|
}
|
|
)
|
|
if (
|
|
self.progress_ref is None
|
|
and options is not None
|
|
and options.reply_to is not None
|
|
and options.notify is False
|
|
):
|
|
self.progress_ref = ref
|
|
if self.progress_ready is not None:
|
|
self.progress_ready.set()
|
|
return ref
|
|
|
|
async def edit(
|
|
self, *, ref: MessageRef, message: RenderedMessage, wait: bool = True
|
|
) -> MessageRef:
|
|
self.edit_calls.append({"ref": ref, "message": message, "wait": wait})
|
|
return ref
|
|
|
|
async def delete(self, *, ref: MessageRef) -> bool:
|
|
self.delete_calls.append(ref)
|
|
return True
|
|
|
|
async def close(self) -> None:
|
|
return None
|
|
|
|
|
|
class _FakeBot:
|
|
def __init__(self) -> None:
|
|
self.command_calls: list[dict] = []
|
|
self.send_calls: list[dict] = []
|
|
self.edit_calls: list[dict] = []
|
|
self.delete_calls: list[dict] = []
|
|
|
|
async def get_updates(
|
|
self,
|
|
offset: int | None,
|
|
timeout_s: int = 50,
|
|
allowed_updates: list[str] | None = None,
|
|
) -> list[dict] | None:
|
|
_ = offset
|
|
_ = timeout_s
|
|
_ = allowed_updates
|
|
return []
|
|
|
|
async def send_message(
|
|
self,
|
|
chat_id: int,
|
|
text: str,
|
|
reply_to_message_id: int | None = None,
|
|
disable_notification: bool | None = False,
|
|
entities: list[dict] | None = None,
|
|
parse_mode: str | None = None,
|
|
*,
|
|
replace_message_id: int | None = None,
|
|
) -> dict:
|
|
self.send_calls.append(
|
|
{
|
|
"chat_id": chat_id,
|
|
"text": text,
|
|
"reply_to_message_id": reply_to_message_id,
|
|
"disable_notification": disable_notification,
|
|
"entities": entities,
|
|
"parse_mode": parse_mode,
|
|
"replace_message_id": replace_message_id,
|
|
}
|
|
)
|
|
return {"message_id": 1}
|
|
|
|
async def edit_message_text(
|
|
self,
|
|
chat_id: int,
|
|
message_id: int,
|
|
text: str,
|
|
entities: list[dict] | None = None,
|
|
parse_mode: str | None = None,
|
|
*,
|
|
wait: bool = True,
|
|
) -> dict:
|
|
self.edit_calls.append(
|
|
{
|
|
"chat_id": chat_id,
|
|
"message_id": message_id,
|
|
"text": text,
|
|
"entities": entities,
|
|
"parse_mode": parse_mode,
|
|
"wait": wait,
|
|
}
|
|
)
|
|
return {"message_id": message_id}
|
|
|
|
async def delete_message(self, chat_id: int, message_id: int) -> bool:
|
|
self.delete_calls.append({"chat_id": chat_id, "message_id": message_id})
|
|
return True
|
|
|
|
async def set_my_commands(
|
|
self,
|
|
commands: list[dict],
|
|
*,
|
|
scope: dict | None = None,
|
|
language_code: str | None = None,
|
|
) -> bool:
|
|
self.command_calls.append(
|
|
{
|
|
"commands": commands,
|
|
"scope": scope,
|
|
"language_code": language_code,
|
|
}
|
|
)
|
|
return True
|
|
|
|
async def get_me(self) -> dict | None:
|
|
return {"id": 1}
|
|
|
|
async def close(self) -> None:
|
|
return None
|
|
|
|
|
|
def _make_cfg(
|
|
transport: _FakeTransport, runner: ScriptRunner | None = None
|
|
) -> TelegramBridgeConfig:
|
|
if runner is None:
|
|
runner = ScriptRunner([Return(answer="ok")], engine=CODEX_ENGINE)
|
|
exec_cfg = ExecBridgeConfig(
|
|
transport=transport,
|
|
presenter=MarkdownPresenter(),
|
|
final_notify=True,
|
|
)
|
|
return TelegramBridgeConfig(
|
|
bot=_FakeBot(),
|
|
router=_make_router(runner),
|
|
chat_id=123,
|
|
startup_msg="",
|
|
exec_cfg=exec_cfg,
|
|
)
|
|
|
|
|
|
def test_strip_engine_command_inline() -> None:
|
|
text, engine = _strip_engine_command(
|
|
"/claude do it", engine_ids=("codex", "claude")
|
|
)
|
|
assert engine == "claude"
|
|
assert text == "do it"
|
|
|
|
|
|
def test_strip_engine_command_newline() -> None:
|
|
text, engine = _strip_engine_command(
|
|
"/codex\nhello", engine_ids=("codex", "claude")
|
|
)
|
|
assert engine == "codex"
|
|
assert text == "hello"
|
|
|
|
|
|
def test_strip_engine_command_ignores_unknown() -> None:
|
|
text, engine = _strip_engine_command("/unknown hi", engine_ids=("codex", "claude"))
|
|
assert engine is None
|
|
assert text == "/unknown hi"
|
|
|
|
|
|
def test_strip_engine_command_bot_suffix() -> None:
|
|
text, engine = _strip_engine_command(
|
|
"/claude@bunny_agent_bot hi", engine_ids=("claude",)
|
|
)
|
|
assert engine == "claude"
|
|
assert text == "hi"
|
|
|
|
|
|
def test_strip_engine_command_only_first_non_empty_line() -> None:
|
|
text, engine = _strip_engine_command(
|
|
"hello\n/claude hi", engine_ids=("codex", "claude")
|
|
)
|
|
assert engine is None
|
|
assert text == "hello\n/claude hi"
|
|
|
|
|
|
def test_build_bot_commands_includes_cancel_and_engine() -> None:
|
|
runner = ScriptRunner(
|
|
[Return(answer="ok")], engine=CODEX_ENGINE, resume_value="sid"
|
|
)
|
|
router = _make_router(runner)
|
|
commands = _build_bot_commands(router)
|
|
|
|
assert {"command": "cancel", "description": "cancel run"} in commands
|
|
assert any(cmd["command"] == "codex" for cmd in commands)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_telegram_transport_passes_replace_and_wait() -> None:
|
|
bot = _FakeBot()
|
|
transport = TelegramTransport(bot)
|
|
reply = MessageRef(channel_id=123, message_id=10)
|
|
replace = MessageRef(channel_id=123, message_id=11)
|
|
|
|
await transport.send(
|
|
channel_id=123,
|
|
message=RenderedMessage(text="hello"),
|
|
options=SendOptions(reply_to=reply, notify=True, replace=replace),
|
|
)
|
|
assert bot.send_calls
|
|
assert bot.send_calls[0]["replace_message_id"] == 11
|
|
|
|
await transport.edit(
|
|
ref=replace,
|
|
message=RenderedMessage(text="edit"),
|
|
wait=False,
|
|
)
|
|
assert bot.edit_calls
|
|
assert bot.edit_calls[0]["wait"] is False
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_telegram_transport_edit_wait_false_returns_ref() -> None:
|
|
class _OutboxBot:
|
|
def __init__(self) -> None:
|
|
self.edit_calls: list[dict[str, object]] = []
|
|
|
|
async def get_updates(
|
|
self,
|
|
offset: int | None,
|
|
timeout_s: int = 50,
|
|
allowed_updates: list[str] | None = None,
|
|
) -> list[dict] | None:
|
|
return None
|
|
|
|
async def send_message(
|
|
self,
|
|
chat_id: int,
|
|
text: str,
|
|
reply_to_message_id: int | None = None,
|
|
disable_notification: bool | None = False,
|
|
entities: list[dict] | None = None,
|
|
parse_mode: str | None = None,
|
|
*,
|
|
replace_message_id: int | None = None,
|
|
) -> dict | None:
|
|
return None
|
|
|
|
async def edit_message_text(
|
|
self,
|
|
chat_id: int,
|
|
message_id: int,
|
|
text: str,
|
|
entities: list[dict] | None = None,
|
|
parse_mode: str | None = None,
|
|
*,
|
|
wait: bool = True,
|
|
) -> dict | None:
|
|
self.edit_calls.append(
|
|
{
|
|
"chat_id": chat_id,
|
|
"message_id": message_id,
|
|
"text": text,
|
|
"entities": entities,
|
|
"parse_mode": parse_mode,
|
|
"wait": wait,
|
|
}
|
|
)
|
|
if not wait:
|
|
return None
|
|
return {"message_id": message_id}
|
|
|
|
async def delete_message(
|
|
self,
|
|
chat_id: int,
|
|
message_id: int,
|
|
) -> bool:
|
|
return False
|
|
|
|
async def set_my_commands(
|
|
self,
|
|
commands: list[dict[str, object]],
|
|
*,
|
|
scope: dict[str, object] | None = None,
|
|
language_code: str | None = None,
|
|
) -> bool:
|
|
return False
|
|
|
|
async def get_me(self) -> dict | None:
|
|
return None
|
|
|
|
async def close(self) -> None:
|
|
return None
|
|
|
|
bot = _OutboxBot()
|
|
transport = TelegramTransport(bot)
|
|
ref = MessageRef(channel_id=123, message_id=1)
|
|
|
|
result = await transport.edit(
|
|
ref=ref,
|
|
message=RenderedMessage(text="edit"),
|
|
wait=False,
|
|
)
|
|
|
|
assert result == ref
|
|
assert bot.edit_calls
|
|
assert bot.edit_calls[0]["wait"] is False
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_cancel_without_reply_prompts_user() -> None:
|
|
transport = _FakeTransport()
|
|
cfg = _make_cfg(transport)
|
|
msg = {"chat": {"id": 123}, "message_id": 10}
|
|
running_tasks: dict = {}
|
|
|
|
await _handle_cancel(cfg, msg, running_tasks)
|
|
|
|
assert len(transport.send_calls) == 1
|
|
assert "reply to the progress message" in transport.send_calls[0]["message"].text
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_cancel_with_no_progress_message_says_nothing_running() -> None:
|
|
transport = _FakeTransport()
|
|
cfg = _make_cfg(transport)
|
|
msg = {
|
|
"chat": {"id": 123},
|
|
"message_id": 10,
|
|
"reply_to_message": {"text": "no message id"},
|
|
}
|
|
running_tasks: dict = {}
|
|
|
|
await _handle_cancel(cfg, msg, running_tasks)
|
|
|
|
assert len(transport.send_calls) == 1
|
|
assert "nothing is currently running" in transport.send_calls[0]["message"].text
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_cancel_with_finished_task_says_nothing_running() -> None:
|
|
transport = _FakeTransport()
|
|
cfg = _make_cfg(transport)
|
|
progress_id = 99
|
|
msg = {
|
|
"chat": {"id": 123},
|
|
"message_id": 10,
|
|
"reply_to_message": {"message_id": progress_id},
|
|
}
|
|
running_tasks: dict = {}
|
|
|
|
await _handle_cancel(cfg, msg, running_tasks)
|
|
|
|
assert len(transport.send_calls) == 1
|
|
assert "nothing is currently running" in transport.send_calls[0]["message"].text
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_cancel_cancels_running_task() -> None:
|
|
transport = _FakeTransport()
|
|
cfg = _make_cfg(transport)
|
|
progress_id = 42
|
|
msg = {
|
|
"chat": {"id": 123},
|
|
"message_id": 10,
|
|
"reply_to_message": {"message_id": progress_id},
|
|
}
|
|
|
|
running_task = RunningTask()
|
|
running_tasks = {MessageRef(channel_id=123, message_id=progress_id): running_task}
|
|
await _handle_cancel(cfg, msg, running_tasks)
|
|
|
|
assert running_task.cancel_requested.is_set() is True
|
|
assert len(transport.send_calls) == 0 # No error message sent
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_cancel_only_cancels_matching_progress_message() -> None:
|
|
transport = _FakeTransport()
|
|
cfg = _make_cfg(transport)
|
|
task_first = RunningTask()
|
|
task_second = RunningTask()
|
|
msg = {
|
|
"chat": {"id": 123},
|
|
"message_id": 10,
|
|
"reply_to_message": {"message_id": 1},
|
|
}
|
|
running_tasks = {
|
|
MessageRef(channel_id=123, message_id=1): task_first,
|
|
MessageRef(channel_id=123, message_id=2): task_second,
|
|
}
|
|
|
|
await _handle_cancel(cfg, msg, running_tasks)
|
|
|
|
assert task_first.cancel_requested.is_set() is True
|
|
assert task_second.cancel_requested.is_set() is False
|
|
assert len(transport.send_calls) == 0
|
|
|
|
|
|
def test_cancel_command_accepts_extra_text() -> None:
|
|
assert _is_cancel_command("/cancel now") is True
|
|
assert _is_cancel_command("/cancel@takopi please") is True
|
|
assert _is_cancel_command("/cancelled") is False
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_send_with_resume_waits_for_token() -> None:
|
|
transport = _FakeTransport()
|
|
cfg = _make_cfg(transport)
|
|
sent: list[tuple[int, int, str, ResumeToken | None]] = []
|
|
|
|
async def enqueue(
|
|
chat_id: int, user_msg_id: int, text: str, resume: ResumeToken
|
|
) -> None:
|
|
sent.append((chat_id, user_msg_id, text, resume))
|
|
|
|
running_task = RunningTask()
|
|
|
|
async def trigger_resume() -> None:
|
|
await anyio.sleep(0)
|
|
running_task.resume = ResumeToken(engine=CODEX_ENGINE, value="abc123")
|
|
running_task.resume_ready.set()
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(trigger_resume)
|
|
await _send_with_resume(
|
|
cfg,
|
|
enqueue,
|
|
running_task,
|
|
123,
|
|
10,
|
|
"hello",
|
|
)
|
|
|
|
assert sent == [
|
|
(123, 10, "hello", ResumeToken(engine=CODEX_ENGINE, value="abc123"))
|
|
]
|
|
assert transport.send_calls == []
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_send_with_resume_reports_when_missing() -> None:
|
|
transport = _FakeTransport()
|
|
cfg = _make_cfg(transport)
|
|
sent: list[tuple[int, int, str, ResumeToken | None]] = []
|
|
|
|
async def enqueue(
|
|
chat_id: int, user_msg_id: int, text: str, resume: ResumeToken
|
|
) -> None:
|
|
sent.append((chat_id, user_msg_id, text, resume))
|
|
|
|
running_task = RunningTask()
|
|
running_task.done.set()
|
|
|
|
await _send_with_resume(
|
|
cfg,
|
|
enqueue,
|
|
running_task,
|
|
123,
|
|
10,
|
|
"hello",
|
|
)
|
|
|
|
assert sent == []
|
|
assert transport.send_calls
|
|
assert "resume token" in transport.send_calls[-1]["message"].text.lower()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_run_main_loop_routes_reply_to_running_resume() -> None:
|
|
progress_ready = anyio.Event()
|
|
stop_polling = anyio.Event()
|
|
reply_ready = anyio.Event()
|
|
hold = anyio.Event()
|
|
|
|
transport = _FakeTransport(progress_ready=progress_ready)
|
|
bot = _FakeBot()
|
|
resume_value = "abc123"
|
|
runner = ScriptRunner(
|
|
[Wait(hold), Sleep(0.05), Return(answer="ok")],
|
|
engine=CODEX_ENGINE,
|
|
resume_value=resume_value,
|
|
)
|
|
exec_cfg = ExecBridgeConfig(
|
|
transport=transport,
|
|
presenter=MarkdownPresenter(),
|
|
final_notify=True,
|
|
)
|
|
cfg = TelegramBridgeConfig(
|
|
bot=bot,
|
|
router=_make_router(runner),
|
|
chat_id=123,
|
|
startup_msg="",
|
|
exec_cfg=exec_cfg,
|
|
)
|
|
|
|
async def poller(_cfg: TelegramBridgeConfig):
|
|
yield {
|
|
"message_id": 1,
|
|
"text": "first",
|
|
"chat": {"id": 123},
|
|
"from": {"id": 123},
|
|
}
|
|
await progress_ready.wait()
|
|
assert transport.progress_ref is not None
|
|
reply_ready.set()
|
|
yield {
|
|
"message_id": 2,
|
|
"text": "followup",
|
|
"chat": {"id": 123},
|
|
"from": {"id": 123},
|
|
"reply_to_message": {"message_id": transport.progress_ref.message_id},
|
|
}
|
|
await stop_polling.wait()
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(run_main_loop, cfg, poller)
|
|
try:
|
|
with anyio.fail_after(2):
|
|
await reply_ready.wait()
|
|
await anyio.sleep(0)
|
|
hold.set()
|
|
with anyio.fail_after(2):
|
|
while len(runner.calls) < 2:
|
|
await anyio.sleep(0)
|
|
assert runner.calls[1][1] == ResumeToken(
|
|
engine=CODEX_ENGINE, value=resume_value
|
|
)
|
|
finally:
|
|
hold.set()
|
|
stop_polling.set()
|
|
tg.cancel_scope.cancel()
|