Add request ID tracking to token validation for support correlation

packages/reflex-hosting-cli/news/6821.feature.md

@@ -0,0 +1 @@
+Token validation requests now send an `X-Request-ID` header, and failed validations print the ID (e.g. `Unable to validate access token: ... (auth request id: ...)`) so it can be quoted to support to correlate the failure with server-side logs.

packages/reflex-hosting-cli/src/reflex_cli/utils/exceptions.py

@@ -9,6 +9,28 @@ class NotAuthenticatedError(ReflexHostingCliError):
     """Raised when the user is not authenticated."""
 
 
+class TokenValidationError(ReflexHostingCliError):
+    """Raised when validating an access token with the control plane fails.
+
+    Carries the X-Request-ID sent with the validation request so callers can
+    surface it for correlation with server-side logs.
+    """
+
+    def __init__(self, message: str, request_id: str = ""):
+        """Initialize the error.
+
+        Args:
+            message: The error message.
+            request_id: The X-Request-ID header sent with the failed request.
+        """
+        super().__init__(message)
+        self.request_id = request_id
+
+
+class TokenAccessDeniedError(TokenValidationError, ValueError):
+    """Raised when the control plane rejects the access token."""
+
+
 class GetAppError(ReflexHostingCliError):
     """Raised when retrieving an app fails."""
 

packages/reflex-hosting-cli/src/reflex_cli/utils/hosting.py

@@ -32,6 +32,8 @@
     ResponseError,
     ScaleAppError,
     ScaleParamError,
+    TokenAccessDeniedError,
+    TokenValidationError,
 )
 
 
@@ -389,6 +391,24 @@ def is_reflex_enterprise_installed() -> bool:
         return True
 
 
+_last_auth_request_id: str = ""
+
+
+def get_auth_request_id() -> str:
+    """Get the request id sent with the most recent token validation request.
+
+    The id is sent to the control plane as the ``X-Request-ID`` header, so it
+    can be quoted to support to correlate a failed authentication with the
+    server-side logs.
+
+    Returns:
+        The request id of the last ``validate_token`` call, or an empty string
+        if no validation request has been made in this process.
+
+    """
+    return _last_auth_request_id
+
+
 def validate_token(token: str) -> dict[str, Any]:
     """Validate the token with the control plane.
 
@@ -399,12 +419,15 @@ def validate_token(token: str) -> dict[str, Any]:
         Information about the user associated with the token.
 
     Raises:
-        ValueError: if access denied.
-        Exception: if runs into timeout, failed requests, unexpected errors. These should be tried again.
+        TokenAccessDeniedError: if access denied.
+        TokenValidationError: if runs into timeout, failed requests, unexpected errors. These should be tried again.
 
     """
     import httpx
 
+    global _last_auth_request_id
+    request_id = _last_auth_request_id = uuid.uuid4().hex
+
     try:
         # Add reflex-enterprise detection flag as query parameter
         params = {
@@ -415,24 +438,28 @@ def validate_token(token: str) -> dict[str, Any]:
 
         response = httpx.post(
             urljoin(constants.Hosting.HOSTING_SERVICE, "/api/v1/authenticate/me"),
-            headers=authorization_header(token),
+            headers={**authorization_header(token), "X-Request-ID": request_id},
             params=params,
             timeout=constants.Hosting.TIMEOUT,
         )
         response.raise_for_status()
         return response.json()
     except httpx.RequestError as re:
-        console.debug(f"Request to auth server failed due to {re}")
-        raise Exception(str(re)) from re
+        console.debug(
+            f"Request to auth server failed due to {re} (request id: {request_id})"
+        )
+        raise TokenValidationError(str(re), request_id=request_id) from re
     except httpx.HTTPError as ex:
-        console.debug(f"Unable to validate the token due to: {ex}")
-        raise Exception("server error") from ex
+        console.debug(
+            f"Unable to validate the token due to: {ex} (request id: {request_id})"
+        )
+        raise TokenValidationError("server error", request_id=request_id) from ex
     except ValueError as ve:
-        console.debug("Access denied")
-        raise ValueError("access denied") from ve
+        console.debug(f"Access denied (request id: {request_id})")
+        raise TokenAccessDeniedError("access denied", request_id=request_id) from ve
     except Exception as ex:
-        console.debug(f"Unexpected error: {ex}")
-        raise Exception("internal errors") from ex
+        console.debug(f"Unexpected error: {ex} (request id: {request_id})")
+        raise TokenValidationError("internal errors", request_id=request_id) from ex
 
 
 def delete_token_from_config():
@@ -2486,11 +2513,16 @@ def validate_token_with_retries(access_token: str) -> dict[str, Any]:
     with console.status("Validating access token ..."):
         try:
             return validate_token(access_token)
-        except ValueError:
-            console.error("Access denied")
+        except ValueError as ex:
+            # getattr: mocks/foreign ValueErrors don't carry a request id.
+            request_id = getattr(ex, "request_id", "") or get_auth_request_id()
+            console.error(f"Access denied (auth request id: {request_id})")
             delete_token_from_config()
         except Exception as ex:
-            console.debug(f"Unable to validate token due to: {ex}")
+            request_id = getattr(ex, "request_id", "") or get_auth_request_id()
+            console.warn(
+                f"Unable to validate access token: {ex} (auth request id: {request_id})"
+            )
     return {}
 
 

tests/units/reflex_cli/utils/test_hosting.py

@@ -7,6 +7,7 @@
 import httpx
 import pytest
 from pytest_mock import MockerFixture, MockFixture
+from reflex_cli.utils.exceptions import TokenValidationError
 from reflex_cli.utils.hosting import (
     AuthenticatedClient,
     ScaleParams,
@@ -18,6 +19,7 @@
     delete_token_from_config,
     gcp_deploy_available,
     get_app_history,
+    get_auth_request_id,
     get_authenticated_client,
     get_existing_access_token,
     get_gcp_provider_status,
@@ -34,6 +36,8 @@
     set_app_provider,
     submit_security_review,
     update_deployment_description,
+    validate_token,
+    validate_token_with_retries,
 )
 
 _CLIENT = AuthenticatedClient(token="fake-token", validated_data={})
@@ -592,3 +596,58 @@ def test_get_app_history_includes_description_and_can_rollback(mocker: MockerFix
     assert history[0]["can rollback"] is True
     assert history[1]["description"] == ""
     assert history[1]["can rollback"] is False
+
+
+def test_validate_token_sends_request_id_header(mocker: MockerFixture):
+    """Each validation request carries a fresh X-Request-ID header."""
+    mock_post = mocker.patch(
+        "httpx.post", return_value=_ok(mocker, {"tier": "enterprise"})
+    )
+
+    assert validate_token("some-token") == {"tier": "enterprise"}
+
+    headers = mock_post.call_args.kwargs["headers"]
+    assert headers["X-API-TOKEN"] == "some-token"
+    assert headers["X-Request-ID"] == get_auth_request_id() != ""
+
+    first_request_id = get_auth_request_id()
+    validate_token("some-token")
+    assert get_auth_request_id() != first_request_id
+
+
+def test_validate_token_with_retries_warns_with_request_id(mocker: MockerFixture):
+    """A failed validation surfaces the request id for support correlation."""
+    mocker.patch("httpx.post", return_value=_error(mocker, 500, "boom"))
+    mock_warn = mocker.patch("reflex_cli.utils.hosting.console.warn")
+
+    assert validate_token_with_retries("some-token") == {}
+
+    request_id = get_auth_request_id()
+    assert request_id
+    assert request_id in mock_warn.call_args.args[0]
+
+
+def test_validate_token_with_retries_access_denied_reports_request_id(
+    mocker: MockerFixture,
+):
+    """The access denied error message includes the auth request id."""
+    response = mocker.Mock()
+    response.raise_for_status.return_value = None
+    response.json.side_effect = ValueError("bad json")
+    mocker.patch("httpx.post", return_value=response)
+    mock_error = mocker.patch("reflex_cli.utils.hosting.console.error")
+    mocker.patch("reflex_cli.utils.hosting.delete_token_from_config")
+
+    assert validate_token_with_retries("some-token") == {}
+
+    assert get_auth_request_id() in mock_error.call_args.args[0]
+
+
+def test_validate_token_failure_carries_request_id_on_exception(mocker: MockerFixture):
+    """Validation errors carry the request id of their own request."""
+    mocker.patch("httpx.post", return_value=_error(mocker, 500, "boom"))
+
+    with pytest.raises(TokenValidationError) as exc_info:
+        validate_token("some-token")
+
+    assert exc_info.value.request_id == get_auth_request_id() != ""