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

> Learn how to add custom tools to the MCP server in call518/mcp-airflow-api. Extend existing modules or create a new custom_tools.py file for seamless integration.

- Repository: [JungJungIn/mcp-airflow-api](https://github.com/call518/mcp-airflow-api)
- Tags: how-to-guide
- Published: 2026-02-26

---

**To add custom tools to the MCP server, either extend the existing version-specific tool modules ([`v1_tools.py`](https://github.com/call518/mcp-airflow-api/blob/main/v1_tools.py) or [`v2_tools.py`](https://github.com/call518/mcp-airflow-api/blob/main/v2_tools.py)) with the `@mcp.tool()` decorator, or create a separate [`custom_tools.py`](https://github.com/call518/mcp-airflow-api/blob/main/custom_tools.py) module and import it in [`mcp_main.py`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/src/mcp_airflow_api/mcp_main.py), the `create_mcp_server()` function instantiates a `FastMCP` object and delegates tool registration to version-specific modules:

```python

# 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`](https://github.com/call518/mcp-airflow-api/blob/main/v1_tools.py) and [`v2_tools.py`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/src/mcp_airflow_api/tools/v2_tools.py) and add your coroutine below the existing tool definitions:

```python

# 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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/src/mcp_airflow_api/tools/custom_tools.py) with a registration function that follows the same pattern as the built-in modules:

```python

# 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`](https://github.com/call518/mcp-airflow-api/blob/main/v1_tools.py) or [`v2_tools.py`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/src/mcp_airflow_api/mcp_main.py) immediately after the version-specific registration:

```python

# 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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/v1_tools.py) or [`v2_tools.py`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/mcp_main.py).
- **Create a separate [`custom_tools.py`](https://github.com/call518/mcp-airflow-api/blob/main/custom_tools.py)** for production-grade, version-agnostic extensions, importing and registering it in [`mcp_main.py`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/custom_tools.py) approach, import and call `register_custom_tools()` in [`mcp_main.py`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/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`](https://github.com/call518/mcp-airflow-api/blob/main/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.