From 9472772ae9069d3338433bbac830b61e4cbc64ba Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 29 Jan 2026 15:16:00 -0500 Subject: [PATCH 1/5] feat: support multiple fxa scopes --- src/mlpa/core/config.py | 3 +++ src/mlpa/core/routers/fxa/fxa.py | 41 ++++++++++++++++++++++---------- 2 files changed, 32 insertions(+), 12 deletions(-) diff --git a/src/mlpa/core/config.py b/src/mlpa/core/config.py index a72f0c2..f154c0f 100644 --- a/src/mlpa/core/config.py +++ b/src/mlpa/core/config.py @@ -96,6 +96,9 @@ def valid_service_types(self) -> list[str]: # FxA CLIENT_ID: str = "default-client-id" CLIENT_SECRET: str = "default-client-secret" + ADDITIONAL_FXA_SCOPE_1: str | None = None + ADDITIONAL_FXA_SCOPE_2: str | None = None + ADDITIONAL_FXA_SCOPE_3: str | None = None # PostgreSQL LITELLM_DB_NAME: str = "litellm" diff --git a/src/mlpa/core/routers/fxa/fxa.py b/src/mlpa/core/routers/fxa/fxa.py index 11af14e..50bfda2 100644 --- a/src/mlpa/core/routers/fxa/fxa.py +++ b/src/mlpa/core/routers/fxa/fxa.py @@ -2,8 +2,8 @@ from typing import Annotated from fastapi import APIRouter, Header, HTTPException -from fastapi.concurrency import run_in_threadpool +from mlpa.core.config import env from mlpa.core.logger import logger from mlpa.core.prometheus_metrics import PrometheusResult, metrics from mlpa.core.utils import get_fxa_client @@ -11,21 +11,38 @@ router = APIRouter() client = get_fxa_client() +FXA_DEFAULT_SCOPE = "profile" +scopes = filter( + None, + [ + FXA_DEFAULT_SCOPE, + env.ADDITIONAL_FXA_SCOPE_1, + env.ADDITIONAL_FXA_SCOPE_2, + env.ADDITIONAL_FXA_SCOPE_3, + ], +) async def fxa_auth(authorization: Annotated[str | None, Header()]): start_time = time.perf_counter() token = authorization.removeprefix("Bearer ").split()[0] result = PrometheusResult.ERROR - - try: - profile = await run_in_threadpool(client.verify_token, token, scope="profile") - result = PrometheusResult.SUCCESS - return profile - except Exception as e: - logger.error(f"FxA auth error: {e}") + success = False + errors = [] + for scope in scopes: + try: + profile = client.verify_token(token, scope=scope) + result = PrometheusResult.SUCCESS + success = True + break + except Exception as e: + errors.append(e) + continue + finally: + metrics.validate_fxa_latency.labels(result=result).observe( + time.time() - start_time + ) + if not success: + logger.error(f"FxA auth error: {errors}") raise HTTPException(status_code=401, detail="Invalid FxA auth") - finally: - metrics.validate_fxa_latency.labels(result=result).observe( - time.perf_counter() - start_time - ) + return profile From 0a6ad4e773fa1961c1b2aa6aaaf716c98d381754 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 29 Jan 2026 15:30:36 -0500 Subject: [PATCH 2/5] add run_in_threadpool, refactor logic flow --- src/mlpa/core/routers/fxa/fxa.py | 40 ++++++++++++++++---------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/src/mlpa/core/routers/fxa/fxa.py b/src/mlpa/core/routers/fxa/fxa.py index 50bfda2..fc0f7ac 100644 --- a/src/mlpa/core/routers/fxa/fxa.py +++ b/src/mlpa/core/routers/fxa/fxa.py @@ -2,6 +2,7 @@ from typing import Annotated from fastapi import APIRouter, Header, HTTPException +from fastapi.concurrency import run_in_threadpool from mlpa.core.config import env from mlpa.core.logger import logger @@ -12,14 +13,15 @@ client = get_fxa_client() FXA_DEFAULT_SCOPE = "profile" -scopes = filter( - None, - [ +FXA_SCOPES = tuple( + scope + for scope in ( FXA_DEFAULT_SCOPE, env.ADDITIONAL_FXA_SCOPE_1, env.ADDITIONAL_FXA_SCOPE_2, env.ADDITIONAL_FXA_SCOPE_3, - ], + ) + if scope ) @@ -27,22 +29,20 @@ async def fxa_auth(authorization: Annotated[str | None, Header()]): start_time = time.perf_counter() token = authorization.removeprefix("Bearer ").split()[0] result = PrometheusResult.ERROR - success = False errors = [] - for scope in scopes: - try: - profile = client.verify_token(token, scope=scope) - result = PrometheusResult.SUCCESS - success = True - break - except Exception as e: - errors.append(e) - continue - finally: - metrics.validate_fxa_latency.labels(result=result).observe( - time.time() - start_time - ) - if not success: + try: + for scope in FXA_SCOPES: + try: + profile = await run_in_threadpool( + client.verify_token, token, scope=scope + ) + result = PrometheusResult.SUCCESS + return profile + except Exception as e: + errors.append(e) logger.error(f"FxA auth error: {errors}") raise HTTPException(status_code=401, detail="Invalid FxA auth") - return profile + finally: + metrics.validate_fxa_latency.labels(result=result).observe( + time.perf_counter() - start_time + ) From 71b97162eb33c7bc175901d6bafa0875b46bbe12 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Fri, 30 Jan 2026 08:37:28 -0500 Subject: [PATCH 3/5] update default scope to profile:uid --- src/mlpa/core/routers/fxa/fxa.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mlpa/core/routers/fxa/fxa.py b/src/mlpa/core/routers/fxa/fxa.py index fc0f7ac..f952fae 100644 --- a/src/mlpa/core/routers/fxa/fxa.py +++ b/src/mlpa/core/routers/fxa/fxa.py @@ -12,7 +12,7 @@ router = APIRouter() client = get_fxa_client() -FXA_DEFAULT_SCOPE = "profile" +FXA_DEFAULT_SCOPE = "profile:uid" FXA_SCOPES = tuple( scope for scope in ( From 84a740a4fdbeeaa8bff94767b1b1f9084f14baad Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Fri, 30 Jan 2026 09:10:47 -0500 Subject: [PATCH 4/5] async multi scope fxa auth --- src/mlpa/core/routers/fxa/fxa.py | 29 ++++++++++++------ src/tests/unit/test_fxa_auth.py | 51 ++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 9 deletions(-) create mode 100644 src/tests/unit/test_fxa_auth.py diff --git a/src/mlpa/core/routers/fxa/fxa.py b/src/mlpa/core/routers/fxa/fxa.py index f952fae..2608c71 100644 --- a/src/mlpa/core/routers/fxa/fxa.py +++ b/src/mlpa/core/routers/fxa/fxa.py @@ -1,3 +1,4 @@ +import asyncio import time from typing import Annotated @@ -31,15 +32,25 @@ async def fxa_auth(authorization: Annotated[str | None, Header()]): result = PrometheusResult.ERROR errors = [] try: - for scope in FXA_SCOPES: - try: - profile = await run_in_threadpool( - client.verify_token, token, scope=scope - ) - result = PrometheusResult.SUCCESS - return profile - except Exception as e: - errors.append(e) + tasks = [ + asyncio.create_task( + run_in_threadpool(client.verify_token, token, scope=scope) + ) + for scope in FXA_SCOPES + ] + try: + for task in asyncio.as_completed(tasks): + try: + profile = await task + result = PrometheusResult.SUCCESS + return profile + except Exception as e: + errors.append(e) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) logger.error(f"FxA auth error: {errors}") raise HTTPException(status_code=401, detail="Invalid FxA auth") finally: diff --git a/src/tests/unit/test_fxa_auth.py b/src/tests/unit/test_fxa_auth.py new file mode 100644 index 0000000..d9d65ba --- /dev/null +++ b/src/tests/unit/test_fxa_auth.py @@ -0,0 +1,51 @@ +import asyncio + +import pytest +from fastapi import HTTPException + +from mlpa.core.prometheus_metrics import PrometheusResult +from mlpa.core.routers.fxa import fxa as fxa_module + + +async def test_fxa_auth_returns_first_successful_scope(mocker): + scopes = ("profile", "scope-a", "scope-b") + mocker.patch.object(fxa_module, "FXA_SCOPES", scopes) + + async def fake_run_in_threadpool(_fn, _token, *, scope): + if scope == "scope-b": + await asyncio.sleep(0.01) + return {"user": "ok"} + await asyncio.sleep(0.02) + raise Exception(f"invalid-{scope}") + + mocker.patch.object(fxa_module, "run_in_threadpool", new=fake_run_in_threadpool) + mock_metrics = mocker.patch.object(fxa_module, "metrics") + + profile = await fxa_module.fxa_auth("Bearer test-token") + + assert profile == {"user": "ok"} + mock_metrics.validate_fxa_latency.labels.assert_called_once_with( + result=PrometheusResult.SUCCESS + ) + mock_metrics.validate_fxa_latency.labels().observe.assert_called_once() + + +async def test_fxa_auth_raises_when_all_scopes_fail(mocker): + scopes = ("profile", "scope-a") + mocker.patch.object(fxa_module, "FXA_SCOPES", scopes) + + async def fake_run_in_threadpool(_fn, _token, *, scope): + await asyncio.sleep(0.01) + raise Exception(f"invalid-{scope}") + + mocker.patch.object(fxa_module, "run_in_threadpool", new=fake_run_in_threadpool) + mock_metrics = mocker.patch.object(fxa_module, "metrics") + + with pytest.raises(HTTPException) as exc_info: + await fxa_module.fxa_auth("Bearer test-token") + + assert exc_info.value.status_code == 401 + mock_metrics.validate_fxa_latency.labels.assert_called_once_with( + result=PrometheusResult.ERROR + ) + mock_metrics.validate_fxa_latency.labels().observe.assert_called_once() From 153b85d81691a5c0cb476680f4bd50db2fa02119 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Fri, 30 Jan 2026 11:28:41 -0500 Subject: [PATCH 5/5] migrate profile -> profile:uid scope values --- CONTRIBUTING.md | 4 ++-- src/tests/consts.py | 2 +- src/tests/integration/test_mock_router_integration.py | 6 +++--- src/tests/mocks.py | 2 +- src/tests/unit/test_fxa_auth.py | 4 ++-- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3ef9625..972ffd3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -139,7 +139,7 @@ from fxa.tools.bearer import get_bearer_token fxa_token: str = get_bearer_token( your_mozilla_account_email, your_mozilla_account_password, - scopes=["profile"], + scopes=["profile:uid"], client_id="5882386c6d801776" # a common client_id for the dev environment, account_server_url="https://api.accounts.firefox.com", oauth_server_url="https://oauth.accounts.firefox.com", @@ -148,7 +148,7 @@ fxa_token: str = get_bearer_token( ## How to update static docs/index.html from redoc -Make sure you have Node installed +Ensure Node is installed 1. `make install` 2. `mlpa` diff --git a/src/tests/consts.py b/src/tests/consts.py index 661d81e..3d165e4 100644 --- a/src/tests/consts.py +++ b/src/tests/consts.py @@ -65,7 +65,7 @@ MOCK_FXA_USER_DATA = { "user": TEST_USER_ID, "client_id": "test-client-id", - "scope": ["profile"], + "scope": ["profile:uid"], "generation": 1, "profile_changed_at": 1234567890, } diff --git a/src/tests/integration/test_mock_router_integration.py b/src/tests/integration/test_mock_router_integration.py index 2b99b59..9732ef3 100644 --- a/src/tests/integration/test_mock_router_integration.py +++ b/src/tests/integration/test_mock_router_integration.py @@ -232,7 +232,7 @@ def test_mock_chat_completions_no_auth_missing_user_in_token( mock_fxa_client._verify_jwt_token.return_value = { "client_id": "test-client-id", - "scope": ["profile"], + "scope": ["profile:uid"], } response = mocked_client_integration.post( @@ -259,7 +259,7 @@ def test_mock_chat_completions_no_auth_blocked_user( mock_fxa_client._verify_jwt_token.return_value = { "user": "blocked-user-id", "client_id": "test-client-id", - "scope": ["profile"], + "scope": ["profile:uid"], } with patch( @@ -295,7 +295,7 @@ def test_mock_chat_completions_latency_simulation(self, mocked_client_integratio mock_fxa_client._verify_jwt_token.return_value = { "user": TEST_USER_ID, "client_id": "test-client-id", - "scope": ["profile"], + "scope": ["profile:uid"], } with patch.dict("os.environ", {"MOCK_LATENCY_MS": "100"}): diff --git a/src/tests/mocks.py b/src/tests/mocks.py index 9ab731c..1230b2f 100644 --- a/src/tests/mocks.py +++ b/src/tests/mocks.py @@ -195,7 +195,7 @@ def __init__(self, client_id: str, client_secret: str, fxa_url: str): self.client_secret = client_secret self.fxa_url = fxa_url - def verify_token(self, token: str, scope: str = "profile"): + def verify_token(self, token: str, scope: str = "profile:uid"): if token == TEST_FXA_TOKEN: return {"user": TEST_USER_ID} raise Exception("Invalid token") diff --git a/src/tests/unit/test_fxa_auth.py b/src/tests/unit/test_fxa_auth.py index d9d65ba..790fd65 100644 --- a/src/tests/unit/test_fxa_auth.py +++ b/src/tests/unit/test_fxa_auth.py @@ -8,7 +8,7 @@ async def test_fxa_auth_returns_first_successful_scope(mocker): - scopes = ("profile", "scope-a", "scope-b") + scopes = ("profile:uid", "scope-a", "scope-b") mocker.patch.object(fxa_module, "FXA_SCOPES", scopes) async def fake_run_in_threadpool(_fn, _token, *, scope): @@ -31,7 +31,7 @@ async def fake_run_in_threadpool(_fn, _token, *, scope): async def test_fxa_auth_raises_when_all_scopes_fail(mocker): - scopes = ("profile", "scope-a") + scopes = ("profile:uid", "scope-a") mocker.patch.object(fxa_module, "FXA_SCOPES", scopes) async def fake_run_in_threadpool(_fn, _token, *, scope):