From f5ae0df9f587ad065be6b57c618efbd667cff871 Mon Sep 17 00:00:00 2001 From: Brian Krabach Date: Wed, 8 Apr 2026 18:50:17 -0700 Subject: [PATCH] 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 --- muxplex/main.py | 48 +++++++++++++++-- muxplex/tests/test_api.py | 111 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 156 insertions(+), 3 deletions(-) diff --git a/muxplex/main.py b/muxplex/main.py index d20a908..4b5a794 100644 --- a/muxplex/main.py +++ b/muxplex/main.py @@ -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: diff --git a/muxplex/tests/test_api.py b/muxplex/tests/test_api.py index eae1575..724d9d8 100644 --- a/muxplex/tests/test_api.py +++ b/muxplex/tests/test_api.py @@ -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.