Fix the client user-agent headers

zenml-io · Dec 20, 2024 · 83564b2 · 83564b2
1 parent eb86cd5
commit 83564b2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/zenml/zen_server/auth.py b/src/zenml/zen_server/auth.py
@@ -775,6 +775,7 @@ def authenticate_external_user(
     # between web clients (i.e. the dashboard) and CLI clients (i.e. the
     # zenml CLI).
     user_agent = request.headers.get("User-Agent", "").lower()
+    logger.info(f"User agent: {user_agent}")
     if "zenml/" in user_agent:
         store.update_onboarding_state(
             completed_steps={OnboardingStep.DEVICE_VERIFIED}

diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py
@@ -4053,7 +4053,12 @@ def get_or_generate_api_token(self) -> str:
                 )
 
             data: Optional[Dict[str, str]] = None
-            headers: Dict[str, str] = {}
+
+            # Use a custom user agent to identify the ZenML client in the server
+            # logs.
+            headers: Dict[str, str] = {
+                "User-Agent": "zenml/" + zenml.__version__,
+            }
 
             # Check if an API key is configured
             api_key = credentials_store.get_api_key(self.url)