Add token_provider for short-lived MCP auth in common.ai

providers/common/ai/docs/connections/mcp.rst

@@ -45,7 +45,9 @@ Host
     Examples: ``http://localhost:3001/mcp``, ``https://mcp.example.com/v1``
 
 Auth Token (Password field)
-    Optional authentication token for the MCP server.
+    Optional authentication token for the MCP server. Sent as a static
+    ``Authorization: Bearer <token>`` header on HTTP/SSE transports. For
+    short-lived or minted tokens, use a ``token_provider`` instead (see below).
 
 Command (Extra field)
     The command to run for ``stdio`` transport. Required when transport is ``stdio``.
@@ -96,3 +98,39 @@ Examples
         "conn_type": "mcp",
         "extra": "{\"transport\": \"stdio\", \"command\": \"python\", \"args\": [\"-m\", \"my_server\"], \"timeout\": 30}"
     }
+
+Short-lived or minted tokens
+----------------------------
+
+Some MCP endpoints require a freshly minted, short-lived token rather than a
+static one. For example, `Snowflake managed MCP servers
+<https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-agents-mcp>`__
+are best authenticated with a `key-pair JWT
+<https://docs.snowflake.com/en/user-guide/key-pair-auth>`__: the private key never
+leaves your environment and the signed JWT expires after about an hour, so it
+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):
+
+.. code-block:: python
+
+    from airflow.providers.common.ai.toolsets.mcp import MCPToolset
+
+
+    def mint_snowflake_jwt() -> str:
+        # Sign a short-lived JWT from the Snowflake connection's key-pair.
+        ...
+
+
+    toolset = MCPToolset(
+        mcp_conn_id="snowflake_managed_mcp",
+        token_provider=mint_snowflake_jwt,
+    )
+
+``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.

providers/common/ai/docs/toolsets.rst

@@ -272,6 +272,12 @@ Parameters
 - ``tool_prefix``: Optional prefix prepended to tool names to avoid
   collisions when using multiple MCP servers (e.g. ``"weather"`` produces
   ``"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
+  :ref:`howto/connection:mcp`.
 
 Using Multiple MCP Servers
 ^^^^^^^^^^^^^^^^^^^^^^^^^^

providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py

@@ -16,10 +16,23 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from airflow.providers.common.compat.sdk import BaseHook
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from airflow.sdk.execution_time.secrets_masker import mask_secret
+else:
+    try:
+        from airflow.sdk.log import mask_secret
+    except ImportError:
+        try:
+            from airflow.sdk.execution_time.secrets_masker import mask_secret
+        except ImportError:
+            from airflow.utils.log.secrets_masker import mask_secret
+
 
 class MCPHook(BaseHook):
     """
@@ -36,9 +49,22 @@ class MCPHook(BaseHook):
         - **Extra.args**: Command arguments for stdio transport (e.g. ``["mcp-run-python"]``)
         - **Extra.timeout**: Connection timeout in seconds for stdio (default: 10)
 
+    For HTTP/SSE transports the ``Authorization`` header is, by default, a static
+    ``Bearer`` token taken from the connection ``password``. Endpoints that require
+    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.
+
     :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.
     """
 
     conn_name_attr = "mcp_conn_id"
@@ -50,11 +76,14 @@ def __init__(
         self,
         mcp_conn_id: str = default_conn_name,
         tool_prefix: str | None = None,
+        *,
+        token_provider: Callable[[], 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._server: Any = None
 
     @staticmethod
@@ -68,6 +97,28 @@ def get_ui_field_behaviour() -> dict[str, Any]:
             },
         }
 
+    def _auth_headers(self, conn: Any) -> dict[str, str] | None:
+        """
+        Build the ``Authorization`` header for HTTP/SSE transports.
+
+        Prefers ``token_provider`` (minted per connection) over the static
+        connection ``password``. Returns ``None`` when neither yields a token.
+        """
+        if self.token_provider is not None:
+            token = self.token_provider()
+            if not isinstance(token, str) or not token:
+                raise ValueError(
+                    f"token_provider for connection {self.mcp_conn_id!r} must return a non-empty "
+                    f"string token, got {type(token).__name__}."
+                )
+            # The static connection password is masked when the connection is
+            # fetched; mask the minted token too so it never leaks into task logs.
+            mask_secret(token)
+            return {"Authorization": f"Bearer {token}"}
+        if conn.password:
+            return {"Authorization": f"Bearer {conn.password}"}
+        return None
+
     def get_conn(self) -> Any:
         """
         Return a configured PydanticAI MCP server instance.
@@ -94,16 +145,19 @@ def get_conn(self) -> Any:
         conn = self.get_connection(self.mcp_conn_id)
         extra = conn.extra_dejson
         transport = extra.get("transport", "http")
-        headers = {"Authorization": f"Bearer {conn.password}"} if conn.password else None
 
         if transport == "http":
             if not conn.host:
                 raise ValueError(f"Connection {self.mcp_conn_id!r} requires a host URL for HTTP transport.")
-            self._server = MCPServerStreamableHTTP(conn.host, headers=headers, tool_prefix=self.tool_prefix)
+            self._server = MCPServerStreamableHTTP(
+                conn.host, headers=self._auth_headers(conn), tool_prefix=self.tool_prefix
+            )
         elif transport == "sse":
             if not conn.host:
                 raise ValueError(f"Connection {self.mcp_conn_id!r} requires a host URL for SSE transport.")
-            self._server = MCPServerSSE(conn.host, headers=headers, tool_prefix=self.tool_prefix)
+            self._server = MCPServerSSE(
+                conn.host, headers=self._auth_headers(conn), tool_prefix=self.tool_prefix
+            )
         elif transport == "stdio":
             command = extra.get("command")
             if not command:

providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py

@@ -24,6 +24,8 @@
 from typing_extensions import Self
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from pydantic_ai._run_context import RunContext
 
 
@@ -46,19 +48,30 @@ class MCPToolset(AbstractToolset[Any]):
     :class:`~pydantic_ai.mcp.MCPServerSSE`, and
     :class:`~pydantic_ai.mcp.MCPServerStdio` all implement ``AbstractToolset``.
 
+    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.
+
     :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.
     """
 
     def __init__(
         self,
         mcp_conn_id: str,
         *,
         tool_prefix: str | None = None,
+        token_provider: Callable[[], str] | None = None,
     ) -> None:
         self._mcp_conn_id = mcp_conn_id
         self._tool_prefix = tool_prefix
+        self._token_provider = token_provider
         self._server: Any = None
 
     @property
@@ -69,7 +82,11 @@ def _get_server(self) -> Any:
         if self._server is None:
             from airflow.providers.common.ai.hooks.mcp import MCPHook
 
-            hook = MCPHook(mcp_conn_id=self._mcp_conn_id, tool_prefix=self._tool_prefix)
+            hook = MCPHook(
+                mcp_conn_id=self._mcp_conn_id,
+                tool_prefix=self._tool_prefix,
+                token_provider=self._token_provider,
+            )
             self._server = hook.get_conn()
         return self._server
 

providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import json
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -246,3 +246,103 @@ def test_hidden_fields(self):
     def test_relabeling(self):
         behaviour = MCPHook.get_ui_field_behaviour()
         assert behaviour["relabeling"]["password"] == "Auth Token"
+
+
+class TestMCPHookTokenProvider:
+    @patch(_MCP_HTTP)
+    def test_http_uses_token_provider(self, mock_server_cls):
+        hook = MCPHook(mcp_conn_id="test_conn", token_provider=lambda: "minted-jwt")
+        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()
+
+        mock_server_cls.assert_called_once_with(
+            "http://localhost:3001/mcp",
+            headers={"Authorization": "Bearer minted-jwt"},
+            tool_prefix=None,
+        )
+
+    @patch(_MCP_HTTP)
+    def test_token_provider_overrides_static_password(self, mock_server_cls):
+        hook = MCPHook(mcp_conn_id="test_conn", token_provider=lambda: "fresh")
+        conn = Connection(
+            conn_id="test_conn",
+            conn_type="mcp",
+            host="http://localhost:3001/mcp",
+            password="static-pat",
+        )
+        with patch.object(hook, "get_connection", return_value=conn):
+            hook.get_conn()
+
+        mock_server_cls.assert_called_once_with(
+            "http://localhost:3001/mcp",
+            headers={"Authorization": "Bearer fresh"},
+            tool_prefix=None,
+        )
+
+    @patch(_MCP_HTTP)
+    def test_token_provider_called_when_establishing_connection(self, mock_server_cls):
+        provider = MagicMock(return_value="tok")
+        hook = MCPHook(mcp_conn_id="test_conn", token_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_called_once_with()
+
+    @patch(_MCP_HTTP)
+    def test_masks_minted_token(self, mock_server_cls):
+        """The minted token must be registered with secret masking, like conn.password."""
+        hook = MCPHook(mcp_conn_id="test_conn", token_provider=lambda: "minted-jwt")
+        conn = Connection(conn_id="test_conn", conn_type="mcp", host="http://localhost:3001/mcp")
+        with (
+            patch.object(hook, "get_connection", return_value=conn),
+            patch("airflow.providers.common.ai.hooks.mcp.mask_secret") as mock_mask,
+        ):
+            hook.get_conn()
+
+        mock_mask.assert_called_once_with("minted-jwt")
+
+    @pytest.mark.parametrize("bad_token", ["", None], ids=["empty", "non_string"])
+    @patch(_MCP_HTTP)
+    def test_invalid_token_raises(self, mock_server_cls, bad_token):
+        """A token_provider returning a non-string or empty value fails loud."""
+        hook = MCPHook(mcp_conn_id="test_conn", token_provider=lambda: bad_token)
+        conn = Connection(conn_id="test_conn", conn_type="mcp", host="http://localhost:3001/mcp")
+        with patch.object(hook, "get_connection", return_value=conn):
+            with pytest.raises(ValueError, match="must return a non-empty string token"):
+                hook.get_conn()
+
+    @patch(_MCP_SSE)
+    def test_sse_uses_token_provider(self, mock_server_cls):
+        hook = MCPHook(mcp_conn_id="test_conn", token_provider=lambda: "minted")
+        conn = Connection(
+            conn_id="test_conn",
+            conn_type="mcp",
+            host="http://localhost:3001/sse",
+            extra=json.dumps({"transport": "sse"}),
+        )
+        with patch.object(hook, "get_connection", return_value=conn):
+            hook.get_conn()
+
+        mock_server_cls.assert_called_once_with(
+            "http://localhost:3001/sse",
+            headers={"Authorization": "Bearer minted"},
+            tool_prefix=None,
+        )
+
+    @patch(_MCP_STDIO)
+    def test_stdio_does_not_invoke_token_provider(self, mock_server_cls):
+        """stdio has no HTTP headers, so the token provider must not be called."""
+        provider = MagicMock(return_value="tok")
+        hook = MCPHook(mcp_conn_id="test_conn", token_provider=provider)
+        conn = Connection(
+            conn_id="test_conn",
+            conn_type="mcp",
+            extra=json.dumps({"transport": "stdio", "command": "uvx", "args": ["mcp-run-python"]}),
+        )
+        with patch.object(hook, "get_connection", return_value=conn):
+            hook.get_conn()
+
+        provider.assert_not_called()
+        mock_server_cls.assert_called_once_with("uvx", args=["mcp-run-python"], timeout=10, tool_prefix=None)

providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py

@@ -43,7 +43,7 @@ def test_delegates_to_hook(self, mock_hook_cls):
         ts = MCPToolset("mcp_conn")
         server = ts._get_server()
 
-        mock_hook_cls.assert_called_once_with(mcp_conn_id="mcp_conn", tool_prefix=None)
+        mock_hook_cls.assert_called_once_with(mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=None)
         mock_hook_cls.return_value.get_conn.assert_called_once()
         assert server is mock_server
 
@@ -54,7 +54,23 @@ def test_passes_tool_prefix_to_hook(self, mock_hook_cls):
         ts = MCPToolset("mcp_conn", tool_prefix="weather")
         ts._get_server()
 
-        mock_hook_cls.assert_called_once_with(mcp_conn_id="mcp_conn", tool_prefix="weather")
+        mock_hook_cls.assert_called_once_with(
+            mcp_conn_id="mcp_conn", tool_prefix="weather", token_provider=None
+        )
+
+    @patch(_HOOK_PATH, autospec=True)
+    def test_passes_token_provider_to_hook(self, mock_hook_cls):
+        mock_hook_cls.return_value.get_conn.return_value = MagicMock()
+
+        def provider() -> str:
+            return "minted"
+
+        ts = MCPToolset("mcp_conn", token_provider=provider)
+        ts._get_server()
+
+        mock_hook_cls.assert_called_once_with(
+            mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=provider
+        )
 
     @patch(_HOOK_PATH, autospec=True)
     def test_caches_server(self, mock_hook_cls):