How to Add Custom Tools to the MCP Server in call518/mcp-airflow-api

To add custom tools to the MCP server, either extend the existing version-specific tool modules (v1_tools.py or v2_tools.py) with the @mcp.tool() decorator, or create a separate custom_tools.py module and import it in mcp_main.py after the version-specific registration.

The call518/mcp-airflow-api repository implements a dynamic MCP (Model Context Protocol) server that automatically exposes Python functions as LLM-callable tools based on the target Airflow API version. Understanding the registration architecture allows you to safely extend the server with domain-specific functionality without breaking the core integration.

Understanding the MCP Server Tool Registration Architecture

The server uses a layered registration pattern that separates common tools from version-specific implementations.

Server Bootstrap and Version Detection

In src/mcp_airflow_api/mcp_main.py, the create_mcp_server() function instantiates a FastMCP object and delegates tool registration to version-specific modules:

# src/mcp_airflow_api/mcp_main.py (lines 41-50)

if api_version == "v1":
    from mcp_airflow_api.tools import v1_tools
    v1_tools.register_tools(mcp_instance)     # v1 registration

elif api_version == "v2":
    from mcp_airflow_api.tools import v2_tools
    v2_tools.register_tools(mcp_instance)     # v2 registration

Version-Specific Registration Flow

Both v1_tools.py and v2_tools.py follow an identical three-step pattern:

1. Set the request function: Assign the appropriate airflow_request variant (v1 or v2) to the global variable used by common_tools.
2. Register common tools: Call common_tools.register_common_tools(mcp) to register the 43 core tools shared across versions.
3. Add version-specific tools: Define additional tools using @mcp.tool() that are unique to that API version.

For example, v2_tools.py (lines 14-25) implements this pattern to expose Airflow 3.x-specific asset management tools.

Method 1: Extend Existing Version-Specific Modules

The fastest way to add custom tools to the MCP server is to append them directly to the existing version modules. This approach requires no changes to mcp_main.py because the register_tools() function already processes any tool defined within its scope.

Adding Tools to v2_tools.py

Open src/mcp_airflow_api/tools/v2_tools.py and add your coroutine below the existing tool definitions:

# src/mcp_airflow_api/tools/v2_tools.py

@mcp.tool()
async def get_dag_last_success(dag_id: str) -> Dict[str, Any]:
    """
    Returns the most recent successful DAG run for the given DAG.
    """
    # Re-use the version-specific request function already set above

    resp = await airflow_request_v2(
        "GET",
        f"/dags/{dag_id}/dagRuns?state=success&order_by=-execution_date&limit=1"
    )
    resp.raise_for_status()
    runs = resp.json().get("dag_runs", [])
    return runs[0] if runs else {"message": "No successful runs found"}

Because register_tools() executes after the common tools registration, any function decorated with @mcp.tool() inside this file is automatically exposed to the LLM when the server starts.

Method 2: Create a Dedicated Custom Tools Module

For production deployments or when you need version-agnostic utilities, create a separate custom_tools.py module. This maintains clean separation between upstream code and your internal extensions.

Creating custom_tools.py

Create src/mcp_airflow_api/tools/custom_tools.py with a registration function that follows the same pattern as the built-in modules:

# src/mcp_airflow_api/tools/custom_tools.py

from typing import Dict, Any
import logging

logger = logging.getLogger(__name__)

def register_custom_tools(mcp):
    """Register user-defined tools."""
    logger.info("Registering custom MCP tools")

    @mcp.tool()
    async def list_my_connections(limit: int = 10) -> Dict[str, Any]:
        """
        Retrieves a subset of Airflow connections.
        """
        resp = await airflow_request("GET", f"/connections?limit={limit}")
        resp.raise_for_status()
        return resp.json()

Note that airflow_request is not imported here; it relies on the global variable set by v1_tools.py or v2_tools.py during their registration phase.

Registering in mcp_main.py

Import your custom module and invoke its registration function in src/mcp_airflow_api/mcp_main.py immediately after the version-specific registration:

# src/mcp_airflow_api/mcp_main.py

# ... after the version-specific registration block

from mcp_airflow_api.tools import custom_tools   # new import

# Inside create_mcp_server(), after v1/v2 registration:

if api_version == "v1":
    from mcp_airflow_api.tools import v1_tools
    v1_tools.register_tools(mcp_instance)
elif api_version == "v2":
    from mcp_airflow_api.tools import v2_tools
    v2_tools.register_tools(mcp_instance)

# Register user tools

custom_tools.register_custom_tools(mcp_instance)   # register custom tools

This ensures your custom tools are loaded with the correct Airflow API context already established.

Key Files and Architecture Reference

File	Role	Critical Sections
src/mcp_airflow_api/mcp_main.py	Server bootstrap, version detection, orchestration	create_mcp_server() (lines 29-34), version-specific imports (lines 41-50)
src/mcp_airflow_api/tools/v1_tools.py	V1-specific tool registration (Airflow 2.x)	register_tools() (lines 13-24)
src/mcp_airflow_api/tools/v2_tools.py	V2-specific tool registration (Airflow 3.x)	register_tools() (lines 14-25), asset tools (lines 27-63)
src/mcp_airflow_api/tools/common_tools.py	Shared 43 core tools	register_common_tools(mcp) (starts at line 21)
src/mcp_airflow_api/functions.py	HTTP helpers	airflow_request, get_api_version

Summary

• The MCP server dynamically loads tools based on the configured Airflow API version (v1 or v2), delegating registration to v1_tools.py or v2_tools.py.
• Use @mcp.tool() to expose any Python coroutine as an LLM-callable tool; the decorator is provided by the FastMCP instance.
• Extend existing modules for quick, version-specific additions without touching mcp_main.py.
• Create a separate custom_tools.py for production-grade, version-agnostic extensions, importing and registering it in mcp_main.py after the version-specific registration.
• Leverage airflow_request (set globally by the version modules) to ensure your custom tools use the correct API endpoint and authentication.

Frequently Asked Questions

Do I need to modify mcp_main.py when adding tools to v1_tools.py or v2_tools.py?

No. The register_tools() function in each version module automatically registers any tool decorated with @mcp.tool() when it is called by create_mcp_server(). Simply append your new coroutine to the existing file and restart the server.

Can I use the same custom tools for both Airflow API v1 and v2?

Yes, but you must ensure the tools are registered after the version-specific module sets the global airflow_request function. When using a separate custom_tools.py approach, import and call register_custom_tools() in mcp_main.py immediately after the v1 or v2 registration block. This ensures airflow_request points to the correct versioned endpoint.

Where should I place helper functions that are not exposed as tools?

Place non-tool helper functions in src/mcp_airflow_api/tools/common_tools.py (if shared across versions) or create a new utility module such as src/mcp_airflow_api/tools/utils.py. Import these helpers into your custom tool modules as needed. Do not decorate helper functions with @mcp.tool() unless you want them exposed to the LLM.

How does the server know which Airflow API version to use?

The server reads the AIRFLOW_API_VERSION environment variable (or configuration) during startup in create_mcp_server(). Based on this value (v1 or v2), it imports the corresponding tool module (v1_tools or v2_tools), which sets the appropriate airflow_request function and registers version-specific tools before loading any custom extensions.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}