LCORE-1420: Fixing MCP Authorization#1278
Open
jrobertboos wants to merge 3 commits intolightspeed-core:mainfrom
Open
LCORE-1420: Fixing MCP Authorization#1278jrobertboos wants to merge 3 commits intolightspeed-core:mainfrom
jrobertboos wants to merge 3 commits intolightspeed-core:mainfrom
Conversation
- Removed unused AuthenticationError handling in tools endpoint. - Updated `probe_mcp_oauth_and_raise_401` to accept an authorization header. - Enhanced OAuth probing logic to check for 401 status and include authorization in requests. - Cleaned up e2e tests by removing unnecessary assertions and skips related to MCP authorization. - Added response handling for authorized status in mock MCP server.
Contributor
WalkthroughCentralized MCP OAuth checks were introduced: endpoints now call a new async check_mcp_auth after configuration validation. probe_mcp was refactored to accept optional Authorization and to raise HTTPException (with/without WWW-Authenticate). Supporting code and tests were updated to use the new flow. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant Endpoint as App Endpoint
participant Config as AppConfig
participant MCP as MCP Server
Client->>Endpoint: HTTP request
Endpoint->>Config: validate & load configuration
Endpoint->>Endpoint: await check_mcp_auth(configuration, mcp_headers)
Endpoint->>MCP: probe_mcp GET (optional Authorization)
MCP-->>Endpoint: 200 / 401 (+ optional WWW-Authenticate)
Endpoint-->>Client: proceed or raise HTTPException(401, maybe WWW-Authenticate)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 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 |
- Added `check_mcp_auth` function to probe MCP servers for OAuth requirements. - Updated `query_endpoint_handler`, `streaming_query_endpoint_handler`, and `tools_endpoint_handler` to call `check_mcp_auth` for OAuth validation. - Removed deprecated `probe_mcp_oauth_and_raise_401` references from the codebase. - Enhanced documentation for the new `check_mcp_auth` function to clarify its purpose and usage.
Contributor
Author
|
In library mode for the This is likely what is also causing the |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/query.py`:
- Around line 126-127: The non-streaming query endpoint must catch
asyncio.CancelledError around the blocking calls to get_topic_summary(...) and
the call chain that invokes responses.create(...) inside retrieve_response(...)
to avoid 500s on client cancellation; update the handler in query.py to wrap
those await calls (or the higher-level retrieve_response call) with an except
asyncio.CancelledError: block that logs or silently returns an appropriate
response/cleanup and re-raises or returns gracefully (mirroring the streaming
handler pattern), ensuring you reference get_topic_summary, retrieve_response,
and responses.create to locate the spots to add the cancellation handling.
In `@src/app/endpoints/tools.py`:
- Around line 148-150: The code mutates the request-scoped mcp_headers by
calling headers = mcp_headers.get(toolgroup.identifier, {}) followed by
headers.pop("Authorization", None); change this to operate on a local copy:
retrieve the dict via mcp_headers.get(toolgroup.identifier) (or {}), create a
shallow copy (e.g., dict(headers) or headers.copy()) into a new variable (e.g.,
local_headers), then extract the Authorization value from that copy into
authorization without modifying the original mcp_headers; update subsequent uses
to reference local_headers instead of headers (look for occurrences around
toolgroup.identifier, headers, and authorization in src/app/endpoints/tools.py).
In `@src/utils/mcp_oauth_probe.py`:
- Around line 42-43: The code currently retrieves the Authorization token using
headers.get("Authorization") which is case-sensitive; update the lookup in
mcp_oauth_probe.py to be case-insensitive by normalizing header keys (e.g.,
build a temporary dict with lowercased keys from headers and then use lowercased
"authorization") or check both "Authorization" and "authorization" before
falling back to None so the token is found regardless of header casing for
mcp_headers, headers, authorization, and mcp_server.name.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 82a1905f-4294-4403-b48a-0d0edf768690
📒 Files selected for processing (10)
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.pysrc/app/endpoints/tools.pysrc/utils/mcp_oauth_probe.pysrc/utils/responses.pytests/e2e/features/mcp.featuretests/e2e/mock_mcp_server/server.pytests/integration/endpoints/test_tools_integration.pytests/unit/app/endpoints/test_tools.pytests/unit/utils/test_responses.py
💤 Files with no reviewable changes (2)
- src/utils/responses.py
- tests/unit/utils/test_responses.py
Comment on lines
+126
to
+127
| await check_mcp_auth(configuration, mcp_headers) | ||
|
|
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect cancellation handling vs. response/topic-summary call sites.
rg -n "responses\.create|get_topic_summary|CancelledError" src/app/endpoints/query.py src/app/endpoints/streaming_query.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 932
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read query.py around the critical lines
echo "=== query.py around line 126-127 (auth check) ==="
sed -n '120,135p' src/app/endpoints/query.py
echo ""
echo "=== query.py around line 222 (get_topic_summary call) ==="
sed -n '200,240p' src/app/endpoints/query.py
echo ""
echo "=== query.py around line 318 (responses.create call) ==="
sed -n '300,330p' src/app/endpoints/query.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 3859
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Get the function definition and overall structure
echo "=== query.py function signature and structure (lines 1-150) ==="
sed -n '80,150p' src/app/endpoints/query.py
echo ""
echo "=== Check for outer try-except blocks in query handler ==="
rg -n "try:|except|CancelledError" src/app/endpoints/query.py | head -20Repository: lightspeed-core/lightspeed-stack
Length of output: 3403
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# See lines 200-350 to understand exception handling scope
echo "=== query.py lines 200-350 (full context) ==="
sed -n '200,350p' src/app/endpoints/query.py | cat -n
echo ""
echo "=== streaming_query.py exception handling around line 490 ==="
sed -n '480,520p' src/app/endpoints/streaming_query.py | cat -nRepository: lightspeed-core/lightspeed-stack
Length of output: 8827
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check get_topic_summary implementation
echo "=== get_topic_summary implementation ==="
rg -A 15 "^async def get_topic_summary" src/Repository: lightspeed-core/lightspeed-stack
Length of output: 1031
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Get the full get_topic_summary implementation
echo "=== get_topic_summary full implementation ==="
rg -A 40 "^async def get_topic_summary" src/utils/responses.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 1623
Add CancelledError handling to non-streaming query path to prevent 500 errors during request cancellation.
The non-streaming query endpoint lacks explicit cancellation handling for get_topic_summary() (line 222) and responses.create() (within retrieve_response()). If the client cancels the request during either call, the unhandled CancelledError propagates to the ASGI layer, resulting in a 500 error. The streaming endpoint handles this correctly with an explicit except asyncio.CancelledError: handler (see streaming_query.py:490), but the non-streaming path needs equivalent protection to prevent 500 responses when authenticated MCP requests are interrupted.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/app/endpoints/query.py` around lines 126 - 127, The non-streaming query
endpoint must catch asyncio.CancelledError around the blocking calls to
get_topic_summary(...) and the call chain that invokes responses.create(...)
inside retrieve_response(...) to avoid 500s on client cancellation; update the
handler in query.py to wrap those await calls (or the higher-level
retrieve_response call) with an except asyncio.CancelledError: block that logs
or silently returns an appropriate response/cleanup and re-raises or returns
gracefully (mirroring the streaming handler pattern), ensuring you reference
get_topic_summary, retrieve_response, and responses.create to locate the spots
to add the cancellation handling.
Comment on lines
148
to
+150
| headers = mcp_headers.get(toolgroup.identifier, {}) | ||
| authorization = headers.pop("Authorization", None) | ||
|
|
Contributor
There was a problem hiding this comment.
Avoid mutating request-scoped MCP headers in-place.
Line 149 uses pop, which mutates the dict sourced from mcp_headers. Use a local copy/extraction so downstream code sees stable input.
♻️ Proposed fix
- headers = mcp_headers.get(toolgroup.identifier, {})
- authorization = headers.pop("Authorization", None)
+ original_headers = mcp_headers.get(toolgroup.identifier, {})
+ authorization = next(
+ (v for k, v in original_headers.items() if k.lower() == "authorization"),
+ None,
+ )
+ headers = {
+ k: v for k, v in original_headers.items() if k.lower() != "authorization"
+ }As per coding guidelines "Avoid in-place parameter modification anti-patterns; return new data structures instead."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/app/endpoints/tools.py` around lines 148 - 150, The code mutates the
request-scoped mcp_headers by calling headers =
mcp_headers.get(toolgroup.identifier, {}) followed by
headers.pop("Authorization", None); change this to operate on a local copy:
retrieve the dict via mcp_headers.get(toolgroup.identifier) (or {}), create a
shallow copy (e.g., dict(headers) or headers.copy()) into a new variable (e.g.,
local_headers), then extract the Authorization value from that copy into
authorization without modifying the original mcp_headers; update subsequent uses
to reference local_headers instead of headers (look for occurrences around
toolgroup.identifier, headers, and authorization in src/app/endpoints/tools.py).
Comment on lines
+42
to
+43
| headers = mcp_headers.get(mcp_server.name, {}) | ||
| authorization = headers.get("Authorization", None) |
Contributor
There was a problem hiding this comment.
Handle Authorization lookup case-insensitively.
Line 43 only checks "Authorization" with exact casing. If MCP-HEADERS sends "authorization", the token is ignored and OAuth probing can incorrectly return 401.
♻️ Proposed fix
- authorization = headers.get("Authorization", None)
+ authorization = next(
+ (value for key, value in headers.items() if key.lower() == "authorization"),
+ None,
+ )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/utils/mcp_oauth_probe.py` around lines 42 - 43, The code currently
retrieves the Authorization token using headers.get("Authorization") which is
case-sensitive; update the lookup in mcp_oauth_probe.py to be case-insensitive
by normalizing header keys (e.g., build a temporary dict with lowercased keys
from headers and then use lowercased "authorization") or check both
"Authorization" and "authorization" before falling back to None so the token is
found regardless of header casing for mcp_headers, headers, authorization, and
mcp_server.name.
Contributor
Author
|
If this PR wants to be merged now I can add @skip-in-library-mode tags to the failing tests |
Contributor
|
@max-svistunov please take a look. TYVM |
Contributor
max-svistunov
left a comment
max-svistunov
left a comment
There was a problem hiding this comment.
MCP auth checks moved to early in the pipeline is much cleaner 👌. I left just two comments to consider.
| except BadRequestError: | ||
| logger.error("Toolgroup %s is not found", toolgroup.identifier) | ||
| continue | ||
| except (AuthenticationError, AuthenticationRequiredError) as e: |
Contributor
There was a problem hiding this comment.
Here we handled auth from llamastack itself as well, not just from MCP Oath. Now the check_mcp_auth only covers the MCP Oauth scenario. So if client.tools.list() now gives AuthenticationError for any problem that is not MCP-related, it now puts forward a 500 error instead of 401.
A simple fix could be: re-add catch for AuthenticationError (to return 401 again) -- the deleted lines 157, 162, and 163:
except AuthenticationError as e:
error_response = UnauthorizedResponse(cause=str(e))
raise HTTPException(**error_response.model_dump()) from e
Contributor
Author
There was a problem hiding this comment.
The only reason I got rid of the llama stack authentication error catching in /tools is because it results in an asymmetrical implementation between tools and the query endpoints. I can not think of a scenario where client.tools.list() would return a 401 that has not already been caught (because only mcp servers are able to be authd against in llamastack)
Contributor
Author
There was a problem hiding this comment.
Actually I might be able to add an exception handler in the handle_known_apistatus_errors() and extend that to be used in /tools wdyt?
Contributor
Author
There was a problem hiding this comment.
Actually I might be able to add an exception handler in the handle_known_apistatus_errors() and extend that to be used in /tools wdyt?
But should I add this if its basically an impossible to reach branch?
Contributor
There was a problem hiding this comment.
Agreed that currently only MCP servers trigger auth in llamastack. OTOH, it might be worth it as defense-in-depth -- if llamastack adds auth for non-MCP reason we would get 500 errors again. It's a low risk though, so up to you 🤷
Contributor
Author
There was a problem hiding this comment.
I ended up not implementing this error catching mainly to reduce size of this PR as well as not add any (as of right now) unneeded complexity. If this becomes an issue we will address it at that time :). However with the toolgroup depreciation large changes will likely begin to emerge that will likely make any work done with this exception catching irrelevant so for the time being I don't think its necessary.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/e2e/features/mcp.feature (1)
67-67: Looser assertion is acceptable but less precise.Changing from an exact JSON body assertion to "contains mcp-oauth" provides flexibility but reduces specificity. This may mask regressions in the response structure.
Consider whether a more targeted assertion (e.g., checking for specific fields while ignoring dynamic values) would better balance flexibility and coverage.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/e2e/features/mcp.feature` at line 67, The step "And The body of the response contains mcp-oauth" is too loose; replace it with a targeted assertion that verifies the response JSON contains the expected keys/fields (e.g., the mcp-oauth object and its required properties) or validate against a small JSON schema instead of a substring match; update the step definition that implements "The body of the response contains mcp-oauth" to parse the response body as JSON and assert presence and types of specific fields (e.g., client_id, access_token, expires_in) so tests remain flexible for dynamic values but still catch structural regressions.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/e2e/features/mcp.feature`:
- Line 67: The step "And The body of the response contains mcp-oauth" is too
loose; replace it with a targeted assertion that verifies the response JSON
contains the expected keys/fields (e.g., the mcp-oauth object and its required
properties) or validate against a small JSON schema instead of a substring
match; update the step definition that implements "The body of the response
contains mcp-oauth" to parse the response body as JSON and assert presence and
types of specific fields (e.g., client_id, access_token, expires_in) so tests
remain flexible for dynamic values but still catch structural regressions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 8a194b1f-023c-4eb7-8bcd-1ecfe9f35d66
📒 Files selected for processing (2)
src/utils/mcp_oauth_probe.pytests/e2e/features/mcp.feature
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Refactored the way that MCP auth is checked. Instead of having it built into utility functions the authentication checking was extracted into a separate function, which is reused across all affected endpoints (
/tools,/query,/streaming_query).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
Summary by CodeRabbit
Bug Fixes
Behavior Changes
Tests