Add telemetry API ID for long-running OBO. (#4105)

* Add telemetry API ID for long-running OBO. Add tests.

* Add separate Api IDs for both long-running OBO methods.

* Apply suggestions from code review

Co-authored-by: Gladwin Johnson <90415114+gladjohn@users.noreply.github.com>

* Fix comments.

---------

Co-authored-by: Gladwin Johnson <90415114+gladjohn@users.noreply.github.com>

src/client/Microsoft.Identity.Client/ApiConfig/AcquireTokenOnBehalfOfParameterBuilder.cs

@@ -180,8 +180,22 @@ protected override void Validate()
         }
         /// <inheritdoc />
         internal override ApiEvent.ApiIds CalculateApiEventId()
-        {
-            return ApiEvent.ApiIds.AcquireTokenOnBehalfOf;
+        {
+            if (string.IsNullOrEmpty(Parameters.LongRunningOboCacheKey))
+            {
+                return ApiEvent.ApiIds.AcquireTokenOnBehalfOf;
+            }
+            else
+            {
+                if (Parameters.UserAssertion != null)
+                {
+                    return ApiEvent.ApiIds.InitiateLongRunningObo;
+                }
+                else
+                {
+                    return ApiEvent.ApiIds.AcquireTokenInLongRunningObo;
+                }
+            }
         }
     }
 }

src/client/Microsoft.Identity.Client/Cache/CacheKeyFactory.cs

@@ -1,11 +1,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-using System;
 using System.Globalization;
 using Microsoft.Identity.Client.Cache.Items;
-using Microsoft.Identity.Client.Cache.Keys;
 using Microsoft.Identity.Client.Internal.Requests;
+using Microsoft.Identity.Client.TelemetryCore.Internal.Events;
 
 namespace Microsoft.Identity.Client.Cache
 {
@@ -26,13 +25,13 @@ public static string GetKeyFromRequest(AuthenticationRequestParameters requestPa
                 return key;
             }
 
-            if (requestParameters.ApiId == TelemetryCore.Internal.Events.ApiEvent.ApiIds.AcquireTokenSilent ||
-                requestParameters.ApiId == TelemetryCore.Internal.Events.ApiEvent.ApiIds.RemoveAccount)
+            if (requestParameters.ApiId == ApiEvent.ApiIds.AcquireTokenSilent ||
+                requestParameters.ApiId == ApiEvent.ApiIds.RemoveAccount)
             {
                 return requestParameters.Account?.HomeAccountId?.Identifier;
             }
 
-            if (requestParameters.ApiId == TelemetryCore.Internal.Events.ApiEvent.ApiIds.GetAccountById)
+            if (requestParameters.ApiId == ApiEvent.ApiIds.GetAccountById)
             {
                 return requestParameters.HomeAccountId;
             }
@@ -50,7 +49,7 @@ public static string GetKeyFromRequest(AuthenticationRequestParameters requestPa
             }
 
             if (requestParameters.AppConfig.IsConfidentialClient ||
-                requestParameters.ApiId == TelemetryCore.Internal.Events.ApiEvent.ApiIds.AcquireTokenSilent)
+                requestParameters.ApiId == ApiEvent.ApiIds.AcquireTokenSilent)
             {
                 return homeAccountIdFromResponse;
             }
@@ -68,15 +67,15 @@ public static string GetKeyFromRequest(AuthenticationRequestParameters requestPa
 
         private static bool GetOboOrAppKey(AuthenticationRequestParameters requestParameters, out string key)
         {
-            if (requestParameters.ApiId == TelemetryCore.Internal.Events.ApiEvent.ApiIds.AcquireTokenOnBehalfOf)
+            if (ApiEvent.IsOnBehalfOfRequest(requestParameters.ApiId))
             {
                 key = GetOboKey(requestParameters.LongRunningOboCacheKey, requestParameters.UserAssertion);
                 return true;
             }
 
-            if (requestParameters.ApiId == TelemetryCore.Internal.Events.ApiEvent.ApiIds.AcquireTokenForClient || 
-                requestParameters.ApiId == TelemetryCore.Internal.Events.ApiEvent.ApiIds.AcquireTokenForSystemAssignedManagedIdentity ||
-                requestParameters.ApiId == TelemetryCore.Internal.Events.ApiEvent.ApiIds.AcquireTokenForUserAssignedManagedIdentity)
+            if (requestParameters.ApiId == ApiEvent.ApiIds.AcquireTokenForClient || 
+                requestParameters.ApiId == ApiEvent.ApiIds.AcquireTokenForSystemAssignedManagedIdentity ||
+                requestParameters.ApiId == ApiEvent.ApiIds.AcquireTokenForUserAssignedManagedIdentity)
             {
                 string tenantId = requestParameters.Authority.TenantId ?? "";
                 key = GetClientCredentialKey(requestParameters.AppConfig.ClientId, tenantId, requestParameters.AuthenticationScheme?.KeyId);

src/client/Microsoft.Identity.Client/Internal/Requests/OnBehalfOfRequest.cs

@@ -44,7 +44,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
             CacheRefreshReason cacheInfoTelemetry = CacheRefreshReason.NotApplicable;
 
             //Check if initiating a long running process
-            if (IsLongOboInitialize())
+            if (AuthenticationRequestParameters.ApiId == ApiEvent.ApiIds.InitiateLongRunningObo)
             {
                 //Long running process should not use cached tokens
                 logger.Info("[OBO Request] Initiating long running process. Fetching OBO token from ESTS.");
@@ -135,18 +135,14 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
             }
         }
 
-        private bool IsLongOboInitialize()
-        {
-            return AuthenticationRequestParameters.UserAssertion != null && !string.IsNullOrEmpty(AuthenticationRequestParameters.LongRunningOboCacheKey);
-        }
-
         private async Task<AuthenticationResult> RefreshRtOrFetchNewAccessTokenAsync(CancellationToken cancellationToken)
         {
             var logger = AuthenticationRequestParameters.RequestContext.Logger;
 
-            if (IsLongRunningObo())
+            if (ApiEvent.IsLongRunningObo(AuthenticationRequestParameters.ApiId))
             {
-                AuthenticationRequestParameters.RequestContext.Logger.Info("[OBO request] Long-running OBO flow, trying to refresh using an refresh token flow.");
+                AuthenticationRequestParameters.RequestContext.Logger.Info("[OBO request] Long-running OBO flow, trying to refresh using a refresh token flow.");
+
 
                 // Look for a refresh token
                 MsalRefreshTokenCacheItem cachedRefreshToken = await CacheManager.FindRefreshTokenAsync().ConfigureAwait(false);
@@ -174,7 +170,7 @@ private async Task<AuthenticationResult> RefreshRtOrFetchNewAccessTokenAsync(Can
                     return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse).ConfigureAwait(false);
                 }
 
-                if (AcquireTokenInLongRunningOboWasCalled())
+                if (AuthenticationRequestParameters.ApiId == ApiEvent.ApiIds.AcquireTokenInLongRunningObo)
                 {
                     AuthenticationRequestParameters.RequestContext.Logger.Error("[OBO request] AcquireTokenInLongRunningProcess was called and no access or refresh tokens were found in the cache.");
                     throw new MsalClientException(MsalError.OboCacheKeyNotInCacheError, MsalErrorMessage.OboCacheKeyNotInCache);
@@ -190,19 +186,12 @@ private async Task<AuthenticationResult> RefreshRtOrFetchNewAccessTokenAsync(Can
             return await FetchNewAccessTokenAsync(cancellationToken).ConfigureAwait(false);
         }
 
-        // Returns whether AcquireTokenInLongRunningProcess was called (user assertion is null in this case)
-        private bool AcquireTokenInLongRunningOboWasCalled()
-        {
-            return AuthenticationRequestParameters.UserAssertion == null &&
-                !string.IsNullOrEmpty(AuthenticationRequestParameters.LongRunningOboCacheKey);
-        }
-
         private async Task<AuthenticationResult> FetchNewAccessTokenAsync(CancellationToken cancellationToken)
         {
             var msalTokenResponse = await SendTokenRequestAsync(GetBodyParameters(), cancellationToken).ConfigureAwait(false);
 
             // We always retrieve a refresh token in OBO but we don't want to cache it for normal OBO flow, only for long-running OBO
-            if (!IsLongRunningObo())
+            if (!ApiEvent.IsLongRunningObo(AuthenticationRequestParameters.ApiId))
             {
                 msalTokenResponse.RefreshToken = null;
             }
@@ -238,10 +227,5 @@ private async Task<AuthenticationResult> FetchNewAccessTokenAsync(CancellationTo
 
             return new KeyValuePair<string, string>(Constants.CcsRoutingHintHeader, _ccsRoutingHint);
         }
-
-        private bool IsLongRunningObo()
-        {
-            return !string.IsNullOrEmpty(AuthenticationRequestParameters.LongRunningOboCacheKey);
-        }
     }
 }

src/client/Microsoft.Identity.Client/TelemetryCore/Internal/Events/ApiEvent.cs

@@ -26,11 +26,17 @@ public enum ApiIds
             GetAccountById = 1011,
             GetAccountsByUserFlow = 1012,
             RemoveAccount = 1013,
+
+            // MSAL 4.51.0+
             RemoveOboTokens = 1014,
-            // The API Ids for managed identity will not be found in Http telemetry,
+            // The API IDs for managed identity will not be found in HTTP telemetry,
             // as we don't hit eSTS for managed identity calls.
             AcquireTokenForSystemAssignedManagedIdentity = 1015,
-            AcquireTokenForUserAssignedManagedIdentity = 1016
+            AcquireTokenForUserAssignedManagedIdentity = 1016,
+
+            // MSAL 4.54.0+
+            InitiateLongRunningObo = 1017,
+            AcquireTokenInLongRunningObo = 1018,
         }
 
         public ApiEvent(Guid correlationId)
@@ -120,5 +126,10 @@ public string TokenTypeString
         {
             get => TokenType.HasValue ? TokenType.Value.ToString("D") : null;
         }
+
+        public static bool IsLongRunningObo(ApiIds apiId) => apiId == ApiIds.InitiateLongRunningObo || apiId == ApiIds.AcquireTokenInLongRunningObo;
+
+        public static bool IsOnBehalfOfRequest(ApiIds apiId) => apiId == ApiIds.AcquireTokenOnBehalfOf || IsLongRunningObo(apiId);
+
     }
 }

src/client/Microsoft.Identity.Client/TokenCache.ITokenCacheInternal.cs

@@ -538,7 +538,7 @@ private void MergeWamAccountIds(MsalAccountCacheItem msalAccountCacheItem)
             string requestTenantId = requestParams.Authority.TenantId;
             bool filterByTenantId = true;
 
-            if (requestParams.ApiId == ApiEvent.ApiIds.AcquireTokenOnBehalfOf) // OBO
+            if (ApiEvent.IsOnBehalfOfRequest(requestParams.ApiId))
             {
                 tokenCacheItems.FilterWithLogging(item =>
                         !string.IsNullOrEmpty(item.OboCacheKey) &&
@@ -574,9 +574,9 @@ private void MergeWamAccountIds(MsalAccountCacheItem msalAccountCacheItem)
 
             // Only AcquireTokenSilent has an IAccount in the request that can be used for filtering
             if (requestParams.ApiId != ApiEvent.ApiIds.AcquireTokenForClient &&
-                requestParams.ApiId != ApiEvent.ApiIds.AcquireTokenForSystemAssignedManagedIdentity && 
+                requestParams.ApiId != ApiEvent.ApiIds.AcquireTokenForSystemAssignedManagedIdentity &&
                 requestParams.ApiId != ApiEvent.ApiIds.AcquireTokenForUserAssignedManagedIdentity &&
-                requestParams.ApiId != ApiEvent.ApiIds.AcquireTokenOnBehalfOf)
+                !ApiEvent.IsOnBehalfOfRequest(requestParams.ApiId))
             {
                 tokenCacheItems.FilterWithLogging(item => item.HomeAccountId.Equals(
                                 requestParams.Account.HomeAccountId?.Identifier, StringComparison.OrdinalIgnoreCase),
@@ -835,7 +835,7 @@ internal async Task ExpireAllAccessTokensForTestAsync()
             AuthenticationRequestParameters requestParams,
             string familyId)
         {
-            if (requestParams.ApiId == ApiEvent.ApiIds.AcquireTokenOnBehalfOf) // OBO
+            if (ApiEvent.IsOnBehalfOfRequest(requestParams.ApiId))
             {
                 cacheItems.FilterWithLogging(item =>
                                 !string.IsNullOrEmpty(item.OboCacheKey) &&

tests/Microsoft.Identity.Test.Unit/CacheTests/CacheKeyFactoryTests.cs

@@ -217,7 +217,7 @@ public void TestCacheKeyForObo_WithCacheKey()
             var tenantAuthority = AuthorityInfo.FromAadAuthority(AzureCloudInstance.AzurePublic, tenant: TestConstants.AadTenantId, validateAuthority: false);
             var acquireTokenCommonParameters = new AcquireTokenCommonParameters
             {
-                ApiId = ApiEvent.ApiIds.AcquireTokenOnBehalfOf,
+                ApiId = ApiEvent.ApiIds.InitiateLongRunningObo,
                 AuthorityOverride = tenantAuthority
             };
 

tests/Microsoft.Identity.Test.Unit/RequestsTests/IntegratedWindowsAuthUsernamePasswordTests.cs

@@ -19,6 +19,8 @@
 using Microsoft.Identity.Test.Common.Core.Helpers;
 using Microsoft.Identity.Test.Common.Core.Mocks;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Microsoft.Identity.Client.TelemetryCore.Internal.Events;
+using Microsoft.Identity.Client.TelemetryCore;
 
 namespace Microsoft.Identity.Test.Unit.RequestsTests
 {
@@ -323,6 +325,7 @@ public async Task AcquireTokenByIntegratedWindowsAuth3rdPartyIDPTestAsync(string
             }
         }
 
+        [TestMethod]
         [DeploymentItem(@"Resources\TestMex.xml")]
         [DeploymentItem(@"Resources\WsTrustResponse13.xml")]
         public async Task AcquireTokenByIntegratedWindowsAuthMetadataTestAsync()
@@ -362,6 +365,9 @@ public async Task AcquireTokenByIntegratedWindowsAuthMetadataTestAsync()
                 StringAssert.Contains(realmDiscoveryHandler.ActualRequestMessage.Headers.ToString(), TestConstants.XClientVer,
                     "Client info header should contain " + TestConstants.XClientVer,
                     StringComparison.OrdinalIgnoreCase);
+
+                // Assert telemetry ApiId
+                Assert.AreEqual(ApiEvent.ApiIds.AcquireTokenByIntegratedWindowsAuth.ToString("D"), mockTokenRequestHttpHandler.ActualRequestMessage.Headers.GetValues(TelemetryConstants.XClientCurrentTelemetry).Single().Split('|')[1].Split(',')[0]);
             }
         }