fix: propagate CancellationToken through async MCP tooling and observability call chains

src/Observability/Hosting/Caching/AgenticTokenCache.cs

@@ -96,7 +96,11 @@ public void RegisterObservability(string agentId, string tenantId, AgenticTokenS
         /// <returns>
         /// The observability token if available; otherwise, <c>null</c>.
         /// </returns>
-        public async Task<string?> GetObservabilityToken(string agentId, string tenantId)
+        public Task<string?> GetObservabilityToken(string agentId, string tenantId)
+            => GetObservabilityToken(agentId, tenantId, CancellationToken.None);
+
+        /// <inheritdoc/>
+        public async Task<string?> GetObservabilityToken(string agentId, string tenantId, CancellationToken cancellationToken)
         {
             if (!_map.TryGetValue($"{agentId}:{tenantId}", out var entry))
                 return null;

src/Observability/Hosting/Caching/IExporterTokenCache.cs

@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.Agents.A365.Observability.Hosting.Caching
@@ -18,5 +19,10 @@ public interface IExporterTokenCache<T> where T : class
         /// Returns an observability token (cached inside the credential) or null on failure/not registered.
         /// </summary>
         Task<string?> GetObservabilityToken(string agentId, string tenantId);
+
+        /// <summary>
+        /// Returns an observability token (cached inside the credential) or null on failure/not registered, with cancellation support.
+        /// </summary>
+        Task<string?> GetObservabilityToken(string agentId, string tenantId, CancellationToken cancellationToken);
     }
 }

src/Observability/Hosting/Caching/ServiceTokenCache.cs

@@ -122,15 +122,19 @@ public void RegisterObservability(string agentId, string tenantId, string token,
         /// <param name="agentId">The agent identifier.</param>
         /// <param name="tenantId">The tenant identifier.</param>
         /// <returns>The observability token if valid; otherwise, null.</returns>
-        public async Task<string?> GetObservabilityToken(string agentId, string tenantId)
+        public Task<string?> GetObservabilityToken(string agentId, string tenantId)
+            => GetObservabilityToken(agentId, tenantId, CancellationToken.None);
+
+        /// <inheritdoc/>
+        public Task<string?> GetObservabilityToken(string agentId, string tenantId, CancellationToken cancellationToken)
         {
             if (string.IsNullOrWhiteSpace(agentId) || string.IsNullOrWhiteSpace(tenantId))
-                return null;
+                return Task.FromResult<string?>(null);
 
             var key = GetKey(agentId, tenantId);
 
             if (!_map.TryGetValue(key, out var entry))
-                return null;
+                return Task.FromResult<string?>(null);
 
             // Check if token has expired
             if (DateTimeOffset.UtcNow >= entry.ExpiresAt)
@@ -140,10 +144,10 @@ public void RegisterObservability(string agentId, string tenantId, string token,
                 {
                     removedEntry.ClearToken();
                 }
-                return null;
+                return Task.FromResult<string?>(null);
             }
 
-            return await Task.FromResult(entry.Token).ConfigureAwait(false);
+            return Task.FromResult(entry.Token);
         }
 
         /// <summary>

src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterAsyncE2ETests.cs

@@ -17,6 +17,7 @@ public class Agent365ExporterAsyncE2ETests
     {
         private TestHttpMessageHandler? _handler;
         private ServiceProvider? _provider;
+        private HttpClient? _httpClient;
         private bool _receivedRequest;
         private string? _receivedContent;
 
@@ -408,8 +409,17 @@ private void SetupExporterTest()
                 req.Headers.Authorization.Should().NotBeNull();
                 return new HttpResponseMessage(System.Net.HttpStatusCode.OK);
             });
-            var httpClient = new HttpClient(this._handler);
-            this._provider = this.CreateTestServiceProvider(httpClient);
+            this._httpClient = new HttpClient(this._handler);
+            this._provider = this.CreateTestServiceProvider(this._httpClient);
+        }
+
+        [TestCleanup]
+        public void Cleanup()
+        {
+            this._provider?.Dispose();
+            this._provider = null;
+            this._httpClient?.Dispose();
+            this._httpClient = null;
         }
     }
 }

src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/McpToolServerConfigurationService_ToolEnumerationTests.cs

@@ -10,6 +10,7 @@
 using ModelContextProtocol.Client;
 using Moq;
 using System.Net.Http;
+using System.Threading;
 using Xunit;
 
 namespace Microsoft.Agents.A365.Tooling.Core.Tests;
@@ -53,7 +54,7 @@ public async Task EnumerateToolsFromServersAsync_WhenListServersFails_ReturnsEmp
         // Arrange
         var toolOptions = new ToolOptions();
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ThrowsAsync(new Exception("Network error"));
 
         // Act
@@ -74,7 +75,7 @@ public async Task EnumerateToolsFromServersAsync_WhenNoServersConfigured_Returns
         // Arrange
         var toolOptions = new ToolOptions();
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new List<MCPServerConfig>());
 
         // Act
@@ -101,15 +102,16 @@ public async Task EnumerateToolsFromServersAsync_FiltersInvalidServers_WithMissi
         };
 
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(servers);
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "valid-server"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .ReturnsAsync(new List<McpClientTool>());
 
         // Act
@@ -137,15 +139,16 @@ public async Task EnumerateToolsFromServersAsync_FiltersInvalidServers_WithMissi
         };
 
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(servers);
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "valid-server"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .ReturnsAsync(new List<McpClientTool>());
 
         // Act
@@ -175,23 +178,25 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesToolsFromMultipleServ
         var tools2 = new List<McpClientTool>();
 
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(servers);
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "server1"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .ReturnsAsync(tools1);
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "server2"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .ReturnsAsync(tools2);
 
         // Act
@@ -223,23 +228,25 @@ public async Task EnumerateToolsFromServersAsync_HandlesIndividualServerFailures
         var workingTools = new List<McpClientTool>();
 
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(servers);
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "failing-server"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .ThrowsAsync(new Exception("Server connection failed"));
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "working-server"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .ReturnsAsync(workingTools);
 
         // Act
@@ -274,15 +281,16 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesInParallel()
         var tcs3 = new TaskCompletionSource<IList<McpClientTool>>();
 
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(servers);
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "server1"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .Returns(async () =>
             {
                 lock (callOrder) callOrder.Add("server1-start");
@@ -296,7 +304,8 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesInParallel()
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "server2"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .Returns(async () =>
             {
                 lock (callOrder) callOrder.Add("server2-start");
@@ -310,7 +319,8 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesInParallel()
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "server3"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .Returns(async () =>
             {
                 lock (callOrder) callOrder.Add("server3-start");
@@ -372,23 +382,25 @@ public async Task EnumerateAllToolsAsync_ReturnsFlatListOfAllTools()
         var tools2 = new List<McpClientTool>();
 
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(servers);
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "server1"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .ReturnsAsync(tools1);
 
         _mockService
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "server2"),
                 It.IsAny<string>(),
-                It.IsAny<ToolOptions>()))
+                It.IsAny<ToolOptions>(),
+                It.IsAny<CancellationToken>()))
             .ReturnsAsync(tools2);
 
         // Act
@@ -409,7 +421,7 @@ public async Task EnumerateAllToolsAsync_WhenNoServers_ReturnsEmptyList()
         // Arrange
         var toolOptions = new ToolOptions();
         _mockService
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new List<MCPServerConfig>());
 
         // Act

src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/McpToolServerConfigurationService_V2TokenTests.cs

@@ -58,7 +58,7 @@ private static FakeTokenProvider TokenProvider(params (string scope, string toke
 
     private void SetupListServers(IEnumerable<MCPServerConfig> servers) =>
         _service
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(servers.ToList());
 
     // ─── AttachPerAudienceTokens (via ListToolServersWithTokensAsync) ─────────
@@ -246,16 +246,16 @@ public async Task EnumerateToolsFromServersAsync_WithTokenProvider_AttachesToken
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "v1"),
-                It.IsAny<string>(), It.IsAny<ToolOptions>()))
-            .Callback<ITurnContext, MCPServerConfig, string, ToolOptions>((_, s, _, _) => capturedV1 = s)
+                It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
+            .Callback<ITurnContext, MCPServerConfig, string, ToolOptions, CancellationToken>((_, s, _, _, _) => capturedV1 = s)
             .ReturnsAsync(new List<McpClientTool>());
 
         _service
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "v2"),
-                It.IsAny<string>(), It.IsAny<ToolOptions>()))
-            .Callback<ITurnContext, MCPServerConfig, string, ToolOptions>((_, s, _, _) => capturedV2 = s)
+                It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
+            .Callback<ITurnContext, MCPServerConfig, string, ToolOptions, CancellationToken>((_, s, _, _, _) => capturedV2 = s)
             .ReturnsAsync(new List<McpClientTool>());
 
         // Act
@@ -280,7 +280,7 @@ public async Task EnumerateToolsFromServersAsync_WithTokenProvider_WhenListFails
     {
         // Arrange
         _service
-            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+            .Setup(x => x.ListToolServersAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ThrowsAsync(new Exception("gateway unreachable"));
         var provider = TokenProvider();
 
@@ -308,7 +308,7 @@ public async Task EnumerateToolsFromServersAsync_WithTokenProvider_SkipsServersW
             .Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(),
                 It.Is<MCPServerConfig>(s => s.mcpServerName == "valid"),
-                It.IsAny<string>(), It.IsAny<ToolOptions>()))
+                It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new List<McpClientTool>());
 
         // Act
@@ -338,7 +338,7 @@ public async Task EnumerateToolsFromServersAsync_LegacyPath_V2Server_ThrowsWithM
             new Mock<IHttpClientFactory>().Object) { CallBase = true };
 
         service.Setup(x => x.ListToolServersAsync(
-                It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+                It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new List<MCPServerConfig> { V2Server("mail") });
 
         var act = () => service.Object.EnumerateToolsFromServersAsync(
@@ -369,11 +369,11 @@ public async Task EnumerateToolsFromServersAsync_LegacyPath_ApiPrefixedAtgAudien
             audience = $"api://{AtgAppId}"  // equivalent ATG audience form
         };
         service.Setup(x => x.ListToolServersAsync(
-                It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>()))
+                It.IsAny<string>(), It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new List<MCPServerConfig> { v1WithApiAudience });
         service.Setup(x => x.GetMcpClientToolsAsync(
                 It.IsAny<ITurnContext>(), It.IsAny<MCPServerConfig>(),
-                It.IsAny<string>(), It.IsAny<ToolOptions>()))
+                It.IsAny<string>(), It.IsAny<ToolOptions>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new List<McpClientTool>());
 
         // Act — should NOT throw; api://<AtgAppId> is a V1 server

src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/AddToolServersToAgent_Tests.cs

@@ -67,7 +67,8 @@ await service.AddToolServersToAgent(
                 TestAuthToken,
                 It.IsAny<IMcpTokenProvider>(),
                 mockTurnContext.Object,
-                It.Is<ToolOptions>(o => o.UserAgentConfiguration == Agent365AgentFrameworkSdkUserAgentConfiguration.Instance)),
+                It.Is<ToolOptions>(o => o.UserAgentConfiguration == Agent365AgentFrameworkSdkUserAgentConfiguration.Instance),
+                It.IsAny<CancellationToken>()),
             Times.Once);
     }
 

src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/GetMcpToolsAsync_Tests.cs

@@ -39,7 +39,8 @@ await service.GetMcpToolsAsync(
                 TestAuthToken,
                 It.IsAny<IMcpTokenProvider>(),
                 mockTurnContext.Object,
-                It.Is<ToolOptions>(o => o.UserAgentConfiguration == Agent365AgentFrameworkSdkUserAgentConfiguration.Instance)),
+                It.Is<ToolOptions>(o => o.UserAgentConfiguration == Agent365AgentFrameworkSdkUserAgentConfiguration.Instance),
+                It.IsAny<CancellationToken>()),
             Times.Once);
     }