LCORE-1246: Add OAuth authentication method for MCP servers#1156
LCORE-1246: Add OAuth authentication method for MCP servers#1156tisnik merged 4 commits intolightspeed-core:mainfrom
Conversation
- Introduced a new OAuth authentication method for MCP servers, allowing clients to use the "oauth" keyword in authorization headers. - Updated documentation to reflect the addition of OAuth as a supported authentication method. - Enhanced error handling in the tools endpoint to return a 401 status with a WWW-Authenticate header when OAuth is required but no token is provided. - Modified relevant tests to cover the new OAuth functionality and ensure proper error responses.
WalkthroughAdds OAuth as an MCP authentication option: new constant, preserved keyword in header resolution, async MCP OAuth probe that can raise 401 with WWW-Authenticate, conversion of get_mcp_tools to async (callsites updated), tools endpoint OAuth-aware 401 handling, docs and tests updated. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant Endpoint as Tools Endpoint
participant MCP as MCP Server
Client->>Endpoint: GET /tools (no token)
Endpoint->>MCP: GET /tools (no Authorization)
MCP-->>Endpoint: 401 Unauthorized + WWW-Authenticate
Endpoint->>Endpoint: probe_mcp_oauth_and_raise_401(url) -> raises HTTP 401 (propagate WWW-Authenticate)
Endpoint-->>Client: 401 Unauthorized + WWW-Authenticate
Client->>Client: Obtain token via OAuth flow
Client->>Endpoint: GET /tools (Authorization: Bearer <token>)
Endpoint->>MCP: GET /tools (with token)
MCP-->>Endpoint: 200 OK + tools
Endpoint-->>Client: 200 OK + tools
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Fix all issues with AI agents
In `@docs/config.md`:
- Line 374: In the authorization_headers description update the phrase "client
provided token" to the hyphenated form "client-provided token" (the clause that
explains using the string 'client' instead of a file path); edit the docs entry
for authorization_headers to use "client-provided" wherever it appears to
correct the compound adjective usage.
In `@src/app/endpoints/tools.py`:
- Line 5: The module imports the synchronous requests library (requests) but
defines async endpoint handlers, which will block the event loop; replace the
requests import with an async HTTP client (e.g., aiohttp.ClientSession or
httpx.AsyncClient), update any call sites that use requests.get/post/etc. to use
the async client with await (for example create or accept an
aiohttp.ClientSession and call await session.get(...)), and ensure proper
session lifecycle (use async context manager or dependency-injected session) so
all network I/O in your async endpoint functions remains non-blocking.
- Around line 93-104: The block catching AuthenticationError is performing a
blocking synchronous probe and not guarding against probe failures or missing
headers; update the probe to use an async HTTP client (e.g., aiohttp) and await
the request instead of requests.get, wrap the probe in try/except to catch
connection, timeout, and client errors (e.g., aiohttp.ClientError,
asyncio.TimeoutError) and fall back to a safe UnauthorizedResponse when the
probe fails, and when constructing the HTTPException use
resp.headers.get("WWW-Authenticate") with a sensible default (or omit the header
if absent) to avoid KeyError; apply the same async/non-blocking error-handling
pattern to the get_mcp_tools function in src/utils/responses.py.
In `@src/models/config.py`:
- Around line 512-514: The concatenated description string for the special cases
is missing a space between the end of the "client" sentence and the start of "3.
Usage of the oauth..." sentence; update the string fragments (the literal
containing "To specify this use a string 'client' instead of the file path." or
the immediately following fragment starting "3. Usage of the oauth token in the
header.") so there is a space either appended to the first fragment or prefixed
to the second fragment to ensure the final description reads "...file path. 3.
Usage of the oauth..." (refer to the config description string that includes the
"'client'" and "'oauth'" special-case text).
In `@src/utils/responses.py`:
- Around line 389-405: The probe using requests.get in the get_mcp_tools logic
blocks the async event loop and lacks error handling for connection/timeouts and
missing WWW-Authenticate header; replace the synchronous probe with an async
HTTP call (e.g., httpx.AsyncClient) or run requests.get in an executor, wrap the
call in try/except to catch request exceptions (ConnectionError, Timeout,
HTTPError) and handle them gracefully, and when building the HTTPException use
resp.headers.get("WWW-Authenticate") (only include the header if present) so you
don't KeyError; update the code paths in get_mcp_tools and callers
(prepare_tools) accordingly to await the async probe or handle executor futures.
🧹 Nitpick comments (6)
README.md (1)
462-462: Consider whether OAuth servers should appear in the/v1/mcp-auth/client-optionsendpoint response.This line states that only
"client"servers are returned by the discovery endpoint. Since"oauth"servers also require client-provided tokens (viaMCP-HEADERS), clients would benefit from discovering them here too. If intentionally excluded, a brief note explaining why would help.tests/unit/utils/test_mcp_auth_headers.py (1)
96-112: Consider adding"oauth"to the mixed-types test.The mixed-types test covers
client, file, andkubernetesbut notoauth. Adding it here would validate that all four header types work together.Suggested addition
headers = { "Authorization": "client", "X-API-Key": str(secret_file), "X-K8s-Token": "kubernetes", + "X-OAuth-Token": "oauth", } result = resolve_authorization_headers(headers) # Special keywords should be preserved, file should be resolved assert result["Authorization"] == "client" assert result["X-API-Key"] == "file-secret" assert result["X-K8s-Token"] == "kubernetes" + assert result["X-OAuth-Token"] == "oauth"tests/unit/utils/test_responses.py (1)
513-537: Test covers the happy path well, but consider adding edge-case tests.The test correctly validates the OAuth 401 flow. However, given the issues flagged in the production code, consider adding tests for:
- MCP server probe returning a response without
WWW-Authenticateheader (currently would causeKeyError)requests.getraising aConnectionErrorduring the probeThese would serve as regression tests once the production code is hardened.
src/utils/mcp_auth_headers.py (1)
19-22: Docstring updated correctly for OAuth.The parameter documentation now covers all three special keywords. Consider also adding an OAuth example in the
Examplessection (lines 28-39) for completeness — entirely optional.src/models/config.py (1)
535-551: Docstring doesn't mention'oauth'as a preserved special value.The
resolve_auth_headersdocstring (Line 542) states only'kubernetes'and'client'are preserved for later substitution. Now that'oauth'is also a pass-through keyword, consider updating the docstring to include it.src/utils/responses.py (1)
354-361: OAuth token resolution is an exact duplicate of theMCP_AUTH_CLIENTcase.The
MCP_AUTH_OAUTHbranch (lines 354-361) is identical toMCP_AUTH_CLIENT(lines 346-353). Consider merging them to avoid duplication.Proposed fix
- case constants.MCP_AUTH_CLIENT: - # use client provided 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 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 constants.MCP_AUTH_CLIENT | constants.MCP_AUTH_OAUTH: + # use client or oauth provided 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)
tisnik
left a comment
There was a problem hiding this comment.
LGTM, but please look at CodeRabbitAI comments, as some of them are relevant
- Updated the `get_mcp_tools` function to be asynchronous, allowing for non-blocking calls to MCP servers. - Introduced a new utility function `probe_mcp_oauth_and_raise_401` to handle OAuth probing and error responses. - Enhanced the tools endpoint to utilize the new OAuth probing method for improved error handling. - Updated documentation to clarify the usage of client-provided tokens in authorization headers. - Modified tests to accommodate the asynchronous nature of `get_mcp_tools` and ensure proper functionality of the new OAuth handling.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/utils/mcp_oauth_probe.py`:
- Around line 20-34: The function docstring in src/utils/mcp_oauth_probe.py that
starts "Probe MCP endpoint and raise 401..." is missing a Returns section; add a
Returns section following Google Python docstring conventions that states the
function returns None (it does not return a value) and documents the side-effect
that it raises an HTTPException on success or failure (refer to the existing
Raises section). Ensure the Returns section is placed after Args and before
Raises and clearly indicates "None" as the return value.
- Around line 3-19: Replace the old typing and FastAPI import patterns: in the
imports drop unused FastAPI symbols and keep only HTTPException, and change the
probe_mcp_oauth_and_raise_401 signature to use PEP 604 union syntax
(BaseException | None) instead of Optional[BaseException]; update any other
typing imports accordingly so only necessary imports remain and the function
parameter uses the modern union type.
…andling - Removed unnecessary import of `Optional` and simplified type hinting for `chain_from`. - Updated function documentation to clarify the behavior and return value. - Enhanced clarity in the argument description for `chain_from`.
| 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 |
There was a problem hiding this comment.
does not it leak any sensitive data into logs?
There was a problem hiding this comment.
what sensitive data would it leak? All I am logging is the endpoint that failed to deliver WWW-authenticate headers.
| 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( |
| ) 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 |
|
|
||
|
|
||
| def get_mcp_tools( | ||
| async def get_mcp_tools( # pylint: disable=too-many-return-statements,too-many-locals |
There was a problem hiding this comment.
the pylint is right - it's getting too complicated
There was a problem hiding this comment.
Yeah, this function has become pretty wild. But if the only problem is the decomposition, I would accept this for now.
asimurka
left a comment
There was a problem hiding this comment.
LGTM... if you are OK with the Critical suggestion.
|
|
||
|
|
||
| def get_mcp_tools( | ||
| async def get_mcp_tools( # pylint: disable=too-many-return-statements,too-many-locals |
There was a problem hiding this comment.
Yeah, this function has become pretty wild. But if the only problem is the decomposition, I would accept this for now.
|
Thanks for the review! I can work on splitting up the Also @asimurka wdym by "the critical suggestion"? I thought that I had resolved all the code rabbit issues, did I miss something? |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/unit/utils/test_responses.py (1)
527-564: Missing test: OAuth success path (whenmcp_headersprovides the token).The suite covers OAuth with no headers (→ 401) but not the happy path where
mcp_headerscontains the per-server entry and the token is successfully resolved. Add a test analogous totest_get_mcp_tools_with_mcp_headersbut withauthorization_headers={"Authorization": "oauth"}.🧪 Suggested test skeleton
`@pytest.mark.asyncio` async def test_get_mcp_tools_oauth_with_headers_returns_tool(self) -> None: """Test get_mcp_tools resolves OAuth token from mcp_headers when provided.""" servers = [ ModelContextProtocolServer( name="oauth-server", url="http://localhost:3000", authorization_headers={"Authorization": "oauth"}, ), ] mcp_headers = {"oauth-server": {"Authorization": "Bearer my-oauth-token"}} tools = await get_mcp_tools(servers, token=None, mcp_headers=mcp_headers) assert len(tools) == 1 assert tools[0]["headers"] == {"Authorization": "Bearer my-oauth-token"}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/utils/test_responses.py` around lines 527 - 564, Add a new async test that mirrors the existing OAuth-no-headers case but supplies mcp_headers so the OAuth token is resolved: create a servers list with ModelContextProtocolServer(name="oauth-server", url="http://localhost:3000", authorization_headers={"Authorization": "oauth"}), set mcp_headers = {"oauth-server": {"Authorization": "Bearer my-oauth-token"}}, call await get_mcp_tools(servers, token=None, mcp_headers=mcp_headers) and assert the returned tools list has length 1 and that tools[0]["headers"] == {"Authorization": "Bearer my-oauth-token"} (name the test e.g. test_get_mcp_tools_oauth_with_headers_returns_tool).src/utils/responses.py (1)
364-379:MCP_AUTH_CLIENTandMCP_AUTH_OAUTHcases are identical — collapse with an OR pattern.Both branches return a token from
mcp_headersusing the same logic; only the label differs. Python structural pattern matching supportscase A | B:.♻️ Proposed refactor
- case constants.MCP_AUTH_CLIENT: - # use client provided 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 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 constants.MCP_AUTH_CLIENT | constants.MCP_AUTH_OAUTH: + # use client-provided or OAuth token from per-request mcp_headers + 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)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/utils/responses.py` around lines 364 - 379, The two match arms for constants.MCP_AUTH_CLIENT and constants.MCP_AUTH_OAUTH in the pattern match are identical; collapse them into a single branch using an OR pattern (case constants.MCP_AUTH_CLIENT | constants.MCP_AUTH_OAUTH:) and keep the shared logic that checks mcp_headers, looks up mcp_server.name, and returns c_headers.get(header, None); update the existing duplicate blocks (referenced by constants.MCP_AUTH_CLIENT, constants.MCP_AUTH_OAUTH, mcp_headers, mcp_server, header) to a single implementation to avoid repetition.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/app/endpoints/rlsapi_v1.py`:
- Line 308: The call to get_mcp_tools(configuration.mcp_servers) can probe
OAuth-only MCPs and raise HTTPException(401), breaking the CLA /infer flow; wrap
the call in a try/except that catches HTTPException with status_code == 401
(and/or message indicating "requires OAuth"), log a warning, and set mcp_tools =
[] to silently skip OAuth-protected servers so CLA continues to work;
alternatively, adjust get_mcp_tools to skip probing when no token/mcp_headers
are provided (i.e., only probe servers if a token or mcp_headers exist) and
return an empty list when all servers require OAuth.
---
Duplicate comments:
In `@src/app/endpoints/tools.py`:
- Around line 93-100: The AuthenticationError except-block should await
probe_mcp_oauth_and_raise_401 and chain the original exception, then fall back
to raising an HTTPException built from UnauthorizedResponse(cause=str(e));
ensure the except handler for AuthenticationError still calls await
probe_mcp_oauth_and_raise_401(toolgroup.mcp_endpoint.uri, chain_from=e) when
toolgroup.mcp_endpoint is present and otherwise constructs error_response =
UnauthorizedResponse(cause=str(e)) and raises
HTTPException(**error_response.model_dump()) from e so the async probe and
exception chaining are preserved.
In `@src/utils/responses.py`:
- Line 31: The async import probe_mcp_oauth_and_raise_401 in
src/utils/responses.py is correct; ensure you actually call
probe_mcp_oauth_and_raise_401 from the async request-handling paths in this
module (or remove the import if it ends up unused) so the aiohttp-based probe
replaces the old synchronous requests.get logic and provides proper
WWW-Authenticate error handling.
---
Nitpick comments:
In `@src/utils/responses.py`:
- Around line 364-379: The two match arms for constants.MCP_AUTH_CLIENT and
constants.MCP_AUTH_OAUTH in the pattern match are identical; collapse them into
a single branch using an OR pattern (case constants.MCP_AUTH_CLIENT |
constants.MCP_AUTH_OAUTH:) and keep the shared logic that checks mcp_headers,
looks up mcp_server.name, and returns c_headers.get(header, None); update the
existing duplicate blocks (referenced by constants.MCP_AUTH_CLIENT,
constants.MCP_AUTH_OAUTH, mcp_headers, mcp_server, header) to a single
implementation to avoid repetition.
In `@tests/unit/utils/test_responses.py`:
- Around line 527-564: Add a new async test that mirrors the existing
OAuth-no-headers case but supplies mcp_headers so the OAuth token is resolved:
create a servers list with ModelContextProtocolServer(name="oauth-server",
url="http://localhost:3000", authorization_headers={"Authorization": "oauth"}),
set mcp_headers = {"oauth-server": {"Authorization": "Bearer my-oauth-token"}},
call await get_mcp_tools(servers, token=None, mcp_headers=mcp_headers) and
assert the returned tools list has length 1 and that tools[0]["headers"] ==
{"Authorization": "Bearer my-oauth-token"} (name the test e.g.
test_get_mcp_tools_oauth_with_headers_returns_tool).
| 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) |
There was a problem hiding this comment.
OAuth probe in get_mcp_tools will break the CLA endpoint when OAuth MCP servers are configured.
get_mcp_tools is called with no token and no mcp_headers. With this PR, if any MCP server in the configuration uses authorization_headers: {"Authorization": "oauth"}, the function will probe the MCP endpoint and raise HTTPException(401). Because this call sits outside the try/except block that wraps inference (Line 314), FastAPI propagates the 401 directly to the CLA client.
Impact: Any configuration that includes even one OAuth-configured MCP server renders the entire /infer endpoint non-functional for CLA users — they receive a 401 with "MCP server at … requires OAuth" for every request, with no way to recover since CLA has no OAuth flow. Before this PR, unresolvable auth servers were silently skipped.
Recommended mitigation — catch the OAuth 401 and fall back to an empty tools list:
🛡️ Proposed fix
- mcp_tools = await get_mcp_tools(configuration.mcp_servers)
+ try:
+ mcp_tools = await get_mcp_tools(configuration.mcp_servers)
+ except HTTPException:
+ logger.warning(
+ "MCP tools unavailable for rlsapi request %s (OAuth MCP servers require "
+ "per-request headers not supported by this endpoint); proceeding without MCP tools.",
+ request_id,
+ )
+ mcp_tools = []📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| mcp_tools = await get_mcp_tools(configuration.mcp_servers) | |
| try: | |
| mcp_tools = await get_mcp_tools(configuration.mcp_servers) | |
| except HTTPException: | |
| logger.warning( | |
| "MCP tools unavailable for rlsapi request %s (OAuth MCP servers require " | |
| "per-request headers not supported by this endpoint); proceeding without MCP tools.", | |
| request_id, | |
| ) | |
| mcp_tools = [] |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/app/endpoints/rlsapi_v1.py` at line 308, The call to
get_mcp_tools(configuration.mcp_servers) can probe OAuth-only MCPs and raise
HTTPException(401), breaking the CLA /infer flow; wrap the call in a try/except
that catches HTTPException with status_code == 401 (and/or message indicating
"requires OAuth"), log a warning, and set mcp_tools = [] to silently skip
OAuth-protected servers so CLA continues to work; alternatively, adjust
get_mcp_tools to skip probing when no token/mcp_headers are provided (i.e., only
probe servers if a token or mcp_headers exist) and return an empty list when all
servers require OAuth.
I was talking about that one in |
Description
MCP servers can use the "oauth" keyword in authorization_headers. If no token is provided, the service returns 401 with a WWW-Authenticate header (from the MCP server) so clients can run OAuth and retry with a token in MCP-HEADERS.
Type of change
Tools used to create PR
Identify any AI code assistants used in this PR (for transparency and review context)
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
New Features
Bug Fixes
Documentation