Fix repeated Kusto token acquisition latency (#4027)

* Fix repeated Kusto token acquisition

Cache access tokens per Kusto client, refresh before expiration, and serialize concurrent acquisition. Add focused regression tests and document repeated-query behavior.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* Clarify Kusto token cache lifetime across transports

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* Document token acquisition guidance for HTTP clients

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

CONTRIBUTING.md

@@ -42,6 +42,7 @@ If you are contributing significant changes, or if the issue is already assigned
     - [Code Style](#code-style)
       - [Spelling Check](#spelling-check)
       - [Requirements](#requirements)
+    - [Token Acquisition for HTTP Clients](#token-acquisition-for-http-clients)
     - [AOT Compatibility Analysis](#aot-compatibility-analysis)
       - [Running the Analysis](#running-the-analysis)
       - [Installing Git Hooks](#installing-git-hooks)
@@ -670,6 +671,13 @@ This will check all files for spelling errors using the project's dictionary. Ad
 - XML documentation for public APIs
 - Follow Model Context Protocol (MCP) patterns
 
+### Token Acquisition for HTTP Clients
+
+- Obtain credentials through the existing injected authentication providers. Prefer Azure SDK clients or a reusable Azure Core bearer-token authentication pipeline to manage token caching and refresh.
+- For direct HTTP calls, do not assume repeated `GetTokenAsync` calls are inexpensive: credentials such as `AzureCliCredential` can invoke an external process even when a token is cached. Measure token acquisition separately from HTTP execution when investigating latency.
+- If explicit token caching is necessary, bind it to a fixed credential identity and scope; never share tokens across users, tenants, or audiences. Respect expiration and earlier refresh times, synchronize acquisition, and propagate cancellation and failures without preventing subsequent attempts.
+- Keep clients transport-agnostic and respect the configured cache lifetime and user-isolation guarantees. Test token reuse, refresh, concurrent acquisition, cancellation, failure recovery, and identity isolation.
+
 ### AOT Compatibility Analysis
 
 The AOT compatibility analysis helps identify potential issues that might prevent the Azure MCP Server from working correctly when compiled with AOT or when trimming is enabled.

servers/Azure.Mcp.Server/changelog-entries/1792790170892.yaml

@@ -0,0 +1,3 @@
+changes:
+  - section: "Bugs Fixed"
+    description: "Reduced repeated Kusto query latency by caching access tokens per client and refreshing them before expiration."

servers/Azure.Mcp.Server/docs/azmcp-commands.md

@@ -2847,6 +2847,13 @@ azmcp kusto query [--cluster-uri <cluster-uri> | --subscription <subscription> -
 
 ```
 
+Access tokens are reused for the lifetime of a cached Kusto client and refreshed
+five minutes before expiration or at the credential's earlier refresh time.
+With the current cache implementations, cross-call reuse occurs in stdio mode
+when caching is enabled; HTTP mode does not retain clients across calls.
+Tokens remain scoped to the client's credential and cloud; they are not shared
+across tenants.
+
 ### Azure Database for MySQL Operations
 
 ```bash

servers/Azure.Mcp.Server/docs/e2eTestPrompts.md

@@ -681,6 +681,7 @@ The `Interaction` column describes whether a prompt can invoke its tool immediat
 | kusto_database_list | List all databases in the Data Explorer cluster <cluster_name> | none |
 | kusto_database_list | Show me the databases in the Data Explorer cluster <cluster_name> | none |
 | kusto_query | Show me all items that contain the word <search_term> in the Data Explorer table <table_name> in cluster <cluster_name> | none |
+| kusto_query | Run `print value=1` three times in database <database_name> at cluster URI <cluster_uri> using the same server session and report each call's elapsed time | none |
 | kusto_sample | Show me a data sample from the Data Explorer table <table_name> in cluster <cluster_name> | none |
 | kusto_table_list | List all tables in the Data Explorer database <database_name> in cluster <cluster_name> | none |
 | kusto_table_list | Show me the tables in the Data Explorer database <database_name> in cluster <cluster_name> | none |

tools/Azure.Mcp.Tools.Kusto/src/Services/KustoClient.cs

@@ -76,9 +76,15 @@ public sealed class KustoClient(
 
     private readonly string _clusterUri = ValidateAndNormalizeClusterUri(clusterUri);
     private readonly TokenCredential _tokenCredential = tokenCredential;
+    // Tokens are cached only within this client instance. Cross-call reuse depends
+    // on KustoService's ICacheService: stdio retains clients when caching is enabled;
+    // the current HTTP cache is a no-op, so clients are not reused across calls.
+    private readonly SemaphoreSlim _tokenLock = new(1, 1);
+    private AccessToken _accessToken;
     private readonly string _userAgent = userAgent;
     private readonly IAzureService _azureService = azureService;
     private static readonly TimeSpan s_httpClientTimeout = TimeSpan.FromSeconds(240);
+    private static readonly TimeSpan s_tokenRefreshBuffer = TimeSpan.FromMinutes(5);
     private static readonly string s_application = "AzureMCP";
     private static readonly string s_clientRequestIdPrefix = "AzMcp";
 
@@ -234,14 +240,9 @@ public sealed class KustoClient(
 
     private async Task<HttpRequestMessage> GenerateRequestAsync(string uri, string database, string text, CancellationToken cancellationToken)
     {
-        var httpRequest = new HttpRequestMessage(HttpMethod.Post, uri);
-        var scopes = new string[]
-        {
-            GetKustoScope()
-        };
         var clientRequestId = s_clientRequestIdPrefix + Guid.NewGuid().ToString();
-        var tokenRequestContext = new TokenRequestContext(scopes, clientRequestId);
-        var accessToken = await _tokenCredential.GetTokenAsync(tokenRequestContext, cancellationToken);
+        var accessToken = await GetAccessTokenAsync(clientRequestId, cancellationToken).ConfigureAwait(false);
+        var httpRequest = new HttpRequestMessage(HttpMethod.Post, uri);
         httpRequest.Headers.Authorization = new("bearer", accessToken.Token);
         httpRequest.Headers.Add("User-Agent", _userAgent);
         httpRequest.Headers.Add("x-ms-client-request-id", clientRequestId);
@@ -265,6 +266,27 @@ public sealed class KustoClient(
         return httpRequest;
     }
 
+    private async Task<AccessToken> GetAccessTokenAsync(string clientRequestId, CancellationToken cancellationToken)
+    {
+        await _tokenLock.WaitAsync(cancellationToken).ConfigureAwait(false);
+        try
+        {
+            var now = DateTimeOffset.UtcNow;
+            if (_accessToken.ExpiresOn <= now + s_tokenRefreshBuffer ||
+                _accessToken.RefreshOn <= now)
+            {
+                var tokenRequestContext = new TokenRequestContext([GetKustoScope()], clientRequestId);
+                _accessToken = await _tokenCredential.GetTokenAsync(tokenRequestContext, cancellationToken).ConfigureAwait(false);
+            }
+
+            return _accessToken;
+        }
+        finally
+        {
+            _tokenLock.Release();
+        }
+    }
+
     private static async Task<KustoResult> SendRequestAsync(HttpClient httpClient, HttpRequestMessage httpRequest, CancellationToken cancellationToken)
     {
         var httpResponse = await httpClient.SendAsync(httpRequest, HttpCompletionOption.ResponseContentRead, cancellationToken);

tools/Azure.Mcp.Tools.Kusto/tests/Azure.Mcp.Tools.Kusto.Tests/KustoClientTests.cs

@@ -47,6 +47,172 @@ public sealed class KustoClientTests
         Assert.NotNull(result);
     }
 
+    [Theory]
+    [InlineData("https://test.kusto.windows.net", "https://kusto.kusto.windows.net/.default")]
+    [InlineData("https://test.kusto.chinacloudapi.cn", "https://kusto.kusto.chinacloudapi.cn/.default")]
+    [InlineData("https://test.kusto.usgovcloudapi.net", "https://kusto.kusto.usgovcloudapi.net/.default")]
+    public async Task ExecuteCommandAsync_ReusesTokenForQueryAndControlRequests(string clusterUri, string scope)
+    {
+        var requestIds = new HashSet<string>();
+        using var handler = new MockHttpMessageHandler(request =>
+        {
+            Assert.Equal("noop-token", request.Headers.Authorization?.Parameter);
+            Assert.True(requestIds.Add(Assert.Single(request.Headers.GetValues("x-ms-client-request-id"))));
+        });
+        _azureService.GetClient(Arg.Any<string>()).Returns(_ => new HttpClient(handler, disposeHandler: false));
+        var client = new KustoClient(clusterUri, _tokenCredential, "azmcp", _azureService);
+        var cancellationToken = TestContext.Current.CancellationToken;
+
+        await client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+        await client.ExecuteQueryCommandAsync("testdb", "print value=2", cancellationToken);
+        await client.ExecuteControlCommandAsync("testdb", ".show tables", cancellationToken);
+
+        Assert.Equal(3, requestIds.Count);
+        await _tokenCredential.Received(1).GetTokenAsync(
+            Arg.Is<TokenRequestContext>(context => context.Scopes.SequenceEqual(new[] { scope })),
+            cancellationToken);
+    }
+
+    [Theory]
+    [InlineData(-1, false)]
+    [InlineData(2, false)]
+    [InlineData(60, true)]
+    public async Task ExecuteCommandAsync_RefreshesExpiredOrExpiringToken(int expiresInMinutes, bool refreshRequested)
+    {
+        _tokenCredential.GetTokenAsync(Arg.Any<TokenRequestContext>(), Arg.Any<CancellationToken>())
+            .Returns(
+                new AccessToken(
+                    "old-token",
+                    DateTimeOffset.UtcNow.AddMinutes(expiresInMinutes),
+                    refreshRequested ? DateTimeOffset.UtcNow.AddMinutes(-1) : null),
+                new AccessToken("new-token", DateTimeOffset.UtcNow.AddHours(1)));
+        var tokens = new List<string?>();
+        using var handler = new MockHttpMessageHandler(request => tokens.Add(request.Headers.Authorization?.Parameter));
+        _azureService.GetClient(Arg.Any<string>()).Returns(_ => new HttpClient(handler, disposeHandler: false));
+        var client = new KustoClient("https://test.kusto.windows.net", _tokenCredential, "azmcp", _azureService);
+        var cancellationToken = TestContext.Current.CancellationToken;
+
+        await client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+        await client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+        await client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+
+        Assert.Equal(new[] { "old-token", "new-token", "new-token" }, tokens);
+        await _tokenCredential.Received(2).GetTokenAsync(Arg.Any<TokenRequestContext>(), cancellationToken);
+    }
+
+    [Fact]
+    public async Task ExecuteCommandAsync_ConcurrentRequestsShareTokenAcquisition()
+    {
+        var tokenSource = new TaskCompletionSource<AccessToken>(TaskCreationOptions.RunContinuationsAsynchronously);
+        _tokenCredential.GetTokenAsync(Arg.Any<TokenRequestContext>(), Arg.Any<CancellationToken>())
+            .Returns(_ => new ValueTask<AccessToken>(tokenSource.Task));
+        using var handler = new MockHttpMessageHandler(request =>
+            Assert.Equal("shared-token", request.Headers.Authorization?.Parameter));
+        _azureService.GetClient(Arg.Any<string>()).Returns(_ => new HttpClient(handler, disposeHandler: false));
+        var client = new KustoClient("https://test.kusto.windows.net", _tokenCredential, "azmcp", _azureService);
+        var cancellationToken = TestContext.Current.CancellationToken;
+
+        var requests = Enumerable.Range(0, 10)
+            .Select(_ => client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken))
+            .ToArray();
+        tokenSource.SetResult(new AccessToken("shared-token", DateTimeOffset.UtcNow.AddHours(1)));
+        await Task.WhenAll(requests);
+
+        await _tokenCredential.Received(1).GetTokenAsync(Arg.Any<TokenRequestContext>(), cancellationToken);
+    }
+
+    [Fact]
+    public async Task ExecuteCommandAsync_TokenFailurePropagatesAndAllowsRetry()
+    {
+        var failure = new InvalidOperationException("Token acquisition failed.");
+        _tokenCredential.GetTokenAsync(Arg.Any<TokenRequestContext>(), Arg.Any<CancellationToken>())
+            .Returns(
+                _ => throw failure,
+                _ => new AccessToken("retry-token", DateTimeOffset.UtcNow.AddHours(1)));
+        using var handler = new MockHttpMessageHandler(request =>
+            Assert.Equal("retry-token", request.Headers.Authorization?.Parameter));
+        _azureService.GetClient(Arg.Any<string>()).Returns(_ => new HttpClient(handler, disposeHandler: false));
+        var client = new KustoClient("https://test.kusto.windows.net", _tokenCredential, "azmcp", _azureService);
+        var cancellationToken = TestContext.Current.CancellationToken;
+
+        Assert.Same(failure, await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken)));
+        _azureService.DidNotReceive().GetClient(Arg.Any<string>());
+        await client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+
+        await _tokenCredential.Received(2).GetTokenAsync(Arg.Any<TokenRequestContext>(), cancellationToken);
+    }
+
+    [Fact]
+    public async Task ExecuteCommandAsync_CancelledWaitDoesNotCancelSharedAcquisition()
+    {
+        var tokenSource = new TaskCompletionSource<AccessToken>(TaskCreationOptions.RunContinuationsAsynchronously);
+        _tokenCredential.GetTokenAsync(Arg.Any<TokenRequestContext>(), Arg.Any<CancellationToken>())
+            .Returns(_ => new ValueTask<AccessToken>(tokenSource.Task));
+        using var handler = new MockHttpMessageHandler();
+        _azureService.GetClient(Arg.Any<string>()).Returns(_ => new HttpClient(handler, disposeHandler: false));
+        var client = new KustoClient("https://test.kusto.windows.net", _tokenCredential, "azmcp", _azureService);
+        var cancellationToken = TestContext.Current.CancellationToken;
+        using var cancelled = new CancellationTokenSource();
+
+        var firstRequest = client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+        var waitingRequest = client.ExecuteQueryCommandAsync("testdb", "print value=1", cancelled.Token);
+        await cancelled.CancelAsync();
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => waitingRequest);
+        tokenSource.SetResult(new AccessToken("shared-token", DateTimeOffset.UtcNow.AddHours(1)));
+        await firstRequest;
+        await client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+
+        await _tokenCredential.Received(1).GetTokenAsync(Arg.Any<TokenRequestContext>(), cancellationToken);
+    }
+
+    [Fact]
+    public async Task ExecuteCommandAsync_CancelledTokenAcquisitionAllowsRetry()
+    {
+        using var cancelled = new CancellationTokenSource();
+        _tokenCredential.GetTokenAsync(Arg.Any<TokenRequestContext>(), cancelled.Token)
+            .Returns(async call =>
+            {
+                await cancelled.CancelAsync();
+                call.Arg<CancellationToken>().ThrowIfCancellationRequested();
+                return default;
+            });
+        using var handler = new MockHttpMessageHandler();
+        _azureService.GetClient(Arg.Any<string>()).Returns(_ => new HttpClient(handler, disposeHandler: false));
+        var client = new KustoClient("https://test.kusto.windows.net", _tokenCredential, "azmcp", _azureService);
+        var cancellationToken = TestContext.Current.CancellationToken;
+
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(() =>
+            client.ExecuteQueryCommandAsync("testdb", "print value=1", cancelled.Token));
+        await client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+        await client.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+
+        await _tokenCredential.Received(1).GetTokenAsync(Arg.Any<TokenRequestContext>(), cancelled.Token);
+        await _tokenCredential.Received(1).GetTokenAsync(Arg.Any<TokenRequestContext>(), cancellationToken);
+    }
+
+    [Fact]
+    public async Task ExecuteCommandAsync_DifferentClientsDoNotShareTokens()
+    {
+        var otherCredential = Substitute.For<TokenCredential>();
+        otherCredential.GetTokenAsync(Arg.Any<TokenRequestContext>(), Arg.Any<CancellationToken>())
+            .Returns(new AccessToken("other-token", DateTimeOffset.UtcNow.AddHours(1)));
+        var tokens = new List<string?>();
+        using var handler = new MockHttpMessageHandler(request => tokens.Add(request.Headers.Authorization?.Parameter));
+        _azureService.GetClient(Arg.Any<string>()).Returns(_ => new HttpClient(handler, disposeHandler: false));
+        var firstClient = new KustoClient("https://test.kusto.windows.net", _tokenCredential, "azmcp", _azureService);
+        var secondClient = new KustoClient("https://test.kusto.windows.net", otherCredential, "azmcp", _azureService);
+        var cancellationToken = TestContext.Current.CancellationToken;
+
+        await firstClient.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+        await secondClient.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+        await firstClient.ExecuteQueryCommandAsync("testdb", "print value=1", cancellationToken);
+
+        Assert.Equal(new[] { "noop-token", "other-token", "noop-token" }, tokens);
+        await _tokenCredential.Received(1).GetTokenAsync(Arg.Any<TokenRequestContext>(), cancellationToken);
+        await otherCredential.Received(1).GetTokenAsync(Arg.Any<TokenRequestContext>(), cancellationToken);
+    }
+
     [Fact]
     public async Task ExecuteQueryCommandAsync_KeepsPathAndUsesKustoScope()
     {
@@ -222,12 +388,13 @@ public sealed class KustoClientTests
 
     #endregion
 
-    private sealed class MockHttpMessageHandler : HttpMessageHandler
+    private sealed class MockHttpMessageHandler(Action<HttpRequestMessage>? onRequest = null) : HttpMessageHandler
     {
         public Uri? LastRequestUri { get; private set; }
 
         protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
         {
+            onRequest?.Invoke(request);
             LastRequestUri = request.RequestUri;
 
             // The caller (HttpClient.SendAsync) takes ownership and is responsible for disposal HttpResponseMessage.