feat: support passing token to mcp in header - #1397
Conversation
✅ Deploy Preview for holmes-docs ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
WalkthroughThreads an optional Changes
Sequence Diagram(s)sequenceDiagram
participant Client as HTTP Client
participant Server as Server (server.py)
participant Investigation as Investigation (holmes/core/investigation.py)
participant LLMCaller as LLM/Tool Caller (holmes/core/tool_calling_llm.py)
participant ToolInvoke as ToolInvokeContext (holmes/core/tools.py)
participant MCPToolset as Remote MCP Toolset (holmes/plugins/toolsets/mcp/toolset_mcp.py)
Client->>Server: POST /investigate (with headers)
Server->>Server: extract_passthrough_headers(request) -> request_context
Server->>Investigation: investigate_issues(..., request_context)
Investigation->>LLMCaller: prompt_call(..., request_context)
LLMCaller->>LLMCaller: process_tool_decisions(..., request_context)
LLMCaller->>ToolInvoke: _directly_invoke_tool_call(..., request_context)
ToolInvoke->>MCPToolset: get_initialized_mcp_session(request_context)
MCPToolset->>MCPToolset: _render_headers(request_context) // render {{ request_context }} & {{ env.VAR }}
MCPToolset->>MCPToolset: _invoke_async(..., request_context) with rendered headers
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/test_approval_workflow.py (1)
135-142: Silence Ruff ARG001/ARG005 for unused mock args.
These new params are required for signature compatibility but unused in tests.🧹 Suggested lint-safe tweak
- def mock_invoke_tool( + def mock_invoke_tool( tool_name: str, tool_params: dict, user_approved: bool, tool_call_id: str, tool_number: Optional[int] = None, - request_context: Optional[dict] = None, - ) -> StructuredToolResult: + request_context: Optional[dict] = None, + ) -> StructuredToolResult: # noqa: ARG001 @@ - ai.process_tool_decisions = MagicMock( - side_effect=lambda messages, tool_decisions, request_context=None: ( + ai.process_tool_decisions = MagicMock( + side_effect=lambda messages, tool_decisions, request_context=None: ( # noqa: ARG005 messages @@ - ai.process_tool_decisions = MagicMock( - side_effect=lambda messages, tool_decisions, request_context=None: ( + ai.process_tool_decisions = MagicMock( + side_effect=lambda messages, tool_decisions, request_context=None: ( # noqa: ARG005 messagesAlso applies to: 268-270, 415-417
server.py (1)
254-303: Merge request bodycontextfield into request context for all non-investigate endpoints.All four request classes (
WorkloadHealthRequest,WorkloadHealthChatRequest,IssueChatRequest,ChatRequest) define acontextfield, but the endpoints extract only passthrough headers and never merge the request body's context. This silently drops client-provided context data.The issue exists in all four endpoints:
/api/workload_health_check(line 297)/api/workload_health_chat(line 338)/api/issue_chat(line 370)/api/chat(line 432)✅ Suggested fix (apply to each endpoint)
- request_context = extract_passthrough_headers(http_request) + request_context = extract_passthrough_headers(http_request) + if request.context: + request_context.update(request.context)Replace
requestwith the actual parameter name for each endpoint (issue_chat_request,chat_request, etc.).
🤖 Fix all issues with AI agents
In `@design/aks-multi-cluster-cross-cluster-investigation.md`:
- Around line 117-123: Replace the JWT-looking example values in the HTTP
request sample (the Authorization and X-Azure-Token header values) with explicit
placeholders like <HOLMES_SERVER_TOKEN> and <AZURE_ACCESS_TOKEN> (and similarly
update the other occurrences noted at lines 191-197 and 242-246); edit the
request block that contains "Authorization: Bearer <holmes-server-token>" and
"X-Azure-Token: eyJ0eXAiOiJKV1QiLCJhbG..." to use non-secret placeholders to
avoid secret-scanner triggering.
- Line 19: The markdown headings and fenced code blocks in the document (e.g.,
the heading "Headers Pass-through with Template-based Configuration" and other
similar headings noted) do not follow markdownlint rules; update each heading to
use proper Markdown heading syntax (consistent `#/`## levels) and add explicit
fenced-code languages to all code blocks (examples: ```http, ```json, ```yaml,
```bash, ```python) so that linting and syntax highlighting work correctly;
apply the same fixes to the other occurrences referenced in the review (the
similar headings/code fences elsewhere in the doc).
In `@holmes/core/tools.py`:
- Around line 165-175: The model_dump override currently accesses
data["request_context"] directly; to satisfy RUF019 replace those direct index
accesses with dict.get to avoid KeyError warnings: use request_ctx =
data.get("request_context") and then if request_ctx: set data["request_context"]
= {k: "***REDACTED***" for k in request_ctx.keys()} (keep the same behavior and
return data) — update the model_dump function to use data.get("request_context")
instead of data["request_context"].
In `@holmes/plugins/toolsets/mcp/toolset_mcp.py`:
- Around line 125-131: The request_context is being set on the toolset before
acquiring the per-server lock causing races; move the call to
self.toolset.set_request_context(context.request_context) inside the with lock:
block (near where get_server_lock(...) is used) so the assignment is serialized
with the lock, then wrap the asyncio.run(self._invoke_async(params)) call in a
try/finally and clear the request context in the finally (e.g., call
self.toolset.set_request_context(None) or a provided clear method) to ensure
_render_headers() never reads a stale context after the request completes.
- Around line 231-310: In _render_template, the current re.sub calls use raw
replacement strings which can be misinterpreted (backrefs, escapes); update both
re.sub calls (the one that replaces request_context header patterns and the one
that replaces env variable patterns) to pass a replacement function/lambda that
returns the replacement string (ensuring it's a str) so replacements are treated
literally and backslashes/backreferences aren't processed.
In `@server.py`:
- Around line 398-414: extract_passthrough_headers currently title-cases header
names (e.g., X-Api-Key) but _render_template performs case-sensitive lookups
(e.g., X-API-Key) causing misses; make header handling case-insensitive by
normalizing keys to a single canonical form (recommend lowercase) in
extract_passthrough_headers (produce headers dict with lowercase keys) and
update _render_template to normalize any lookup key to lowercase (or wrap
request_context.headers in a case-insensitive mapping) so templates can use any
case and still resolve values; change references in extract_passthrough_headers
and _render_template accordingly.
In `@tests/test_mcp_toolset.py`:
- Around line 1322-1326: The test currently assigns rendered =
mock_toolset._render_headers() but never uses it (F841); update the assertion to
also validate rendered so the value is used and the test is stronger — for
example, assert that rendered contains the expected result or error indicator
returned by mock_toolset._render_headers() (in addition to checking
caplog.text), referencing the call to mock_toolset._render_headers(), the
rendered variable, and the existing log check for "Header key 'X-Nonexistent'
not found in request_context".
- Around line 142-170: Replace the unused fixture parameter usage in the test
functions by applying pytest.mark.usefixtures("suppress_migration_warnings") to
the test functions instead of accepting suppress_migration_warnings as a
parameter; specifically, update test_toolset_returns_configured_extra_headers
and test_toolset_without_extra_headers_returns_none (and other tests mentioned)
to remove the suppress_migration_warnings argument and add
`@pytest.mark.usefixtures`("suppress_migration_warnings") above each test so Ruff
ARG002 is resolved while preserving the fixture's side effects.
gossion
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@tests/test_mcp_toolset.py`:
- Line 10: Remove the unused import of LLM from the test file: delete the line
importing "LLM" from "holmes.core.llm" so that the test module no longer has an
unused symbol; ensure no other references to LLM exist in
tests/test_mcp_toolset.py and run tests to confirm.
♻️ Duplicate comments (4)
tests/test_mcp_toolset.py (2)
1140-1143: Apply@pytest.mark.usefixturesat class level to fix ARG002 warnings.All tests in
TestHeaderTemplateRenderingusesuppress_migration_warningsfor side effects only. Apply the decorator at class level and remove the parameter from each test method to satisfy Ruff's unused argument check.Suggested fix
+@pytest.mark.usefixtures("suppress_migration_warnings") class TestHeaderTemplateRendering: """Test header template rendering functionality""" - def test_render_hardcoded_headers(self, suppress_migration_warnings): + def test_render_hardcoded_headers(self): """Test rendering of hardcoded header values"""Remove the
suppress_migration_warningsparameter from all test methods in this class.
1322-1326: Userenderedvariable to strengthen the test and fix F841.The
renderedvariable is assigned but never used. Assert on its value to strengthen the test and resolve the Ruff warning.Suggested fix
mock_toolset.set_request_context({"headers": {"X-Other": "value"}}) rendered = mock_toolset._render_headers() + # The header with missing key should still be in the result with unrendered template + assert rendered is not None + assert "X-Missing" in rendered assert "Header key 'X-Nonexistent' not found in request_context" in caplog.textholmes/plugins/toolsets/mcp/toolset_mcp.py (2)
124-129: Moverequest_contextassignment inside the lock to prevent race conditions.The
_request_contextis instance state shared across concurrent requests. Setting it outside the lock allows multiple requests to overwrite each other's context before acquiring the per-server lock, causing wrong headers to be rendered. The assignment must be serialized inside the lock, with cleanup in a finally block.🔒 Suggested fix
- # Set request_context in toolset for header template rendering - self.toolset.set_request_context(context.request_context) - - with lock: - return asyncio.run(self._invoke_async(params)) + with lock: + # Set request_context in toolset for header template rendering + self.toolset.set_request_context(context.request_context) + try: + return asyncio.run(self._invoke_async(params)) + finally: + self.toolset.set_request_context(None)
286-287: Use lambda replacements inre.subto prevent data corruption.
re.sub()interprets backslashes and backreference patterns (\1,\n) in the replacement string. Header values from request context or environment variables can contain these characters, causing corruption or errors. Use a lambda to ensure literal replacement.Suggested fix
- result = re.sub(pattern_regex, replacement, result) + result = re.sub(pattern_regex, lambda _: replacement, result)Apply to both occurrences at lines 287 and 300.
Also applies to: 299-300
🧹 Nitpick comments (1)
holmes/plugins/toolsets/mcp/toolset_mcp.py (1)
48-55: Add backward compatibility migration forheaders→extra_headersrename.Per coding guidelines, maintain backward compatibility when renaming config fields. Users with existing
headersconfiguration will silently lose their settings. Add a@model_validatorto migrate the old field name and log a deprecation warning.Suggested approach
class MCPConfig(BaseModel): model_config = {"extra": "allow"} url: AnyUrl mode: MCPMode = MCPMode.SSE extra_headers: Optional[Dict[str, str]] = None `@model_validator`(mode="after") def migrate_headers(self) -> "MCPConfig": # Access extra fields via model_extra if hasattr(self, "model_extra") and self.model_extra and "headers" in self.model_extra: if self.extra_headers is None: self.extra_headers = self.model_extra["headers"] logging.warning( "MCPConfig: 'headers' field is deprecated, use 'extra_headers' instead." ) return selfApply similar migration to
StdioMCPConfig. Based on learnings.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@holmes/plugins/toolsets/mcp/toolset_mcp.py`:
- Around line 248-260: The bare except in the header-rendering loop should be
narrowed to the actual exceptions raised by _render_template to satisfy Ruff
BLE001: catch re.error, AttributeError, and TypeError instead of Exception (or
alternatively add a justified "# noqa: BLE001" comment); update the except
clause in the loop that iterates self._mcp_config.extra_headers in
toolset_mcp.py (referencing self._mcp_config, _render_template, and
final_headers) to catch re.error, AttributeError, and TypeError and keep the
existing logging inside the handler.
In `@server.py`:
- Around line 197-211: The current merge in investigate_issues lets
investigate_request.context override filtered passthrough headers by calling
request_context.update(investigate_request.context); change this to ignore any
"headers" key (or re-filter that nested headers map against BLOCKED_HEADERS)
before merging so body-supplied headers cannot bypass protections—modify the
investigate_issues handler to strip or validate
investigate_request.context.get("headers") (or remove the "headers" key) prior
to request_context.update(...) and ensure any remaining keys are safe to merge.
♻️ Duplicate comments (2)
server.py (1)
224-241: Same header-override risk here as in/api/investigate.holmes/plugins/toolsets/mcp/toolset_mcp.py (1)
277-315: Use literal replacements inre.subto avoid backslash/backref corruption.Direct replacements interpret
\1,\n, etc. in header/env values. Use a replacement function to keep values literal.🐛 Proposed fix
- result = re.sub(pattern_regex, matching_value, result) + result = re.sub(pattern_regex, lambda _m: matching_value, result) ... - result = re.sub(pattern_regex, replacement, result) + result = re.sub(pattern_regex, lambda _m: replacement, result)
3b50ed6 to
45118c1
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@tests/test_mcp_toolset.py`:
- Around line 1414-1437: Rename the unused callback parameters in
capture_sse_client_call by prefixing them with underscores (e.g., change url and
sse_read_timeout to _url and _sse_read_timeout) to satisfy Ruff ARG001; also
avoid assigning the session variable in the first test block by changing the
async context manager usage from "as session:" to "as _:" when calling
get_initialized_mcp_session in the run_test coroutine, and for the second
occurrence where mcp_tool._invoke_async is invoked without a context manager,
remove or rename any unused local variable named session to _ (or eliminate the
assignment) so no unused "session" identifier remains.
♻️ Duplicate comments (2)
holmes/plugins/toolsets/mcp/toolset_mcp.py (2)
259-271: Narrow the broadexcept Exceptionin header rendering (BLE001).Catching only expected exceptions keeps behavior the same while satisfying Ruff.
♻️ Suggested fix
- except Exception as e: + except (re.error, AttributeError, TypeError) as e: logging.warning( f"MCP toolset '{self.name}': Failed to render header template " f"'{header_name}': {e}" )
288-326: Use literal replacements inre.subto avoid backreference surprises.Header/env values containing backslashes or
\1-like sequences can be corrupted by direct replacement strings.♻️ Suggested fix
- result = re.sub(pattern_regex, matching_value, result) + result = re.sub( + pattern_regex, lambda _m: str(matching_value), result + ) @@ - result = re.sub(pattern_regex, replacement, result) + result = re.sub(pattern_regex, lambda _m: replacement, result)
🧹 Nitpick comments (1)
holmes/core/tool_calling_llm.py (1)
279-313: Avoid constructingDummySpan()in default args.Default args are evaluated once; use
Noneand instantiate inside to satisfy Ruff B008 and avoid shared instances. Also apply the same pattern to other entrypoints for consistency.♻️ Suggested refactor
def prompt_call( self, system_prompt: str, user_prompt: str, response_format: Optional[Union[dict, Type[BaseModel]]] = None, sections: Optional[InputSectionsDataType] = None, - trace_span=DummySpan(), + trace_span: Optional[DummySpan] = None, request_context: Optional[Dict[str, Any]] = None, ) -> LLMResult: + trace_span = trace_span or DummySpan() messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ] return self.call( messages, response_format=response_format, user_prompt=user_prompt, sections=sections, trace_span=trace_span, request_context=request_context, ) def messages_call( self, messages: List[Dict[str, str]], response_format: Optional[Union[dict, Type[BaseModel]]] = None, - trace_span=DummySpan(), + trace_span: Optional[DummySpan] = None, request_context: Optional[Dict[str, Any]] = None, ) -> LLMResult: + trace_span = trace_span or DummySpan() return self.call( messages, response_format=response_format, trace_span=trace_span, request_context=request_context, )
02e5f64 to
bb9dcdb
Compare
5e6e631 to
de87af3
Compare
arikalon1
left a comment
arikalon1
left a comment
There was a problem hiding this comment.
nice work @gossion !
Left 2 small comments.
In addition, can we add a little note on this on the MCP docs? Might be useful for other users as well. (https://holmesgpt.dev/data-sources/remote-mcp-servers/#1-run-stdio-mcp-as-http-endpoint)
de87af3 to
dbabce1
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
holmes/core/tool_calling_llm.py (1)
863-872: Missingrequest_contextin the "approval no longer needed" re-invocation path.When approval is no longer needed and the tool is re-invoked at lines 864-870,
request_contextis not passed to_directly_invoke_tool_call. This is inconsistent with the "user approved" branch at line 892 which correctly passesrequest_context. MCP headers won't be rendered for this code path.🐛 Proposed fix
with trace_span.start_span(type="tool") as tool_span: tool_call_result.result = self._directly_invoke_tool_call( tool_name=tool_call_result.tool_name, tool_params=tool_call_result.result.params or {}, user_approved=False, tool_number=tool_number, tool_call_id=tool_call_result.tool_call_id, + request_context=request_context, ) ToolCallingLLM._log_tool_call_result(tool_span, tool_call_result)
🧹 Nitpick comments (2)
holmes/plugins/toolsets/mcp/toolset_mcp.py (1)
390-397: Add# noqa: BLE001for consistency with line 355.The exception handling here follows the same pattern as
_render_headers(line 355) where a broadExceptioncatch is used as a fail-safe during template rendering. For consistency, add the noqa comment to suppress the Ruff warning.♻️ Suggested fix
- except Exception as e: + except Exception as e: # noqa: BLE001tests/test_mcp_toolset.py (1)
1416-1422: Prefix unusedsse_read_timeoutparameter with underscore.The
_urlparameter is correctly prefixed, butsse_read_timeoutis also unused and should follow the same convention to satisfy Ruff ARG001.♻️ Suggested fix
- def capture_sse_client_call(_url, headers, *, sse_read_timeout): + def capture_sse_client_call(_url, headers, *, _sse_read_timeout):Also applies to line 1493.
|
had some conflicts, I fixup all commits and rebased it. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@docs/data-sources/remote-mcp-servers.md`:
- Around line 53-73: Update the docs for mcp_servers to state that both headers
and extra_headers may be defined together and are merged for the outbound MCP
request, with extra_headers taking precedence on name collisions; specifically
mention the fields shown in the snippet (mcp_servers -> my_server -> config ->
headers and extra_headers) and that template variables like
request_context.headers[...] in extra_headers will override any same-named keys
from headers when the final request headers are constructed.
🧹 Nitpick comments (1)
docs/data-sources/remote-mcp-servers.md (1)
93-97: Avoid secret-scanner false positives in curl example.
Even placeholder tokens can trigger scanners. Consider using an env var placeholder to keep docs clean.♻️ Proposed doc tweak
-curl -X POST http://holmes-server/api/investigate \ - -H "X-Auth-Token: your-auth-token-here" \ +curl -X POST http://holmes-server/api/investigate \ + -H "X-Auth-Token: ${X_AUTH_TOKEN}" \ -H "Content-Type: application/json" \ -d '{"question": "Check system status"}'
mainred
left a comment
mainred
left a comment
There was a problem hiding this comment.
This PR overall looks good to me. Thanks.
cff758d to
7815abc
Compare
Signed-off-by: Guoxun Wei <guwe@microsoft.com>
Signed-off-by: Guoxun Wei <guwe@microsoft.com>
Signed-off-by: Guoxun Wei <guwe@microsoft.com>
089cf34 to
d38842c
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
holmes/core/tool_calling_llm.py (1)
858-870: Propagate request_context when approval is already satisfied.The fast-path re-invocation skips request_context, so MCP headers can be missing for auto-approved tools.
🐛 Suggested fix
tool_call_result.result = self._directly_invoke_tool_call( tool_name=tool_call_result.tool_name, tool_params=tool_call_result.result.params or {}, user_approved=False, tool_number=tool_number, tool_call_id=tool_call_result.tool_call_id, + request_context=request_context, )
🤖 Fix all issues with AI agents
In `@docs/data-sources/remote-mcp-servers.md`:
- Around line 91-97: Replace the literal header value in the curl example to
avoid gitleaks false positives: update the "X-Auth-Token: your-auth-token-here"
header to use a clearly-placeholder variable (e.g., X-Auth-Token: ${AUTH_TOKEN}
or X-Auth-Token: <AUTH_TOKEN>) in the HolmesGPT request example so scanners
won’t treat it as a real secret; modify the example that constructs the curl
POST to /api/investigate accordingly.
In `@holmes/plugins/toolsets/mcp/toolset_mcp.py`:
- Around line 390-397: The current broad except in the template rendering block
should be narrowed to catch Jinja2-specific errors: replace the generic "except
Exception as e" around Template(template_str) / template.render(context) with
"except jinja2.TemplateError as e" (or "except TemplateError as e" if
TemplateError is imported) so only Jinja2 template errors are handled; keep the
logging using self.name and template_str and return template_str on that
specific except. Ensure jinja2.TemplateError is imported or referenced fully.
In `@tests/test_mcp_toolset.py`:
- Around line 1418-1421: The local callback function capture_sse_client_call
defines unused parameters sse_read_timeout and httpx_client_factory; rename them
to _sse_read_timeout and _httpx_client_factory (and do the same in the other
occurrence at the second instance) so they are prefixed with underscores to
satisfy the Ruff ARG rule while keeping the function body unchanged; update the
function signature in capture_sse_client_call accordingly wherever it appears.
♻️ Duplicate comments (1)
docs/data-sources/remote-mcp-servers.md (1)
49-73: Clarify thatextra_headerscan be used alongsideheaders.The text currently says “instead of
headers”, but the implementation merges both (withextra_headerstaking precedence).✏️ Suggested wording
-Use the `extra_headers` field (instead of `headers`) with template variables to reference headers from the incoming request: +Use `extra_headers` for templated values, optionally alongside `headers` for static values. When both are present, `extra_headers` takes precedence on key collisions:
🧹 Nitpick comments (1)
tests/test_mcp_toolset.py (1)
1145-1329: Consider stubbing_get_server_toolsin header-rendering unit tests.Calling
prerequisites_callablecan attempt real network I/O; using a local stub (or setting_mcp_configdirectly) keeps these tests deterministic and fast.
3 participants
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests
✏️ Tip: You can customize this high-level summary in your review settings.