izackp/takopi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: queue telegram requests with rate limits (#54)

Browse Source
This commit is contained in:
banteg
2026-01-05 12:00:37 +04:00
committed by GitHub
parent c64913ed6d
commit 2d8fbc8a5a
9 changed files with 898 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_exec_bridge.py
+54 -40
View File
@@ -8,6 +8,7 @@ from takopi.model import EngineId, ResumeToken, TakopiEvent
from takopi.render import MarkdownParts, prepare_telegram
from takopi.router import AutoRouter, RunnerEntry
from takopi.runners.codex import CodexRunner
from takopi.telegram import TelegramClient
from takopi.runners.mock import Advance, Emit, Raise, Return, ScriptRunner, Sleep, Wait
from tests.factories import action_completed, action_started
@@ -189,7 +190,10 @@ class _FakeBot:
disable_notification: bool | None = False,
entities: list[dict] | None = None,
parse_mode: str | None = None,
*,
replace_message_id: int | None = None,
) -> dict:
_ = replace_message_id
self.send_calls.append(
{
"chat_id": chat_id,
@@ -211,7 +215,10 @@ class _FakeBot:
text: str,
entities: list[dict] | None = None,
parse_mode: str | None = None,
*,
wait: bool = True,
) -> dict:
_ = wait
self.edit_calls.append(
{
"chat_id": chat_id,
@@ -223,7 +230,11 @@ class _FakeBot:
)
return {"message_id": message_id}
async def delete_message(self, chat_id: int, message_id: int) -> bool:
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
@@ -281,15 +292,33 @@ class _FakeClock:
self._sleep_event = None
async def sleep(self, delay: float) -> None:
self.sleep_calls += 1
if delay <= 0:
await anyio.sleep(0)
return
self.sleep_calls += 1
self._sleep_until = self._now + delay
self._sleep_event = anyio.Event()
await self._sleep_event.wait()
def _queued_bot(
bot: "_FakeBot", *, clock: "_FakeClock | None" = None
) -> TelegramClient:
if clock is None:
return TelegramClient(
client=bot,
private_chat_rps=0.0,
group_chat_rps=0.0,
)
return TelegramClient(
client=bot,
clock=clock,
sleep=clock.sleep,
private_chat_rps=0.0,
group_chat_rps=0.0,
)
def _return_runner(
*, answer: str = "ok", resume_value: str | None = None
) -> ScriptRunner:
@@ -307,7 +336,7 @@ async def test_final_notify_sends_loud_final_message() -> None:
bot = _FakeBot()
runner = _return_runner(answer="ok")
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -335,7 +364,7 @@ async def test_handle_message_strips_resume_line_from_prompt() -> None:
bot = _FakeBot()
runner = ScriptRunner([Return(answer="ok")], engine=CODEX_ENGINE)
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -366,7 +395,7 @@ async def test_long_final_message_edits_progress_message() -> None:
bot = _FakeBot()
runner = _return_runner(answer="x" * 10_000)
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=False,
@@ -384,7 +413,8 @@ async def test_long_final_message_edits_progress_message() -> None:
assert len(bot.send_calls) == 1
assert bot.send_calls[0]["disable_notification"] is True
assert len(bot.edit_calls) == 1
assert bot.edit_calls
assert "done" in bot.edit_calls[-1]["text"].lower()
@pytest.mark.anyio
@@ -408,7 +438,7 @@ async def test_progress_edits_are_rate_limited() -> None:
advance=clock.set,
)
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot, clock=clock),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -423,12 +453,10 @@ async def test_progress_edits_are_rate_limited() -> None:
text="hi",
resume_token=None,
clock=clock,
sleep=clock.sleep,
progress_edit_every=1.0,
)
assert len(bot.edit_calls) == 1
assert "echo 2" in bot.edit_calls[0]["text"]
assert bot.edit_calls
assert "working" in bot.edit_calls[-1]["text"].lower()
@pytest.mark.anyio
@@ -453,7 +481,7 @@ async def test_progress_edits_do_not_sleep_again_without_new_events() -> None:
advance=clock.set,
)
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot, clock=clock),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -469,33 +497,18 @@ async def test_progress_edits_do_not_sleep_again_without_new_events() -> None:
text="hi",
resume_token=None,
clock=clock,
sleep=clock.sleep,
progress_edit_every=1.0,
)
async with anyio.create_task_group() as tg:
tg.start_soon(run_handle_message)
for _ in range(100):
if clock._sleep_until is not None:
break
await anyio.sleep(0)
assert clock._sleep_until == pytest.approx(1.0)
clock.set(1.0)
for _ in range(100):
if bot.edit_calls:
break
await anyio.sleep(0)
assert len(bot.edit_calls) == 1
for _ in range(5):
await anyio.sleep(0)
assert clock.sleep_calls == 1
assert bot.edit_calls
assert clock.sleep_calls == 0
assert clock._sleep_until is None
hold.set()
@@ -529,7 +542,7 @@ async def test_bridge_flow_sends_progress_edits_and_final_resume() -> None:
resume_value=session_id,
)
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot, clock=clock),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -544,8 +557,6 @@ async def test_bridge_flow_sends_progress_edits_and_final_resume() -> None:
text="do it",
resume_token=None,
clock=clock,
sleep=clock.sleep,
progress_edit_every=1.0,
)
assert bot.send_calls[0]["reply_to_message_id"] == 42
@@ -564,7 +575,7 @@ async def test_handle_cancel_without_reply_prompts_user() -> None:
bot = _FakeBot()
runner = _return_runner(answer="ok")
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -586,7 +597,7 @@ async def test_handle_cancel_with_no_progress_message_says_nothing_running() ->
bot = _FakeBot()
runner = _return_runner(answer="ok")
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -612,7 +623,7 @@ async def test_handle_cancel_with_finished_task_says_nothing_running() -> None:
bot = _FakeBot()
runner = _return_runner(answer="ok")
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -639,7 +650,7 @@ async def test_handle_cancel_cancels_running_task() -> None:
bot = _FakeBot()
runner = _return_runner(answer="ok")
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -669,7 +680,7 @@ async def test_handle_cancel_only_cancels_matching_progress_message() -> None:
bot = _FakeBot()
runner = _return_runner(answer="ok")
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -714,7 +725,7 @@ async def test_handle_message_cancelled_renders_cancelled_state() -> None:
resume_value=session_id,
)
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -764,7 +775,7 @@ async def test_handle_message_error_preserves_resume_token() -> None:
resume_value=session_id,
)
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
@@ -873,6 +884,8 @@ async def test_run_main_loop_routes_reply_to_running_resume() -> None:
disable_notification: bool | None = False,
entities: list[dict] | None = None,
parse_mode: str | None = None,
*,
replace_message_id: int | None = None,
) -> dict:
msg = await super().send_message(
chat_id=chat_id,
@@ -881,6 +894,7 @@ async def test_run_main_loop_routes_reply_to_running_resume() -> None:
disable_notification=disable_notification,
entities=entities,
parse_mode=parse_mode,
replace_message_id=replace_message_id,
)
if self.progress_id is None and reply_to_message_id is not None:
self.progress_id = int(msg["message_id"])
@@ -895,7 +909,7 @@ async def test_run_main_loop_routes_reply_to_running_resume() -> None:
resume_value=resume_value,
)
cfg = BridgeConfig(
bot=bot,
bot=_queued_bot(bot),
router=_make_router(runner),
chat_id=123,
final_notify=True,
tests/test_telegram_client.py
+6 -5
View File
@@ -2,7 +2,7 @@ import httpx
import pytest
from takopi.logging import setup_logging
from takopi.telegram import TelegramClient
from takopi.telegram import TelegramClient, TelegramRetryAfter
@pytest.mark.anyio
@@ -25,12 +25,13 @@ async def test_telegram_429_no_retry() -> None:
client = httpx.AsyncClient(transport=transport)
try:
tg = TelegramClient("123:abcDEF_ghij", client=client)
result = await tg._post("sendMessage", {"chat_id": 1, "text": "hi"})
tg = TelegramClient("123:abcDEF_ghij", http_client=client)
with pytest.raises(TelegramRetryAfter) as exc:
await tg._post("sendMessage", {"chat_id": 1, "text": "hi"})
finally:
await client.aclose()
assert result is None
assert exc.value.retry_after == 3
assert len(calls) == 1
@@ -48,7 +49,7 @@ async def test_no_token_in_logs_on_http_error(
client = httpx.AsyncClient(transport=transport)
try:
tg = TelegramClient(token, client=client)
tg = TelegramClient(token, http_client=client)
await tg._post("getUpdates", {"timeout": 1})
finally:
await client.aclose()
tests/test_telegram_queue.py
+251
View File
@@ -0,0 +1,251 @@
import anyio
import pytest
from takopi.telegram import TelegramClient, TelegramRetryAfter
class _FakeBot:
def __init__(self) -> None:
self.calls: list[str] = []
self.edit_calls: list[str] = []
self.delete_calls: list[tuple[int, int]] = []
self._edit_attempts = 0
self._updates_attempts = 0
self.retry_after: float | None = None
self.updates_retry_after: float | None = 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:
_ = reply_to_message_id
_ = disable_notification
_ = entities
_ = parse_mode
_ = replace_message_id
self.calls.append("send_message")
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:
_ = chat_id
_ = message_id
_ = entities
_ = parse_mode
_ = wait
self.calls.append("edit_message_text")
self.edit_calls.append(text)
if self.retry_after is not None and self._edit_attempts == 0:
self._edit_attempts += 1
raise TelegramRetryAfter(self.retry_after)
self._edit_attempts += 1
return {"message_id": message_id}
async def delete_message(
self,
chat_id: int,
message_id: int,
) -> bool:
self.calls.append("delete_message")
self.delete_calls.append((chat_id, message_id))
return True
async def set_my_commands(
self,
commands: list[dict],
*,
scope: dict | None = None,
language_code: str | None = None,
) -> bool:
_ = commands
_ = scope
_ = language_code
return True
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
if self.updates_retry_after is not None and self._updates_attempts == 0:
self._updates_attempts += 1
raise TelegramRetryAfter(self.updates_retry_after)
self._updates_attempts += 1
return []
async def close(self) -> None:
return None
async def get_me(self) -> dict | None:
return {"id": 1}
@pytest.mark.anyio
async def test_edits_coalesce_latest() -> None:
class _BlockingBot(_FakeBot):
def __init__(self) -> None:
super().__init__()
self.edit_started = anyio.Event()
self.release = anyio.Event()
self._block_first = True
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:
if self._block_first:
self._block_first = False
self.edit_started.set()
await self.release.wait()
return await super().edit_message_text(
chat_id=chat_id,
message_id=message_id,
text=text,
entities=entities,
parse_mode=parse_mode,
wait=wait,
)
bot = _BlockingBot()
client = TelegramClient(client=bot, private_chat_rps=0.0, group_chat_rps=0.0)
await client.edit_message_text(
chat_id=1,
message_id=1,
text="first",
wait=False,
)
with anyio.fail_after(1):
await bot.edit_started.wait()
await client.edit_message_text(
chat_id=1,
message_id=1,
text="second",
wait=False,
)
await client.edit_message_text(
chat_id=1,
message_id=1,
text="third",
wait=False,
)
bot.release.set()
with anyio.fail_after(1):
while len(bot.edit_calls) < 2:
await anyio.sleep(0)
assert bot.edit_calls == ["first", "third"]
@pytest.mark.anyio
async def test_send_preempts_pending_edit() -> None:
bot = _FakeBot()
client = TelegramClient(client=bot, private_chat_rps=10.0, group_chat_rps=10.0)
await client.edit_message_text(
chat_id=1,
message_id=1,
text="first",
)
await client.edit_message_text(
chat_id=1,
message_id=1,
text="progress",
wait=False,
)
with anyio.fail_after(1):
await client.send_message(chat_id=1, text="final")
await anyio.sleep(0.2)
assert bot.calls[0] == "edit_message_text"
assert bot.calls[1] == "send_message"
assert bot.calls[-1] == "edit_message_text"
@pytest.mark.anyio
async def test_delete_drops_pending_edits() -> None:
bot = _FakeBot()
client = TelegramClient(client=bot, private_chat_rps=10.0, group_chat_rps=10.0)
await client.edit_message_text(
chat_id=1,
message_id=1,
text="first",
)
await client.edit_message_text(
chat_id=1,
message_id=1,
text="progress",
wait=False,
)
with anyio.fail_after(1):
await client.delete_message(
chat_id=1,
message_id=1,
)
await anyio.sleep(0.2)
assert bot.delete_calls == [(1, 1)]
assert bot.edit_calls == ["first"]
@pytest.mark.anyio
async def test_retry_after_retries_once() -> None:
bot = _FakeBot()
bot.retry_after = 0.0
client = TelegramClient(client=bot, private_chat_rps=0.0, group_chat_rps=0.0)
result = await client.edit_message_text(
chat_id=1,
message_id=1,
text="retry",
)
assert result == {"message_id": 1}
assert bot._edit_attempts == 2
@pytest.mark.anyio
async def test_get_updates_retries_on_retry_after() -> None:
bot = _FakeBot()
bot.updates_retry_after = 0.0
client = TelegramClient(client=bot, private_chat_rps=0.0, group_chat_rps=0.0)
with anyio.fail_after(1):
updates = await client.get_updates(offset=None, timeout_s=0)
assert updates == []
assert bot._updates_attempts == 2