From 2b652c592634e9b992f46e46fd94e622dc41cfcc Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 02:29:16 +0000
Subject: [PATCH 1/9] support cookies

---
 pyproject.toml                                |  3 ++-
 .../client/airflow_client.py                  | 21 ++++++++++------
 src/airflow_mcp_server/server.py              | 13 +++++++---
 src/airflow_mcp_server/server_safe.py         | 13 +++++++---
 src/airflow_mcp_server/server_unsafe.py       | 13 +++++++---
 src/airflow_mcp_server/tools/tool_manager.py  | 24 ++++++++++++++----
 tests/client/test_airflow_client.py           | 25 +++++++++++++++++++
 7 files changed, 90 insertions(+), 22 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e3f292a..b8dc0ed 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,7 +52,8 @@ build-backend = "hatchling.build"
 exclude = [
     "*",
     "!src/**",
-    "!pyproject.toml"
+    "!pyproject.toml",
+    "!assets/**"
 ]
 
 [tool.hatch.build.targets.wheel]
diff --git a/src/airflow_mcp_server/client/airflow_client.py b/src/airflow_mcp_server/client/airflow_client.py
index e9f247c..1014a3e 100644
--- a/src/airflow_mcp_server/client/airflow_client.py
+++ b/src/airflow_mcp_server/client/airflow_client.py
@@ -35,18 +35,22 @@ class AirflowClient:
         self,
         spec_path: Path | str | dict | bytes | BinaryIO | TextIO,
         base_url: str,
-        auth_token: str,
+        auth_token: str | None = None,
+        cookie: str | None = None,
     ) -> None:
         """Initialize Airflow client.
 
         Args:
             spec_path: OpenAPI spec as file path, dict, bytes, or file object
             base_url: Base URL for API
-            auth_token: Authentication token
+            auth_token: Authentication token (optional if cookie is provided)
+            cookie: Session cookie (optional if auth_token is provided)
 
         Raises:
-            ValueError: If spec_path is invalid or spec cannot be loaded
+            ValueError: If spec_path is invalid or spec cannot be loaded or if neither auth_token nor cookie is provided
         """
+        if not auth_token and not cookie:
+            raise ValueError("Either auth_token or cookie must be provided")
         try:
             # Load and parse OpenAPI spec
             if isinstance(spec_path, dict):
@@ -96,10 +100,13 @@ class AirflowClient:
 
             # API configuration
             self.base_url = base_url.rstrip("/")
-            self.headers = {
-                "Authorization": f"Basic {auth_token}",
-                "Accept": "application/json",
-            }
+            self.headers = {"Accept": "application/json"}
+
+            # Set authentication header based on what was provided
+            if auth_token:
+                self.headers["Authorization"] = f"Basic {auth_token}"
+            elif cookie:
+                self.headers["Cookie"] = cookie
 
         except Exception as e:
             logger.error("Failed to initialize AirflowClient: %s", e)
diff --git a/src/airflow_mcp_server/server.py b/src/airflow_mcp_server/server.py
index 09933d2..87d3a94 100644
--- a/src/airflow_mcp_server/server.py
+++ b/src/airflow_mcp_server/server.py
@@ -22,9 +22,16 @@ logger = logging.getLogger(__name__)
 
 async def serve() -> None:
     """Start MCP server."""
-    required_vars = ["AIRFLOW_BASE_URL", "AUTH_TOKEN"]
-    if not all(var in os.environ for var in required_vars):
-        raise ValueError(f"Missing required environment variables: {required_vars}")
+    # Check for AIRFLOW_BASE_URL which is always required
+    if "AIRFLOW_BASE_URL" not in os.environ:
+        raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
+
+    # Check for either AUTH_TOKEN or COOKIE
+    has_auth_token = "AUTH_TOKEN" in os.environ
+    has_cookie = "COOKIE" in os.environ
+
+    if not has_auth_token and not has_cookie:
+        raise ValueError("Either AUTH_TOKEN or COOKIE environment variable must be provided")
 
     server = Server("airflow-mcp-server")
 
diff --git a/src/airflow_mcp_server/server_safe.py b/src/airflow_mcp_server/server_safe.py
index bf81b3e..6be1f36 100644
--- a/src/airflow_mcp_server/server_safe.py
+++ b/src/airflow_mcp_server/server_safe.py
@@ -13,9 +13,16 @@ logger = logging.getLogger(__name__)
 
 async def serve() -> None:
     """Start MCP server in safe mode (read-only operations)."""
-    required_vars = ["AIRFLOW_BASE_URL", "AUTH_TOKEN"]
-    if not all(var in os.environ for var in required_vars):
-        raise ValueError(f"Missing required environment variables: {required_vars}")
+    # Check for AIRFLOW_BASE_URL which is always required
+    if "AIRFLOW_BASE_URL" not in os.environ:
+        raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
+
+    # Check for either AUTH_TOKEN or COOKIE
+    has_auth_token = "AUTH_TOKEN" in os.environ
+    has_cookie = "COOKIE" in os.environ
+
+    if not has_auth_token and not has_cookie:
+        raise ValueError("Either AUTH_TOKEN or COOKIE environment variable must be provided")
 
     server = Server("airflow-mcp-server-safe")
 
diff --git a/src/airflow_mcp_server/server_unsafe.py b/src/airflow_mcp_server/server_unsafe.py
index bcc7932..b33b46c 100644
--- a/src/airflow_mcp_server/server_unsafe.py
+++ b/src/airflow_mcp_server/server_unsafe.py
@@ -13,9 +13,16 @@ logger = logging.getLogger(__name__)
 
 async def serve() -> None:
     """Start MCP server in unsafe mode (all operations)."""
-    required_vars = ["AIRFLOW_BASE_URL", "AUTH_TOKEN"]
-    if not all(var in os.environ for var in required_vars):
-        raise ValueError(f"Missing required environment variables: {required_vars}")
+    # Check for AIRFLOW_BASE_URL which is always required
+    if "AIRFLOW_BASE_URL" not in os.environ:
+        raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
+
+    # Check for either AUTH_TOKEN or COOKIE
+    has_auth_token = "AUTH_TOKEN" in os.environ
+    has_cookie = "COOKIE" in os.environ
+
+    if not has_auth_token and not has_cookie:
+        raise ValueError("Either AUTH_TOKEN or COOKIE environment variable must be provided")
 
     server = Server("airflow-mcp-server-unsafe")
 
diff --git a/src/airflow_mcp_server/tools/tool_manager.py b/src/airflow_mcp_server/tools/tool_manager.py
index ee70e86..9eb9cbe 100644
--- a/src/airflow_mcp_server/tools/tool_manager.py
+++ b/src/airflow_mcp_server/tools/tool_manager.py
@@ -32,12 +32,26 @@ def _initialize_client() -> AirflowClient:
         except Exception as e:
             raise ValueError("Default OpenAPI spec not found in package resources") from e
 
-    required_vars = ["AIRFLOW_BASE_URL", "AUTH_TOKEN"]
-    missing_vars = [var for var in required_vars if var not in os.environ]
-    if missing_vars:
-        raise ValueError(f"Missing required environment variables: {missing_vars}")
+    # Check for base URL
+    if "AIRFLOW_BASE_URL" not in os.environ:
+        raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
 
-    return AirflowClient(spec_path=spec_path, base_url=os.environ["AIRFLOW_BASE_URL"], auth_token=os.environ["AUTH_TOKEN"])
+    # Check for either AUTH_TOKEN or COOKIE
+    has_auth_token = "AUTH_TOKEN" in os.environ
+    has_cookie = "COOKIE" in os.environ
+
+    if not has_auth_token and not has_cookie:
+        raise ValueError("Either AUTH_TOKEN or COOKIE environment variable must be provided")
+
+    # Initialize client with appropriate authentication method
+    client_args = {"spec_path": spec_path, "base_url": os.environ["AIRFLOW_BASE_URL"]}
+
+    if has_auth_token:
+        client_args["auth_token"] = os.environ["AUTH_TOKEN"]
+    elif has_cookie:
+        client_args["cookie"] = os.environ["COOKIE"]
+
+    return AirflowClient(**client_args)
 
 
 async def _initialize_tools() -> None:
diff --git a/tests/client/test_airflow_client.py b/tests/client/test_airflow_client.py
index 3fdc03d..0b4460c 100644
--- a/tests/client/test_airflow_client.py
+++ b/tests/client/test_airflow_client.py
@@ -32,6 +32,31 @@ def test_init_client_initialization(client: AirflowClient) -> None:
     assert isinstance(client.spec, OpenAPI)
     assert client.base_url == "http://localhost:8080/api/v1"
     assert client.headers["Authorization"] == "Basic test-token"
+    assert "Cookie" not in client.headers
+
+
+def test_init_client_with_cookie() -> None:
+    with resources.files("airflow_mcp_server.resources").joinpath("v1.yaml").open("rb") as f:
+        spec = yaml.safe_load(f)
+    client = AirflowClient(
+        spec_path=spec,
+        base_url="http://localhost:8080/api/v1",
+        cookie="session=b18e8c5e-92f5-4d1e-a8f2-7c1b62110cae.vmX5kqDq5TdvT9BzTlypMVclAwM",
+    )
+    assert isinstance(client.spec, OpenAPI)
+    assert client.base_url == "http://localhost:8080/api/v1"
+    assert "Authorization" not in client.headers
+    assert client.headers["Cookie"] == "session=b18e8c5e-92f5-4d1e-a8f2-7c1b62110cae.vmX5kqDq5TdvT9BzTlypMVclAwM"
+
+
+def test_init_client_missing_auth() -> None:
+    with resources.files("airflow_mcp_server.resources").joinpath("v1.yaml").open("rb") as f:
+        spec = yaml.safe_load(f)
+    with pytest.raises(ValueError, match="Either auth_token or cookie must be provided"):
+        AirflowClient(
+            spec_path=spec,
+            base_url="http://localhost:8080/api/v1",
+        )
 
 
 def test_init_load_spec_from_bytes() -> None:

From 5663f5662175ca40eb1de6aef718e2d1732304a6 Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 02:29:45 +0000
Subject: [PATCH 2/9] Tests for PRs

---
 .github/workflows/pytest.yml | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)
 create mode 100644 .github/workflows/pytest.yml

diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
new file mode 100644
index 0000000..7e53814
--- /dev/null
+++ b/.github/workflows/pytest.yml
@@ -0,0 +1,29 @@
+name: Run Tests
+
+on:
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
+
+    steps:
+    - uses: actions/checkout@v3
+    
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+        
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install .[test]
+        
+    - name: Run pytest
+      run: |
+        pytest tests/ -v

From 8b38a26e8a617a7c0b344422d3dae10c1eea2d74 Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 02:33:48 +0000
Subject: [PATCH 3/9] Updates for using Cookies

---
 README.md | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index fa81cf0..f35aebf 100644
--- a/README.md
+++ b/README.md
@@ -29,7 +29,10 @@ https://github.com/user-attachments/assets/f3e60fff-8680-4dd9-b08e-fa7db655a705
       ],
       "env": {
         "AIRFLOW_BASE_URL": "http://<host:port>/api/v1",
-        "AUTH_TOKEN": "<base64_encoded_username_password>"
+        // Either use AUTH_TOKEN for basic auth
+        "AUTH_TOKEN": "<base64_encoded_username_password>",
+        // Or use COOKIE for cookie-based auth
+        "COOKIE": "<session_cookie>"
       }
     }
   }
@@ -57,10 +60,17 @@ airflow-mcp-server --unsafe
 
 The MCP Server expects environment variables to be set:
 - `AIRFLOW_BASE_URL`: The base URL of the Airflow API
-- `AUTH_TOKEN`: The token to use for authorization (_This should be base64 encoded username:password_)
+- `AUTH_TOKEN`: The token to use for basic auth (_This should be base64 encoded username:password_) (_Optional if COOKIE is provided_)
+- `COOKIE`: The session cookie to use for authentication (_Optional if AUTH_TOKEN is provided_)
 - `OPENAPI_SPEC`: The path to the OpenAPI spec file (_Optional_) (_defaults to latest stable release_)
 
-*Currently, only Basic Auth is supported.*
+**Authentication**
+
+The server supports two authentication methods:
+- **Basic Auth**: Using base64 encoded username:password via `AUTH_TOKEN` environment variable
+- **Cookie**: Using session cookie via `COOKIE` environment variable
+
+At least one of these authentication methods must be provided.
 
 **Page Limit**
 

From 355fb55bdb71cd03c35aea4d4917d288fa02141e Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 06:10:04 +0000
Subject: [PATCH 4/9] precedence implementation

---
 src/airflow_mcp_server/client/airflow_client.py |  8 ++++----
 src/airflow_mcp_server/server.py                | 12 +++++++++++-
 src/airflow_mcp_server/server_safe.py           | 12 +++++++++++-
 src/airflow_mcp_server/server_unsafe.py         | 12 +++++++++++-
 src/airflow_mcp_server/tools/tool_manager.py    |  8 +++++---
 5 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/src/airflow_mcp_server/client/airflow_client.py b/src/airflow_mcp_server/client/airflow_client.py
index 1014a3e..e1e5117 100644
--- a/src/airflow_mcp_server/client/airflow_client.py
+++ b/src/airflow_mcp_server/client/airflow_client.py
@@ -102,11 +102,11 @@ class AirflowClient:
             self.base_url = base_url.rstrip("/")
             self.headers = {"Accept": "application/json"}
 
-            # Set authentication header based on what was provided
-            if auth_token:
-                self.headers["Authorization"] = f"Basic {auth_token}"
-            elif cookie:
+            # Set authentication header based on precedence (cookie > auth_token)
+            if cookie:
                 self.headers["Cookie"] = cookie
+            elif auth_token:
+                self.headers["Authorization"] = f"Basic {auth_token}"
 
         except Exception as e:
             logger.error("Failed to initialize AirflowClient: %s", e)
diff --git a/src/airflow_mcp_server/server.py b/src/airflow_mcp_server/server.py
index 87d3a94..87f35e2 100644
--- a/src/airflow_mcp_server/server.py
+++ b/src/airflow_mcp_server/server.py
@@ -21,7 +21,17 @@ logger = logging.getLogger(__name__)
 
 
 async def serve() -> None:
-    """Start MCP server."""
+    """Start MCP server.
+
+    Configuration precedence:
+    1. Environment variables (highest)
+    2. Command line arguments (if applicable)
+    3. Default values (lowest)
+
+    For authentication:
+    1. Cookie authentication (highest)
+    2. Auth token authentication (secondary)
+    """
     # Check for AIRFLOW_BASE_URL which is always required
     if "AIRFLOW_BASE_URL" not in os.environ:
         raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
diff --git a/src/airflow_mcp_server/server_safe.py b/src/airflow_mcp_server/server_safe.py
index 6be1f36..70a0ecf 100644
--- a/src/airflow_mcp_server/server_safe.py
+++ b/src/airflow_mcp_server/server_safe.py
@@ -12,7 +12,17 @@ logger = logging.getLogger(__name__)
 
 
 async def serve() -> None:
-    """Start MCP server in safe mode (read-only operations)."""
+    """Start MCP server in safe mode (read-only operations).
+
+    Configuration precedence:
+    1. Environment variables (highest)
+    2. Command line arguments (if applicable)
+    3. Default values (lowest)
+
+    For authentication:
+    1. Cookie authentication (highest)
+    2. Auth token authentication (secondary)
+    """
     # Check for AIRFLOW_BASE_URL which is always required
     if "AIRFLOW_BASE_URL" not in os.environ:
         raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
diff --git a/src/airflow_mcp_server/server_unsafe.py b/src/airflow_mcp_server/server_unsafe.py
index b33b46c..a47cfc1 100644
--- a/src/airflow_mcp_server/server_unsafe.py
+++ b/src/airflow_mcp_server/server_unsafe.py
@@ -12,7 +12,17 @@ logger = logging.getLogger(__name__)
 
 
 async def serve() -> None:
-    """Start MCP server in unsafe mode (all operations)."""
+    """Start MCP server in unsafe mode (all operations).
+
+    Configuration precedence:
+    1. Environment variables (highest)
+    2. Command line arguments (if applicable)
+    3. Default values (lowest)
+
+    For authentication:
+    1. Cookie authentication (highest)
+    2. Auth token authentication (secondary)
+    """
     # Check for AIRFLOW_BASE_URL which is always required
     if "AIRFLOW_BASE_URL" not in os.environ:
         raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
diff --git a/src/airflow_mcp_server/tools/tool_manager.py b/src/airflow_mcp_server/tools/tool_manager.py
index 9eb9cbe..8106cd3 100644
--- a/src/airflow_mcp_server/tools/tool_manager.py
+++ b/src/airflow_mcp_server/tools/tool_manager.py
@@ -46,10 +46,12 @@ def _initialize_client() -> AirflowClient:
     # Initialize client with appropriate authentication method
     client_args = {"spec_path": spec_path, "base_url": os.environ["AIRFLOW_BASE_URL"]}
 
-    if has_auth_token:
-        client_args["auth_token"] = os.environ["AUTH_TOKEN"]
-    elif has_cookie:
+    # Apply cookie auth first if available (highest precedence)
+    if has_cookie:
         client_args["cookie"] = os.environ["COOKIE"]
+    # Otherwise use auth token if available
+    elif has_auth_token:
+        client_args["auth_token"] = os.environ["AUTH_TOKEN"]
 
     return AirflowClient(**client_args)
 

From 420b6fc68f8692774fee61acffeed6bef1a5741a Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 06:10:28 +0000
Subject: [PATCH 5/9] safe and unsafe are mutually exclusive

---
 src/airflow_mcp_server/__init__.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/airflow_mcp_server/__init__.py b/src/airflow_mcp_server/__init__.py
index 8decf7f..3b90bcb 100644
--- a/src/airflow_mcp_server/__init__.py
+++ b/src/airflow_mcp_server/__init__.py
@@ -22,12 +22,15 @@ def main(verbose: int, safe: bool, unsafe: bool) -> None:
 
     logging.basicConfig(level=logging_level, stream=sys.stderr)
 
+    # Determine server mode with proper precedence
     if safe and unsafe:
+        # CLI argument validation
         raise click.UsageError("Options --safe and --unsafe are mutually exclusive")
-
-    if safe:
+    elif safe:
+        # CLI argument for safe mode
         asyncio.run(serve_safe())
-    else:  # Default to unsafe mode
+    else:
+        # Default to unsafe mode
         asyncio.run(serve_unsafe())
 
 

From ea60acd54a80497600c70229e1387407368c1271 Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 08:53:37 +0000
Subject: [PATCH 6/9] Add AirflowConfig class for centralized configuration
 management

---
 src/airflow_mcp_server/__init__.py           | 26 ++++++++--
 src/airflow_mcp_server/config.py             | 25 +++++++++
 src/airflow_mcp_server/server.py             | 29 +++--------
 src/airflow_mcp_server/server_safe.py        | 29 +++--------
 src/airflow_mcp_server/server_unsafe.py      | 29 +++--------
 src/airflow_mcp_server/tools/tool_manager.py | 53 +++++++++-----------
 6 files changed, 91 insertions(+), 100 deletions(-)
 create mode 100644 src/airflow_mcp_server/config.py

diff --git a/src/airflow_mcp_server/__init__.py b/src/airflow_mcp_server/__init__.py
index 3b90bcb..fe6efa2 100644
--- a/src/airflow_mcp_server/__init__.py
+++ b/src/airflow_mcp_server/__init__.py
@@ -1,9 +1,11 @@
 import asyncio
 import logging
+import os
 import sys
 
 import click
 
+from airflow_mcp_server.config import AirflowConfig
 from airflow_mcp_server.server_safe import serve as serve_safe
 from airflow_mcp_server.server_unsafe import serve as serve_unsafe
 
@@ -12,7 +14,11 @@ from airflow_mcp_server.server_unsafe import serve as serve_unsafe
 @click.option("-v", "--verbose", count=True, help="Increase verbosity")
 @click.option("--safe", "-s", is_flag=True, help="Use only read-only tools")
 @click.option("--unsafe", "-u", is_flag=True, help="Use all tools (default)")
-def main(verbose: int, safe: bool, unsafe: bool) -> None:
+@click.option("--base-url", help="Airflow API base URL")
+@click.option("--spec-path", help="Path to OpenAPI spec file")
+@click.option("--auth-token", help="Authentication token")
+@click.option("--cookie", help="Session cookie")
+def main(verbose: int, safe: bool, unsafe: bool, base_url: str = None, spec_path: str = None, auth_token: str = None, cookie: str = None) -> None:
     """MCP server for Airflow"""
     logging_level = logging.WARN
     if verbose == 1:
@@ -22,16 +28,30 @@ def main(verbose: int, safe: bool, unsafe: bool) -> None:
 
     logging.basicConfig(level=logging_level, stream=sys.stderr)
 
+    # Read environment variables with proper precedence
+    # Environment variables take precedence over CLI arguments
+    config_base_url = os.environ.get("AIRFLOW_BASE_URL") or base_url
+    config_spec_path = os.environ.get("AIRFLOW_SPEC_PATH") or spec_path
+    config_auth_token = os.environ.get("AIRFLOW_AUTH_TOKEN") or auth_token
+    config_cookie = os.environ.get("AIRFLOW_COOKIE") or cookie
+
+    # Initialize configuration
+    try:
+        config = AirflowConfig(base_url=config_base_url, spec_path=config_spec_path, auth_token=config_auth_token, cookie=config_cookie)
+    except ValueError as e:
+        click.echo(f"Configuration error: {e}", err=True)
+        sys.exit(1)
+
     # Determine server mode with proper precedence
     if safe and unsafe:
         # CLI argument validation
         raise click.UsageError("Options --safe and --unsafe are mutually exclusive")
     elif safe:
         # CLI argument for safe mode
-        asyncio.run(serve_safe())
+        asyncio.run(serve_safe(config))
     else:
         # Default to unsafe mode
-        asyncio.run(serve_unsafe())
+        asyncio.run(serve_unsafe(config))
 
 
 if __name__ == "__main__":
diff --git a/src/airflow_mcp_server/config.py b/src/airflow_mcp_server/config.py
new file mode 100644
index 0000000..f3bb9d2
--- /dev/null
+++ b/src/airflow_mcp_server/config.py
@@ -0,0 +1,25 @@
+class AirflowConfig:
+    """Centralized configuration for Airflow MCP server."""
+
+    def __init__(self, base_url: str | None = None, spec_path: str | None = None, auth_token: str | None = None, cookie: str | None = None) -> None:
+        """Initialize configuration with provided values.
+
+        Args:
+            base_url: Airflow API base URL
+            spec_path: Path to OpenAPI spec file
+            auth_token: Authentication token
+            cookie: Session cookie
+
+        Raises:
+            ValueError: If required configuration is missing
+        """
+        self.base_url = base_url
+        if not self.base_url:
+            raise ValueError("Missing required configuration: base_url")
+
+        self.spec_path = spec_path
+        self.auth_token = auth_token
+        self.cookie = cookie
+
+        if not self.auth_token and not self.cookie:
+            raise ValueError("Either auth_token or cookie must be provided")
diff --git a/src/airflow_mcp_server/server.py b/src/airflow_mcp_server/server.py
index 87f35e2..c9063b4 100644
--- a/src/airflow_mcp_server/server.py
+++ b/src/airflow_mcp_server/server.py
@@ -1,11 +1,11 @@
 import logging
-import os
 from typing import Any
 
 from mcp.server import Server
 from mcp.server.stdio import stdio_server
 from mcp.types import TextContent, Tool
 
+from airflow_mcp_server.config import AirflowConfig
 from airflow_mcp_server.tools.tool_manager import get_airflow_tools, get_tool
 
 # ===========THIS IS FOR DEBUGGING WITH MCP INSPECTOR===================
@@ -20,35 +20,18 @@ from airflow_mcp_server.tools.tool_manager import get_airflow_tools, get_tool
 logger = logging.getLogger(__name__)
 
 
-async def serve() -> None:
+async def serve(config: AirflowConfig) -> None:
     """Start MCP server.
 
-    Configuration precedence:
-    1. Environment variables (highest)
-    2. Command line arguments (if applicable)
-    3. Default values (lowest)
-
-    For authentication:
-    1. Cookie authentication (highest)
-    2. Auth token authentication (secondary)
+    Args:
+        config: Configuration object with auth and URL settings
     """
-    # Check for AIRFLOW_BASE_URL which is always required
-    if "AIRFLOW_BASE_URL" not in os.environ:
-        raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
-
-    # Check for either AUTH_TOKEN or COOKIE
-    has_auth_token = "AUTH_TOKEN" in os.environ
-    has_cookie = "COOKIE" in os.environ
-
-    if not has_auth_token and not has_cookie:
-        raise ValueError("Either AUTH_TOKEN or COOKIE environment variable must be provided")
-
     server = Server("airflow-mcp-server")
 
     @server.list_tools()
     async def list_tools() -> list[Tool]:
         try:
-            return await get_airflow_tools()
+            return await get_airflow_tools(config)
         except Exception as e:
             logger.error("Failed to list tools: %s", e)
             raise
@@ -56,7 +39,7 @@ async def serve() -> None:
     @server.call_tool()
     async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
         try:
-            tool = await get_tool(name)
+            tool = await get_tool(config, name)
             async with tool.client:
                 result = await tool.run(body=arguments)
             return [TextContent(type="text", text=str(result))]
diff --git a/src/airflow_mcp_server/server_safe.py b/src/airflow_mcp_server/server_safe.py
index 70a0ecf..543b3b5 100644
--- a/src/airflow_mcp_server/server_safe.py
+++ b/src/airflow_mcp_server/server_safe.py
@@ -1,45 +1,28 @@
 import logging
-import os
 from typing import Any
 
 from mcp.server import Server
 from mcp.server.stdio import stdio_server
 from mcp.types import TextContent, Tool
 
+from airflow_mcp_server.config import AirflowConfig
 from airflow_mcp_server.tools.tool_manager import get_airflow_tools, get_tool
 
 logger = logging.getLogger(__name__)
 
 
-async def serve() -> None:
+async def serve(config: AirflowConfig) -> None:
     """Start MCP server in safe mode (read-only operations).
 
-    Configuration precedence:
-    1. Environment variables (highest)
-    2. Command line arguments (if applicable)
-    3. Default values (lowest)
-
-    For authentication:
-    1. Cookie authentication (highest)
-    2. Auth token authentication (secondary)
+    Args:
+        config: Configuration object with auth and URL settings
     """
-    # Check for AIRFLOW_BASE_URL which is always required
-    if "AIRFLOW_BASE_URL" not in os.environ:
-        raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
-
-    # Check for either AUTH_TOKEN or COOKIE
-    has_auth_token = "AUTH_TOKEN" in os.environ
-    has_cookie = "COOKIE" in os.environ
-
-    if not has_auth_token and not has_cookie:
-        raise ValueError("Either AUTH_TOKEN or COOKIE environment variable must be provided")
-
     server = Server("airflow-mcp-server-safe")
 
     @server.list_tools()
     async def list_tools() -> list[Tool]:
         try:
-            return await get_airflow_tools(mode="safe")
+            return await get_airflow_tools(config, mode="safe")
         except Exception as e:
             logger.error("Failed to list tools: %s", e)
             raise
@@ -49,7 +32,7 @@ async def serve() -> None:
         try:
             if not name.startswith("get_"):
                 raise ValueError("Only GET operations allowed in safe mode")
-            tool = await get_tool(name)
+            tool = await get_tool(config, name)
             async with tool.client:
                 result = await tool.run(body=arguments)
             return [TextContent(type="text", text=str(result))]
diff --git a/src/airflow_mcp_server/server_unsafe.py b/src/airflow_mcp_server/server_unsafe.py
index a47cfc1..347db5c 100644
--- a/src/airflow_mcp_server/server_unsafe.py
+++ b/src/airflow_mcp_server/server_unsafe.py
@@ -1,45 +1,28 @@
 import logging
-import os
 from typing import Any
 
 from mcp.server import Server
 from mcp.server.stdio import stdio_server
 from mcp.types import TextContent, Tool
 
+from airflow_mcp_server.config import AirflowConfig
 from airflow_mcp_server.tools.tool_manager import get_airflow_tools, get_tool
 
 logger = logging.getLogger(__name__)
 
 
-async def serve() -> None:
+async def serve(config: AirflowConfig) -> None:
     """Start MCP server in unsafe mode (all operations).
 
-    Configuration precedence:
-    1. Environment variables (highest)
-    2. Command line arguments (if applicable)
-    3. Default values (lowest)
-
-    For authentication:
-    1. Cookie authentication (highest)
-    2. Auth token authentication (secondary)
+    Args:
+        config: Configuration object with auth and URL settings
     """
-    # Check for AIRFLOW_BASE_URL which is always required
-    if "AIRFLOW_BASE_URL" not in os.environ:
-        raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
-
-    # Check for either AUTH_TOKEN or COOKIE
-    has_auth_token = "AUTH_TOKEN" in os.environ
-    has_cookie = "COOKIE" in os.environ
-
-    if not has_auth_token and not has_cookie:
-        raise ValueError("Either AUTH_TOKEN or COOKIE environment variable must be provided")
-
     server = Server("airflow-mcp-server-unsafe")
 
     @server.list_tools()
     async def list_tools() -> list[Tool]:
         try:
-            return await get_airflow_tools(mode="unsafe")
+            return await get_airflow_tools(config, mode="unsafe")
         except Exception as e:
             logger.error("Failed to list tools: %s", e)
             raise
@@ -47,7 +30,7 @@ async def serve() -> None:
     @server.call_tool()
     async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
         try:
-            tool = await get_tool(name)
+            tool = await get_tool(config, name)
             async with tool.client:
                 result = await tool.run(body=arguments)
             return [TextContent(type="text", text=str(result))]
diff --git a/src/airflow_mcp_server/tools/tool_manager.py b/src/airflow_mcp_server/tools/tool_manager.py
index 8106cd3..3588cc7 100644
--- a/src/airflow_mcp_server/tools/tool_manager.py
+++ b/src/airflow_mcp_server/tools/tool_manager.py
@@ -1,10 +1,10 @@
 import logging
-import os
 from importlib import resources
 
 from mcp.types import Tool
 
 from airflow_mcp_server.client.airflow_client import AirflowClient
+from airflow_mcp_server.config import AirflowConfig
 from airflow_mcp_server.parser.operation_parser import OperationParser
 from airflow_mcp_server.tools.airflow_tool import AirflowTool
 
@@ -13,16 +13,19 @@ logger = logging.getLogger(__name__)
 _tools_cache: dict[str, AirflowTool] = {}
 
 
-def _initialize_client() -> AirflowClient:
-    """Initialize Airflow client with environment variables or embedded spec.
+def _initialize_client(config: AirflowConfig) -> AirflowClient:
+    """Initialize Airflow client with configuration.
+
+    Args:
+        config: Configuration object with auth and URL settings
 
     Returns:
         AirflowClient instance
 
     Raises:
-        ValueError: If required environment variables are missing or default spec is not found
+        ValueError: If default spec is not found
     """
-    spec_path = os.environ.get("OPENAPI_SPEC")
+    spec_path = config.spec_path
     if not spec_path:
         # Fallback to embedded v1.yaml
         try:
@@ -32,41 +35,33 @@ def _initialize_client() -> AirflowClient:
         except Exception as e:
             raise ValueError("Default OpenAPI spec not found in package resources") from e
 
-    # Check for base URL
-    if "AIRFLOW_BASE_URL" not in os.environ:
-        raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")
-
-    # Check for either AUTH_TOKEN or COOKIE
-    has_auth_token = "AUTH_TOKEN" in os.environ
-    has_cookie = "COOKIE" in os.environ
-
-    if not has_auth_token and not has_cookie:
-        raise ValueError("Either AUTH_TOKEN or COOKIE environment variable must be provided")
-
     # Initialize client with appropriate authentication method
-    client_args = {"spec_path": spec_path, "base_url": os.environ["AIRFLOW_BASE_URL"]}
+    client_args = {"spec_path": spec_path, "base_url": config.base_url}
 
     # Apply cookie auth first if available (highest precedence)
-    if has_cookie:
-        client_args["cookie"] = os.environ["COOKIE"]
+    if config.cookie:
+        client_args["cookie"] = config.cookie
     # Otherwise use auth token if available
-    elif has_auth_token:
-        client_args["auth_token"] = os.environ["AUTH_TOKEN"]
+    elif config.auth_token:
+        client_args["auth_token"] = config.auth_token
 
     return AirflowClient(**client_args)
 
 
-async def _initialize_tools() -> None:
+async def _initialize_tools(config: AirflowConfig) -> None:
     """Initialize tools cache with Airflow operations.
 
+    Args:
+        config: Configuration object with auth and URL settings
+
     Raises:
         ValueError: If initialization fails
     """
     global _tools_cache
 
     try:
-        client = _initialize_client()
-        spec_path = os.environ.get("OPENAPI_SPEC")
+        client = _initialize_client(config)
+        spec_path = config.spec_path
         if not spec_path:
             with resources.files("airflow_mcp_server.resources").joinpath("v1.yaml").open("rb") as f:
                 spec_path = f.name
@@ -84,10 +79,11 @@ async def _initialize_tools() -> None:
         raise ValueError(f"Failed to initialize tools: {e}") from e
 
 
-async def get_airflow_tools(mode: str = "unsafe") -> list[Tool]:
+async def get_airflow_tools(config: AirflowConfig, mode: str = "unsafe") -> list[Tool]:
     """Get list of available Airflow tools based on mode.
 
     Args:
+        config: Configuration object with auth and URL settings
         mode: "safe" for GET operations only, "unsafe" for all operations (default)
 
     Returns:
@@ -97,7 +93,7 @@ async def get_airflow_tools(mode: str = "unsafe") -> list[Tool]:
         ValueError: If initialization fails
     """
     if not _tools_cache:
-        await _initialize_tools()
+        await _initialize_tools(config)
 
     tools = []
     for operation_id, tool in _tools_cache.items():
@@ -120,10 +116,11 @@ async def get_airflow_tools(mode: str = "unsafe") -> list[Tool]:
     return tools
 
 
-async def get_tool(name: str) -> AirflowTool:
+async def get_tool(config: AirflowConfig, name: str) -> AirflowTool:
     """Get specific tool by name.
 
     Args:
+        config: Configuration object with auth and URL settings
         name: Tool/operation name
 
     Returns:
@@ -134,7 +131,7 @@ async def get_tool(name: str) -> AirflowTool:
         ValueError: If tool initialization fails
     """
     if not _tools_cache:
-        await _initialize_tools()
+        await _initialize_tools(config)
 
     if name not in _tools_cache:
         raise KeyError(f"Tool {name} not found")

From 492e79ef2a720a84460bb5eb9520e1fd2b1607ab Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 11:12:45 +0000
Subject: [PATCH 7/9] uv lock update

---
 uv.lock | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/uv.lock b/uv.lock
index 15e19b1..362c5e4 100644
--- a/uv.lock
+++ b/uv.lock
@@ -111,7 +111,7 @@ wheels = [
 
 [[package]]
 name = "airflow-mcp-server"
-version = "0.2.0"
+version = "0.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles" },

From 679523a7c6b81c6abf9258294aa155e35c1914ea Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 11:18:45 +0000
Subject: [PATCH 8/9] Using older env variables

---
 src/airflow_mcp_server/__init__.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/airflow_mcp_server/__init__.py b/src/airflow_mcp_server/__init__.py
index fe6efa2..4f860f2 100644
--- a/src/airflow_mcp_server/__init__.py
+++ b/src/airflow_mcp_server/__init__.py
@@ -31,9 +31,9 @@ def main(verbose: int, safe: bool, unsafe: bool, base_url: str = None, spec_path
     # Read environment variables with proper precedence
     # Environment variables take precedence over CLI arguments
     config_base_url = os.environ.get("AIRFLOW_BASE_URL") or base_url
-    config_spec_path = os.environ.get("AIRFLOW_SPEC_PATH") or spec_path
-    config_auth_token = os.environ.get("AIRFLOW_AUTH_TOKEN") or auth_token
-    config_cookie = os.environ.get("AIRFLOW_COOKIE") or cookie
+    config_spec_path = os.environ.get("OPENAPI_SPEC") or spec_path
+    config_auth_token = os.environ.get("AUTH_TOKEN") or auth_token
+    config_cookie = os.environ.get("COOKIE") or cookie
 
     # Initialize configuration
     try:

From d8887d3a2bf886bd3784add0e69bd97cc6b85198 Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Tue, 25 Feb 2025 11:20:09 +0000
Subject: [PATCH 9/9] Task update

---
 README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/README.md b/README.md
index f35aebf..bc7de51 100644
--- a/README.md
+++ b/README.md
@@ -81,6 +81,7 @@ The default is 100 items, but you can change it using `maximum_page_limit` optio
 - [x] First API
 - [x] Parse OpenAPI Spec
 - [x] Safe/Unsafe mode implementation
+- [x] Allow session auth
 - [ ] Parse proper description with list_tools.
 - [ ] Airflow config fetch (_specifically for page limit_)
 - [ ] Env variables optional (_env variables might not be ideal for airflow plugins_)