LCORE-1463: Fixed MCP E2E Tests#1560
Conversation
WalkthroughMCP authentication validation is refactored across endpoints and utilities. The Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
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 |
|
Please Review |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/utils/mcp_oauth_probe.py (1)
58-65:⚠️ Potential issue | 🟠 MajorMake authorization header handling case-insensitive.
headers.get("Authorization")misses valid lowercase/mixed-case header names, andstartswith("Bearer ")double-prefixes valid lowercase schemes likebearer token. That can skip the probe or send an invalid probe header for otherwise valid MCP auth inputs.Suggested fix
logger = get_logger(__name__) +def _get_authorization_header(headers: Mapping[str, str]) -> Optional[str]: + """Return the Authorization header value using case-insensitive matching.""" + for header_name, header_value in headers.items(): + if header_name.lower() == "authorization": + return header_value + return None + + async def check_mcp_auth( @@ probes = [] for mcp_server in configuration.mcp_servers: headers = complete_headers.get(mcp_server.name, {}) - auth_header = headers.get("Authorization") + auth_header = _get_authorization_header(headers) if auth_header is not None: authorization = ( auth_header - if auth_header.startswith("Bearer ") + if auth_header.lower().startswith("bearer ") else f"Bearer {auth_header}" )🤖 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 58 - 65, The header lookup and scheme-check in the MCP probe are case-sensitive: change the logic around complete_headers.get(mcp_server.name, {}), auth_header = headers.get("Authorization"), and the startswith("Bearer ") check so header names and the "Bearer" scheme are treated case-insensitively; e.g., normalize header keys (or perform a case-insensitive search) to find any variant like "authorization" or "AUTHORIZATION", normalize auth_header's scheme to lowercase before checking and only add the "Bearer " prefix when the scheme is absent (avoid double-prefixing), and assign the result back to authorization so downstream code uses a correctly formatted bearer header.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@src/utils/mcp_oauth_probe.py`:
- Around line 58-65: The header lookup and scheme-check in the MCP probe are
case-sensitive: change the logic around complete_headers.get(mcp_server.name,
{}), auth_header = headers.get("Authorization"), and the startswith("Bearer ")
check so header names and the "Bearer" scheme are treated case-insensitively;
e.g., normalize header keys (or perform a case-insensitive search) to find any
variant like "authorization" or "AUTHORIZATION", normalize auth_header's scheme
to lowercase before checking and only add the "Bearer " prefix when the scheme
is absent (avoid double-prefixing), and assign the result back to authorization
so downstream code uses a correctly formatted bearer header.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 3bb85e58-d86e-4776-a5bd-7c0caded81ea
📒 Files selected for processing (6)
src/app/endpoints/query.pysrc/app/endpoints/responses.pysrc/app/endpoints/streaming_query.pysrc/app/endpoints/tools.pysrc/utils/mcp_oauth_probe.pytests/e2e/features/mcp.feature
💤 Files with no reviewable changes (1)
- tests/e2e/features/mcp.feature
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
🧰 Additional context used
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Use absolute imports for internal modules:from authentication import get_auth_dependency
Import FastAPI dependencies with:from fastapi import APIRouter, HTTPException, Request, status, Depends
Import Llama Stack client with:from llama_stack_client import AsyncLlamaStackClient
Checkconstants.pyfor shared constants before defining new ones
All modules start with descriptive docstrings explaining purpose
Uselogger = get_logger(__name__)fromlog.pyfor module logging
Type aliases defined at module level for clarity
Use Final[type] as type hint for all constants
All functions require docstrings with brief descriptions
Complete type annotations for parameters and return types in functions
Usetyping_extensions.Selffor model validators in Pydantic models
Use modern union type syntaxstr | intinstead ofUnion[str, int]
UseOptional[Type]for optional type hints
Use snake_case with descriptive, action-oriented function names (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead
Useasync deffor I/O operations and external API calls
HandleAPIConnectionErrorfrom Llama Stack in error handling
Use standard log levels with clear purposes: debug, info, warning, error
All classes require descriptive docstrings explaining purpose
Use PascalCase for class names with standard suffixes: Configuration, Error/Exception, Resolver, Interface
Use ABC for abstract base classes with@abstractmethoddecorators
Use@model_validatorand@field_validatorfor Pydantic model validation
Complete type annotations for all class attributes; use specific types, notAny
Follow Google Python docstring conventions with Parameters, Returns, Raises, and Attributes sections
Files:
src/app/endpoints/query.pysrc/app/endpoints/responses.pysrc/app/endpoints/tools.pysrc/app/endpoints/streaming_query.pysrc/utils/mcp_oauth_probe.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use FastAPI
HTTPExceptionwith appropriate status codes for API endpoints
Files:
src/app/endpoints/query.pysrc/app/endpoints/responses.pysrc/app/endpoints/tools.pysrc/app/endpoints/streaming_query.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Pydantic models extend
ConfigurationBasefor config,BaseModelfor data models
Files:
src/app/endpoints/query.pysrc/app/endpoints/responses.pysrc/app/endpoints/tools.pysrc/app/endpoints/streaming_query.pysrc/utils/mcp_oauth_probe.py
🧠 Learnings (2)
📚 Learning: 2026-01-14T09:37:51.612Z
Learnt from: asimurka
Repo: lightspeed-core/lightspeed-stack PR: 988
File: src/app/endpoints/query.py:319-339
Timestamp: 2026-01-14T09:37:51.612Z
Learning: In the lightspeed-stack repository, when provider_id == "azure", the Azure provider with provider_type "remote::azure" is guaranteed to be present in the providers list. Therefore, avoid defensive StopIteration handling for next() when locating the Azure provider in providers within src/app/endpoints/query.py. This change applies specifically to this file (or nearby provider lookup code) and relies on the invariant that the Azure provider exists; if the invariant could be violated, keep the existing StopIteration handling.
Applied to files:
src/app/endpoints/query.py
📚 Learning: 2026-04-06T20:18:07.852Z
Learnt from: major
Repo: lightspeed-core/lightspeed-stack PR: 1463
File: src/app/endpoints/rlsapi_v1.py:266-271
Timestamp: 2026-04-06T20:18:07.852Z
Learning: In the lightspeed-stack codebase, within `src/app/endpoints/` inference/MCP endpoints, treat `tools: Optional[list[Any]]` in MCP tool definitions as an intentional, consistent typing pattern (used across `query`, `responses`, `streaming_query`, `rlsapi_v1`). Do not raise or suggest this as a typing issue during code review; changing it in isolation could break endpoint typing consistency across the codebase.
Applied to files:
src/app/endpoints/query.pysrc/app/endpoints/responses.pysrc/app/endpoints/tools.pysrc/app/endpoints/streaming_query.py
🔇 Additional comments (5)
src/utils/mcp_oauth_probe.py (1)
23-54: Good centralization of MCP header construction.The updated signature gives
check_mcp_auth()enough context to resolve file-based, Kubernetes, client, OAuth, and propagated headers consistently before probing.src/app/endpoints/query.py (1)
133-136: LGTM.The updated call passes both the extracted bearer token and incoming request headers, so
check_mcp_auth()can build the same complete MCP header set used later in request preparation.src/app/endpoints/responses.py (1)
256-258: LGTM.The unpacking matches the
AuthTuplecontract, and the probe now receives the token plus request headers needed for consistent MCP header resolution.src/app/endpoints/tools.py (1)
140-141: LGTM.This keeps the auth probe aligned with the new
check_mcp_auth()contract while preservingcomplete_mcp_headersfor the later tools-list calls.src/app/endpoints/streaming_query.py (1)
187-191: LGTM.The MCP auth probe now runs after token extraction and receives the request headers, matching the downstream MCP header-building path.
3 participants
Description
Fixes MCP authentication probing by consolidating header building logic inside
check_mcp_auth(). Previously, endpoints inconsistently built MCP headers before calling the probe function, causing authentication failures in some scenarios.Changes:
check_mcp_auth()to accepttokenandrequest_headersparameters and build complete MCP headers internally usingbuild_mcp_headers()Type of change
Tools used to create PR
Related Tickets & Documents
Checklist before requesting a review
Testing
"Check if query endpoint reports error when MCP file-based invalid auth token is passed"
"Check if streaming_query endpoint reports error when MCP file-based invalid auth token is passed"
"Check if query endpoint reports error when MCP kubernetes invalid auth token is passed"
"Check if streaming_query endpoint reports error when MCP kubernetes invalid auth token is passed"
Summary by CodeRabbit
Bug Fixes
Tests