From f58a10f0997f5f415989711f0d323cdc40d2eb2e Mon Sep 17 00:00:00 2001
From: jiasli <4003950+jiasli@users.noreply.github.com>
Date: Tue, 15 Apr 2025 17:19:02 +0800
Subject: [PATCH] enable_broker_on_wsl

---
 src/azure-cli-core/azure/cli/core/_profile.py      | 3 +++
 src/azure-cli-core/azure/cli/core/auth/identity.py | 5 ++++-
 src/azure-cli/requirements.py3.Linux.txt           | 1 +
 src/azure-cli/requirements.py3.windows.txt         | 2 +-
 4 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py
index cd977628510..f0eb1897e2c 100644
--- a/src/azure-cli-core/azure/cli/core/_profile.py
+++ b/src/azure-cli-core/azure/cli/core/_profile.py
@@ -1009,6 +1009,8 @@ def _create_identity_instance(cli_ctx, authority, tenant_id=None, client_id=None
     from .telemetry import set_broker_info
     set_broker_info(enable_broker_on_windows)
 
+    # On WSL, use core.enable_broker_on_wsl=true to use broker (WAM)
+    enable_broker_on_wsl = cli_ctx.config.getboolean('core', 'enable_broker_on_wsl', fallback=False)
     # PREVIEW: In Azure Stack environment, use core.instance_discovery=false to disable MSAL's instance discovery.
     instance_discovery = cli_ctx.config.getboolean('core', 'instance_discovery', True)
 
@@ -1016,6 +1018,7 @@ def _create_identity_instance(cli_ctx, authority, tenant_id=None, client_id=None
                     encrypt=encrypt,
                     use_msal_http_cache=use_msal_http_cache,
                     enable_broker_on_windows=enable_broker_on_windows,
+                    enable_broker_on_wsl=enable_broker_on_wsl,
                     instance_discovery=instance_discovery)
 
 
diff --git a/src/azure-cli-core/azure/cli/core/auth/identity.py b/src/azure-cli-core/azure/cli/core/auth/identity.py
index 6f86bee804c..2305f62d72c 100644
--- a/src/azure-cli-core/azure/cli/core/auth/identity.py
+++ b/src/azure-cli-core/azure/cli/core/auth/identity.py
@@ -58,7 +58,8 @@ class Identity:  # pylint: disable=too-many-instance-attributes
     _service_principal_store_instance = None
 
     def __init__(self, authority, tenant_id=None, client_id=None, encrypt=False, use_msal_http_cache=True,
-                 enable_broker_on_windows=None, instance_discovery=None):
+                 enable_broker_on_windows=None, enable_broker_on_wsl=None,
+                 instance_discovery=None):
         """
         :param authority: Authentication authority endpoint. For example,
             - AAD: https://login.microsoftonline.com
@@ -74,6 +75,7 @@ def __init__(self, authority, tenant_id=None, client_id=None, encrypt=False, use
         self._encrypt = encrypt
         self._use_msal_http_cache = use_msal_http_cache
         self._enable_broker_on_windows = enable_broker_on_windows
+        self._enable_broker_on_wsl = enable_broker_on_wsl
         self._instance_discovery = instance_discovery
 
         # Build the authority in MSAL style
@@ -114,6 +116,7 @@ def _msal_public_app_kwargs(self):
         # enable_broker_on_windows can only be used on PublicClientApplication.
         return {**self._msal_app_kwargs,
                 "enable_broker_on_windows": self._enable_broker_on_windows,
+                "enable_broker_on_wsl": self._enable_broker_on_wsl,
                 "enable_pii_log": True}
 
     @property
diff --git a/src/azure-cli/requirements.py3.Linux.txt b/src/azure-cli/requirements.py3.Linux.txt
index 0eadb119cfd..6156dabdd4d 100644
--- a/src/azure-cli/requirements.py3.Linux.txt
+++ b/src/azure-cli/requirements.py3.Linux.txt
@@ -117,6 +117,7 @@ psutil==6.1.0
 pycomposefile==0.0.32
 PyGithub==1.55
 PyJWT==2.4.0
+pymsalruntime==0.18.1
 PyNaCl==1.5.0
 pyOpenSSL==25.0.0
 python-dateutil==2.8.0
diff --git a/src/azure-cli/requirements.py3.windows.txt b/src/azure-cli/requirements.py3.windows.txt
index bfd1e18ec3b..e2f97b99564 100644
--- a/src/azure-cli/requirements.py3.windows.txt
+++ b/src/azure-cli/requirements.py3.windows.txt
@@ -116,7 +116,7 @@ psutil==6.1.0
 pycomposefile==0.0.32
 PyGithub==1.55
 PyJWT==2.4.0
-pymsalruntime==0.16.2
+pymsalruntime==0.18.1
 PyNaCl==1.5.0
 pyOpenSSL==25.0.0
 python-dateutil==2.8.0