package name header update for tool name in logs

modelcontextprotocol/middleware.py

@@ -100,7 +100,6 @@ async def on_call_tool(self, context: MiddlewareContext, call_next):
 
             # Tool is allowed, proceed with execution
             logger.debug(f"Tool access granted: {tool_name}", tool=tool_name)
-
             return await call_next(context)
 
         except ToolError:

modelcontextprotocol/settings.py

@@ -13,6 +13,7 @@ class Settings(BaseSettings):
     ATLAN_AGENT_ID: str = "NA"
     ATLAN_AGENT: str = "atlan-mcp"
     ATLAN_MCP_USER_AGENT: str = f"Atlan MCP Server {MCP_VERSION}"
+    ATLAN_TOOL_NAME: str = "X-Atlan-Mcp-Tool"
     MCP_TRANSPORT: str = "stdio"
     MCP_HOST: str = "0.0.0.0"
     MCP_PORT: int = 8000

modelcontextprotocol/tools/assets.py

@@ -10,6 +10,7 @@
 )
 from pyatlan.model.assets import Readme, AtlasGlossaryTerm
 from pyatlan.model.fluent_search import CompoundQuery, FluentSearch
+from settings import get_settings
 
 # Initialize logging
 logger = logging.getLogger(__name__)
@@ -47,7 +48,6 @@ def update_assets(
         logger.info(
             f"Updating {len(updatable_assets)} assets with attribute '{attribute_name}'"
         )
-
         # Validate attribute values
         if len(updatable_assets) != len(attribute_values):
             error_msg = "Number of asset GUIDs must match number of attribute values"
@@ -67,7 +67,8 @@ def update_assets(
 
         # Get Atlan client
         client = get_atlan_client()
-
+        settings = get_settings()
+        client.update_headers({settings.ATLAN_TOOL_NAME: "update_assets_tool"})
         # Create assets with updated values
         assets = []
         # readme_update_parent_assets: Assets that were updated with readme.

modelcontextprotocol/tools/dsl.py

@@ -5,6 +5,8 @@
 from client import get_atlan_client
 from pyatlan.model.search import DSL, IndexSearchRequest
 from utils.search import SearchUtils
+from settings import get_settings
+
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -19,6 +21,7 @@ def get_assets_by_dsl(dsl_query: Union[str, Dict[str, Any]]) -> Dict[str, Any]:
         Dict[str, Any]: A dictionary containing the results and aggregations
     """
     logger.info("Starting DSL-based asset search")
+
     try:
         # Parse string to dict if needed
         if isinstance(dsl_query, str):
@@ -47,6 +50,8 @@ def get_assets_by_dsl(dsl_query: Union[str, Dict[str, Any]]) -> Dict[str, Any]:
 
         logger.info("Executing DSL search request")
         client = get_atlan_client()
+        settings = get_settings()
+        client.update_headers({settings.ATLAN_TOOL_NAME: "get_assets_by_dsl_tool"})
         search_response = client.asset.search(index_request)
         processed_results = SearchUtils.process_results(search_response)
         return processed_results

modelcontextprotocol/tools/glossary.py

@@ -16,11 +16,12 @@
     GlossaryCategory,
     GlossaryTerm,
 )
+from settings import get_settings
 
 logger = logging.getLogger(__name__)
 
 
-def save_assets(assets: List[Asset]) -> List[Dict[str, Any]]:
+def save_assets(assets: List[Asset], tool_name: str) -> List[Dict[str, Any]]:
     """
     Common bulk save and response processing for any asset type.
 
@@ -34,7 +35,10 @@ def save_assets(assets: List[Asset]) -> List[Dict[str, Any]]:
         Exception: If there's an error saving the assets.
     """
     logger.info("Starting bulk save operation")
+
     client = get_atlan_client()
+    settings = get_settings()
+    client.update_headers({settings.ATLAN_TOOL_NAME: tool_name})
     try:
         response = client.asset.save(assets)
     except Exception as e:
@@ -85,6 +89,7 @@ def create_glossary_assets(
     """
     data = glossaries if isinstance(glossaries, list) else [glossaries]
     logger.info(f"Creating {len(data)} glossary asset(s)")
+
     logger.debug(f"Glossary specifications: {data}")
 
     specs = [Glossary(**item) for item in data]
@@ -104,7 +109,7 @@ def create_glossary_assets(
             logger.debug(f"Set certificate status for {spec.name}: {cs.value}")
         assets.append(glossary)
 
-    return save_assets(assets)
+    return save_assets(assets, "create_glossaries")
 
 
 def create_glossary_category_assets(
@@ -137,6 +142,7 @@ def create_glossary_category_assets(
     """
     data = categories if isinstance(categories, list) else [categories]
     logger.info(f"Creating {len(data)} glossary category asset(s)")
+
     logger.debug(f"Category specifications: {data}")
 
     specs = [GlossaryCategory(**item) for item in data]
@@ -166,7 +172,7 @@ def create_glossary_category_assets(
 
         assets.append(category)
 
-    return save_assets(assets)
+    return save_assets(assets, "create_glossary_categories")
 
 
 def create_glossary_term_assets(
@@ -222,4 +228,4 @@ def create_glossary_term_assets(
             term.certificate_status = cs.value
         assets.append(term)
 
-    return save_assets(assets)
+    return save_assets(assets, "create_glossary_terms")

modelcontextprotocol/tools/lineage.py

@@ -7,6 +7,7 @@
 from pyatlan.model.fields.atlan_fields import AtlanField
 from utils.search import SearchUtils
 from utils.constants import DEFAULT_SEARCH_ATTRIBUTES
+from settings import get_settings
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -89,6 +90,8 @@ def traverse_lineage(
 
         logger.info("Executing lineage request")
         client = get_atlan_client()
+        settings = get_settings()
+        client.update_headers({settings.ATLAN_TOOL_NAME: "traverse_lineage_tool"})
         response = client.asset.get_lineage_list(request)
 
         # Process results using same pattern as search

modelcontextprotocol/tools/query.py

@@ -10,6 +10,7 @@
 
 from client import get_atlan_client
 from pyatlan.model.query import QueryRequest
+from settings import get_settings
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -78,7 +79,8 @@ def query_asset(
         # Get Atlan client
         logger.debug("Getting Atlan client")
         client = get_atlan_client()
-
+        settings = get_settings()
+        client.update_headers({settings.ATLAN_TOOL_NAME: "query_asset_tool"})
         # Build query request
         logger.debug("Building QueryRequest object")
         query_request = QueryRequest(

modelcontextprotocol/tools/search.py

@@ -7,6 +7,7 @@
 from pyatlan.model.fields.atlan_fields import AtlanField
 from utils.search import SearchUtils
 from utils.constants import DEFAULT_SEARCH_ATTRIBUTES
+from settings import get_settings
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -73,6 +74,8 @@ def search_assets(
         f"Starting asset search with parameters: asset_type={asset_type}, "
         f"limit={limit}, include_archived={include_archived}"
     )
+
+    # Set tool-specific headers
     logger.debug(
         f"Full search parameters: conditions={conditions}, "
         f"negative_conditions={negative_conditions}, some_conditions={some_conditions}, "
@@ -291,6 +294,8 @@ def search_assets(
 
         logger.info("Executing search request")
         client = get_atlan_client()
+        settings = get_settings()
+        client.update_headers({settings.ATLAN_TOOL_NAME: "search_assets_tool"})
         search_response = client.asset.search(request)
         processed_results = SearchUtils.process_results(search_response)
         logger.info(