Merge pull request #1156 from jrobertboos/lcore-1246

LCORE-1246: Add OAuth authentication method for MCP servers

Author: tisnik

README.md

@@ -34,6 +34,7 @@ The service includes comprehensive user data collection capabilities for various
                 * [1. Static Tokens from Files (Recommended for Service Credentials)](#1-static-tokens-from-files-recommended-for-service-credentials)
                 * [2. Kubernetes Service Account Tokens (For K8s Deployments)](#2-kubernetes-service-account-tokens-for-k8s-deployments)
                 * [3. Client-Provided Tokens (For Per-User Authentication)](#3-client-provided-tokens-for-per-user-authentication)
+                * [4. OAuth (For MCP Servers Requiring OAuth)](#4-oauth-for-mcp-servers-requiring-oauth)
                 * [Client-Authenticated MCP Servers Discovery](#client-authenticated-mcp-servers-discovery)
                 * [Combining Authentication Methods](#combining-authentication-methods)
                 * [Authentication Method Comparison](#authentication-method-comparison)
@@ -355,7 +356,7 @@ In addition to the basic configuration above, you can configure authentication h
 
 #### Configuring MCP Server Authentication
 
-Lightspeed Core Stack supports three methods for authenticating with MCP servers, each suited for different use cases:
+Lightspeed Core Stack supports four methods for authenticating with MCP servers, each suited for different use cases:
 
 ##### 1. Static Tokens from Files (Recommended for Service Credentials)
 
@@ -392,7 +393,7 @@ mcp_servers:
       Authorization: "kubernetes"    # Uses user's k8s token from request auth
 ```
 
-**Note:** Kubernetes token-based MCP authorization only works when Lightspeed Core Stack is configured with Kubernetes authentication (`authentication.k8s`). For any other authentication types, MCP servers configured with `Authorization: "kubernetes"` are removed from the available MCP servers list.
+**Note:** Kubernetes token-based MCP authorization only works when Lightspeed Core Stack is configured with Kubernetes authentication (`authentication.module` is `k8s`) or `noop-with-token`. For any other authentication types, MCP servers configured with `Authorization: "kubernetes"` are removed from the available MCP servers list.
 
 ##### 3. Client-Provided Tokens (For Per-User Authentication)
 
@@ -420,6 +421,20 @@ curl -X POST "http://localhost:8080/v1/query" \
 
 **Structure**: `MCP-HEADERS: {"<server-name>": {"<header-name>": "<header-value>", ...}, ...}`
 
+##### 4. OAuth (For MCP Servers Requiring OAuth)
+
+Use the special `"oauth"` keyword when the MCP server requires OAuth and the client will supply a token (e.g. via `MCP-HEADERS` after obtaining it from an OAuth flow):
+
+```yaml
+mcp_servers:
+  - name: "oauth-protected-service"
+    url: "https://mcp.example.com"
+    authorization_headers:
+      Authorization: "oauth"    # Token provided via MCP-HEADERS (from OAuth flow)
+```
+
+When no token is provided for an OAuth-configured server, the service may respond with **401 Unauthorized** and a **`WWW-Authenticate`** header (probed from the MCP server). Clients can use this to drive an OAuth flow and then retry with the token in `MCP-HEADERS`.
+
 ##### Client-Authenticated MCP Servers Discovery
 
 To help clients determine which MCP servers require client-provided tokens, use the **MCP Client Auth Options** endpoint:
@@ -481,6 +496,7 @@ mcp_servers:
 | **Static File** | Service tokens, API keys | File path in config | Global (all users) | `"/var/secrets/token"` |
 | **Kubernetes** | K8s service accounts | `"kubernetes"` keyword | Per-user (from auth) | `"kubernetes"` |
 | **Client** | User-specific tokens | `"client"` keyword + HTTP header | Per-request | `"client"` |
+| **OAuth** | OAuth-protected MCP servers | `"oauth"` keyword + HTTP header | Per-request (from OAuth flow) | `"oauth"` |
 
 ##### Important: Automatic Server Skipping
 
@@ -489,6 +505,7 @@ mcp_servers:
 **Examples:**
 - A server with `Authorization: "kubernetes"` will be skipped if the user's request doesn't include a Kubernetes token
 - A server with `Authorization: "client"` will be skipped if no `MCP-HEADERS` are provided in the request
+- A server with `Authorization: "oauth"` and no token in `MCP-HEADERS` may cause the API to return **401 Unauthorized** with a **`WWW-Authenticate`** header (so the client can perform OAuth and retry)
 - A server with multiple headers will be skipped if **any** required header cannot be resolved
 
 Skipped servers are logged as warnings. Check Lightspeed Core logs to see which servers were skipped and why.

docs/config.md

@@ -372,7 +372,7 @@ Useful resources:
 | name | string | MCP server name that must be unique |
 | provider_id | string | MCP provider identification |
 | url | string | URL of the MCP server |
-| authorization_headers | object | Headers to send to the MCP server. The map contains the header name and the path to a file containing the header value (secret). There are 2 special cases: 1. Usage of the kubernetes token in the header. To specify this use a string 'kubernetes' instead of the file path. 2. Usage of the client provided token in the header. To specify this use a string 'client' instead of the file path. |
+| authorization_headers | object | Headers to send to the MCP server. The map contains the header name and the path to a file containing the header value (secret). There are 3 special cases: 1. Usage of the kubernetes token in the header — use the string 'kubernetes' instead of the file path. 2. Usage of the client-provided token in the header — use the string 'client' instead of the file path. 3. Usage of OAuth token (resolved at request time or 401 with WWW-Authenticate) — use the string 'oauth' instead of the file path. |
 | timeout | integer | Timeout in seconds for requests to the MCP server. If not specified, the default timeout from Llama Stack will be used. Note: This field is reserved for future use when Llama Stack adds timeout support. |
 
 

src/app/endpoints/rlsapi_v1.py

@@ -305,7 +305,7 @@ async def infer_endpoint(
     input_source = infer_request.get_input_source()
     instructions = _build_instructions(infer_request.context.systeminfo)
     model_id = _get_default_model_id()
-    mcp_tools = get_mcp_tools(configuration.mcp_servers)
+    mcp_tools = await get_mcp_tools(configuration.mcp_servers)
     logger.debug(
         "Request %s: Combined input source length: %d", request_id, len(input_source)
     )

src/app/endpoints/tools.py

@@ -3,7 +3,7 @@
 from typing import Annotated, Any
 
 from fastapi import APIRouter, Depends, HTTPException, Request
-from llama_stack_client import APIConnectionError, BadRequestError
+from llama_stack_client import APIConnectionError, BadRequestError, AuthenticationError
 
 from authentication import get_auth_dependency
 from authentication.interface import AuthTuple
@@ -19,6 +19,7 @@
     UnauthorizedResponse,
 )
 from utils.endpoints import check_configuration_loaded
+from utils.mcp_oauth_probe import probe_mcp_oauth_and_raise_401
 from utils.tool_formatter import format_tools_list
 from log import get_logger
 
@@ -39,7 +40,7 @@
 
 @router.get("/tools", responses=tools_responses)
 @authorize(Action.GET_TOOLS)
-async def tools_endpoint_handler(  # pylint: disable=too-many-locals
+async def tools_endpoint_handler(  # pylint: disable=too-many-locals,too-many-statements
     request: Request,
     auth: Annotated[AuthTuple, Depends(get_auth_dependency())],
 ) -> ToolsResponse:
@@ -89,6 +90,14 @@ async def tools_endpoint_handler(  # pylint: disable=too-many-locals
         except BadRequestError:
             logger.error("Toolgroup %s is not found", toolgroup.identifier)
             continue
+        except AuthenticationError as e:
+            logger.error("Authentication error: %s", e)
+            if toolgroup.mcp_endpoint:
+                await probe_mcp_oauth_and_raise_401(
+                    toolgroup.mcp_endpoint.uri, chain_from=e
+                )
+            error_response = UnauthorizedResponse(cause=str(e))
+            raise HTTPException(**error_response.model_dump()) from e
         except APIConnectionError as e:
             logger.error("Unable to connect to Llama Stack: %s", e)
             response = ServiceUnavailableResponse(

src/constants.py

@@ -125,6 +125,7 @@
 # MCP authorization header special values
 MCP_AUTH_KUBERNETES = "kubernetes"
 MCP_AUTH_CLIENT = "client"
+MCP_AUTH_OAUTH = "oauth"
 
 # default RAG tool value
 DEFAULT_RAG_TOOL = "file_search"

src/models/config.py

@@ -505,11 +505,13 @@ class ModelContextProtocolServer(ConfigurationBase):
             "Headers to send to the MCP server. "
             "The map contains the header name and the path to a file containing "
             "the header value (secret). "
-            "There are 2 special cases: "
+            "There are 3 special cases: "
             "1. Usage of the kubernetes token in the header. "
             "To specify this use a string 'kubernetes' instead of the file path. "
-            "2. Usage of the client provided token in the header. "
-            "To specify this use a string 'client' instead of the file path."
+            "2. Usage of the client-provided token in the header. "
+            "To specify this use a string 'client' instead of the file path. "
+            "3. Usage of the oauth token in the header. "
+            "To specify this use a string 'oauth' instead of the file path. "
         ),
     )
 

src/utils/mcp_auth_headers.py

@@ -16,8 +16,10 @@ def resolve_authorization_headers(
 
     Parameters:
         authorization_headers: Map of header names to secret locations or special keywords.
-            - If value is "kubernetes": leave is unchanged. We substitute it during request.
-            - If value is "client": leave it unchanged. . We substitute it during request.
+            - If value is "kubernetes": leave unchanged. We substitute it during request.
+            - If value is "client": leave unchanged. We substitute it during request.
+            - If value is "oauth": leave unchanged; if no token is provided, a 401 with
+              WWW-Authenticate may be forwarded from the MCP server.
             - Otherwise: Treat as file path and read the secret from that file
 
     Returns:
@@ -55,6 +57,14 @@ def resolve_authorization_headers(
                     "Header %s will use client-provided token (resolved at request time)",
                     header_name,
                 )
+            elif value == constants.MCP_AUTH_OAUTH:
+                # Special case: Keep oauth keyword; if no token provided, probe endpoint
+                # and forward 401 WWW-Authenticate for client-driven OAuth flow
+                resolved[header_name] = constants.MCP_AUTH_OAUTH
+                logger.debug(
+                    "Header %s will use OAuth token (resolved at request time or 401)",
+                    header_name,
+                )
             else:
                 # Regular case: Read secret from file path
                 secret_path = Path(value).expanduser()

src/utils/mcp_oauth_probe.py

@@ -0,0 +1,51 @@
+"""Probe MCP server for OAuth and raise 401 with WWW-Authenticate when required."""
+
+import aiohttp
+from fastapi import HTTPException
+
+from models.responses import UnauthorizedResponse
+
+from log import get_logger
+
+logger = get_logger(__name__)
+
+
+async def probe_mcp_oauth_and_raise_401(
+    url: str,
+    chain_from: BaseException | None = None,
+) -> None:
+    """Probe MCP endpoint and raise 401 so the client can perform OAuth.
+
+    Performs an async GET to the given URL to obtain a WWW-Authenticate header,
+    then raises HTTPException with status 401 and that header. If the probe
+    fails (connection error, timeout), raises 401 without the header.
+
+    Args:
+        url: MCP server URL to probe.
+        chain_from: Exception to chain the HTTPException from when
+            the probe succeeds (e.g. the original AuthenticationError).
+
+    Returns:
+        None. Always raises an HTTPException.
+
+    Raises:
+        HTTPException: 401 with WWW-Authenticate when the probe succeeds, or
+            401 without the header when the probe fails.
+    """
+    cause = f"MCP server at {url} requires OAuth"
+    error_response = UnauthorizedResponse(cause=cause)
+    try:
+        timeout = aiohttp.ClientTimeout(total=10)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.get(url) as resp:
+                www_auth = resp.headers.get("WWW-Authenticate")
+                if www_auth is None:
+                    logger.warning("No WWW-Authenticate header received from %s", url)
+                    raise HTTPException(**error_response.model_dump()) from chain_from
+                raise HTTPException(
+                    **error_response.model_dump(),
+                    headers={"WWW-Authenticate": www_auth},
+                ) from chain_from
+    except (aiohttp.ClientError, TimeoutError) as probe_err:
+        logger.warning("OAuth probe failed for %s: %s", url, probe_err)
+        raise HTTPException(**error_response.model_dump()) from probe_err

src/utils/responses.py

@@ -28,6 +28,7 @@
     InternalServerErrorResponse,
     ServiceUnavailableResponse,
 )
+from utils.mcp_oauth_probe import probe_mcp_oauth_and_raise_401
 from utils.prompts import get_system_prompt, get_topic_summary_system_prompt
 from utils.query import (
     evaluate_model_hints,
@@ -183,7 +184,7 @@ async def prepare_tools(
         toolgroups.extend(rag_tools)
 
     # Add MCP server tools
-    mcp_tools = get_mcp_tools(config.mcp_servers, token, mcp_headers)
+    mcp_tools = await get_mcp_tools(config.mcp_servers, token, mcp_headers)
     if mcp_tools:
         toolgroups.extend(mcp_tools)
         logger.debug(
@@ -332,7 +333,7 @@ def get_rag_tools(vector_store_ids: list[str]) -> Optional[list[dict[str, Any]]]
     ]
 
 
-def get_mcp_tools(
+async def get_mcp_tools(  # pylint: disable=too-many-return-statements,too-many-locals
     mcp_servers: list[ModelContextProtocolServer],
     token: str | None = None,
     mcp_headers: Optional[McpHeaders] = None,
@@ -346,6 +347,10 @@ def get_mcp_tools(
 
     Returns:
         List of MCP tool definitions with server details and optional auth headers
+
+    Raises:
+        HTTPException: 401 with WWW-Authenticate header when an MCP server uses OAuth,
+            no headers are passed, and the server responds with 401 and WWW-Authenticate.
     """
 
     def _get_token_value(original: str, header: str) -> str | None:
@@ -364,6 +369,14 @@ def _get_token_value(original: str, header: str) -> str | None:
                 if c_headers is None:
                     return None
                 return c_headers.get(header, None)
+            case constants.MCP_AUTH_OAUTH:
+                # use oauth token
+                if mcp_headers is None:
+                    return None
+                c_headers = mcp_headers.get(mcp_server.name, None)
+                if c_headers is None:
+                    return None
+                return c_headers.get(header, None)
             case _:
                 # use provided
                 return original
@@ -391,6 +404,16 @@ def _get_token_value(original: str, header: str) -> str | None:
         if mcp_server.authorization_headers and len(headers) != len(
             mcp_server.authorization_headers
         ):
+            # If OAuth was required and no headers passed, probe endpoint and forward
+            # 401 with WWW-Authenticate so the client can perform OAuth
+            uses_oauth = (
+                constants.MCP_AUTH_OAUTH
+                in mcp_server.resolved_authorization_headers.values()
+            )
+            if uses_oauth and (
+                mcp_headers is None or not mcp_headers.get(mcp_server.name)
+            ):
+                await probe_mcp_oauth_and_raise_401(mcp_server.url)
             logger.warning(
                 "Skipping MCP server %s: required %d auth headers but only resolved %d",
                 mcp_server.name,

tests/integration/endpoints/test_rlsapi_v1_integration.py

@@ -389,6 +389,7 @@ async def test_rlsapi_v1_infer_no_mcp_servers_passes_empty_tools(
 
     mocker.patch(
         "app.endpoints.rlsapi_v1.get_mcp_tools",
+        new_callable=mocker.AsyncMock,
         return_value=[],
     )
 
@@ -437,6 +438,7 @@ async def test_rlsapi_v1_infer_mcp_tools_passed_to_llm(
     ]
     mocker.patch(
         "app.endpoints.rlsapi_v1.get_mcp_tools",
+        new_callable=mocker.AsyncMock,
         return_value=mcp_tools,
     )