mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-05 09:12:39 +00:00
feat: add environment variable filtering hook to StdioTransport
- Introduced an optional `_env_filter_hook` to allow extensions to modify the environment variables passed to MCP servers, enabling features like credential stripping. - Updated tests to ensure the filtering hook is applied correctly after merging user-supplied and default environment variables.
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
"""Stdio transport for MCP servers running as local processes."""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
import subprocess
|
||||
from typing import Any
|
||||
|
||||
@@ -9,6 +10,16 @@ from typing_extensions import Self
|
||||
from crewai.mcp.transports.base import BaseTransport, TransportType
|
||||
|
||||
|
||||
_env_filter_hook: Callable[[dict[str, str]], dict[str, str]] | None = None
|
||||
"""Optional hook to post-process the environment passed to stdio MCP servers.
|
||||
|
||||
Extensions (e.g., enterprise policy) can set this to enforce org-wide rules such
|
||||
as stripping credentials from `env` before the subprocess is spawned. The hook
|
||||
receives the merged env (SDK defaults + user-supplied `env=`) and returns the
|
||||
filtered env. Set to None to disable.
|
||||
"""
|
||||
|
||||
|
||||
class StdioTransport(BaseTransport):
|
||||
"""Stdio transport for connecting to local MCP servers.
|
||||
|
||||
@@ -75,6 +86,9 @@ class StdioTransport(BaseTransport):
|
||||
process_env = get_default_environment()
|
||||
process_env.update(self.env)
|
||||
|
||||
if _env_filter_hook is not None:
|
||||
process_env = _env_filter_hook(process_env)
|
||||
|
||||
server_params = StdioServerParameters(
|
||||
command=self.command,
|
||||
args=self.args,
|
||||
|
||||
@@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import crewai.mcp.transports.stdio as stdio_transport_module
|
||||
from crewai.mcp.transports.stdio import StdioTransport
|
||||
|
||||
|
||||
@@ -79,3 +80,45 @@ async def test_user_env_overrides_default_environment():
|
||||
await transport.connect()
|
||||
|
||||
assert captured["env"]["PATH"] == "/custom/bin"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_env_filter_hook_runs_after_merge():
|
||||
"""An extension-supplied env_filter_hook must be applied to the final env."""
|
||||
transport = StdioTransport(
|
||||
command="python",
|
||||
args=["server.py"],
|
||||
env={"OPENAI_API_KEY": "sk-test", "AWS_SECRET_ACCESS_KEY": "should-strip"},
|
||||
)
|
||||
|
||||
captured: dict[str, dict[str, str] | None] = {}
|
||||
|
||||
fake_ctx = MagicMock()
|
||||
fake_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock()))
|
||||
fake_ctx.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
def fake_stdio_client(server_params):
|
||||
captured["env"] = server_params.env
|
||||
return fake_ctx
|
||||
|
||||
def drop_aws(env):
|
||||
return {k: v for k, v in env.items() if not k.startswith("AWS_")}
|
||||
|
||||
original_hook = stdio_transport_module._env_filter_hook
|
||||
stdio_transport_module._env_filter_hook = drop_aws
|
||||
try:
|
||||
with (
|
||||
patch("mcp.client.stdio.stdio_client", side_effect=fake_stdio_client),
|
||||
patch(
|
||||
"mcp.client.stdio.get_default_environment",
|
||||
return_value={"PATH": "/usr/bin"},
|
||||
),
|
||||
):
|
||||
await transport.connect()
|
||||
finally:
|
||||
stdio_transport_module._env_filter_hook = original_hook
|
||||
|
||||
env = captured["env"]
|
||||
assert "AWS_SECRET_ACCESS_KEY" not in env
|
||||
assert env.get("OPENAI_API_KEY") == "sk-test"
|
||||
assert env.get("PATH") == "/usr/bin"
|
||||
|
||||
Reference in New Issue
Block a user