airflow.providers.common.ai.hooks.mcp

Classes

MCPHook

Hook for connecting to MCP (Model Context Protocol) servers.

Module Contents

class airflow.providers.common.ai.hooks.mcp.MCPHook(mcp_conn_id=default_conn_name, tool_prefix=None, *, token_provider=None, **kwargs)[source]

Bases: airflow.providers.common.compat.sdk.BaseHook

Hook for connecting to MCP (Model Context Protocol) servers.

Manages connection configuration for MCP servers. Supports three transport types: HTTP (Streamable HTTP), SSE, and stdio.

Connection fields:

  • host: Server URL for HTTP/SSE transports (e.g. http://localhost:3001/mcp)

  • password: Auth token (optional)

  • 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.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.

Parameters:

  • mcp_conn_id (str) – Airflow connection ID for the MCP server.

  • tool_prefix (str | None) – Optional prefix prepended to tool names (e.g. "weather" → "weather_get_forecast").

  • token_provider (collections.abc.Callable[[], str] | None) – 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'[source]
default_conn_name = 'mcp_default'[source]
conn_type = 'mcp'[source]
hook_name = 'MCP Server'[source]
mcp_conn_id = 'mcp_default'[source]
tool_prefix = None[source]
token_provider = None[source]
static get_ui_field_behaviour()[source]

Return custom field behaviour for the Airflow connection form.

get_conn()[source]

Return a configured PydanticAI MCP server instance.

Creates the appropriate MCP server based on the transport type in the connection’s extra field:

  • http (default): MCPServerStreamableHTTP

  • sse: MCPServerSSE

  • stdio: MCPServerStdio

The result is cached for the lifetime of this hook instance.

test_connection()[source]

Test connection by verifying configuration is valid.

Validates that the connection has the required fields for the configured transport type. Does NOT connect to the MCP server — that requires an async context manager.

Was this entry helpful?