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:
+45
-3
@@ -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:
|
||||||
|
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user