diff --git a/apps/device-host-agent/host_agent/mcp_lock.py b/apps/device-host-agent/host_agent/mcp_lock.py new file mode 100644 index 0000000..a3bd5ab --- /dev/null +++ b/apps/device-host-agent/host_agent/mcp_lock.py @@ -0,0 +1,157 @@ +"""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 = 60.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] \ No newline at end of file diff --git a/apps/device-host-agent/tests/test_mcp_lock.py b/apps/device-host-agent/tests/test_mcp_lock.py new file mode 100644 index 0000000..c02f9e9 --- /dev/null +++ b/apps/device-host-agent/tests/test_mcp_lock.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import threading +from datetime import UTC, datetime + +from host_agent.mcp_lock import McpBusyTracker + + +def _tracker_with_now() -> tuple[McpBusyTracker, list[datetime]]: + times: list[datetime] = [] + + def now() -> datetime: + return times[-1] if times else datetime(2026, 1, 1, tzinfo=UTC) + + tracker = McpBusyTracker(ttl_seconds=60.0, now=now) + return tracker, times + + +def test_acquire_succeeds_on_empty() -> None: + tracker, _ = _tracker_with_now() + assert tracker.acquire("phone-1", "sess-a") is True + assert "phone-1" in tracker.busy_device_ids() + + +def test_acquire_fails_when_held_by_other_session() -> None: + tracker, _ = _tracker_with_now() + assert tracker.acquire("phone-1", "sess-a") is True + assert tracker.acquire("phone-1", "sess-b") is False + + +def test_acquire_is_idempotent_for_same_session() -> None: + tracker, _ = _tracker_with_now() + assert tracker.acquire("phone-1", "sess-a") is True + # Same session re-acquiring is allowed (acts as renew). + assert tracker.acquire("phone-1", "sess-a") is True + + +def test_renew_refreshes_last_seen() -> None: + tracker, times = _tracker_with_now() + times.append(datetime(2026, 1, 1, 12, 0, tzinfo=UTC)) + tracker.acquire("phone-1", "sess-a") + initial = tracker.snapshot()[0] + times.append(datetime(2026, 1, 1, 12, 0, 30, tzinfo=UTC)) + assert tracker.renew("phone-1", "sess-a") is True + refreshed = tracker.snapshot()[0] + assert refreshed.last_seen_at > initial.last_seen_at + + +def test_renew_fails_when_held_by_other() -> None: + tracker, _ = _tracker_with_now() + tracker.acquire("phone-1", "sess-a") + assert tracker.renew("phone-1", "sess-b") is False + + +def test_release_returns_freed_device_ids() -> None: + tracker, _ = _tracker_with_now() + tracker.acquire("phone-1", "sess-a") + tracker.acquire("phone-2", "sess-a") + freed = tracker.release("sess-a") + assert sorted(freed) == ["phone-1", "phone-2"] + assert tracker.busy_device_ids() == [] + + +def test_release_only_frees_caller_session() -> None: + tracker, _ = _tracker_with_now() + tracker.acquire("phone-1", "sess-a") + tracker.acquire("phone-1", "sess-b") # fails + freed = tracker.release("sess-b") + assert freed == [] + assert "phone-1" in tracker.busy_device_ids() + + +def test_ttl_sweeps_expired_leases() -> None: + tracker, times = _tracker_with_now() + times.append(datetime(2026, 1, 1, 12, 0, tzinfo=UTC)) + tracker.acquire("phone-1", "sess-a") + # Advance past TTL without renew. + times.append(datetime(2026, 1, 1, 12, 1, 1, tzinfo=UTC)) # 61s later + assert tracker.busy_device_ids() == [] + + +def test_renew_after_ttl_tolerates_same_session() -> None: + """Scene 10: lease expired but session_id matches -> re-acquire.""" + tracker, times = _tracker_with_now() + times.append(datetime(2026, 1, 1, 12, 0, tzinfo=UTC)) + tracker.acquire("phone-1", "sess-a") + times.append(datetime(2026, 1, 1, 12, 1, 1, tzinfo=UTC)) # expired + # renew from the same session should succeed (re-acquire). + assert tracker.renew("phone-1", "sess-a") is True + assert "phone-1" in tracker.busy_device_ids() + + +def test_snapshot_matches_busy_device_ids() -> None: + tracker, _ = _tracker_with_now() + tracker.acquire("phone-1", "sess-a") + tracker.acquire("phone-2", "sess-a") + snap = tracker.snapshot() + assert {lease.device_id for lease in snap} == set(tracker.busy_device_ids()) + + +def test_wait_until_usable_succeeds_when_free() -> None: + tracker, _ = _tracker_with_now() + ok = tracker.wait_until_usable( + "phone-1", "sess-a", timeout=1.0, poll_interval=0.01 + ) + assert ok is True + assert "phone-1" in tracker.busy_device_ids() + + +def test_wait_until_usable_returns_false_on_timeout() -> None: + tracker, _ = _tracker_with_now() + tracker.acquire("phone-1", "sess-a") + ok = tracker.wait_until_usable( + "phone-1", "sess-b", timeout=0.1, poll_interval=0.02 + ) + assert ok is False + + +def test_wait_until_usable_blocks_then_succeeds_when_released() -> None: + tracker, _ = _tracker_with_now() + tracker.acquire("phone-1", "sess-a") + + def releaser() -> None: + import time + + time.sleep(0.05) + tracker.release("sess-a") + + t = threading.Thread(target=releaser) + t.start() + try: + ok = tracker.wait_until_usable( + "phone-1", "sess-b", timeout=2.0, poll_interval=0.02 + ) + assert ok is True + finally: + t.join() + + +def test_wait_until_usable_blocks_then_fails_when_cloud_remains_busy() -> None: + tracker, _ = _tracker_with_now() + ok = tracker.wait_until_usable( + "phone-1", + "sess-a", + timeout=0.1, + poll_interval=0.02, + cloud_busy_check=lambda: True, + ) + assert ok is False + assert tracker.busy_device_ids() == [] \ No newline at end of file