feat(host-agent): add McpBusyTracker for per-device session locks

2026-07-21 13:58:37 +08:00
parent c7faee8da3
commit 61c923b92b
2 changed files with 307 additions and 0 deletions
@@ -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]