feat(telegram): add file transfer support (#83)

This commit is contained in:
banteg
2026-01-11 05:32:31 +04:00
committed by GitHub
parent 22d010ece3
commit ab1ecc277d
15 changed files with 2047 additions and 48 deletions
+112
View File
@@ -0,0 +1,112 @@
from __future__ import annotations
from pathlib import Path
from typing import Any
import pytest
from takopi.config import ConfigError, empty_projects_config
from takopi.model import EngineId
from takopi.router import AutoRouter, RunnerEntry
from takopi.runners.mock import Return, ScriptRunner
from takopi.telegram import backend as telegram_backend
from takopi.transport_runtime import TransportRuntime
def test_build_startup_message_includes_missing_engines(tmp_path: Path) -> None:
codex = EngineId("codex")
pi = EngineId("pi")
runner = ScriptRunner([Return(answer="ok")], engine=codex)
missing = ScriptRunner([Return(answer="ok")], engine=pi)
router = AutoRouter(
entries=[
RunnerEntry(engine=codex, runner=runner, available=True),
RunnerEntry(engine=pi, runner=missing, available=False, issue="missing"),
],
default_engine=codex,
)
runtime = TransportRuntime(router=router, projects=empty_projects_config())
message = telegram_backend._build_startup_message(
runtime, startup_pwd=str(tmp_path)
)
assert "takopi is ready" in message
assert "agents: `codex (not installed: pi)`" in message
assert "projects: `none`" in message
def test_telegram_backend_build_and_run_wires_config(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
config_path = tmp_path / "takopi.toml"
config_path.write_text(
'watch_config = true\ntransport = "telegram"\n\n'
"[transports.telegram]\n"
'bot_token = "token"\n'
"chat_id = 321\n",
encoding="utf-8",
)
codex = EngineId("codex")
runner = ScriptRunner([Return(answer="ok")], engine=codex)
router = AutoRouter(
entries=[RunnerEntry(engine=codex, runner=runner, available=True)],
default_engine=codex,
)
runtime = TransportRuntime(router=router, projects=empty_projects_config())
captured: dict[str, Any] = {}
async def fake_run_main_loop(cfg, **kwargs) -> None:
captured["cfg"] = cfg
captured["kwargs"] = kwargs
class _FakeClient:
def __init__(self, token: str) -> None:
self.token = token
async def close(self) -> None:
return None
monkeypatch.setattr(telegram_backend, "run_main_loop", fake_run_main_loop)
monkeypatch.setattr(telegram_backend, "TelegramClient", _FakeClient)
transport_config = {
"bot_token": "token",
"chat_id": 321,
"voice_transcription": True,
"files": {"enabled": True, "allowed_user_ids": [1, 2]},
"topics": {"enabled": True, "scope": "main"},
}
telegram_backend.TelegramBackend().build_and_run(
transport_config=transport_config,
config_path=config_path,
runtime=runtime,
final_notify=False,
default_engine_override=None,
)
cfg = captured["cfg"]
kwargs = captured["kwargs"]
assert cfg.chat_id == 321
assert cfg.voice_transcription is not None
assert cfg.voice_transcription.enabled is True
assert cfg.files.enabled is True
assert cfg.files.allowed_user_ids == frozenset({1, 2})
assert cfg.topics.enabled is True
assert cfg.bot.token == "token"
assert kwargs["watch_config"] is True
assert kwargs["transport_id"] == "telegram"
def test_build_files_config_rejects_non_dict(tmp_path: Path) -> None:
config_path = tmp_path / "takopi.toml"
transport_config: dict[str, object] = {"files": ["nope"]}
with pytest.raises(ConfigError, match="transports.telegram.files"):
telegram_backend._build_files_config(
transport_config,
config_path=config_path,
)
+292 -1
View File
@@ -10,6 +10,7 @@ import takopi.telegram.bridge as bridge
from takopi.directives import parse_directives
from takopi.telegram.bridge import (
TelegramBridgeConfig,
TelegramFilesConfig,
TelegramPresenter,
TelegramTransport,
_build_bot_commands,
@@ -30,7 +31,11 @@ from takopi.progress import ProgressTracker
from takopi.router import AutoRouter, RunnerEntry
from takopi.transport_runtime import TransportRuntime
from takopi.runners.mock import Return, ScriptRunner, Sleep, Wait
from takopi.telegram.types import TelegramCallbackQuery, TelegramIncomingMessage
from takopi.telegram.types import (
TelegramCallbackQuery,
TelegramDocument,
TelegramIncomingMessage,
)
from takopi.transport import MessageRef, RenderedMessage, SendOptions
from tests.plugin_fixtures import FakeEntryPoint, install_entrypoints
@@ -100,6 +105,7 @@ class _FakeBot(BotClient):
self.command_calls: list[dict] = []
self.callback_calls: list[dict] = []
self.send_calls: list[dict] = []
self.document_calls: list[dict] = []
self.edit_calls: list[dict] = []
self.edit_topic_calls: list[dict[str, Any]] = []
self.delete_calls: list[dict] = []
@@ -151,6 +157,29 @@ class _FakeBot(BotClient):
)
return {"message_id": 1}
async def send_document(
self,
chat_id: int,
filename: str,
content: bytes,
reply_to_message_id: int | None = None,
message_thread_id: int | None = None,
disable_notification: bool | None = False,
caption: str | None = None,
) -> dict[str, Any]:
self.document_calls.append(
{
"chat_id": chat_id,
"filename": filename,
"content": content,
"reply_to_message_id": reply_to_message_id,
"message_thread_id": message_thread_id,
"disable_notification": disable_notification,
"caption": caption,
}
)
return {"message_id": 2}
async def edit_message_text(
self,
chat_id: int,
@@ -331,6 +360,7 @@ def test_build_bot_commands_includes_cancel_and_engine() -> None:
commands = _build_bot_commands(runtime)
assert {"command": "cancel", "description": "cancel run"} in commands
assert {"command": "file", "description": "upload or fetch files"} in commands
assert any(cmd["command"] == "codex" for cmd in commands)
@@ -529,6 +559,27 @@ async def test_telegram_transport_edit_wait_false_returns_ref() -> None:
_ = reply_markup
return None
async def send_document(
self,
chat_id: int,
filename: str,
content: bytes,
reply_to_message_id: int | None = None,
message_thread_id: int | None = None,
disable_notification: bool | None = False,
caption: str | None = None,
) -> dict | None:
_ = (
chat_id,
filename,
content,
reply_to_message_id,
message_thread_id,
disable_notification,
caption,
)
return None
async def edit_message_text(
self,
chat_id: int,
@@ -715,6 +766,130 @@ async def test_handle_cancel_only_cancels_matching_progress_message() -> None:
assert len(transport.send_calls) == 0
@pytest.mark.anyio
async def test_handle_file_put_writes_file(tmp_path: Path) -> None:
payload = b"hello"
class _FileBot(_FakeBot):
async def get_file(self, file_id: str) -> dict[str, Any] | None:
_ = file_id
return {"file_path": "files/hello.txt"}
async def download_file(self, file_path: str) -> bytes | None:
_ = file_path
return payload
transport = _FakeTransport()
bot = _FileBot()
runner = ScriptRunner([Return(answer="ok")], engine=CODEX_ENGINE)
projects = ProjectsConfig(
projects={
"proj": ProjectConfig(
alias="proj",
path=tmp_path,
worktrees_dir=Path(".worktrees"),
)
},
default_project=None,
)
runtime = TransportRuntime(router=_make_router(runner), projects=projects)
exec_cfg = ExecBridgeConfig(
transport=transport,
presenter=MarkdownPresenter(),
final_notify=True,
)
cfg = TelegramBridgeConfig(
bot=bot,
runtime=runtime,
chat_id=123,
startup_msg="",
exec_cfg=exec_cfg,
files=TelegramFilesConfig(enabled=True),
)
msg = TelegramIncomingMessage(
transport="telegram",
chat_id=123,
message_id=10,
text="",
reply_to_message_id=None,
reply_to_text=None,
sender_id=321,
chat_type="private",
document=TelegramDocument(
file_id="doc-id",
file_name="hello.txt",
mime_type="text/plain",
file_size=len(payload),
raw={"file_id": "doc-id"},
),
)
await bridge._handle_file_put(cfg, msg, "/proj uploads/hello.txt", None, None)
target = tmp_path / "uploads" / "hello.txt"
assert target.read_bytes() == payload
assert transport.send_calls
text = transport.send_calls[-1]["message"].text
assert "saved uploads/hello.txt" in text
assert "(5 b)" in text
@pytest.mark.anyio
async def test_handle_file_get_sends_document_for_allowed_user(
tmp_path: Path,
) -> None:
payload = b"fetch"
target = tmp_path / "hello.txt"
target.write_bytes(payload)
transport = _FakeTransport()
bot = _FakeBot()
runner = ScriptRunner([Return(answer="ok")], engine=CODEX_ENGINE)
projects = ProjectsConfig(
projects={
"proj": ProjectConfig(
alias="proj",
path=tmp_path,
worktrees_dir=Path(".worktrees"),
)
},
default_project=None,
)
runtime = TransportRuntime(router=_make_router(runner), projects=projects)
exec_cfg = ExecBridgeConfig(
transport=transport,
presenter=MarkdownPresenter(),
final_notify=True,
)
cfg = TelegramBridgeConfig(
bot=bot,
runtime=runtime,
chat_id=123,
startup_msg="",
exec_cfg=exec_cfg,
files=TelegramFilesConfig(
enabled=True,
allowed_user_ids=frozenset({42}),
),
)
msg = TelegramIncomingMessage(
transport="telegram",
chat_id=-100,
message_id=10,
text="",
reply_to_message_id=None,
reply_to_text=None,
sender_id=42,
chat_type="supergroup",
)
await bridge._handle_file_get(cfg, msg, "/proj hello.txt", None, None)
assert bot.document_calls
assert bot.document_calls[0]["filename"] == "hello.txt"
assert bot.document_calls[0]["content"] == payload
@pytest.mark.anyio
async def test_handle_callback_cancel_cancels_running_task() -> None:
transport = _FakeTransport()
@@ -1169,6 +1344,122 @@ async def test_run_main_loop_replies_in_same_thread() -> None:
assert all(call["options"].thread_id == 77 for call in reply_calls)
@pytest.mark.anyio
async def test_run_main_loop_batches_media_group_upload(
tmp_path: Path,
) -> None:
payloads = {
"photos/file_1.jpg": b"one",
"photos/file_2.jpg": b"two",
}
file_map = {
"doc-1": "photos/file_1.jpg",
"doc-2": "photos/file_2.jpg",
}
class _MediaBot(_FakeBot):
async def get_file(self, file_id: str) -> dict[str, Any] | None:
file_path = file_map.get(file_id)
if file_path is None:
return None
return {"file_path": file_path}
async def download_file(self, file_path: str) -> bytes | None:
return payloads.get(file_path)
transport = _FakeTransport()
bot = _MediaBot()
runner = ScriptRunner([Return(answer="ok")], engine=CODEX_ENGINE)
projects = ProjectsConfig(
projects={
"proj": ProjectConfig(
alias="proj",
path=tmp_path,
worktrees_dir=Path(".worktrees"),
)
},
default_project=None,
)
runtime = TransportRuntime(router=_make_router(runner), projects=projects)
exec_cfg = ExecBridgeConfig(
transport=transport,
presenter=MarkdownPresenter(),
final_notify=True,
)
cfg = TelegramBridgeConfig(
bot=bot,
runtime=runtime,
chat_id=123,
startup_msg="",
exec_cfg=exec_cfg,
files=TelegramFilesConfig(enabled=True, auto_put=True),
)
msg1 = TelegramIncomingMessage(
transport="telegram",
chat_id=123,
message_id=1,
text="/file put /proj incoming/test1",
reply_to_message_id=None,
reply_to_text=None,
sender_id=321,
chat_type="private",
media_group_id="grp-1",
document=TelegramDocument(
file_id="doc-1",
file_name=None,
mime_type="image/jpeg",
file_size=len(payloads["photos/file_1.jpg"]),
raw={"file_id": "doc-1"},
),
)
msg2 = TelegramIncomingMessage(
transport="telegram",
chat_id=123,
message_id=2,
text="",
reply_to_message_id=None,
reply_to_text=None,
sender_id=321,
chat_type="private",
media_group_id="grp-1",
document=TelegramDocument(
file_id="doc-2",
file_name=None,
mime_type="image/jpeg",
file_size=len(payloads["photos/file_2.jpg"]),
raw={"file_id": "doc-2"},
),
)
stop_polling = anyio.Event()
async def poller(_cfg: TelegramBridgeConfig):
yield msg1
yield msg2
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(3):
while len(transport.send_calls) < 1:
await anyio.sleep(0.05)
assert len(transport.send_calls) == 1
text = transport.send_calls[0]["message"].text
assert "saved file_1.jpg, file_2.jpg" in text
assert "to incoming/test1/" in text
target_dir = tmp_path / "incoming" / "test1"
assert (target_dir / "file_1.jpg").read_bytes() == payloads[
"photos/file_1.jpg"
]
assert (target_dir / "file_2.jpg").read_bytes() == payloads[
"photos/file_2.jpg"
]
finally:
stop_polling.set()
tg.cancel_scope.cancel()
@pytest.mark.anyio
async def test_run_main_loop_handles_command_plugins(monkeypatch) -> None:
class _Command:
+148 -1
View File
@@ -32,6 +32,7 @@ def test_parse_incoming_update_maps_fields() -> None:
assert msg.chat_type == "supergroup"
assert msg.is_forum is True
assert msg.voice is None
assert msg.document is None
assert msg.raw == update["message"]
@@ -51,7 +52,11 @@ def test_parse_incoming_update_filters_non_matching_chat() -> None:
def test_parse_incoming_update_filters_non_text_and_non_voice() -> None:
update = {
"update_id": 1,
"message": {"message_id": 10, "chat": {"id": 123}, "photo": []},
"message": {
"message_id": 10,
"chat": {"id": 123},
"location": {"latitude": 1.0, "longitude": 2.0},
},
}
assert parse_incoming_update(update, chat_id=123) is None
@@ -84,6 +89,148 @@ def test_parse_incoming_update_voice_message() -> None:
assert msg.voice.duration == 3
def test_parse_incoming_update_document_message() -> None:
update = {
"update_id": 1,
"message": {
"message_id": 10,
"caption": "/file put incoming/doc.txt",
"chat": {"id": 123},
"document": {
"file_id": "doc-id",
"file_unique_id": "uniq",
"file_name": "doc.txt",
"mime_type": "text/plain",
"file_size": 4321,
},
},
}
msg = parse_incoming_update(update, chat_id=123)
assert msg is not None
assert isinstance(msg, TelegramIncomingMessage)
assert msg.text == "/file put incoming/doc.txt"
assert msg.document is not None
assert msg.document.file_id == "doc-id"
assert msg.document.file_name == "doc.txt"
assert msg.document.mime_type == "text/plain"
assert msg.document.file_size == 4321
def test_parse_incoming_update_photo_message() -> None:
update = {
"update_id": 1,
"message": {
"message_id": 10,
"caption": "/file put incoming/photo.jpg",
"chat": {"id": 123},
"photo": [
{
"file_id": "small",
"file_unique_id": "uniq-small",
"file_size": 100,
"width": 90,
"height": 90,
},
{
"file_id": "large",
"file_unique_id": "uniq-large",
"file_size": 1000,
"width": 800,
"height": 600,
},
],
},
}
msg = parse_incoming_update(update, chat_id=123)
assert msg is not None
assert isinstance(msg, TelegramIncomingMessage)
assert msg.text == "/file put incoming/photo.jpg"
assert msg.document is not None
assert msg.document.file_id == "large"
assert msg.document.file_name is None
assert msg.document.file_size == 1000
def test_parse_incoming_update_media_group_id() -> None:
update = {
"update_id": 1,
"message": {
"message_id": 10,
"chat": {"id": 123},
"media_group_id": "group-1",
"photo": [
{
"file_id": "large",
"file_unique_id": "uniq-large",
"file_size": 1000,
"width": 800,
"height": 600,
}
],
},
}
msg = parse_incoming_update(update, chat_id=123)
assert msg is not None
assert isinstance(msg, TelegramIncomingMessage)
assert msg.media_group_id == "group-1"
def test_parse_incoming_update_video_message() -> None:
update = {
"update_id": 1,
"message": {
"message_id": 10,
"caption": "/file put incoming/video.mp4",
"chat": {"id": 123},
"video": {
"file_id": "video-id",
"file_unique_id": "uniq",
"file_name": "video.mp4",
"mime_type": "video/mp4",
"file_size": 4242,
},
},
}
msg = parse_incoming_update(update, chat_id=123)
assert msg is not None
assert isinstance(msg, TelegramIncomingMessage)
assert msg.text == "/file put incoming/video.mp4"
assert msg.document is not None
assert msg.document.file_id == "video-id"
assert msg.document.file_name == "video.mp4"
assert msg.document.mime_type == "video/mp4"
assert msg.document.file_size == 4242
def test_parse_incoming_update_sticker_message() -> None:
update = {
"update_id": 1,
"message": {
"message_id": 10,
"chat": {"id": 123},
"sticker": {
"file_id": "sticker-id",
"file_unique_id": "uniq",
"file_size": 2468,
},
},
}
msg = parse_incoming_update(update, chat_id=123)
assert msg is not None
assert isinstance(msg, TelegramIncomingMessage)
assert msg.text == ""
assert msg.document is not None
assert msg.document.file_id == "sticker-id"
assert msg.document.file_name is None
assert msg.document.mime_type is None
assert msg.document.file_size == 2468
def test_parse_incoming_update_callback_query() -> None:
update = {
"update_id": 1,
+22
View File
@@ -40,6 +40,28 @@ class _FakeBot(BotClient):
self.calls.append("send_message")
return {"message_id": 1}
async def send_document(
self,
chat_id: int,
filename: str,
content: bytes,
reply_to_message_id: int | None = None,
message_thread_id: int | None = None,
disable_notification: bool | None = False,
caption: str | None = None,
) -> dict[str, Any]:
_ = (
chat_id,
filename,
content,
reply_to_message_id,
message_thread_id,
disable_notification,
caption,
)
self.calls.append("send_document")
return {"message_id": 1}
async def edit_message_text(
self,
chat_id: int,