Improve live test playback running (#2333)

* Improve live test playback running

* Add LRO helper method for playback

* Fix linting

* Add comment and fix cancellation behavior

Author: alzimmermsft

core/Azure.Mcp.Core/src/Services/Azure/BaseAzureResourceService.cs

@@ -196,10 +196,9 @@ protected async Task<ArmClient> CreateArmClientWithApiVersionAsync(string resour
     /// <param name="cancellationToken">Cancellation token.</param>
     /// <returns>The <see cref="GenericResource"/> instance for the requested resource.</returns>
     /// <exception cref="ArgumentNullException">Thrown when a required parameter is null.</exception>
-    protected async Task<GenericResource> GetGenericResourceAsync(ArmClient armClient, ResourceIdentifier resourceIdentifier, CancellationToken cancellationToken = default)
+    protected static async Task<GenericResource> GetGenericResourceAsync(ArmClient armClient, ResourceIdentifier resourceIdentifier, CancellationToken cancellationToken = default)
     {
-        if (armClient == null)
-            throw new ArgumentNullException(nameof(armClient));
+        ArgumentNullException.ThrowIfNull(armClient);
 
         var genericResources = armClient.GetGenericResources();
         var response = await genericResources.GetAsync(resourceIdentifier, cancellationToken).ConfigureAwait(false);
@@ -224,10 +223,15 @@ protected async Task<GenericResource> GetGenericResourceAsync(ArmClient armClien
     /// <returns>The <see cref="GenericResource"/> instance for the requested resource.</returns>
     /// <exception cref="ArgumentNullException">Thrown when a required parameter is null.</exception>
     /// <exception cref="InvalidOperationException">Thrown when the content is invalid.</exception>
-    protected async Task<GenericResource> CreateOrUpdateGenericResourceAsync<T>(ArmClient armClient, ResourceIdentifier resourceIdentifier, AzureLocation azureLocation, T content, JsonTypeInfo<T> jsonTypeInfo, CancellationToken cancellationToken)
+    protected static async Task<GenericResource> CreateOrUpdateGenericResourceAsync<T>(
+        ArmClient armClient,
+        ResourceIdentifier resourceIdentifier,
+        AzureLocation azureLocation,
+        T content,
+        JsonTypeInfo<T> jsonTypeInfo,
+        CancellationToken cancellationToken)
     {
-        if (armClient == null)
-            throw new ArgumentNullException(nameof(armClient));
+        ArgumentNullException.ThrowIfNull(armClient);
 
         // Convert from T to GenericResourceData
         byte[] jsonBytes = JsonSerializer.SerializeToUtf8Bytes(content, jsonTypeInfo);
@@ -236,7 +240,8 @@ protected async Task<GenericResource> CreateOrUpdateGenericResourceAsync<T>(ArmC
         GenericResourceData data = dataModel.Create(ref reader, new ModelReaderWriterOptions("W"))
             ?? throw new InvalidOperationException("Failed to create deployment data");
         // Create the resource
-        var result = await armClient.GetGenericResources().CreateOrUpdateAsync(WaitUntil.Completed, resourceIdentifier, data, cancellationToken);
+        var result = await armClient.GetGenericResources().CreateOrUpdateAsync(WaitUntil.Started, resourceIdentifier, data, cancellationToken);
+        await WaitForLroCompletionAsync(result, cancellationToken);
         return result.Value;
     }
 }

core/Azure.Mcp.Core/src/Services/Azure/BaseAzureService.cs

@@ -7,6 +7,7 @@
 using Azure.Core.Pipeline;
 using Azure.Mcp.Core.Services.Azure.Tenant;
 using Azure.ResourceManager;
+using Microsoft.Mcp.Core.Helpers;
 using Microsoft.Mcp.Core.Options;
 using Microsoft.Mcp.Core.Services.Azure;
 
@@ -30,6 +31,7 @@ public abstract class BaseAzureService
     private static readonly string s_framework;
     private static readonly string s_platform;
     private static readonly string s_defaultUserAgent;
+    private static readonly TimeSpan? s_defaultPollInterval;
 
     static BaseAzureService()
     {
@@ -41,6 +43,13 @@ static BaseAzureService()
         // Initialize the default user agent policy without transport type
         s_defaultUserAgent = $"azmcp/{s_version} ({s_framework}; {s_platform})";
         s_sharedUserAgentPolicy = new UserAgentPolicy(s_defaultUserAgent);
+
+#if DEBUG
+        if (EnvironmentHelpers.IsPlaybackTesting())
+        {
+            s_defaultPollInterval = TimeSpan.FromMilliseconds(1);
+        }
+#endif
     }
 
     /// <summary>
@@ -76,20 +85,12 @@ public static void InitializeUserAgentPolicy(string transportType)
     /// Disables upper bounds enforcement on retry policy values (delays, timeouts, max retries).
     /// This method should be called once during application startup when the --dangerously-disable-retry-limits flag is set.
     /// </summary>
-    public static void DisableRetryLimits()
-    {
-        s_retryLimitsDisabled = true;
-    }
+    public static void DisableRetryLimits() => s_retryLimitsDisabled = true;
 
     /// <summary>
     /// Resets the retry limits flag. For testing only.
     /// </summary>
-    internal static void ResetRetryLimits()
-    {
-        s_retryLimitsDisabled = false;
-    }
-
-
+    internal static void ResetRetryLimits() => s_retryLimitsDisabled = false;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="BaseAzureService"/> class.
@@ -175,7 +176,7 @@ protected async Task<TokenCredential> GetCredential(string? tenant, Cancellation
 
         try
         {
-            return await TenantService!.GetTokenCredentialAsync(tenantId, cancellationToken);
+            return await TenantService.GetTokenCredentialAsync(tenantId, cancellationToken);
         }
         catch (Exception ex)
         {
@@ -295,4 +296,45 @@ protected static void ValidateRequiredParameters(params (string name, string? va
                 $"Required parameter{(missingParams.Length > 1 ? "s are" : " is")} null or empty: {string.Join(", ", missingParams)}");
         }
     }
+
+    /// <summary>
+    /// Waits for the completion of a long-running operation, periodically polling the operation status until it completes.
+    /// </summary>
+    /// <typeparam name="T">The return type.</typeparam>
+    /// <param name="operation">The long-running operation.</param>
+    /// <param name="cancellationToken">The cancellation token that can cancel the request.</param>
+    /// <returns>The response once the long-running operation completes.</returns>
+    protected static async Task WaitForLroCompletionAsync<T>(Operation<T> operation, CancellationToken cancellationToken = default) where T : notnull
+    {
+        ArgumentNullException.ThrowIfNull(operation);
+
+        if (s_defaultPollInterval.HasValue)
+        {
+            await operation.WaitForCompletionAsync(s_defaultPollInterval.Value, cancellationToken);
+        }
+        else
+        {
+            await operation.WaitForCompletionAsync(cancellationToken);
+        }
+    }
+
+    /// <summary>
+    /// Waits for the completion of a long-running operation, periodically polling the operation status until it completes.
+    /// </summary>
+    /// <param name="operation">The long-running operation.</param>
+    /// <param name="cancellationToken">The cancellation token that can cancel the request.</param>
+    /// <returns>The response once the long-running operation completes.</returns>
+    protected static async Task WaitForLroCompletionAsync(Operation operation, CancellationToken cancellationToken = default)
+    {
+        ArgumentNullException.ThrowIfNull(operation);
+
+        if (s_defaultPollInterval.HasValue)
+        {
+            await operation.WaitForCompletionResponseAsync(s_defaultPollInterval.Value, cancellationToken);
+        }
+        else
+        {
+            await operation.WaitForCompletionResponseAsync(cancellationToken);
+        }
+    }
 }

core/Azure.Mcp.Core/tests/Azure.Mcp.Core.LiveTests/RecordingFramework/RecordedCommandTestHarness.cs

@@ -13,7 +13,8 @@ namespace Azure.Mcp.Core.LiveTests.RecordingFramework;
 /// </summary>
 /// <param name="output"></param>
 /// <param name="fixture"></param>
-internal sealed class RecordedCommandTestHarness(ITestOutputHelper output, TestProxyFixture fixture, LiveServerFixture liveServerFixture) : RecordedCommandTestsBase(output, fixture, liveServerFixture)
+internal sealed class RecordedCommandTestHarness(ITestOutputHelper output, TestProxyFixture fixture, LiveServerFixture liveServerFixture)
+    : RecordedCommandTestsBase(output, fixture, liveServerFixture)
 {
     public TestMode DesiredMode { get; set; } = TestMode.Record;
 
@@ -37,13 +38,7 @@ protected override ValueTask LoadSettingsAsync()
         return ValueTask.CompletedTask;
     }
 
-    public void ResetVariables()
-    {
-        TestVariables.Clear();
-    }
+    public void ResetVariables() => TestVariables.Clear();
 
-    public string GetRecordingId()
-    {
-        return RecordingId;
-    }
+    public string GetRecordingId() => RecordingId;
 }

core/Azure.Mcp.Core/tests/Azure.Mcp.Core.LiveTests/RecordingFramework/RecordedCommandTestsBaseTests.cs

@@ -1,8 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System;
-using System.IO;
 using System.Reflection;
 using System.Text;
 using System.Text.Json;

core/Microsoft.Mcp.Core/src/Helpers/EnvironmentVariableHelpers.cs

@@ -37,4 +37,16 @@ public static void SetAzureSubscriptionId(string? subscriptionId)
     {
         Environment.SetEnvironmentVariable(AzureSubscriptionIdEnvironmentVariable, subscriptionId);
     }
+
+    public static bool IsPlaybackTesting()
+    {
+#if DEBUG
+        // In debug builds, check for the presence of an environment variable to determine if we're in playback testing mode.
+        var testModeEnv = Environment.GetEnvironmentVariable("TEST_MODE");
+        return string.Equals(testModeEnv, "Playback", StringComparison.OrdinalIgnoreCase);
+#else
+        // In non-debug builds, never consider ourselves to be in playback testing mode.
+        return false;
+#endif
+    }
 }

core/Microsoft.Mcp.Core/src/Services/Azure/Authentication/CustomChainedCredential.cs

@@ -130,9 +130,8 @@ private static TokenCredential CreateCredential(string? tenantId, ILogger<Custom
         bool hasExplicitCredentialSetting = !string.IsNullOrEmpty(tokenCredentials);
 
 #if DEBUG
-        bool isPlaybackMode = string.Equals(tokenCredentials, "PlaybackTokenCredential", StringComparison.OrdinalIgnoreCase);
         // Short-circuit for playback to avoid any real auth & interactive prompts.
-        if (isPlaybackMode)
+        if (EnvironmentHelpers.IsPlaybackTesting())
         {
             logger?.LogDebug("Playback mode detected: using PlaybackTokenCredential.");
             return new PlaybackTokenCredential();

core/Microsoft.Mcp.Core/src/Services/Http/RecordingRedirectHandler.cs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Net.Http.Headers;
+using Microsoft.Mcp.Core.Helpers;
+
 namespace Microsoft.Mcp.Core.Services.Http;
 
 /// <summary>
@@ -10,26 +13,23 @@ namespace Microsoft.Mcp.Core.Services.Http;
 /// This handler is intended to be injected as the LAST delegating handler (closest to the transport) so
 /// that it rewrites the final outgoing wire request.
 /// </summary>
-internal sealed class RecordingRedirectHandler : DelegatingHandler
+/// <param name="proxyUri">The URI of the recording/replace proxy to redirect requests to.</param>
+internal sealed class RecordingRedirectHandler(Uri proxyUri) : DelegatingHandler
 {
     private const string CosmosSerializationHeader = "x-ms-cosmos-supported-serialization-formats";
-    private readonly Uri _proxyUri;
-
-    public RecordingRedirectHandler(Uri proxyUri)
-    {
-        _proxyUri = proxyUri ?? throw new ArgumentNullException(nameof(proxyUri));
-    }
+    private readonly Uri _proxyUri = proxyUri ?? throw new ArgumentNullException(nameof(proxyUri));
+    private readonly bool _playbackTesting = EnvironmentHelpers.IsPlaybackTesting();
 
     protected override HttpResponseMessage Send(HttpRequestMessage request, CancellationToken cancellationToken)
     {
         Redirect(request);
-        return base.Send(request, cancellationToken)!;
+        return StripRetryAfter(base.Send(request, cancellationToken));
     }
 
     protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
     {
         Redirect(request);
-        return await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
+        return StripRetryAfter(await base.SendAsync(request, cancellationToken).ConfigureAwait(false));
     }
 
     private void Redirect(HttpRequestMessage message)
@@ -60,4 +60,24 @@ private void Redirect(HttpRequestMessage message)
 
         message.RequestUri = builder.Uri;
     }
+
+    private HttpResponseMessage StripRetryAfter(HttpResponseMessage response)
+    {
+        if (_playbackTesting)
+        {
+            if (response.Headers.Remove("Retry-After"))
+            {
+                response.Headers.RetryAfter = new RetryConditionHeaderValue(TimeSpan.Zero);
+            }
+            if (response.Headers.Remove("x-ms-retry-after-ms"))
+            {
+                response.Headers.Add("x-ms-retry-after-ms", "0");
+            }
+            if (response.Headers.Remove("retry-after-ms"))
+            {
+                response.Headers.Add("retry-after-ms", "0");
+            }
+        }
+        return response;
+    }
 }

core/Microsoft.Mcp.Core/tests/Microsoft.Mcp.Tests/Client/CommandTestsBase.cs

@@ -12,7 +12,8 @@
 
 namespace Microsoft.Mcp.Tests.Client;
 
-public abstract class CommandTestsBase(ITestOutputHelper output, LiveServerFixture liveServerFixture) : IAsyncLifetime, IDisposable, IClassFixture<LiveServerFixture>
+public abstract class CommandTestsBase(ITestOutputHelper output, LiveServerFixture liveServerFixture)
+    : IAsyncLifetime, IDisposable, IClassFixture<LiveServerFixture>
 {
     protected const string TenantNameReason = "Service principals cannot use TenantName for lookup";
 
@@ -65,7 +66,7 @@ protected virtual async ValueTask LoadSettingsAsync()
         TestMode = Settings.TestMode;
     }
 
-    private async Task<LiveTestSettings?> TryLoadLiveSettingsAsync()
+    private static async Task<LiveTestSettings?> TryLoadLiveSettingsAsync()
     {
         try
         {
@@ -90,22 +91,25 @@ protected virtual async ValueTask LoadSettingsAsync()
             // { "AZURE_SUBSCRIPTION_ID", Settings.SubscriptionId }
         ];
 
-        if (proxy != null && proxy.Proxy != null)
+        // Add any custom environment variables from settings
+        if (Settings?.EnvironmentVariables != null)
         {
-            envVarDictionary.Add("TEST_PROXY_URL", proxy.Proxy.BaseUri);
-
-            if (TestMode is TestMode.Playback)
+            foreach (var kvp in Settings.EnvironmentVariables)
             {
-                envVarDictionary.Add("AZURE_TOKEN_CREDENTIALS", "PlaybackTokenCredential");
+                envVarDictionary[kvp.Key] = kvp.Value;
             }
         }
 
-        // Add any custom environment variables from settings
-        if (Settings?.EnvironmentVariables != null)
+        if (proxy != null && proxy.Proxy != null)
         {
-            foreach (var kvp in Settings.EnvironmentVariables)
+            envVarDictionary["TEST_PROXY_URL"] = proxy.Proxy.BaseUri;
+
+            if (TestMode is TestMode.Playback)
             {
-                envVarDictionary[kvp.Key] = kvp.Value;
+                // AZURE_TOKEN_CREDENTIALS=PlaybackTokenCredential tells the server to use a special credential that
+                // returns fake tokens in playback mode, which prevents any accidental live calls if a test is misconfigured
+                envVarDictionary["AZURE_TOKEN_CREDENTIALS"] = "PlaybackTokenCredential";
+                envVarDictionary["TEST_MODE"] = "Playback";
             }
         }
 
@@ -262,8 +266,5 @@ protected virtual void Dispose(bool disposing)
 
     // subclasses should override this method to dispose async resources
     // overrides should still call base.DisposeAsyncCore()
-    protected virtual ValueTask DisposeAsyncCore()
-    {
-        return ValueTask.CompletedTask;
-    }
+    protected virtual ValueTask DisposeAsyncCore() => ValueTask.CompletedTask;
 }

core/Microsoft.Mcp.Core/tests/Microsoft.Mcp.Tests/Client/Helpers/LiveTestSettings.cs

@@ -12,6 +12,12 @@ public class LiveTestSettings
 {
     public const string TestSettingsFileName = ".testsettings.json";
 
+    private static readonly JsonSerializerOptions s_jsonOptions = new()
+    {
+        PropertyNameCaseInsensitive = true,
+        Converters = { new JsonStringEnumConverter() }
+    };
+
     public string PrincipalName { get; set; } = string.Empty;
     public bool IsServicePrincipal { get; set; }
     public string TenantId { get; set; } = string.Empty;
@@ -53,11 +59,7 @@ public static bool TryLoadTestSettings([NotNullWhen(true)] out LiveTestSettings?
         {
             var json = File.ReadAllText(path);
 
-            settings = JsonSerializer.Deserialize<LiveTestSettings>(json, new JsonSerializerOptions()
-            {
-                PropertyNameCaseInsensitive = true,
-                Converters = { new JsonStringEnumConverter() }
-            });
+            settings = JsonSerializer.Deserialize<LiveTestSettings>(json, s_jsonOptions);
 
             if (settings != null)
             {