fix(server): ensure resource protected metadata RFC compliance (#1200)

Signed-off-by: Tomas Pilar <thomas7pilar@gmail.com>

Author: pilartomas

apps/beeai-cli/src/beeai_cli/commands/server.py

@@ -126,7 +126,7 @@ async def server_login(server: typing.Annotated[str | None, typer.Argument()] =
     else:
         console.info("No authentication tokens found for this server. Proceeding to log in.")
         async with httpx.AsyncClient(verify=await get_verify_option(server)) as client:
-            resp = await client.get(f"{server}/api/v1/.well-known/oauth-protected-resource")
+            resp = await client.get(f"{server}/.well-known/oauth-protected-resource")
             if resp.is_error:
                 console.error("This server does not appear to run a compatible version of BeeAI Platform.")
                 sys.exit(1)

apps/beeai-server/src/beeai_server/api/dependencies.py

@@ -5,7 +5,7 @@
 from typing import Annotated
 from uuid import UUID
 
-from fastapi import Depends, HTTPException, Path, Query, Security, status
+from fastapi import Depends, HTTPException, Path, Query, Request, Security, status
 from fastapi.security import APIKeyCookie, HTTPAuthorizationCredentials, HTTPBasic, HTTPBasicCredentials, HTTPBearer
 from jwt import PyJWTError
 from kink import di
@@ -56,6 +56,7 @@ async def authenticate_oauth_user(
     cookie_auth: str | None,
     user_service: UserServiceDependency,
     configuration: ConfigurationDependency,
+    request: Request,
 ) -> AuthorizedUser:
     """
     Authenticate using an OIDC/OAuth2 JWT bearer token with JWKS.
@@ -69,8 +70,9 @@ async def authenticate_oauth_user(
             detail=f"Invalid Authorization header: {e}",
         ) from e
 
+    expected_audience = str(request.url.replace(path="/"))
     claims, issuer = await decode_oauth_jwt_or_introspect(
-        token=token, jwks_dict=di["JWKS_CACHE"], aud="beeai-server", configuration=configuration
+        token=token, jwks_dict=di["JWKS_CACHE"], aud=expected_audience, configuration=configuration
     )
     if not claims:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or expired token")
@@ -112,6 +114,7 @@ async def authorized_user(
     basic_auth: Annotated[HTTPBasicCredentials | None, Depends(HTTPBasic(auto_error=False))],
     bearer_auth: Annotated[HTTPAuthorizationCredentials | None, Depends(HTTPBearer(auto_error=False))],
     cookie_auth: Annotated[str | None, Security(api_key_cookie)],
+    request: Request,
 ) -> AuthorizedUser:
     if bearer_auth:
         # Check Bearer token first - locally this allows for "checking permissions" for development purposes
@@ -128,12 +131,12 @@ async def authorized_user(
             return token
         except PyJWTError:
             if configuration.auth.oidc.enabled:
-                return await authenticate_oauth_user(bearer_auth, cookie_auth, user_service, configuration)
+                return await authenticate_oauth_user(bearer_auth, cookie_auth, user_service, configuration, request)
             # TODO: update agents
             logger.warning("Bearer token is invalid, agent is not probably not using llm extension correctly")
 
     if configuration.auth.oidc.enabled and cookie_auth:
-        return await authenticate_oauth_user(bearer_auth, cookie_auth, user_service, configuration)
+        return await authenticate_oauth_user(bearer_auth, cookie_auth, user_service, configuration, request)
 
     if configuration.auth.basic.enabled:
         if basic_auth and basic_auth.password == configuration.auth.basic.admin_password.get_secret_value():

apps/beeai-server/src/beeai_server/api/routes/auth.py

@@ -2,15 +2,19 @@
 # SPDX-License-Identifier: Apache-2.0
 import logging
 
-from fastapi import APIRouter
+from fastapi import APIRouter, Request
 
 from beeai_server.api.dependencies import AuthServiceDependency
 
 logger = logging.getLogger(__name__)
 
-router = APIRouter()
+well_known_router = APIRouter()
 
 
-@router.get("/.well-known/oauth-protected-resource")
-def protected_resource_metadata(auth_servide: AuthServiceDependency):
-    return auth_servide.protected_resource_metadata()
+@well_known_router.get("/oauth-protected-resource/{resource:path}")
+def protected_resource_metadata(
+    request: Request,
+    auth_servide: AuthServiceDependency,
+    resource: str = "",
+):
+    return auth_servide.protected_resource_metadata(resource=str(request.url.replace(path=resource)))

apps/beeai-server/src/beeai_server/application.py

@@ -12,10 +12,10 @@
 from kink import Container, di, inject
 from opentelemetry.metrics import CallbackOptions, Observation, get_meter
 from starlette.requests import Request
-from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
+from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_500_INTERNAL_SERVER_ERROR
 
 from beeai_server.api.routes.a2a import router as a2a_router
-from beeai_server.api.routes.auth import router as auth_router
+from beeai_server.api.routes.auth import well_known_router as auth_well_known_router
 from beeai_server.api.routes.configurations import router as configuration_router
 from beeai_server.api.routes.contexts import router as contexts_router
 from beeai_server.api.routes.files import router as files_router
@@ -72,12 +72,18 @@ async def custom_http_exception_handler(request: Request, exc):
                 exception = exc
             case _:
                 exception = HTTPException(HTTP_500_INTERNAL_SERVER_ERROR, detail=repr(extract_messages(exc)))
+
+        if isinstance(exception, HTTPException) and exc.status_code == HTTP_401_UNAUTHORIZED:
+            exception.headers = exception.headers or {}
+            exception.headers |= {
+                "WWW-Authenticate": f'Bearer resource_metadata="{request.url.replace(path="/.well-known/oauth-protected-resource")}"'  # We don't define multiple resource domains at the moment
+            }
+
         return await http_exception_handler(request, exception)
 
 
 def mount_routes(app: FastAPI):
     server_router = APIRouter()
-    server_router.include_router(auth_router, prefix="", tags=["auth"])
     server_router.include_router(a2a_router, prefix="/a2a")
     server_router.include_router(mcp_router, prefix="/mcp")
     server_router.include_router(provider_router, prefix="/providers", tags=["providers"])
@@ -89,7 +95,11 @@ def mount_routes(app: FastAPI):
     server_router.include_router(vector_stores_router, prefix="/vector_stores", tags=["vector_stores"])
     server_router.include_router(user_feedback_router, prefix="/user_feedback", tags=["user_feedback"])
 
+    well_known_router = APIRouter()
+    well_known_router.include_router(auth_well_known_router, prefix="")
+
     app.include_router(server_router, prefix="/api/v1", tags=["provider"])
+    app.include_router(well_known_router, prefix="/.well-known", tags=["well-known"])
 
     @app.get("/healthcheck")
     async def healthcheck():

apps/beeai-server/src/beeai_server/service_layer/services/auth.py

@@ -15,13 +15,9 @@ class AuthService:
     def __init__(self, configuration: Configuration):
         self._config = configuration
 
-    def protected_resource_metadata(self) -> dict:
-        resource_id = f"http://localhost:{self._config.port}"  # TODO
-        providers = self._config.auth.oidc.providers
-        authorization_server = [str(p.issuer) for p in providers if p.issuer is not None]
-
+    def protected_resource_metadata(self, *, resource: str) -> dict:
         return {
-            "resource_id": resource_id,
-            "authorization_servers": authorization_server,
+            "resource": resource,
+            "authorization_servers": [str(p.issuer) for p in self._config.auth.oidc.providers if p.issuer is not None],
             "scopes_supported": list(self._config.auth.oidc.scope),
         }