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..756c7f4 100644 --- a/src/mlpa/core/routers/health/health.py +++ b/src/mlpa/core/routers/health/health.py @@ -2,14 +2,35 @@ 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 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"} @@ -20,13 +41,13 @@ 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 ) - data = response.json() - litellm_status = data + litellm_status = response.json() + current_litellm_version = await get_litellm_version(client) + return { "status": "connected", "mlpa_version": mlpa_version, @@ -34,5 +55,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": current_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..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 @@ -101,7 +106,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 +117,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 +128,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 +144,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 +194,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 +204,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