diff --git a/src/litellm_admin_mcp/gateway.py b/src/litellm_admin_mcp/gateway.py index a4de0f3..a047705 100644 --- a/src/litellm_admin_mcp/gateway.py +++ b/src/litellm_admin_mcp/gateway.py @@ -117,7 +117,12 @@ async def _request(self, method: str, path: str, credential: str, *, raise AdminError("Gateway returned an invalid JSON response." + uncertain, 502) from None async def authorize(self, credential: str) -> str: - data = await self._request("GET", "/user/info", credential) + try: + data = await self._request("GET", "/user/info", credential) + except AdminError as exc: + if exc.status != 404: + raise + data = None user = data.get("user_info") if isinstance(data, dict) else None if (not isinstance(user, dict) or user.get("user_role") != "proxy_admin" or not isinstance(user.get("user_id"), str) or not user["user_id"] diff --git a/tests/conftest.py b/tests/conftest.py index 4d572c5..5a49a2f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -57,6 +57,8 @@ def __init__(self): def handle(self, request): self.requests.append(request) token = request.headers.get("authorization", "").removeprefix("Bearer ") + if request.url.path == "/user/info" and token == "team-key": + return httpx2.Response(404, json={"error": {"message": "User None not found"}}) if request.url.path == "/user/info": role = "proxy_admin" if token in {"alice-admin", "bob-admin"} and token not in self.revoked else "internal_user" return httpx2.Response(200, json={"user_id": token, "user_info": {"user_id": token, "user_role": role}}) diff --git a/tests/test_transports.py b/tests/test_transports.py index 0129518..a818e16 100644 --- a/tests/test_transports.py +++ b/tests/test_transports.py @@ -44,6 +44,7 @@ async def test_streamable_http_rejects_anonymous_and_keeps_callers_separate(stub assert (await anonymous.get(base + "/healthz")).status_code == 200 assert (await anonymous.post(base + "/mcp", json={})).status_code == 401 assert (await anonymous.post(base + "/mcp", json={}, headers={"Authorization": "Bearer viewer"})).status_code == 403 + assert (await anonymous.post(base + "/mcp", json={}, headers={"Authorization": "Bearer team-key"})).status_code == 403 for credential in ("alice-admin", "bob-admin"): async with httpx2.AsyncClient(headers={"Authorization": "Bearer " + credential}) as client: async with streamable_http_client(base + "/mcp", http_client=client) as (read, write):