agentic-mobile-control/apps/device-host-agent/host_agent/mcp_lock.py

"""Per-device MCP session-level busy tracker.

The cloud-side assignment path and the MCP-driven path both drive devices
through the same in-process ``DeviceManager``. This tracker records which
devices are currently held by an MCP session so that:

- MCP tool calls against a device held by another session (or by a cloud
  assignment — checked separately by the caller via ``AgentStatusTracker``)
  can fail fast with a busy error.
- The heartbeat payload can advertise ``mcp_busy_device_ids`` so the cloud
  scheduler won't dispatch conflicting assignments to the same device.

Leases expire ``ttl_seconds`` after the last ``renew()`` call (set on every
tool call from the holding session). Expired leases are lazy-swept on read.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from datetime import UTC, datetime
from threading import Lock
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Callable


@dataclass(frozen=True)
class McpDeviceLease:
    device_id: str
    session_id: str
    acquired_at: datetime
    last_seen_at: datetime


class McpBusyTracker:
    def __init__(
        self,
        *,
        ttl_seconds: float = 20.0,
        now: Callable[[], datetime] | None = None,
    ) -> None:
        self._ttl = float(ttl_seconds)
        self._now = now or (lambda: datetime.now(UTC))
        self._lock = Lock()
        # device_id -> McpDeviceLease
        self._leases: dict[str, McpDeviceLease] = {}

    def acquire(self, device_id: str, session_id: str) -> bool:
        with self._lock:
            self._sweep_locked()
            existing = self._leases.get(device_id)
            if existing is not None and existing.session_id != session_id:
                return False
            now = self._now()
            lease = McpDeviceLease(
                device_id=device_id,
                session_id=session_id,
                acquired_at=(existing.acquired_at if existing is not None else now),
                last_seen_at=now,
            )
            self._leases[device_id] = lease
            return True

    def renew(self, device_id: str, session_id: str) -> bool:
        with self._lock:
            self._sweep_locked()
            existing = self._leases.get(device_id)
            # Tolerate boundary: lease may have been swept, but if the caller
            # is the legitimate previous holder, re-acquire on their behalf.
            if existing is None:
                now = self._now()
                self._leases[device_id] = McpDeviceLease(
                    device_id=device_id,
                    session_id=session_id,
                    acquired_at=now,
                    last_seen_at=now,
                )
                return True
            if existing.session_id != session_id:
                return False
            self._leases[device_id] = McpDeviceLease(
                device_id=device_id,
                session_id=session_id,
                acquired_at=existing.acquired_at,
                last_seen_at=self._now(),
            )
            return True

    def release(self, session_id: str) -> list[str]:
        with self._lock:
            freed = [
                device_id
                for device_id, lease in self._leases.items()
                if lease.session_id == session_id
            ]
            for device_id in freed:
                del self._leases[device_id]
            return freed

    def release_device(self, device_id: str, session_id: str) -> bool:
        with self._lock:
            existing = self._leases.get(device_id)
            if existing is None or existing.session_id != session_id:
                return False
            del self._leases[device_id]
            return True

    def busy_device_ids(self) -> list[str]:
        with self._lock:
            self._sweep_locked()
            return sorted(self._leases)

    def snapshot(self) -> list[McpDeviceLease]:
        with self._lock:
            self._sweep_locked()
            return sorted(self._leases.values(), key=lambda lease: lease.device_id)

    def wait_until_usable(
        self,
        device_id: str,
        session_id: str,
        *,
        timeout: float,
        poll_interval: float = 1.0,
        cloud_busy_check: Callable[[], bool] | None = None,
    ) -> bool:
        """Block until ``device_id`` is acquirable by ``session_id`` or timeout.

        Reserved capability. MVP callers use try-acquire (``acquire`` -> False
        means busy). This method exists for future wiring where the cloud
        assignment path or an explicit MCP tool may opt to wait.
        """
        deadline = time.monotonic() + timeout
        while True:
            cloud_busy = cloud_busy_check() if cloud_busy_check else False
            if not cloud_busy:
                if self.acquire(device_id, session_id):
                    return True
            if time.monotonic() >= deadline:
                return False
            remaining = deadline - time.monotonic()
            time.sleep(max(0.0, min(poll_interval, remaining)))

    def _sweep_locked(self) -> None:
        """Caller holds ``self._lock``. Drops leases past their TTL."""
        cutoff = self._now()
        expired = [
            device_id
            for device_id, lease in self._leases.items()
            if (cutoff - lease.last_seen_at).total_seconds() > self._ttl
        ]
        for device_id in expired:
            del self._leases[device_id]