From 21b669b5476b42001cde00adcf61baa31fde71fd Mon Sep 17 00:00:00 2001
From: Brian Krabach <brkrabac@microsoft.com>
Date: Wed, 1 Apr 2026 10:33:32 -0700
Subject: [PATCH] feat: wire federation key into AuthMiddleware at startup

---
 muxplex/main.py           |  4 +++-
 muxplex/tests/test_api.py | 48 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/muxplex/main.py b/muxplex/main.py
index a5b9c7c..49f177a 100644
--- a/muxplex/main.py
+++ b/muxplex/main.py
@@ -60,7 +60,7 @@ from muxplex.state import (
     save_state,
     state_lock,
 )
-from muxplex.settings import load_settings, patch_settings
+from muxplex.settings import load_federation_key, load_settings, patch_settings
 from muxplex.ttyd import kill_orphan_ttyd, kill_ttyd, spawn_ttyd, TTYD_PORT
 
 # ---------------------------------------------------------------------------
@@ -244,6 +244,7 @@ def _resolve_auth() -> tuple[str, str]:
 _auth_mode, _auth_password = _resolve_auth()
 _auth_secret = load_or_create_secret()
 _auth_ttl = int(os.environ.get("MUXPLEX_SESSION_TTL", "604800"))
+_federation_key = load_federation_key()
 
 app.add_middleware(
     AuthMiddleware,
@@ -251,6 +252,7 @@ app.add_middleware(
     secret=_auth_secret,
     ttl_seconds=_auth_ttl,
     password=_auth_password,
+    federation_key=_federation_key,
 )
 
 # CORS: allow_origins=["*"] with allow_credentials=True is intentional for
diff --git a/muxplex/tests/test_api.py b/muxplex/tests/test_api.py
index 46b898e..71a9d0b 100644
--- a/muxplex/tests/test_api.py
+++ b/muxplex/tests/test_api.py
@@ -1387,3 +1387,51 @@ def test_get_auth_token_returns_401_when_not_authenticated(monkeypatch):
         # No cookie set — endpoint must return 401 with application/json accept
         response = c.get("/api/auth/token", headers={"Accept": "application/json"})
     assert response.status_code == 401
+
+
+# ---------------------------------------------------------------------------
+# Federation Bearer token auth
+# ---------------------------------------------------------------------------
+
+
+def test_federation_bearer_auth_accepted(monkeypatch):
+    """Request with valid Bearer token gets 200 on /api/sessions when federation key is set.
+
+    Patches the federation key on the AuthMiddleware instance (since the key is
+    loaded once at module startup) and verifies that a Bearer-authenticated
+    request reaches /api/sessions with HTTP 200.
+
+    Before implementation: fails with ImportError — _federation_key not in main.py.
+    After implementation: _federation_key exists, middleware is found and patched,
+    Bearer request is accepted.
+    """
+    from muxplex.main import _federation_key  # ImportError before implementation
+    from muxplex.auth import AuthMiddleware
+    import muxplex.main as main_module
+
+    federation_key = "test-federation-key-abc123"
+    monkeypatch.setenv("MUXPLEX_PASSWORD", "test-password")
+
+    with TestClient(main_module.app) as c:
+        # Traverse the compiled middleware stack to find the AuthMiddleware instance
+        stack = main_module.app.middleware_stack
+        auth_mw = None
+        for _ in range(20):
+            if isinstance(stack, AuthMiddleware):
+                auth_mw = stack
+                break
+            stack = getattr(stack, "app", None)
+            if stack is None:
+                break
+
+        assert auth_mw is not None, "AuthMiddleware not found in middleware stack"
+        # Patch the federation key so Bearer token auth is enabled
+        auth_mw.federation_key = federation_key
+
+        # A request with a matching Bearer token must pass auth and get 200
+        response = c.get(
+            "/api/sessions",
+            headers={"Authorization": f"Bearer {federation_key}"},
+        )
+
+    assert response.status_code == 200