This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e29d36f527b Add `env_provider` and `Extra.env` support to MCP stdio
transport (#69225)
e29d36f527b is described below
commit e29d36f527b62a96040992d10e826934e4dddf71
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Jul 2 05:31:42 2026 +0530
Add `env_provider` and `Extra.env` support to MCP stdio transport (#69225)
---
providers/common/ai/docs/connections/mcp.rst | 67 ++++++-
providers/common/ai/docs/toolsets.rst | 13 +-
.../common/ai/example_dags/example_mcp.py | 41 ++++
.../src/airflow/providers/common/ai/hooks/mcp.py | 86 ++++++++-
.../airflow/providers/common/ai/toolsets/mcp.py | 21 ++-
.../ai/tests/unit/common/ai/hooks/test_mcp.py | 207 ++++++++++++++++++++-
.../ai/tests/unit/common/ai/toolsets/test_mcp.py | 22 ++-
7 files changed, 435 insertions(+), 22 deletions(-)
diff --git a/providers/common/ai/docs/connections/mcp.rst
b/providers/common/ai/docs/connections/mcp.rst
index 0eb5f7bfc60..07be621e9a1 100644
--- a/providers/common/ai/docs/connections/mcp.rst
+++ b/providers/common/ai/docs/connections/mcp.rst
@@ -59,6 +59,14 @@ Arguments (Extra field)
Examples: ``["mcp-run-python"]``, ``["-m", "my_mcp_server"]``
+Environment (Extra field)
+ JSON object of environment variables for the ``stdio`` subprocess. Ignored
+ for ``http``/``sse``. For a secret that lives in a different connection or
is
+ minted fresh per call, use an ``env_provider`` instead of storing it here
+ (see below).
+
+ Examples: ``{"MY_SERVER_MODE": "readonly"}``
+
Examples
--------
@@ -99,6 +107,15 @@ Examples
"extra": "{\"transport\": \"stdio\", \"command\": \"python\",
\"args\": [\"-m\", \"my_server\"], \"timeout\": 30}"
}
+**Stdio with subprocess environment variables**
+
+.. code-block:: json
+
+ {
+ "conn_type": "mcp",
+ "extra": "{\"transport\": \"stdio\", \"command\": \"my-mcp-server\",
\"args\": [], \"env\": {\"SERVER_MODE\": \"readonly\"}}"
+ }
+
Short-lived or minted tokens
----------------------------
@@ -112,9 +129,11 @@ cannot be stored as a static connection ``password``. The
same applies to OAuth
refresh tokens, Workload Identity Federation, and GitHub App installation
tokens.
For these, pass a ``token_provider`` callable to ``MCPHook`` or ``MCPToolset``
-instead of a static token. It is called each time the connection is established
-and its return value is used as the bearer token, so a fresh token is minted
(and
-registered with secret masking so it does not leak into task logs):
+instead of a static token. It is called once, the first time a given hook or
+toolset instance establishes a connection (the result is then cached for that
+instance's lifetime), and its return value is used as the bearer token, so a
+fresh token is minted (and registered with secret masking so it does not leak
+into task logs) without ever being written to the connection:
.. code-block:: python
@@ -134,3 +153,45 @@ registered with secret masking so it does not leak into
task logs):
``token_provider`` is resolved in DAG code (it is a Python callable, not a
stored
connection field), so the signing key stays in your environment and is never
baked
into the serialized DAG.
+
+Secrets in stdio subprocess environments
+-----------------------------------------
+
+The ``stdio`` transport runs the MCP server as a local subprocess, and many
such
+servers read credentials from their own environment rather than accepting them
+as arguments -- for example, a server that reaches Splunk needs a Splunk API
key
+in ``SPLUNK_API_KEY``. ``Extra.env`` (like the rest of ``extra``) is
Fernet-encrypted
+at rest, the same as ``password``, so it is a fine place for a static value
that
+genuinely belongs to *this* connection.
+
+Use ``env_provider`` instead when the credential has no stable static form to
+store here at all -- the same situation ``token_provider`` exists for on
+HTTP/SSE:
+
+- **It lives in a different connection.** A server that reaches Splunk needs a
+ Splunk credential, not a credential for the MCP server itself; duplicating
+ it into this connection's ``Extra.env`` means two places to rotate and a
+ real chance the copies drift.
+- **It's minted fresh per call** -- an OAuth token, a Vault lease, an
+ STS-assumed role -- so there is no fixed value to store anywhere, on this
+ connection or any other.
+
+``env_provider`` is called once, the first time a given hook or toolset
instance
+establishes a connection (the result is then cached for that instance's
+lifetime). Its return value is merged over ``Extra.env`` (``env_provider`` keys
+win on conflicts), and -- as a secondary benefit -- every value it returns is
+explicitly registered with secret masking regardless of key name, unlike
+``Extra.env``, which is only masked in the Connections UI/API and task logs if
+the key name happens to match a fixed set of sensitive-looking names
+(``api_key``, ``token``, ``secret``, ``password``, etc.).
+
+The example below covers the first case -- the Splunk credential is already
+managed as its own Airflow connection:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
+ :language: python
+ :start-after: [START howto_toolset_mcp_env_provider]
+ :end-before: [END howto_toolset_mcp_env_provider]
+
+Like ``token_provider``, ``env_provider`` is resolved in DAG code, so the
secret is
+fetched at task-execution time and never baked into the serialized DAG.
diff --git a/providers/common/ai/docs/toolsets.rst
b/providers/common/ai/docs/toolsets.rst
index fd2927187c9..d97a8fece2b 100644
--- a/providers/common/ai/docs/toolsets.rst
+++ b/providers/common/ai/docs/toolsets.rst
@@ -315,10 +315,17 @@ Parameters
``"weather_get_forecast"``).
- ``token_provider``: Optional zero-argument callable returning a bearer token.
When set, it overrides the connection's static ``password`` for the
- ``Authorization`` header and is called each time the server connection is
- established -- use it for short-lived or minted tokens (e.g. a Snowflake
- managed MCP server authenticated with a key-pair JWT). See
+ ``Authorization`` header. Called once, the first time this toolset
+ establishes a connection -- use it for short-lived or minted tokens (e.g. a
+ Snowflake managed MCP server authenticated with a key-pair JWT). See
:ref:`howto/connection:mcp`.
+- ``env_provider``: Optional zero-argument callable returning a
+ ``dict[str, str]`` merged over the connection's ``Extra.env`` (winning on key
+ conflicts) for the ``stdio`` subprocess environment -- use it when the
+ credential a local stdio MCP server needs lives in a different connection, or
+ is minted fresh per call (e.g. a Splunk/Vault token), rather than storing it
+ statically here. Called once, the first time this toolset establishes a
+ connection. See :ref:`howto/connection:mcp`.
Using Multiple MCP Servers
^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
index f2e9a423679..023360b3516 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
@@ -91,3 +91,44 @@ example_mcp_multiple_servers()
# MCPToolset(StdioTransport(command="uvx", args=["mcp-run-python"])),
# ],
# )
+
+
+# ---------------------------------------------------------------------------
+# 4. Stdio server with a minted secret in the subprocess environment
+# ---------------------------------------------------------------------------
+# For local stdio MCP servers that read credentials from their own environment
+# (e.g. a server that needs a Splunk API key), pass env_provider instead of
+# storing the secret in the connection's static Extra.env. It is resolved at
+# task-execution time -- never baked into the serialized DAG -- and merged
+# over Extra.env, with env_provider's keys winning on conflicts. Here the
+# secret lives in a *different* connection (the Splunk one, not this MCP
+# server's own connection) -- something a static Extra.env cannot express.
+
+
+# [START howto_toolset_mcp_env_provider]
+def _mint_splunk_env() -> dict[str, str]:
+ from airflow.providers.common.compat.sdk import BaseHook
+
+ password = BaseHook.get_connection("splunk_default").password
+ if not password:
+ raise ValueError("splunk_default connection has no password set")
+ return {"SPLUNK_API_KEY": password}
+
+
+@dag(tags=["example"])
+def example_mcp_stdio_env_provider():
+ """Use a stdio MCP server whose subprocess needs a secret from another
connection."""
+ AgentOperator(
+ task_id="stdio_env_agent",
+ prompt="Investigate the ticket and summarize findings.",
+ llm_conn_id="pydanticai_default",
+ system_prompt="You are a support triage agent with access to MCP
tools.",
+ toolsets=[
+ MCPToolset(mcp_conn_id="spacefarer_mcp",
env_provider=_mint_splunk_env),
+ ],
+ )
+
+
+# [END howto_toolset_mcp_env_provider]
+
+example_mcp_stdio_env_provider()
diff --git a/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
index b53ea712e8c..1cc816b3649 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
@@ -47,6 +47,8 @@ class MCPHook(BaseHook):
- **Extra.transport**: Transport type — ``http`` (default), ``sse``,
or ``stdio``
- **Extra.command**: Command to run for stdio transport (e.g. ``uvx``)
- **Extra.args**: Command arguments for stdio transport (e.g.
``["mcp-run-python"]``)
+ - **Extra.env**: Environment variables for the stdio subprocess (e.g.
+ ``{"API_KEY": "..."}``). Ignored for HTTP/SSE.
- **Extra.timeout**: Connection init timeout in seconds for stdio
(default: 10)
For HTTP/SSE transports the ``Authorization`` header is, by default, a
static
@@ -54,17 +56,33 @@ class MCPHook(BaseHook):
a freshly minted or short-lived token (e.g. a Snowflake managed MCP server
authenticated with a key-pair JWT, OAuth/refresh tokens, Workload Identity
Federation, or GitHub App installation tokens) can pass a
``token_provider``
- callable instead. It is invoked each time the server connection is
established
- and its return value is used as the bearer token, so a fresh token is
minted
- without storing a long-lived secret in the connection.
+ callable instead. It is invoked once, the first time this hook establishes
a
+ connection, and its return value is used as the bearer token, so a fresh
token
+ is minted without storing a long-lived secret in the connection.
+
+ For the ``stdio`` transport, the subprocess environment is, by default, the
+ static ``Extra.env`` mapping -- a fine place for a value that genuinely
+ belongs to this connection. When the credential has no stable static form
to
+ store here at all -- it lives in a different connection (e.g. a Splunk
+ credential a Splunk connection already manages), or is minted fresh per
call
+ (e.g. a Vault lease) -- pass an ``env_provider`` callable instead. Its
return
+ value is merged over ``Extra.env`` (``env_provider`` keys win on
conflicts),
+ invoked once, the first time this hook establishes a connection.
:param mcp_conn_id: Airflow connection ID for the MCP server.
:param tool_prefix: Optional prefix prepended to tool names
(e.g. ``"weather"`` → ``"weather_get_forecast"``).
:param token_provider: Optional zero-argument callable returning a bearer
token string. When set, it overrides the connection ``password`` for
the
- ``Authorization`` header on HTTP/SSE transports and is called each
time the
- server connection is established. Ignored for the ``stdio`` transport.
+ ``Authorization`` header on HTTP/SSE transports. Called once, the first
+ time this hook establishes a connection (the result is then cached for
+ the hook's lifetime). Ignored for the ``stdio`` transport.
+ :param env_provider: Optional zero-argument callable returning a
+ ``dict[str, str]`` of environment variables for the ``stdio``
subprocess.
+ Merged over ``Extra.env`` (``env_provider`` wins on key conflicts).
Called
+ once, the first time this hook establishes a connection (the result is
+ then cached for the hook's lifetime). Ignored for the ``http``/``sse``
+ transports.
"""
conn_name_attr = "mcp_conn_id"
@@ -78,12 +96,14 @@ class MCPHook(BaseHook):
tool_prefix: str | None = None,
*,
token_provider: Callable[[], str] | None = None,
+ env_provider: Callable[[], dict[str, str]] | None = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.mcp_conn_id = mcp_conn_id
self.tool_prefix = tool_prefix
self.token_provider = token_provider
+ self.env_provider = env_provider
self._server: Any = None
@staticmethod
@@ -119,6 +139,57 @@ class MCPHook(BaseHook):
return {"Authorization": f"Bearer {conn.password}"}
return None
+ def _validate_env_dict(self, env: dict[Any, Any], source: str) -> None:
+ """Raise if ``env`` (from ``source``, for error messages) has
non-string keys/values."""
+ for key, value in env.items():
+ if not isinstance(key, str) or not key:
+ raise ValueError(
+ f"{source} for connection {self.mcp_conn_id!r} must have
non-empty string keys, "
+ f"got {type(key).__name__}."
+ )
+ if not isinstance(value, str) or not value:
+ raise ValueError(
+ f"{source} for connection {self.mcp_conn_id!r} must have
non-empty string values, "
+ f"got {type(value).__name__} for key {key!r}."
+ )
+
+ def _stdio_env(self, extra: dict[str, Any]) -> dict[str, str] | None:
+ """
+ Build the subprocess environment for the ``stdio`` transport.
+
+ Merges the static ``Extra.env`` mapping with ``env_provider`` (minted
per
+ connection), with ``env_provider`` keys taking precedence. Returns
``None``
+ when neither yields anything, in which case the MCP stdio client falls
back
+ to a small allowlist of inherited variables (see the ``mcp`` package's
+ ``get_default_environment()``), not the full parent process
environment.
+ """
+ extra_env = extra.get("env")
+ if extra_env is None:
+ extra_env = {}
+ if not isinstance(extra_env, dict):
+ raise ValueError(
+ f"'env' in extra for connection {self.mcp_conn_id!r} must be
an object of "
+ f"string keys/values, got {type(extra_env).__name__}."
+ )
+ self._validate_env_dict(extra_env, "'env' in extra")
+
+ env: dict[str, str] = dict(extra_env)
+ if self.env_provider is not None:
+ provided = self.env_provider()
+ if not isinstance(provided, dict):
+ raise ValueError(
+ f"env_provider for connection {self.mcp_conn_id!r} must
return a dict[str, str], "
+ f"got {type(provided).__name__}."
+ )
+ self._validate_env_dict(provided, "env_provider")
+ # The static connection extra is masked when the connection is
fetched;
+ # mask minted values too so they never leak into task logs.
+ for value in provided.values():
+ mask_secret(value)
+ env.update(provided)
+
+ return env or None
+
def get_conn(self) -> Any:
"""
Return a configured PydanticAI MCP toolset instance.
@@ -169,7 +240,10 @@ class MCPHook(BaseHook):
if isinstance(args, str):
args = [args]
timeout = extra.get("timeout", 10)
- toolset = MCPToolset(StdioTransport(command=command, args=args),
init_timeout=timeout)
+ toolset = MCPToolset(
+ StdioTransport(command=command, args=args,
env=self._stdio_env(extra)),
+ init_timeout=timeout,
+ )
else:
raise ValueError(
f"Unknown transport {transport!r} in connection
{self.mcp_conn_id!r}. "
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py
index 5a35cc3f54d..158d3bd7c35 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py
@@ -50,15 +50,27 @@ class MCPToolset(AbstractToolset[Any]):
For MCP endpoints that need a freshly minted or short-lived token (e.g. a
Snowflake managed MCP server authenticated with a key-pair JWT, or OAuth /
Workload Identity / GitHub App tokens), pass a ``token_provider`` callable.
- It is invoked each time the connection is established and its return value
is
- used as the bearer token, so a fresh token is minted rather than storing a
- long-lived secret in the connection.
+ It is invoked once, the first time this toolset establishes a connection,
and
+ its return value is used as the bearer token, so a fresh token is minted
+ rather than storing a long-lived secret in the connection.
+
+ For ``stdio`` servers whose subprocess needs a credential that lives in a
+ different connection, or one minted fresh per call (e.g. a Splunk/Vault
+ token), pass an ``env_provider`` callable instead. It is invoked once, the
+ first time this toolset establishes a connection, and its return value is
+ merged over the connection's static ``Extra.env`` (``env_provider`` wins on
+ key conflicts).
:param mcp_conn_id: Airflow connection ID for the MCP server.
:param tool_prefix: Optional prefix prepended to tool names
(e.g. ``"weather"`` → ``"weather_get_forecast"``).
:param token_provider: Optional zero-argument callable returning a bearer
token string, overriding the connection ``password`` for HTTP/SSE auth.
+ Called once, the first time this toolset establishes a connection.
+ :param env_provider: Optional zero-argument callable returning a
+ ``dict[str, str]`` merged over the connection's ``Extra.env`` (winning
on
+ key conflicts) for the ``stdio`` subprocess environment. Called once,
the
+ first time this toolset establishes a connection.
"""
def __init__(
@@ -67,10 +79,12 @@ class MCPToolset(AbstractToolset[Any]):
*,
tool_prefix: str | None = None,
token_provider: Callable[[], str] | None = None,
+ env_provider: Callable[[], dict[str, str]] | None = None,
) -> None:
self._mcp_conn_id = mcp_conn_id
self._tool_prefix = tool_prefix
self._token_provider = token_provider
+ self._env_provider = env_provider
self._server: Any = None
@property
@@ -85,6 +99,7 @@ class MCPToolset(AbstractToolset[Any]):
mcp_conn_id=self._mcp_conn_id,
tool_prefix=self._tool_prefix,
token_provider=self._token_provider,
+ env_provider=self._env_provider,
)
self._server = hook.get_conn()
return self._server
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
index 6e7a5629c3f..d226d585dd7 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
@@ -143,7 +143,7 @@ class TestMCPHookGetConn:
with patch.object(hook, "get_connection", return_value=conn):
result = hook.get_conn()
- mock_transport_cls.assert_called_once_with(command="uvx",
args=["mcp-run-python"])
+ mock_transport_cls.assert_called_once_with(command="uvx",
args=["mcp-run-python"], env=None)
mock_toolset_cls.assert_called_once_with(mock_transport_cls.return_value,
init_timeout=10)
assert result is mock_toolset_cls.return_value
@@ -161,7 +161,7 @@ class TestMCPHookGetConn:
with patch.object(hook, "get_connection", return_value=conn):
hook.get_conn()
- mock_transport_cls.assert_called_once_with(command="python",
args=["-m", "server"])
+ mock_transport_cls.assert_called_once_with(command="python",
args=["-m", "server"], env=None)
mock_toolset_cls.assert_called_once_with(mock_transport_cls.return_value,
init_timeout=30)
@patch(_MCP_TOOLSET, autospec=True)
@@ -176,7 +176,7 @@ class TestMCPHookGetConn:
with patch.object(hook, "get_connection", return_value=conn):
hook.get_conn()
- mock_transport_cls.assert_called_once_with(command="uvx",
args=["mcp-run-python"])
+ mock_transport_cls.assert_called_once_with(command="uvx",
args=["mcp-run-python"], env=None)
def test_http_without_host_raises(self):
hook = MCPHook(mcp_conn_id="test_conn")
@@ -364,5 +364,204 @@ class TestMCPHookTokenProvider:
hook.get_conn()
provider.assert_not_called()
- mock_transport_cls.assert_called_once_with(command="uvx",
args=["mcp-run-python"])
+ mock_transport_cls.assert_called_once_with(command="uvx",
args=["mcp-run-python"], env=None)
mock_toolset_cls.assert_called_once_with(mock_transport_cls.return_value,
init_timeout=10)
+
+
+class TestMCPHookEnvProvider:
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_STDIO_TRANSPORT, autospec=True)
+ def test_stdio_uses_env_provider(self, mock_transport_cls,
mock_toolset_cls):
+ hook = MCPHook(mcp_conn_id="test_conn", env_provider=lambda:
{"SPLUNK_API_KEY": "minted"})
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command":
"spacefarer-mcp", "args": []}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.get_conn()
+
+ mock_transport_cls.assert_called_once_with(
+ command="spacefarer-mcp", args=[], env={"SPLUNK_API_KEY": "minted"}
+ )
+
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_STDIO_TRANSPORT, autospec=True)
+ def test_static_extra_env_used_when_no_provider(self, mock_transport_cls,
mock_toolset_cls):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps(
+ {
+ "transport": "stdio",
+ "command": "spacefarer-mcp",
+ "args": [],
+ "env": {"SPACEFARER_MCP_ALLOW_EXECUTION": "0"},
+ }
+ ),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.get_conn()
+
+ mock_transport_cls.assert_called_once_with(
+ command="spacefarer-mcp", args=[],
env={"SPACEFARER_MCP_ALLOW_EXECUTION": "0"}
+ )
+
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_STDIO_TRANSPORT, autospec=True)
+ def test_env_provider_merges_over_static_extra_env(self,
mock_transport_cls, mock_toolset_cls):
+ """env_provider keys win over Extra.env on conflicts; other static
keys survive."""
+ hook = MCPHook(mcp_conn_id="test_conn", env_provider=lambda:
{"SPLUNK_API_KEY": "fresh"})
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps(
+ {
+ "transport": "stdio",
+ "command": "spacefarer-mcp",
+ "args": [],
+ "env": {"SPLUNK_API_KEY": "stale",
"SPACEFARER_MCP_ALLOW_EXECUTION": "0"},
+ }
+ ),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.get_conn()
+
+ mock_transport_cls.assert_called_once_with(
+ command="spacefarer-mcp",
+ args=[],
+ env={"SPLUNK_API_KEY": "fresh", "SPACEFARER_MCP_ALLOW_EXECUTION":
"0"},
+ )
+
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_STDIO_TRANSPORT, autospec=True)
+ def test_env_provider_called_when_establishing_connection(self,
mock_transport_cls, mock_toolset_cls):
+ provider = MagicMock(return_value={"SPLUNK_API_KEY": "minted"})
+ hook = MCPHook(mcp_conn_id="test_conn", env_provider=provider)
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command":
"spacefarer-mcp", "args": []}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.get_conn()
+
+ provider.assert_called_once_with()
+
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_STDIO_TRANSPORT, autospec=True)
+ def test_masks_provided_env_values(self, mock_transport_cls,
mock_toolset_cls):
+ """Minted env values must be registered with secret masking, like the
minted token."""
+ hook = MCPHook(mcp_conn_id="test_conn", env_provider=lambda:
{"SPLUNK_API_KEY": "minted"})
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command":
"spacefarer-mcp", "args": []}),
+ )
+ with (
+ patch.object(hook, "get_connection", return_value=conn),
+ patch("airflow.providers.common.ai.hooks.mcp.mask_secret",
autospec=True) as mock_mask,
+ ):
+ hook.get_conn()
+
+ mock_mask.assert_called_once_with("minted")
+
+ @pytest.mark.parametrize("bad_value", ["", None, 123], ids=["empty",
"none", "non_string"])
+ def test_invalid_env_provider_value_raises(self, bad_value):
+ hook = MCPHook(mcp_conn_id="test_conn", env_provider=lambda:
{"SPLUNK_API_KEY": bad_value})
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command":
"spacefarer-mcp", "args": []}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="must have non-empty string
values"):
+ hook.get_conn()
+
+ def test_invalid_env_provider_key_raises(self):
+ hook = MCPHook(mcp_conn_id="test_conn", env_provider=lambda: {1:
"value"})
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command":
"spacefarer-mcp", "args": []}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="must have non-empty string
keys"):
+ hook.get_conn()
+
+ def test_invalid_env_provider_return_type_raises(self):
+ hook = MCPHook(mcp_conn_id="test_conn", env_provider=lambda:
"not-a-dict")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command":
"spacefarer-mcp", "args": []}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="must return a dict\\[str,
str\\]"):
+ hook.get_conn()
+
+ def test_invalid_extra_env_type_raises(self):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps(
+ {"transport": "stdio", "command": "spacefarer-mcp", "args":
[], "env": ["not", "a", "dict"]}
+ ),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="must be an object of"):
+ hook.get_conn()
+
+ @pytest.mark.parametrize("falsy_bad_env", [[], 0, False, ""], ids=["list",
"zero", "false", "str"])
+ def test_falsy_invalid_extra_env_type_raises(self, falsy_bad_env):
+ """A falsy-but-wrong-type ``env`` (``[]``, ``0``, ...) must not
silently become ``{}``."""
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps(
+ {
+ "transport": "stdio",
+ "command": "spacefarer-mcp",
+ "args": [],
+ "env": falsy_bad_env,
+ }
+ ),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="must be an object of"):
+ hook.get_conn()
+
+ def test_invalid_extra_env_value_raises(self):
+ """A non-string value in the static Extra.env must fail with a clear
message, not a bare
+ TypeError deep inside subprocess.Popen when the transport actually
spawns."""
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps(
+ {
+ "transport": "stdio",
+ "command": "spacefarer-mcp",
+ "args": [],
+ "env": {"SPACEFARER_MCP_ALLOW_EXECUTION": 0},
+ }
+ ),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="must have non-empty string
values"):
+ hook.get_conn()
+
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_HTTP_TRANSPORT, autospec=True)
+ def test_http_does_not_invoke_env_provider(self, mock_transport_cls,
mock_toolset_cls):
+ """HTTP/SSE have no subprocess environment, so the env provider must
not be called."""
+ provider = MagicMock(return_value={"SPLUNK_API_KEY": "minted"})
+ hook = MCPHook(mcp_conn_id="test_conn", env_provider=provider)
+ conn = Connection(conn_id="test_conn", conn_type="mcp",
host="http://localhost:3001/mcp";)
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.get_conn()
+
+ provider.assert_not_called()
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py
index bbc3e079d58..b5fa0a56253 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py
@@ -43,7 +43,9 @@ class TestMCPToolsetGetServer:
ts = MCPToolset("mcp_conn")
server = ts._get_server()
- mock_hook_cls.assert_called_once_with(mcp_conn_id="mcp_conn",
tool_prefix=None, token_provider=None)
+ mock_hook_cls.assert_called_once_with(
+ mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=None,
env_provider=None
+ )
mock_hook_cls.return_value.get_conn.assert_called_once()
assert server is mock_server
@@ -55,7 +57,7 @@ class TestMCPToolsetGetServer:
ts._get_server()
mock_hook_cls.assert_called_once_with(
- mcp_conn_id="mcp_conn", tool_prefix="weather", token_provider=None
+ mcp_conn_id="mcp_conn", tool_prefix="weather",
token_provider=None, env_provider=None
)
@patch(_HOOK_PATH, autospec=True)
@@ -69,7 +71,21 @@ class TestMCPToolsetGetServer:
ts._get_server()
mock_hook_cls.assert_called_once_with(
- mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=provider
+ mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=provider,
env_provider=None
+ )
+
+ @patch(_HOOK_PATH, autospec=True)
+ def test_passes_env_provider_to_hook(self, mock_hook_cls):
+ mock_hook_cls.return_value.get_conn.return_value = MagicMock()
+
+ def env_provider() -> dict[str, str]:
+ return {"SPLUNK_API_KEY": "minted"}
+
+ ts = MCPToolset("mcp_conn", env_provider=env_provider)
+ ts._get_server()
+
+ mock_hook_cls.assert_called_once_with(
+ mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=None,
env_provider=env_provider
)
@patch(_HOOK_PATH, autospec=True)
