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
2026-04-08 18:50:17 -07:00
parent 5eb727fe86
commit f5ae0df9f5
2 changed files with 156 additions and 3 deletions
+45 -3
View File
@@ -81,6 +81,7 @@ _log = logging.getLogger(__name__)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
_poll_task: asyncio.Task | None = None _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) # 9. Apply bell clear rule (acknowledge bells when device is watching fullscreen)
apply_bell_clear_rule(state) 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 # 10. Prune devices that haven't sent a heartbeat recently
prune_devices(state) prune_devices(state)
@@ -161,6 +202,7 @@ async def _poll_loop() -> None:
@contextlib.asynccontextmanager @contextlib.asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
global _poll_task global _poll_task
global _federation_client
# Startup: kill any orphaned ttyd from a previous muxplex run, then # Startup: kill any orphaned ttyd from a previous muxplex run, then
# start the background poll loop. # start the background poll loop.
@@ -187,6 +229,7 @@ async def lifespan(app: FastAPI):
# Bearer token auth handles authorization. Users who need cert verification # Bearer token auth handles authorization. Users who need cert verification
# should use mkcert (CA-trusted) or Tailscale (LE-trusted) certs. # should use mkcert (CA-trusted) or Tailscale (LE-trusted) certs.
) )
_federation_client = app.state.federation_client
yield yield
@@ -194,6 +237,7 @@ async def lifespan(app: FastAPI):
client = getattr(app.state, "federation_client", None) client = getattr(app.state, "federation_client", None)
if client is not None: if client is not None:
await client.aclose() await client.aclose()
_federation_client = None
except Exception: except Exception:
_log.exception("federation_client aclose error") _log.exception("federation_client aclose error")
finally: finally:
@@ -871,9 +915,7 @@ async def federation_terminal_ws_proxy(websocket: WebSocket, remote_id: int) ->
async with websockets.connect( async with websockets.connect(
ws_url, ws_url,
subprotocols=[Subprotocol("tty")], subprotocols=[Subprotocol("tty")],
additional_headers={"Authorization": f"Bearer {remote_key}"} additional_headers={"Authorization": f"Bearer {remote_key}"} if remote_key else {},
if remote_key
else {},
ssl=ssl_context, ssl=ssl_context,
) as remote_ws: ) as remote_ws:
+111
View File
@@ -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(): def test_federation_auth_headers_guard_empty_key():
"""Every federation Authorization header construction must guard against empty key. """Every federation Authorization header construction must guard against empty key.