LCORE-1368: Enable MCP Headers in /tools#1216
LCORE-1368: Enable MCP Headers in /tools#1216jrobertboos wants to merge 1 commit intolightspeed-core:mainfrom
/tools#1216Conversation
…on and additional request parameters
WalkthroughThe changes add MCP headers support to the tools endpoint by injecting McpHeaders via dependency, extracting per-toolgroup headers from the MCP context, and propagating them to client.tools.list calls along with separated authorization handling. Tests are updated to match the new endpoint signature and verify header/query argument passing. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/unit/app/endpoints/test_tools.py (1)
192-225: No test covers the actual feature: MCP headers (includingAuthorization) passed to toolgroup callsAll tests pass
{}asmcp_headers, so the new code path at lines 91–97 oftools.py— extracting per-toolgroup headers, separatingAuthorization, and forwarding both — is never exercised. At minimum, a test like the following is needed:`@pytest.mark.asyncio` async def test_tools_endpoint_with_mcp_headers( mocker: MockerFixture, mock_configuration: Configuration, mock_tools_response: list[MockType], ) -> None: """Test that per-toolgroup MCP headers and Authorization are propagated correctly.""" ... mcp_headers = { "filesystem-tools": { "Authorization": "Bearer secret-token", "X-Custom-Header": "custom-value", } } response = await tools.tools_endpoint_handler.__wrapped__( mock_request, mock_auth, mcp_headers ) mock_client.tools.list.assert_any_call( toolgroup_id="filesystem-tools", extra_headers={"X-Custom-Header": "custom-value"}, # Authorization stripped extra_query={"authorization": "Bearer secret-token"}, )This is the primary new behaviour introduced by the PR and needs direct coverage.
Would you like me to generate the full test, including a case for a toolgroup absent from
mcp_headers(empty headers) and one verifying that only the matching toolgroup receives its headers?🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/app/endpoints/test_tools.py` around lines 192 - 225, Add a unit test that exercises the MCP header propagation path in tools.tools_endpoint_handler: call tools.tools_endpoint_handler.__wrapped__(mock_request, mock_auth, mcp_headers) where mcp_headers includes an entry for "filesystem-tools" with Authorization and another custom header; assert mock_client.tools.list was called for filesystem-tools with extra_headers containing only the custom header (Authorization removed) and extra_query containing authorization set to the token, and also include a case where a toolgroup not present in mcp_headers yields empty extra_headers/extra_query; reference tools.tools_endpoint_handler and the header-extraction code paths around the new logic in tools.py to locate where to invoke and assert behavior.src/app/endpoints/tools.py (2)
91-92:headers.pop()mutates the dict stored insidemcp_headers— guideline violation
mcp_headers.get(toolgroup.identifier, {})returns a direct reference to the value dict in themcp_headersmapping (not a copy). Calling.pop("Authorization", None)on it mutates themcp_headersdata structure in place. Any post-loop access tomcp_headers(e.g., debug logging, future retry logic) would see a partially stripped state.♻️ Proposed fix — copy before mutating
- headers = mcp_headers.get(toolgroup.identifier, {}) - authorization = headers.pop("Authorization", None) + all_headers = mcp_headers.get(toolgroup.identifier, {}) + authorization = all_headers.get("Authorization") + headers = {k: v for k, v in all_headers.items() if k != "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 91 - 92, The code mutates the dict stored in mcp_headers by calling headers.pop(...); instead, make a shallow copy of the headers dict returned by mcp_headers.get(toolgroup.identifier, {}) before modifying it (e.g., use dict(...) or .copy()), then pop "Authorization" from that copy so mcp_headers remains unchanged; update the logic around mcp_headers, headers, and authorization in the same scope (referencing mcp_headers and toolgroup.identifier) to use the copied dict.
47-47: PreferAnnotatedstyle for dependency injection to match the rest of the signature
authon line 46 uses the modernAnnotated[T, Depends(...)]form;mcp_headersuses the legacy default-value style, creating an inconsistency in the same function signature.♻️ Proposed refactor
- mcp_headers: McpHeaders = Depends(mcp_headers_dependency), + mcp_headers: Annotated[McpHeaders, Depends(mcp_headers_dependency)],🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/app/endpoints/tools.py` at line 47, The mcp_headers parameter uses the legacy default-value dependency style instead of the modern Annotated form used by auth; update the function signature to use Annotated[McpHeaders, Depends(mcp_headers_dependency)] for mcp_headers (replace the current mcp_headers: McpHeaders = Depends(mcp_headers_dependency) pattern), ensure Annotated is imported (from typing or typing_extensions as your project uses), and keep the rest of the signature (including auth) unchanged so that mcp_headers, mcp_headers_dependency, Annotated, Depends and McpHeaders are consistently used.
🤖 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/app/endpoints/tools.py`:
- Around line 91-92: The code mutates the dict stored in mcp_headers by calling
headers.pop(...); instead, make a shallow copy of the headers dict returned by
mcp_headers.get(toolgroup.identifier, {}) before modifying it (e.g., use
dict(...) or .copy()), then pop "Authorization" from that copy so mcp_headers
remains unchanged; update the logic around mcp_headers, headers, and
authorization in the same scope (referencing mcp_headers and
toolgroup.identifier) to use the copied dict.
- Line 47: The mcp_headers parameter uses the legacy default-value dependency
style instead of the modern Annotated form used by auth; update the function
signature to use Annotated[McpHeaders, Depends(mcp_headers_dependency)] for
mcp_headers (replace the current mcp_headers: McpHeaders =
Depends(mcp_headers_dependency) pattern), ensure Annotated is imported (from
typing or typing_extensions as your project uses), and keep the rest of the
signature (including auth) unchanged so that mcp_headers,
mcp_headers_dependency, Annotated, Depends and McpHeaders are consistently used.
In `@tests/unit/app/endpoints/test_tools.py`:
- Around line 192-225: Add a unit test that exercises the MCP header propagation
path in tools.tools_endpoint_handler: call
tools.tools_endpoint_handler.__wrapped__(mock_request, mock_auth, mcp_headers)
where mcp_headers includes an entry for "filesystem-tools" with Authorization
and another custom header; assert mock_client.tools.list was called for
filesystem-tools with extra_headers containing only the custom header
(Authorization removed) and extra_query containing authorization set to the
token, and also include a case where a toolgroup not present in mcp_headers
yields empty extra_headers/extra_query; reference tools.tools_endpoint_handler
and the header-extraction code paths around the new logic in tools.py to locate
where to invoke and assert behavior.
| tools_response = await client.tools.list( | ||
| toolgroup_id=toolgroup.identifier, | ||
| extra_headers=headers, | ||
| extra_query={"authorization": authorization}, |
There was a problem hiding this comment.
I am getting slightly weird behavior with this implementation ...
It works for authorizing the MCP server however once I have sent a valid authorization token it does not need the token again. Basically if I access /tools with MCP-HEADERS set it will successfully authorize however if I access /tools again without passing MCP-HEADERS it will still successfully authorize and not return a 401. This behavior persists until I restart the llama-stack server.
Im guessing that this is because the tool list is either cached or the authorization header is persisted but I am not sure just yet.
tisnik
left a comment
tisnik
left a comment
There was a problem hiding this comment.
Does not it break API compatibility. Is the OpenAPI spec affected by this change?
|
I don't think it breaks API compatibility, because the headers are optional... when looking at the query/streaming_query openapi nothing was added regarding the MCP headers. I could be missing something but I don't think so. |
2 participants
Description
Adds the ability to pass
MCP-HEADERSto the MCP servers accessed by the/toolsendpoint.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
Release Notes