Skip to content

[LEADS-25] Support new MCP header#175

Open
Anxhela21 wants to merge 5 commits intolightspeed-core:mainfrom
Anxhela21:anx/mcp-header
Open

[LEADS-25] Support new MCP header#175
Anxhela21 wants to merge 5 commits intolightspeed-core:mainfrom
Anxhela21:anx/mcp-header

Conversation

@Anxhela21
Copy link
Collaborator

@Anxhela21 Anxhela21 commented Feb 26, 2026

Description

Update new MCP header to be compatible with Lightspeed-Stack API implementation.

Type of change

  • Refactor
  • New feature
  • Bug fix
  • CVE fix
  • Optimization
  • Documentation Update
  • Configuration Update
  • Bump-up service version
  • Bump-up dependent library
  • Bump-up library or tool used for development (does not change the final image)
  • CI configuration change
  • Unit tests improvement

Tools used to create PR

Identify any AI code assistants used in this PR (for transparency and review context)

  • Assisted-by: (e.g., Claude, CodeRabbit, Ollama, etc., N/A if not used)
  • Generated by: (e.g., tool name and version; N/A if not used)

Related Tickets & Documents

  • Related Issue #
  • Closes #

Checklist before requesting a review

  • I have performed a self-review of my code.
  • PR has passed all pre-merge test jobs.
  • If it is a core feature, I have added thorough tests.

Testing

  • Please provide detailed steps to perform tests related to this code change.
  • How were the fix/results from this change verified? Please provide relevant screenshots or results.

Summary by CodeRabbit

  • New Features

    • Configurable per-server authentication headers in API settings (enable, server list, auth type, env var, custom header name).
    • Validation to ensure auth type and required header fields are present.
  • Bug Fixes

    • Client now builds and sends configured per-server headers, warns on missing credentials, and falls back to legacy API key behavior when no MCP headers are available.

@coderabbitai
Copy link
Contributor

coderabbitai bot commented Feb 26, 2026

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

Client authentication header handling was extended to support per-MCP-server headers. New MCP header models and config were added; the client now builds an MCP-HEADERS JSON payload from config + environment (supports bearer, api_key, custom) and falls back to legacy API_KEY Authorization when none resolved.

Changes

Cohort / File(s) Summary
API client logic
src/lightspeed_evaluation/core/api/client.py
Added private _build_mcp_headers(self) -> dict[str, dict[str, str]]; _setup_client now calls it and injects MCP-HEADERS (JSON) assembled per-server from config + env vars. Falls back to legacy API_KEY Authorization if no MCP headers resolved.
System models
src/lightspeed_evaluation/core/models/system.py
Added MCPServerConfig and MCPHeadersConfig models with validation (auth_type ∈ {bearer,api_key,custom}; header_name required when custom); extended APIConfig with optional mcp_headers: Optional[MCPHeadersConfig].
Configuration
config/system.yaml
Introduced api.mcp_headers block (enabled, servers → per-server auth_type, env_var, optional header_name) to configure per-server header sources; legacy API_KEY comment retained.

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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title '[LEADS-25] Support new MCP header' clearly references the main change: adding support for MCP (Model Context Protocol) headers in the API client for Lightspeed-Stack compatibility, which aligns with all three modified files.
Docstring Coverage ✅ Passed Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@Anxhela21 Anxhela21 marked this pull request as ready for review March 2, 2026 04:12
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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-HEADERS is set with the expected JSON shape when API_KEY is present, and absent when API_KEY is 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.

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e395459 and 90473fc.

📒 Files selected for processing (1)
  • src/lightspeed_evaluation/core/api/client.py

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 90473fc and 22ed169.

📒 Files selected for processing (3)
  • config/system.yaml
  • src/lightspeed_evaluation/core/api/client.py
  • src/lightspeed_evaluation/core/models/system.py

Comment on lines +58 to 59
# Legacy authentication (fallback when mcp_headers.enabled is false)
# Authentication via API_KEY environment variable only for MCP server
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

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.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 22ed169 and 9dd70d6.

📒 Files selected for processing (2)
  • src/lightspeed_evaluation/core/api/client.py
  • src/lightspeed_evaluation/core/models/system.py

Comment on lines +160 to +163
auth_type: str = Field(
...,
description="Authentication type: bearer, api_key",
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

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.

Suggested change
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.

Comment on lines +198 to +205
enabled: bool = Field(
default=True,
description="Enable MCP headers functionality",
)
servers: dict[str, MCPServerConfig] = Field(
default_factory=dict,
description="MCP server configurations",
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

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.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

♻️ Duplicate comments (3)
config/system.yaml (1)

58-59: ⚠️ Potential issue | 🟡 Minor

Broaden the fallback comment to match actual runtime conditions.

Line 58 currently implies fallback only when mcp_headers.enabled is false, but fallback also happens when mcp_headers is 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 | 🟡 Minor

Update auth_type description to include custom.

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 | 🟠 Major

Validate enabled=True requires at least one configured MCP server.

Without model-level validation, invalid MCP config can pass and fall through to legacy API_KEY behavior in src/lightspeed_evaluation/core/api/client.py, masking misconfiguration.

🔧 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
As per coding guidelines: "Use custom exceptions from core.system.exceptions module for error handling".
🤖 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 actual MCP-HEADERS payload 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

📥 Commits

Reviewing files that changed from the base of the PR and between 9dd70d6 and 8fb5ef3.

📒 Files selected for processing (3)
  • config/system.yaml
  • src/lightspeed_evaluation/core/api/client.py
  • src/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>
Signed-off-by: Anxhela Coba <acoba@redhat.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant