feat: add heartbeat-driven bell clearing for remote sessions in poll cycle
- Add module-level _federation_client reference for use in background poll task - Assign _federation_client in lifespan startup; clear on shutdown - Add step 12 to _run_poll_cycle: iterate devices viewing a remote session in fullscreen with recent interaction (<60s), fire POST bell/clear to the active remote instance with Bearer auth (fire-and-forget, errors logged) - Add test_poll_cycle_fires_federation_bell_clear_for_remote_session test Closes task-3 of federation-state-propagation-plan
This commit is contained in:
Brian Krabach
+45
-3
@@ -81,6 +81,7 @@ _log = logging.getLogger(__name__)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_poll_task: asyncio.Task | None = None
|
||||
_federation_client: httpx.AsyncClient | None = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -131,6 +132,46 @@ async def _run_poll_cycle() -> None:
|
||||
# 9. Apply bell clear rule (acknowledge bells when device is watching fullscreen)
|
||||
apply_bell_clear_rule(state)
|
||||
|
||||
# 12. Fire bell/clear to the active remote for any device viewing a remote
|
||||
# session in fullscreen with recent interaction. Fire-and-forget: errors
|
||||
# are logged and do not abort the rest of the poll cycle.
|
||||
if _federation_client is not None:
|
||||
active_remote_id = state.get("active_remote_id")
|
||||
if active_remote_id is not None:
|
||||
settings = load_settings()
|
||||
remote_instances = settings.get("remote_instances", [])
|
||||
if (
|
||||
isinstance(active_remote_id, int)
|
||||
and 0 <= active_remote_id < len(remote_instances)
|
||||
):
|
||||
remote = remote_instances[active_remote_id]
|
||||
remote_url: str = remote.get("url", "").rstrip("/")
|
||||
remote_key: str = remote.get("key", "")
|
||||
now = time.time()
|
||||
for device in state.get("devices", {}).values():
|
||||
viewing_session = device.get("viewing_session")
|
||||
view_mode = device.get("view_mode")
|
||||
last_interaction_at = device.get("last_interaction_at", 0)
|
||||
if (
|
||||
viewing_session
|
||||
and view_mode == "fullscreen"
|
||||
and (now - last_interaction_at) < 60
|
||||
):
|
||||
bell_clear_url = (
|
||||
f"{remote_url}/api/sessions/{viewing_session}/bell/clear"
|
||||
)
|
||||
try:
|
||||
await _federation_client.post(
|
||||
bell_clear_url,
|
||||
headers={"Authorization": f"Bearer {remote_key}"} if remote_key else {},
|
||||
)
|
||||
except Exception:
|
||||
_log.warning(
|
||||
"federation bell clear failed for %s: %s",
|
||||
viewing_session,
|
||||
bell_clear_url,
|
||||
)
|
||||
|
||||
# 10. Prune devices that haven't sent a heartbeat recently
|
||||
prune_devices(state)
|
||||
|
||||
@@ -161,6 +202,7 @@ async def _poll_loop() -> None:
|
||||
@contextlib.asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
global _poll_task
|
||||
global _federation_client
|
||||
|
||||
# Startup: kill any orphaned ttyd from a previous muxplex run, then
|
||||
# start the background poll loop.
|
||||
@@ -187,6 +229,7 @@ async def lifespan(app: FastAPI):
|
||||
# Bearer token auth handles authorization. Users who need cert verification
|
||||
# should use mkcert (CA-trusted) or Tailscale (LE-trusted) certs.
|
||||
)
|
||||
_federation_client = app.state.federation_client
|
||||
|
||||
yield
|
||||
|
||||
@@ -194,6 +237,7 @@ async def lifespan(app: FastAPI):
|
||||
client = getattr(app.state, "federation_client", None)
|
||||
if client is not None:
|
||||
await client.aclose()
|
||||
_federation_client = None
|
||||
except Exception:
|
||||
_log.exception("federation_client aclose error")
|
||||
finally:
|
||||
@@ -871,9 +915,7 @@ async def federation_terminal_ws_proxy(websocket: WebSocket, remote_id: int) ->
|
||||
async with websockets.connect(
|
||||
ws_url,
|
||||
subprotocols=[Subprotocol("tty")],
|
||||
additional_headers={"Authorization": f"Bearer {remote_key}"}
|
||||
if remote_key
|
||||
else {},
|
||||
additional_headers={"Authorization": f"Bearer {remote_key}"} if remote_key else {},
|
||||
ssl=ssl_context,
|
||||
) as remote_ws:
|
||||
|
||||
|
||||
@@ -2691,6 +2691,117 @@ def test_federation_create_session_returns_502_when_remote_returns_error(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Poll cycle federation bell clear (task-3)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_poll_cycle_fires_federation_bell_clear_for_remote_session(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""_run_poll_cycle() fires POST bell/clear to remote when a device is viewing a remote session.
|
||||
|
||||
Sets up state with active_remote_id=0, one device viewing 'build' in fullscreen
|
||||
with a recent interaction timestamp, mocks the module-level _federation_client,
|
||||
runs one poll cycle, and verifies mock_client.post was called with the correct
|
||||
URL and Bearer auth header.
|
||||
"""
|
||||
import json
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import muxplex.main as main_mod
|
||||
import muxplex.settings as settings_mod
|
||||
from muxplex.state import save_state
|
||||
|
||||
# Set up settings with one remote instance at index 0
|
||||
settings_path = tmp_path / "settings.json"
|
||||
settings_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"remote_instances": [
|
||||
{
|
||||
"url": "http://remote-host:8088",
|
||||
"key": "test-key",
|
||||
"name": "remote-host",
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(settings_mod, "SETTINGS_PATH", settings_path)
|
||||
|
||||
# Set up state with active_remote_id=0, one device viewing 'build' in fullscreen
|
||||
state = {
|
||||
"active_session": None,
|
||||
"active_remote_id": 0,
|
||||
"session_order": ["build"],
|
||||
"sessions": {
|
||||
"build": {
|
||||
"bell": {"last_fired_at": None, "seen_at": None, "unseen_count": 0}
|
||||
}
|
||||
},
|
||||
"devices": {
|
||||
"dev-1": {
|
||||
"label": "My Device",
|
||||
"viewing_session": "build",
|
||||
"view_mode": "fullscreen",
|
||||
"last_interaction_at": time.time(),
|
||||
"last_heartbeat_at": time.time(),
|
||||
}
|
||||
},
|
||||
}
|
||||
save_state(state)
|
||||
|
||||
# Mock all poll-cycle dependencies so the cycle completes without real tmux
|
||||
async def mock_enumerate():
|
||||
return ["build"]
|
||||
|
||||
async def mock_snapshot_all(names):
|
||||
return {"build": "pane text"}
|
||||
|
||||
async def mock_process_bell_flags(names, state):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr("muxplex.main.enumerate_sessions", mock_enumerate)
|
||||
monkeypatch.setattr("muxplex.main.snapshot_all", mock_snapshot_all)
|
||||
monkeypatch.setattr(
|
||||
"muxplex.main.update_session_cache", lambda names, snapshots: None
|
||||
)
|
||||
monkeypatch.setattr("muxplex.main.apply_bell_clear_rule", lambda state: None)
|
||||
monkeypatch.setattr("muxplex.main.prune_devices", lambda state: None)
|
||||
monkeypatch.setattr("muxplex.main.process_bell_flags", mock_process_bell_flags)
|
||||
|
||||
# Capture POST calls from the mocked federation client
|
||||
post_calls: list[dict] = []
|
||||
|
||||
async def mock_post(url, **kwargs):
|
||||
post_calls.append({"url": url, "kwargs": kwargs})
|
||||
resp = MagicMock()
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.post = mock_post
|
||||
monkeypatch.setattr(main_mod, "_federation_client", mock_client)
|
||||
|
||||
# Run one poll cycle
|
||||
await main_mod._run_poll_cycle()
|
||||
|
||||
# Verify mock_client.post was called exactly once with the correct URL and auth
|
||||
assert len(post_calls) == 1, (
|
||||
f"Expected exactly 1 POST call to remote bell/clear, got {len(post_calls)}: {post_calls}"
|
||||
)
|
||||
call = post_calls[0]
|
||||
assert "/api/sessions/build/bell/clear" in call["url"], (
|
||||
f"Expected URL to contain '/api/sessions/build/bell/clear', got: {call['url']}"
|
||||
)
|
||||
headers = call["kwargs"].get("headers", {})
|
||||
assert headers.get("Authorization") == "Bearer test-key", (
|
||||
f"Expected 'Authorization: Bearer test-key' header, got: {headers}"
|
||||
)
|
||||
|
||||
|
||||
def test_federation_auth_headers_guard_empty_key():
|
||||
"""Every federation Authorization header construction must guard against empty key.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user