diff --git a/src/Observability/Hosting/Caching/AgenticTokenCache.cs b/src/Observability/Hosting/Caching/AgenticTokenCache.cs
index bfe964e4..86fe99d5 100644
--- a/src/Observability/Hosting/Caching/AgenticTokenCache.cs
+++ b/src/Observability/Hosting/Caching/AgenticTokenCache.cs
@@ -96,7 +96,11 @@ public void RegisterObservability(string agentId, string tenantId, AgenticTokenS
///
/// The observability token if available; otherwise, null.
///
- public async Task GetObservabilityToken(string agentId, string tenantId)
+ public Task GetObservabilityToken(string agentId, string tenantId)
+ => GetObservabilityToken(agentId, tenantId, CancellationToken.None);
+
+ ///
+ public async Task GetObservabilityToken(string agentId, string tenantId, CancellationToken cancellationToken)
{
if (!_map.TryGetValue($"{agentId}:{tenantId}", out var entry))
return null;
diff --git a/src/Observability/Hosting/Caching/IExporterTokenCache.cs b/src/Observability/Hosting/Caching/IExporterTokenCache.cs
index f2ac5e5c..f0e45eca 100644
--- a/src/Observability/Hosting/Caching/IExporterTokenCache.cs
+++ b/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 where T : class
/// Returns an observability token (cached inside the credential) or null on failure/not registered.
///
Task GetObservabilityToken(string agentId, string tenantId);
+
+ ///
+ /// Returns an observability token (cached inside the credential) or null on failure/not registered, with cancellation support.
+ ///
+ Task GetObservabilityToken(string agentId, string tenantId, CancellationToken cancellationToken);
}
}
\ No newline at end of file
diff --git a/src/Observability/Hosting/Caching/ServiceTokenCache.cs b/src/Observability/Hosting/Caching/ServiceTokenCache.cs
index fe94a0bc..07424c86 100644
--- a/src/Observability/Hosting/Caching/ServiceTokenCache.cs
+++ b/src/Observability/Hosting/Caching/ServiceTokenCache.cs
@@ -122,15 +122,19 @@ public void RegisterObservability(string agentId, string tenantId, string token,
/// The agent identifier.
/// The tenant identifier.
/// The observability token if valid; otherwise, null.
- public async Task GetObservabilityToken(string agentId, string tenantId)
+ public Task GetObservabilityToken(string agentId, string tenantId)
+ => GetObservabilityToken(agentId, tenantId, CancellationToken.None);
+
+ ///
+ public Task GetObservabilityToken(string agentId, string tenantId, CancellationToken cancellationToken)
{
if (string.IsNullOrWhiteSpace(agentId) || string.IsNullOrWhiteSpace(tenantId))
- return null;
+ return Task.FromResult(null);
var key = GetKey(agentId, tenantId);
if (!_map.TryGetValue(key, out var entry))
- return null;
+ return Task.FromResult(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(null);
}
- return await Task.FromResult(entry.Token).ConfigureAwait(false);
+ return Task.FromResult(entry.Token);
}
///
diff --git a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterAsyncE2ETests.cs b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterAsyncE2ETests.cs
index d0897a4d..27da9626 100644
--- a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterAsyncE2ETests.cs
+++ b/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;
}
}
}
diff --git a/src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/McpToolServerConfigurationService_ToolEnumerationTests.cs b/src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/McpToolServerConfigurationService_ToolEnumerationTests.cs
index 600ef40d..d9b76f15 100644
--- a/src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/McpToolServerConfigurationService_ToolEnumerationTests.cs
+++ b/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(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.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(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(new List());
// Act
@@ -101,7 +102,7 @@ public async Task EnumerateToolsFromServersAsync_FiltersInvalidServers_WithMissi
};
_mockService
- .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(servers);
_mockService
@@ -109,7 +110,8 @@ public async Task EnumerateToolsFromServersAsync_FiltersInvalidServers_WithMissi
It.IsAny(),
It.Is(s => s.mcpServerName == "valid-server"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.ReturnsAsync(new List());
// Act
@@ -137,7 +139,7 @@ public async Task EnumerateToolsFromServersAsync_FiltersInvalidServers_WithMissi
};
_mockService
- .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(servers);
_mockService
@@ -145,7 +147,8 @@ public async Task EnumerateToolsFromServersAsync_FiltersInvalidServers_WithMissi
It.IsAny(),
It.Is(s => s.mcpServerName == "valid-server"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.ReturnsAsync(new List());
// Act
@@ -175,7 +178,7 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesToolsFromMultipleServ
var tools2 = new List();
_mockService
- .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(servers);
_mockService
@@ -183,7 +186,8 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesToolsFromMultipleServ
It.IsAny(),
It.Is(s => s.mcpServerName == "server1"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.ReturnsAsync(tools1);
_mockService
@@ -191,7 +195,8 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesToolsFromMultipleServ
It.IsAny(),
It.Is(s => s.mcpServerName == "server2"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.ReturnsAsync(tools2);
// Act
@@ -223,7 +228,7 @@ public async Task EnumerateToolsFromServersAsync_HandlesIndividualServerFailures
var workingTools = new List();
_mockService
- .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(servers);
_mockService
@@ -231,7 +236,8 @@ public async Task EnumerateToolsFromServersAsync_HandlesIndividualServerFailures
It.IsAny(),
It.Is(s => s.mcpServerName == "failing-server"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.ThrowsAsync(new Exception("Server connection failed"));
_mockService
@@ -239,7 +245,8 @@ public async Task EnumerateToolsFromServersAsync_HandlesIndividualServerFailures
It.IsAny(),
It.Is(s => s.mcpServerName == "working-server"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.ReturnsAsync(workingTools);
// Act
@@ -274,7 +281,7 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesInParallel()
var tcs3 = new TaskCompletionSource>();
_mockService
- .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(servers);
_mockService
@@ -282,7 +289,8 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesInParallel()
It.IsAny(),
It.Is(s => s.mcpServerName == "server1"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.Returns(async () =>
{
lock (callOrder) callOrder.Add("server1-start");
@@ -296,7 +304,8 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesInParallel()
It.IsAny(),
It.Is(s => s.mcpServerName == "server2"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.Returns(async () =>
{
lock (callOrder) callOrder.Add("server2-start");
@@ -310,7 +319,8 @@ public async Task EnumerateToolsFromServersAsync_EnumeratesInParallel()
It.IsAny(),
It.Is(s => s.mcpServerName == "server3"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.Returns(async () =>
{
lock (callOrder) callOrder.Add("server3-start");
@@ -372,7 +382,7 @@ public async Task EnumerateAllToolsAsync_ReturnsFlatListOfAllTools()
var tools2 = new List();
_mockService
- .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(servers);
_mockService
@@ -380,7 +390,8 @@ public async Task EnumerateAllToolsAsync_ReturnsFlatListOfAllTools()
It.IsAny(),
It.Is(s => s.mcpServerName == "server1"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.ReturnsAsync(tools1);
_mockService
@@ -388,7 +399,8 @@ public async Task EnumerateAllToolsAsync_ReturnsFlatListOfAllTools()
It.IsAny(),
It.Is(s => s.mcpServerName == "server2"),
It.IsAny(),
- It.IsAny()))
+ It.IsAny(),
+ It.IsAny()))
.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(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(new List());
// Act
diff --git a/src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/McpToolServerConfigurationService_V2TokenTests.cs b/src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/McpToolServerConfigurationService_V2TokenTests.cs
index a894de30..81fb0b9d 100644
--- a/src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/McpToolServerConfigurationService_V2TokenTests.cs
+++ b/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 servers) =>
_service
- .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(servers.ToList());
// ─── AttachPerAudienceTokens (via ListToolServersWithTokensAsync) ─────────
@@ -246,16 +246,16 @@ public async Task EnumerateToolsFromServersAsync_WithTokenProvider_AttachesToken
.Setup(x => x.GetMcpClientToolsAsync(
It.IsAny(),
It.Is(s => s.mcpServerName == "v1"),
- It.IsAny(), It.IsAny()))
- .Callback((_, s, _, _) => capturedV1 = s)
+ It.IsAny(), It.IsAny(), It.IsAny()))
+ .Callback((_, s, _, _, _) => capturedV1 = s)
.ReturnsAsync(new List());
_service
.Setup(x => x.GetMcpClientToolsAsync(
It.IsAny(),
It.Is(s => s.mcpServerName == "v2"),
- It.IsAny(), It.IsAny()))
- .Callback((_, s, _, _) => capturedV2 = s)
+ It.IsAny(), It.IsAny(), It.IsAny()))
+ .Callback((_, s, _, _, _) => capturedV2 = s)
.ReturnsAsync(new List());
// Act
@@ -280,7 +280,7 @@ public async Task EnumerateToolsFromServersAsync_WithTokenProvider_WhenListFails
{
// Arrange
_service
- .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny()))
+ .Setup(x => x.ListToolServersAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ThrowsAsync(new Exception("gateway unreachable"));
var provider = TokenProvider();
@@ -308,7 +308,7 @@ public async Task EnumerateToolsFromServersAsync_WithTokenProvider_SkipsServersW
.Setup(x => x.GetMcpClientToolsAsync(
It.IsAny(),
It.Is(s => s.mcpServerName == "valid"),
- It.IsAny(), It.IsAny()))
+ It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(new List());
// Act
@@ -338,7 +338,7 @@ public async Task EnumerateToolsFromServersAsync_LegacyPath_V2Server_ThrowsWithM
new Mock().Object) { CallBase = true };
service.Setup(x => x.ListToolServersAsync(
- It.IsAny(), It.IsAny(), It.IsAny()))
+ It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(new List { 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(), It.IsAny(), It.IsAny()))
+ It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(new List { v1WithApiAudience });
service.Setup(x => x.GetMcpClientToolsAsync(
It.IsAny(), It.IsAny(),
- It.IsAny(), It.IsAny()))
+ It.IsAny(), It.IsAny(), It.IsAny()))
.ReturnsAsync(new List());
// Act — should NOT throw; api:// is a V1 server
diff --git a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/AddToolServersToAgent_Tests.cs b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/AddToolServersToAgent_Tests.cs
index bad16f30..78b94d58 100644
--- a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/AddToolServersToAgent_Tests.cs
+++ b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/AddToolServersToAgent_Tests.cs
@@ -67,7 +67,8 @@ await service.AddToolServersToAgent(
TestAuthToken,
It.IsAny(),
mockTurnContext.Object,
- It.Is(o => o.UserAgentConfiguration == Agent365AgentFrameworkSdkUserAgentConfiguration.Instance)),
+ It.Is(o => o.UserAgentConfiguration == Agent365AgentFrameworkSdkUserAgentConfiguration.Instance),
+ It.IsAny()),
Times.Once);
}
diff --git a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/GetMcpToolsAsync_Tests.cs b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/GetMcpToolsAsync_Tests.cs
index 856cbb03..0ee2475d 100644
--- a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/GetMcpToolsAsync_Tests.cs
+++ b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/GetMcpToolsAsync_Tests.cs
@@ -39,7 +39,8 @@ await service.GetMcpToolsAsync(
TestAuthToken,
It.IsAny(),
mockTurnContext.Object,
- It.Is(o => o.UserAgentConfiguration == Agent365AgentFrameworkSdkUserAgentConfiguration.Instance)),
+ It.Is(o => o.UserAgentConfiguration == Agent365AgentFrameworkSdkUserAgentConfiguration.Instance),
+ It.IsAny()),
Times.Once);
}
diff --git a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/McpToolRegistrationServiceTestBase.cs b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/McpToolRegistrationServiceTestBase.cs
index c6a35368..99c0aa59 100644
--- a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/McpToolRegistrationServiceTestBase.cs
+++ b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/McpToolRegistrationServiceTestBase.cs
@@ -79,11 +79,12 @@ protected void SetupMocksForAddToolServers(Action? captureToolOptio
It.IsAny(),
It.IsAny(),
It.IsAny(),
- It.IsAny()));
+ It.IsAny(),
+ It.IsAny()));
if (captureToolOptions != null)
{
- setup.Callback((_, _, _, _, options) => captureToolOptions(options));
+ setup.Callback((_, _, _, _, options, _) => captureToolOptions(options));
}
setup.ReturnsAsync((new List(), new Dictionary>()));
@@ -104,11 +105,12 @@ protected void SetupMocksForGetMcpTools(Action? captureToolOptions
It.IsAny(),
It.IsAny(),
It.IsAny(),
- It.IsAny()));
+ It.IsAny(),
+ It.IsAny()));
if (captureToolOptions != null)
{
- setup.Callback((_, _, _, _, options) => captureToolOptions(options));
+ setup.Callback((_, _, _, _, options, _) => captureToolOptions(options));
}
setup.ReturnsAsync((new List(), new Dictionary>()));
diff --git a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AzureAIFoundry.Tests/McpToolRegistrationServiceTests.cs b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AzureAIFoundry.Tests/McpToolRegistrationServiceTests.cs
index d3ef8ef6..3c7d65cb 100644
--- a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AzureAIFoundry.Tests/McpToolRegistrationServiceTests.cs
+++ b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AzureAIFoundry.Tests/McpToolRegistrationServiceTests.cs
@@ -94,11 +94,12 @@ private void SetupMocksForEmptyToolEnumeration(Action? captureToolO
It.IsAny(),
It.IsAny(),
It.IsAny(),
- It.IsAny()));
+ It.IsAny