feat(host-agent): add McpBusyTracker for per-device session locks
This commit is contained in:
@@ -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]
|
||||
@@ -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() == []
|
||||
Reference in New Issue
Block a user