From 56eb784625589783ba3404cf0ec7cd8ac902db48 Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Fri, 2 Oct 2026 23:03:28 +0800 Subject: [PATCH 1/9] feat: add desktop control takeover prototype Track physical user activity, gate MCP tool calls, and show visible AI ownership. Document the asynchronous release and in-flight command limits pending further safety work. --- README.md | 1 + manifest.json | 4 + src/windows_mcp/__main__.py | 92 ++- src/windows_mcp/desktop/control.py | 581 +++++++++++++++++ src/windows_mcp/desktop/control_context.py | 26 + src/windows_mcp/desktop/control_ledger.py | 82 +++ src/windows_mcp/desktop/control_overlay.py | 431 +++++++++++++ src/windows_mcp/desktop/control_win32.py | 314 ++++++++++ src/windows_mcp/desktop/flash_overlay.py | 61 +- src/windows_mcp/desktop/service.py | 187 ++++-- src/windows_mcp/infrastructure/config.py | 25 + src/windows_mcp/tools/__init__.py | 2 + .../tools/control_notifications.py | 251 ++++++++ src/windows_mcp/tools/control_status.py | 21 + tests/test_control.py | 584 ++++++++++++++++++ tests/test_control_config.py | 32 + tests/test_control_mcp_gate.py | 516 ++++++++++++++++ tests/test_control_overlay.py | 380 ++++++++++++ tests/test_deterministic_drag_desktop.py | 37 +- tests/test_flash_overlay.py | 32 + tests/test_stdio_handshake.py | 7 +- 21 files changed, 3598 insertions(+), 68 deletions(-) create mode 100644 src/windows_mcp/desktop/control.py create mode 100644 src/windows_mcp/desktop/control_context.py create mode 100644 src/windows_mcp/desktop/control_ledger.py create mode 100644 src/windows_mcp/desktop/control_overlay.py create mode 100644 src/windows_mcp/desktop/control_win32.py create mode 100644 src/windows_mcp/tools/control_notifications.py create mode 100644 src/windows_mcp/tools/control_status.py create mode 100644 tests/test_control.py create mode 100644 tests/test_control_config.py create mode 100644 tests/test_control_mcp_gate.py create mode 100644 tests/test_control_overlay.py diff --git a/README.md b/README.md index d9f65af9..01423cba 100755 --- a/README.md +++ b/README.md @@ -719,6 +719,7 @@ MCP Client can access the following tools to interact with Windows: - `Shortcut`: Press keyboard shortcuts (`Ctrl+c`, `Alt+Tab`, etc). - `Wait`: Pause for a defined duration. - `WaitFor`: Wait until text, an active window, an element, or a focused element appears by polling UI state inside one tool call. +- `ControlStatus`: Read the desktop control owner and remaining user-idle wait. It stays available when user takeover blocks new tool calls. During AI control, a visible edge and cursor glow indicate ownership; moving the physical mouse substantially or pressing `Ctrl+Alt+Shift+Backspace` returns control to the user until 10 seconds of inactivity. Stepwise desktop input stops at its next checkpoint, but an already-running external command (for example `PowerShell` or an `App` launch) may continue after takeover; this control gate does not cancel processes or undo completed side effects. On a hook fault, injected keys and mouse buttons are released asynchronously, so physical input can briefly overlap a still-held AI key or button. - `DisplayInventory`: Read display layout, work areas, effective DPI, and scale metadata. - `Screenshot`: Fast screenshot-first desktop capture with cursor position, active/open windows, and an image. Skips UI tree extraction for speed and should be the default first call when you mainly need visual context. Supports `display=[0]` or `display=[0,1]` using zero-based active Windows display indices, and `region=[left, top, right, bottom]` (virtual-desktop pixel coordinates) to capture just that rectangle instead of the whole screen — cheaper on tokens when you already know which area matters. `region` takes precedence over `display` when both are given; an invalid or out-of-bounds region raises an error. After capture, a brief orange-red glowing border is drawn inside the captured area as a visual confirmation (set `WINDOWS_MCP_DISABLE_FLASH=1` to disable). - `Snapshot`: Full desktop state capture for workflows that need interactive element ids, scrollable regions, or `use_dom=True` browser extraction. Supports `use_vision=True` for including screenshots, `display=[0]` or `display=[0,1]` using zero-based active Windows display indices, and `region=[left, top, right, bottom]` (virtual-desktop pixel coordinates) to inspect just that rectangle instead of the whole screen; `region` takes precedence over `display` when both are given, and an invalid or out-of-bounds region raises an error. diff --git a/manifest.json b/manifest.json index ee16861f..ac608fd4 100755 --- a/manifest.json +++ b/manifest.json @@ -79,6 +79,10 @@ } }, "tools": [ + { + "name": "ControlStatus", + "description": "Reports whether AI or the user currently controls the desktop and the remaining user-idle wait before AI control can resume. This read-only status remains available during user takeover." + }, { "name": "App", "description": "Manages Windows applications with three modes: 'launch' (opens the prescibed application), 'resize' (adjusts the size/position of a named window or the active window if name is omitted), 'switch' (brings specific window into focus)." diff --git a/src/windows_mcp/__main__.py b/src/windows_mcp/__main__.py index b2d06a92..3db34477 100755 --- a/src/windows_mcp/__main__.py +++ b/src/windows_mcp/__main__.py @@ -1,5 +1,6 @@ from contextlib import asynccontextmanager from windows_mcp.config import enable_debug +from windows_mcp.infrastructure.config import ControlConfig from windows_mcp.infrastructure import ( AuthKeyMiddleware, OAuthOnlyMiddleware, @@ -29,6 +30,7 @@ import shlex import secrets import subprocess +import threading import click import os import sys @@ -49,11 +51,14 @@ def _echo_section(title: str) -> None: Setting PYTHONIOENCODING=utf-8 is a workaround; this is the fix. """ enc = getattr(sys.stdout, "encoding", None) or "ascii" + # Older Ruff releases drop the required parentheses around this exception tuple. + # fmt: off try: "─".encode(enc) bar = "─" * 3 except (UnicodeEncodeError, LookupError): bar = "===" + # fmt: on click.echo(f"\n{bar} {title} {bar}") @@ -62,6 +67,8 @@ def _echo_section(title: str) -> None: analytics: Any | None = None screen_size: Any | None = None _mcp: FastMCP | None = None +_control_config = ControlConfig() +_control_notifier: Any | None = None instructions = dedent(""" Windows MCP server provides tools to interact directly with the Windows desktop, @@ -240,7 +247,7 @@ def _start_watchdog(desktop): def _build_mcp() -> FastMCP: """Create the MCP server instance.""" - global _mcp + global _mcp, _control_notifier if _mcp is not None: return _mcp @@ -251,11 +258,57 @@ def _build_mcp() -> FastMCP: from windows_mcp.tools import register_all except ModuleNotFoundError as exc: _exit_missing_dependency(exc) + from windows_mcp.tools.control_notifications import ControlNotifier, ControlToolGate + from windows_mcp.desktop.control import get_controller + from windows_mcp.desktop import control_overlay + + controller = get_controller() + notifier = ControlNotifier(controller) + _control_notifier = notifier + control_loop: asyncio.AbstractEventLoop | None = None + + def apply_control_state(status: dict) -> None: + try: + # Queued input-thread updates can arrive after a later takeover. + # Never let an old AI event re-show the indicator in user mode. + if status.get("generation") is not None and status[ + "generation" + ] != controller.status().get("generation"): + return + if status["state"] == "ai": + control_overlay.set_active(True) + elif status["state"] == "takeover_pending": + control_overlay.set_pending(True) + else: + control_overlay.set_active(False) + except Exception: + # A missing visual indicator invalidates the AI control lease. + controller._fail_open() + logger.exception("AI control indicator failed; releasing physical input") + + def show_control_state(status: dict) -> None: + if threading.current_thread() is controller._thread: + # Never wait for a window-owner acknowledgement on the hook thread. + # Windows can silently remove a low level hook after a slow callback. + try: + if control_loop is None or control_loop.is_closed(): + raise RuntimeError("server event loop unavailable") + control_loop.call_soon_threadsafe(apply_control_state, status.copy()) + except RuntimeError: + controller._fail_open() + logger.exception("Could not schedule AI control indicator update") + else: + # begin_call waits for the visible AI indicator before any tool runs. + apply_control_state(status) + + controller.subscribe(show_control_state) @asynccontextmanager async def lifespan(app: FastMCP): """Runs initialization code before the server starts and cleanup code after it shuts down.""" global desktop, watchdog, analytics, screen_size + nonlocal control_loop + control_loop = asyncio.get_running_loop() if os.getenv("ANONYMIZED_TELEMETRY", "true").lower() != "false": analytics = PostHogAnalytics() @@ -264,17 +317,34 @@ async def lifespan(app: FastMCP): watchdog = _start_watchdog(desktop) + # Set thresholds before installing input hooks; invalid TOML was rejected at load. + controller.mouse_takeover_units = _control_config.mouse_takeover_units + controller.mouse_takeover_pixels = _control_config.mouse_takeover_pixels + try: + control_overlay.start() + controller.set_health_probe(control_overlay.is_healthy) + controller.start() + notifier.start() logger.debug("Server started, entering main loop") yield finally: - logger.debug("Shutting down: stopping watchdog and analytics") - if watchdog: - watchdog.stop() - if analytics: - await analytics.close() + await notifier.close() + try: + controller.stop() + finally: + try: + control_overlay.stop() + finally: + logger.debug("Shutting down: stopping watchdog and analytics") + if watchdog: + watchdog.stop() + if analytics: + await analytics.close() + control_loop = None _mcp = FastMCP(name="windows-mcp", instructions=instructions, lifespan=lifespan) + _mcp.add_middleware(ControlToolGate(controller, notifier)) register_all(_mcp, get_desktop=_get_desktop, get_analytics=_get_analytics) return _mcp @@ -385,11 +455,12 @@ def _remove(name): if explicit_tools: keep = {t for t in explicit_tools if t in registered} + keep.add("ControlStatus") for name in registered - keep: _remove(name) elif exclude_tools: for name in exclude_tools: - if name in registered: + if name in registered and name != "ControlStatus": _remove(name) logger.debug("Tool filter applied: explicit=%s exclude=%s", explicit_tools, exclude_tools) @@ -410,6 +481,11 @@ def _run_server( stateless_http: bool = False, ) -> None: mcp = _build_mcp() + if _control_notifier is not None: + # stdio has one session until process exit; stateless HTTP has no + # durable session to which a later ownership event can be pushed. + _control_notifier.ttl = float("inf") if transport == "stdio" else 60.0 + _control_notifier.push_enabled = not (transport == "streamable-http" and stateless_http) if explicit_tools or exclude_tools: _apply_tool_filter(mcp, explicit_tools, exclude_tools) match transport: @@ -588,6 +664,7 @@ def serve( oauth_client_secret, stateless_http, ): + global _control_config asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) install_selfpipe_guard() if transport == Transport.STDIO.value: @@ -604,6 +681,7 @@ def serve( cfg = load_config(config_path) except (FileNotFoundError, ValueError) as exc: raise click.ClickException(str(exc)) + _control_config = cfg.control transport = _choose_value(ctx, "transport", transport, cfg.server.transport, "stdio") host = _choose_value(ctx, "host", host, cfg.server.host, "localhost") diff --git a/src/windows_mcp/desktop/control.py b/src/windows_mcp/desktop/control.py new file mode 100644 index 00000000..bbe1d06a --- /dev/null +++ b/src/windows_mcp/desktop/control.py @@ -0,0 +1,581 @@ +"""Local desktop ownership and physical input interception. + +The low level hooks never wait for the coordinator. A stale control heartbeat +or a failed event queue releases physical input in the hook itself. +""" + +from __future__ import annotations + +import ctypes +import math +import queue +import threading +import time +from typing import Callable +from windows_mcp.desktop.control_context import current_token, current_steps, get_step_count +from windows_mcp.desktop.control_ledger import InputLedger +from windows_mcp.desktop.control_win32 import ( + _CHORD, + _MouseHookData, + _KeyHookData, + _user32, + _key, + _mouse_button, + _read_raw_mouse, + _interactive_desktop, + run_input_monitor, +) + + +class ControlBlocked(RuntimeError): + def __init__(self, code: str, status: dict): + self.code = code + self.status = status + super().__init__(f"{code}: {status['state']}") + + +class ControlCoordinator: + def __init__(self, mouse_takeover_units: int = 120, mouse_takeover_pixels: int = 120): + self.mouse_takeover_units = mouse_takeover_units + self.mouse_takeover_pixels = mouse_takeover_pixels + self._lock = threading.Lock() + self._state = "unavailable" + self._generation = 0 + self._active_calls = 0 + self._lease_until = 0.0 + self._last_user = 0.0 + self._last_physical_event = 0.0 # Written before hook events enter the queue. + self._last_move = 0.0 + self._point_origin: tuple[int, int] | None = None + self._raw_origin: dict[int, tuple[int, int]] = {} + self._raw_device: int | None = None + self._listeners: list[Callable[[dict], None]] = [] + self.input_ledger = InputLedger() + self._health_probe: Callable[[], bool] | None = None + self._events: queue.Queue[tuple] = queue.Queue(maxsize=4096) + self._stop = threading.Event() + self._release_requested = threading.Event() + self._ready = threading.Event() + self._thread: threading.Thread | None = None + self._watchdog: threading.Thread | None = None + self._startup_error: Exception | None = None + self._thread_beat = 0.0 + self._suppress = False # Atomic snapshot read by the hook callbacks. + self._deadline = 0.0 + self._emergency = False + self._fast_takeover = False + self._fast_pending = False + self._rotating = False + self._pressed: set[int] = set() + self._mouse_down: set[int] = set() + self._delivered_keys: set[int] = set() + self._delivered_mouse: set[int] = set() + self._quarantine: set[int] = set() + self._mouse_hook = None + self._key_hook = None + self._mouse_callback = None + self._key_callback = None + self._window_callback = None + self._rotate = False + + def subscribe(self, callback: Callable[[dict], None]) -> None: + with self._lock: + self._listeners.append(callback) + + def set_health_probe(self, probe: Callable[[], bool]) -> None: + """Require the visual indicator to stay alive during AI control.""" + self._health_probe = probe + + def _snapshot_locked(self, now: float) -> dict: + state = ( + "user" + if self._fast_takeover + else "takeover_pending" + if self._fast_pending and self._state == "ai" + else self._state + ) + wait = ( + max(0.0, max(self._last_user, self._last_physical_event) + 10.0 - now) + if state == "user" + else 0.0 + ) + return {"state": state, "generation": self._generation, "wait_seconds": wait} + + def status(self) -> dict: + if self._emergency: + self._mark_unavailable() + now = time.monotonic() + with self._lock: + changed = self._set_locked("unavailable") if self._emergency else self._tick_locked(now) + result = self._snapshot_locked(now) + if changed: + self._notify(result) + return result + + def _notify(self, result: dict) -> None: + if result["state"] != "ai": + # Serialize hold cleanup against a new lease, before any listener + # can wait on a visual effect or an MCP transport. + with self._lock: + if result["generation"] != self._generation: + return + try: + self.input_ledger.release_all() + except Exception: + self._fail_open() + self._set_locked("unavailable") + result = self._snapshot_locked(time.monotonic()) + for listener in tuple(self._listeners): + try: + listener(result) + except Exception: + pass # A disconnected MCP session must not hold desktop control. + + def _set_locked(self, state: str) -> bool: + if state == self._state: + return False + self._state = state + if state != "ai": + self.input_ledger.block_new() + self._generation += 1 + self._point_origin = None + self._raw_origin.clear() + self._raw_device = None + if state != "takeover_pending": + self._fast_pending = False + self._suppress = state in ("ai", "takeover_pending") and not self._emergency + self._deadline = time.monotonic() + 1.0 if self._suppress else 0.0 + return True + + def _tick_locked(self, now: float) -> bool: + if self._state == "ready" and now - self._last_physical_event < 10.0: + self._last_user = self._last_physical_event + return self._set_locked("user") + if ( + self._state == "user" + and not self._pressed + and not self._mouse_down + and now - max(self._last_user, self._last_physical_event) >= 10.0 + ): + return self._set_locked("ready") + if ( + self._state == "takeover_pending" + and not self._mouse_down + and now - self._last_move >= 0.3 + ): + return self._set_locked( + "ai" if self._active_calls or now < self._lease_until else "ready" + ) + if self._state == "ai" and not self._active_calls and now >= self._lease_until: + if now - self._last_physical_event < 10.0: + self._last_user = self._last_physical_event + return self._set_locked("user") + return self._set_locked("ready") + return False + + def begin_call(self, name: str) -> int: + if self._emergency: + self._mark_unavailable() + now = time.monotonic() + with self._lock: + changed = self._set_locked("unavailable") if self._emergency else self._tick_locked(now) + if self._pressed or self._mouse_down: + raise ControlBlocked("PHYSICAL_INPUT_HELD", self._snapshot_locked(now)) + if ( + self._emergency + or self._fast_takeover + or self._fast_pending + or self._rotating + or self._state not in ("ready", "ai") + ): + status = self._snapshot_locked(now) + code = {"user": "USER_CONTROL", "takeover_pending": "TAKEOVER_PENDING"}.get( + status["state"], "CONTROL_UNAVAILABLE" + ) + raise ControlBlocked(code, status) + try: + self.input_ledger.enable() # Any failed release keeps the lease closed. + except RuntimeError: + self._fail_open() + self._set_locked("unavailable") + raise ControlBlocked("CONTROL_UNAVAILABLE", self._snapshot_locked(now)) from None + changed = self._set_locked("ai") or changed + self._active_calls += 1 + token = self._generation + result = self._snapshot_locked(now) + if changed: + self._notify(result) + return token + + def checkpoint(self, token: int) -> None: + with self._lock: + if ( + self._state != "ai" + or self._generation != token + or self._fast_takeover + or self._fast_pending + or self._rotating + or self._emergency + or not self._suppress + ): + status = self._snapshot_locked(time.monotonic()) + status["executed_steps"] = get_step_count() + raise ControlBlocked("CONTROL_PREEMPTED", status) + + def checkpoint_current(self) -> None: + """Validate each desktop input step inside a guarded MCP call.""" + token = current_token.get() + if token is not None: + self.checkpoint(token) + + def record_step_current(self) -> None: + counter = current_steps.get() + if current_token.get() is not None and counter is not None: + counter.value += 1 + + def physical_key_down(self, vk: int) -> bool: + # A swallowed physical down did not reach the app, so the AI must + # release its own injected down when its tool is preempted. + return _key(vk) in self._delivered_keys + + def physical_mouse_down(self, button: str) -> bool: + return {"left": 1, "right": 2, "middle": 3}[button] in self._delivered_mouse + + def end_call(self, token: int) -> None: + with self._lock: + self._active_calls = max(0, self._active_calls - 1) + if self._state in ("ai", "takeover_pending") and not self._active_calls: + self._lease_until = time.monotonic() + 15.0 + + def start(self) -> None: + if self._thread and self._thread.is_alive(): + if self._stop.is_set(): + raise RuntimeError("Previous physical input monitor has not stopped") + return + self._stop.clear() + self._release_requested.clear() + self._ready.clear() + self._startup_error = None + if not _interactive_desktop(): + raise RuntimeError("No active interactive desktop for physical input control") + self.input_ledger.release_all() + self.input_ledger.enable() + self._emergency = False + self._fast_takeover = False + self._fast_pending = False + self._rotating = False + self._thread = threading.Thread(target=self._run, name="windows-mcp-input", daemon=True) + self._thread.start() + if not self._ready.wait(3.0) or self._startup_error: + self.stop() + raise RuntimeError("Physical input monitor could not start") from self._startup_error + self._watchdog = threading.Thread( + target=self._watch, name="windows-mcp-input-watch", daemon=True + ) + self._watchdog.start() + with self._lock: + now = time.monotonic() + # A new server must not impose a cooldown without physical input. + recent_input = self._last_physical_event > 0 and now - self._last_physical_event < 10.0 + if recent_input: + self._last_user = self._last_physical_event + self._set_locked("user" if recent_input else "ready") + result = self._snapshot_locked(now) + self._notify(result) + + def stop(self) -> None: + self._suppress = False + self._deadline = 0.0 + self.input_ledger.block_new() + self._stop.set() + self._release_requested.set() # Wake the watchdog before joining it. + if self._thread: + self._thread.join(timeout=2.0) + if self._watchdog: + self._watchdog.join(timeout=1.0) + with self._lock: + changed = self._set_locked("unavailable") + result = self._snapshot_locked(time.monotonic()) + try: + self.input_ledger.release_all() + finally: + if changed: + self._notify(result) + if (self._thread and self._thread.is_alive()) or ( + self._watchdog and self._watchdog.is_alive() + ): + raise RuntimeError("Physical input monitor did not stop cleanly") + self._events = queue.Queue(maxsize=4096) + self._pressed.clear() + self._mouse_down.clear() + self._delivered_keys.clear() + self._delivered_mouse.clear() + self._quarantine.clear() + self._thread = self._watchdog = None + + def _fail_open(self) -> None: + self._suppress = False + self._deadline = 0.0 + self.input_ledger.block_new() # Atomic flag; hooks must never wait for AI injection. + self._emergency = True + self._rotate = True + # Hook callbacks must not wait on SendInput or the ledger lock. Wake + # the existing watchdog to clear AI holds immediately off-hook. + self._release_requested.set() + + def _mark_unavailable(self) -> None: + # Clear injected holds before any visual or MCP listener can delay us. + self.input_ledger.block_new() + try: + self.input_ledger.release_all() + except Exception: + pass # Keep failed holds for the next watchdog/recovery attempt. + if not self._lock.acquire(blocking=False): + return + try: + changed = self._set_locked("unavailable") + result = self._snapshot_locked(time.monotonic()) + finally: + self._lock.release() + if changed: + self._notify(result) + + def _recover_after_rehook(self) -> None: + """Reopen the lease only after hooks, display and held inputs are safe.""" + if not self._emergency or self._stop.is_set(): + return + try: + self.input_ledger.release_all() + if self.input_ledger.pending(): + return + if ( + self._health_probe is not None and not self._health_probe() + ) or not _interactive_desktop(): + return + with self._lock: + self.input_ledger.enable() + self._emergency = False + self._fast_takeover = False + self._fast_pending = False + self._last_user = time.monotonic() + changed = self._set_locked("user") + result = self._snapshot_locked(self._last_user) + except Exception: + return # A failed release or health check keeps control unavailable. + if changed: + self._notify(result) + + def _queue(self, event: tuple) -> None: + try: + self._events.put_nowait(event) + except queue.Full: + self._fail_open() + + def _physical_mouse(self, code, wparam, lparam): + if code < 0: + return _user32.CallNextHookEx(self._mouse_hook, code, wparam, lparam) + valid = False + button = None + try: + data = ctypes.cast(lparam, ctypes.POINTER(_MouseHookData)).contents + if data.flags & 1: # LL mouse injection flags. + return _user32.CallNextHookEx(self._mouse_hook, code, wparam, lparam) + valid = True + self._last_physical_event = time.monotonic() + if wparam == 0x200 and self._suppress: + self._fast_pending = True # Pause AI before queued movement is processed. + self.input_ledger.block_new() + button = _mouse_button(wparam, data.mouseData) + if button: + if button[1]: + self._mouse_down.add(button[0]) + else: + self._mouse_down.discard(button[0]) + self._queue(("point", data.pt.x, data.pt.y, int(wparam))) + if self._suppress and not self._emergency: + if time.monotonic() <= self._deadline: + return 1 + self._fail_open() + except Exception: + self._fail_open() + if valid and button: + if button[1]: + self._delivered_mouse.add(button[0]) + else: + self._delivered_mouse.discard(button[0]) + return _user32.CallNextHookEx(self._mouse_hook, code, wparam, lparam) + + def _physical_key(self, code, wparam, lparam): + if code < 0: + return _user32.CallNextHookEx(self._key_hook, code, wparam, lparam) + vk = None + down = False + try: + data = ctypes.cast(lparam, ctypes.POINTER(_KeyHookData)).contents + if data.flags & 0x10: # LLKHF_INJECTED: AI SendInput stays usable. + return _user32.CallNextHookEx(self._key_hook, code, wparam, lparam) + self._last_physical_event = time.monotonic() + vk = _key(data.vkCode) + down = wparam in (0x100, 0x104) + was_down = vk in self._pressed + if down: + self._pressed.add(vk) + else: + self._pressed.discard(vk) + if vk in self._quarantine: + if not down: + self._quarantine.discard(vk) + self._queue(("key",)) # Idle starts after the last chord key is released. + return 1 + if ( + down + and vk == 0x08 + and not was_down + and not self._fast_takeover + and self._suppress + and _CHORD.issubset(self._pressed) + ): + self._quarantine.update(self._pressed & _CHORD) + self._fast_takeover = True + self._suppress = False # Release first, then tell the coordinator. + self.input_ledger.block_new() + self._queue(("hotkey",)) + return 1 + self._queue(("key",)) + if self._suppress and not self._emergency: + if time.monotonic() <= self._deadline: + return 1 + self._fail_open() + except Exception: + self._fail_open() + if vk is not None and down: + self._delivered_keys.add(vk) + elif vk is not None: + self._delivered_keys.discard(vk) + return _user32.CallNextHookEx(self._key_hook, code, wparam, lparam) + + def _raw_input(self, lparam) -> None: + raw = _read_raw_mouse(lparam) + if raw is None: + return # Null devices can be touchpads; do not infer source. + self._last_physical_event = time.monotonic() + if self._suppress: + self._fast_pending = True + self.input_ledger.block_new() + self._queue(("raw", *raw)) + + def _handle(self, event: tuple) -> None: + now = time.monotonic() + with self._lock: + if self._emergency: + changed = self._set_locked("unavailable") + else: + changed = self._tick_locked(now) + kind = event[0] + if kind == "hotkey" and self._state in ("ai", "takeover_pending"): + self._last_user = now + changed = self._set_locked("user") or changed + self._fast_takeover = False + elif kind in ("key", "point", "raw") and self._state in ("ready", "user"): + self._last_user = now + changed = self._set_locked("user") or changed + elif kind == "point" and self._state in ("ai", "takeover_pending"): + if event[3] == 0x200: # WM_MOUSEMOVE + changed = self._candidate_locked(now) or changed + if self._point_origin is None: + self._point_origin = (event[1], event[2]) + ox, oy = self._point_origin + if math.hypot(event[1] - ox, event[2] - oy) >= self.mouse_takeover_pixels: + self._last_user = now + changed = self._set_locked("user") or changed + elif kind == "raw" and self._state in ("ai", "takeover_pending"): + changed = self._candidate_locked(now) or changed + device, dx, dy = event[1:] + if self._raw_device is not None and device != self._raw_device: + self._raw_origin.clear() # A new device starts a new segment. + self._raw_device = device + x, y = self._raw_origin.get(device, (0, 0)) + x, y = x + dx, y + dy + self._raw_origin[device] = (x, y) + if math.hypot(x, y) >= self.mouse_takeover_units: + self._last_user = now + changed = self._set_locked("user") or changed + result = self._snapshot_locked(now) + if changed: + self._notify(result) + + def _candidate_locked(self, now: float) -> bool: + if self._state == "ai": + changed = self._set_locked("takeover_pending") + else: + changed = False + self._last_move = now + return changed + + def _install_hooks(self, instance) -> None: + mouse = _user32.SetWindowsHookExW(14, self._mouse_callback, instance, 0) + key = _user32.SetWindowsHookExW(13, self._key_callback, instance, 0) + if not mouse or not key: + if mouse: + _user32.UnhookWindowsHookEx(mouse) + if key: + _user32.UnhookWindowsHookEx(key) + raise ctypes.WinError(ctypes.get_last_error()) + old_mouse, old_key = self._mouse_hook, self._key_hook + self._mouse_hook, self._key_hook = mouse, key + for old in (old_mouse, old_key): + if old: + _user32.UnhookWindowsHookEx(old) + + def _run(self) -> None: + # Keep the Win32 message pump alongside its raw-input structures. + run_input_monitor(self) + + def _watch(self) -> None: + while True: + self._release_requested.wait(0.25) + self._release_requested.clear() + if self._stop.is_set(): + return + # A hook fault wakes this thread specifically to release AI-held + # input. Do that before potentially slow display/WTS probes. + if self._emergency: + self._mark_unavailable() + continue + now = time.monotonic() + try: + indicator_ok = self._health_probe is None or self._health_probe() + desktop_ok = _interactive_desktop() + except Exception: + indicator_ok = False + desktop_ok = False + # The watchdog cannot renew suppression without a live coordinator + # proof: a wedged lock, stopped message loop or emergency fails open. + if ( + not self._thread + or not self._thread.is_alive() + or now - self._thread_beat > 1.0 + or not indicator_ok + or not desktop_ok + ): + self._fail_open() + self._mark_unavailable() + continue + if not self._lock.acquire(blocking=False): + self._fail_open() + continue + try: + if self._state in ("ai", "takeover_pending") and not self._fast_takeover: + self._deadline = now + 1.0 + self._suppress = True + else: + self._suppress = False + finally: + self._lock.release() + + +_controller = ControlCoordinator() + + +def get_controller() -> ControlCoordinator: + return _controller diff --git a/src/windows_mcp/desktop/control_context.py b/src/windows_mcp/desktop/control_context.py new file mode 100644 index 00000000..03295fb1 --- /dev/null +++ b/src/windows_mcp/desktop/control_context.py @@ -0,0 +1,26 @@ +"""Request-local desktop control token and cross-thread input-step counter.""" + +from __future__ import annotations + +import contextvars +from dataclasses import dataclass + + +@dataclass +class StepCounter: + # AnyIO copies context variables into worker threads. The value object is + # shared, so a synchronous tool's executed steps remain visible to middleware. + value: int = 0 + + +current_token: contextvars.ContextVar[int | None] = contextvars.ContextVar( + "windows_mcp_control_token", default=None +) +current_steps: contextvars.ContextVar[StepCounter | None] = contextvars.ContextVar( + "windows_mcp_control_steps", default=None +) + + +def get_step_count() -> int: + counter = current_steps.get() + return 0 if counter is None else counter.value diff --git a/src/windows_mcp/desktop/control_ledger.py b/src/windows_mcp/desktop/control_ledger.py new file mode 100644 index 00000000..4bc8c923 --- /dev/null +++ b/src/windows_mcp/desktop/control_ledger.py @@ -0,0 +1,82 @@ +"""Track AI-held keys and mouse buttons until each injected down is cleared.""" + +from __future__ import annotations + +import threading +from collections.abc import Callable + +from fastmcp.exceptions import ToolError + + +class InputUnavailable(ToolError): + """A physical takeover or fault blocked a new injected key/button down.""" + + +class InputLedger: + def __init__(self) -> None: + self._lock = threading.RLock() + self._holds: dict[str, tuple[Callable[[], None], Callable[[], bool]]] = {} + self._blocked = False + + def block_new(self) -> None: + self._blocked = True + + def enable(self) -> None: + with self._lock: + if self._holds: + raise RuntimeError("AI input remains held") + self._blocked = False + + def press( + self, + name: str, + send_down: Callable[[], None], + send_up: Callable[[], None], + physical_down_delivered: Callable[[], bool], + ) -> None: + """Record before injection: a send routine may inject and then raise.""" + with self._lock: + if self._blocked: + raise InputUnavailable("AI input is unavailable") + if name in self._holds: + raise RuntimeError(f"AI input already held: {name}") + self._holds[name] = (send_up, physical_down_delivered) + try: + # Keep the lock through injection so a watchdog cannot send + # up before the matching down has actually reached Windows. + send_down() + except BaseException: + try: + self.release(name) + except Exception: + pass # Keep the entry for shutdown's second attempt. + raise + + def release(self, name: str) -> None: + with self._lock: + entry = self._holds.get(name) + if entry is None: + return + send_up, physical_down_delivered = entry + # If a physical down reached the app after takeover, its own up + # will clear the state. Injecting another up would steal the key. + if not physical_down_delivered(): + send_up() + self._holds.pop(name, None) + + def release_all(self) -> None: + with self._lock: + names = tuple(self._holds) + first_error: Exception | None = None + for name in names: + try: + self.release(name) + except Exception as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + + def pending(self) -> tuple[str, ...]: + with self._lock: + return tuple(self._holds) diff --git a/src/windows_mcp/desktop/control_overlay.py b/src/windows_mcp/desktop/control_overlay.py new file mode 100644 index 00000000..b713df49 --- /dev/null +++ b/src/windows_mcp/desktop/control_overlay.py @@ -0,0 +1,431 @@ +"""Small, click-through Win32 indicators for an active AI desktop lease. + +The display thread owns every HWND. Calls from MCP threads change desired state +and wait for an acknowledgement, so a screenshot cannot race a visible frame. +""" + +from contextlib import contextmanager +import ctypes +import logging +import threading +import time +from ctypes import wintypes + +from PIL import Image, ImageDraw, ImageFilter + +from windows_mcp.desktop import flash_overlay +from windows_mcp import uia + +logger = logging.getLogger(__name__) + +_BORDER = 28 +_CURSOR_SIZE = 88 +_REFRESH_SECONDS = 0.05 +_SW_HIDE = 0 +_WDA_EXCLUDEFROMCAPTURE = 0x00000011 +_BLUE = (45, 145, 255) +_AMBER = (255, 170, 55) + +_user32 = ctypes.windll.user32 +_user32.GetCursorPos.argtypes = [ctypes.POINTER(flash_overlay._POINT)] +_user32.GetCursorPos.restype = wintypes.BOOL +_user32.SetWindowDisplayAffinity.argtypes = [wintypes.HWND, wintypes.DWORD] +_user32.SetWindowDisplayAffinity.restype = wintypes.BOOL +_user32.IsWindowVisible.argtypes = [wintypes.HWND] +_user32.IsWindowVisible.restype = wintypes.BOOL + + +def _edge_bitmap(width: int, height: int, side: str, color: tuple[int, int, int]) -> bytes: + """Pre-render one inward gradient strip; cost scales with perimeter.""" + image = Image.new("RGBA", (width, height)) + draw = ImageDraw.Draw(image) + for depth in range(min(_BORDER, height if side in ("top", "bottom") else width)): + alpha = int(190 * (1 - depth / _BORDER) ** 2) + if side == "top": + draw.line((0, depth, width - 1, depth), fill=(*color, alpha)) + elif side == "bottom": + y = height - 1 - depth + draw.line((0, y, width - 1, y), fill=(*color, alpha)) + elif side == "left": + draw.line((depth, 0, depth, height - 1), fill=(*color, alpha)) + else: + x = width - 1 - depth + draw.line((x, 0, x, height - 1), fill=(*color, alpha)) + return flash_overlay._premultiplied_bgra(image, 1.0) + + +def _cursor_bitmap(color: tuple[int, int, int]) -> bytes: + image = Image.new("RGBA", (_CURSOR_SIZE, _CURSOR_SIZE)) + draw = ImageDraw.Draw(image) + inset = 17 + box = (inset, inset, _CURSOR_SIZE - inset - 1, _CURSOR_SIZE - inset - 1) + draw.ellipse(box, outline=(*color, 210), width=5) + glow = image.filter(ImageFilter.GaussianBlur(9)) + return flash_overlay._premultiplied_bgra(Image.alpha_composite(glow, image), 1.0) + + +class _Layer: + def __init__(self, x: int, y: int, width: int, height: int, bgra: bytes, name: str): + self.x, self.y, self.width, self.height = x, y, width, height + self.class_name = f"WindowsMCPControl_{name}_{id(self):x}" + self.hwnd, self.instance = flash_overlay._create_layered_window( + self.class_name, x, y, width, height + ) + try: + flash_overlay._push_bitmap(self.hwnd, x, y, width, height, bgra) + if not _user32.SetWindowDisplayAffinity(self.hwnd, _WDA_EXCLUDEFROMCAPTURE): + # Own screenshots are protected by suspend_for_capture. Affinity + # varies by Windows compositor and only affects other capturers. + logger.warning("display affinity unavailable for AI control indicator") + except BaseException: + self.close() + raise + + def show(self) -> None: + _user32.ShowWindow(self.hwnd, flash_overlay._SW_SHOWNA) + if not _user32.SetWindowPos( + self.hwnd, + flash_overlay._HWND_TOPMOST, + 0, + 0, + 0, + 0, + flash_overlay._SWP_NOSIZE + | flash_overlay._SWP_NOMOVE + | flash_overlay._SWP_NOACTIVATE + | flash_overlay._SWP_SHOWWINDOW, + ) or not _user32.IsWindowVisible(self.hwnd): + raise RuntimeError("AI control indicator window could not be shown") + + def hide(self) -> None: + _user32.ShowWindow(self.hwnd, _SW_HIDE) + if _user32.IsWindowVisible(self.hwnd): + raise RuntimeError("AI control indicator window could not be hidden") + + def move(self, x: int, y: int) -> None: + if (x, y) == (self.x, self.y): + return + if not _user32.SetWindowPos( + self.hwnd, + flash_overlay._HWND_TOPMOST, + x, + y, + 0, + 0, + flash_overlay._SWP_NOSIZE | flash_overlay._SWP_NOACTIVATE, + ): + raise RuntimeError("AI control cursor indicator could not move") + self.x, self.y = x, y + + def close(self) -> None: + if self.hwnd: + _user32.DestroyWindow(self.hwnd) + _user32.UnregisterClassW(self.class_name, self.instance) + self.hwnd = None + + +def _monitor_rects() -> tuple[tuple[int, int, int, int], ...]: + return tuple((r.left, r.top, r.right, r.bottom) for r in uia.GetMonitorsRect()) + + +def _build_layers( + rects: tuple[tuple[int, int, int, int], ...], pending: bool +) -> tuple[list[_Layer], _Layer]: + if not rects: + raise RuntimeError("no display available for AI control indicator") + color = _AMBER if pending else _BLUE + layers: list[_Layer] = [] + try: + for index, (left, top, right, bottom) in enumerate(rects): + width, height = right - left, bottom - top + if width <= 0 or height <= 0: + raise RuntimeError("invalid display geometry for AI control indicator") + border = min(_BORDER, width // 2, height // 2) + if border <= 0: + raise RuntimeError("display too small for AI control indicator") + strips = ( + (left, top, width, border, "top"), + (left, bottom - border, width, border, "bottom"), + (left, top + border, border, height - 2 * border, "left"), + (right - border, top + border, border, height - 2 * border, "right"), + ) + for x, y, w, h, side in strips: + if w and h: + layers.append( + _Layer(x, y, w, h, _edge_bitmap(w, h, side, color), f"{index}_{side}") + ) + point = flash_overlay._POINT() + if not _user32.GetCursorPos(ctypes.byref(point)): + raise RuntimeError("cannot locate cursor for AI control indicator") + ring = _Layer( + point.x - _CURSOR_SIZE // 2, + point.y - _CURSOR_SIZE // 2, + _CURSOR_SIZE, + _CURSOR_SIZE, + _cursor_bitmap(color), + "cursor", + ) + return layers, ring + except BaseException: + for layer in reversed(layers): + layer.close() + raise + + +class _Indicator: + def __init__(self) -> None: + self.condition = threading.Condition() + self.thread = threading.Thread( + target=self._run, name="windows-mcp-control-indicator", daemon=True + ) + self.active = False + self.pending = False + self.suspended = 0 + self.stopping = False + self.version = 0 + self.applied = 0 + self.error: BaseException | None = None + self.heartbeat = 0.0 + self.started = threading.Event() + + def start(self) -> None: + self.thread.start() + if not self.started.wait(timeout=3.0): + raise RuntimeError("AI control indicator thread did not start") + if self.error or not self.thread.is_alive(): + raise RuntimeError("AI control indicator thread failed at startup") from self.error + + def change( + self, *, active: bool | None = None, pending: bool | None = None, suspended: int = 0 + ) -> None: + with self.condition: + if self.stopping: + raise RuntimeError("AI control indicator is stopping") + if self.error: + raise RuntimeError("AI control indicator failed") from self.error + if not self.thread.is_alive(): + raise RuntimeError("AI control indicator stopped") + if active is not None: + self.active = active + if pending is not None: + self.pending = pending + self.suspended += suspended + if self.suspended < 0: + self.suspended = 0 + raise RuntimeError("unbalanced screenshot indicator suspension") + self.version += 1 + version = self.version + self.condition.notify_all() + if not self.condition.wait_for( + lambda: ( + self.applied >= version + or self.stopping + or self.error is not None + or not self.thread.is_alive() + ), + timeout=3.0, + ): + # The owner may later recover and apply this hidden state. + # Make the failure sticky so health checks can never grant a + # new AI lease with an invisible indicator. + self.error = RuntimeError("AI control indicator did not acknowledge state change") + self.condition.notify_all() + raise self.error + if self.stopping: + raise RuntimeError("AI control indicator is stopping") + if self.error: + raise RuntimeError("AI control indicator failed") from self.error + + def stop(self) -> None: + with self.condition: + self.stopping = True + self.condition.notify_all() + if self.thread.is_alive(): + self.thread.join(timeout=3.0) + if self.thread.is_alive(): + raise RuntimeError("AI control indicator did not stop") + + def _run(self) -> None: + layers: list[_Layer] = [] + ring: _Layer | None = None + rects: tuple[tuple[int, int, int, int], ...] = () + mode: bool | None = None + visible = False + try: + self.heartbeat = time.monotonic() + self.started.set() + while True: + with self.condition: + self.condition.wait(timeout=_REFRESH_SECONDS) + if self.stopping: + return + active, pending = self.active, self.pending + suspended, version = self.suspended, self.version + if active: + current_rects = _monitor_rects() + if ring is None or current_rects != rects or mode != pending: + for layer in reversed(layers): + layer.close() + layers = [] + if ring: + ring.close() + ring = None + layers, ring = _build_layers(current_rects, pending) + rects, mode = current_rects, pending + visible = False + if suspended: + if visible: + for layer in layers: + layer.hide() + ring.hide() + visible = False + else: + if not visible: + for layer in layers: + layer.show() + ring.show() + visible = True + if not pending: + point = flash_overlay._POINT() + if _user32.GetCursorPos(ctypes.byref(point)): + ring.move(point.x - _CURSOR_SIZE // 2, point.y - _CURSOR_SIZE // 2) + else: + ring.hide() + elif visible: + for layer in layers: + layer.hide() + if ring: + ring.hide() + visible = False + # Pump messages on the owning thread, including display events. + for layer in layers: + flash_overlay._pump_messages(layer.hwnd) + if ring: + flash_overlay._pump_messages(ring.hwnd) + with self.condition: + self.heartbeat = time.monotonic() + self.applied = version + self.condition.notify_all() + except BaseException as exc: + logger.exception("AI control indicator failed") + with self.condition: + self.error = exc + self.condition.notify_all() + finally: + self.started.set() + for layer in reversed(layers): + layer.close() + if ring: + ring.close() + + +_lock = threading.Lock() +_capture_lock = threading.RLock() +_instance: _Indicator | None = None + + +def start() -> None: + """Start the window owner thread, without displaying AI control yet.""" + global _instance + with _capture_lock: + with _lock: + current = _instance + if current is None: + # Retain the reference before starting: a failed startup may + # leave a live thread that must be stopped before retrying. + indicator = _Indicator() + _instance = indicator + else: + indicator = current + if current is not None: + if not is_healthy() or indicator.stopping: + raise RuntimeError("Existing AI control indicator requires stop before restart") + return + try: + indicator.start() + except BaseException as startup_error: + try: + indicator.stop() + except BaseException as cleanup_error: + # A later stop() can retry the join; never stack a new HWND + # owner on top of a potentially live failed one. + raise RuntimeError( + "AI control indicator startup failed; cleanup is pending" + ) from cleanup_error + with _lock: + if _instance is indicator: + _instance = None + raise startup_error + + +def stop() -> None: + """Destroy all indication windows and join the owner thread.""" + global _instance + with _capture_lock: + with _lock: + indicator = _instance + if indicator: + indicator.stop() + # The thread has exited and destroyed its windows. Only now may a + # future start allocate another owner thread. + with _lock: + if _instance is indicator: + _instance = None + + +def set_active(active: bool) -> None: + """Display or hide AI indication; activation failure raises.""" + + # Deactivation must remain immediate during a capture: the screenshot's + # final resume will then observe active=False and leave the window hidden. + def apply() -> None: + with _lock: + indicator = _instance + if indicator is None: + if not active: + return + raise RuntimeError("AI control indicator has not started") + indicator.change(active=active, pending=False) + + if active: + with _capture_lock: + apply() + else: + apply() + + +def set_pending(pending: bool) -> None: + """Use amber edge indication without cursor following during takeover check.""" + with _capture_lock: + with _lock: + indicator = _instance + if indicator: + indicator.change(pending=pending) + + +def is_healthy() -> bool: + """Reject an alive but stalled window owner after one second.""" + with _lock: + indicator = _instance + return bool( + indicator + and indicator.thread.is_alive() + and not indicator.stopping + and indicator.error is None + and 0 <= time.monotonic() - indicator.heartbeat <= 1.0 + ) + + +@contextmanager +def suspend_for_capture(): + """Serialize captures and wait for the indicator to disappear first.""" + with _capture_lock: + with _lock: + indicator = _instance + if indicator: + indicator.change(suspended=1) + try: + yield + finally: + if indicator: + indicator.change(suspended=-1) diff --git a/src/windows_mcp/desktop/control_win32.py b/src/windows_mcp/desktop/control_win32.py new file mode 100644 index 00000000..ee99f91e --- /dev/null +++ b/src/windows_mcp/desktop/control_win32.py @@ -0,0 +1,314 @@ +"""Win32 structures, API bindings and desktop health for input ownership.""" + +import ctypes +import queue +import threading +import time +from ctypes import wintypes +from typing import Any + +_LRESULT = ctypes.c_ssize_t +_HOOKPROC = ctypes.WINFUNCTYPE(_LRESULT, ctypes.c_int, wintypes.WPARAM, wintypes.LPARAM) +_WNDPROC = ctypes.WINFUNCTYPE( + _LRESULT, wintypes.HWND, wintypes.UINT, wintypes.WPARAM, wintypes.LPARAM +) + + +class _WndClass(ctypes.Structure): + _fields_ = [ + ("cbSize", wintypes.UINT), + ("style", wintypes.UINT), + ("lpfnWndProc", _WNDPROC), + ("cbClsExtra", ctypes.c_int), + ("cbWndExtra", ctypes.c_int), + ("hInstance", wintypes.HINSTANCE), + ("hIcon", wintypes.HICON), + ("hCursor", wintypes.HANDLE), + ("hbrBackground", wintypes.HBRUSH), + ("lpszMenuName", wintypes.LPCWSTR), + ("lpszClassName", wintypes.LPCWSTR), + ("hIconSm", wintypes.HICON), + ] + + +class _MouseHookData(ctypes.Structure): + _fields_ = [ + ("pt", wintypes.POINT), + ("mouseData", wintypes.DWORD), + ("flags", wintypes.DWORD), + ("time", wintypes.DWORD), + ("dwExtraInfo", ctypes.c_size_t), + ] + + +class _KeyHookData(ctypes.Structure): + _fields_ = [ + ("vkCode", wintypes.DWORD), + ("scanCode", wintypes.DWORD), + ("flags", wintypes.DWORD), + ("time", wintypes.DWORD), + ("dwExtraInfo", ctypes.c_size_t), + ] + + +class _RawDevice(ctypes.Structure): + _fields_ = [ + ("usUsagePage", wintypes.USHORT), + ("usUsage", wintypes.USHORT), + ("dwFlags", wintypes.DWORD), + ("hwndTarget", wintypes.HWND), + ] + + +class _RawHeader(ctypes.Structure): + _fields_ = [ + ("dwType", wintypes.DWORD), + ("dwSize", wintypes.DWORD), + ("hDevice", wintypes.HANDLE), + ("wParam", wintypes.WPARAM), + ] + + +class _RawMouse(ctypes.Structure): + _fields_ = [ + ("usFlags", wintypes.USHORT), + ("usButtonFlags", wintypes.USHORT), + ("usButtonData", wintypes.USHORT), + ("ulRawButtons", wintypes.DWORD), + ("lLastX", ctypes.c_long), + ("lLastY", ctypes.c_long), + ("ulExtraInformation", wintypes.DWORD), + ] + + +_user32 = ctypes.WinDLL("user32", use_last_error=True) +_kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) +_user32.SetWindowsHookExW.argtypes = [ctypes.c_int, _HOOKPROC, wintypes.HINSTANCE, wintypes.DWORD] +_user32.SetWindowsHookExW.restype = wintypes.HHOOK +_user32.CallNextHookEx.argtypes = [wintypes.HHOOK, ctypes.c_int, wintypes.WPARAM, wintypes.LPARAM] +_user32.CallNextHookEx.restype = _LRESULT +_user32.UnhookWindowsHookEx.argtypes = [wintypes.HHOOK] +_user32.RegisterClassExW.argtypes = [ctypes.POINTER(_WndClass)] +_user32.RegisterClassExW.restype = ctypes.c_ushort +_user32.CreateWindowExW.argtypes = [ + wintypes.DWORD, + wintypes.LPCWSTR, + wintypes.LPCWSTR, + wintypes.DWORD, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + wintypes.HWND, + wintypes.HANDLE, + wintypes.HINSTANCE, + ctypes.c_void_p, +] +_user32.CreateWindowExW.restype = wintypes.HWND +_user32.DestroyWindow.argtypes = [wintypes.HWND] +_user32.UnregisterClassW.argtypes = [wintypes.LPCWSTR, wintypes.HINSTANCE] +_user32.OpenInputDesktop.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] +_user32.OpenInputDesktop.restype = wintypes.HANDLE +_user32.CloseDesktop.argtypes = [wintypes.HANDLE] +_user32.GetUserObjectInformationW.argtypes = [ + wintypes.HANDLE, + ctypes.c_int, + ctypes.c_void_p, + wintypes.DWORD, + ctypes.POINTER(wintypes.DWORD), +] +_user32.GetUserObjectInformationW.restype = wintypes.BOOL +_user32.RegisterRawInputDevices.argtypes = [ + ctypes.POINTER(_RawDevice), + wintypes.UINT, + wintypes.UINT, +] +_user32.GetRawInputData.argtypes = [ + wintypes.HANDLE, + wintypes.UINT, + ctypes.c_void_p, + ctypes.POINTER(wintypes.UINT), + wintypes.UINT, +] +_user32.GetRawInputData.restype = wintypes.UINT +_user32.DefWindowProcW.argtypes = [wintypes.HWND, wintypes.UINT, wintypes.WPARAM, wintypes.LPARAM] +_user32.DefWindowProcW.restype = _LRESULT +_kernel32.GetModuleHandleW.argtypes = [wintypes.LPCWSTR] +_kernel32.GetModuleHandleW.restype = wintypes.HINSTANCE +_wts = ctypes.WinDLL("wtsapi32", use_last_error=True) +_wts.WTSQuerySessionInformationW.argtypes = [ + wintypes.HANDLE, + wintypes.DWORD, + ctypes.c_int, + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(wintypes.DWORD), +] +_wts.WTSQuerySessionInformationW.restype = wintypes.BOOL +_wts.WTSFreeMemory.argtypes = [ctypes.c_void_p] + +_CHORD = {0x10, 0x11, 0x12, 0x08} + + +def _key(vk: int) -> int: + if vk in (0xA0, 0xA1): + return 0x10 + if vk in (0xA2, 0xA3): + return 0x11 + if vk in (0xA4, 0xA5): + return 0x12 + return vk + + +def _mouse_button(message: int, data: int) -> tuple[int, bool] | None: + down = {0x201: 1, 0x204: 2, 0x207: 3} + up = {0x202: 1, 0x205: 2, 0x208: 3} + if message in down: + return down[message], True + if message in up: + return up[message], False + if message in (0x20B, 0x20C): # WM_XBUTTONDOWN / WM_XBUTTONUP. + xbutton = (data >> 16) & 0xFFFF + if xbutton in (1, 2): + return xbutton + 3, message == 0x20B + return None + + +def _read_raw_mouse(lparam: int) -> tuple[int, int, int] | None: + """Read a relative physical mouse event; null devices are ambiguous.""" + size = wintypes.UINT() + header_size = ctypes.sizeof(_RawHeader) + if ( + _user32.GetRawInputData(lparam, 0x10000003, None, ctypes.byref(size), header_size) + == 0xFFFFFFFF + ): + return None + if size.value < header_size + ctypes.sizeof(_RawMouse): + return None + buffer = ctypes.create_string_buffer(size.value) + if ( + _user32.GetRawInputData(lparam, 0x10000003, buffer, ctypes.byref(size), header_size) + == 0xFFFFFFFF + ): + return None + header = ctypes.cast(buffer, ctypes.POINTER(_RawHeader)).contents + if header.dwType != 0 or not header.hDevice: + return None + mouse = _RawMouse.from_buffer_copy(buffer, header_size) + if mouse.usFlags & 1: + return None + return int(header.hDevice), mouse.lLastX, mouse.lLastY + + +def _interactive_desktop() -> bool: + """Require an active session on the normal user input desktop.""" + buffer = ctypes.c_void_p() + length = wintypes.DWORD() + # WTS_CURRENT_SESSION, WTSConnectState=8; WTSActive=0. + if not _wts.WTSQuerySessionInformationW( + None, 0xFFFFFFFF, 8, ctypes.byref(buffer), ctypes.byref(length) + ): + return False + try: + if length.value < ctypes.sizeof(ctypes.c_int): + return False + if ctypes.cast(buffer, ctypes.POINTER(ctypes.c_int)).contents.value != 0: + return False + finally: + _wts.WTSFreeMemory(buffer) + # The input desktop changes to Winlogon for UAC/lock screen. + desktop = _user32.OpenInputDesktop(0, False, 0x0001) # DESKTOP_READOBJECTS. + if not desktop: + return False + try: + name = ctypes.create_unicode_buffer(64) + needed = wintypes.DWORD() + return ( + bool( + _user32.GetUserObjectInformationW( + desktop, 2, name, ctypes.sizeof(name), ctypes.byref(needed) + ) + ) + and name.value.casefold() == "default" + ) + finally: + _user32.CloseDesktop(desktop) + + +def run_input_monitor(owner: Any) -> None: + """Pump the Win32 hooks and raw input for one control coordinator.""" + name = f"WindowsMCPControl{threading.get_ident()}" + instance = _kernel32.GetModuleHandleW(None) + owner._mouse_callback = _HOOKPROC(owner._physical_mouse) + owner._key_callback = _HOOKPROC(owner._physical_key) + + @_WNDPROC + def window_proc(hwnd, message, wparam, lparam): + if message == 0x00FF: # WM_INPUT + try: + owner._raw_input(lparam) + except Exception: + owner._fail_open() + return _user32.DefWindowProcW(hwnd, message, wparam, lparam) + + owner._window_callback = window_proc + wc = _WndClass() + wc.cbSize = ctypes.sizeof(wc) + wc.lpfnWndProc = window_proc + wc.hInstance = instance + wc.lpszClassName = name + atom = _user32.RegisterClassExW(ctypes.byref(wc)) + hwnd = None + registered = False + try: + if not atom: + raise ctypes.WinError(ctypes.get_last_error()) + hwnd = _user32.CreateWindowExW(0, name, name, 0, 0, 0, 0, 0, None, None, instance, None) + if not hwnd: + raise ctypes.WinError(ctypes.get_last_error()) + device = _RawDevice(1, 2, 0x100, hwnd) # Mouse, RIDEV_INPUTSINK. + if not _user32.RegisterRawInputDevices(ctypes.byref(device), 1, ctypes.sizeof(device)): + raise ctypes.WinError(ctypes.get_last_error()) + registered = True + owner._install_hooks(instance) + owner._ready.set() + last_rotate = time.monotonic() + msg = wintypes.MSG() + while not owner._stop.is_set(): + owner._thread_beat = time.monotonic() + while _user32.PeekMessageW(ctypes.byref(msg), None, 0, 0, 1): + _user32.TranslateMessage(ctypes.byref(msg)) + _user32.DispatchMessageW(ctypes.byref(msg)) + for _ in range(128): + try: + event = owner._events.get_nowait() + except queue.Empty: + break + owner._handle(event) + owner._handle(("tick",)) + if owner._rotate or time.monotonic() - last_rotate >= 5.0: + owner._rotating = True # Pause AI without letting physical down leak. + try: + owner._install_hooks(instance) + owner._rotate = False + last_rotate = time.monotonic() + owner._recover_after_rehook() + finally: + owner._rotating = False + time.sleep(0.01) + except Exception as exc: + owner._startup_error = exc + owner._fail_open() + owner._ready.set() + finally: + owner._suppress = False + for hook in (owner._mouse_hook, owner._key_hook): + if hook: + _user32.UnhookWindowsHookEx(hook) + owner._mouse_hook = owner._key_hook = None + if registered: + removed = _RawDevice(1, 2, 1, None) # RIDEV_REMOVE. + _user32.RegisterRawInputDevices(ctypes.byref(removed), 1, ctypes.sizeof(removed)) + if hwnd: + _user32.DestroyWindow(hwnd) + if atom: + _user32.UnregisterClassW(name, instance) diff --git a/src/windows_mcp/desktop/flash_overlay.py b/src/windows_mcp/desktop/flash_overlay.py index ca3907ae..c5bb89ae 100644 --- a/src/windows_mcp/desktop/flash_overlay.py +++ b/src/windows_mcp/desktop/flash_overlay.py @@ -31,6 +31,7 @@ _lock = threading.Lock() _active_overlay: "_Overlay | None" = None +_all_overlays: set["_Overlay"] = set() def _flash_disabled() -> bool: @@ -45,16 +46,25 @@ def __init__(self) -> None: self.thread: threading.Thread | None = None -def cancel_active_flash(timeout: float = 0.25) -> None: - """Tear down any flash overlay currently on screen.""" +def cancel_active_flash(timeout: float = 0.25) -> bool: + """Stop every pending flash and report whether all windows have closed. + + A replaced flash can still be winding down while a newer one is active; + screenshot capture must wait for both rather than just the latest overlay. + """ global _active_overlay with _lock: - ov = _active_overlay + overlays = set(_all_overlays) + if _active_overlay is not None: + overlays.add(_active_overlay) _active_overlay = None - if ov is None: - return - ov.stop_event.set() - ov.closed_event.wait(timeout=timeout) + for overlay in overlays: + overlay.stop_event.set() + deadline = time.monotonic() + timeout + return all( + overlay.closed_event.wait(timeout=max(0.0, deadline - time.monotonic())) + for overlay in overlays + ) def show_capture_flash(capture_rect: "object | None" = None) -> None: @@ -106,9 +116,18 @@ def show_capture_flash(capture_rect: "object | None" = None) -> None: global _active_overlay prev = _active_overlay _active_overlay = overlay + _all_overlays.add(overlay) if prev is not None: prev.stop_event.set() - overlay.thread.start() + try: + overlay.thread.start() + except BaseException: + with _lock: + _all_overlays.discard(overlay) + if _active_overlay is overlay: + _active_overlay = None + overlay.closed_event.set() + raise # --------------------------------------------------------------------------- @@ -210,7 +229,26 @@ class _WNDCLASSEX(ctypes.Structure): _user32.CreateWindowExW.restype = wintypes.HWND +_user32.CreateWindowExW.argtypes = [ + wintypes.DWORD, + wintypes.LPCWSTR, + wintypes.LPCWSTR, + wintypes.DWORD, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + wintypes.HWND, + wintypes.HMENU, + wintypes.HINSTANCE, + wintypes.LPVOID, +] _user32.RegisterClassExW.restype = ctypes.c_ushort +_user32.RegisterClassExW.argtypes = [ctypes.POINTER(_WNDCLASSEX)] +_user32.UnregisterClassW.argtypes = [wintypes.LPCWSTR, wintypes.HINSTANCE] +_user32.UnregisterClassW.restype = wintypes.BOOL +_kernel32.GetModuleHandleW.argtypes = [wintypes.LPCWSTR] +_kernel32.GetModuleHandleW.restype = wintypes.HMODULE _user32.DefWindowProcW.restype = _LRESULT _user32.DefWindowProcW.argtypes = [ wintypes.HWND, @@ -478,6 +516,11 @@ def _run_overlay( from PIL import Image # noqa: F401 — fail fast if Pillow missing except Exception: logger.debug("Pillow unavailable; skipping screenshot flash") + with _lock: + global _active_overlay + if _active_overlay is overlay: + _active_overlay = None + _all_overlays.discard(overlay) overlay.closed_event.set() return @@ -561,7 +604,7 @@ def _run_overlay( except Exception: pass with _lock: - global _active_overlay if _active_overlay is overlay: _active_overlay = None + _all_overlays.discard(overlay) overlay.closed_event.set() diff --git a/src/windows_mcp/desktop/service.py b/src/windows_mcp/desktop/service.py index 29d2045f..aab1ef39 100644 --- a/src/windows_mcp/desktop/service.py +++ b/src/windows_mcp/desktop/service.py @@ -14,8 +14,10 @@ from PIL import ImageFont, ImageDraw, Image from windows_mcp.tree.service import Tree from windows_mcp.desktop import screenshot as screenshot_capture +from windows_mcp.desktop.control import get_controller from windows_mcp.desktop import flash_overlay from windows_mcp.infrastructure import safe_get +from windows_mcp.desktop import control_overlay from urllib.parse import urljoin from locale import getpreferredencoding from typing import Literal @@ -718,12 +720,15 @@ def _validate_screen_point(x: int, y: int) -> None: ) def click(self, loc: tuple[int, int] | list[int], button: str = "left", clicks: int = 1): + owner = get_controller() + owner.checkpoint_current() if isinstance(loc, list): x, y = loc[0], loc[1] else: x, y = loc self._validate_screen_point(x, y) if clicks == 0: + owner.checkpoint_current() uia.SetCursorPos(x, y) return match button: @@ -731,15 +736,23 @@ def click(self, loc: tuple[int, int] | list[int], button: str = "left", clicks: if clicks >= 2: dbl_wait = uia.GetDoubleClickTime() / 2000.0 for i in range(clicks): + owner.checkpoint_current() uia.Click(x, y, waitTime=dbl_wait if i < clicks - 1 else 0.5) + owner.record_step_current() else: + owner.checkpoint_current() uia.Click(x, y) + owner.record_step_current() case "right": for _ in range(clicks): + owner.checkpoint_current() uia.RightClick(x, y) + owner.record_step_current() case "middle": for _ in range(clicks): + owner.checkpoint_current() uia.MiddleClick(x, y) + owner.record_step_current() # Strings longer than this typed via clipboard paste instead of # per-key SendKeys. SendKeys at high cadence loses keystrokes on @@ -761,29 +774,44 @@ def type( ): clear = as_bool(clear, "clear") press_enter = as_bool(press_enter, "press_enter") + owner = get_controller() + owner.checkpoint_current() x, y = loc self._validate_screen_point(x, y) uia.Click(x, y) + owner.record_step_current() if caret_position == "start": + owner.checkpoint_current() uia.SendKeys("{Home}", waitTime=0.05) + owner.record_step_current() elif caret_position == "end": + owner.checkpoint_current() uia.SendKeys("{End}", waitTime=0.05) + owner.record_step_current() if clear: sleep(0.5) + owner.checkpoint_current() uia.SendKeys("{Ctrl}a", waitTime=0.05) + owner.record_step_current() + owner.checkpoint_current() uia.SendKeys("{Back}", waitTime=0.05) + owner.record_step_current() # Per-key SendKeys for short text (so escape sequences keep working); # clipboard paste for long text (so the scan-code queue can't race). has_control_chars = any(c in text for c in ("\n", "\t", "{", "}")) if len(text) >= self._LONG_TEXT_PASTE_THRESHOLD and not has_control_chars: self._paste_text(text) else: - escaped_text = _escape_text_for_sendkeys(text) # Bump interval from 0.02 → 0.04. Keeps short-text speed acceptable - # while reducing key-loss on slower systems. - uia.SendKeys(escaped_text, interval=0.04, waitTime=0.05) + # while reducing key-loss; each character is a cancellation point. + for ch in text: + owner.checkpoint_current() + uia.SendKeys(_escape_text_for_sendkeys(ch), interval=0.04, waitTime=0.05) + owner.record_step_current() if press_enter: + owner.checkpoint_current() uia.SendKeys("{Enter}", waitTime=0.05) + owner.record_step_current() def _paste_text(self, text: str): """Stash text on the clipboard, Ctrl+V, restore prior clipboard. @@ -791,6 +819,8 @@ def _paste_text(self, text: str): route through SendKeys instead so escape sequences are honored. """ prior = None + owner = get_controller() + owner.checkpoint_current() try: prior = uia.GetClipboardText() except Exception: @@ -798,14 +828,18 @@ def _paste_text(self, text: str): uia.SetClipboardText(text) # Tiny pause so the OS clipboard write settles before Ctrl+V reads. sleep(0.05) - uia.SendKeys("{Ctrl}v", waitTime=0.05) - # Restore prior clipboard so we don't surprise other tools. - if prior is not None: - sleep(0.05) - try: - uia.SetClipboardText(prior) - except Exception: - pass + try: + owner.checkpoint_current() + uia.SendKeys("{Ctrl}v", waitTime=0.05) + owner.record_step_current() + finally: + # Do not overwrite a clipboard value changed by the user. + if prior is not None: + try: + if uia.GetClipboardText() == text: + uia.SetClipboardText(prior) + except Exception: + pass def scroll( self, @@ -814,29 +848,55 @@ def scroll( direction: Literal["up", "down", "left", "right"] = "down", wheel_times: int = 1, ) -> str | None: + owner = get_controller() + owner.checkpoint_current() if loc: self.move(loc) match type: case "vertical": match direction: case "up": - uia.WheelUp(wheel_times) + for _ in range(wheel_times): + owner.checkpoint_current() + uia.WheelUp(1) + owner.record_step_current() case "down": - uia.WheelDown(wheel_times) + for _ in range(wheel_times): + owner.checkpoint_current() + uia.WheelDown(1) + owner.record_step_current() case _: return 'Invalid direction. Use "up" or "down".' case "horizontal": match direction: case "left": - uia.PressKey(uia.Keys.VK_SHIFT, waitTime=0.05) - uia.WheelUp(wheel_times) - sleep(0.05) - uia.ReleaseKey(uia.Keys.VK_SHIFT, waitTime=0.05) + try: + owner.input_ledger.press( + "key:shift", + lambda: uia.PressKey(uia.Keys.VK_SHIFT, waitTime=0.05), + lambda: uia.ReleaseKey(uia.Keys.VK_SHIFT, waitTime=0.05), + lambda: owner.physical_key_down(uia.Keys.VK_SHIFT), + ) + for _ in range(wheel_times): + owner.checkpoint_current() + uia.WheelUp(1) + owner.record_step_current() + finally: + owner.input_ledger.release("key:shift") case "right": - uia.PressKey(uia.Keys.VK_SHIFT, waitTime=0.05) - uia.WheelDown(wheel_times) - sleep(0.05) - uia.ReleaseKey(uia.Keys.VK_SHIFT, waitTime=0.05) + try: + owner.input_ledger.press( + "key:shift", + lambda: uia.PressKey(uia.Keys.VK_SHIFT, waitTime=0.05), + lambda: uia.ReleaseKey(uia.Keys.VK_SHIFT, waitTime=0.05), + lambda: owner.physical_key_down(uia.Keys.VK_SHIFT), + ) + for _ in range(wheel_times): + owner.checkpoint_current() + uia.WheelDown(1) + owner.record_step_current() + finally: + owner.input_ledger.release("key:shift") case _: return 'Invalid direction. Use "left" or "right".' case _: @@ -873,28 +933,61 @@ def drag( from_loc: tuple[int, int] | list[int] | None = None, duration: float | int | str | None = None, ) -> dict[str, object]: + owner = get_controller() x, y = self._normalize_drag_point(loc, "loc") normalized_from_loc = ( None if from_loc is None else self._normalize_drag_point(from_loc, "from_loc") ) effective_duration = self._normalize_drag_duration(duration) + owner.checkpoint_current() sleep(0.5) + owner.checkpoint_current() if normalized_from_loc is None: cx, cy = uia.GetCursorPos() else: cx, cy = normalized_from_loc - uia.DragDrop(cx, cy, x, y, moveSpeed=1, duration=effective_duration) - return { - "start": [cx, cy], - "end": [x, y], - "duration": effective_duration, - } + try: + owner.input_ledger.press( + "mouse:left", + lambda: uia.PressMouse(cx, cy, waitTime=0.05), + lambda: uia.ReleaseMouse(waitTime=0.05), + lambda: owner.physical_mouse_down("left"), + ) + owner.record_step_current() + # Keep both the legacy drag and duration-based drag interruptible. + steps = ( + max(1, min(80, max(abs(x - cx), abs(y - cy)) // 20)) + if effective_duration is None + else max(1, min(200, math.ceil(effective_duration / 0.01))) + ) + interval = 0.01 if effective_duration is None else effective_duration / steps + for step in range(1, steps + 1): + owner.checkpoint_current() + uia.SetCursorPos( + cx + round((x - cx) * step / steps), cy + round((y - cy) * step / steps) + ) + owner.record_step_current() + if interval: + sleep(interval) + finally: + owner.input_ledger.release("mouse:left") + return {"start": [cx, cy], "end": [x, y], "duration": effective_duration} def move(self, loc: tuple[int, int]): + owner = get_controller() x, y = loc - uia.MoveTo(x, y, moveSpeed=10) + owner.checkpoint_current() + cx, cy = uia.GetCursorPos() + steps = max(1, min(80, max(abs(x - cx), abs(y - cy)) // 20)) + for step in range(1, steps + 1): + owner.checkpoint_current() + uia.SetCursorPos(cx + (x - cx) * step // steps, cy + (y - cy) * step // steps) + owner.record_step_current() + sleep(0.003) def shortcut(self, shortcut: str): + owner = get_controller() + owner.checkpoint_current() keys = shortcut.split("+") sendkeys_str = "" for key in keys: @@ -905,21 +998,35 @@ def shortcut(self, shortcut: str): name = _KEY_ALIASES.get(key.lower(), key) sendkeys_str += "{" + name + "}" uia.SendKeys(sendkeys_str, interval=0.01) + owner.record_step_current() def multi_select(self, press_ctrl: bool | str = False, locs: list[tuple[int, int]] = []): + owner = get_controller() + owner.checkpoint_current() press_ctrl = press_ctrl is True or ( isinstance(press_ctrl, str) and press_ctrl.lower() == "true" ) - if press_ctrl: - uia.PressKey(uia.Keys.VK_CONTROL, waitTime=0.05) - for loc in locs: - x, y = loc - uia.Click(x, y, waitTime=0.2) - sleep(0.5) - uia.ReleaseKey(uia.Keys.VK_CONTROL, waitTime=0.05) + try: + if press_ctrl: + owner.input_ledger.press( + "key:ctrl", + lambda: uia.PressKey(uia.Keys.VK_CONTROL, waitTime=0.05), + lambda: uia.ReleaseKey(uia.Keys.VK_CONTROL, waitTime=0.05), + lambda: owner.physical_key_down(uia.Keys.VK_CONTROL), + ) + for loc in locs: + owner.checkpoint_current() + x, y = loc + uia.Click(x, y, waitTime=0.2) + owner.record_step_current() + sleep(0.5) + finally: + if press_ctrl: + owner.input_ledger.release("key:ctrl") def multi_edit(self, locs: list[tuple[int, int, str]]): for loc in locs: + get_controller().checkpoint_current() x, y, text = loc self.type((x, y), text=text, clear=True) @@ -1245,9 +1352,13 @@ def get_display_union_rect( ) def get_screenshot(self, capture_rect: uia.Rect | None = None) -> Image.Image: - flash_overlay.cancel_active_flash() - image, used_backend = screenshot_capture.capture(capture_rect) - self._last_screenshot_backend = used_backend + # The persistent AI indicator is hidden and acknowledged before any + # screenshot backend reads pixels; it is restored even if capture fails. + with control_overlay.suspend_for_capture(): + if not flash_overlay.cancel_active_flash(): + raise RuntimeError("Previous screenshot flash did not close before capture") + image, used_backend = screenshot_capture.capture(capture_rect) + self._last_screenshot_backend = used_backend flash_overlay.show_capture_flash(capture_rect) return image diff --git a/src/windows_mcp/infrastructure/config.py b/src/windows_mcp/infrastructure/config.py index e2c5124f..c6c408ce 100644 --- a/src/windows_mcp/infrastructure/config.py +++ b/src/windows_mcp/infrastructure/config.py @@ -38,11 +38,19 @@ class ToolsConfig: exclude: list[str] = field(default_factory=list) +@dataclass +class ControlConfig: + # Raw relative device units and hook screen pixels are intentionally separate. + mouse_takeover_units: int = 120 + mouse_takeover_pixels: int = 120 + + @dataclass class WindowsMCPConfig: server: ServerConfig = field(default_factory=ServerConfig) security: SecurityConfig = field(default_factory=SecurityConfig) tools: ToolsConfig = field(default_factory=ToolsConfig) + control: ControlConfig = field(default_factory=ControlConfig) source_path: Path | None = None @@ -104,6 +112,7 @@ def load_config(path: Path | None) -> WindowsMCPConfig: server = data.get("server", {}) security = data.get("security", {}) tools = data.get("tools", {}) + control = data.get("control", {}) _VALID_TRANSPORTS = {"stdio", "sse", "streamable-http"} if "transport" in server: @@ -154,6 +163,13 @@ def load_config(path: Path | None) -> WindowsMCPConfig: if "exclude" in tools: cfg.tools.exclude = _list_of_strings(tools["exclude"], "tools.exclude") + for key in ("mouse_takeover_units", "mouse_takeover_pixels"): + if key in control: + value = _strict_int(control[key], f"control.{key}") + if not 40 <= value <= 2000: + raise ValueError(f"control.{key} must be between 40 and 2000") + setattr(cfg.control, key, value) + cfg.source_path = path return cfg @@ -202,4 +218,13 @@ def write_config(cfg: WindowsMCPConfig, path: Path) -> None: items = ", ".join(f'"{t}"' for t in cfg.tools.exclude) lines += ["[tools]", f"exclude = [{items}]", ""] + ctrl, default_ctrl = cfg.control, ControlConfig() + control_lines = [ + f"{key} = {getattr(ctrl, key)}" + for key in ("mouse_takeover_units", "mouse_takeover_pixels") + if getattr(ctrl, key) != getattr(default_ctrl, key) + ] + if control_lines: + lines += ["[control]", *control_lines, ""] + path.write_text("\n".join(lines), encoding="utf-8") diff --git a/src/windows_mcp/tools/__init__.py b/src/windows_mcp/tools/__init__.py index b4b5c95c..01d9a8f7 100644 --- a/src/windows_mcp/tools/__init__.py +++ b/src/windows_mcp/tools/__init__.py @@ -3,6 +3,7 @@ from windows_mcp.tools import ( app, clipboard, + control_status, display, filesystem, input, @@ -16,6 +17,7 @@ ) _MODULES = [ + control_status, app, display, shell, diff --git a/src/windows_mcp/tools/control_notifications.py b/src/windows_mcp/tools/control_notifications.py new file mode 100644 index 00000000..7655a774 --- /dev/null +++ b/src/windows_mcp/tools/control_notifications.py @@ -0,0 +1,251 @@ +"""Global MCP ownership gate and best-effort session state notifications.""" + +from __future__ import annotations + +import asyncio +import json +import logging +import time +from dataclasses import dataclass +from typing import Any, Callable + +from fastmcp.exceptions import ToolError +from fastmcp.server.middleware import Middleware, MiddlewareContext + +logger = logging.getLogger(__name__) + + +@dataclass +class _Session: + session: Any + loop: asyncio.AbstractEventLoop + last_seen: float + + +class ControlNotifier: + """Send state transitions to live MCP sessions without holding control locks.""" + + def __init__(self, controller: Any, *, ttl: float = 60.0) -> None: + self.controller = controller + self.ttl = ttl + self._sessions: dict[int, _Session] = {} + self._loop: asyncio.AbstractEventLoop | None = None + self._reaper: asyncio.Task | None = None + self._state = "unavailable" + self._generation = -1 + self._active = False + self._subscribed = False + self.push_enabled = True + self._pending: list[dict] = [] + self._flush_handle: asyncio.TimerHandle | None = None + self._send_tasks: set[asyncio.Task] = set() + + def start(self) -> None: + self._loop = asyncio.get_running_loop() + self._active = True + initial = self.controller.status() + self._state = initial["state"] + self._generation = initial["generation"] + if not self._subscribed: + self.controller.subscribe(self._on_state) + self._subscribed = True + self._reaper = self._loop.create_task(self._reap()) + + async def close(self) -> None: + self._active = False + if self._flush_handle: + self._flush_handle.cancel() + self._flush_handle = None + self._pending.clear() + if self._reaper: + self._reaper.cancel() + try: + await self._reaper + except asyncio.CancelledError: + pass + self._reaper = None + tasks = tuple(self._send_tasks) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._send_tasks.clear() + self._sessions.clear() + self._loop = None + + def remember(self, context: MiddlewareContext) -> None: + """Refresh only sessions seen in a real request; stateless sessions expire.""" + if not self._active or not self.push_enabled or context.fastmcp_context is None: + return + request_loop = asyncio.get_running_loop() + if request_loop is not self._loop: + # A session on another event loop cannot be sent or awaited safely + # by this server's notifier; the next tool response remains the fallback. + return + try: + session = context.fastmcp_context.session + except RuntimeError: + return + self._sessions[id(session)] = _Session(session, request_loop, time.monotonic()) + + async def _reap(self) -> None: + while self._active: + await asyncio.sleep(10) + cutoff = time.monotonic() - self.ttl + for key, value in tuple(self._sessions.items()): + if value.last_seen < cutoff or value.loop.is_closed(): + self._sessions.pop(key, None) + + def _on_state(self, status: dict) -> None: + loop = self._loop + if not self._active or loop is None or loop.is_closed(): + return + try: + loop.call_soon_threadsafe(self._enqueue, status.copy()) + except RuntimeError: + # Shutdown can close the event loop between the check and scheduling. + logger.debug("Control event loop closed before notification dispatch") + + def _enqueue(self, status: dict) -> None: + if not self._active: + return + self._pending.append(status) + if self._flush_handle is None: + # State callbacks can race after releasing the coordinator lock. + # A short batch sorts by generation before notifying clients. + self._flush_handle = asyncio.get_running_loop().call_later(0.02, self._flush) + + def _flush(self) -> None: + self._flush_handle = None + pending, self._pending = self._pending, [] + for status in sorted(pending, key=lambda item: item["generation"]): + self._dispatch(status) + + def _dispatch(self, status: dict) -> None: + if not self._active: + return + if status["generation"] <= self._generation: + return # Cross-thread delivery can reorder old state events. + self._generation = status["generation"] + previous, current = self._state, status["state"] + self._state = current + event = None + if current == "user" and previous in ("ready", "ai", "takeover_pending"): + event = "USER_CONTROL" + elif current == "ready" and previous == "user": + event = "AI_CONTROL_AVAILABLE" + if event is None: + return + cutoff = time.monotonic() - self.ttl + for key, value in tuple(self._sessions.items()): + if value.last_seen < cutoff or value.loop.is_closed(): + self._sessions.pop(key, None) + continue + task = value.loop.create_task(self._send(key, value, event, status)) + self._send_tasks.add(task) + task.add_done_callback(self._send_tasks.discard) + + async def _send(self, key: int, entry: _Session, event: str, status: dict) -> None: + if not self._active: + return + try: + await entry.session.send_log_message( + level="notice", + data={"event": event, "status": status}, + logger="windows-mcp.control", + ) + except Exception: + self._sessions.pop(key, None) + logger.debug("MCP control notification failed; session removed", exc_info=True) + + +class ControlToolGate(Middleware): + """All present and future MCP tools use the same ownership decision.""" + + def __init__(self, controller: Any, notifier: ControlNotifier) -> None: + self.controller = controller + self.notifier = notifier + self._call_lock = asyncio.Lock() + + @staticmethod + def _blocked(exc: Exception) -> ToolError: + return ToolError( + json.dumps( + {"code": exc.code, "status": exc.status}, + ensure_ascii=False, + ) + ) + + async def on_call_tool(self, context: MiddlewareContext, call_next: Callable) -> Any: + self.notifier.remember(context) + name = context.message.name + if name == "ControlStatus": + return await call_next(context) + # Waiting behind an already running external command must not hide a + # user takeover. Poll ownership while waiting, then recheck under lock. + from windows_mcp.desktop.control import ControlBlocked + + def reject_user_state() -> None: + status = self.controller.status() + if status["state"] not in ("ready", "ai"): + raise self._blocked( + ControlBlocked( + {"user": "USER_CONTROL", "takeover_pending": "TAKEOVER_PENDING"}.get( + status["state"], "CONTROL_UNAVAILABLE" + ), + status, + ) + ) + + reject_user_state() + while True: + try: + await asyncio.wait_for(self._call_lock.acquire(), timeout=0.1) + break + except asyncio.TimeoutError: + reject_user_state() + try: + try: + token = self.controller.begin_call(name) + except Exception as exc: + from windows_mcp.desktop.control import ControlBlocked + + if not isinstance(exc, ControlBlocked): + raise + raise self._blocked(exc) from exc + from windows_mcp.desktop.control_context import ( + current_token, + current_steps, + StepCounter, + get_step_count, + ) + + context_token = current_token.set(token) + steps_token = current_steps.set(StepCounter()) + try: + # A synchronous indicator failure can invalidate the lease + # during begin_call's state notification. + self.controller.checkpoint(token) + return await call_next(context) + except Exception as exc: + from windows_mcp.desktop.control import ControlBlocked + from windows_mcp.desktop.control_ledger import InputUnavailable + + if isinstance(exc, ControlBlocked): + raise self._blocked(exc) from exc + if isinstance(exc, InputUnavailable): + status = self.controller.status() + status["executed_steps"] = get_step_count() + code = ( + "CONTROL_UNAVAILABLE" + if status["state"] == "unavailable" + else "CONTROL_PREEMPTED" + ) + raise self._blocked(ControlBlocked(code, status)) from exc + raise + finally: + current_steps.reset(steps_token) + current_token.reset(context_token) + self.controller.end_call(token) + finally: + self._call_lock.release() diff --git a/src/windows_mcp/tools/control_status.py b/src/windows_mcp/tools/control_status.py new file mode 100644 index 00000000..a7546106 --- /dev/null +++ b/src/windows_mcp/tools/control_status.py @@ -0,0 +1,21 @@ +"""Read-only MCP control ownership status tool.""" + +from mcp.types import ToolAnnotations + +from windows_mcp.desktop.control import get_controller + + +def register(mcp, *, get_desktop, get_analytics): + @mcp.tool( + name="ControlStatus", + description="Get desktop control owner and remaining user idle wait in seconds.", + annotations=ToolAnnotations( + title="ControlStatus", + readOnlyHint=True, + destructiveHint=False, + idempotentHint=True, + openWorldHint=False, + ), + ) + def control_status_tool() -> dict: + return get_controller().status() diff --git a/tests/test_control.py b/tests/test_control.py new file mode 100644 index 00000000..6ad5aa9c --- /dev/null +++ b/tests/test_control.py @@ -0,0 +1,584 @@ +"""State and fail-open tests for physical desktop ownership.""" + +import ctypes +import queue +import threading +from unittest.mock import Mock + +import pytest + +from windows_mcp.desktop import control +from windows_mcp.desktop import service +from windows_mcp.desktop.control_context import StepCounter +from windows_mcp.desktop.control_ledger import InputLedger, InputUnavailable + + +def ready_controller() -> control.ControlCoordinator: + owner = control.ControlCoordinator() + with owner._lock: + owner._set_locked("ready") + return owner + + +def test_start_is_ready_without_physical_input(monkeypatch): + owner = control.ControlCoordinator() + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + + def run(): + owner._thread_beat = control.time.monotonic() + owner._ready.set() + owner._stop.wait() + + monkeypatch.setattr(owner, "_run", run) + owner.start() + try: + assert owner.status()["state"] == "ready" + token = owner.begin_call("Wait") + owner.checkpoint(token) + owner.end_call(token) + finally: + owner.stop() + + +def test_start_preserves_cooldown_after_real_physical_input(monkeypatch): + owner = control.ControlCoordinator() + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + + def run(): + owner._thread_beat = control.time.monotonic() + owner._last_physical_event = control.time.monotonic() + owner._ready.set() + owner._stop.wait() + + monkeypatch.setattr(owner, "_run", run) + owner.start() + try: + assert owner.status()["state"] == "user" + with pytest.raises(control.ControlBlocked, match="USER_CONTROL"): + owner.begin_call("Wait") + finally: + owner.stop() + + +def test_physical_mouse_pixel_threshold_preempts_only_at_boundary(): + owner = ready_controller() + owner.begin_call("Click") + owner._handle(("point", 100, 100, 0x200)) + owner._handle(("point", 219, 100, 0x200)) + assert owner.status()["state"] == "takeover_pending" + owner._handle(("point", 220, 100, 0x200)) + assert owner.status()["state"] == "user" + with pytest.raises(control.ControlBlocked) as exc: + owner.begin_call("Click") + assert exc.value.code == "USER_CONTROL" + + +def test_raw_mouse_is_per_device_and_relative(): + owner = ready_controller() + owner.begin_call("Move") + owner._handle(("raw", 1, 119, 0)) + owner._handle(("raw", 2, 119, 0)) + assert owner.status()["state"] == "takeover_pending" + owner._handle(("raw", 2, 1, 0)) + assert owner.status()["state"] == "user" + + +def test_raw_device_switch_resets_previous_candidate(): + owner = ready_controller() + owner.begin_call("Move") + owner._handle(("raw", 1, 100, 0)) + owner._handle(("raw", 2, 100, 0)) + owner._handle(("raw", 1, 100, 0)) + assert owner.status()["state"] == "takeover_pending" + owner._handle(("raw", 1, 20, 0)) + assert owner.status()["state"] == "user" + + +def test_user_idle_resumes_at_ten_seconds(monkeypatch): + clock = [100.0] + monkeypatch.setattr(control.time, "monotonic", lambda: clock[0]) + owner = ready_controller() + owner._handle(("key",)) + assert owner.status()["state"] == "user" + clock[0] = 109.999 + assert owner.status()["state"] == "user" + owner._handle(("key",)) + clock[0] = 119.998 + assert owner.status()["state"] == "user" + clock[0] = 119.999 + assert owner.status()["state"] == "ready" + + +def test_user_idle_waits_until_held_mouse_is_released(monkeypatch): + clock = [100.0] + monkeypatch.setattr(control.time, "monotonic", lambda: clock[0]) + owner = ready_controller() + owner._handle(("key",)) + owner._mouse_down.add(1) + clock[0] = 111.0 + assert owner.status()["state"] == "user" + owner._mouse_down.clear() + owner._handle(("point", 10, 10, 0x202)) + clock[0] = 120.999 + assert owner.status()["state"] == "user" + clock[0] = 121.0 + assert owner.status()["state"] == "ready" + + +def test_user_idle_counts_physical_up_before_queue_consumption(monkeypatch): + clock = [100.0] + monkeypatch.setattr(control.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + owner = ready_controller() + owner._handle(("key",)) + owner._pressed.add(0x41) + clock[0] = 111.0 + event = control._KeyHookData() + event.vkCode = 0x41 + assert owner._physical_key(0, 0x101, ctypes.addressof(event)) == 7 + # The event remains queued, but the hook timestamp already blocks resume. + assert owner.status()["state"] == "user" + clock[0] = 120.999 + assert owner.status()["state"] == "user" + clock[0] = 121.0 + assert owner.status()["state"] == "ready" + + +def test_pending_does_not_expire_while_mouse_button_is_held(monkeypatch): + clock = [100.0] + monkeypatch.setattr(control.time, "monotonic", lambda: clock[0]) + owner = ready_controller() + owner.begin_call("Click") + owner._handle(("point", 10, 10, 0x200)) + owner._mouse_down.add(1) + clock[0] = 101.0 + assert owner.status()["state"] == "takeover_pending" + owner._mouse_down.clear() + assert owner.status()["state"] == "ai" + + +def test_pending_cancels_without_replaying_stale_call(monkeypatch): + clock = [100.0] + monkeypatch.setattr(control.time, "monotonic", lambda: clock[0]) + owner = ready_controller() + token = owner.begin_call("Type") + owner._handle(("point", 10, 10, 0x200)) + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + owner.end_call(token) + clock[0] = 100.301 + assert owner.status()["state"] == "ai" + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + + +@pytest.mark.parametrize("press_order", [(0x11, 0x12, 0x10, 0x08), (0x10, 0x12, 0x11, 0x08)]) +def test_hotkey_fast_path_releases_before_coordinator_and_quarantines_up(monkeypatch, press_order): + owner = ready_controller() + owner.begin_call("Type") + owner._deadline = control.time.monotonic() + 1 + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + + def send(vk, message): + data = control._KeyHookData() + data.vkCode = vk + return owner._physical_key(0, message, ctypes.addressof(data)) + + for vk in press_order: + assert send(vk, 0x100) == 1 + assert not owner._suppress + assert owner._fast_takeover + assert send(0x08, 0x100) == 1 # Repeated Backspace is still swallowed. + for vk in (0x08, 0x10, 0x12, 0x11): + assert send(vk, 0x101) == 1 + assert not owner._quarantine + owner._handle(("hotkey",)) + assert owner.status()["state"] == "user" + assert send(0x41, 0x100) == 7 + + +def test_backspace_repeat_after_modifiers_does_not_trigger_hotkey(monkeypatch): + owner = ready_controller() + owner.begin_call("Type") + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + + def down(vk): + data = control._KeyHookData() + data.vkCode = vk + return owner._physical_key(0, 0x100, ctypes.addressof(data)) + + for vk in (0x08, 0x11, 0x12, 0x10, 0x08): + assert down(vk) == 1 + assert not owner._fast_takeover + assert owner._suppress + + +def test_injected_key_passes_and_queue_full_fails_open(monkeypatch): + owner = ready_controller() + owner.begin_call("Type") + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + injected = control._KeyHookData() + injected.vkCode = 0x41 + injected.flags = 0x10 + assert owner._physical_key(0, 0x100, ctypes.addressof(injected)) == 7 + owner._events = queue.Queue(maxsize=1) + owner._queue(("first",)) + owner._queue(("overflow",)) + assert owner._emergency and not owner._suppress + assert owner._physical_key(-1, 0x100, ctypes.addressof(injected)) == 7 + + +def test_physical_move_pauses_ai_before_event_queue_is_processed(monkeypatch): + owner = ready_controller() + token = owner.begin_call("Move") + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + data = control._MouseHookData() + data.pt.x, data.pt.y = 10, 20 + assert owner._physical_mouse(0, 0x200, ctypes.addressof(data)) == 1 + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + owner._handle(owner._events.get_nowait()) + assert owner.status()["state"] == "takeover_pending" + + +def test_swallowed_physical_down_does_not_skip_ai_release(monkeypatch): + owner = ready_controller() + owner.begin_call("Drag") + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + data = control._MouseHookData() + assert owner._physical_mouse(0, 0x201, ctypes.addressof(data)) == 1 + assert not owner.physical_mouse_down("left") + owner._handle(("hotkey",)) + assert owner._physical_mouse(0, 0x201, ctypes.addressof(data)) == 7 + assert owner.physical_mouse_down("left") + assert owner._physical_mouse(0, 0x202, ctypes.addressof(data)) == 7 + assert not owner.physical_mouse_down("left") + + +def test_fail_open_updates_status_and_notifies_without_input_thread(): + owner = ready_controller() + owner.begin_call("Type") + events = [] + owner.subscribe(events.append) + owner._fail_open() + assert owner.status()["state"] == "unavailable" + assert events[-1]["state"] == "unavailable" + + +def test_stop_rejects_still_running_hook_thread(): + owner = ready_controller() + owner._thread = Mock() + owner._thread.is_alive.return_value = True + owner._watchdog = None + with pytest.raises(RuntimeError, match="did not stop cleanly"): + owner.stop() + with pytest.raises(RuntimeError, match="has not stopped"): + owner.start() + + +def test_watchdog_cannot_refresh_if_coordinator_lock_is_held(monkeypatch): + owner = ready_controller() + owner.begin_call("Type") + owner._thread = Mock() + owner._thread.is_alive.return_value = True + owner._thread_beat = control.time.monotonic() + owner._lock.acquire() + try: + owner._stop.is_set = Mock(side_effect=[False, True]) + owner._release_requested.wait = Mock() + owner._watch() + finally: + owner._lock.release() + assert owner._emergency and not owner._suppress + + +def test_type_stops_before_next_character_after_preemption(monkeypatch): + owner = Mock() + checks = [None, None, control.ControlBlocked("CONTROL_PREEMPTED", {"state": "user"})] + owner.checkpoint_current.side_effect = checks + monkeypatch.setattr(service, "get_controller", lambda: owner) + click = Mock() + send = Mock() + monkeypatch.setattr(service.uia, "Click", click) + monkeypatch.setattr(service.uia, "SendKeys", send) + desktop = service.Desktop.__new__(service.Desktop) + with pytest.raises(control.ControlBlocked): + desktop.type((1, 2), "abc") + assert send.call_count == 1 + send.assert_any_call("a", interval=0.04, waitTime=0.05) + + +def test_drag_releases_its_mouse_button_after_preemption(monkeypatch): + owner = Mock() + owner.input_ledger = InputLedger() + owner.checkpoint_current.side_effect = [ + None, + None, + control.ControlBlocked("CONTROL_PREEMPTED", {"state": "user"}), + ] + owner.physical_mouse_down.return_value = False + monkeypatch.setattr(service, "get_controller", lambda: owner) + monkeypatch.setattr(service, "sleep", lambda _: None) + monkeypatch.setattr(service.uia, "GetCursorPos", lambda: (0, 0)) + press = Mock() + release = Mock() + monkeypatch.setattr(service.uia, "PressMouse", press) + monkeypatch.setattr(service.uia, "ReleaseMouse", release) + desktop = service.Desktop.__new__(service.Desktop) + with pytest.raises(control.ControlBlocked): + desktop.drag((100, 0)) + press.assert_called_once() + release.assert_called_once() + + +def test_ai_hold_is_released_if_press_raises_after_injection(): + ledger = InputLedger() + sent = [] + + def down_then_raise(): + sent.append("down") + raise RuntimeError("wait failed after injection") + + with pytest.raises(RuntimeError, match="wait failed"): + ledger.press("key:shift", down_then_raise, lambda: sent.append("up"), lambda: False) + assert sent == ["down", "up"] + assert ledger.pending() == () + + +def test_ai_hold_defers_up_only_when_physical_down_reached_app(): + ledger = InputLedger() + sent = [] + ledger.press("mouse:left", lambda: sent.append("down"), lambda: sent.append("up"), lambda: True) + ledger.release("mouse:left") + assert sent == ["down"] + assert ledger.pending() == () + + +def test_fail_open_blocks_new_press_and_clears_before_notification(): + owner = ready_controller() + owner.begin_call("Drag") + actions = [] + owner.input_ledger.press( + "mouse:left", lambda: actions.append("down"), lambda: actions.append("up"), lambda: False + ) + owner.subscribe(lambda _: actions.append("notified")) + owner._fail_open() + with pytest.raises(InputUnavailable, match="unavailable"): + owner.input_ledger.press( + "key:shift", lambda: actions.append("late down"), lambda: None, lambda: False + ) + owner._mark_unavailable() + assert actions == ["down", "up", "notified"] + + +def test_fail_open_wakes_existing_watchdog_to_release_ai_hold(monkeypatch): + owner = ready_controller() + owner.begin_call("Drag") + released = threading.Event() + owner.input_ledger.press("mouse:left", lambda: None, released.set, lambda: False) + owner._thread = Mock() + owner._thread.is_alive.return_value = True + owner._thread_beat = control.time.monotonic() + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + watcher = threading.Thread(target=owner._watch) + watcher.start() + try: + owner._fail_open() + assert owner._release_requested.is_set() or released.is_set() + assert released.wait(1.0) + assert owner.input_ledger.pending() == () + finally: + owner._stop.set() + owner._release_requested.set() + watcher.join(2) + assert not watcher.is_alive() + + +def test_fail_open_releases_before_slow_health_probe(monkeypatch): + owner = ready_controller() + owner.begin_call("Type") + released, probe_entered, probe_continue = ( + threading.Event(), + threading.Event(), + threading.Event(), + ) + owner.input_ledger.press("key:ctrl", lambda: None, released.set, lambda: False) + owner._thread = Mock() + owner._thread.is_alive.return_value = True + owner._thread_beat = control.time.monotonic() + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + + def slow_probe(): + probe_entered.set() + probe_continue.wait(2) + return True + + owner.set_health_probe(slow_probe) + # The hook path signals cleanup; it never waits for a UI health probe. + owner._fail_open() + watcher = threading.Thread(target=owner._watch) + watcher.start() + try: + assert released.wait(1) + assert not probe_entered.is_set() + finally: + probe_continue.set() + owner._stop.set() + owner._release_requested.set() + watcher.join(2) + assert not watcher.is_alive() + + +def test_takeover_after_checkpoint_blocks_new_ai_down(monkeypatch): + owner = ready_controller() + token = owner.begin_call("Drag") + owner.checkpoint(token) + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + data = control._MouseHookData() + assert owner._physical_mouse(0, 0x200, ctypes.addressof(data)) == 1 + assert owner._fast_pending + with pytest.raises(InputUnavailable, match="unavailable"): + owner.input_ledger.press( + "mouse:left", lambda: pytest.fail("late AI down"), lambda: None, lambda: False + ) + owner._handle(owner._events.get_nowait()) + assert owner.status()["state"] == "takeover_pending" + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + + +@pytest.mark.parametrize("action", ["scroll", "drag", "multi_select"]) +def test_takeover_between_checkpoint_and_press_injects_no_down(monkeypatch, action): + owner = Mock() + owner.input_ledger = InputLedger() + calls = 0 + + def checkpoint(): + nonlocal calls + calls += 1 + if calls == (2 if action == "drag" else 1): + owner.input_ledger.block_new() + + owner.checkpoint_current.side_effect = checkpoint + monkeypatch.setattr(service, "get_controller", lambda: owner) + monkeypatch.setattr(service, "sleep", lambda _: None) + monkeypatch.setattr(service.uia, "GetCursorPos", lambda: (0, 0)) + key_down, mouse_down = Mock(), Mock() + monkeypatch.setattr(service.uia, "PressKey", key_down) + monkeypatch.setattr(service.uia, "PressMouse", mouse_down) + desktop = service.Desktop.__new__(service.Desktop) + with pytest.raises(InputUnavailable): + if action == "scroll": + desktop.scroll(type="horizontal", direction="left") + elif action == "drag": + desktop.drag((100, 0)) + else: + desktop.multi_select(press_ctrl=True, locs=[(100, 100)]) + key_down.assert_not_called() + mouse_down.assert_not_called() + assert owner.input_ledger.pending() == () + + +def test_new_lease_reenables_ledger_after_cancelled_mouse_candidate(): + owner = ready_controller() + token = owner.begin_call("Move") + owner.input_ledger.block_new() # Physical move's immediate guard. + owner._handle(("point", 0, 0, 0x200)) + assert owner.status()["state"] == "takeover_pending" + owner._last_move -= 0.4 + assert owner.status()["state"] == "ai" + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + owner.begin_call("Click") + actions = [] + owner.input_ledger.press( + "mouse:left", lambda: actions.append("down"), lambda: actions.append("up"), lambda: False + ) + owner.input_ledger.release("mouse:left") + assert actions == ["down", "up"] + + +def test_release_waits_for_inflight_press_then_clears_hold(): + ledger = InputLedger() + entered, finish = threading.Event(), threading.Event() + actions = [] + + def down(): + entered.set() + assert finish.wait(2) + actions.append("down") + + pressing = threading.Thread( + target=lambda: ledger.press("key:ctrl", down, lambda: actions.append("up"), lambda: False) + ) + pressing.start() + assert entered.wait(2) + ledger.block_new() + releasing = threading.Thread(target=ledger.release_all) + releasing.start() + finish.set() + pressing.join(2) + releasing.join(2) + assert not pressing.is_alive() and not releasing.is_alive() + assert actions == ["down", "up"] + assert ledger.pending() == () + + +def test_release_all_attempts_other_holds_after_one_failure(): + ledger = InputLedger() + events = [] + + def fail_up(): + raise RuntimeError("up failed") + + ledger.press("key:ctrl", lambda: None, fail_up, lambda: False) + ledger.press("mouse:left", lambda: None, lambda: events.append("mouse up"), lambda: False) + with pytest.raises(RuntimeError, match="up failed"): + ledger.release_all() + assert events == ["mouse up"] + assert ledger.pending() == ("key:ctrl",) + + +def test_recovery_requires_ledger_clear_health_and_interactive_desktop(monkeypatch): + owner = ready_controller() + owner.begin_call("Drag") + attempts = 0 + + def up(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("temporarily blocked") + + owner.input_ledger.press("mouse:left", lambda: None, up, lambda: False) + owner._fail_open() + owner._recover_after_rehook() + assert owner.status()["state"] == "unavailable" + owner.set_health_probe(lambda: False) + owner._recover_after_rehook() + assert owner.status()["state"] == "unavailable" + owner.set_health_probe(lambda: True) + monkeypatch.setattr(control, "_interactive_desktop", lambda: False) + owner._recover_after_rehook() + assert owner.status()["state"] == "unavailable" + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + owner._recover_after_rehook() + assert owner.status()["state"] == "user" + assert owner.input_ledger.pending() == () + + +def test_preemption_reports_executed_steps(): + owner = ready_controller() + token = owner.begin_call("Type") + context = control.current_token.set(token) + step_context = control.current_steps.set(StepCounter()) + try: + owner.record_step_current() + owner.record_step_current() + owner._handle(("hotkey",)) + with pytest.raises(control.ControlBlocked) as exc: + owner.checkpoint_current() + assert exc.value.status["executed_steps"] == 2 + finally: + control.current_steps.reset(step_context) + control.current_token.reset(context) diff --git a/tests/test_control_config.py b/tests/test_control_config.py new file mode 100644 index 00000000..474640cd --- /dev/null +++ b/tests/test_control_config.py @@ -0,0 +1,32 @@ +from pathlib import Path + +import pytest + +from windows_mcp.infrastructure.config import ( + ControlConfig, + WindowsMCPConfig, + load_config, + write_config, +) + + +def test_default_control_thresholds(): + assert load_config(None).control == ControlConfig(120, 120) + + +def test_control_thresholds_round_trip(tmp_path: Path): + config = WindowsMCPConfig() + config.control.mouse_takeover_units = 240 + config.control.mouse_takeover_pixels = 180 + path = tmp_path / "config.toml" + write_config(config, path) + assert load_config(path).control == config.control + + +@pytest.mark.parametrize("key", ["mouse_takeover_units", "mouse_takeover_pixels"]) +@pytest.mark.parametrize("value", ["true", '"120"', "39", "2001"]) +def test_control_thresholds_reject_invalid(tmp_path: Path, key: str, value: str): + path = tmp_path / "config.toml" + path.write_text(f"[control]\n{key} = {value}\n", encoding="utf-8") + with pytest.raises(ValueError, match=f"control.{key}"): + load_config(path) diff --git a/tests/test_control_mcp_gate.py b/tests/test_control_mcp_gate.py new file mode 100644 index 00000000..de1241ba --- /dev/null +++ b/tests/test_control_mcp_gate.py @@ -0,0 +1,516 @@ +"""Ownership gate and notification regressions at the MCP boundary.""" + +from __future__ import annotations + +import asyncio +import socket +import threading +import time +from types import SimpleNamespace + +import pytest +from fastmcp import Client, FastMCP + +from windows_mcp.desktop.control import ControlBlocked, ControlCoordinator, current_token +from windows_mcp.desktop.control_ledger import InputLedger +from windows_mcp.tools.control_notifications import ControlNotifier, ControlToolGate + + +class FakeController: + def __init__(self): + self.state = "ready" + self.generation = 1 + self.calls = [] + self.listeners = [] + + def status(self): + return {"state": self.state, "generation": self.generation, "wait_seconds": 0.0} + + def subscribe(self, callback): + self.listeners.append(callback) + + def emit(self, state): + self.state = state + self.generation += 1 + for callback in self.listeners: + callback(self.status()) + + def begin_call(self, name): + if self.state not in ("ready", "ai"): + raise ControlBlocked("USER_CONTROL", self.status()) + self.calls.append(("begin", name)) + self.state = "ai" + return 7 + + def end_call(self, token): + self.calls.append(("end", token)) + + def checkpoint(self, token): + if self.state != "ai": + raise ControlBlocked("CONTROL_PREEMPTED", self.status()) + + +@pytest.mark.asyncio +async def test_gate_blocks_dynamic_tools_but_status_remains_available(): + controller = FakeController() + mcp = FastMCP("test") + notifier = ControlNotifier(controller) + mcp.add_middleware(ControlToolGate(controller, notifier)) + effects = [] + + @mcp.tool(name="DynamicTool") + def dynamic_tool(): + effects.append(current_token.get()) + return "executed" + + @mcp.tool(name="ControlStatus") + def control_status(): + return controller.status() + + async with Client(mcp) as client: + controller.state = "user" + blocked = await client.call_tool("DynamicTool", raise_on_error=False) + assert blocked.is_error + assert "USER_CONTROL" in str(blocked.content) + assert effects == [] + status = await client.call_tool("ControlStatus") + assert "user" in str(status.content) + controller.state = "ready" + result = await client.call_tool("DynamicTool") + assert "executed" in str(result.content) + assert effects == [7] + assert current_token.get() is None + assert controller.calls == [("begin", "DynamicTool"), ("end", 7)] + + +@pytest.mark.asyncio +async def test_ledger_preemption_is_structured_and_releases_tool_lock(): + controller = FakeController() + step_owner = ControlCoordinator() + ledger = InputLedger() + mcp = FastMCP("test") + mcp.add_middleware(ControlToolGate(controller, ControlNotifier(controller))) + effects = [] + + @mcp.tool(name="Interrupted") + def interrupted(): + for _ in range(3): + step_owner.record_step_current() + ledger.block_new() + controller.state = "user" + ledger.press("key:ctrl", lambda: effects.append("down"), lambda: None, lambda: False) + + @mcp.tool(name="After") + def after(): + effects.append("after") + + async with Client(mcp) as client: + result = await client.call_tool("Interrupted", raise_on_error=False) + assert result.is_error + assert '"code": "CONTROL_PREEMPTED"' in str(result.content) + assert '"executed_steps": 3' in str(result.content) + assert ledger.pending() == () + controller.state = "ready" + ledger.enable() + await asyncio.wait_for(client.call_tool("After"), 1) + assert effects == ["after"] + + +@pytest.mark.asyncio +async def test_queued_call_rechecks_after_takeover(): + controller = FakeController() + mcp = FastMCP("test") + mcp.add_middleware(ControlToolGate(controller, ControlNotifier(controller))) + started = asyncio.Event() + release = asyncio.Event() + effects = [] + + @mcp.tool(name="Hold") + async def hold(): + started.set() + await release.wait() + effects.append("hold") + + @mcp.tool(name="Next") + def next_tool(): + effects.append("next") + + async with Client(mcp) as client: + first = asyncio.create_task(client.call_tool("Hold")) + await started.wait() + second = asyncio.create_task(client.call_tool("Next", raise_on_error=False)) + await asyncio.sleep(0) + controller.state = "user" + release.set() + await first + result = await second + assert result.is_error + assert effects == ["hold"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tool_name", ["App", "PowerShell"]) +async def test_already_running_external_command_is_not_claimed_cancelled(tool_name): + controller = FakeController() + mcp = FastMCP("test") + mcp.add_middleware(ControlToolGate(controller, ControlNotifier(controller))) + started, release = asyncio.Event(), asyncio.Event() + effects = [] + + @mcp.tool(name=tool_name) + async def command(): + started.set() + await release.wait() + effects.append("external effect") + + async with Client(mcp) as client: + running = asyncio.create_task(client.call_tool(tool_name)) + await asyncio.wait_for(started.wait(), 1) + controller.state = "user" + blocked = await client.call_tool(tool_name, raise_on_error=False) + assert blocked.is_error + release.set() + await asyncio.wait_for(running, 1) + assert effects == ["external effect"] + + +@pytest.mark.asyncio +async def test_lease_invalidated_during_begin_does_not_run_tool(): + class Invalidated(FakeController): + def begin_call(self, name): + token = super().begin_call(name) + self.state = "user" # Equivalent to indicator failure/takeover in callback. + return token + + controller = Invalidated() + mcp = FastMCP("test") + mcp.add_middleware(ControlToolGate(controller, ControlNotifier(controller))) + effects = [] + + @mcp.tool(name="Action") + def action(): + effects.append("ran") + + async with Client(mcp) as client: + result = await client.call_tool("Action", raise_on_error=False) + assert result.is_error + assert "CONTROL_PREEMPTED" in str(result.content) + assert effects == [] + assert controller.calls == [("begin", "Action"), ("end", 7)] + + +class FakeSession: + def __init__(self, *, fail=False): + self.messages = [] + self.fail = fail + + async def send_log_message(self, **kwargs): + if self.fail: + raise RuntimeError("closed") + self.messages.append(kwargs) + + +@pytest.mark.asyncio +async def test_notifier_sends_takeover_and_recovery_and_prunes_failed_session(): + controller = FakeController() + notifier = ControlNotifier(controller) + notifier.start() + good, bad = FakeSession(), FakeSession(fail=True) + loop = asyncio.get_running_loop() + now = time.monotonic() + from windows_mcp.tools.control_notifications import _Session + + notifier._sessions[id(good)] = _Session(good, loop, now) + notifier._sessions[id(bad)] = _Session(bad, loop, now) + controller.emit("ai") + controller.emit("user") + await asyncio.sleep(0.05) + assert [m["data"]["event"] for m in good.messages] == ["USER_CONTROL"] + assert id(bad) not in notifier._sessions + controller.emit("ready") + await asyncio.sleep(0.05) + assert [m["data"]["event"] for m in good.messages] == [ + "USER_CONTROL", + "AI_CONTROL_AVAILABLE", + ] + await notifier.close() + controller.emit("user") + await asyncio.sleep(0) + assert len(good.messages) == 2 + + +@pytest.mark.asyncio +async def test_notifier_orders_cross_thread_events_by_generation(): + controller = FakeController() + notifier = ControlNotifier(controller) + notifier.start() + session = FakeSession() + from windows_mcp.tools.control_notifications import _Session + + notifier._sessions[id(session)] = _Session( + session, asyncio.get_running_loop(), time.monotonic() + ) + notifier._on_state({"state": "user", "generation": 3, "wait_seconds": 10.0}) + notifier._on_state({"state": "ai", "generation": 2, "wait_seconds": 0.0}) + await asyncio.sleep(0.06) + assert [m["data"]["event"] for m in session.messages] == ["USER_CONTROL"] + await notifier.close() + + +@pytest.mark.asyncio +async def test_notifier_close_cancels_and_awaits_blocked_send(): + controller = FakeController() + notifier = ControlNotifier(controller) + notifier.start() + started = asyncio.Event() + release = asyncio.Event() + cancelled = asyncio.Event() + delivered = [] + + class BlockingSession: + async def send_log_message(self, **kwargs): + started.set() + try: + await release.wait() + except asyncio.CancelledError: + cancelled.set() + raise + delivered.append(kwargs) + + session = BlockingSession() + from windows_mcp.tools.control_notifications import _Session + + notifier._sessions[id(session)] = _Session( + session, asyncio.get_running_loop(), time.monotonic() + ) + controller.emit("ai") + controller.emit("user") + await asyncio.wait_for(started.wait(), 1) + assert len(notifier._send_tasks) == 1 + + await asyncio.wait_for(notifier.close(), 1) + assert cancelled.is_set() + assert notifier._send_tasks == set() + assert delivered == [] + release.set() + await asyncio.sleep(0) + assert delivered == [] + + +@pytest.mark.asyncio +async def test_stateless_mode_registers_no_push_session(): + controller = FakeController() + notifier = ControlNotifier(controller) + notifier.push_enabled = False + notifier.start() + try: + context = SimpleNamespace(fastmcp_context=SimpleNamespace(session=FakeSession())) + notifier.remember(context) + assert notifier._sessions == {} + finally: + await notifier.close() + + +@pytest.mark.asyncio +async def test_expired_sessions_receive_no_push_and_status_filter_is_reserved(): + from windows_mcp.__main__ import _apply_tool_filter + + controller = FakeController() + notifier = ControlNotifier(controller, ttl=60) + notifier.start() + stale = FakeSession() + from windows_mcp.tools.control_notifications import _Session + + notifier._sessions[id(stale)] = _Session( + stale, asyncio.get_running_loop(), time.monotonic() - 61 + ) + controller.emit("ai") + controller.emit("user") + await asyncio.sleep(0.05) + assert stale.messages == [] + assert id(stale) not in notifier._sessions + await notifier.close() + + mcp = FastMCP("test") + + @mcp.tool(name="ControlStatus") + def status(): + return "status" + + @mcp.tool(name="Other") + def other(): + return "other" + + _apply_tool_filter(mcp, explicit_tools=["Other"], exclude_tools=None) + tools = await mcp.list_tools() + assert {tool.name for tool in tools} == {"ControlStatus", "Other"} + _apply_tool_filter(mcp, explicit_tools=None, exclude_tools=["ControlStatus", "Other"]) + tools = await mcp.list_tools() + assert {tool.name for tool in tools} == {"ControlStatus"} + + +@pytest.mark.asyncio +async def test_overlay_callback_never_waits_on_hook_thread(monkeypatch): + from windows_mcp import __main__ as wm + from windows_mcp.desktop import control as control_module, control_overlay + from windows_mcp.desktop import service as desktop_service + from windows_mcp.watchdog import service as watchdog_service + from windows_mcp import tools as tool_registry + + class Controller(FakeController): + _thread = None + fail_count = 0 + + def start(self): + self.state = "user" + + def stop(self): + self.state = "unavailable" + + def _fail_open(self): + self.fail_count += 1 + + def set_health_probe(self, probe): + self.health_probe = probe + + class Desktop: + tree = type("Tree", (), {"on_focus_change": lambda *args: None})() + + def get_screen_size(self): + return (100, 100) + + class WatchDog: + def set_focus_callback(self, callback): + pass + + def start(self): + pass + + def stop(self): + pass + + controller = Controller() + calls = [] + + def slow_show(active): + time.sleep(0.15) + calls.append(active) + + monkeypatch.setenv("ANONYMIZED_TELEMETRY", "false") + monkeypatch.setattr(wm, "_mcp", None) + monkeypatch.setattr(control_module, "get_controller", lambda: controller) + monkeypatch.setattr(tool_registry, "register_all", lambda *a, **k: None) + monkeypatch.setattr(desktop_service, "Desktop", Desktop) + monkeypatch.setattr(watchdog_service, "WatchDog", WatchDog) + monkeypatch.setattr(control_overlay, "start", lambda: None) + monkeypatch.setattr(control_overlay, "stop", lambda: None) + monkeypatch.setattr(control_overlay, "set_active", slow_show) + monkeypatch.setattr(control_overlay, "set_pending", lambda active: calls.append(active)) + + mcp = wm._build_mcp() + callback = controller.listeners[0] + async with mcp._lifespan(mcp): + assert controller.health_probe is control_overlay.is_healthy + # begin_call runs on the server loop and must wait for the indicator. + start = time.monotonic() + callback({"state": "ai"}) + assert time.monotonic() - start >= 0.14 + assert calls == [True] + + elapsed = [] + + def from_hook_thread(): + start = time.monotonic() + callback({"state": "user"}) + elapsed.append(time.monotonic() - start) + + thread = threading.Thread(target=from_hook_thread) + controller._thread = thread + thread.start() + thread.join(timeout=1) + assert elapsed and elapsed[0] < 0.1 + await asyncio.sleep(0.2) + assert calls == [True, False] + + monkeypatch.setattr( + control_overlay, + "set_active", + lambda active: (_ for _ in ()).throw(RuntimeError("lost")), + ) + callback({"state": "ai"}) + assert controller.fail_count == 1 + + +@pytest.mark.asyncio +async def test_live_mcp_client_receives_unsolicited_state_log(): + controller = FakeController() + notifier = ControlNotifier(controller) + mcp = FastMCP("test") + mcp.add_middleware(ControlToolGate(controller, notifier)) + received = [] + delivered = asyncio.Event() + + @mcp.tool(name="ControlStatus") + def status(): + return controller.status() + + async def on_log(message): + received.append(message) + delivered.set() + + notifier.start() + try: + async with Client(mcp, log_handler=on_log) as client: + await client.call_tool("ControlStatus") # Registers the live session. + controller.emit("ai") + controller.emit("user") + await asyncio.wait_for(delivered.wait(), timeout=2) + assert any("USER_CONTROL" in str(m.data) for m in received) + finally: + await notifier.close() + + +@pytest.mark.asyncio +async def test_stateful_http_client_receives_unsolicited_state_log(): + """Exercise the actual HTTP session transport, not only in-memory Client.""" + import uvicorn + + controller = FakeController() + notifier = ControlNotifier(controller) + mcp = FastMCP("test") + mcp.add_middleware(ControlToolGate(controller, notifier)) + received = [] + delivered = asyncio.Event() + + @mcp.tool(name="ControlStatus") + def status(): + return controller.status() + + async def on_log(message): + received.append(message) + delivered.set() + + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(128) + port = listener.getsockname()[1] + app = mcp.http_app(stateless_http=False) + server = uvicorn.Server(uvicorn.Config(app, log_level="error", lifespan="on")) + server_task = asyncio.create_task(server.serve(sockets=[listener])) + try: + for _ in range(100): + if server.started: + break + await asyncio.sleep(0.01) + assert server.started + notifier.start() + async with Client(f"http://127.0.0.1:{port}/mcp", log_handler=on_log) as client: + await client.call_tool("ControlStatus") + controller.emit("ai") + controller.emit("user") + await asyncio.wait_for(delivered.wait(), 2) + assert any("USER_CONTROL" in str(message.data) for message in received) + finally: + await notifier.close() + server.should_exit = True + await asyncio.wait_for(server_task, 3) diff --git a/tests/test_control_overlay.py b/tests/test_control_overlay.py new file mode 100644 index 00000000..68d5bd8a --- /dev/null +++ b/tests/test_control_overlay.py @@ -0,0 +1,380 @@ +"""Indicator geometry, lifecycle, and screenshot exclusion checks.""" + +import threading +import time +from contextlib import contextmanager +from unittest.mock import Mock + +from PIL import Image +import pytest + +from windows_mcp.desktop import control_overlay, flash_overlay +from windows_mcp.desktop.service import Desktop + + +@pytest.fixture(autouse=True) +def _stop_indicator(): + control_overlay.stop() + yield + control_overlay.stop() + + +def test_edge_and_cursor_are_translucent(): + edge = control_overlay._edge_bitmap(20, 28, "left", (45, 145, 255)) + assert len(edge) == 20 * 28 * 4 + assert edge[3] > edge[-1] + ring = control_overlay._cursor_bitmap((45, 145, 255)) + assert len(ring) == control_overlay._CURSOR_SIZE**2 * 4 + assert ring[3] == 0 # no opaque square over the desktop + + +def test_multimonitor_creates_narrow_clickthrough_layers(monkeypatch): + made = [] + + class FakeLayer: + def __init__(self, x, y, width, height, bgra, name): + made.append((x, y, width, height, name)) + + def close(self): + pass + + monkeypatch.setattr(control_overlay, "_Layer", FakeLayer) + monkeypatch.setattr(control_overlay._user32, "GetCursorPos", lambda ptr: True) + rects = ((-1920, 0, 0, 1080), (0, 0, 2560, 1440)) + edges, ring = control_overlay._build_layers(rects, pending=False) + assert len(edges) == 8 + assert ring is not None + assert made[0][:4] == (-1920, 0, 1920, 28) + assert made[4][:4] == (0, 0, 2560, 28) + assert all(w <= 28 or h <= 28 for _, _, w, h, name in made if name != "cursor") + + +def test_layer_failure_closes_prior_windows(monkeypatch): + closed = [] + + class FailingLayer: + count = 0 + + def __init__(self, *args): + self.count = FailingLayer.count + FailingLayer.count += 1 + if self.count == 2: + raise OSError("window creation failed") + + def close(self): + closed.append(self.count) + + monkeypatch.setattr(control_overlay, "_Layer", FailingLayer) + with pytest.raises(OSError, match="window creation failed"): + control_overlay._build_layers(((0, 0, 800, 600),), pending=False) + assert closed == [1, 0] + + +def test_capture_suspends_and_restores_including_error(monkeypatch): + calls = [] + + class FakeIndicator: + def change(self, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(control_overlay, "_instance", FakeIndicator()) + with pytest.raises(ValueError, match="capture failed"): + with control_overlay.suspend_for_capture(): + calls.append("capture") + raise ValueError("capture failed") + assert calls == [{"suspended": 1}, "capture", {"suspended": -1}] + + +def test_indicator_lifecycle_and_no_restore_after_takeover(monkeypatch): + events = [] + + class FakeLayer: + hwnd = 1 + + def show(self): + events.append("show") + + def hide(self): + events.append("hide") + + def move(self, x, y): + pass + + def close(self): + events.append("close") + + monkeypatch.setattr(control_overlay, "_monitor_rects", lambda: ((0, 0, 800, 600),)) + monkeypatch.setattr( + control_overlay, "_build_layers", lambda rects, pending: ([FakeLayer()], FakeLayer()) + ) + monkeypatch.setattr(control_overlay._user32, "GetCursorPos", lambda ptr: True) + monkeypatch.setattr(flash_overlay, "_pump_messages", lambda hwnd: None) + assert control_overlay.is_healthy() is False + control_overlay.start() + assert control_overlay.is_healthy() is True + control_overlay.set_active(True) + assert events.count("show") == 2 + with control_overlay.suspend_for_capture(): + assert events.count("hide") == 2 + control_overlay.set_active(False) + assert events.count("show") == 2 + control_overlay.stop() + assert control_overlay.is_healthy() is False + assert events.count("close") == 2 + + +def test_failed_indicator_reports_unhealthy(monkeypatch): + class FailedIndicator: + thread = threading.current_thread() + stopping = False + error = RuntimeError("window failed") + + def stop(self): + pass + + monkeypatch.setattr(control_overlay, "_instance", FailedIndicator()) + assert control_overlay.is_healthy() is False + + +def test_alive_but_stalled_indicator_reports_unhealthy(monkeypatch): + class StalledIndicator: + thread = threading.current_thread() # still alive, but no heartbeat refresh + stopping = False + error = None + heartbeat = time.monotonic() - 2.0 + + def stop(self): + pass + + monkeypatch.setattr(control_overlay, "_instance", StalledIndicator()) + assert control_overlay.is_healthy() is False + + +def test_capture_ack_timeout_cannot_hide_recovered_indicator(monkeypatch): + indicator = control_overlay._Indicator() + indicator.thread = Mock() + indicator.thread.is_alive.return_value = True + indicator.heartbeat = time.monotonic() + indicator.stop = Mock() + monkeypatch.setattr(control_overlay, "_instance", indicator) + monkeypatch.setattr(threading.Condition, "wait_for", lambda self, predicate, timeout: False) + + with pytest.raises(RuntimeError, match="did not acknowledge"): + with control_overlay.suspend_for_capture(): + pytest.fail("Capture must not run before the indicator is hidden") + assert indicator.suspended == 1 + # Even if the window owner later beats again, the incomplete transition + # must never be accepted as a healthy, visible AI lease. + indicator.heartbeat = time.monotonic() + assert control_overlay.is_healthy() is False + with pytest.raises(RuntimeError, match="indicator failed"): + control_overlay.set_active(True) + + +def test_actual_owner_thread_stall_expires_heartbeat(monkeypatch): + entered = threading.Event() + release = threading.Event() + + def stalled_monitor_lookup(): + entered.set() + release.wait(timeout=3.0) + raise RuntimeError("simulated window owner failure") + + monkeypatch.setattr(control_overlay, "_monitor_rects", stalled_monitor_lookup) + control_overlay.start() + indicator = control_overlay._instance + assert indicator is not None + with indicator.condition: + indicator.active = True + indicator.condition.notify_all() + assert entered.wait(timeout=1.0) + try: + time.sleep(1.1) + assert indicator.thread.is_alive() + assert control_overlay.is_healthy() is False + finally: + release.set() + control_overlay.stop() + + +def test_fresh_indicator_heartbeat_reports_healthy(monkeypatch): + class FreshIndicator: + thread = threading.current_thread() + stopping = False + error = None + heartbeat = time.monotonic() + + def stop(self): + pass + + monkeypatch.setattr(control_overlay, "_instance", FreshIndicator()) + assert control_overlay.is_healthy() is True + + +def test_stop_failure_retains_owner_for_retry(monkeypatch): + class SlowOwner: + calls = 0 + + def stop(self): + self.calls += 1 + if self.calls == 1: + raise RuntimeError("join timed out") + + owner = SlowOwner() + monkeypatch.setattr(control_overlay, "_instance", owner) + with pytest.raises(RuntimeError, match="join timed out"): + control_overlay.stop() + assert control_overlay._instance is owner + control_overlay.stop() + assert control_overlay._instance is None + assert owner.calls == 2 + + +def test_failed_start_keeps_uncleaned_owner_and_prevents_stacking(monkeypatch): + created = [] + + class FailedOwner: + def __init__(self): + self.stopping = False + self.error = None + self.heartbeat = time.monotonic() + self.thread = threading.current_thread() + self.stop_calls = 0 + created.append(self) + + def start(self): + raise RuntimeError("startup failed") + + def stop(self): + self.stopping = True + self.stop_calls += 1 + if self.stop_calls == 1: + raise RuntimeError("join timed out") + + monkeypatch.setattr(control_overlay, "_Indicator", FailedOwner) + with pytest.raises(RuntimeError, match="cleanup is pending"): + control_overlay.start() + assert control_overlay._instance is created[0] + with pytest.raises(RuntimeError, match="requires stop"): + control_overlay.start() + assert len(created) == 1 + control_overlay.stop() + assert control_overlay._instance is None + + +def test_failed_start_clears_owner_after_successful_cleanup(monkeypatch): + class FailedOwner: + def start(self): + raise RuntimeError("startup failed") + + def stop(self): + pass + + monkeypatch.setattr(control_overlay, "_Indicator", FailedOwner) + with pytest.raises(RuntimeError, match="startup failed"): + control_overlay.start() + assert control_overlay._instance is None + + +def test_two_captures_are_serialized(monkeypatch): + active = 0 + peak = 0 + barrier = threading.Barrier(3) + + class FakeIndicator: + def change(self, **kwargs): + nonlocal active, peak + active += kwargs["suspended"] + peak = max(peak, active) + + monkeypatch.setattr(control_overlay, "_instance", FakeIndicator()) + + def capture(): + barrier.wait() + with control_overlay.suspend_for_capture(): + threading.Event().wait(0.02) + + threads = [threading.Thread(target=capture) for _ in range(2)] + for thread in threads: + thread.start() + barrier.wait() + for thread in threads: + thread.join() + assert active == 0 + assert peak == 1 + + +@pytest.mark.parametrize("backend", ["dxcam", "mss", "pillow"]) +def test_screenshot_capture_runs_only_while_indicator_hidden(monkeypatch, backend): + events = [] + + @contextmanager + def suspended(): + events.append("hide") + try: + yield + finally: + events.append("restore") + + monkeypatch.setattr(control_overlay, "suspend_for_capture", suspended) + monkeypatch.setattr( + flash_overlay, "cancel_active_flash", lambda: events.append("cancel flash") or True + ) + monkeypatch.setattr( + flash_overlay, "show_capture_flash", lambda rect: events.append("show flash") + ) + monkeypatch.setattr( + "windows_mcp.desktop.service.screenshot_capture.capture", + lambda rect: (events.append("capture") or Image.new("RGB", (2, 2)), backend), + ) + desktop = Desktop.__new__(Desktop) + desktop.get_screenshot() + assert desktop._last_screenshot_backend == backend + assert events == ["hide", "cancel flash", "capture", "restore", "show flash"] + + +def test_screenshot_error_restores_indicator_without_flash(monkeypatch): + events = [] + + @contextmanager + def suspended(): + events.append("hide") + try: + yield + finally: + events.append("restore") + + monkeypatch.setattr(control_overlay, "suspend_for_capture", suspended) + monkeypatch.setattr(flash_overlay, "cancel_active_flash", lambda: True) + monkeypatch.setattr(flash_overlay, "show_capture_flash", lambda rect: events.append("flash")) + + def failing_capture(rect): + events.append("capture") + raise OSError("capture failed") + + monkeypatch.setattr("windows_mcp.desktop.service.screenshot_capture.capture", failing_capture) + with pytest.raises(OSError, match="capture failed"): + Desktop.__new__(Desktop).get_screenshot() + assert events == ["hide", "capture", "restore"] + + +def test_unclosed_flash_prevents_new_capture(monkeypatch): + events = [] + + @contextmanager + def suspended(): + events.append("hide") + try: + yield + finally: + events.append("restore") + + monkeypatch.setattr(control_overlay, "suspend_for_capture", suspended) + monkeypatch.setattr(flash_overlay, "cancel_active_flash", lambda: False) + monkeypatch.setattr( + "windows_mcp.desktop.service.screenshot_capture.capture", + lambda rect: events.append("capture"), + ) + with pytest.raises(RuntimeError, match="did not close"): + Desktop.__new__(Desktop).get_screenshot() + assert events == ["hide", "restore"] diff --git a/tests/test_deterministic_drag_desktop.py b/tests/test_deterministic_drag_desktop.py index 65bc945c..23162d54 100644 --- a/tests/test_deterministic_drag_desktop.py +++ b/tests/test_deterministic_drag_desktop.py @@ -1,6 +1,8 @@ import pytest +from unittest.mock import Mock from windows_mcp.desktop import service +from windows_mcp.desktop.control_ledger import InputLedger from windows_mcp.desktop.service import Desktop @@ -10,20 +12,27 @@ def _desktop() -> Desktop: return desktop +def _stub_control(monkeypatch: pytest.MonkeyPatch) -> None: + """Give drag tests an isolated AI-input ledger without real mouse injection.""" + owner = Mock() + owner.input_ledger = InputLedger() + owner.physical_mouse_down.return_value = False + monkeypatch.setattr(service, "get_controller", lambda: owner) + + def test_desktop_drag_uses_explicit_start_and_duration( monkeypatch: pytest.MonkeyPatch, ) -> None: - calls: list[tuple[int, int, int, int, int, float | None]] = [] + calls: list[tuple[str, tuple[int, ...]]] = [] desktop = _desktop() + _stub_control(monkeypatch) monkeypatch.setattr(service, "sleep", lambda seconds: None) monkeypatch.setattr( - service.uia, - "DragDrop", - lambda x1, y1, x2, y2, moveSpeed=1, duration=None: calls.append( - (x1, y1, x2, y2, moveSpeed, duration) - ), + service.uia, "PressMouse", lambda x, y, **kw: calls.append(("down", (x, y))) ) + monkeypatch.setattr(service.uia, "SetCursorPos", lambda x, y: calls.append(("move", (x, y)))) + monkeypatch.setattr(service.uia, "ReleaseMouse", lambda **kw: calls.append(("up", ()))) result = desktop.drag( [100, 200], @@ -31,27 +40,31 @@ def test_desktop_drag_uses_explicit_start_and_duration( duration="0.25", ) - assert calls == [(10, 20, 100, 200, 1, 0.25)] + assert calls[0] == ("down", (10, 20)) + assert calls[-2:] == [("move", (100, 200)), ("up", ())] + assert len([call for call in calls if call[0] == "move"]) == 25 assert result["start"] == [10, 20] assert result["end"] == [100, 200] assert result["duration"] == 0.25 def test_desktop_drag_legacy_start_uses_current_cursor(monkeypatch: pytest.MonkeyPatch) -> None: - calls: list[tuple[int, int, int, int]] = [] + calls: list[tuple[str, tuple[int, ...]]] = [] desktop = _desktop() + _stub_control(monkeypatch) monkeypatch.setattr(service, "sleep", lambda seconds: None) monkeypatch.setattr(service.uia, "GetCursorPos", lambda: (7, 8)) monkeypatch.setattr( - service.uia, - "DragDrop", - lambda x1, y1, x2, y2, **kwargs: calls.append((x1, y1, x2, y2)), + service.uia, "PressMouse", lambda x, y, **kw: calls.append(("down", (x, y))) ) + monkeypatch.setattr(service.uia, "SetCursorPos", lambda x, y: calls.append(("move", (x, y)))) + monkeypatch.setattr(service.uia, "ReleaseMouse", lambda **kw: calls.append(("up", ()))) result = desktop.drag((30, 40)) - assert calls == [(7, 8, 30, 40)] + assert calls[0] == ("down", (7, 8)) + assert calls[-2:] == [("move", (30, 40)), ("up", ())] assert result["start"] == [7, 8] assert result["duration"] is None diff --git a/tests/test_flash_overlay.py b/tests/test_flash_overlay.py index 1393f14c..e8d975d6 100644 --- a/tests/test_flash_overlay.py +++ b/tests/test_flash_overlay.py @@ -7,6 +7,8 @@ import sys import threading +import ctypes +from ctypes import wintypes from unittest.mock import patch import pytest @@ -19,12 +21,17 @@ def _reset_active_overlay(): """Each test starts and ends with no overlay registered.""" with flash_overlay._lock: flash_overlay._active_overlay = None + flash_overlay._all_overlays.clear() yield with flash_overlay._lock: ov = flash_overlay._active_overlay flash_overlay._active_overlay = None + all_overlays = tuple(flash_overlay._all_overlays) + flash_overlay._all_overlays.clear() if ov is not None: ov.stop_event.set() + for pending in all_overlays: + pending.stop_event.set() class TestFlashDisabled: @@ -43,6 +50,15 @@ def test_falsy_values_keep_enabled(self, monkeypatch, value): assert flash_overlay._flash_disabled() is False +def test_window_api_signatures_preserve_x64_handles(): + """Without argtypes ctypes narrows Python integers to 32-bit c_int.""" + assert ctypes.sizeof(flash_overlay._kernel32.GetModuleHandleW.restype) == ctypes.sizeof( + ctypes.c_void_p + ) + assert flash_overlay._user32.CreateWindowExW.argtypes[10] is wintypes.HINSTANCE + assert flash_overlay._user32.UnregisterClassW.argtypes[1] is wintypes.HINSTANCE + + class _FakeRect: """Stand-in for ``uia.Rect`` with the four corner attributes.""" @@ -210,6 +226,22 @@ def test_signals_stop_and_clears_active(self, monkeypatch): assert overlay.stop_event.is_set() assert flash_overlay._active_overlay is None + def test_rejects_capture_when_replaced_flash_has_not_closed(self): + old = flash_overlay._Overlay() + latest = flash_overlay._Overlay() + latest.closed_event.set() + with flash_overlay._lock: + flash_overlay._all_overlays.update((old, latest)) + flash_overlay._active_overlay = latest + + assert flash_overlay.cancel_active_flash(timeout=0) is False + assert old.stop_event.is_set() + assert latest.stop_event.is_set() + assert flash_overlay._active_overlay is None + + old.closed_event.set() + assert flash_overlay.cancel_active_flash(timeout=0) is True + class TestIntensityCurve: def test_full_screen_is_bell_curve(self): diff --git a/tests/test_stdio_handshake.py b/tests/test_stdio_handshake.py index 2b0e2808..847ea7a7 100644 --- a/tests/test_stdio_handshake.py +++ b/tests/test_stdio_handshake.py @@ -47,6 +47,7 @@ "App", "Click", "Clipboard", + "ControlStatus", "DisplayInventory", "FileSystem", "Move", @@ -101,12 +102,14 @@ async def test_handshake_completes_and_lists_tools() -> None: async def test_tool_call_round_trips_over_stdio() -> None: """A tool call survives the full client -> pipe -> server -> pipe path. - Uses Wait, the only state-changing-free tool with a deterministic result. + Uses ControlStatus, which remains available during user takeover and + when the test host cannot show the required desktop indicator. """ async with Client(_transport()) as client: - result = await asyncio.wait_for(client.call_tool("Wait", {"duration": 1}), STARTUP_TIMEOUT) + result = await asyncio.wait_for(client.call_tool("ControlStatus", {}), STARTUP_TIMEOUT) assert result.is_error is False + assert result.data["state"] in {"user", "ready", "ai", "unavailable"} async def test_request_before_initialize_is_rejected() -> None: From eace53bd6981264c479202be84217314f314911c Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Sat, 3 Oct 2026 01:55:29 +0800 Subject: [PATCH 2/9] Guard desktop takeover with visible indicator and capture pause --- src/windows_mcp/__main__.py | 11 +- src/windows_mcp/desktop/control.py | 166 +++++++++------------ src/windows_mcp/desktop/control_hooks.py | 110 ++++++++++++++ src/windows_mcp/desktop/control_overlay.py | 5 +- src/windows_mcp/desktop/flash_overlay.py | 28 +--- src/windows_mcp/desktop/overlay_bitmap.py | 27 ++++ src/windows_mcp/desktop/service.py | 20 ++- tests/test_control.py | 108 ++++++++++++++ tests/test_control_mcp_gate.py | 30 ++++ tests/test_control_overlay.py | 48 +++++- 10 files changed, 421 insertions(+), 132 deletions(-) create mode 100644 src/windows_mcp/desktop/control_hooks.py create mode 100644 src/windows_mcp/desktop/overlay_bitmap.py diff --git a/src/windows_mcp/__main__.py b/src/windows_mcp/__main__.py index 3db34477..7a471d32 100755 --- a/src/windows_mcp/__main__.py +++ b/src/windows_mcp/__main__.py @@ -277,6 +277,8 @@ def apply_control_state(status: dict) -> None: return if status["state"] == "ai": control_overlay.set_active(True) + if status.get("generation") is not None: + controller.arm_visible(status["generation"]) elif status["state"] == "takeover_pending": control_overlay.set_pending(True) else: @@ -289,11 +291,16 @@ def apply_control_state(status: dict) -> None: def show_control_state(status: dict) -> None: if threading.current_thread() is controller._thread: # Never wait for a window-owner acknowledgement on the hook thread. - # Windows can silently remove a low level hook after a slow callback. + # A capture may hold the overlay lock for an unbounded backend call; + # offload the wait so ControlStatus stays responsive on the loop. try: if control_loop is None or control_loop.is_closed(): raise RuntimeError("server event loop unavailable") - control_loop.call_soon_threadsafe(apply_control_state, status.copy()) + control_loop.call_soon_threadsafe( + lambda: asyncio.create_task( + asyncio.to_thread(apply_control_state, status.copy()) + ) + ) except RuntimeError: controller._fail_open() logger.exception("Could not schedule AI control indicator update") diff --git a/src/windows_mcp/desktop/control.py b/src/windows_mcp/desktop/control.py index bbe1d06a..19da0666 100644 --- a/src/windows_mcp/desktop/control.py +++ b/src/windows_mcp/desktop/control.py @@ -14,14 +14,12 @@ from typing import Callable from windows_mcp.desktop.control_context import current_token, current_steps, get_step_count from windows_mcp.desktop.control_ledger import InputLedger +from windows_mcp.desktop import control_hooks from windows_mcp.desktop.control_win32 import ( - _CHORD, - _MouseHookData, - _KeyHookData, + _MouseHookData as _MouseHookData, # Keep hook test helpers on this module. + _KeyHookData as _KeyHookData, _user32, _key, - _mouse_button, - _read_raw_mouse, _interactive_desktop, run_input_monitor, ) @@ -52,6 +50,9 @@ def __init__(self, mouse_takeover_units: int = 120, mouse_takeover_pixels: int = self._listeners: list[Callable[[dict], None]] = [] self.input_ledger = InputLedger() self._health_probe: Callable[[], bool] | None = None + self._visual_armed = False + self._visual_paused = False + self._visual_event_stamp = 0.0 self._events: queue.Queue[tuple] = queue.Queue(maxsize=4096) self._stop = threading.Event() self._release_requested = threading.Event() @@ -86,6 +87,53 @@ def set_health_probe(self, probe: Callable[[], bool]) -> None: """Require the visual indicator to stay alive during AI control.""" self._health_probe = probe + def arm_visible(self, generation: int) -> None: + """Suppress physical input only after the AI indicator is visible.""" + with self._lock: + if self._state != "ai" or self._generation != generation or self._visual_paused: + return + if ( + self._last_physical_event != self._visual_event_stamp + or self._pressed + or self._mouse_down + or self._fast_takeover + or self._fast_pending + or self._emergency + ): + return + self.input_ledger.enable() + self._visual_armed = True + self._deadline = time.monotonic() + 1.0 + self._suppress = True + if self._last_physical_event != self._visual_event_stamp: + self._suppress = False + self._visual_armed = False + self.input_ledger.block_new() + + def pause_for_capture(self) -> int | None: + """Release physical input before an AI indicator is hidden for capture.""" + with self._lock: + if self._state != "ai" or not self._visual_armed: + return None + self._suppress = False + self._visual_armed = False + self._visual_paused = True + self._visual_event_stamp = self._last_physical_event + self.input_ledger.block_new() + return self._generation + + def resume_after_capture(self, generation: int | None, *, restored: bool) -> None: + """Re-arm only after the indicator is restored and no user event occurred.""" + if generation is None: + return + if not restored or (self._health_probe is not None and not self._health_probe()): + self._fail_open() + return + with self._lock: + if generation == self._generation: + self._visual_paused = False + self.arm_visible(generation) + def _snapshot_locked(self, now: float) -> dict: state = ( "user" @@ -138,12 +186,15 @@ def _set_locked(self, state: str) -> bool: if state != "ai": self.input_ledger.block_new() self._generation += 1 + self._visual_armed = state == "ai" and self._health_probe is None + self._visual_paused = False + self._visual_event_stamp = self._last_physical_event self._point_origin = None self._raw_origin.clear() self._raw_device = None if state != "takeover_pending": self._fast_pending = False - self._suppress = state in ("ai", "takeover_pending") and not self._emergency + self._suppress = self._visual_armed and not self._emergency self._deadline = time.monotonic() + 1.0 if self._suppress else 0.0 return True @@ -186,6 +237,7 @@ def begin_call(self, name: str) -> int: or self._fast_takeover or self._fast_pending or self._rotating + or self._visual_paused or self._state not in ("ready", "ai") ): status = self._snapshot_locked(now) @@ -236,7 +288,7 @@ def record_step_current(self) -> None: def physical_key_down(self, vk: int) -> bool: # A swallowed physical down did not reach the app, so the AI must # release its own injected down when its tool is preempted. - return _key(vk) in self._delivered_keys + return _key(vk) in {_key(held) for held in self._delivered_keys} def physical_mouse_down(self, button: str) -> bool: return {"left": 1, "right": 2, "middle": 3}[button] in self._delivered_mouse @@ -372,97 +424,13 @@ def _queue(self, event: tuple) -> None: self._fail_open() def _physical_mouse(self, code, wparam, lparam): - if code < 0: - return _user32.CallNextHookEx(self._mouse_hook, code, wparam, lparam) - valid = False - button = None - try: - data = ctypes.cast(lparam, ctypes.POINTER(_MouseHookData)).contents - if data.flags & 1: # LL mouse injection flags. - return _user32.CallNextHookEx(self._mouse_hook, code, wparam, lparam) - valid = True - self._last_physical_event = time.monotonic() - if wparam == 0x200 and self._suppress: - self._fast_pending = True # Pause AI before queued movement is processed. - self.input_ledger.block_new() - button = _mouse_button(wparam, data.mouseData) - if button: - if button[1]: - self._mouse_down.add(button[0]) - else: - self._mouse_down.discard(button[0]) - self._queue(("point", data.pt.x, data.pt.y, int(wparam))) - if self._suppress and not self._emergency: - if time.monotonic() <= self._deadline: - return 1 - self._fail_open() - except Exception: - self._fail_open() - if valid and button: - if button[1]: - self._delivered_mouse.add(button[0]) - else: - self._delivered_mouse.discard(button[0]) - return _user32.CallNextHookEx(self._mouse_hook, code, wparam, lparam) + return control_hooks.physical_mouse(self, code, wparam, lparam) def _physical_key(self, code, wparam, lparam): - if code < 0: - return _user32.CallNextHookEx(self._key_hook, code, wparam, lparam) - vk = None - down = False - try: - data = ctypes.cast(lparam, ctypes.POINTER(_KeyHookData)).contents - if data.flags & 0x10: # LLKHF_INJECTED: AI SendInput stays usable. - return _user32.CallNextHookEx(self._key_hook, code, wparam, lparam) - self._last_physical_event = time.monotonic() - vk = _key(data.vkCode) - down = wparam in (0x100, 0x104) - was_down = vk in self._pressed - if down: - self._pressed.add(vk) - else: - self._pressed.discard(vk) - if vk in self._quarantine: - if not down: - self._quarantine.discard(vk) - self._queue(("key",)) # Idle starts after the last chord key is released. - return 1 - if ( - down - and vk == 0x08 - and not was_down - and not self._fast_takeover - and self._suppress - and _CHORD.issubset(self._pressed) - ): - self._quarantine.update(self._pressed & _CHORD) - self._fast_takeover = True - self._suppress = False # Release first, then tell the coordinator. - self.input_ledger.block_new() - self._queue(("hotkey",)) - return 1 - self._queue(("key",)) - if self._suppress and not self._emergency: - if time.monotonic() <= self._deadline: - return 1 - self._fail_open() - except Exception: - self._fail_open() - if vk is not None and down: - self._delivered_keys.add(vk) - elif vk is not None: - self._delivered_keys.discard(vk) - return _user32.CallNextHookEx(self._key_hook, code, wparam, lparam) + return control_hooks.physical_key(self, code, wparam, lparam) def _raw_input(self, lparam) -> None: - raw = _read_raw_mouse(lparam) - if raw is None: - return # Null devices can be touchpads; do not infer source. - self._last_physical_event = time.monotonic() - if self._suppress: - self._fast_pending = True - self.input_ledger.block_new() - self._queue(("raw", *raw)) + control_hooks.raw_input(self, lparam) def _handle(self, event: tuple) -> None: now = time.monotonic() @@ -472,6 +440,14 @@ def _handle(self, event: tuple) -> None: else: changed = self._tick_locked(now) kind = event[0] + if ( + kind in ("key", "point", "raw") + and self._state == "ai" + and not self._visual_armed + ): + # Physical input was delivered while the indicator was not armed. + self._last_user = now + changed = self._set_locked("user") or changed if kind == "hotkey" and self._state in ("ai", "takeover_pending"): self._last_user = now changed = self._set_locked("user") or changed @@ -565,7 +541,7 @@ def _watch(self) -> None: self._fail_open() continue try: - if self._state in ("ai", "takeover_pending") and not self._fast_takeover: + if self._state == "ai" and self._visual_armed and not self._fast_takeover: self._deadline = now + 1.0 self._suppress = True else: diff --git a/src/windows_mcp/desktop/control_hooks.py b/src/windows_mcp/desktop/control_hooks.py new file mode 100644 index 00000000..2f2b7d64 --- /dev/null +++ b/src/windows_mcp/desktop/control_hooks.py @@ -0,0 +1,110 @@ +"""Fast physical-input hook callbacks; never wait for UI or MCP work.""" + +import ctypes +import time + +from windows_mcp.desktop.control_win32 import ( + _CHORD, + _KeyHookData, + _MouseHookData, + _key, + _mouse_button, + _read_raw_mouse, + _user32, +) + + +def physical_mouse(owner, code, wparam, lparam): + if code < 0: + return _user32.CallNextHookEx(owner._mouse_hook, code, wparam, lparam) + valid = False + button = None + try: + data = ctypes.cast(lparam, ctypes.POINTER(_MouseHookData)).contents + if data.flags & 1: # LL mouse injection flags. + return _user32.CallNextHookEx(owner._mouse_hook, code, wparam, lparam) + valid = True + owner._last_physical_event = time.monotonic() + if wparam == 0x200 and owner._suppress: + owner._fast_pending = True # Pause AI before queued movement is processed. + owner.input_ledger.block_new() + button = _mouse_button(wparam, data.mouseData) + if button: + if button[1]: + owner._mouse_down.add(button[0]) + else: + owner._mouse_down.discard(button[0]) + owner._queue(("point", data.pt.x, data.pt.y, int(wparam))) + if owner._suppress and not owner._emergency: + if time.monotonic() <= owner._deadline: + return 1 + owner._fail_open() + except Exception: + owner._fail_open() + if valid and button: + if button[1]: + owner._delivered_mouse.add(button[0]) + else: + owner._delivered_mouse.discard(button[0]) + return _user32.CallNextHookEx(owner._mouse_hook, code, wparam, lparam) + + +def physical_key(owner, code, wparam, lparam): + if code < 0: + return _user32.CallNextHookEx(owner._key_hook, code, wparam, lparam) + vk = None + down = False + try: + data = ctypes.cast(lparam, ctypes.POINTER(_KeyHookData)).contents + if data.flags & 0x10: # LLKHF_INJECTED: AI SendInput stays usable. + return _user32.CallNextHookEx(owner._key_hook, code, wparam, lparam) + owner._last_physical_event = time.monotonic() + vk = int(data.vkCode) + down = wparam in (0x100, 0x104) + was_down = vk in owner._pressed + if down: + owner._pressed.add(vk) + else: + owner._pressed.discard(vk) + if vk in owner._quarantine: + if not down: + owner._quarantine.discard(vk) + owner._queue(("key",)) # Idle starts after the last chord key is released. + return 1 + if ( + down + and vk == 0x08 + and not was_down + and not owner._fast_takeover + and owner._suppress + and _CHORD.issubset({_key(held) for held in owner._pressed}) + ): + owner._quarantine.update(held for held in owner._pressed if _key(held) in _CHORD) + owner._fast_takeover = True + owner._suppress = False # Release first, then tell the coordinator. + owner.input_ledger.block_new() + owner._queue(("hotkey",)) + return 1 + owner._queue(("key",)) + if owner._suppress and not owner._emergency: + if time.monotonic() <= owner._deadline: + return 1 + owner._fail_open() + except Exception: + owner._fail_open() + if vk is not None and down: + owner._delivered_keys.add(vk) + elif vk is not None: + owner._delivered_keys.discard(vk) + return _user32.CallNextHookEx(owner._key_hook, code, wparam, lparam) + + +def raw_input(owner, lparam) -> None: + raw = _read_raw_mouse(lparam) + if raw is None: + return # Null devices can be touchpads; do not infer source. + owner._last_physical_event = time.monotonic() + if owner._suppress: + owner._fast_pending = True + owner.input_ledger.block_new() + owner._queue(("raw", *raw)) diff --git a/src/windows_mcp/desktop/control_overlay.py b/src/windows_mcp/desktop/control_overlay.py index b713df49..e35591a5 100644 --- a/src/windows_mcp/desktop/control_overlay.py +++ b/src/windows_mcp/desktop/control_overlay.py @@ -14,6 +14,7 @@ from PIL import Image, ImageDraw, ImageFilter from windows_mcp.desktop import flash_overlay +from windows_mcp.desktop.overlay_bitmap import premultiplied_bgra from windows_mcp import uia logger = logging.getLogger(__name__) @@ -51,7 +52,7 @@ def _edge_bitmap(width: int, height: int, side: str, color: tuple[int, int, int] else: x = width - 1 - depth draw.line((x, 0, x, height - 1), fill=(*color, alpha)) - return flash_overlay._premultiplied_bgra(image, 1.0) + return premultiplied_bgra(image, 1.0) def _cursor_bitmap(color: tuple[int, int, int]) -> bytes: @@ -61,7 +62,7 @@ def _cursor_bitmap(color: tuple[int, int, int]) -> bytes: box = (inset, inset, _CURSOR_SIZE - inset - 1, _CURSOR_SIZE - inset - 1) draw.ellipse(box, outline=(*color, 210), width=5) glow = image.filter(ImageFilter.GaussianBlur(9)) - return flash_overlay._premultiplied_bgra(Image.alpha_composite(glow, image), 1.0) + return premultiplied_bgra(Image.alpha_composite(glow, image), 1.0) class _Layer: diff --git a/src/windows_mcp/desktop/flash_overlay.py b/src/windows_mcp/desktop/flash_overlay.py index c5bb89ae..b755729b 100644 --- a/src/windows_mcp/desktop/flash_overlay.py +++ b/src/windows_mcp/desktop/flash_overlay.py @@ -18,6 +18,8 @@ import time from ctypes import wintypes +from windows_mcp.desktop.overlay_bitmap import premultiplied_bgra as _premultiplied_bgra + logger = logging.getLogger(__name__) _FLASH_RGB = (0xFF, 0x45, 0x00) @@ -344,32 +346,6 @@ def _render_glow_rgba( return Image.alpha_composite(blurred, sharp) -def _premultiplied_bgra(rgba_image, intensity: float) -> bytes: - """Convert PIL RGBA to BGRA premultiplied bytes scaled by ``intensity``.""" - bgra = bytearray(rgba_image.tobytes("raw", "BGRA")) - if intensity >= 1.0: - for i in range(0, len(bgra), 4): - a = bgra[i + 3] - if a == 0: - continue - bgra[i] = (bgra[i] * a) // 255 - bgra[i + 1] = (bgra[i + 1] * a) // 255 - bgra[i + 2] = (bgra[i + 2] * a) // 255 - else: - for i in range(0, len(bgra), 4): - a = (bgra[i + 3] * int(intensity * 255)) // 255 - bgra[i + 3] = a - if a == 0: - bgra[i] = 0 - bgra[i + 1] = 0 - bgra[i + 2] = 0 - continue - bgra[i] = (bgra[i] * a) // 255 - bgra[i + 1] = (bgra[i + 1] * a) // 255 - bgra[i + 2] = (bgra[i + 2] * a) // 255 - return bytes(bgra) - - # --------------------------------------------------------------------------- # Window management # --------------------------------------------------------------------------- diff --git a/src/windows_mcp/desktop/overlay_bitmap.py b/src/windows_mcp/desktop/overlay_bitmap.py new file mode 100644 index 00000000..5a154440 --- /dev/null +++ b/src/windows_mcp/desktop/overlay_bitmap.py @@ -0,0 +1,27 @@ +"""Shared pixel conversion for the control indicator and capture flash.""" + + +def premultiplied_bgra(rgba_image, intensity: float) -> bytes: + """Convert a PIL RGBA image to scaled, premultiplied BGRA bytes.""" + bgra = bytearray(rgba_image.tobytes("raw", "BGRA")) + if intensity >= 1.0: + for i in range(0, len(bgra), 4): + a = bgra[i + 3] + if a == 0: + continue + bgra[i] = (bgra[i] * a) // 255 + bgra[i + 1] = (bgra[i + 1] * a) // 255 + bgra[i + 2] = (bgra[i + 2] * a) // 255 + else: + for i in range(0, len(bgra), 4): + a = (bgra[i + 3] * int(intensity * 255)) // 255 + bgra[i + 3] = a + if a == 0: + bgra[i] = 0 + bgra[i + 1] = 0 + bgra[i + 2] = 0 + continue + bgra[i] = (bgra[i] * a) // 255 + bgra[i + 1] = (bgra[i + 1] * a) // 255 + bgra[i + 2] = (bgra[i + 2] * a) // 255 + return bytes(bgra) diff --git a/src/windows_mcp/desktop/service.py b/src/windows_mcp/desktop/service.py index aab1ef39..ae3bd2ba 100644 --- a/src/windows_mcp/desktop/service.py +++ b/src/windows_mcp/desktop/service.py @@ -1353,12 +1353,20 @@ def get_display_union_rect( def get_screenshot(self, capture_rect: uia.Rect | None = None) -> Image.Image: # The persistent AI indicator is hidden and acknowledged before any - # screenshot backend reads pixels; it is restored even if capture fails. - with control_overlay.suspend_for_capture(): - if not flash_overlay.cancel_active_flash(): - raise RuntimeError("Previous screenshot flash did not close before capture") - image, used_backend = screenshot_capture.capture(capture_rect) - self._last_screenshot_backend = used_backend + # screenshot backend reads pixels. Physical input is released first so + # a stalled capture cannot leave the user blocked by an invisible AI. + controller = get_controller() + capture_generation = controller.pause_for_capture() + restored = False + try: + with control_overlay.suspend_for_capture(): + if not flash_overlay.cancel_active_flash(): + raise RuntimeError("Previous screenshot flash did not close before capture") + image, used_backend = screenshot_capture.capture(capture_rect) + self._last_screenshot_backend = used_backend + restored = True + finally: + controller.resume_after_capture(capture_generation, restored=restored) flash_overlay.show_capture_flash(capture_rect) return image diff --git a/tests/test_control.py b/tests/test_control.py index 6ad5aa9c..193e0e1e 100644 --- a/tests/test_control.py +++ b/tests/test_control.py @@ -40,6 +40,114 @@ def run(): owner.stop() +def test_indicator_ack_precedes_suppression_and_physical_input_wins(monkeypatch): + owner = ready_controller() + owner.set_health_probe(lambda: True) + observed = [] + owner.subscribe(lambda status: observed.append((status["state"], owner._suppress))) + token = owner.begin_call("Type") + assert observed == [("ai", False)] + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + event = control._KeyHookData() + event.vkCode = 0x41 + assert owner._physical_key(0, 0x100, ctypes.addressof(event)) == 7 + owner.arm_visible(token) # A late visual acknowledgement cannot reclaim the key. + assert not owner._suppress + owner._handle(owner._events.get_nowait()) + assert owner.status()["state"] == "user" + + +def test_indicator_ack_arms_and_capture_pause_releases_physical_input(): + owner = ready_controller() + owner.set_health_probe(lambda: True) + owner.subscribe( + lambda status: owner.arm_visible(status["generation"]) if status["state"] == "ai" else None + ) + token = owner.begin_call("Screenshot") + owner.checkpoint(token) + assert owner._suppress + + generation = owner.pause_for_capture() + assert generation == token + assert not owner._suppress + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + owner.resume_after_capture(generation, restored=True) + owner.checkpoint(token) + assert owner._suppress + + +def test_capture_restore_failure_never_rearms(): + owner = ready_controller() + owner.set_health_probe(lambda: True) + owner.subscribe( + lambda status: owner.arm_visible(status["generation"]) if status["state"] == "ai" else None + ) + owner.begin_call("Screenshot") + generation = owner.pause_for_capture() + owner.resume_after_capture(generation, restored=False) + assert owner._emergency and not owner._suppress + + +def test_stale_capture_generation_never_rearms_after_takeover(): + owner = ready_controller() + owner.set_health_probe(lambda: True) + owner.subscribe( + lambda status: owner.arm_visible(status["generation"]) if status["state"] == "ai" else None + ) + owner.begin_call("Screenshot") + generation = owner.pause_for_capture() + owner._handle(("key",)) + assert owner.status()["state"] == "user" + owner.resume_after_capture(generation, restored=True) + assert not owner._suppress + + +def test_watchdog_does_not_rearm_hidden_indicator(monkeypatch): + owner = ready_controller() + owner.set_health_probe(lambda: True) + owner.subscribe( + lambda status: owner.arm_visible(status["generation"]) if status["state"] == "ai" else None + ) + owner.begin_call("Screenshot") + owner.pause_for_capture() + owner._thread = Mock() + owner._thread.is_alive.return_value = True + owner._thread_beat = control.time.monotonic() + owner._stop.is_set = Mock(side_effect=[False, True]) + owner._release_requested.wait = Mock() + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + owner._watch() + assert not owner._suppress + + +@pytest.mark.parametrize("left,right", [(0xA0, 0xA1), (0xA2, 0xA3), (0xA4, 0xA5)]) +def test_opposite_modifier_release_does_not_clear_held_key(monkeypatch, left, right): + clock = [100.0] + monkeypatch.setattr(control.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + owner = ready_controller() + owner._handle(("key",)) + + def send(vk, message): + data = control._KeyHookData() + data.vkCode = vk + return owner._physical_key(0, message, ctypes.addressof(data)) + + send(left, 0x100) + send(right, 0x100) + send(left, 0x101) # Releasing one side must not clear the other. + assert owner._pressed == {right} + clock[0] = 120.0 + assert owner.status()["state"] == "user" + send(right, 0x101) + clock[0] = 130.0 + assert owner.status()["state"] == "ready" + + def test_start_preserves_cooldown_after_real_physical_input(monkeypatch): owner = control.ControlCoordinator() monkeypatch.setattr(control, "_interactive_desktop", lambda: True) diff --git a/tests/test_control_mcp_gate.py b/tests/test_control_mcp_gate.py index de1241ba..94a6bab0 100644 --- a/tests/test_control_mcp_gate.py +++ b/tests/test_control_mcp_gate.py @@ -148,6 +148,36 @@ def next_tool(): assert effects == ["hold"] +@pytest.mark.asyncio +async def test_cancelled_active_and_queued_calls_release_gate_and_lease(): + controller = FakeController() + gate = ControlToolGate(controller, ControlNotifier(controller)) + context = SimpleNamespace(message=SimpleNamespace(name="Action"), fastmcp_context=None) + started = asyncio.Event() + never = asyncio.Event() + + async def hold(_context): + started.set() + await never.wait() + + active = asyncio.create_task(gate.on_call_tool(context, hold)) + await asyncio.wait_for(started.wait(), 1) + queued = asyncio.create_task(gate.on_call_tool(context, hold)) + await asyncio.sleep(0) + queued.cancel() + with pytest.raises(asyncio.CancelledError): + await queued + active.cancel() + with pytest.raises(asyncio.CancelledError): + await active + assert not gate._call_lock.locked() + assert controller.calls == [("begin", "Action"), ("end", 7)] + + controller.state = "ready" + assert await gate.on_call_tool(context, lambda _: asyncio.sleep(0, result="ok")) == "ok" + assert controller.calls[-2:] == [("begin", "Action"), ("end", 7)] + + @pytest.mark.asyncio @pytest.mark.parametrize("tool_name", ["App", "PowerShell"]) async def test_already_running_external_command_is_not_claimed_cancelled(tool_name): diff --git a/tests/test_control_overlay.py b/tests/test_control_overlay.py index 68d5bd8a..e16295aa 100644 --- a/tests/test_control_overlay.py +++ b/tests/test_control_overlay.py @@ -8,7 +8,7 @@ from PIL import Image import pytest -from windows_mcp.desktop import control_overlay, flash_overlay +from windows_mcp.desktop import control, control_overlay, flash_overlay from windows_mcp.desktop.service import Desktop @@ -358,6 +358,52 @@ def failing_capture(rect): assert events == ["hide", "capture", "restore"] +def test_stalled_capture_releases_physical_input_and_keeps_status_available(monkeypatch): + owner = control.ControlCoordinator() + with owner._lock: + owner._set_locked("ready") + owner.set_health_probe(lambda: True) + owner.subscribe( + lambda status: owner.arm_visible(status["generation"]) if status["state"] == "ai" else None + ) + token = owner.begin_call("Screenshot") + started, release = threading.Event(), threading.Event() + monkeypatch.setattr("windows_mcp.desktop.service.get_controller", lambda: owner) + monkeypatch.setattr(flash_overlay, "cancel_active_flash", lambda: True) + monkeypatch.setattr(flash_overlay, "show_capture_flash", lambda rect: None) + + @contextmanager + def hidden(): + assert not owner._suppress + yield + + def capture(rect): + started.set() + assert release.wait(2.0) + return Image.new("RGB", (2, 2)), "test" + + monkeypatch.setattr(control_overlay, "suspend_for_capture", hidden) + monkeypatch.setattr("windows_mcp.desktop.service.screenshot_capture.capture", capture) + result = [] + worker = threading.Thread( + target=lambda: result.append(Desktop.__new__(Desktop).get_screenshot()) + ) + worker.start() + try: + assert started.wait(1.0) + assert not owner._suppress + assert owner.status()["state"] == "ai" + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + owner._handle(("key",)) # Input delivered during the hidden interval wins. + assert owner.status()["state"] == "user" + finally: + release.set() + worker.join(timeout=2.0) + assert not worker.is_alive() + assert result and not owner._suppress + + def test_unclosed_flash_prevents_new_capture(monkeypatch): events = [] From c08020834c6f1e0a5518e6c72982a725e9c05428 Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Sat, 3 Oct 2026 02:07:19 +0800 Subject: [PATCH 3/9] Avoid concurrent iteration of physical modifier state --- README.md | 2 +- src/windows_mcp/desktop/control.py | 12 ++++++++++-- src/windows_mcp/desktop/control_hooks.py | 10 +++++++--- src/windows_mcp/desktop/overlay_bitmap.py | 14 ++++++++++++-- tests/test_control.py | 13 +++++++++++++ 5 files changed, 43 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 01423cba..6cfc2251 100755 --- a/README.md +++ b/README.md @@ -719,7 +719,7 @@ MCP Client can access the following tools to interact with Windows: - `Shortcut`: Press keyboard shortcuts (`Ctrl+c`, `Alt+Tab`, etc). - `Wait`: Pause for a defined duration. - `WaitFor`: Wait until text, an active window, an element, or a focused element appears by polling UI state inside one tool call. -- `ControlStatus`: Read the desktop control owner and remaining user-idle wait. It stays available when user takeover blocks new tool calls. During AI control, a visible edge and cursor glow indicate ownership; moving the physical mouse substantially or pressing `Ctrl+Alt+Shift+Backspace` returns control to the user until 10 seconds of inactivity. Stepwise desktop input stops at its next checkpoint, but an already-running external command (for example `PowerShell` or an `App` launch) may continue after takeover; this control gate does not cancel processes or undo completed side effects. On a hook fault, injected keys and mouse buttons are released asynchronously, so physical input can briefly overlap a still-held AI key or button. +- `ControlStatus`: Read the desktop control owner and remaining user-idle wait. It stays available when user takeover blocks new tool calls. During AI control, a visible edge and cursor glow indicate ownership; moving the physical mouse substantially or pressing `Ctrl+Alt+Shift+Backspace` returns control to the user until 10 seconds of inactivity. During screenshot capture, the indicator is temporarily hidden and physical input is released; any physical input during this interval preempts AI control, although `ControlStatus` can briefly still report `ai` until that event is processed. Stepwise desktop input stops at its next checkpoint, but an already-running external command (for example `PowerShell` or an `App` launch) may continue after takeover; this control gate does not cancel processes or undo completed side effects. On a hook fault, injected keys and mouse buttons are released asynchronously, so physical input can briefly overlap a still-held AI key or button. - `DisplayInventory`: Read display layout, work areas, effective DPI, and scale metadata. - `Screenshot`: Fast screenshot-first desktop capture with cursor position, active/open windows, and an image. Skips UI tree extraction for speed and should be the default first call when you mainly need visual context. Supports `display=[0]` or `display=[0,1]` using zero-based active Windows display indices, and `region=[left, top, right, bottom]` (virtual-desktop pixel coordinates) to capture just that rectangle instead of the whole screen — cheaper on tokens when you already know which area matters. `region` takes precedence over `display` when both are given; an invalid or out-of-bounds region raises an error. After capture, a brief orange-red glowing border is drawn inside the captured area as a visual confirmation (set `WINDOWS_MCP_DISABLE_FLASH=1` to disable). - `Snapshot`: Full desktop state capture for workflows that need interactive element ids, scrollable regions, or `use_dom=True` browser extraction. Supports `use_vision=True` for including screenshots, `display=[0]` or `display=[0,1]` using zero-based active Windows display indices, and `region=[left, top, right, bottom]` (virtual-desktop pixel coordinates) to inspect just that rectangle instead of the whole screen; `region` takes precedence over `display` when both are given, and an invalid or out-of-bounds region raises an error. diff --git a/src/windows_mcp/desktop/control.py b/src/windows_mcp/desktop/control.py index 19da0666..11582a6c 100644 --- a/src/windows_mcp/desktop/control.py +++ b/src/windows_mcp/desktop/control.py @@ -287,8 +287,16 @@ def record_step_current(self) -> None: def physical_key_down(self, vk: int) -> bool: # A swallowed physical down did not reach the app, so the AI must - # release its own injected down when its tool is preempted. - return _key(vk) in {_key(held) for held in self._delivered_keys} + # release its own injected down when its tool is preempted. Hook + # callbacks mutate this set concurrently: use atomic membership checks + # rather than iterating it while resolving left/right modifiers. + normalized = _key(vk) + aliases = { + 0x10: (0x10, 0xA0, 0xA1), + 0x11: (0x11, 0xA2, 0xA3), + 0x12: (0x12, 0xA4, 0xA5), + } + return any(key in self._delivered_keys for key in aliases.get(normalized, (vk,))) def physical_mouse_down(self, button: str) -> bool: return {"left": 1, "right": 2, "middle": 3}[button] in self._delivered_mouse diff --git a/src/windows_mcp/desktop/control_hooks.py b/src/windows_mcp/desktop/control_hooks.py index 2f2b7d64..764590e9 100644 --- a/src/windows_mcp/desktop/control_hooks.py +++ b/src/windows_mcp/desktop/control_hooks.py @@ -2,6 +2,7 @@ import ctypes import time +from typing import Any from windows_mcp.desktop.control_win32 import ( _CHORD, @@ -14,7 +15,8 @@ ) -def physical_mouse(owner, code, wparam, lparam): +def physical_mouse(owner: Any, code: int, wparam: int, lparam: int) -> int: + """Handle a physical mouse hook callback without waiting on other threads.""" if code < 0: return _user32.CallNextHookEx(owner._mouse_hook, code, wparam, lparam) valid = False @@ -49,7 +51,8 @@ def physical_mouse(owner, code, wparam, lparam): return _user32.CallNextHookEx(owner._mouse_hook, code, wparam, lparam) -def physical_key(owner, code, wparam, lparam): +def physical_key(owner: Any, code: int, wparam: int, lparam: int) -> int: + """Track each physical key independently and honor the takeover chord.""" if code < 0: return _user32.CallNextHookEx(owner._key_hook, code, wparam, lparam) vk = None @@ -99,7 +102,8 @@ def physical_key(owner, code, wparam, lparam): return _user32.CallNextHookEx(owner._key_hook, code, wparam, lparam) -def raw_input(owner, lparam) -> None: +def raw_input(owner: Any, lparam: int) -> None: + """Queue relative mouse data from a device, independent of injected input.""" raw = _read_raw_mouse(lparam) if raw is None: return # Null devices can be touchpads; do not infer source. diff --git a/src/windows_mcp/desktop/overlay_bitmap.py b/src/windows_mcp/desktop/overlay_bitmap.py index 5a154440..aef1075d 100644 --- a/src/windows_mcp/desktop/overlay_bitmap.py +++ b/src/windows_mcp/desktop/overlay_bitmap.py @@ -1,8 +1,18 @@ """Shared pixel conversion for the control indicator and capture flash.""" +from PIL import Image -def premultiplied_bgra(rgba_image, intensity: float) -> bytes: - """Convert a PIL RGBA image to scaled, premultiplied BGRA bytes.""" + +def premultiplied_bgra(rgba_image: Image.Image, intensity: float) -> bytes: + """Convert a PIL RGBA image to scaled, premultiplied BGRA bytes. + + Args: + rgba_image: Source RGBA bitmap. + intensity: Opacity multiplier for the layered window. + + Returns: + Premultiplied BGRA pixels in row-major order. + """ bgra = bytearray(rgba_image.tobytes("raw", "BGRA")) if intensity >= 1.0: for i in range(0, len(bgra), 4): diff --git a/tests/test_control.py b/tests/test_control.py index 193e0e1e..488439a4 100644 --- a/tests/test_control.py +++ b/tests/test_control.py @@ -148,6 +148,19 @@ def send(vk, message): assert owner.status()["state"] == "ready" +@pytest.mark.parametrize("generic,side", [(0x10, 0xA0), (0x11, 0xA3), (0x12, 0xA4)]) +def test_delivered_modifier_lookup_never_iterates_live_hook_set(generic, side): + class NoIteration(set): + def __iter__(self): + raise AssertionError("hook-owned set must not be iterated") + + owner = ready_controller() + owner._delivered_keys = NoIteration({side}) + assert owner.physical_key_down(generic) + assert owner.physical_key_down(side) + assert not owner.physical_key_down(0x41) + + def test_start_preserves_cooldown_after_real_physical_input(monkeypatch): owner = control.ControlCoordinator() monkeypatch.setattr(control, "_interactive_desktop", lambda: True) From e1bb9791808ed10022cd33d2f27fc11c0cc031d7 Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Sat, 3 Oct 2026 03:38:13 +0800 Subject: [PATCH 4/9] Guard takeover indicator generations and document hook limits --- README.md | 2 +- src/windows_mcp/__main__.py | 14 ++- src/windows_mcp/desktop/control.py | 11 +- src/windows_mcp/desktop/control_overlay.py | 35 ++++-- tests/test_control_hook_rotation.py | 125 +++++++++++++++++++++ tests/test_control_mcp_gate.py | 26 ++++- tests/test_control_overlay.py | 39 +++++++ 7 files changed, 230 insertions(+), 22 deletions(-) create mode 100644 tests/test_control_hook_rotation.py diff --git a/README.md b/README.md index 6cfc2251..d7360a79 100755 --- a/README.md +++ b/README.md @@ -719,7 +719,7 @@ MCP Client can access the following tools to interact with Windows: - `Shortcut`: Press keyboard shortcuts (`Ctrl+c`, `Alt+Tab`, etc). - `Wait`: Pause for a defined duration. - `WaitFor`: Wait until text, an active window, an element, or a focused element appears by polling UI state inside one tool call. -- `ControlStatus`: Read the desktop control owner and remaining user-idle wait. It stays available when user takeover blocks new tool calls. During AI control, a visible edge and cursor glow indicate ownership; moving the physical mouse substantially or pressing `Ctrl+Alt+Shift+Backspace` returns control to the user until 10 seconds of inactivity. During screenshot capture, the indicator is temporarily hidden and physical input is released; any physical input during this interval preempts AI control, although `ControlStatus` can briefly still report `ai` until that event is processed. Stepwise desktop input stops at its next checkpoint, but an already-running external command (for example `PowerShell` or an `App` launch) may continue after takeover; this control gate does not cancel processes or undo completed side effects. On a hook fault, injected keys and mouse buttons are released asynchronously, so physical input can briefly overlap a still-held AI key or button. +- `ControlStatus`: Read the desktop control owner and remaining user-idle wait. It stays available when user takeover blocks new tool calls. During AI control, a visible edge and cursor glow indicate ownership; moving the physical mouse substantially or pressing `Ctrl+Alt+Shift+Backspace` returns control to the user until 10 seconds of inactivity. During screenshot capture, the indicator is temporarily hidden and physical input is released; any physical input during this interval preempts AI control, although `ControlStatus` can briefly still report `ai` until that event is processed. Stepwise desktop input stops at its next checkpoint, but an already-running external command (for example `PowerShell` or an `App` launch) may continue after takeover; this control gate does not cancel processes or undo completed side effects. On a detected hook fault, injected keys and mouse buttons are released asynchronously, so physical input can briefly overlap a still-held AI key or button. Windows may also silently remove a low-level hook without notifying this process; periodic hook replacement reduces but does not eliminate that undetected window. User takeover is therefore best-effort, not a security boundary or a guarantee of exclusive control. - `DisplayInventory`: Read display layout, work areas, effective DPI, and scale metadata. - `Screenshot`: Fast screenshot-first desktop capture with cursor position, active/open windows, and an image. Skips UI tree extraction for speed and should be the default first call when you mainly need visual context. Supports `display=[0]` or `display=[0,1]` using zero-based active Windows display indices, and `region=[left, top, right, bottom]` (virtual-desktop pixel coordinates) to capture just that rectangle instead of the whole screen — cheaper on tokens when you already know which area matters. `region` takes precedence over `display` when both are given; an invalid or out-of-bounds region raises an error. After capture, a brief orange-red glowing border is drawn inside the captured area as a visual confirmation (set `WINDOWS_MCP_DISABLE_FLASH=1` to disable). - `Snapshot`: Full desktop state capture for workflows that need interactive element ids, scrollable regions, or `use_dom=True` browser extraction. Supports `use_vision=True` for including screenshots, `display=[0]` or `display=[0,1]` using zero-based active Windows display indices, and `region=[left, top, right, bottom]` (virtual-desktop pixel coordinates) to inspect just that rectangle instead of the whole screen; `region` takes precedence over `display` when both are given, and an invalid or out-of-bounds region raises an error. diff --git a/src/windows_mcp/__main__.py b/src/windows_mcp/__main__.py index 7a471d32..661ced1c 100755 --- a/src/windows_mcp/__main__.py +++ b/src/windows_mcp/__main__.py @@ -275,21 +275,25 @@ def apply_control_state(status: dict) -> None: "generation" ] != controller.status().get("generation"): return + generation = status.get("generation") if status["state"] == "ai": - control_overlay.set_active(True) - if status.get("generation") is not None: + applied = control_overlay.set_active(True, generation=generation) + if applied and generation is not None: controller.arm_visible(status["generation"]) elif status["state"] == "takeover_pending": - control_overlay.set_pending(True) + control_overlay.set_pending(True, generation=generation) else: - control_overlay.set_active(False) + control_overlay.set_active(False, generation=generation) except Exception: # A missing visual indicator invalidates the AI control lease. controller._fail_open() logger.exception("AI control indicator failed; releasing physical input") def show_control_state(status: dict) -> None: - if threading.current_thread() is controller._thread: + if threading.current_thread() is controller._thread or ( + status["state"] in ("ai", "takeover_pending") + and not getattr(controller, "visible_ack_required", lambda: True)() + ): # Never wait for a window-owner acknowledgement on the hook thread. # A capture may hold the overlay lock for an unbounded backend call; # offload the wait so ControlStatus stays responsive on the loop. diff --git a/src/windows_mcp/desktop/control.py b/src/windows_mcp/desktop/control.py index 11582a6c..b1ccb7f5 100644 --- a/src/windows_mcp/desktop/control.py +++ b/src/windows_mcp/desktop/control.py @@ -50,6 +50,7 @@ def __init__(self, mouse_takeover_units: int = 120, mouse_takeover_pixels: int = self._listeners: list[Callable[[dict], None]] = [] self.input_ledger = InputLedger() self._health_probe: Callable[[], bool] | None = None + self._begin_notification = threading.local() self._visual_armed = False self._visual_paused = False self._visual_event_stamp = 0.0 @@ -87,6 +88,10 @@ def set_health_probe(self, probe: Callable[[], bool]) -> None: """Require the visual indicator to stay alive during AI control.""" self._health_probe = probe + def visible_ack_required(self) -> bool: + """Only a new MCP lease must wait for its visual acknowledgement.""" + return bool(getattr(self._begin_notification, "active", False)) + def arm_visible(self, generation: int) -> None: """Suppress physical input only after the AI indicator is visible.""" with self._lock: @@ -256,7 +261,11 @@ def begin_call(self, name: str) -> int: token = self._generation result = self._snapshot_locked(now) if changed: - self._notify(result) + self._begin_notification.active = True + try: + self._notify(result) + finally: + self._begin_notification.active = False return token def checkpoint(self, token: int) -> None: diff --git a/src/windows_mcp/desktop/control_overlay.py b/src/windows_mcp/desktop/control_overlay.py index e35591a5..6ed50f52 100644 --- a/src/windows_mcp/desktop/control_overlay.py +++ b/src/windows_mcp/desktop/control_overlay.py @@ -185,6 +185,7 @@ def __init__(self) -> None: self.stopping = False self.version = 0 self.applied = 0 + self.generation = -1 self.error: BaseException | None = None self.heartbeat = 0.0 self.started = threading.Event() @@ -197,15 +198,26 @@ def start(self) -> None: raise RuntimeError("AI control indicator thread failed at startup") from self.error def change( - self, *, active: bool | None = None, pending: bool | None = None, suspended: int = 0 - ) -> None: + self, + *, + active: bool | None = None, + pending: bool | None = None, + suspended: int = 0, + generation: int | None = None, + ) -> bool: with self.condition: + # A newer takeover can hide immediately while an older AI show + # waits behind screenshot capture. Never apply that stale show. + if generation is not None and generation < self.generation: + return False if self.stopping: raise RuntimeError("AI control indicator is stopping") if self.error: raise RuntimeError("AI control indicator failed") from self.error if not self.thread.is_alive(): raise RuntimeError("AI control indicator stopped") + if generation is not None: + self.generation = generation if active is not None: self.active = active if pending is not None: @@ -236,6 +248,7 @@ def change( raise RuntimeError("AI control indicator is stopping") if self.error: raise RuntimeError("AI control indicator failed") from self.error + return True def stop(self) -> None: with self.condition: @@ -374,34 +387,34 @@ def stop() -> None: _instance = None -def set_active(active: bool) -> None: +def set_active(active: bool, *, generation: int | None = None) -> bool: """Display or hide AI indication; activation failure raises.""" # Deactivation must remain immediate during a capture: the screenshot's # final resume will then observe active=False and leave the window hidden. - def apply() -> None: + def apply() -> bool: with _lock: indicator = _instance if indicator is None: if not active: - return + return True raise RuntimeError("AI control indicator has not started") - indicator.change(active=active, pending=False) + return indicator.change(active=active, pending=False, generation=generation) if active: with _capture_lock: - apply() - else: - apply() + return apply() + return apply() -def set_pending(pending: bool) -> None: +def set_pending(pending: bool, *, generation: int | None = None) -> bool: """Use amber edge indication without cursor following during takeover check.""" with _capture_lock: with _lock: indicator = _instance if indicator: - indicator.change(pending=pending) + return indicator.change(pending=pending, generation=generation) + return True def is_healthy() -> bool: diff --git a/tests/test_control_hook_rotation.py b/tests/test_control_hook_rotation.py new file mode 100644 index 00000000..d660601e --- /dev/null +++ b/tests/test_control_hook_rotation.py @@ -0,0 +1,125 @@ +"""Hook replacement and fail-open behavior at the Win32 monitor boundary.""" + +from collections import deque +from types import SimpleNamespace + +import pytest + +from windows_mcp.desktop import control, control_win32 +from windows_mcp.desktop.control_ledger import InputUnavailable + + +def _fake_user32(handles, events): + hooks = deque(handles) + + def install(kind, callback, instance, thread_id): + handle = hooks.popleft() + events.append(("install", kind, handle)) + return handle + + def unhook(handle): + events.append(("unhook", handle)) + return 1 + + def register_raw(devices, count, size): + events.append(("raw", count)) + return 1 + + return SimpleNamespace( + SetWindowsHookExW=install, + UnhookWindowsHookEx=unhook, + RegisterClassExW=lambda wc: 1, + CreateWindowExW=lambda *args: 777, + RegisterRawInputDevices=register_raw, + PeekMessageW=lambda *args: 0, + TranslateMessage=lambda *args: 1, + DispatchMessageW=lambda *args: 1, + DefWindowProcW=lambda *args: 0, + DestroyWindow=lambda hwnd: events.append(("destroy", hwnd)), + UnregisterClassW=lambda *args: events.append(("unregister",)), + ) + + +def _run_monitor(monkeypatch, handles, *, active=False): + events = [] + user32 = _fake_user32(handles, events) + monkeypatch.setattr(control, "_user32", user32) + monkeypatch.setattr(control_win32, "_user32", user32) + clock = [100.0] + monkeypatch.setattr(control_win32.time, "monotonic", lambda: clock[0]) + owner = control.ControlCoordinator() + releases = [] + if active: + with owner._lock: + owner._set_locked("ready") + owner.begin_call("Smoke") + owner.input_ledger.press( + "key:shift", + lambda: releases.append("down"), + lambda: releases.append("up"), + lambda: False, + ) + sleeps = [] + + def advance(seconds): + sleeps.append(seconds) + clock[0] += 5.01 + if len(sleeps) == 2: + owner._stop.set() + + monkeypatch.setattr(control_win32.time, "sleep", advance) + control_win32.run_input_monitor(owner) + return owner, events, sleeps, releases + + +def test_hook_rotation_replaces_only_after_both_installations(monkeypatch): + events = [] + monkeypatch.setattr(control, "_user32", _fake_user32([11, 12, 21, 22], events)) + owner = control.ControlCoordinator() + owner._mouse_callback = owner._key_callback = object() + owner._install_hooks(None) + owner._install_hooks(None) + assert events == [ + ("install", 14, 11), + ("install", 13, 12), + ("install", 14, 21), + ("install", 13, 22), + ("unhook", 11), + ("unhook", 12), + ] + assert (owner._mouse_hook, owner._key_hook) == (21, 22) + + +def test_input_monitor_rotates_after_interval_and_cleans_hooks(monkeypatch): + owner, events, sleeps, _ = _run_monitor(monkeypatch, [11, 12, 21, 22]) + assert owner._startup_error is None + assert len(sleeps) == 2 + assert events.index(("install", 14, 21)) < events.index(("unhook", 11)) + assert events.index(("install", 13, 22)) < events.index(("unhook", 12)) + assert events[-5:] == [ + ("unhook", 21), + ("unhook", 22), + ("raw", 1), # RIDEV_REMOVE + ("destroy", 777), + ("unregister",), + ] + + +def test_hook_install_partial_failure_preserves_old_and_fails_open(monkeypatch): + owner, events, _, releases = _run_monitor(monkeypatch, [11, 12, 21, 0], active=True) + assert owner._startup_error is not None + assert events.index(("unhook", 21)) < events.index(("unhook", 11)) + assert events.index(("unhook", 21)) < events.index(("unhook", 12)) + assert not owner._suppress + assert owner.status()["state"] == "unavailable" + assert releases == ["down", "up"] + + +def test_input_monitor_rotation_failure_releases_physical_input(monkeypatch): + owner, events, _, releases = _run_monitor(monkeypatch, [11, 12, 0, 0], active=True) + assert owner._emergency and not owner._suppress + assert ("unhook", 11) in events and ("unhook", 12) in events + owner.status() # The watchdog/status path clears outstanding AI holds. + assert releases == ["down", "up"] + with pytest.raises(InputUnavailable): + owner.input_ledger.press("key:shift", lambda: None, lambda: None, lambda: False) diff --git a/tests/test_control_mcp_gate.py b/tests/test_control_mcp_gate.py index 94a6bab0..241a9f67 100644 --- a/tests/test_control_mcp_gate.py +++ b/tests/test_control_mcp_gate.py @@ -390,6 +390,7 @@ async def test_overlay_callback_never_waits_on_hook_thread(monkeypatch): class Controller(FakeController): _thread = None fail_count = 0 + sync_visible = True def start(self): self.state = "user" @@ -403,6 +404,9 @@ def _fail_open(self): def set_health_probe(self, probe): self.health_probe = probe + def visible_ack_required(self): + return self.sync_visible + class Desktop: tree = type("Tree", (), {"on_focus_change": lambda *args: None})() @@ -422,9 +426,10 @@ def stop(self): controller = Controller() calls = [] - def slow_show(active): + def slow_show(active, *, generation=None): time.sleep(0.15) calls.append(active) + return True monkeypatch.setenv("ANONYMIZED_TELEMETRY", "false") monkeypatch.setattr(wm, "_mcp", None) @@ -435,7 +440,11 @@ def slow_show(active): monkeypatch.setattr(control_overlay, "start", lambda: None) monkeypatch.setattr(control_overlay, "stop", lambda: None) monkeypatch.setattr(control_overlay, "set_active", slow_show) - monkeypatch.setattr(control_overlay, "set_pending", lambda active: calls.append(active)) + monkeypatch.setattr( + control_overlay, + "set_pending", + lambda active, *, generation=None: calls.append(active), + ) mcp = wm._build_mcp() callback = controller.listeners[0] @@ -447,6 +456,14 @@ def slow_show(active): assert time.monotonic() - start >= 0.14 assert calls == [True] + # A ControlStatus-triggered visual update must not block the server loop. + controller.sync_visible = False + start = time.monotonic() + callback({"state": "ai"}) + assert time.monotonic() - start < 0.1 + await asyncio.sleep(0.2) + assert calls == [True, True] + elapsed = [] def from_hook_thread(): @@ -460,13 +477,14 @@ def from_hook_thread(): thread.join(timeout=1) assert elapsed and elapsed[0] < 0.1 await asyncio.sleep(0.2) - assert calls == [True, False] + assert calls == [True, True, False] monkeypatch.setattr( control_overlay, "set_active", - lambda active: (_ for _ in ()).throw(RuntimeError("lost")), + lambda active, *, generation=None: (_ for _ in ()).throw(RuntimeError("lost")), ) + controller.sync_visible = True callback({"state": "ai"}) assert controller.fail_count == 1 diff --git a/tests/test_control_overlay.py b/tests/test_control_overlay.py index e16295aa..1833b50d 100644 --- a/tests/test_control_overlay.py +++ b/tests/test_control_overlay.py @@ -171,6 +171,45 @@ def test_capture_ack_timeout_cannot_hide_recovered_indicator(monkeypatch): control_overlay.set_active(True) +def test_newer_user_state_rejects_ai_show_waiting_behind_capture(monkeypatch): + indicator = control_overlay._Indicator() + indicator.thread = Mock() + indicator.thread.is_alive.return_value = True + monkeypatch.setattr(control_overlay, "_instance", indicator) + stop_ack = threading.Event() + + def acknowledge(): + while not stop_ack.is_set(): + with indicator.condition: + indicator.condition.wait(timeout=0.01) + indicator.applied = indicator.version + indicator.condition.notify_all() + + ack_thread = threading.Thread(target=acknowledge) + ack_thread.start() + result = [] + old_ai = threading.Thread( + target=lambda: result.append(control_overlay.set_active(True, generation=1)) + ) + try: + with control_overlay._capture_lock: + old_ai.start() + assert control_overlay.set_active(False, generation=2) + old_ai.join(timeout=1.0) + assert not old_ai.is_alive() + assert result == [False] + assert indicator.generation == 2 + assert not indicator.active + finally: + stop_ack.set() + with indicator.condition: + indicator.condition.notify_all() + ack_thread.join(timeout=1.0) + if old_ai.is_alive(): + old_ai.join(timeout=1.0) + monkeypatch.setattr(control_overlay, "_instance", None) + + def test_actual_owner_thread_stall_expires_heartbeat(monkeypatch): entered = threading.Event() release = threading.Event() From 76a42300e8485bcd50b729f1859e9ef16bc84cd3 Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Sat, 3 Oct 2026 03:47:40 +0800 Subject: [PATCH 5/9] Keep takeover status updates off the event loop --- src/windows_mcp/__main__.py | 8 +- tests/test_control.py | 296 ------------------------------- tests/test_control_input.py | 310 +++++++++++++++++++++++++++++++++ tests/test_control_mcp_gate.py | 12 +- 4 files changed, 325 insertions(+), 301 deletions(-) create mode 100644 tests/test_control_input.py diff --git a/src/windows_mcp/__main__.py b/src/windows_mcp/__main__.py index 661ced1c..0ef5877c 100755 --- a/src/windows_mcp/__main__.py +++ b/src/windows_mcp/__main__.py @@ -290,13 +290,13 @@ def apply_control_state(status: dict) -> None: logger.exception("AI control indicator failed; releasing physical input") def show_control_state(status: dict) -> None: - if threading.current_thread() is controller._thread or ( - status["state"] in ("ai", "takeover_pending") - and not getattr(controller, "visible_ack_required", lambda: True)() + if threading.current_thread() is controller._thread or not ( + status["state"] == "ai" and controller.visible_ack_required() ): # Never wait for a window-owner acknowledgement on the hook thread. # A capture may hold the overlay lock for an unbounded backend call; - # offload the wait so ControlStatus stays responsive on the loop. + # only a new MCP lease may wait for the visible AI acknowledgement. + # Other state transitions must not block the MCP event loop. try: if control_loop is None or control_loop.is_closed(): raise RuntimeError("server event loop unavailable") diff --git a/tests/test_control.py b/tests/test_control.py index 488439a4..06d0687e 100644 --- a/tests/test_control.py +++ b/tests/test_control.py @@ -2,15 +2,11 @@ import ctypes import queue -import threading from unittest.mock import Mock import pytest from windows_mcp.desktop import control -from windows_mcp.desktop import service -from windows_mcp.desktop.control_context import StepCounter -from windows_mcp.desktop.control_ledger import InputLedger, InputUnavailable def ready_controller() -> control.ControlCoordinator: @@ -411,295 +407,3 @@ def test_watchdog_cannot_refresh_if_coordinator_lock_is_held(monkeypatch): finally: owner._lock.release() assert owner._emergency and not owner._suppress - - -def test_type_stops_before_next_character_after_preemption(monkeypatch): - owner = Mock() - checks = [None, None, control.ControlBlocked("CONTROL_PREEMPTED", {"state": "user"})] - owner.checkpoint_current.side_effect = checks - monkeypatch.setattr(service, "get_controller", lambda: owner) - click = Mock() - send = Mock() - monkeypatch.setattr(service.uia, "Click", click) - monkeypatch.setattr(service.uia, "SendKeys", send) - desktop = service.Desktop.__new__(service.Desktop) - with pytest.raises(control.ControlBlocked): - desktop.type((1, 2), "abc") - assert send.call_count == 1 - send.assert_any_call("a", interval=0.04, waitTime=0.05) - - -def test_drag_releases_its_mouse_button_after_preemption(monkeypatch): - owner = Mock() - owner.input_ledger = InputLedger() - owner.checkpoint_current.side_effect = [ - None, - None, - control.ControlBlocked("CONTROL_PREEMPTED", {"state": "user"}), - ] - owner.physical_mouse_down.return_value = False - monkeypatch.setattr(service, "get_controller", lambda: owner) - monkeypatch.setattr(service, "sleep", lambda _: None) - monkeypatch.setattr(service.uia, "GetCursorPos", lambda: (0, 0)) - press = Mock() - release = Mock() - monkeypatch.setattr(service.uia, "PressMouse", press) - monkeypatch.setattr(service.uia, "ReleaseMouse", release) - desktop = service.Desktop.__new__(service.Desktop) - with pytest.raises(control.ControlBlocked): - desktop.drag((100, 0)) - press.assert_called_once() - release.assert_called_once() - - -def test_ai_hold_is_released_if_press_raises_after_injection(): - ledger = InputLedger() - sent = [] - - def down_then_raise(): - sent.append("down") - raise RuntimeError("wait failed after injection") - - with pytest.raises(RuntimeError, match="wait failed"): - ledger.press("key:shift", down_then_raise, lambda: sent.append("up"), lambda: False) - assert sent == ["down", "up"] - assert ledger.pending() == () - - -def test_ai_hold_defers_up_only_when_physical_down_reached_app(): - ledger = InputLedger() - sent = [] - ledger.press("mouse:left", lambda: sent.append("down"), lambda: sent.append("up"), lambda: True) - ledger.release("mouse:left") - assert sent == ["down"] - assert ledger.pending() == () - - -def test_fail_open_blocks_new_press_and_clears_before_notification(): - owner = ready_controller() - owner.begin_call("Drag") - actions = [] - owner.input_ledger.press( - "mouse:left", lambda: actions.append("down"), lambda: actions.append("up"), lambda: False - ) - owner.subscribe(lambda _: actions.append("notified")) - owner._fail_open() - with pytest.raises(InputUnavailable, match="unavailable"): - owner.input_ledger.press( - "key:shift", lambda: actions.append("late down"), lambda: None, lambda: False - ) - owner._mark_unavailable() - assert actions == ["down", "up", "notified"] - - -def test_fail_open_wakes_existing_watchdog_to_release_ai_hold(monkeypatch): - owner = ready_controller() - owner.begin_call("Drag") - released = threading.Event() - owner.input_ledger.press("mouse:left", lambda: None, released.set, lambda: False) - owner._thread = Mock() - owner._thread.is_alive.return_value = True - owner._thread_beat = control.time.monotonic() - monkeypatch.setattr(control, "_interactive_desktop", lambda: True) - watcher = threading.Thread(target=owner._watch) - watcher.start() - try: - owner._fail_open() - assert owner._release_requested.is_set() or released.is_set() - assert released.wait(1.0) - assert owner.input_ledger.pending() == () - finally: - owner._stop.set() - owner._release_requested.set() - watcher.join(2) - assert not watcher.is_alive() - - -def test_fail_open_releases_before_slow_health_probe(monkeypatch): - owner = ready_controller() - owner.begin_call("Type") - released, probe_entered, probe_continue = ( - threading.Event(), - threading.Event(), - threading.Event(), - ) - owner.input_ledger.press("key:ctrl", lambda: None, released.set, lambda: False) - owner._thread = Mock() - owner._thread.is_alive.return_value = True - owner._thread_beat = control.time.monotonic() - monkeypatch.setattr(control, "_interactive_desktop", lambda: True) - - def slow_probe(): - probe_entered.set() - probe_continue.wait(2) - return True - - owner.set_health_probe(slow_probe) - # The hook path signals cleanup; it never waits for a UI health probe. - owner._fail_open() - watcher = threading.Thread(target=owner._watch) - watcher.start() - try: - assert released.wait(1) - assert not probe_entered.is_set() - finally: - probe_continue.set() - owner._stop.set() - owner._release_requested.set() - watcher.join(2) - assert not watcher.is_alive() - - -def test_takeover_after_checkpoint_blocks_new_ai_down(monkeypatch): - owner = ready_controller() - token = owner.begin_call("Drag") - owner.checkpoint(token) - monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) - data = control._MouseHookData() - assert owner._physical_mouse(0, 0x200, ctypes.addressof(data)) == 1 - assert owner._fast_pending - with pytest.raises(InputUnavailable, match="unavailable"): - owner.input_ledger.press( - "mouse:left", lambda: pytest.fail("late AI down"), lambda: None, lambda: False - ) - owner._handle(owner._events.get_nowait()) - assert owner.status()["state"] == "takeover_pending" - with pytest.raises(control.ControlBlocked): - owner.checkpoint(token) - - -@pytest.mark.parametrize("action", ["scroll", "drag", "multi_select"]) -def test_takeover_between_checkpoint_and_press_injects_no_down(monkeypatch, action): - owner = Mock() - owner.input_ledger = InputLedger() - calls = 0 - - def checkpoint(): - nonlocal calls - calls += 1 - if calls == (2 if action == "drag" else 1): - owner.input_ledger.block_new() - - owner.checkpoint_current.side_effect = checkpoint - monkeypatch.setattr(service, "get_controller", lambda: owner) - monkeypatch.setattr(service, "sleep", lambda _: None) - monkeypatch.setattr(service.uia, "GetCursorPos", lambda: (0, 0)) - key_down, mouse_down = Mock(), Mock() - monkeypatch.setattr(service.uia, "PressKey", key_down) - monkeypatch.setattr(service.uia, "PressMouse", mouse_down) - desktop = service.Desktop.__new__(service.Desktop) - with pytest.raises(InputUnavailable): - if action == "scroll": - desktop.scroll(type="horizontal", direction="left") - elif action == "drag": - desktop.drag((100, 0)) - else: - desktop.multi_select(press_ctrl=True, locs=[(100, 100)]) - key_down.assert_not_called() - mouse_down.assert_not_called() - assert owner.input_ledger.pending() == () - - -def test_new_lease_reenables_ledger_after_cancelled_mouse_candidate(): - owner = ready_controller() - token = owner.begin_call("Move") - owner.input_ledger.block_new() # Physical move's immediate guard. - owner._handle(("point", 0, 0, 0x200)) - assert owner.status()["state"] == "takeover_pending" - owner._last_move -= 0.4 - assert owner.status()["state"] == "ai" - with pytest.raises(control.ControlBlocked): - owner.checkpoint(token) - owner.begin_call("Click") - actions = [] - owner.input_ledger.press( - "mouse:left", lambda: actions.append("down"), lambda: actions.append("up"), lambda: False - ) - owner.input_ledger.release("mouse:left") - assert actions == ["down", "up"] - - -def test_release_waits_for_inflight_press_then_clears_hold(): - ledger = InputLedger() - entered, finish = threading.Event(), threading.Event() - actions = [] - - def down(): - entered.set() - assert finish.wait(2) - actions.append("down") - - pressing = threading.Thread( - target=lambda: ledger.press("key:ctrl", down, lambda: actions.append("up"), lambda: False) - ) - pressing.start() - assert entered.wait(2) - ledger.block_new() - releasing = threading.Thread(target=ledger.release_all) - releasing.start() - finish.set() - pressing.join(2) - releasing.join(2) - assert not pressing.is_alive() and not releasing.is_alive() - assert actions == ["down", "up"] - assert ledger.pending() == () - - -def test_release_all_attempts_other_holds_after_one_failure(): - ledger = InputLedger() - events = [] - - def fail_up(): - raise RuntimeError("up failed") - - ledger.press("key:ctrl", lambda: None, fail_up, lambda: False) - ledger.press("mouse:left", lambda: None, lambda: events.append("mouse up"), lambda: False) - with pytest.raises(RuntimeError, match="up failed"): - ledger.release_all() - assert events == ["mouse up"] - assert ledger.pending() == ("key:ctrl",) - - -def test_recovery_requires_ledger_clear_health_and_interactive_desktop(monkeypatch): - owner = ready_controller() - owner.begin_call("Drag") - attempts = 0 - - def up(): - nonlocal attempts - attempts += 1 - if attempts == 1: - raise RuntimeError("temporarily blocked") - - owner.input_ledger.press("mouse:left", lambda: None, up, lambda: False) - owner._fail_open() - owner._recover_after_rehook() - assert owner.status()["state"] == "unavailable" - owner.set_health_probe(lambda: False) - owner._recover_after_rehook() - assert owner.status()["state"] == "unavailable" - owner.set_health_probe(lambda: True) - monkeypatch.setattr(control, "_interactive_desktop", lambda: False) - owner._recover_after_rehook() - assert owner.status()["state"] == "unavailable" - monkeypatch.setattr(control, "_interactive_desktop", lambda: True) - owner._recover_after_rehook() - assert owner.status()["state"] == "user" - assert owner.input_ledger.pending() == () - - -def test_preemption_reports_executed_steps(): - owner = ready_controller() - token = owner.begin_call("Type") - context = control.current_token.set(token) - step_context = control.current_steps.set(StepCounter()) - try: - owner.record_step_current() - owner.record_step_current() - owner._handle(("hotkey",)) - with pytest.raises(control.ControlBlocked) as exc: - owner.checkpoint_current() - assert exc.value.status["executed_steps"] == 2 - finally: - control.current_steps.reset(step_context) - control.current_token.reset(context) diff --git a/tests/test_control_input.py b/tests/test_control_input.py new file mode 100644 index 00000000..32b93ed7 --- /dev/null +++ b/tests/test_control_input.py @@ -0,0 +1,310 @@ +"""AI input checkpoints, ledger cleanup, and recovery tests.""" + +import ctypes +import threading +from unittest.mock import Mock + +import pytest + +from windows_mcp.desktop import control, service +from windows_mcp.desktop.control_context import StepCounter +from windows_mcp.desktop.control_ledger import InputLedger, InputUnavailable + + +def ready_controller() -> control.ControlCoordinator: + owner = control.ControlCoordinator() + with owner._lock: + owner._set_locked("ready") + return owner + + +def test_type_stops_before_next_character_after_preemption(monkeypatch): + owner = Mock() + checks = [None, None, control.ControlBlocked("CONTROL_PREEMPTED", {"state": "user"})] + owner.checkpoint_current.side_effect = checks + monkeypatch.setattr(service, "get_controller", lambda: owner) + click = Mock() + send = Mock() + monkeypatch.setattr(service.uia, "Click", click) + monkeypatch.setattr(service.uia, "SendKeys", send) + desktop = service.Desktop.__new__(service.Desktop) + with pytest.raises(control.ControlBlocked): + desktop.type((1, 2), "abc") + assert send.call_count == 1 + send.assert_any_call("a", interval=0.04, waitTime=0.05) + + +def test_drag_releases_its_mouse_button_after_preemption(monkeypatch): + owner = Mock() + owner.input_ledger = InputLedger() + owner.checkpoint_current.side_effect = [ + None, + None, + control.ControlBlocked("CONTROL_PREEMPTED", {"state": "user"}), + ] + owner.physical_mouse_down.return_value = False + monkeypatch.setattr(service, "get_controller", lambda: owner) + monkeypatch.setattr(service, "sleep", lambda _: None) + monkeypatch.setattr(service.uia, "GetCursorPos", lambda: (0, 0)) + press = Mock() + release = Mock() + monkeypatch.setattr(service.uia, "PressMouse", press) + monkeypatch.setattr(service.uia, "ReleaseMouse", release) + desktop = service.Desktop.__new__(service.Desktop) + with pytest.raises(control.ControlBlocked): + desktop.drag((100, 0)) + press.assert_called_once() + release.assert_called_once() + + +def test_ai_hold_is_released_if_press_raises_after_injection(): + ledger = InputLedger() + sent = [] + + def down_then_raise(): + sent.append("down") + raise RuntimeError("wait failed after injection") + + with pytest.raises(RuntimeError, match="wait failed"): + ledger.press("key:shift", down_then_raise, lambda: sent.append("up"), lambda: False) + assert sent == ["down", "up"] + assert ledger.pending() == () + + +def test_ai_hold_defers_up_only_when_physical_down_reached_app(): + ledger = InputLedger() + sent = [] + ledger.press("mouse:left", lambda: sent.append("down"), lambda: sent.append("up"), lambda: True) + ledger.release("mouse:left") + assert sent == ["down"] + assert ledger.pending() == () + + +def test_fail_open_blocks_new_press_and_clears_before_notification(): + owner = ready_controller() + owner.begin_call("Drag") + actions = [] + owner.input_ledger.press( + "mouse:left", lambda: actions.append("down"), lambda: actions.append("up"), lambda: False + ) + owner.subscribe(lambda _: actions.append("notified")) + owner._fail_open() + with pytest.raises(InputUnavailable, match="unavailable"): + owner.input_ledger.press( + "key:shift", lambda: actions.append("late down"), lambda: None, lambda: False + ) + owner._mark_unavailable() + assert actions == ["down", "up", "notified"] + + +def test_fail_open_wakes_existing_watchdog_to_release_ai_hold(monkeypatch): + owner = ready_controller() + owner.begin_call("Drag") + released = threading.Event() + owner.input_ledger.press("mouse:left", lambda: None, released.set, lambda: False) + owner._thread = Mock() + owner._thread.is_alive.return_value = True + owner._thread_beat = control.time.monotonic() + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + watcher = threading.Thread(target=owner._watch) + watcher.start() + try: + owner._fail_open() + assert owner._release_requested.is_set() or released.is_set() + assert released.wait(1.0) + assert owner.input_ledger.pending() == () + finally: + owner._stop.set() + owner._release_requested.set() + watcher.join(2) + assert not watcher.is_alive() + + +def test_fail_open_releases_before_slow_health_probe(monkeypatch): + owner = ready_controller() + owner.begin_call("Type") + released, probe_entered, probe_continue = ( + threading.Event(), + threading.Event(), + threading.Event(), + ) + owner.input_ledger.press("key:ctrl", lambda: None, released.set, lambda: False) + owner._thread = Mock() + owner._thread.is_alive.return_value = True + owner._thread_beat = control.time.monotonic() + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + + def slow_probe(): + probe_entered.set() + probe_continue.wait(2) + return True + + owner.set_health_probe(slow_probe) + # The hook path signals cleanup; it never waits for a UI health probe. + owner._fail_open() + watcher = threading.Thread(target=owner._watch) + watcher.start() + try: + assert released.wait(1) + assert not probe_entered.is_set() + finally: + probe_continue.set() + owner._stop.set() + owner._release_requested.set() + watcher.join(2) + assert not watcher.is_alive() + + +def test_takeover_after_checkpoint_blocks_new_ai_down(monkeypatch): + owner = ready_controller() + token = owner.begin_call("Drag") + owner.checkpoint(token) + monkeypatch.setattr(control._user32, "CallNextHookEx", lambda *args: 7) + data = control._MouseHookData() + assert owner._physical_mouse(0, 0x200, ctypes.addressof(data)) == 1 + assert owner._fast_pending + with pytest.raises(InputUnavailable, match="unavailable"): + owner.input_ledger.press( + "mouse:left", lambda: pytest.fail("late AI down"), lambda: None, lambda: False + ) + owner._handle(owner._events.get_nowait()) + assert owner.status()["state"] == "takeover_pending" + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + + +@pytest.mark.parametrize("action", ["scroll", "drag", "multi_select"]) +def test_takeover_between_checkpoint_and_press_injects_no_down(monkeypatch, action): + owner = Mock() + owner.input_ledger = InputLedger() + calls = 0 + + def checkpoint(): + nonlocal calls + calls += 1 + if calls == (2 if action == "drag" else 1): + owner.input_ledger.block_new() + + owner.checkpoint_current.side_effect = checkpoint + monkeypatch.setattr(service, "get_controller", lambda: owner) + monkeypatch.setattr(service, "sleep", lambda _: None) + monkeypatch.setattr(service.uia, "GetCursorPos", lambda: (0, 0)) + key_down, mouse_down = Mock(), Mock() + monkeypatch.setattr(service.uia, "PressKey", key_down) + monkeypatch.setattr(service.uia, "PressMouse", mouse_down) + desktop = service.Desktop.__new__(service.Desktop) + with pytest.raises(InputUnavailable): + if action == "scroll": + desktop.scroll(type="horizontal", direction="left") + elif action == "drag": + desktop.drag((100, 0)) + else: + desktop.multi_select(press_ctrl=True, locs=[(100, 100)]) + key_down.assert_not_called() + mouse_down.assert_not_called() + assert owner.input_ledger.pending() == () + + +def test_new_lease_reenables_ledger_after_cancelled_mouse_candidate(): + owner = ready_controller() + token = owner.begin_call("Move") + owner.input_ledger.block_new() # Physical move's immediate guard. + owner._handle(("point", 0, 0, 0x200)) + assert owner.status()["state"] == "takeover_pending" + owner._last_move -= 0.4 + assert owner.status()["state"] == "ai" + with pytest.raises(control.ControlBlocked): + owner.checkpoint(token) + owner.begin_call("Click") + actions = [] + owner.input_ledger.press( + "mouse:left", lambda: actions.append("down"), lambda: actions.append("up"), lambda: False + ) + owner.input_ledger.release("mouse:left") + assert actions == ["down", "up"] + + +def test_release_waits_for_inflight_press_then_clears_hold(): + ledger = InputLedger() + entered, finish = threading.Event(), threading.Event() + actions = [] + + def down(): + entered.set() + assert finish.wait(2) + actions.append("down") + + pressing = threading.Thread( + target=lambda: ledger.press("key:ctrl", down, lambda: actions.append("up"), lambda: False) + ) + pressing.start() + assert entered.wait(2) + ledger.block_new() + releasing = threading.Thread(target=ledger.release_all) + releasing.start() + finish.set() + pressing.join(2) + releasing.join(2) + assert not pressing.is_alive() and not releasing.is_alive() + assert actions == ["down", "up"] + assert ledger.pending() == () + + +def test_release_all_attempts_other_holds_after_one_failure(): + ledger = InputLedger() + events = [] + + def fail_up(): + raise RuntimeError("up failed") + + ledger.press("key:ctrl", lambda: None, fail_up, lambda: False) + ledger.press("mouse:left", lambda: None, lambda: events.append("mouse up"), lambda: False) + with pytest.raises(RuntimeError, match="up failed"): + ledger.release_all() + assert events == ["mouse up"] + assert ledger.pending() == ("key:ctrl",) + + +def test_recovery_requires_ledger_clear_health_and_interactive_desktop(monkeypatch): + owner = ready_controller() + owner.begin_call("Drag") + attempts = 0 + + def up(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("temporarily blocked") + + owner.input_ledger.press("mouse:left", lambda: None, up, lambda: False) + owner._fail_open() + owner._recover_after_rehook() + assert owner.status()["state"] == "unavailable" + owner.set_health_probe(lambda: False) + owner._recover_after_rehook() + assert owner.status()["state"] == "unavailable" + owner.set_health_probe(lambda: True) + monkeypatch.setattr(control, "_interactive_desktop", lambda: False) + owner._recover_after_rehook() + assert owner.status()["state"] == "unavailable" + monkeypatch.setattr(control, "_interactive_desktop", lambda: True) + owner._recover_after_rehook() + assert owner.status()["state"] == "user" + assert owner.input_ledger.pending() == () + + +def test_preemption_reports_executed_steps(): + owner = ready_controller() + token = owner.begin_call("Type") + context = control.current_token.set(token) + step_context = control.current_steps.set(StepCounter()) + try: + owner.record_step_current() + owner.record_step_current() + owner._handle(("hotkey",)) + with pytest.raises(control.ControlBlocked) as exc: + owner.checkpoint_current() + assert exc.value.status["executed_steps"] == 2 + finally: + control.current_steps.reset(step_context) + control.current_token.reset(context) diff --git a/tests/test_control_mcp_gate.py b/tests/test_control_mcp_gate.py index 241a9f67..88d1751c 100644 --- a/tests/test_control_mcp_gate.py +++ b/tests/test_control_mcp_gate.py @@ -464,6 +464,16 @@ def slow_show(active, *, generation=None): await asyncio.sleep(0.2) assert calls == [True, True] + # A normal tool can advance user -> ready on the event loop. Its + # visual acknowledgement must not delay other MCP calls. + controller.sync_visible = True + for state in ("user", "ready"): + start = time.monotonic() + callback({"state": state}) + assert time.monotonic() - start < 0.1 + await asyncio.sleep(0.2) + assert calls == [True, True, False, False] + elapsed = [] def from_hook_thread(): @@ -477,7 +487,7 @@ def from_hook_thread(): thread.join(timeout=1) assert elapsed and elapsed[0] < 0.1 await asyncio.sleep(0.2) - assert calls == [True, True, False] + assert calls == [True, True, False, False, False] monkeypatch.setattr( control_overlay, From b15a549f9005a15246215317a8e7c6ed9588674a Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Sat, 3 Oct 2026 03:59:02 +0800 Subject: [PATCH 6/9] Keep visual acknowledgement off the MCP event loop --- .../tools/control_notifications.py | 26 ++++- tests/test_control_mcp_ack.py | 99 +++++++++++++++++++ 2 files changed, 124 insertions(+), 1 deletion(-) create mode 100644 tests/test_control_mcp_ack.py diff --git a/src/windows_mcp/tools/control_notifications.py b/src/windows_mcp/tools/control_notifications.py index 7655a774..c9a7c2ac 100644 --- a/src/windows_mcp/tools/control_notifications.py +++ b/src/windows_mcp/tools/control_notifications.py @@ -176,6 +176,30 @@ def _blocked(exc: Exception) -> ToolError: ) ) + async def _begin_call(self, name: str) -> int: + """Wait for the visible lease off-loop, retaining ownership on cancel.""" + worker = asyncio.create_task(asyncio.to_thread(self.controller.begin_call, name)) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + # The worker cannot be stopped safely mid-ack. Keep the gate lock + # until it finishes, then balance a lease created after cancellation. + while not worker.done(): + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + continue + except Exception: + break + if not worker.cancelled(): + try: + token = worker.result() + except Exception: + pass # A failed begin_call never created a usable lease. + else: + self.controller.end_call(token) + raise + async def on_call_tool(self, context: MiddlewareContext, call_next: Callable) -> Any: self.notifier.remember(context) name = context.message.name @@ -206,7 +230,7 @@ def reject_user_state() -> None: reject_user_state() try: try: - token = self.controller.begin_call(name) + token = await self._begin_call(name) except Exception as exc: from windows_mcp.desktop.control import ControlBlocked diff --git a/tests/test_control_mcp_ack.py b/tests/test_control_mcp_ack.py new file mode 100644 index 00000000..8326d6fe --- /dev/null +++ b/tests/test_control_mcp_ack.py @@ -0,0 +1,99 @@ +"""A slow indicator acknowledgement must not stall status or leak a lease.""" + +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from windows_mcp.desktop.control import ControlBlocked +from windows_mcp.tools.control_notifications import ControlNotifier, ControlToolGate + + +class FakeController: + def __init__(self): + self.state = "ready" + self.generation = 1 + self.calls = [] + self.listeners = [] + + def status(self): + return {"state": self.state, "generation": self.generation, "wait_seconds": 0.0} + + def subscribe(self, callback): + self.listeners.append(callback) + + def emit(self, state): + self.state = state + self.generation += 1 + for callback in self.listeners: + callback(self.status()) + + def begin_call(self, name): + if self.state not in ("ready", "ai"): + raise ControlBlocked("USER_CONTROL", self.status()) + self.calls.append(("begin", name)) + self.state = "ai" + return 7 + + def end_call(self, token): + self.calls.append(("end", token)) + + def checkpoint(self, token): + if self.state != "ai": + raise ControlBlocked("CONTROL_PREEMPTED", self.status()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_during_ack", [False, True]) +async def test_slow_visible_ack_keeps_status_responsive_and_balances_cancel( + cancel_during_ack, +): + started = threading.Event() + release = threading.Event() + + class SlowController(FakeController): + def begin_call(self, name): + started.set() + if not release.wait(timeout=2): + raise RuntimeError("test indicator acknowledgement timed out") + return super().begin_call(name) + + controller = SlowController() + gate = ControlToolGate(controller, ControlNotifier(controller)) + action = SimpleNamespace(message=SimpleNamespace(name="Action"), fastmcp_context=None) + status = SimpleNamespace(message=SimpleNamespace(name="ControlStatus"), fastmcp_context=None) + effects = [] + + async def run_action(_context): + effects.append("ran") + return "ran" + + active = asyncio.create_task(gate.on_call_tool(action, run_action)) + try: + assert await asyncio.wait_for(asyncio.to_thread(started.wait, 1), 1.5) + observed = await asyncio.wait_for( + gate.on_call_tool(status, lambda _: asyncio.sleep(0, result=controller.status())), + 0.5, + ) + assert observed["state"] == "ready" + if cancel_during_ack: + active.cancel() + await asyncio.sleep(0) + assert not active.done() + release.set() + if cancel_during_ack: + with pytest.raises(asyncio.CancelledError): + await active + assert effects == [] + else: + assert await asyncio.wait_for(active, 1) == "ran" + assert effects == ["ran"] + assert controller.calls == [("begin", "Action"), ("end", 7)] + assert not gate._call_lock.locked() + finally: + release.set() + if not active.done(): + active.cancel() + with pytest.raises(asyncio.CancelledError): + await active From 904d6af1491ec0fb88e92789df415932754f650d Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Sat, 3 Oct 2026 05:15:01 +0800 Subject: [PATCH 7/9] Test takeover notifications on legacy MCP sessions --- README.md | 1 + tests/test_control_mcp_gate.py | 14 ++++++++++++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index d7360a79..a855ca90 100755 --- a/README.md +++ b/README.md @@ -720,6 +720,7 @@ MCP Client can access the following tools to interact with Windows: - `Wait`: Pause for a defined duration. - `WaitFor`: Wait until text, an active window, an element, or a focused element appears by polling UI state inside one tool call. - `ControlStatus`: Read the desktop control owner and remaining user-idle wait. It stays available when user takeover blocks new tool calls. During AI control, a visible edge and cursor glow indicate ownership; moving the physical mouse substantially or pressing `Ctrl+Alt+Shift+Backspace` returns control to the user until 10 seconds of inactivity. During screenshot capture, the indicator is temporarily hidden and physical input is released; any physical input during this interval preempts AI control, although `ControlStatus` can briefly still report `ai` until that event is processed. Stepwise desktop input stops at its next checkpoint, but an already-running external command (for example `PowerShell` or an `App` launch) may continue after takeover; this control gate does not cancel processes or undo completed side effects. On a detected hook fault, injected keys and mouse buttons are released asynchronously, so physical input can briefly overlap a still-held AI key or button. Windows may also silently remove a low-level hook without notifying this process; periodic hook replacement reduces but does not eliminate that undetected window. User takeover is therefore best-effort, not a security boundary or a guarantee of exclusive control. + Sessionless MCP 2026-07-28 clients cannot receive unsolicited control-state logs; they should query `ControlStatus` when they need the current owner. - `DisplayInventory`: Read display layout, work areas, effective DPI, and scale metadata. - `Screenshot`: Fast screenshot-first desktop capture with cursor position, active/open windows, and an image. Skips UI tree extraction for speed and should be the default first call when you mainly need visual context. Supports `display=[0]` or `display=[0,1]` using zero-based active Windows display indices, and `region=[left, top, right, bottom]` (virtual-desktop pixel coordinates) to capture just that rectangle instead of the whole screen — cheaper on tokens when you already know which area matters. `region` takes precedence over `display` when both are given; an invalid or out-of-bounds region raises an error. After capture, a brief orange-red glowing border is drawn inside the captured area as a visual confirmation (set `WINDOWS_MCP_DISABLE_FLASH=1` to disable). - `Snapshot`: Full desktop state capture for workflows that need interactive element ids, scrollable regions, or `use_dom=True` browser extraction. Supports `use_vision=True` for including screenshots, `display=[0]` or `display=[0,1]` using zero-based active Windows display indices, and `region=[left, top, right, bottom]` (virtual-desktop pixel coordinates) to inspect just that rectangle instead of the whole screen; `region` takes precedence over `display` when both are given, and an invalid or out-of-bounds region raises an error. diff --git a/tests/test_control_mcp_gate.py b/tests/test_control_mcp_gate.py index 88d1751c..0708f127 100644 --- a/tests/test_control_mcp_gate.py +++ b/tests/test_control_mcp_gate.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import inspect import socket import threading import time @@ -16,6 +17,13 @@ from windows_mcp.tools.control_notifications import ControlNotifier, ControlToolGate +def _legacy_notification_client(transport, **kwargs): + """Use the sessionful era required for unsolicited logs on FastMCP 4.""" + if "mode" in inspect.signature(Client).parameters: + kwargs["mode"] = "legacy" + return Client(transport, **kwargs) + + class FakeController: def __init__(self): self.state = "ready" @@ -518,7 +526,7 @@ async def on_log(message): notifier.start() try: - async with Client(mcp, log_handler=on_log) as client: + async with _legacy_notification_client(mcp, log_handler=on_log) as client: await client.call_tool("ControlStatus") # Registers the live session. controller.emit("ai") controller.emit("user") @@ -562,7 +570,9 @@ async def on_log(message): await asyncio.sleep(0.01) assert server.started notifier.start() - async with Client(f"http://127.0.0.1:{port}/mcp", log_handler=on_log) as client: + async with _legacy_notification_client( + f"http://127.0.0.1:{port}/mcp", log_handler=on_log + ) as client: await client.call_tool("ControlStatus") controller.emit("ai") controller.emit("user") From 229dfda2a1e5d1af2786581bdaa44745d8615474 Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Sat, 3 Oct 2026 05:40:36 +0800 Subject: [PATCH 8/9] Restore accepted desktop control glow on takeover branch --- src/windows_mcp/desktop/control_overlay.py | 112 ++++--- .../desktop/control_overlay_art.py | 167 +++++++++++ src/windows_mcp/desktop/flash_overlay.py | 7 +- src/windows_mcp/desktop/overlay_bitmap.py | 4 + tests/test_control_overlay.py | 55 +--- tests/test_control_overlay_visuals.py | 282 ++++++++++++++++++ 6 files changed, 532 insertions(+), 95 deletions(-) create mode 100644 src/windows_mcp/desktop/control_overlay_art.py create mode 100644 tests/test_control_overlay_visuals.py diff --git a/src/windows_mcp/desktop/control_overlay.py b/src/windows_mcp/desktop/control_overlay.py index 6ed50f52..fb5de330 100644 --- a/src/windows_mcp/desktop/control_overlay.py +++ b/src/windows_mcp/desktop/control_overlay.py @@ -11,16 +11,21 @@ import time from ctypes import wintypes -from PIL import Image, ImageDraw, ImageFilter - from windows_mcp.desktop import flash_overlay -from windows_mcp.desktop.overlay_bitmap import premultiplied_bgra +from windows_mcp.desktop.control_overlay_art import ( + _BORDER, + _CURSOR_SIZE, + _NOTICE_GLOW_PAD, + _breath_opacity, + _cursor_bitmap, + _edge_bitmap, + _notice_bitmap, + _notice_glow_bitmap, +) from windows_mcp import uia logger = logging.getLogger(__name__) -_BORDER = 28 -_CURSOR_SIZE = 88 _REFRESH_SECONDS = 0.05 _SW_HIDE = 0 _WDA_EXCLUDEFROMCAPTURE = 0x00000011 @@ -36,42 +41,18 @@ _user32.IsWindowVisible.restype = wintypes.BOOL -def _edge_bitmap(width: int, height: int, side: str, color: tuple[int, int, int]) -> bytes: - """Pre-render one inward gradient strip; cost scales with perimeter.""" - image = Image.new("RGBA", (width, height)) - draw = ImageDraw.Draw(image) - for depth in range(min(_BORDER, height if side in ("top", "bottom") else width)): - alpha = int(190 * (1 - depth / _BORDER) ** 2) - if side == "top": - draw.line((0, depth, width - 1, depth), fill=(*color, alpha)) - elif side == "bottom": - y = height - 1 - depth - draw.line((0, y, width - 1, y), fill=(*color, alpha)) - elif side == "left": - draw.line((depth, 0, depth, height - 1), fill=(*color, alpha)) - else: - x = width - 1 - depth - draw.line((x, 0, x, height - 1), fill=(*color, alpha)) - return premultiplied_bgra(image, 1.0) - - -def _cursor_bitmap(color: tuple[int, int, int]) -> bytes: - image = Image.new("RGBA", (_CURSOR_SIZE, _CURSOR_SIZE)) - draw = ImageDraw.Draw(image) - inset = 17 - box = (inset, inset, _CURSOR_SIZE - inset - 1, _CURSOR_SIZE - inset - 1) - draw.ellipse(box, outline=(*color, 210), width=5) - glow = image.filter(ImageFilter.GaussianBlur(9)) - return premultiplied_bgra(Image.alpha_composite(glow, image), 1.0) - - class _Layer: - def __init__(self, x: int, y: int, width: int, height: int, bgra: bytes, name: str): + def __init__( + self, x: int, y: int, width: int, height: int, bgra: bytes, name: str, *, breathes: bool = True + ): self.x, self.y, self.width, self.height = x, y, width, height + self.breathes = breathes self.class_name = f"WindowsMCPControl_{name}_{id(self):x}" self.hwnd, self.instance = flash_overlay._create_layered_window( self.class_name, x, y, width, height ) + self.bgra = bgra + self.opacity: int | None = None try: flash_overlay._push_bitmap(self.hwnd, x, y, width, height, bgra) if not _user32.SetWindowDisplayAffinity(self.hwnd, _WDA_EXCLUDEFROMCAPTURE): @@ -82,6 +63,17 @@ def __init__(self, x: int, y: int, width: int, height: int, bgra: bytes, name: s self.close() raise + def set_opacity(self, opacity: int) -> None: + """Modulate an existing bitmap without rebuilding its per-pixel glow.""" + if self.opacity == opacity: + return + # Use the bitmap API for every frame. SetLayeredWindowAttributes can + # report success here while replacing the visible per-pixel bitmap. + flash_overlay._push_bitmap( + self.hwnd, self.x, self.y, self.width, self.height, self.bgra, opacity=opacity + ) + self.opacity = opacity + def show(self) -> None: _user32.ShowWindow(self.hwnd, flash_overlay._SW_SHOWNA) if not _user32.SetWindowPos( @@ -129,9 +121,7 @@ def _monitor_rects() -> tuple[tuple[int, int, int, int], ...]: return tuple((r.left, r.top, r.right, r.bottom) for r in uia.GetMonitorsRect()) -def _build_layers( - rects: tuple[tuple[int, int, int, int], ...], pending: bool -) -> tuple[list[_Layer], _Layer]: +def _build_layers(rects: tuple[tuple[int, int, int, int], ...], pending: bool) -> tuple[list[_Layer], _Layer]: if not rects: raise RuntimeError("no display available for AI control indicator") color = _AMBER if pending else _BLUE @@ -147,14 +137,40 @@ def _build_layers( strips = ( (left, top, width, border, "top"), (left, bottom - border, width, border, "bottom"), + # The horizontal bitmaps include corner glow, leaving these + # vertical strips separate so corners are not blended twice. (left, top + border, border, height - 2 * border, "left"), (right - border, top + border, border, height - 2 * border, "right"), ) for x, y, w, h, side in strips: if w and h: - layers.append( - _Layer(x, y, w, h, _edge_bitmap(w, h, side, color), f"{index}_{side}") - ) + layers.append(_Layer(x, y, w, h, _edge_bitmap(w, h, side, color), f"{index}_{side}")) + notice = _notice_bitmap(width) + if notice is not None: + notice_width, notice_height, bitmap = notice + # Narrow monitors retain the aura without letting its window spill onto a neighbor. + glow_pad = min(_NOTICE_GLOW_PAD, (width - notice_width) // 2) + if height <= border + notice_height + 2 * glow_pad: + continue + notice_x = left + (width - notice_width) // 2 + notice_y = top + border + glow_pad + layers.append(_Layer( + notice_x - glow_pad, + notice_y - glow_pad, + notice_width + 2 * glow_pad, + notice_height + 2 * glow_pad, + _notice_glow_bitmap(notice_width, notice_height, color, glow_pad), + f"{index}_notice_glow", + )) + layers.append(_Layer( + notice_x, + notice_y, + notice_width, + notice_height, + bitmap, + f"{index}_notice", + breathes=False, + )) point = flash_overlay._POINT() if not _user32.GetCursorPos(ctypes.byref(point)): raise RuntimeError("cannot locate cursor for AI control indicator") @@ -265,6 +281,8 @@ def _run(self) -> None: rects: tuple[tuple[int, int, int, int], ...] = () mode: bool | None = None visible = False + was_active = False + breath_started = 0.0 try: self.heartbeat = time.monotonic() self.started.set() @@ -276,6 +294,8 @@ def _run(self) -> None: active, pending = self.active, self.pending suspended, version = self.suspended, self.version if active: + if not was_active: + breath_started = time.monotonic() current_rects = _monitor_rects() if ring is None or current_rects != rects or mode != pending: for layer in reversed(layers): @@ -294,10 +314,17 @@ def _run(self) -> None: ring.hide() visible = False else: + opacity = _breath_opacity(time.monotonic() - breath_started) + for layer in layers: + if layer.breathes: + layer.set_opacity(opacity) + if not pending: + ring.set_opacity(opacity) if not visible: for layer in layers: layer.show() - ring.show() + if not pending: + ring.show() visible = True if not pending: point = flash_overlay._POINT() @@ -311,6 +338,7 @@ def _run(self) -> None: if ring: ring.hide() visible = False + was_active = active # Pump messages on the owning thread, including display events. for layer in layers: flash_overlay._pump_messages(layer.hwnd) diff --git a/src/windows_mcp/desktop/control_overlay_art.py b/src/windows_mcp/desktop/control_overlay_art.py new file mode 100644 index 00000000..96723011 --- /dev/null +++ b/src/windows_mcp/desktop/control_overlay_art.py @@ -0,0 +1,167 @@ +"""Pure bitmap artwork shared by the active AI-control indicator.""" + +import math +import os +from pathlib import Path + +from PIL import Image, ImageChops, ImageDraw, ImageFilter, ImageFont + +from windows_mcp.desktop.overlay_bitmap import premultiplied_bgra + +_BORDER = 56 +_CURSOR_SIZE = 88 +_NOTICE_GLOW_PAD = 32 +_BREATH_PERIOD_SECONDS = 3.0 +_BREATH_MIN_ALPHA = 140 +_NOTICE_TITLE = "AI is controlling this computer" +_NOTICE_SHORTCUT = "Ctrl + Alt + Shift + Backspace" +_NOTICE_HINT = f"Press {_NOTICE_SHORTCUT} to take over" + + +def _breath_opacity(elapsed: float) -> int: + """Start bright and cycle smoothly while keeping the control cue visible.""" + strength = (1 + math.cos(2 * math.pi * elapsed / _BREATH_PERIOD_SECONDS)) / 2 + return round(_BREATH_MIN_ALPHA + (255 - _BREATH_MIN_ALPHA) * strength) + + +def _glow_alpha(depth: int, extent: int = _BORDER) -> int: + """Peak at the screen edge and fade continuously to transparency inward.""" + if depth >= extent - 1: + return 0 + return round(240 * (1 - depth / (extent - 1)) ** 1.4) + + +def _edge_bitmap(width: int, height: int, side: str, color: tuple[int, int, int]) -> bytes: + """Pre-render feathered edges, joining corners at one glow strength.""" + image = Image.new("RGBA", (width, height)) + draw = ImageDraw.Draw(image) + extent = min(_BORDER, height if side in ("top", "bottom") else width) + corner_width = min(extent, width // 2) + for depth in range(extent): + # Let the gradient itself reach the edge; do not draw a separate outline. + alpha = _glow_alpha(depth, extent) + if side in ("top", "bottom"): + y = depth if side == "top" else height - 1 - depth + draw.line((0, y, width - 1, y), fill=(*color, alpha)) + # The side strips start below/above these corner squares. Use the + # nearest screen edge as the depth, matching the straight strips + # without blending two windows or drawing a diagonal bright seam. + for side_depth in range(corner_width): + corner_alpha = _glow_alpha(min(depth, side_depth), extent) + draw.point((side_depth, y), fill=(*color, corner_alpha)) + draw.point((width - 1 - side_depth, y), fill=(*color, corner_alpha)) + elif side == "left": + draw.line((depth, 0, depth, height - 1), fill=(*color, alpha)) + else: + x = width - 1 - depth + draw.line((x, 0, x, height - 1), fill=(*color, alpha)) + return premultiplied_bgra(image, 1.0) + + +def _cursor_bitmap(color: tuple[int, int, int]) -> bytes: + """Build a broad cursor aura with no drawn ellipse or sharp contour.""" + half = (_CURSOR_SIZE - 1) / 2 + pixels = [] + for y in range(_CURSOR_SIZE): + for x in range(_CURSOR_SIZE): + radius = math.hypot(x - half, y - half) + clear_center = 1 - math.exp(-0.5 * (max(0.0, radius - 1) / 5) ** 2) + outer_glow = math.exp(-0.5 * (radius / 20) ** 2) + edge_fade = min(1.0, max(0.0, (half - radius) / 8)) + alpha = round(230 * clear_center * outer_glow * edge_fade) + pixels.append((*color, alpha)) + image = Image.new("RGBA", (_CURSOR_SIZE, _CURSOR_SIZE)) + image.putdata(pixels) + return premultiplied_bgra(image, 1.0) + + +def _notice_shape_mask(width: int, height: int, radius: int = 22) -> Image.Image: + """Downsample a rounded mask so all four corners have smooth pixel coverage.""" + scale = 8 + mask = Image.new("L", (width * scale, height * scale)) + ImageDraw.Draw(mask).rounded_rectangle( + (0, 0, width * scale - 1, height * scale - 1), radius=radius * scale, fill=255 + ) + return mask.resize((width, height), Image.Resampling.LANCZOS).filter( + ImageFilter.GaussianBlur(3) + ) + + +def _notice_panel_mask(width: int, height: int) -> Image.Image: + """Fade the rounded panel through its outer eight pixels.""" + edge_fade = [round(255 * (step / 8) ** 2 * (3 - 2 * step / 8)) for step in range(9)] + horizontal = [edge_fade[min(8, x, width - 1 - x)] for x in range(width)] + vertical = [edge_fade[min(8, y, height - 1 - y)] for y in range(height)] + fade_mask = Image.new("L", (width, height)) + fade_mask.putdata([min(horizontal[x], vertical[y]) for y in range(height) for x in range(width)]) + return ImageChops.multiply(_notice_shape_mask(width, height), fade_mask) + + +def _notice_bitmap(screen_width: int) -> tuple[int, int, bytes] | None: + """Render a centered two-line prompt below the upper glow.""" + available_width = screen_width - 32 + if available_width < 200: + return None + horizontal_padding = 32 if screen_width >= 640 else 20 + font_dir = Path(os.environ.get("WINDIR", r"C:\Windows")) / "Fonts" + # Prefer presentation-sized text, then shrink only when a monitor is narrow. + for hint_size in range(24, 7, -1): + try: + title_font = ImageFont.truetype(str(font_dir / "seguisb.ttf"), hint_size + 8) + hint_font = ImageFont.truetype(str(font_dir / "segoeui.ttf"), hint_size) + except OSError: + title_font = hint_font = ImageFont.load_default() + measure = ImageDraw.Draw(Image.new("RGBA", (1, 1))) + title_box = measure.textbbox((0, 0), _NOTICE_TITLE, font=title_font) + hint_box = measure.textbbox((0, 0), _NOTICE_HINT, font=hint_font) + hint_prefix, _, hint_suffix = _NOTICE_HINT.partition(_NOTICE_SHORTCUT) + prefix_width = math.ceil(measure.textlength(hint_prefix, font=hint_font)) + shortcut_width = math.ceil(measure.textlength(_NOTICE_SHORTCUT, font=hint_font)) + suffix_width = math.ceil(measure.textlength(hint_suffix, font=hint_font)) + title_width, title_height = title_box[2] - title_box[0], title_box[3] - title_box[1] + hint_width, hint_height = prefix_width + shortcut_width + suffix_width + 16, hint_box[3] - hint_box[1] + width = max(title_width, hint_width) + 2 * horizontal_padding + if width <= available_width: + break + else: + return None + + height = title_height + hint_height + 46 + image = Image.new("RGBA", (width, height), (8, 42, 88, 0)) + # Fade every edge from transparent to 80%; the rounded mask softens the corners too. + image.putalpha(_notice_panel_mask(width, height).point(lambda alpha: round(alpha * 0.8))) + draw = ImageDraw.Draw(image) + title_y = 18 - title_box[1] + hint_y = 28 + title_height - hint_box[1] + draw.text(((width - title_width) // 2 - title_box[0], title_y), _NOTICE_TITLE, + font=title_font, fill=(248, 251, 255, 255)) + hint_x = (width - hint_width) // 2 + badge_left = hint_x + prefix_width + # A solid black key label separates the takeover shortcut from its sentence. + draw.rounded_rectangle( + (badge_left, 24 + title_height, badge_left + shortcut_width + 15, + 31 + title_height + hint_height), + radius=6, fill=(0, 0, 0, 255), + ) + draw.text((hint_x, hint_y), hint_prefix, font=hint_font, fill=(225, 240, 255, 255)) + draw.text((badge_left + 8, hint_y), _NOTICE_SHORTCUT, font=hint_font, fill=(255, 255, 255, 255)) + draw.text((badge_left + shortcut_width + 16, hint_y), hint_suffix, + font=hint_font, fill=(225, 240, 255, 255)) + return width, height, premultiplied_bgra(image, 1.0) + + +def _notice_glow_bitmap( + width: int, height: int, color: tuple[int, int, int], pad: int = _NOTICE_GLOW_PAD +) -> bytes: + """Anchor the aura at the solid start of the panel's outer fade.""" + size = (width + 2 * pad, height + 2 * pad) + shape = Image.new("L", size) + # The inset contour meets the panel where its eight-pixel fade begins. + shape.paste(_notice_shape_mask(width - 16, height - 16, 14), (pad + 8, pad + 8)) + outside = ImageChops.subtract(shape.filter(ImageFilter.GaussianBlur(min(15, pad / 2))), shape) + # Feather the inner cutoff as well, so the halo cannot end as a bright line. + outside = outside.filter(ImageFilter.GaussianBlur(3)) + aura = Image.new("RGBA", size, (*color, 0)) + peak = outside.getextrema()[1] + aura.putalpha(outside.point(lambda alpha: round(alpha * 211 / peak))) + return premultiplied_bgra(aura, 1.0) diff --git a/src/windows_mcp/desktop/flash_overlay.py b/src/windows_mcp/desktop/flash_overlay.py index b755729b..f2922722 100644 --- a/src/windows_mcp/desktop/flash_overlay.py +++ b/src/windows_mcp/desktop/flash_overlay.py @@ -402,7 +402,10 @@ def _create_layered_window(class_name: str, x: int, y: int, w: int, h: int): return hwnd, h_instance -def _push_bitmap(hwnd, x: int, y: int, w: int, h: int, bgra: bytes) -> None: +def _push_bitmap( + hwnd, x: int, y: int, w: int, h: int, bgra: bytes, *, opacity: int = 255 +) -> None: + """Upload premultiplied pixels with optional whole-window opacity.""" screen_dc = _user32.GetDC(None) if not screen_dc: raise OSError("GetDC failed") @@ -437,7 +440,7 @@ def _push_bitmap(hwnd, x: int, y: int, w: int, h: int, bgra: bytes) -> None: pos = _POINT(x, y) size = _SIZE(w, h) src_pos = _POINT(0, 0) - blend = _BLENDFUNCTION(_AC_SRC_OVER, 0, 255, _AC_SRC_ALPHA) + blend = _BLENDFUNCTION(_AC_SRC_OVER, 0, opacity, _AC_SRC_ALPHA) ok = _user32.UpdateLayeredWindow( hwnd, screen_dc, diff --git a/src/windows_mcp/desktop/overlay_bitmap.py b/src/windows_mcp/desktop/overlay_bitmap.py index aef1075d..229cf4e0 100644 --- a/src/windows_mcp/desktop/overlay_bitmap.py +++ b/src/windows_mcp/desktop/overlay_bitmap.py @@ -18,6 +18,10 @@ def premultiplied_bgra(rgba_image: Image.Image, intensity: float) -> bytes: for i in range(0, len(bgra), 4): a = bgra[i + 3] if a == 0: + # Fully transparent layered pixels must have zero RGB too. + bgra[i] = 0 + bgra[i + 1] = 0 + bgra[i + 2] = 0 continue bgra[i] = (bgra[i] * a) // 255 bgra[i + 1] = (bgra[i + 1] * a) // 255 diff --git a/tests/test_control_overlay.py b/tests/test_control_overlay.py index 1833b50d..3b59c238 100644 --- a/tests/test_control_overlay.py +++ b/tests/test_control_overlay.py @@ -19,57 +19,6 @@ def _stop_indicator(): control_overlay.stop() -def test_edge_and_cursor_are_translucent(): - edge = control_overlay._edge_bitmap(20, 28, "left", (45, 145, 255)) - assert len(edge) == 20 * 28 * 4 - assert edge[3] > edge[-1] - ring = control_overlay._cursor_bitmap((45, 145, 255)) - assert len(ring) == control_overlay._CURSOR_SIZE**2 * 4 - assert ring[3] == 0 # no opaque square over the desktop - - -def test_multimonitor_creates_narrow_clickthrough_layers(monkeypatch): - made = [] - - class FakeLayer: - def __init__(self, x, y, width, height, bgra, name): - made.append((x, y, width, height, name)) - - def close(self): - pass - - monkeypatch.setattr(control_overlay, "_Layer", FakeLayer) - monkeypatch.setattr(control_overlay._user32, "GetCursorPos", lambda ptr: True) - rects = ((-1920, 0, 0, 1080), (0, 0, 2560, 1440)) - edges, ring = control_overlay._build_layers(rects, pending=False) - assert len(edges) == 8 - assert ring is not None - assert made[0][:4] == (-1920, 0, 1920, 28) - assert made[4][:4] == (0, 0, 2560, 28) - assert all(w <= 28 or h <= 28 for _, _, w, h, name in made if name != "cursor") - - -def test_layer_failure_closes_prior_windows(monkeypatch): - closed = [] - - class FailingLayer: - count = 0 - - def __init__(self, *args): - self.count = FailingLayer.count - FailingLayer.count += 1 - if self.count == 2: - raise OSError("window creation failed") - - def close(self): - closed.append(self.count) - - monkeypatch.setattr(control_overlay, "_Layer", FailingLayer) - with pytest.raises(OSError, match="window creation failed"): - control_overlay._build_layers(((0, 0, 800, 600),), pending=False) - assert closed == [1, 0] - - def test_capture_suspends_and_restores_including_error(monkeypatch): calls = [] @@ -90,6 +39,10 @@ def test_indicator_lifecycle_and_no_restore_after_takeover(monkeypatch): class FakeLayer: hwnd = 1 + breathes = True + + def set_opacity(self, opacity): + events.append(("opacity", opacity)) def show(self): events.append("show") diff --git a/tests/test_control_overlay_visuals.py b/tests/test_control_overlay_visuals.py new file mode 100644 index 00000000..f1a80f0b --- /dev/null +++ b/tests/test_control_overlay_visuals.py @@ -0,0 +1,282 @@ +"""Regression checks for the accepted AI-control glow and notice.""" + +import pytest + +from windows_mcp.desktop import control_overlay, control_overlay_art, flash_overlay + + +def _assert_premultiplied(bgra: bytes) -> None: + for offset in range(0, len(bgra), 4): + blue, green, red, alpha = bgra[offset : offset + 4] + assert blue <= alpha and green <= alpha and red <= alpha + + +def _alpha_at(bgra: bytes, width: int, x: int, y: int) -> int: + return bgra[(y * width + x) * 4 + 3] + + +def test_notice_uses_clear_english_and_transparent_corners(): + assert control_overlay_art._NOTICE_TITLE == "AI is controlling this computer" + assert control_overlay_art._NOTICE_HINT == "Press Ctrl + Alt + Shift + Backspace to take over" + assert control_overlay_art._NOTICE_SHORTCUT == "Ctrl + Alt + Shift + Backspace" + notice = control_overlay_art._notice_bitmap(1920) + assert notice is not None + width, height, bgra = notice + assert 540 <= width < 1920 and height < 120 # No third line remains. + assert len(bgra) == width * height * 4 + # A supersampled rounded mask gives each corner several partial-alpha pixels. + for x_start in (0, width - 22): + for y_start in (0, height - 22): + corner_alpha = { + _alpha_at(bgra, width, x_start + dx, y_start + dy) + for dx in range(22) + for dy in range(22) + } + assert len({alpha for alpha in corner_alpha if 0 < alpha < 204}) >= 8 + for x in (0, width - 1): + for y in (0, height - 1): + assert _alpha_at(bgra, width, x, y) == 0 + # Each straight side now tapers from transparent to the 80%-opaque center. + for edge in ( + [_alpha_at(bgra, width, d, height // 2) for d in range(9)], + [_alpha_at(bgra, width, width - 1 - d, height // 2) for d in range(9)], + [_alpha_at(bgra, width, width // 2, d) for d in range(9)], + [_alpha_at(bgra, width, width // 2, height - 1 - d) for d in range(9)], + ): + assert edge[0] == 0 < edge[1] < edge[4] < edge[8] == 204 + assert all(a <= b for a, b in zip(edge, edge[1:])) + assert _alpha_at(bgra, width, 10, height // 2) == 204 # 204/255 is 0.8. + assert max(bgra[3::4]) == 255 # Text stays fully legible while the glow breathes. + assert any(bgra[offset : offset + 4] == b"\x00\x00\x00\xff" for offset in range(0, len(bgra), 4)) + narrow = control_overlay_art._notice_bitmap(300) + assert narrow is not None and narrow[0] <= 268 + + +def test_notice_outer_glow_fades_without_changing_the_panel(): + width, height = 400, 100 + pad = control_overlay_art._NOTICE_GLOW_PAD + glow = control_overlay_art._notice_glow_bitmap(width, height, (45, 145, 255)) + glow_width = width + 2 * pad + _assert_premultiplied(glow) + assert len(glow) == glow_width * (height + 2 * pad) * 4 + middle_y = pad + height // 2 + assert _alpha_at(glow, glow_width, pad + width // 2, middle_y) == 0 + assert _alpha_at(glow, glow_width, 0, 0) == 0 + assert 208 <= max(glow[3::4]) <= 214 # Keep the accepted aura brightness after shifting it inward. + assert _alpha_at(glow, glow_width, pad - 1, middle_y) > _alpha_at( + glow, glow_width, pad - 15, middle_y + ) > 0 + + +def test_notice_glow_meets_the_feathered_background(): + notice = control_overlay_art._notice_bitmap(1920) + assert notice is not None + width, height, panel = notice + pad = control_overlay_art._NOTICE_GLOW_PAD + halo = control_overlay_art._notice_glow_bitmap(width, height, (45, 145, 255)) + halo_width = width + 2 * pad + for edge in ( + [(depth, height // 2) for depth in range(13)], + [(width - 1 - depth, height // 2) for depth in range(13)], + [(width // 2, depth) for depth in range(13)], + [(width // 2, height - 1 - depth) for depth in range(13)], + ): + combined = [] + for depth, (panel_x, panel_y) in enumerate(edge): + panel_alpha = _alpha_at(panel, width, panel_x, panel_y) + halo_alpha = _alpha_at(halo, halo_width, pad + panel_x, pad + panel_y) + if depth < 8: + assert halo_alpha > 0 + if depth == 8: + assert panel_alpha == 204 + assert halo_alpha > 0 # The aura now passes the fade's start. + combined.append(round(panel_alpha + halo_alpha * (255 - panel_alpha) / 255)) + assert max(abs(a - b) for a, b in zip(combined, combined[1:])) <= 15 + for x_edge, y_edge, x_step, y_step in ( + (0, 0, 1, 1), + (width - 1, 0, -1, 1), + (0, height - 1, 1, -1), + (width - 1, height - 1, -1, -1), + ): + # The rounded silhouette also needs a gradual combined transition. + combined = [] + for depth in range(16): + x, y = x_edge + depth * x_step, y_edge + depth * y_step + panel_alpha = _alpha_at(panel, width, x, y) + halo_alpha = _alpha_at(halo, halo_width, pad + x, pad + y) + combined.append(round(panel_alpha + halo_alpha * (255 - panel_alpha) / 255)) + assert max(abs(a - b) for a, b in zip(combined, combined[1:])) <= 30 + + +def test_breath_opacity_is_smooth_periodic_and_never_disappears(): + opacity = control_overlay_art._breath_opacity + period = control_overlay_art._BREATH_PERIOD_SECONDS + assert opacity(0) == opacity(period) == 255 + assert opacity(period / 2) == control_overlay_art._BREATH_MIN_ALPHA == 140 + assert opacity(period / 4) == opacity(3 * period / 4) + frames = [opacity(index * control_overlay._REFRESH_SECONDS) for index in range(61)] + assert all(140 <= value <= 255 for value in frames) + assert max(abs(a - b) for a, b in zip(frames, frames[1:])) <= 7 + + +def test_layer_opacity_updates_only_on_change_and_fails_closed(monkeypatch): + layer = object.__new__(control_overlay._Layer) + layer.hwnd = 123 + layer.x, layer.y, layer.width, layer.height = 10, 20, 1, 1 + layer.bgra = bytes((0, 0, 0, 0)) + layer.opacity = None + calls = [] + + def push_bitmap(*args, **kwargs): + calls.append((args, kwargs)) + + monkeypatch.setattr(flash_overlay, "_push_bitmap", push_bitmap) + layer.set_opacity(200) + layer.set_opacity(200) + layer.set_opacity(140) + assert calls == [ + ((123, 10, 20, 1, 1, layer.bgra), {"opacity": 200}), + ((123, 10, 20, 1, 1, layer.bgra), {"opacity": 140}), + ] + + def failed_upload(*args, **kwargs): + raise OSError("upload failed") + + monkeypatch.setattr(flash_overlay, "_push_bitmap", failed_upload) + with pytest.raises(OSError, match="upload failed"): + layer.set_opacity(255) + assert layer.opacity == 140 + + +def test_edge_and_cursor_have_soft_glow_without_a_solid_contour(): + size = control_overlay_art._BORDER + assert size == 56 # The requested doubling of the former 28-pixel glow. + for side in ("left", "right", "top", "bottom"): + width = size * 3 if side in ("top", "bottom") else size + height = size if side in ("top", "bottom") else size * 3 + edge = control_overlay_art._edge_bitmap(width, height, side, (45, 145, 255)) + assert len(edge) == width * height * 4 + _assert_premultiplied(edge) + alphas = ( + [_alpha_at(edge, width, x, height // 2) for x in range(width)] + if side in ("left", "right") + else [_alpha_at(edge, width, width // 2, y) for y in range(height)] + ) + if side in ("right", "bottom"): + alphas.reverse() + assert alphas[0] == max(alphas) == 240 # Twice the previous edge opacity. + assert alphas[-1] == 0 + assert all(a >= b for a, b in zip(alphas, alphas[1:])) + assert max(a - b for a, b in zip(alphas, alphas[1:])) <= 8 + assert 30 <= alphas[40] <= 50 # The glow remains broad but fades inward. + + glow = control_overlay_art._cursor_bitmap((45, 145, 255)) + width = control_overlay_art._CURSOR_SIZE + assert len(glow) == width**2 * 4 + _assert_premultiplied(glow) + alphas = [glow[((width // 2) * width + x) * 4 + 3] for x in range(width // 2, width)] + assert glow[3] == 0 # Transparent square corners. + assert alphas[0] == 0 < alphas[10] # Cursor center is fully transparent. + assert alphas[10] > alphas[25] > alphas[-1] + assert alphas[-1] == 0 + assert max(alphas) == 176 # Twice the previous cursor glow peak of 88. + assert max(abs(a - b) for a, b in zip(alphas, alphas[1:])) < 30 + + +def test_corner_glow_joins_edges_without_double_opacity(): + border = control_overlay_art._BORDER + width = border * 4 + top = control_overlay_art._edge_bitmap(width, border, "top", (45, 145, 255)) + bottom = control_overlay_art._edge_bitmap(width, border, "bottom", (45, 145, 255)) + left = control_overlay_art._edge_bitmap(border, border, "left", (45, 145, 255)) + assert _alpha_at(top, width, 0, 0) == control_overlay_art._glow_alpha(0) + assert _alpha_at(top, width, width - 1, 0) == control_overlay_art._glow_alpha(0) + for horizontal in (top, bottom): + for x_depth in (1, 5, 10, 20, 40, 55): + side_alpha = _alpha_at(left, border, x_depth, border // 2) + for y_depth in (1, 5, 10, 20, 40, 55): + y = y_depth if horizontal is top else border - 1 - y_depth + straight_alpha = _alpha_at(horizontal, width, width // 2, y) + expected = control_overlay_art._glow_alpha(min(x_depth, y_depth)) + assert expected <= max(side_alpha, straight_alpha) + assert _alpha_at(horizontal, width, x_depth, y) == expected + assert _alpha_at(horizontal, width, width - 1 - x_depth, y) == expected + + +def test_multimonitor_creates_narrow_clickthrough_layers(monkeypatch): + made = [] + + class FakeLayer: + def __init__(self, x, y, width, height, bgra, name, *, breathes=True): + made.append((x, y, width, height, name, breathes)) + + def close(self): + pass + + monkeypatch.setattr(control_overlay, "_Layer", FakeLayer) + monkeypatch.setattr(control_overlay._user32, "GetCursorPos", lambda ptr: True) + rects = ((-1920, 0, 0, 1080), (0, 0, 2560, 1440)) + layers, ring = control_overlay._build_layers(rects, pending=False) + assert len(layers) == 12 + assert ring is not None + border = control_overlay_art._BORDER + assert made[0][:4] == (-1920, 0, 1920, border) + assert made[2][:4] == (-1920, border, border, 1080 - 2 * border) + assert made[3][:4] == (-border, border, border, 1080 - 2 * border) + assert made[4][1] == border and made[4][4:] == ("0_notice_glow", True) + assert made[5][1] == border + control_overlay_art._NOTICE_GLOW_PAD + assert made[5][4:] == ("0_notice", False) + assert made[4][0] == made[5][0] - control_overlay_art._NOTICE_GLOW_PAD + assert abs((made[5][0] + made[5][2] / 2) - (-1920 / 2)) <= 1 + assert made[6][:4] == (0, 0, 2560, border) + assert made[8][:4] == (0, border, border, 1440 - 2 * border) + assert made[9][:4] == (2560 - border, border, border, 1440 - 2 * border) + assert made[10][1] == border and made[10][4:] == ("1_notice_glow", True) + assert made[11][1] == border + control_overlay_art._NOTICE_GLOW_PAD + assert made[11][4:] == ("1_notice", False) + assert abs((made[11][0] + made[11][2] / 2) - 1280) <= 1 + assert all( + w <= border or h <= border + for _, _, w, h, name, _ in made + if name != "cursor" and not name.endswith(("notice", "notice_glow")) + ) + + +@pytest.mark.parametrize("width", [300, 400]) +def test_notice_glow_stays_on_its_monitor_when_narrow(monkeypatch, width): + made = [] + + class FakeLayer: + def __init__(self, x, y, layer_width, height, bgra, name, *, breathes=True): + made.append((x, y, layer_width, height, name)) + + def close(self): + pass + + monkeypatch.setattr(control_overlay, "_Layer", FakeLayer) + monkeypatch.setattr(control_overlay._user32, "GetCursorPos", lambda ptr: True) + control_overlay._build_layers(((100, 0, 100 + width, 300),), pending=False) + glow = next(layer for layer in made if layer[4] == "0_notice_glow") + assert 100 <= glow[0] + assert glow[0] + glow[2] <= 100 + width + + +def test_layer_failure_closes_prior_windows(monkeypatch): + closed = [] + + class FailingLayer: + count = 0 + + def __init__(self, *args): + self.count = FailingLayer.count + FailingLayer.count += 1 + if self.count == 2: + raise OSError("window creation failed") + + def close(self): + closed.append(self.count) + + monkeypatch.setattr(control_overlay, "_Layer", FailingLayer) + with pytest.raises(OSError, match="window creation failed"): + control_overlay._build_layers(((0, 0, 800, 600),), pending=False) + assert closed == [1, 0] From d501a7d434c268cf826b9797a853e6778cd191ff Mon Sep 17 00:00:00 2001 From: Zeng Fanfan Date: Sat, 3 Oct 2026 06:38:38 +0800 Subject: [PATCH 9/9] Format restored desktop control visuals --- src/windows_mcp/desktop/control_overlay.py | 56 ++++++++++++------- .../desktop/control_overlay_art.py | 36 +++++++++--- src/windows_mcp/desktop/flash_overlay.py | 4 +- tests/test_control_overlay_visuals.py | 16 ++++-- 4 files changed, 75 insertions(+), 37 deletions(-) diff --git a/src/windows_mcp/desktop/control_overlay.py b/src/windows_mcp/desktop/control_overlay.py index fb5de330..cee71992 100644 --- a/src/windows_mcp/desktop/control_overlay.py +++ b/src/windows_mcp/desktop/control_overlay.py @@ -43,7 +43,15 @@ class _Layer: def __init__( - self, x: int, y: int, width: int, height: int, bgra: bytes, name: str, *, breathes: bool = True + self, + x: int, + y: int, + width: int, + height: int, + bgra: bytes, + name: str, + *, + breathes: bool = True, ): self.x, self.y, self.width, self.height = x, y, width, height self.breathes = breathes @@ -121,7 +129,9 @@ def _monitor_rects() -> tuple[tuple[int, int, int, int], ...]: return tuple((r.left, r.top, r.right, r.bottom) for r in uia.GetMonitorsRect()) -def _build_layers(rects: tuple[tuple[int, int, int, int], ...], pending: bool) -> tuple[list[_Layer], _Layer]: +def _build_layers( + rects: tuple[tuple[int, int, int, int], ...], pending: bool +) -> tuple[list[_Layer], _Layer]: if not rects: raise RuntimeError("no display available for AI control indicator") color = _AMBER if pending else _BLUE @@ -144,7 +154,9 @@ def _build_layers(rects: tuple[tuple[int, int, int, int], ...], pending: bool) - ) for x, y, w, h, side in strips: if w and h: - layers.append(_Layer(x, y, w, h, _edge_bitmap(w, h, side, color), f"{index}_{side}")) + layers.append( + _Layer(x, y, w, h, _edge_bitmap(w, h, side, color), f"{index}_{side}") + ) notice = _notice_bitmap(width) if notice is not None: notice_width, notice_height, bitmap = notice @@ -154,23 +166,27 @@ def _build_layers(rects: tuple[tuple[int, int, int, int], ...], pending: bool) - continue notice_x = left + (width - notice_width) // 2 notice_y = top + border + glow_pad - layers.append(_Layer( - notice_x - glow_pad, - notice_y - glow_pad, - notice_width + 2 * glow_pad, - notice_height + 2 * glow_pad, - _notice_glow_bitmap(notice_width, notice_height, color, glow_pad), - f"{index}_notice_glow", - )) - layers.append(_Layer( - notice_x, - notice_y, - notice_width, - notice_height, - bitmap, - f"{index}_notice", - breathes=False, - )) + layers.append( + _Layer( + notice_x - glow_pad, + notice_y - glow_pad, + notice_width + 2 * glow_pad, + notice_height + 2 * glow_pad, + _notice_glow_bitmap(notice_width, notice_height, color, glow_pad), + f"{index}_notice_glow", + ) + ) + layers.append( + _Layer( + notice_x, + notice_y, + notice_width, + notice_height, + bitmap, + f"{index}_notice", + breathes=False, + ) + ) point = flash_overlay._POINT() if not _user32.GetCursorPos(ctypes.byref(point)): raise RuntimeError("cannot locate cursor for AI control indicator") diff --git a/src/windows_mcp/desktop/control_overlay_art.py b/src/windows_mcp/desktop/control_overlay_art.py index 96723011..9a437445 100644 --- a/src/windows_mcp/desktop/control_overlay_art.py +++ b/src/windows_mcp/desktop/control_overlay_art.py @@ -93,7 +93,9 @@ def _notice_panel_mask(width: int, height: int) -> Image.Image: horizontal = [edge_fade[min(8, x, width - 1 - x)] for x in range(width)] vertical = [edge_fade[min(8, y, height - 1 - y)] for y in range(height)] fade_mask = Image.new("L", (width, height)) - fade_mask.putdata([min(horizontal[x], vertical[y]) for y in range(height) for x in range(width)]) + fade_mask.putdata( + [min(horizontal[x], vertical[y]) for y in range(height) for x in range(width)] + ) return ImageChops.multiply(_notice_shape_mask(width, height), fade_mask) @@ -119,7 +121,10 @@ def _notice_bitmap(screen_width: int) -> tuple[int, int, bytes] | None: shortcut_width = math.ceil(measure.textlength(_NOTICE_SHORTCUT, font=hint_font)) suffix_width = math.ceil(measure.textlength(hint_suffix, font=hint_font)) title_width, title_height = title_box[2] - title_box[0], title_box[3] - title_box[1] - hint_width, hint_height = prefix_width + shortcut_width + suffix_width + 16, hint_box[3] - hint_box[1] + hint_width, hint_height = ( + prefix_width + shortcut_width + suffix_width + 16, + hint_box[3] - hint_box[1], + ) width = max(title_width, hint_width) + 2 * horizontal_padding if width <= available_width: break @@ -133,20 +138,33 @@ def _notice_bitmap(screen_width: int) -> tuple[int, int, bytes] | None: draw = ImageDraw.Draw(image) title_y = 18 - title_box[1] hint_y = 28 + title_height - hint_box[1] - draw.text(((width - title_width) // 2 - title_box[0], title_y), _NOTICE_TITLE, - font=title_font, fill=(248, 251, 255, 255)) + draw.text( + ((width - title_width) // 2 - title_box[0], title_y), + _NOTICE_TITLE, + font=title_font, + fill=(248, 251, 255, 255), + ) hint_x = (width - hint_width) // 2 badge_left = hint_x + prefix_width # A solid black key label separates the takeover shortcut from its sentence. draw.rounded_rectangle( - (badge_left, 24 + title_height, badge_left + shortcut_width + 15, - 31 + title_height + hint_height), - radius=6, fill=(0, 0, 0, 255), + ( + badge_left, + 24 + title_height, + badge_left + shortcut_width + 15, + 31 + title_height + hint_height, + ), + radius=6, + fill=(0, 0, 0, 255), ) draw.text((hint_x, hint_y), hint_prefix, font=hint_font, fill=(225, 240, 255, 255)) draw.text((badge_left + 8, hint_y), _NOTICE_SHORTCUT, font=hint_font, fill=(255, 255, 255, 255)) - draw.text((badge_left + shortcut_width + 16, hint_y), hint_suffix, - font=hint_font, fill=(225, 240, 255, 255)) + draw.text( + (badge_left + shortcut_width + 16, hint_y), + hint_suffix, + font=hint_font, + fill=(225, 240, 255, 255), + ) return width, height, premultiplied_bgra(image, 1.0) diff --git a/src/windows_mcp/desktop/flash_overlay.py b/src/windows_mcp/desktop/flash_overlay.py index f2922722..7c5e2f4b 100644 --- a/src/windows_mcp/desktop/flash_overlay.py +++ b/src/windows_mcp/desktop/flash_overlay.py @@ -402,9 +402,7 @@ def _create_layered_window(class_name: str, x: int, y: int, w: int, h: int): return hwnd, h_instance -def _push_bitmap( - hwnd, x: int, y: int, w: int, h: int, bgra: bytes, *, opacity: int = 255 -) -> None: +def _push_bitmap(hwnd, x: int, y: int, w: int, h: int, bgra: bytes, *, opacity: int = 255) -> None: """Upload premultiplied pixels with optional whole-window opacity.""" screen_dc = _user32.GetDC(None) if not screen_dc: diff --git a/tests/test_control_overlay_visuals.py b/tests/test_control_overlay_visuals.py index f1a80f0b..c5461cda 100644 --- a/tests/test_control_overlay_visuals.py +++ b/tests/test_control_overlay_visuals.py @@ -47,7 +47,9 @@ def test_notice_uses_clear_english_and_transparent_corners(): assert all(a <= b for a, b in zip(edge, edge[1:])) assert _alpha_at(bgra, width, 10, height // 2) == 204 # 204/255 is 0.8. assert max(bgra[3::4]) == 255 # Text stays fully legible while the glow breathes. - assert any(bgra[offset : offset + 4] == b"\x00\x00\x00\xff" for offset in range(0, len(bgra), 4)) + assert any( + bgra[offset : offset + 4] == b"\x00\x00\x00\xff" for offset in range(0, len(bgra), 4) + ) narrow = control_overlay_art._notice_bitmap(300) assert narrow is not None and narrow[0] <= 268 @@ -62,10 +64,14 @@ def test_notice_outer_glow_fades_without_changing_the_panel(): middle_y = pad + height // 2 assert _alpha_at(glow, glow_width, pad + width // 2, middle_y) == 0 assert _alpha_at(glow, glow_width, 0, 0) == 0 - assert 208 <= max(glow[3::4]) <= 214 # Keep the accepted aura brightness after shifting it inward. - assert _alpha_at(glow, glow_width, pad - 1, middle_y) > _alpha_at( - glow, glow_width, pad - 15, middle_y - ) > 0 + assert ( + 208 <= max(glow[3::4]) <= 214 + ) # Keep the accepted aura brightness after shifting it inward. + assert ( + _alpha_at(glow, glow_width, pad - 1, middle_y) + > _alpha_at(glow, glow_width, pad - 15, middle_y) + > 0 + ) def test_notice_glow_meets_the_feathered_background():