From 1f7e4ff34ddba85197935b38e51af897f144f4e9 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Tue, 16 Jun 2026 09:09:21 -0400 Subject: [PATCH 1/4] feat: include litellm version in /readiness --- src/mlpa/core/config.py | 1 + src/mlpa/core/routers/health/health.py | 17 +++++++-- src/tests/conftest.py | 6 +++ src/tests/integration/test_fxa.py | 22 +++++------ src/tests/integration/test_health.py | 38 +++++++++---------- .../integration/test_security_headers.py | 8 +++- src/tests/integration/test_user_signup_cap.py | 12 +++--- 7 files changed, 61 insertions(+), 43 deletions(-) diff --git a/src/mlpa/core/config.py b/src/mlpa/core/config.py index b3de32e..4c26d9d 100644 --- a/src/mlpa/core/config.py +++ b/src/mlpa/core/config.py @@ -324,6 +324,7 @@ def __init__(self): env = Env() LITELLM_READINESS_URL = f"{env.LITELLM_API_BASE}/health/readiness" +LITELLM_INFO_URL = f"{env.LITELLM_API_BASE}/public/model_hub/info" LITELLM_COMPLETIONS_URL = f"{env.LITELLM_API_BASE}/v1/chat/completions" LITELLM_SEARCH_URL = f"{env.LITELLM_API_BASE}/v1/search" LITELLM_MASTER_AUTH_HEADERS = { diff --git a/src/mlpa/core/routers/health/health.py b/src/mlpa/core/routers/health/health.py index 4ef9b21..e31747d 100644 --- a/src/mlpa/core/routers/health/health.py +++ b/src/mlpa/core/routers/health/health.py @@ -2,7 +2,11 @@ from fastapi import APIRouter -from mlpa.core.config import LITELLM_MASTER_AUTH_HEADERS, LITELLM_READINESS_URL +from mlpa.core.config import ( + LITELLM_INFO_URL, + LITELLM_MASTER_AUTH_HEADERS, + LITELLM_READINESS_URL, +) from mlpa.core.http_client import get_http_client from mlpa.core.pg_services.services import app_attest_pg, litellm_pg @@ -25,8 +29,10 @@ async def readiness_probe(): response = await client.get( LITELLM_READINESS_URL, headers=LITELLM_MASTER_AUTH_HEADERS, timeout=3 ) - data = response.json() - litellm_status = data + litellm_status = response.json() + + response = await client.get(LITELLM_INFO_URL, timeout=3) + litellm_info = response.json() return { "status": "connected", "mlpa_version": mlpa_version, @@ -34,5 +40,8 @@ async def readiness_probe(): "postgres": "connected" if pg_status else "offline", "app_attest": "connected" if app_attest_pg_status else "offline", }, - "litellm": litellm_status, + "litellm": { + "litellm_version": litellm_info["litellm_version"], + **litellm_status, + }, } diff --git a/src/tests/conftest.py b/src/tests/conftest.py index 00e3560..87a898a 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -17,6 +17,12 @@ def mock_request(): def _force_mlpa_debug_false(): monkeypatch = pytest.MonkeyPatch() monkeypatch.setenv("MLPA_DEBUG", "false") + monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_1", None) + monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_2", None) + monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_3", None) env.MLPA_DEBUG = False + env.ADDITIONAL_FXA_SCOPE_1 = None + env.ADDITIONAL_FXA_SCOPE_2 = None + env.ADDITIONAL_FXA_SCOPE_3 = None yield monkeypatch.undo() diff --git a/src/tests/integration/test_fxa.py b/src/tests/integration/test_fxa.py index f2bef01..b7793fe 100644 --- a/src/tests/integration/test_fxa.py +++ b/src/tests/integration/test_fxa.py @@ -30,7 +30,7 @@ def test_ai_missing_purpose_is_allowed_by_default(mocked_client_integration): "authorization": f"Bearer {TEST_FXA_TOKEN}", "service-type": "ai", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 200 assert response.json() == SUCCESSFUL_CHAT_RESPONSE @@ -45,7 +45,7 @@ def test_ai_invalid_purpose_returns_400(mocked_client_integration): "service-type": "ai", "purpose": "invalid-purpose", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 400 assert "purpose" in str(response.json().get("detail", "")).lower() @@ -59,7 +59,7 @@ def test_invalid_fxa_auth(mocked_client_integration): "service-type": "ai", "purpose": "chat", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 401 @@ -72,7 +72,7 @@ def test_successful_request_with_mocked_fxa_auth(mocked_client_integration): "service-type": "ai", "purpose": "chat", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code != 401 assert response.status_code != 400 @@ -93,7 +93,7 @@ def test_x_dev_authorization_success(mocked_client_integration): "purpose": "chat", "x-dev-authorization": DEV_TOKEN, }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 200 assert response.json() == SUCCESSFUL_CHAT_RESPONSE @@ -112,7 +112,7 @@ def test_x_dev_authorization_missing_fxa(mocked_client_integration): "purpose": "chat", "x-dev-authorization": DEV_TOKEN, }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 422 @@ -131,7 +131,7 @@ def test_x_dev_authorization_invalid_token(mocked_client_integration): "purpose": "chat", "x-dev-authorization": "wrong-token", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 401 @@ -149,7 +149,7 @@ def test_x_dev_authorization_token_not_configured(mocked_client_integration): "purpose": "chat", "x-dev-authorization": "some-token", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 422 @@ -163,7 +163,7 @@ def test_ai_dev_requires_x_dev_authorization(mocked_client_integration): "service-type": "ai-dev", "purpose": "chat", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 401 assert "x-dev-authorization required" in str(response.json().get("detail", "")) @@ -186,7 +186,7 @@ def test_x_dev_authorization_ignored_for_non_dev_service_type( "purpose": "chat", "x-dev-authorization": DEV_TOKEN, }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 200 assert response.json() == SUCCESSFUL_CHAT_RESPONSE @@ -206,7 +206,7 @@ def test_x_dev_authorization_token_not_configured_with_fxa(mocked_client_integra "purpose": "chat", "x-dev-authorization": "some-token", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert response.status_code == 401 assert "Invalid x-dev-authorization" in str(response.json().get("detail", "")) diff --git a/src/tests/integration/test_health.py b/src/tests/integration/test_health.py index dfdd114..2d5d845 100644 --- a/src/tests/integration/test_health.py +++ b/src/tests/integration/test_health.py @@ -1,3 +1,5 @@ +import importlib.metadata + from mlpa.core.config import env @@ -8,39 +10,33 @@ def test_health_liveness(mocked_client_integration, httpx_mock): def test_health_readiness(mocked_client_integration, httpx_mock): + mlpa_version = importlib.metadata.version("mlpa") httpx_mock.add_response( method="GET", url=f"{env.LITELLM_API_BASE}/health/readiness", status_code=200, json={ - "status": "connected", - "pg_server_dbs": {"postgres": "connected", "app_attest": "connected"}, - "litellm": { - "status": "connected", - "db": "connected", - "cache": None, - "litellm_version": "1.77.3", - "success_callbacks": [ - "sync_deployment_callback_on_success", - "_PROXY_VirtualKeyModelMaxBudgetLimiter", - "_ProxyDBLogger", - "_PROXY_MaxBudgetLimiter", - "_PROXY_MaxParallelRequestsHandler_v3", - "_PROXY_CacheControlCheck", - "_PROXY_LiteLLMManagedFiles", - "ServiceLogging", - ], - "use_aiohttp_transport": True, - "last_updated": "2025-10-10T00:00:00", - }, + "status": "healthy", + "db": "connected", }, ) + httpx_mock.add_response( + method="GET", + url=f"{env.LITELLM_API_BASE}/public/model_hub/info", + status_code=200, + json={"litellm_version": "1.84.4"}, + ) readiness_response = mocked_client_integration.get("/health/readiness") assert readiness_response.status_code == 200 assert readiness_response.json().get("status") == "connected" + assert readiness_response.json().get("mlpa_version") == mlpa_version assert readiness_response.json().get("pg_server_dbs") is not None - assert readiness_response.json().get("litellm") is not None + assert readiness_response.json().get("litellm") == { + "litellm_version": "1.84.4", + "status": "healthy", + "db": "connected", + } def test_metrics_endpoint(mocked_client_integration): diff --git a/src/tests/integration/test_security_headers.py b/src/tests/integration/test_security_headers.py index e11e673..92c6646 100644 --- a/src/tests/integration/test_security_headers.py +++ b/src/tests/integration/test_security_headers.py @@ -50,7 +50,13 @@ def test_security_headers_on_all_endpoints(mocked_client_integration, httpx_mock method="GET", url=f"{env.LITELLM_API_BASE}/health/readiness", status_code=200, - json={"status": "connected"}, + json={"status": "healthy", "db": "connected"}, + ) + httpx_mock.add_response( + method="GET", + url=f"{env.LITELLM_API_BASE}/public/model_hub/info", + status_code=200, + json={"litellm_version": "1.84.4"}, ) endpoints = [ diff --git a/src/tests/integration/test_user_signup_cap.py b/src/tests/integration/test_user_signup_cap.py index 2f55de1..908e50a 100644 --- a/src/tests/integration/test_user_signup_cap.py +++ b/src/tests/integration/test_user_signup_cap.py @@ -101,7 +101,7 @@ def verify_token_side_effect( "authorization": f"Bearer {TEST_FXA_TOKEN}", "service-type": "ai", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert resp1.status_code == 200 assert resp1.json() == SUCCESSFUL_CHAT_RESPONSE @@ -112,7 +112,7 @@ def verify_token_side_effect( "authorization": f"Bearer {TEST_FXA_TOKEN}", "service-type": "s2s", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert resp2.status_code == 200 assert resp2.json() == SUCCESSFUL_CHAT_RESPONSE @@ -123,7 +123,7 @@ def verify_token_side_effect( "authorization": f"Bearer {TEST_FXA_TOKEN}", "service-type": "ai", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert resp3.status_code == 403 assert resp3.json()["detail"]["error"] == 4 @@ -139,7 +139,7 @@ def verify_token_side_effect( "authorization": f"Bearer {TEST_FXA_TOKEN}", "service-type": "memories", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert resp4.status_code == 200 assert resp4.json() == SUCCESSFUL_CHAT_RESPONSE @@ -189,7 +189,7 @@ def verify_token_side_effect( "authorization": f"Bearer {TEST_FXA_TOKEN}", "service-type": "ai", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert resp1.status_code == 500 @@ -199,7 +199,7 @@ def verify_token_side_effect( "authorization": f"Bearer {TEST_FXA_TOKEN}", "service-type": "ai", }, - json=SAMPLE_REQUEST.dict(), + json=SAMPLE_REQUEST.model_dump(), ) assert resp2.status_code == 200 assert resp2.json() == SUCCESSFUL_CHAT_RESPONSE From cdbf4fd9b0ac73718579aa9811b16962cba97fb7 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Tue, 16 Jun 2026 09:09:33 -0400 Subject: [PATCH 2/4] feat: add fixture to mock FXA scopes in user signup tests --- src/tests/integration/test_user_signup_cap.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/tests/integration/test_user_signup_cap.py b/src/tests/integration/test_user_signup_cap.py index 908e50a..533b36d 100644 --- a/src/tests/integration/test_user_signup_cap.py +++ b/src/tests/integration/test_user_signup_cap.py @@ -12,6 +12,11 @@ def use_real_get_or_create_user(): return True +@pytest.fixture(autouse=True) +def single_fxa_scope(mocker): + mocker.patch("mlpa.core.auth.fxa.FXA_SCOPES", ("profile:uid",)) + + class _FakeResponse: def __init__(self, payload: dict): self._payload = payload From 34e0fc1ed45ea4be7e21a2e711bf1afc85cc75a3 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Tue, 16 Jun 2026 10:05:24 -0400 Subject: [PATCH 3/4] addr comments --- src/mlpa/core/routers/health/health.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/mlpa/core/routers/health/health.py b/src/mlpa/core/routers/health/health.py index e31747d..288314a 100644 --- a/src/mlpa/core/routers/health/health.py +++ b/src/mlpa/core/routers/health/health.py @@ -24,7 +24,6 @@ async def readiness_probe(): # todo add check to PG and LiteLLM status here pg_status = litellm_pg.check_status() app_attest_pg_status = app_attest_pg.check_status() - litellm_status = {} client = get_http_client() response = await client.get( LITELLM_READINESS_URL, headers=LITELLM_MASTER_AUTH_HEADERS, timeout=3 @@ -41,7 +40,7 @@ async def readiness_probe(): "app_attest": "connected" if app_attest_pg_status else "offline", }, "litellm": { - "litellm_version": litellm_info["litellm_version"], + "litellm_version": litellm_info.get("litellm_version", "N/A"), **litellm_status, }, } From 39d699c833426817a0f279f69eeedb266f5972d8 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Wed, 17 Jun 2026 09:06:45 -0400 Subject: [PATCH 4/4] cache litellm_version after first request --- src/mlpa/core/routers/health/health.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/src/mlpa/core/routers/health/health.py b/src/mlpa/core/routers/health/health.py index 288314a..756c7f4 100644 --- a/src/mlpa/core/routers/health/health.py +++ b/src/mlpa/core/routers/health/health.py @@ -11,9 +11,26 @@ from mlpa.core.pg_services.services import app_attest_pg, litellm_pg mlpa_version = importlib.metadata.version("mlpa") +litellm_version = "N/A" router = APIRouter() +async def get_litellm_version(client): + global litellm_version + + if litellm_version != "N/A": + return litellm_version + + try: + response = await client.get(LITELLM_INFO_URL, timeout=3) + litellm_info = response.json() + except Exception: + return litellm_version + + litellm_version = litellm_info.get("litellm_version", "N/A") + return litellm_version + + @router.get("/liveness", tags=["Health"]) async def liveness_probe(): return {"status": "alive"} @@ -29,9 +46,8 @@ async def readiness_probe(): LITELLM_READINESS_URL, headers=LITELLM_MASTER_AUTH_HEADERS, timeout=3 ) litellm_status = response.json() + current_litellm_version = await get_litellm_version(client) - response = await client.get(LITELLM_INFO_URL, timeout=3) - litellm_info = response.json() return { "status": "connected", "mlpa_version": mlpa_version, @@ -40,7 +56,7 @@ async def readiness_probe(): "app_attest": "connected" if app_attest_pg_status else "offline", }, "litellm": { - "litellm_version": litellm_info.get("litellm_version", "N/A"), + "litellm_version": current_litellm_version, **litellm_status, }, }