Fix the onboarding state to account for zenml login

src/zenml/zen_server/auth.py

@@ -37,7 +37,12 @@
     LOGIN,
     VERSION_1,
 )
-from zenml.enums import AuthScheme, ExecutionStatus, OAuthDeviceStatus
+from zenml.enums import (
+    AuthScheme,
+    ExecutionStatus,
+    OAuthDeviceStatus,
+    OnboardingStep,
+)
 from zenml.exceptions import (
     AuthorizationException,
     CredentialsNotValid,
@@ -630,12 +635,15 @@ def authenticate_device(client_id: UUID, device_code: str) -> AuthContext:
     return AuthContext(user=device_model.user, device=device_model)
 
 
-def authenticate_external_user(external_access_token: str) -> AuthContext:
+def authenticate_external_user(
+    external_access_token: str, request: Request
+) -> AuthContext:
     """Implement external authentication.
 
     Args:
         external_access_token: The access token used to authenticate the user
             to the external authenticator.
+        request: The request object.
 
     Returns:
         The authentication context reflecting the authenticated user.
@@ -761,6 +769,18 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
             )
             context.alias(user_id=external_user.id, previous_id=user.id)
 
+    # This is the best spot to update the onboarding state to mark the
+    # "zenml login" step as completed for ZenML Pro servers, because the
+    # user has just successfully logged in. However, we need to differentiate
+    # 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}
+        )
+
     return AuthContext(user=user)
 
 

src/zenml/zen_server/routers/auth_endpoints.py

@@ -287,7 +287,8 @@ def token(
             return OAuthRedirectResponse(authorization_url=authorization_url)
 
         auth_context = authenticate_external_user(
-            external_access_token=external_access_token
+            external_access_token=external_access_token,
+            request=request,
         )
 
     else:

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)
@@ -4218,6 +4223,11 @@ def session(self) -> requests.Session:
             self._session.mount("https://", HTTPAdapter(max_retries=retries))
             self._session.mount("http://", HTTPAdapter(max_retries=retries))
             self._session.verify = self.config.verify_ssl
+            # Use a custom user agent to identify the ZenML client in the server
+            # logs.
+            self._session.headers.update(
+                {"User-Agent": "zenml/" + zenml.__version__}
+            )
 
         # Note that we return an unauthenticated session here. An API token
         # is only fetched and set in the authorization header when and if it is