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] 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)