PR feedback

Author: Avery-Dunn

msal/application.py

@@ -2617,12 +2617,18 @@ def acquire_token_by_user_federated_identity_credential(
 
         telemetry_context = self._build_telemetry_context(
             self.ACQUIRE_TOKEN_BY_USER_FIC_ID)
+        headers = telemetry_context.generate_headers()
+        if username:
+            headers["X-AnchorMailbox"] = "upn:{}".format(username)
+        elif user_object_id:
+            headers["X-AnchorMailbox"] = "Oid:{}@{}".format(
+                user_object_id, self.authority.tenant)
         response = _clean_up(self.client.obtain_token_by_user_fic(
             scope=self._decorate_scope(scopes),
             assertion=assertion,
             username=username,
             user_object_id=user_object_id,
-            headers=telemetry_context.generate_headers(),
+            headers=headers,
             data=dict(
                 kwargs.pop("data", {}),
                 claims=_merge_claims_challenge_and_capabilities(

msal/oauth2cli/oauth2.py

@@ -106,9 +106,13 @@ def __init__(
                 It can also be a callable (recommended),
                 so that we will do lazy creation of an assertion.
 
-                The callable may accept zero arguments (legacy) or one argument.
-                When it accepts one argument, it will receive a dict containing
-                ``"client_id"``, ``"token_endpoint"``, and optionally ``"fmi_path"``
+                The callable may accept zero arguments (legacy) or one
+                required positional argument.  Callables whose positional
+                parameters all have default values (e.g.
+                ``lambda token=token: token``) are treated as zero-arg.
+                When the callable declares a required positional parameter,
+                it will receive a dict containing ``"client_id"``,
+                ``"token_endpoint"``, and optionally ``"fmi_path"``
                 (when an FMI path is set on the current request).
             client_assertion_type (str):
                 The type of your :attr:`client_assertion` parameter.

msal/token_cache.py

@@ -416,7 +416,7 @@ def __add(self, event, now=None):
                     }
                 grant_types_that_establish_an_account = (
                     _GRANT_TYPE_BROKER, "authorization_code", "password",
-                    Client.DEVICE_FLOW["GRANT_TYPE"])
+                    Client.DEVICE_FLOW["GRANT_TYPE"], "user_fic")
                 if event.get("grant_type") in grant_types_that_establish_an_account:
                     account["account_source"] = event["grant_type"]
                 self.modify(self.CredentialType.ACCOUNT, account, account)

tests/test_application.py

@@ -1197,6 +1197,37 @@ def mock_post(url, headers=None, data=None, *args, **kwargs):
         self.assertNotIn("username", captured_data,
             "username should NOT be in body when user_object_id is provided")
 
+    def test_ccs_routing_header_with_username(self):
+        app = self._make_app()
+        captured_headers = {}
+
+        def mock_post(url, headers=None, data=None, *args, **kwargs):
+            captured_headers.update(headers or {})
+            return MinimalResponse(status_code=200, text=_build_user_fic_response())
+
+        app.acquire_token_by_user_federated_identity_credential(
+            ["scope"], assertion="t2", username="user@contoso.com", post=mock_post)
+        self.assertEqual("upn:user@contoso.com",
+            captured_headers.get("X-AnchorMailbox"),
+            "CCS routing header should use UPN format for username path")
+
+    def test_ccs_routing_header_with_oid(self):
+        app = self._make_app()
+        captured_headers = {}
+
+        def mock_post(url, headers=None, data=None, *args, **kwargs):
+            captured_headers.update(headers or {})
+            return MinimalResponse(status_code=200, text=_build_user_fic_response())
+
+        app.acquire_token_by_user_federated_identity_credential(
+            ["scope"], assertion="t2",
+            user_object_id="user_oid_123", post=mock_post)
+        self.assertIn("X-AnchorMailbox", captured_headers,
+            "CCS routing header should be present for OID path")
+        self.assertTrue(
+            captured_headers["X-AnchorMailbox"].startswith("Oid:"),
+            "CCS routing header should use Oid format for user_object_id path")
+
 
 @patch(_OIDC_DISCOVERY, new=_OIDC_DISCOVERY_MOCK)
 class TestUserFicCacheBehavior(unittest.TestCase):
@@ -1291,6 +1322,24 @@ def mock_post(url, headers=None, data=None, *args, **kwargs):
         self.assertIn("access_token", silent_result)
         self.assertEqual("oid_fic_at", silent_result["access_token"])
 
+    def test_account_source_is_set_to_user_fic(self):
+        """Accounts created by user_fic should have account_source set."""
+        app = self._make_app()
+
+        def mock_post(url, headers=None, data=None, *args, **kwargs):
+            return MinimalResponse(status_code=200, text=_build_user_fic_response(
+                uid="user_oid", utid="tenant_id"))
+
+        app.acquire_token_by_user_federated_identity_credential(
+            ["https://graph.microsoft.com/.default"],
+            assertion="t2", username="user@contoso.com", post=mock_post)
+
+        accounts = app.get_accounts()
+        self.assertTrue(len(accounts) > 0)
+        self.assertEqual("user_fic", accounts[0].get("account_source"),
+            "FIC accounts should have account_source='user_fic' to avoid "
+            "broker path misrouting")
+
 
 @patch(_OIDC_DISCOVERY, new=_OIDC_DISCOVERY_MOCK)
 class TestUserFicInputValidation(unittest.TestCase):