[LEADS-25] Support new MCP header#175
[LEADS-25] Support new MCP header#175Anxhela21 wants to merge 5 commits intolightspeed-core:mainfrom
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughClient authentication header handling was extended to support per-MCP-server headers. New MCP header models and config were added; the client now builds an Changes
Sequence Diagram(s)sequenceDiagram
actor Client
participant Config as Config
participant Env as Environment
participant Remote as RemoteServer
Client->>Config: read api.mcp_headers
alt mcp_headers.enabled == true
loop for each configured server
Client->>Env: read server.env_var
Env-->>Client: token/key (if present)
Client->>Client: map auth_type -> header name/value
end
Client->>Remote: send request with MCP-HEADERS (JSON)
else no mcp config or no resolved credentials
Client->>Env: read API_KEY
Env-->>Client: api_key (if present)
Client->>Remote: send request with Authorization: Bearer <API_KEY>
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 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 |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/lightspeed_evaluation/core/api/client.py (1)
67-70: Add a regression test for MCP header serialization.Line 67–Line 70 changes a protocol-critical auth format, but the PR includes no verification steps. Please add a unit test asserting
MCP-HEADERSis set with the expected JSON shape whenAPI_KEYis present, and absent whenAPI_KEYis missing.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/lightspeed_evaluation/core/api/client.py` around lines 67 - 70, Add unit tests to verify MCP header serialization: create tests that instantiate the client path which runs the code that builds mcp_headers and calls self.client.headers.update (reference the mcp_headers variable and "MCP-HEADERS" header) and assert that when an API_KEY is provided the client's headers contain "MCP-HEADERS" with json.dumps({"filesystem-tools": {"Authorization": f"Bearer {api_key}"}}) and when API_KEY is absent the "MCP-HEADERS" key is not present; use the client constructor or factory used in the diff to trigger the header logic and assert exact JSON string equality and absence accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@src/lightspeed_evaluation/core/api/client.py`:
- Around line 67-70: Add unit tests to verify MCP header serialization: create
tests that instantiate the client path which runs the code that builds
mcp_headers and calls self.client.headers.update (reference the mcp_headers
variable and "MCP-HEADERS" header) and assert that when an API_KEY is provided
the client's headers contain "MCP-HEADERS" with json.dumps({"filesystem-tools":
{"Authorization": f"Bearer {api_key}"}}) and when API_KEY is absent the
"MCP-HEADERS" key is not present; use the client constructor or factory used in
the diff to trigger the header logic and assert exact JSON string equality and
absence accordingly.
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 `@config/system.yaml`:
- Around line 58-59: Update the comment to reflect the actual fallback
conditions: state that authentication via the API_KEY environment variable is
used for the MCP server whenever header-based auth is unavailable — i.e., when
mcp_headers.enabled is false OR mcp_headers is not configured/usable — rather
than only when mcp_headers.enabled is false; reference the mcp_headers.enabled
flag and API_KEY environment variable so operators understand the broader
runtime behavior.
In `@src/lightspeed_evaluation/core/models/system.py`:
- Around line 160-181: Add a model-level validator on MCPServerConfig (e.g., a
`@model_validator`(mode="after") method named validate_auth_config) that checks if
auth_type == "custom" and header_name is None or empty and raises the custom
exception from core.system.exceptions (use that exception type, not ValueError);
this enforces that header_name must be provided for the "custom" auth_type
during config validation. Ensure you reference the existing auth_type field and
header_name attribute in the validator and raise the core.system.exceptions
custom error with a clear message when validation fails.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
config/system.yamlsrc/lightspeed_evaluation/core/api/client.pysrc/lightspeed_evaluation/core/models/system.py
| # Legacy authentication (fallback when mcp_headers.enabled is false) | ||
| # Authentication via API_KEY environment variable only for MCP server |
There was a problem hiding this comment.
Fallback comment is narrower than runtime behavior.
The comment says fallback happens only when mcp_headers.enabled is false, but the client also falls back in other non-configured cases. Please align the comment with actual behavior to avoid operator confusion.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@config/system.yaml` around lines 58 - 59, Update the comment to reflect the
actual fallback conditions: state that authentication via the API_KEY
environment variable is used for the MCP server whenever header-based auth is
unavailable — i.e., when mcp_headers.enabled is false OR mcp_headers is not
configured/usable — rather than only when mcp_headers.enabled is false;
reference the mcp_headers.enabled flag and API_KEY environment variable so
operators understand the broader runtime behavior.
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/lightspeed_evaluation/core/models/system.py`:
- Around line 160-163: The auth_type Field description is outdated — it lists
"bearer, api_key" but validation also accepts "custom"; update the description
for the auth_type Field (the Field declaration named auth_type) to include
"custom" (e.g., "Authentication type: bearer, api_key, custom") so the docstring
matches allowed values.
- Around line 198-205: Add validation to the Pydantic model that ensures when
the enabled field is True the servers dict is non-empty: implement a
`@root_validator` (or `@validator` for servers with values and access to enabled) in
the same model that checks if enabled and not servers and raises a
ValidationError (e.g., ValueError) with a clear message; reference the enabled
and servers fields in the validator so invalid MCP configuration fails fast and
cannot fall through to legacy auth in client.py.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/lightspeed_evaluation/core/api/client.pysrc/lightspeed_evaluation/core/models/system.py
| auth_type: str = Field( | ||
| ..., | ||
| description="Authentication type: bearer, api_key", | ||
| ) |
There was a problem hiding this comment.
Update auth_type description to match supported values.
Line 162 omits custom, even though it is allowed by validation.
✏️ Proposed fix
auth_type: str = Field(
...,
- description="Authentication type: bearer, api_key",
+ description="Authentication type: bearer, api_key, custom",
)📝 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.
| auth_type: str = Field( | |
| ..., | |
| description="Authentication type: bearer, api_key", | |
| ) | |
| auth_type: str = Field( | |
| ..., | |
| description="Authentication type: bearer, api_key, custom", | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/lightspeed_evaluation/core/models/system.py` around lines 160 - 163, The
auth_type Field description is outdated — it lists "bearer, api_key" but
validation also accepts "custom"; update the description for the auth_type Field
(the Field declaration named auth_type) to include "custom" (e.g.,
"Authentication type: bearer, api_key, custom") so the docstring matches allowed
values.
| enabled: bool = Field( | ||
| default=True, | ||
| description="Enable MCP headers functionality", | ||
| ) | ||
| servers: dict[str, MCPServerConfig] = Field( | ||
| default_factory=dict, | ||
| description="MCP server configurations", | ||
| ) |
There was a problem hiding this comment.
Require non-empty servers when MCP headers are enabled.
Line 198 currently allows enabled=True with an empty servers map. That passes validation and can fall through to legacy auth in src/lightspeed_evaluation/core/api/client.py (Line 134), masking a broken MCP configuration.
🔧 Proposed fix
class MCPHeadersConfig(BaseModel):
"""Configuration for MCP headers functionality."""
@@
servers: dict[str, MCPServerConfig] = Field(
default_factory=dict,
description="MCP server configurations",
)
+
+ `@model_validator`(mode="after")
+ def validate_enabled_servers(self) -> "MCPHeadersConfig":
+ """Ensure enabled MCP headers have at least one configured server."""
+ if self.enabled and not self.servers:
+ raise ConfigurationError(
+ "mcp_headers.enabled is true, but no MCP servers are configured."
+ )
+ return self🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/lightspeed_evaluation/core/models/system.py` around lines 198 - 205, Add
validation to the Pydantic model that ensures when the enabled field is True the
servers dict is non-empty: implement a `@root_validator` (or `@validator` for
servers with values and access to enabled) in the same model that checks if
enabled and not servers and raises a ValidationError (e.g., ValueError) with a
clear message; reference the enabled and servers fields in the validator so
invalid MCP configuration fails fast and cannot fall through to legacy auth in
client.py.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (3)
config/system.yaml (1)
58-59:⚠️ Potential issue | 🟡 MinorBroaden the fallback comment to match actual runtime conditions.
Line 58 currently implies fallback only when
mcp_headers.enabledisfalse, but fallback also happens whenmcp_headersis missing or has no configured servers.✏️ Suggested update
- # Legacy authentication (fallback when mcp_headers.enabled is false) - # Authentication via API_KEY environment variable only for MCP server + # Legacy authentication fallback when MCP headers are unavailable + # (e.g., mcp_headers disabled, missing, or no servers configured): + # uses API_KEY environment variable for MCP server authentication🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@config/system.yaml` around lines 58 - 59, Update the existing comment that reads "Legacy authentication (fallback when mcp_headers.enabled is false)" to accurately list all runtime conditions that trigger the fallback: when mcp_headers.enabled is false, when the mcp_headers block is absent, or when mcp_headers has no configured servers; modify the comment near the "Legacy authentication" / "Authentication via API_KEY environment variable only for MCP server" lines to reflect these three conditions so it matches actual behavior at runtime.src/lightspeed_evaluation/core/models/system.py (2)
168-171:⚠️ Potential issue | 🟡 MinorUpdate
auth_typedescription to includecustom.The field description omits
custom, even though validation explicitly supports it.✏️ Proposed text fix
auth_type: str = Field( ..., - description="Authentication type: bearer, api_key", + description="Authentication type: bearer, api_key, custom", )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/lightspeed_evaluation/core/models/system.py` around lines 168 - 171, The Field declaration for auth_type currently lists allowed values as "bearer, api_key" but validation also accepts "custom"; update the description on the auth_type Field (the Field for auth_type in the model) to include "custom" so the docstring matches validation (e.g., "Authentication type: bearer, api_key, custom").
206-213:⚠️ Potential issue | 🟠 MajorValidate
enabled=Truerequires at least one configured MCP server.Without model-level validation, invalid MCP config can pass and fall through to legacy
API_KEYbehavior insrc/lightspeed_evaluation/core/api/client.py, masking misconfiguration.As per coding guidelines: "Use custom exceptions from core.system.exceptions module for error handling".🔧 Proposed validator
class MCPHeadersConfig(BaseModel): @@ servers: dict[str, MCPServerConfig] = Field( default_factory=dict, description="MCP server configurations", ) + + `@model_validator`(mode="after") + def validate_enabled_servers(self) -> "MCPHeadersConfig": + """Ensure enabled MCP headers include at least one server.""" + if self.enabled and not self.servers: + raise ConfigurationError( + "mcp_headers.enabled is true, but no MCP servers are configured." + ) + return self🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/lightspeed_evaluation/core/models/system.py` around lines 206 - 213, Add a model-level validator on the Pydantic model that defines the enabled and servers fields (the class containing enabled: bool and servers: dict[str, MCPServerConfig]) that checks: if enabled is True then servers must not be empty; if the check fails raise the appropriate custom exception imported from core.system.exceptions. Implement this as a `@root_validator` or `@model_validator` on that model, import the specific exception from core.system.exceptions, and ensure the error message clearly indicates "MCP enabled but no servers configured" so misconfiguration is caught before falling back to legacy API_KEY behavior.
🧹 Nitpick comments (1)
src/lightspeed_evaluation/core/api/client.py (1)
94-99: Add assertions for actualMCP-HEADERSpayload in unit tests.Current coverage only checks that
headers.update()was called, which won’t catch malformed MCP payloads or wrong fallback behavior.🧪 Suggested test strengthening (tests/unit/core/api/test_client.py)
+import json + def test_setup_client_with_api_key( self, basic_api_config_streaming_endpoint: APIConfig, mocker: MockerFixture ) -> None: @@ APIClient(basic_api_config_streaming_endpoint) - # Verify headers were updated (should include Authorization header) - assert mock_client.headers.update.call_count >= 1 + # Verify concrete header content, not only call count + updated_headers = [ + call.args[0] + for call in mock_client.headers.update.call_args_list + if call.args and isinstance(call.args[0], dict) + ] + assert {"Content-Type": "application/json"} in updated_headers + assert any("MCP-HEADERS" in h for h in updated_headers) + + mcp_header_payload = next(h["MCP-HEADERS"] for h in updated_headers if "MCP-HEADERS" in h) + assert json.loads(mcp_header_payload) == { + "filesystem-tools": {"Authorization": "Bearer test_secret_key"} + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/lightspeed_evaluation/core/api/client.py` around lines 94 - 99, The tests currently only assert that headers.update() was called; update or add unit tests for the client setup to call self._build_mcp_headers() and assert the actual payload stored under "MCP-HEADERS" is valid JSON and matches the expected structure/content (check keys/values and fallback values), e.g. verify the value of self.client.headers["MCP-HEADERS"] after initialization (or inspect the argument passed to headers.update via the mock) and validate json.loads(...) equals the expected mcp_headers_dict and that fallbacks are used when _build_mcp_headers() returns partial/missing fields.
🤖 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/lightspeed_evaluation/core/models/system.py`:
- Around line 252-254: The mcp_headers Field in the APIConfig class is missing
its closing parenthesis and trailing comma which causes a syntax error; locate
the mcp_headers declaration (symbol: mcp_headers) in APIConfig and close the
Field call by adding the missing ")" after the description string and a comma
before the next field, ensuring the line reads as a complete Field(...)
expression.
---
Duplicate comments:
In `@config/system.yaml`:
- Around line 58-59: Update the existing comment that reads "Legacy
authentication (fallback when mcp_headers.enabled is false)" to accurately list
all runtime conditions that trigger the fallback: when mcp_headers.enabled is
false, when the mcp_headers block is absent, or when mcp_headers has no
configured servers; modify the comment near the "Legacy authentication" /
"Authentication via API_KEY environment variable only for MCP server" lines to
reflect these three conditions so it matches actual behavior at runtime.
In `@src/lightspeed_evaluation/core/models/system.py`:
- Around line 168-171: The Field declaration for auth_type currently lists
allowed values as "bearer, api_key" but validation also accepts "custom"; update
the description on the auth_type Field (the Field for auth_type in the model) to
include "custom" so the docstring matches validation (e.g., "Authentication
type: bearer, api_key, custom").
- Around line 206-213: Add a model-level validator on the Pydantic model that
defines the enabled and servers fields (the class containing enabled: bool and
servers: dict[str, MCPServerConfig]) that checks: if enabled is True then
servers must not be empty; if the check fails raise the appropriate custom
exception imported from core.system.exceptions. Implement this as a
`@root_validator` or `@model_validator` on that model, import the specific exception
from core.system.exceptions, and ensure the error message clearly indicates "MCP
enabled but no servers configured" so misconfiguration is caught before falling
back to legacy API_KEY behavior.
---
Nitpick comments:
In `@src/lightspeed_evaluation/core/api/client.py`:
- Around line 94-99: The tests currently only assert that headers.update() was
called; update or add unit tests for the client setup to call
self._build_mcp_headers() and assert the actual payload stored under
"MCP-HEADERS" is valid JSON and matches the expected structure/content (check
keys/values and fallback values), e.g. verify the value of
self.client.headers["MCP-HEADERS"] after initialization (or inspect the argument
passed to headers.update via the mock) and validate json.loads(...) equals the
expected mcp_headers_dict and that fallbacks are used when _build_mcp_headers()
returns partial/missing fields.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 6f9aab65-810c-4c88-ab8c-23a014a30125
📒 Files selected for processing (3)
config/system.yamlsrc/lightspeed_evaluation/core/api/client.pysrc/lightspeed_evaluation/core/models/system.py
Signed-off-by: Anxhela Coba <acoba@redhat.com>
Signed-off-by: Anxhela Coba <acoba@redhat.com>
Signed-off-by: Anxhela Coba <acoba@redhat.com>
Signed-off-by: Anxhela Coba <acoba@redhat.com>
10360c1 to
99f8a5b
Compare
Signed-off-by: Anxhela Coba <acoba@redhat.com>
| provider: "openai" # LLM provider for queries | ||
| model: "gpt-4o-mini" # Model to use for queries | ||
| no_tools: null # Whether to bypass tools and MCP servers (optional) | ||
| no_tools: false # Whether to bypass tools and MCP servers (optional) |
There was a problem hiding this comment.
Please keep it is as null..
| enabled: true # Enable MCP headers functionality | ||
| servers: # MCP server configurations | ||
| filesystem-tools: | ||
| auth_type: bearer # Authentication type: only bearer is supported |
There was a problem hiding this comment.
is there a plan to support anything else in future ?
|
|
||
| #### Lightspeed Stack API Compatibility | ||
|
|
||
| **Important Note for lightspeed-stack API users**: To use the MCP headers functionality with the lightspeed-stack API, you need to modify the `llama_stack_api/openai_responses.py` file in your lightspeed-stack installation: |
There was a problem hiding this comment.
Any idea, when will this be fixed in lls ?
Also how did it get merged in lightspeed-stack, if it is not supported yet ?
| # Use API_KEY environment variable for authentication | ||
| api_key = os.getenv("API_KEY") | ||
| if api_key and self.client: | ||
| self.client.headers.update({"Authorization": f"Bearer {api_key}"}) |
There was a problem hiding this comment.
i would suggest retain this but when mcp header is disabled or not present. for backward compatibility
| # Fallback to legacy API_KEY behavior for backward compatibility | ||
| api_key = os.getenv("API_KEY") | ||
| if api_key: | ||
| return {"filesystem-tools": {"Authorization": f"Bearer {api_key}"}} | ||
|
|
||
| return {} |
There was a problem hiding this comment.
We are assuming that the server is filesystem-tools when mcp header is not given.. This is not true. Currently we don't have to set the server name.. I would suggest to simply set empty {} and handle it separately the way it was done earlier (previous comment)
| # Fallback to legacy API_KEY behavior for backward compatibility | |
| api_key = os.getenv("API_KEY") | |
| if api_key: | |
| return {"filesystem-tools": {"Authorization": f"Bearer {api_key}"}} | |
| return {} | |
| return {} |
| env_var: str = Field( | ||
| ..., | ||
| min_length=1, | ||
| description="Environment variable containing the token/key", | ||
| ) |
There was a problem hiding this comment.
Please add validation that env variable is set when mcp header is enabled..
Description
Update new MCP header to be compatible with Lightspeed-Stack API implementation.
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