diff --git a/bot/codex/screenshot_ui.py b/bot/codex/screenshot_ui.py deleted file mode 100644 index c38d61f..0000000 --- a/bot/codex/screenshot_ui.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Telegram /screenshot command handling.""" - -from __future__ import annotations - -from aiogram.types import Message - -from bot.utils.routing import TelegramRoute -from utils.screenshot import send_screenshot_to_chat - - -async def handle_screenshot_command(message: Message) -> None: - """Capture the terminal window and send it as a photo.""" - route = TelegramRoute.from_message(message) - await send_screenshot_to_chat(message, route) diff --git a/bot/handlers/__init__.py b/bot/handlers/__init__.py index 6c09db8..0290f98 100644 --- a/bot/handlers/__init__.py +++ b/bot/handlers/__init__.py @@ -3,11 +3,13 @@ from .help import router as help_router from .history import router as history_router from .messages import router as messages_router +from .screenshot import router as screenshot_router from .send import router as send_router from .toolbar import router as toolbar_router __all__ = [ "messages_router", + "screenshot_router", "callbacks_router", "codex_router", "help_router", diff --git a/bot/handlers/codex.py b/bot/handlers/codex.py index 0869f74..76b2778 100644 --- a/bot/handlers/codex.py +++ b/bot/handlers/codex.py @@ -12,15 +12,14 @@ from typing import Any from aiogram import F, Router +from aiogram.dispatcher.event.bases import SkipHandler from aiogram.filters import Command from aiogram.types import CallbackQuery, Message -from aiogram.dispatcher.event.bases import SkipHandler from bot.codex import ( command_flow, model_ui, reply_ui, - screenshot_ui, session_ui, shell_ui, stop_ui, @@ -48,8 +47,8 @@ from core.orchestrator import Orchestrator from core.session import SessionKey from extensions.codex.daemon import codex_daemon -from storage.context_manager import ContextManager from i18n import resolve_locale, tr +from storage.context_manager import ContextManager router = Router(name="codex") @@ -375,12 +374,3 @@ async def handle_shell_approve_callback( orchestrator: Orchestrator, ) -> None: await shell_ui.handle_shell_approve_callback(callback_query, orchestrator, execute_shell=_execute_shell_telegram) - -# --------------------------------------------------------------------------- -# /screenshot command handler -# --------------------------------------------------------------------------- - - -@router.message(Command("screenshot")) -async def cmd_screenshot(message: Message) -> None: - await screenshot_ui.handle_screenshot_command(message) diff --git a/bot/handlers/screenshot.py b/bot/handlers/screenshot.py new file mode 100644 index 0000000..d79876c --- /dev/null +++ b/bot/handlers/screenshot.py @@ -0,0 +1,179 @@ +"""Telegram-native desktop screenshot commands and monitor picker.""" + +from __future__ import annotations + +import html +from contextlib import suppress + +from aiogram import F, Router +from aiogram.filters import Command +from aiogram.types import ( + BufferedInputFile, + CallbackQuery, + InlineKeyboardButton, + InlineKeyboardMarkup, + Message, +) +from aiogram.types import ( + User as TelegramUser, +) +from loguru import logger +from sqlalchemy import select + +from bot.utils.routing import TelegramRoute +from i18n import resolve_locale, tr +from storage.context_manager import ContextManager +from storage.models import User +from utils.screenshot import ( + DisplayMonitor, + capture_desktop_screenshot, + list_display_monitors, +) + +router = Router(name="screenshot") + +_CALLBACK_PREFIX = "screenshot:" +_CALLBACK_DATA_LIMIT = 64 + + +async def _resolve_user_locale( + context_manager: ContextManager | None, + user: TelegramUser | None, +) -> str: + fallback_code = user.language_code if user else None + if context_manager is None or user is None: + return resolve_locale(None, fallback_code) + + result = await context_manager.session.execute(select(User).where(User.id == user.id)) + stored_user = result.scalar_one_or_none() + return resolve_locale( + getattr(stored_user, "ui_language", None) if stored_user else None, + fallback_code, + ) + + +def _monitor_callback_data(user_id: int, monitor: DisplayMonitor) -> str: + prefix = f"{_CALLBACK_PREFIX}{user_id}:" + available = _CALLBACK_DATA_LIMIT - len(prefix.encode("utf-8")) + identifier = monitor.identifier.encode("utf-8")[:available].decode("utf-8", errors="ignore") + return f"{prefix}{identifier}" + + +def build_monitor_keyboard( + monitors: list[DisplayMonitor], + user_id: int, + locale: str, +) -> InlineKeyboardMarkup: + """Build one full-width row for each available display.""" + rows: list[list[InlineKeyboardButton]] = [] + for monitor in monitors: + key = "bot.screenshot.monitor_button_primary" if monitor.is_primary else "bot.screenshot.monitor_button" + rows.append( + [ + InlineKeyboardButton( + text=tr( + key, + locale, + name=monitor.name, + width=monitor.width, + height=monitor.height, + ), + callback_data=_monitor_callback_data(user_id, monitor), + ) + ] + ) + return InlineKeyboardMarkup(inline_keyboard=rows) + + +async def _send_monitor_screenshot( + message: Message, + route: TelegramRoute, + locale: str, + monitor: DisplayMonitor | None, +) -> bool: + png_bytes = await capture_desktop_screenshot(monitor) + if png_bytes is None: + await message.answer( + tr("bot.screenshot.capture_failed", locale), + **route.send_kwargs(), + ) + return False + + monitor_label = monitor.name if monitor else tr("bot.screenshot.default_monitor", locale) + try: + await message.answer_photo( + photo=BufferedInputFile(png_bytes, filename="screenshot.png"), + caption=tr("bot.screenshot.caption", locale, monitor=monitor_label), + **route.send_kwargs(), + ) + except Exception as exc: + logger.warning(f"Failed to send desktop screenshot: {exc}") + await message.answer( + tr("bot.screenshot.send_failed", locale, error=html.escape(str(exc))), + **route.send_kwargs(), + ) + return False + return True + + +@router.message(Command("screenshot")) +async def cmd_screenshot( + message: Message, + context_manager: ContextManager | None = None, +) -> None: + """Capture the only display or ask the user which display to capture.""" + route = TelegramRoute.from_message(message) + locale = await _resolve_user_locale(context_manager, message.from_user) + monitors = list_display_monitors() + + if len(monitors) > 1: + user_id = message.from_user.id if message.from_user else 0 + await message.answer( + tr("bot.screenshot.select_monitor", locale), + reply_markup=build_monitor_keyboard(monitors, user_id, locale), + **route.send_kwargs(), + ) + return + + await _send_monitor_screenshot( + message, + route, + locale, + monitors[0] if monitors else None, + ) + + +@router.callback_query(F.data.startswith(_CALLBACK_PREFIX)) +async def handle_screenshot_callback( + callback: CallbackQuery, + context_manager: ContextManager | None = None, +) -> None: + """Capture the display selected by the user who opened the picker.""" + locale = await _resolve_user_locale(context_manager, callback.from_user) + data = callback.data or "" + parts = data.split(":", maxsplit=2) + if len(parts) != 3 or not parts[1].isdigit(): + await callback.answer(tr("bot.screenshot.monitor_unavailable", locale), show_alert=True) + return + + owner_id = int(parts[1]) + if callback.from_user.id != owner_id: + await callback.answer(tr("bot.screenshot.picker_owner_mismatch", locale), show_alert=True) + return + + monitor = next( + (candidate for candidate in list_display_monitors() if _monitor_callback_data(owner_id, candidate) == data), + None, + ) + if monitor is None or not isinstance(callback.message, Message): + await callback.answer(tr("bot.screenshot.monitor_unavailable", locale), show_alert=True) + return + + await callback.answer(tr("bot.screenshot.capturing", locale)) + with suppress(Exception): + await callback.message.edit_reply_markup(reply_markup=None) + + route = TelegramRoute.from_message(callback.message) + if await _send_monitor_screenshot(callback.message, route, locale, monitor): + with suppress(Exception): + await callback.message.delete() diff --git a/bot/startup.py b/bot/startup.py index 7869018..d11a945 100644 --- a/bot/startup.py +++ b/bot/startup.py @@ -35,7 +35,7 @@ async def send_message(self, chat_id: int, text: str, **kwargs: Any) -> Any: BotCommand(command="stop", description="Interrupt the active Codex or shell task"), BotCommand(command="live", description="Toggle live session status"), BotCommand(command="last", description="Resend the last assistant reply"), - BotCommand(command="screenshot", description="Capture the current desktop"), + BotCommand(command="screenshot", description="Capture a desktop display"), ) diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 3ee8f95..f16acfe 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,7 +1,7 @@ --- title: Changelog category: changelog -last_updated: 2026-07-03 +last_updated: 2026-07-12 relevance: medium summary: Human-maintained release notes for Telegodex related: [PRODUCT_EXPERIENCE.md, RICH_MESSAGES.md, STARTUP.md] @@ -11,6 +11,7 @@ related: [PRODUCT_EXPERIENCE.md, RICH_MESSAGES.md, STARTUP.md] ## Unreleased +- Moved `/screenshot` to the Bot-native supervision surface and added a Windows multi-display picker with one full-width name-and-resolution button per monitor. - Replaced the single-text `/help` reply with a paginated ebook system: `/help` now opens a table-of-contents inline keyboard (max 6 chapters per page), chapter content is rendered as Rich Message Markdown sourced from per-language `.md` files under `i18n/help//`, with YAML frontmatter, page breaks, JSONL error logging to `i18n/help/log/`, and English as the base fallback language. - Improved Telegram inline keyboard layout so settings and action menus use compact rows with full-width footer buttons. - Added i18n system. diff --git a/docs/USAGE.md b/docs/USAGE.md index ce5e294..fde2217 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -29,7 +29,7 @@ Telegodex runs as a Telegram Workbench for local CLI AI workflows. Its primary r | `/stop` | Interrupt the active Codex turn or shell process | | `/live` | Toggle live session status | | `/last` | Resend the last assistant reply | -| `/screenshot` | Capture the current desktop | +| `/screenshot` | Capture a desktop display; choose one when several are connected | Telegodex syncs this command menu with Telegram on startup through `setMyCommands`. BotFather is still where the bot token and platform-level bot settings live, but users should not need to maintain the slash-command menu manually. @@ -146,4 +146,4 @@ Keep `.env`, local `provider.toml`, database files, and logs out of commits. `.e If a provider call fails, confirm the API key, base URL, model name, and network path. For custom providers, test the endpoint with a minimal OpenAI-compatible chat-completions request. -If `/screenshot` fails with an empty-image warning even though Pillow is installed, focus or unminimize the terminal/window and retry. Telegodex falls back from terminal-window capture to full-screen capture, but Windows can still return an empty image for minimized, protected, or unavailable capture surfaces. +If `/screenshot` cannot capture a display, make sure the Windows desktop session is unlocked and Pillow or pyautogui is installed. When a selected display was disconnected or rearranged, run `/screenshot` again to refresh the display list. diff --git a/i18n/help/en/codex.md b/i18n/help/en/codex.md index 1fda6a0..fa11293 100644 --- a/i18n/help/en/codex.md +++ b/i18n/help/en/codex.md @@ -104,15 +104,6 @@ This is the fastest path from "I want this done" to "it is running". --- -## Screenshots - -`/screenshot` captures the current terminal window where Codex is running -and sends it as a photo. Use it when you want to verify what Codex is -looking at right now — file listings, editor state, build output — without -leaving Telegram. - ---- - ## Configuration Codex Bridge auto-detects the `codex` binary from `PATH` and common install diff --git a/i18n/help/en/getting-started.md b/i18n/help/en/getting-started.md index d3fb5ce..4e61109 100644 --- a/i18n/help/en/getting-started.md +++ b/i18n/help/en/getting-started.md @@ -24,7 +24,7 @@ bottom of the chat. This chapter lists every entry point. | `/history` | View conversation history | | `/codex` | Drive Codex CLI as a subprocess (see 🛠️ Codex Bridge) | | `/shell` | Propose or run a shell command through Codex | -| `/screenshot` | Capture the current terminal window | +| `/screenshot` | Capture a desktop display; choose one when several are connected | Commands are case-insensitive and can be typed anywhere in the chat. diff --git a/i18n/help/zh-cn/codex.md b/i18n/help/zh-cn/codex.md index 4e70f9d..cba4e61 100644 --- a/i18n/help/zh-cn/codex.md +++ b/i18n/help/zh-cn/codex.md @@ -97,14 +97,6 @@ Codex 执行: --- -## 截图 - -`/screenshot` 捕获 Codex 当前所在终端窗口,并以图片形式发送。当你想 -核对 Codex 此刻正在看什么——文件列表、编辑器状态、构建输出——而不 -离开 Telegram 时使用它。 - ---- - ## 配置 Codex 桥接会从 `PATH` 与常见安装路径自动探测 `codex` 可执行文件。仅 diff --git a/i18n/help/zh-cn/getting-started.md b/i18n/help/zh-cn/getting-started.md index b4eea6b..286bd9c 100644 --- a/i18n/help/zh-cn/getting-started.md +++ b/i18n/help/zh-cn/getting-started.md @@ -24,7 +24,7 @@ Telegodex 通过 Telegram 命令与聊天窗口底部的菜单按钮进行操作 | `/history` | 查看对话历史 | | `/codex` | 以子进程方式驱动 Codex CLI(见 🛠️ Codex 桥接) | | `/shell` | 通过 Codex 提议或执行 Shell 命令 | -| `/screenshot` | 捕获当前终端窗口 | +| `/screenshot` | 截取桌面显示器；多屏时可先选择 | 命令不区分大小写,可以在对话任意位置输入。 diff --git a/i18n/locales/en.json b/i18n/locales/en.json index 01f3275..6e6cf6a 100644 --- a/i18n/locales/en.json +++ b/i18n/locales/en.json @@ -85,6 +85,18 @@ "codex_command_unhandled": "Codex commands open or control dedicated Codex topics. Use /codex new or /codex resume <thread-id> to open one.", "codex_topic_command_unsupported": "This is a Codex topic. Use plain messages to continue the Codex thread, or use /codex status for Codex-specific controls. Open general Bot settings and AI chat commands from All or a non-Codex topic." }, + "screenshot": { + "select_monitor": "Choose a display to capture:", + "monitor_button": "{name} · {width}x{height}", + "monitor_button_primary": "{name} · {width}x{height} · Primary", + "capturing": "Capturing display...", + "capture_failed": "Could not capture this display. Make sure the desktop session is unlocked and Pillow or pyautogui is available.", + "send_failed": "Failed to send screenshot: {error}", + "caption": "Desktop screenshot · {monitor}", + "default_monitor": "default display", + "monitor_unavailable": "That display is no longer available. Run /screenshot again.", + "picker_owner_mismatch": "This display picker belongs to another user." + }, "common": { "back": "« Back", "close": "« Close", diff --git a/i18n/locales/zh-cn.json b/i18n/locales/zh-cn.json index a79faaa..68c5d88 100644 --- a/i18n/locales/zh-cn.json +++ b/i18n/locales/zh-cn.json @@ -85,6 +85,18 @@ "codex_command_unhandled": "Codex 指令用于打开或控制专属 Codex Topic。请使用 /codex new 或 /codex resume <thread-id> 打开一个 Codex Topic。", "codex_topic_command_unsupported": "这里是 Codex Topic。请直接发送普通消息继续当前 Codex thread，或使用 /codex status 等 Codex 专属控制。通用 Bot 设置和普通 AI Chat 指令请在 All 或非 Codex Topic 中使用。" }, + "screenshot": { + "select_monitor": "请选择要截图的显示器：", + "monitor_button": "{name} · {width}x{height}", + "monitor_button_primary": "{name} · {width}x{height} · 主屏幕", + "capturing": "正在截取显示器...", + "capture_failed": "无法截取该显示器。请确认桌面会话未锁定，并已安装 Pillow 或 pyautogui。", + "send_failed": "截图发送失败：{error}", + "caption": "桌面截图 · {monitor}", + "default_monitor": "默认显示器", + "monitor_unavailable": "该显示器已不可用，请重新发送 /screenshot。", + "picker_owner_mismatch": "这个显示器选择器属于其他用户。" + }, "common": { "back": "« 返回", "close": "« 关闭", diff --git a/main.py b/main.py index df3999a..4d7abec 100644 --- a/main.py +++ b/main.py @@ -26,6 +26,7 @@ help_router, history_router, messages_router, + screenshot_router, send_router, toolbar_router, ) @@ -1077,9 +1078,10 @@ async def main(): await db.init_db() # Initialize i18n - from i18n import get_i18n_manager from pathlib import Path + from i18n import get_i18n_manager + locales_dir = Path(__file__).parent / "i18n" / "locales" try: i18n_manager = get_i18n_manager() @@ -1165,6 +1167,7 @@ async def main(): # 注册路由 dp.include_router(toolbar_router) dp.include_router(codex_router) + dp.include_router(screenshot_router) dp.include_router(history_router) dp.include_router(send_router) dp.include_router(help_router) diff --git a/tests/unit/test_codex_handlers.py b/tests/unit/test_codex_handlers.py index a302c54..5841d76 100644 --- a/tests/unit/test_codex_handlers.py +++ b/tests/unit/test_codex_handlers.py @@ -670,16 +670,6 @@ async def test_cmd_shell_delegates_to_shell_ui(monkeypatch: pytest.MonkeyPatch) assert kwargs["execute_shell"] is codex._execute_shell_telegram -@pytest.mark.asyncio -async def test_cmd_screenshot_delegates_to_screenshot_ui(monkeypatch: pytest.MonkeyPatch) -> None: - message = _message("/screenshot") - handle_screenshot_command = AsyncMock() - monkeypatch.setattr(codex.screenshot_ui, "handle_screenshot_command", handle_screenshot_command) - - await codex.cmd_screenshot(message) - - handle_screenshot_command.assert_awaited_once_with(message) - @pytest.mark.asyncio async def test_codex_stop_callback_delegates_to_stop_ui(monkeypatch: pytest.MonkeyPatch) -> None: orchestrator = SimpleNamespace() diff --git a/tests/unit/test_codex_screenshot_ui.py b/tests/unit/test_codex_screenshot_ui.py deleted file mode 100644 index 374ac6d..0000000 --- a/tests/unit/test_codex_screenshot_ui.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Unit tests for Codex screenshot command handling.""" - -from __future__ import annotations - -from datetime import datetime -from unittest.mock import AsyncMock - -import pytest -from aiogram.types import Message - -from bot.codex import screenshot_ui - - -def _message( - text: str = "/screenshot", - *, - bot: AsyncMock | None = None, - chat_id: int = 100, - user_id: int = 7, - message_thread_id: int | None = 222, -) -> Message: - payload = { - "message_id": 1, - "date": int(datetime.now().timestamp()), - "chat": {"id": chat_id, "type": "supergroup"}, - "from": { - "id": user_id, - "is_bot": False, - "first_name": "Test", - }, - "text": text, - } - if message_thread_id is not None: - payload["message_thread_id"] = message_thread_id - payload["is_topic_message"] = True - return Message.model_validate(payload).as_(bot or AsyncMock()) - - -@pytest.mark.asyncio -async def test_handle_screenshot_command_preserves_topic_route(monkeypatch: pytest.MonkeyPatch) -> None: - message = _message(message_thread_id=222) - send_screenshot_to_chat = AsyncMock() - monkeypatch.setattr(screenshot_ui, "send_screenshot_to_chat", send_screenshot_to_chat) - - await screenshot_ui.handle_screenshot_command(message) - - send_screenshot_to_chat.assert_awaited_once() - args = send_screenshot_to_chat.await_args.args - assert args[0] is message - assert args[1].chat_id == 100 - assert args[1].message_thread_id == 222 diff --git a/tests/unit/test_screenshot.py b/tests/unit/test_screenshot.py index 9044947..39884f6 100644 --- a/tests/unit/test_screenshot.py +++ b/tests/unit/test_screenshot.py @@ -17,41 +17,95 @@ def save(self, buffer: BytesIO, format: str) -> None: img.save(buffer, format=format) +def test_optional_pyautogui_symbol_is_always_defined() -> None: + assert hasattr(screenshot, "pyautogui") + + @pytest.mark.asyncio -async def test_capture_screenshot_ignores_invalid_bbox_and_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: - calls = [] +async def test_capture_selected_monitor_uses_virtual_desktop_bounds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[tuple[int, int, int, int] | None, bool]] = [] - def grab(*, bbox=None): - calls.append(bbox) + def grab(*, bbox=None, all_screens=False): + calls.append((bbox, all_screens)) return _FakeImage() monkeypatch.setattr(screenshot, "_HAS_PIL", True) monkeypatch.setattr(screenshot, "_HAS_PYAUTOGUI", False) - monkeypatch.setattr(screenshot, "_get_terminal_window_rect", lambda: (10, 10, 10, 20)) + monkeypatch.setattr(screenshot, "_enable_windows_dpi_awareness", lambda: None) + monkeypatch.setattr(screenshot.sys, "platform", "win32") monkeypatch.setattr(screenshot.ImageGrab, "grab", grab) + monitor = screenshot.DisplayMonitor( + identifier="DISPLAY2", + name="DISPLAY2", + left=-1920, + top=0, + right=0, + bottom=1080, + ) - data = await screenshot.capture_terminal_screenshot() + data = await screenshot.capture_desktop_screenshot(monitor) assert data is not None - assert calls == [None] + assert calls == [((-1920, 0, 0, 1080), True)] @pytest.mark.asyncio -async def test_capture_screenshot_retries_fullscreen_after_empty_window_image(monkeypatch: pytest.MonkeyPatch) -> None: - calls = [] - - def grab(*, bbox=None): - calls.append(bbox) - if bbox is not None: - return _FakeImage(size=(0, 0)) - return _FakeImage(size=(2, 2)) +async def test_capture_selected_monitor_falls_back_to_pyautogui_region( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pil_calls = [] + pyautogui_calls = [] + + def grab(*, bbox=None, all_screens=False): + pil_calls.append((bbox, all_screens)) + return _FakeImage(size=(0, 0)) + + class FakePyAutoGui: + @staticmethod + def screenshot(*, region=None): + pyautogui_calls.append(region) + return _FakeImage() monkeypatch.setattr(screenshot, "_HAS_PIL", True) - monkeypatch.setattr(screenshot, "_HAS_PYAUTOGUI", False) - monkeypatch.setattr(screenshot, "_get_terminal_window_rect", lambda: (1, 2, 20, 30)) + monkeypatch.setattr(screenshot, "_HAS_PYAUTOGUI", True) + monkeypatch.setattr(screenshot, "_enable_windows_dpi_awareness", lambda: None) + monkeypatch.setattr(screenshot.sys, "platform", "win32") monkeypatch.setattr(screenshot.ImageGrab, "grab", grab) - - data = await screenshot.capture_terminal_screenshot() + monkeypatch.setattr(screenshot, "pyautogui", FakePyAutoGui()) + monitor = screenshot.DisplayMonitor( + identifier="DISPLAY1", + name="DISPLAY1", + left=10, + top=20, + right=810, + bottom=620, + ) + + data = await screenshot.capture_desktop_screenshot(monitor) assert data is not None - assert calls == [(1, 2, 20, 30), None] + assert pil_calls == [((10, 20, 810, 620), True)] + assert pyautogui_calls == [(10, 20, 800, 600)] + + +@pytest.mark.asyncio +async def test_capture_rejects_invalid_monitor_bounds(monkeypatch: pytest.MonkeyPatch) -> None: + def grab(**_kwargs): + raise AssertionError("capture should not be attempted") + + monkeypatch.setattr(screenshot, "_HAS_PIL", True) + monkeypatch.setattr(screenshot, "_HAS_PYAUTOGUI", False) + monkeypatch.setattr(screenshot, "_enable_windows_dpi_awareness", lambda: None) + monkeypatch.setattr(screenshot.ImageGrab, "grab", grab) + monitor = screenshot.DisplayMonitor( + identifier="DISPLAY1", + name="DISPLAY1", + left=10, + top=10, + right=10, + bottom=100, + ) + + assert await screenshot.capture_desktop_screenshot(monitor) is None diff --git a/tests/unit/test_screenshot_handler.py b/tests/unit/test_screenshot_handler.py new file mode 100644 index 0000000..b822516 --- /dev/null +++ b/tests/unit/test_screenshot_handler.py @@ -0,0 +1,152 @@ +"""Unit tests for the Telegram-native screenshot monitor picker.""" + +from __future__ import annotations + +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from aiogram.types import Message + +from bot.handlers import screenshot +from utils.screenshot import DisplayMonitor + + +def _monitor( + identifier: str, + *, + name: str | None = None, + left: int = 0, + width: int = 1920, + height: int = 1080, + primary: bool = False, +) -> DisplayMonitor: + return DisplayMonitor( + identifier=identifier, + name=name or identifier, + left=left, + top=0, + right=left + width, + bottom=height, + is_primary=primary, + ) + + +def _message(*, bot: AsyncMock | None = None, message_thread_id: int = 222) -> Message: + payload = { + "message_id": 1, + "date": int(datetime.now().timestamp()), + "chat": {"id": 100, "type": "supergroup"}, + "from": { + "id": 7, + "is_bot": False, + "first_name": "Test", + "language_code": "en", + }, + "text": "/screenshot", + "message_thread_id": message_thread_id, + "is_topic_message": True, + } + return Message.model_validate(payload).as_(bot or AsyncMock()) + + +def test_monitor_keyboard_uses_one_row_per_display_and_bounded_callbacks() -> None: + monitors = [ + _monitor("DISPLAY1", primary=True), + _monitor("D" * 80, name="Wide secondary display", left=-2560, width=2560, height=1440), + ] + + keyboard = screenshot.build_monitor_keyboard(monitors, user_id=7, locale="en") + + assert len(keyboard.inline_keyboard) == 2 + assert all(len(row) == 1 for row in keyboard.inline_keyboard) + assert "DISPLAY1" in keyboard.inline_keyboard[0][0].text + assert "1920x1080" in keyboard.inline_keyboard[0][0].text + assert "Primary" in keyboard.inline_keyboard[0][0].text + assert all(len((row[0].callback_data or "").encode("utf-8")) <= 64 for row in keyboard.inline_keyboard) + + +@pytest.mark.asyncio +async def test_command_prompts_when_multiple_monitors_are_available( + monkeypatch: pytest.MonkeyPatch, +) -> None: + message = _message() + monitors = [_monitor("DISPLAY1", primary=True), _monitor("DISPLAY2", left=1920)] + send_monitor_screenshot = AsyncMock() + monkeypatch.setattr(screenshot, "_resolve_user_locale", AsyncMock(return_value="en")) + monkeypatch.setattr(screenshot, "list_display_monitors", lambda: monitors) + monkeypatch.setattr(screenshot, "_send_monitor_screenshot", send_monitor_screenshot) + + await screenshot.cmd_screenshot(message) + + send_monitor_screenshot.assert_not_awaited() + method = message.bot.await_args.args[0] + assert method.message_thread_id == 222 + assert len(method.reply_markup.inline_keyboard) == 2 + + +@pytest.mark.asyncio +async def test_command_captures_single_monitor_in_same_topic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + message = _message() + monitor = _monitor("DISPLAY1", primary=True) + send_monitor_screenshot = AsyncMock(return_value=True) + monkeypatch.setattr(screenshot, "_resolve_user_locale", AsyncMock(return_value="en")) + monkeypatch.setattr(screenshot, "list_display_monitors", lambda: [monitor]) + monkeypatch.setattr(screenshot, "_send_monitor_screenshot", send_monitor_screenshot) + + await screenshot.cmd_screenshot(message) + + send_monitor_screenshot.assert_awaited_once() + args = send_monitor_screenshot.await_args.args + assert args[0] is message + assert args[1].chat_id == 100 + assert args[1].message_thread_id == 222 + assert args[3] is monitor + + +@pytest.mark.asyncio +async def test_monitor_callback_rejects_other_users(monkeypatch: pytest.MonkeyPatch) -> None: + monitor = _monitor("DISPLAY1") + data = screenshot._monitor_callback_data(7, monitor) + callback = SimpleNamespace( + data=data, + from_user=SimpleNamespace(id=8, language_code="en"), + message=None, + answer=AsyncMock(), + ) + monkeypatch.setattr(screenshot, "_resolve_user_locale", AsyncMock(return_value="en")) + + await screenshot.handle_screenshot_callback(callback) + + callback.answer.assert_awaited_once() + assert callback.answer.await_args.kwargs["show_alert"] is True + + +@pytest.mark.asyncio +async def test_monitor_callback_captures_selected_monitor_in_same_topic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + message = _message() + monitor = _monitor("DISPLAY2", left=1920) + data = screenshot._monitor_callback_data(7, monitor) + callback = SimpleNamespace( + data=data, + from_user=message.from_user, + message=message, + answer=AsyncMock(), + ) + send_monitor_screenshot = AsyncMock(return_value=True) + monkeypatch.setattr(screenshot, "_resolve_user_locale", AsyncMock(return_value="en")) + monkeypatch.setattr(screenshot, "list_display_monitors", lambda: [monitor]) + monkeypatch.setattr(screenshot, "_send_monitor_screenshot", send_monitor_screenshot) + + await screenshot.handle_screenshot_callback(callback) + + send_monitor_screenshot.assert_awaited_once() + args = send_monitor_screenshot.await_args.args + assert args[0] is message + assert args[1].message_thread_id == 222 + assert args[3] is monitor diff --git a/utils/screenshot.py b/utils/screenshot.py index cfe8048..911bf20 100644 --- a/utils/screenshot.py +++ b/utils/screenshot.py @@ -1,17 +1,51 @@ -"""Screenshot utilities for capturing the terminal window.""" +"""Desktop screenshot capture and terminal rendering utilities.""" from __future__ import annotations import asyncio import re import sys +from dataclasses import dataclass from io import BytesIO from typing import Any -from aiogram.types import BufferedInputFile, Message from loguru import logger -from bot.utils.routing import TelegramRoute +_DPI_AWARENESS_ATTEMPTED = False + + +def _enable_windows_dpi_awareness() -> None: + """Use physical pixel coordinates for mixed-DPI monitor layouts.""" + global _DPI_AWARENESS_ATTEMPTED + if sys.platform != "win32" or _DPI_AWARENESS_ATTEMPTED: + return + _DPI_AWARENESS_ATTEMPTED = True + + try: + import ctypes + + if ctypes.windll.user32.SetProcessDpiAwarenessContext(ctypes.c_void_p(-4)): + return + except Exception as exc: + logger.debug(f"Per-monitor DPI awareness v2 is unavailable: {exc}") + + try: + import ctypes + + if ctypes.windll.shcore.SetProcessDpiAwareness(2) == 0: + return + except Exception as exc: + logger.debug(f"Per-monitor DPI awareness is unavailable: {exc}") + + try: + import ctypes + + ctypes.windll.user32.SetProcessDPIAware() + except Exception as exc: + logger.debug(f"System DPI awareness is unavailable: {exc}") + + +_enable_windows_dpi_awareness() # Optional dependencies — handled via ImportError try: @@ -29,22 +63,9 @@ _HAS_PYAUTOGUI = True except ImportError: + pyautogui = None # type: ignore[assignment] _HAS_PYAUTOGUI = False -try: - import pygetwindow - - _HAS_PYGETWINDOW = True -except ImportError: - _HAS_PYGETWINDOW = False - -try: - import win32gui - - _HAS_WIN32GUI = True -except ImportError: - _HAS_WIN32GUI = False - try: import pyte from pyte.screens import Char @@ -56,78 +77,112 @@ _HAS_PYTE = False -def _get_terminal_window_rect() -> tuple[int, int, int, int] | None: - """Attempt to locate the terminal window and return its bounding box. +@dataclass(frozen=True) +class DisplayMonitor: + """A physical display and its virtual-desktop bounds.""" - Returns ``(left, top, right, bottom)`` or ``None`` if detection fails. - """ - if sys.platform != "win32": - return None + identifier: str + name: str + left: int + top: int + right: int + bottom: int + is_primary: bool = False - # Primary: use the current console window via ctypes. - try: - import ctypes + @property + def width(self) -> int: + return self.right - self.left - hwnd = ctypes.windll.kernel32.GetConsoleWindow() - if hwnd: - if _HAS_WIN32GUI: - left, top, right, bottom = win32gui.GetWindowRect(hwnd) - return (left, top, right, bottom) - # Fallback via ctypes if win32gui is missing. - rect = ctypes.wintypes.RECT() - ctypes.windll.user32.GetWindowRect(hwnd, ctypes.byref(rect)) - return (rect.left, rect.top, rect.right, rect.bottom) - except Exception as exc: - logger.debug(f"Failed to get console window rect: {exc}") + @property + def height(self) -> int: + return self.bottom - self.top - # Secondary: try pygetwindow with common terminal titles. - if _HAS_PYGETWINDOW: - for keyword in ("PowerShell", "Windows Terminal", "cmd", "Command Prompt"): - try: - windows = pygetwindow.getWindowsWithTitle(keyword) - if windows: - win = windows[0] - return (win.left, win.top, win.right, win.bottom) - except Exception as exc: - logger.debug(f"pygetwindow search failed for '{keyword}': {exc}") + @property + def bbox(self) -> tuple[int, int, int, int]: + return (self.left, self.top, self.right, self.bottom) - # Tertiary: enumerate visible windows with win32gui. - if _HAS_WIN32GUI: - try: - results: list[tuple[int, int, int, int]] = [] - def _enum_handler(hwnd: Any, _: Any) -> None: - if not win32gui.IsWindowVisible(hwnd): - return - title = win32gui.GetWindowText(hwnd) - if any(k in title for k in ("PowerShell", "Windows Terminal", "cmd")): - results.append(win32gui.GetWindowRect(hwnd)) +def list_display_monitors() -> list[DisplayMonitor]: + """Return connected Windows displays in stable desktop order.""" + if sys.platform != "win32": + return [] + _enable_windows_dpi_awareness() - win32gui.EnumWindows(_enum_handler, None) - if results: - return results[0] - except Exception as exc: - logger.debug(f"win32gui enumeration failed: {exc}") + try: + import ctypes + from ctypes import wintypes + + class MonitorInfoEx(ctypes.Structure): + _fields_ = [ + ("cbSize", wintypes.DWORD), + ("rcMonitor", wintypes.RECT), + ("rcWork", wintypes.RECT), + ("dwFlags", wintypes.DWORD), + ("szDevice", wintypes.WCHAR * 32), + ] + + class DisplayDevice(ctypes.Structure): + _fields_ = [ + ("cb", wintypes.DWORD), + ("DeviceName", wintypes.WCHAR * 32), + ("DeviceString", wintypes.WCHAR * 128), + ("StateFlags", wintypes.DWORD), + ("DeviceID", wintypes.WCHAR * 128), + ("DeviceKey", wintypes.WCHAR * 128), + ] + + monitors: list[DisplayMonitor] = [] + callback_type = ctypes.WINFUNCTYPE( + wintypes.BOOL, + wintypes.HANDLE, + wintypes.HDC, + ctypes.POINTER(wintypes.RECT), + wintypes.LPARAM, + ) - return None + @callback_type + def collect_monitor(handle: Any, _dc: Any, _rect: Any, _data: Any) -> bool: + info = MonitorInfoEx() + info.cbSize = ctypes.sizeof(info) + if not ctypes.windll.user32.GetMonitorInfoW(handle, ctypes.byref(info)): + return True + + raw_name = str(info.szDevice) or f"DISPLAY{len(monitors) + 1}" + short_name = raw_name.removeprefix(chr(92) * 2 + "." + chr(92)) + identifier = re.sub(r"[^A-Za-z0-9_-]", "_", short_name)[:32] + device = DisplayDevice() + device.cb = ctypes.sizeof(device) + friendly_name = "" + if ctypes.windll.user32.EnumDisplayDevicesW(raw_name, 0, ctypes.byref(device), 0): + friendly_name = str(device.DeviceString).strip() + name = f"{friendly_name} ({short_name})" if friendly_name else short_name + rect = info.rcMonitor + monitors.append( + DisplayMonitor( + identifier=identifier, + name=name, + left=rect.left, + top=rect.top, + right=rect.right, + bottom=rect.bottom, + is_primary=bool(info.dwFlags & 1), + ) + ) + return True + if not ctypes.windll.user32.EnumDisplayMonitors(None, None, collect_monitor, 0): + return [] + return sorted(monitors, key=lambda monitor: (not monitor.is_primary, monitor.top, monitor.left)) + except Exception as exc: + logger.warning(f"Failed to enumerate desktop monitors: {exc}") + return [] -async def capture_terminal_screenshot() -> bytes | None: - """Capture the current terminal window as a PNG image. - Returns PNG bytes or ``None`` if capture fails. - """ +async def capture_desktop_screenshot(monitor: DisplayMonitor | None = None) -> bytes | None: + """Capture one desktop monitor as PNG, or the default display if omitted.""" + _enable_windows_dpi_awareness() loop = asyncio.get_running_loop() - def _valid_bbox(bbox: tuple[int, int, int, int] | None) -> tuple[int, int, int, int] | None: - if bbox is None: - return None - left, top, right, bottom = bbox - if right <= left or bottom <= top: - logger.warning(f"Ignoring invalid screenshot bbox: {bbox}") - return None - return bbox - def _save_png(img: Any) -> bytes | None: width, height = getattr(img, "size", (0, 0)) if width <= 0 or height <= 0: @@ -144,41 +199,48 @@ def _capture() -> bytes | None: logger.warning("No screenshot library available. Install Pillow or pyautogui.") return None - bbox = _valid_bbox(_get_terminal_window_rect()) - attempts: list[tuple[str, tuple[int, int, int, int] | None]] = [] - if bbox is not None: - attempts.append(("terminal-window", bbox)) - attempts.append(("full-screen", None)) + bbox = monitor.bbox if monitor is not None else None + if bbox is not None and (monitor.width <= 0 or monitor.height <= 0): + logger.warning(f"Ignoring invalid monitor bounds: {bbox}") + return None if _HAS_PIL: - for label, grab_bbox in attempts: - try: - img = ImageGrab.grab(bbox=grab_bbox) if grab_bbox else ImageGrab.grab() - data = _save_png(img) - if data: - return data - except Exception as exc: - logger.warning(f"Pillow screenshot capture failed ({label}): {exc}") + try: + if bbox is None: + img = ImageGrab.grab() + elif sys.platform == "win32": + img = ImageGrab.grab(bbox=bbox, all_screens=True) + else: + img = ImageGrab.grab(bbox=bbox) + data = _save_png(img) + if data: + return data + except Exception as exc: + logger.warning(f"Pillow desktop capture failed: {exc}") if _HAS_PYAUTOGUI: - for label, grab_bbox in attempts: - try: - if grab_bbox: - left, top, right, bottom = grab_bbox - img = pyautogui.screenshot(region=(left, top, right - left, bottom - top)) - else: - img = pyautogui.screenshot() - data = _save_png(img) - if data: - return data - except Exception as exc: - logger.warning(f"pyautogui screenshot capture failed ({label}): {exc}") + try: + if bbox: + left, top, right, bottom = bbox + img = pyautogui.screenshot(region=(left, top, right - left, bottom - top)) + else: + img = pyautogui.screenshot() + data = _save_png(img) + if data: + return data + except Exception as exc: + logger.warning(f"pyautogui desktop capture failed: {exc}") return None return await loop.run_in_executor(None, _capture) +async def capture_terminal_screenshot() -> bytes | None: + """Capture the default display for compatibility with older callers.""" + return await capture_desktop_screenshot() + + # --------------------------------------------------------------------------- # ANSI color helpers # --------------------------------------------------------------------------- @@ -385,32 +447,3 @@ def ansi_to_html(ansi_text: str) -> str: return ( f'
{text}
' ) - - -async def send_screenshot_to_chat( - message: Message, - route: TelegramRoute, -) -> None: - """Capture a screenshot and send it to the chat.""" - png_bytes = await capture_terminal_screenshot() - if png_bytes is None: - await message.answer( - "Failed to capture screenshot. Pillow/pyautogui may be missing, " - "or the current desktop/window returned an empty image. " - "Try focusing or unminimizing the terminal and retry.", - **route.send_kwargs(), - ) - return - - try: - await message.answer_photo( - photo=BufferedInputFile(png_bytes, filename="screenshot.png"), - caption="Terminal screenshot", - **route.send_kwargs(), - ) - except Exception as exc: - logger.warning(f"Failed to send screenshot: {exc}") - await message.answer( - f"Failed to send screenshot: {exc}", - **route.send_kwargs(), - )