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/Observability/Hosting/ObservabilityServiceCollectionExtensions.cs b/src/Observability/Hosting/ObservabilityServiceCollectionExtensions.cs index 7e3875c0..04569ceb 100644 --- a/src/Observability/Hosting/ObservabilityServiceCollectionExtensions.cs +++ b/src/Observability/Hosting/ObservabilityServiceCollectionExtensions.cs @@ -27,7 +27,7 @@ public static IServiceCollection AddAgenticTracingExporter(this IServiceCollecti return new Agent365ExporterOptions { ClusterCategory = clusterCategory ?? "production", - TokenResolver = async (agentId, tenantId) => await cache.GetObservabilityToken(agentId, tenantId) + TokenResolver = (agentId, tenantId, ct) => cache.GetObservabilityToken(agentId, tenantId, ct) }; }); @@ -51,7 +51,7 @@ public static IServiceCollection AddServiceTracingExporter(this IServiceCollecti return new Agent365ExporterOptions { ClusterCategory = clusterCategory ?? "production", - TokenResolver = async (agentId, tenantId) => await cache.GetObservabilityToken(agentId, tenantId).ConfigureAwait(false), + TokenResolver = (agentId, tenantId, ct) => cache.GetObservabilityToken(agentId, tenantId, ct), UseS2SEndpoint = true // Service-to-service uses S2S endpoint }; }); diff --git a/src/Observability/Runtime/Tracing/Exporters/Agent365Exporter.cs b/src/Observability/Runtime/Tracing/Exporters/Agent365Exporter.cs index 4ad62bde..a389a959 100644 --- a/src/Observability/Runtime/Tracing/Exporters/Agent365Exporter.cs +++ b/src/Observability/Runtime/Tracing/Exporters/Agent365Exporter.cs @@ -72,7 +72,7 @@ public override ExportResult Export(in Batch batch) groups: groups, resource: _resource, options: _options, - tokenResolver: (agentId, tenantId) => _options.TokenResolver!(agentId, tenantId), + tokenResolver: (agentId, tenantId, ct) => _options.TokenResolver!(agentId, tenantId, ct), sendAsync: request => _httpClient.SendAsync(request) ).GetAwaiter().GetResult(); } diff --git a/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterAsync.cs b/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterAsync.cs index 3eccfcab..989ce349 100644 --- a/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterAsync.cs +++ b/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterAsync.cs @@ -74,8 +74,9 @@ await _core.ExportBatchCoreAsync( groups: groups, resource: this._resource, options: this._options, - tokenResolver: (agentId, tenantId) => this._options.TokenResolver!(agentId, tenantId), - sendAsync: request => this._httpClient.SendAsync(request, cancellationToken) + tokenResolver: (agentId, tenantId, ct) => this._options.TokenResolver!(agentId, tenantId, ct), + sendAsync: request => this._httpClient.SendAsync(request, cancellationToken), + cancellationToken: cancellationToken ).ConfigureAwait(false); } catch (OperationCanceledException) diff --git a/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterCore.cs b/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterCore.cs index 387e7db7..43187cab 100644 --- a/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterCore.cs +++ b/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterCore.cs @@ -13,6 +13,7 @@ using System.Net.Http; using System.Net.Http.Headers; using System.Text; +using System.Threading; using System.Threading.Tasks; namespace Microsoft.Agents.A365.Observability.Runtime.Tracing.Exporters @@ -127,13 +128,15 @@ public string BuildRequestUri(string endpoint, string endpointPath) /// /// /// + /// A cancellation token to cancel the operation. /// public async Task ExportBatchCoreAsync( IEnumerable<(string TenantId, string AgentId, List Activities)> groups, Resource resource, Agent365ExporterOptions options, - Func> tokenResolver, - Func> sendAsync) + Func> tokenResolver, + Func> sendAsync, + CancellationToken cancellationToken = default) { foreach (var g in groups) { @@ -157,7 +160,7 @@ public async Task ExportBatchCoreAsync( string? token = null; try { - token = await tokenResolver(agentId, tenantId).ConfigureAwait(false); + token = await tokenResolver(agentId, tenantId, cancellationToken).ConfigureAwait(false); this._logger?.LogDebug("Agent365ExporterCore: Obtained token for agent {AgentId} tenant {TenantId}.", agentId, tenantId); } catch (Exception ex) diff --git a/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterOptions.cs b/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterOptions.cs index e692576f..38d5c862 100644 --- a/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterOptions.cs +++ b/src/Observability/Runtime/Tracing/Exporters/Agent365ExporterOptions.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Threading; using System.Threading.Tasks; namespace Microsoft.Agents.A365.Observability.Runtime.Tracing.Exporters @@ -10,7 +11,7 @@ namespace Microsoft.Agents.A365.Observability.Runtime.Tracing.Exporters /// Must be fast and non-blocking (use internal caching elsewhere). /// Return null/empty to omit the Authorization header. /// - public delegate Task AsyncAuthTokenResolver(string agentId, string tenantId); + public delegate Task AsyncAuthTokenResolver(string agentId, string tenantId, CancellationToken cancellationToken = default); /// /// Delegate used by the exporter to resolve the endpoint host or URL for a given tenant id. diff --git a/src/Tests/Microsoft.Agents.A365.Observability.Hosting.Tests/Extensions/ObservabilityServiceCollectionExtensionsTests.cs b/src/Tests/Microsoft.Agents.A365.Observability.Hosting.Tests/Extensions/ObservabilityServiceCollectionExtensionsTests.cs index b8671c5f..9df03fe3 100644 --- a/src/Tests/Microsoft.Agents.A365.Observability.Hosting.Tests/Extensions/ObservabilityServiceCollectionExtensionsTests.cs +++ b/src/Tests/Microsoft.Agents.A365.Observability.Hosting.Tests/Extensions/ObservabilityServiceCollectionExtensionsTests.cs @@ -121,7 +121,7 @@ public void AddServiceTracingExporter_TokenResolver_CanBeCalled() var options = serviceProvider.GetRequiredService(); // Act - var token = options.TokenResolver!("test-agent", "test-tenant"); + var token = options.TokenResolver!("test-agent", "test-tenant", default); // Assert // Token resolver should not throw (actual token retrieval logic is in the cache) @@ -139,8 +139,7 @@ public void AddAgenticTracingExporter_TokenResolver_CanBeCalled() var options = serviceProvider.GetRequiredService(); // Act - var token = options.TokenResolver!("test-agent", "test-tenant"); - + var token = options.TokenResolver!("test-agent", "test-tenant", default); // Assert // Token resolver should not throw (actual token retrieval logic is in the cache) // This just verifies the resolver is wired up diff --git a/src/Tests/Microsoft.Agents.A365.Observability.Hosting.Tests/ObservabilityBuilderExtensionsTests.cs b/src/Tests/Microsoft.Agents.A365.Observability.Hosting.Tests/ObservabilityBuilderExtensionsTests.cs index 6161d7e9..93751f4d 100644 --- a/src/Tests/Microsoft.Agents.A365.Observability.Hosting.Tests/ObservabilityBuilderExtensionsTests.cs +++ b/src/Tests/Microsoft.Agents.A365.Observability.Hosting.Tests/ObservabilityBuilderExtensionsTests.cs @@ -83,7 +83,7 @@ public void AddA365Tracing_WithOpenTelemetryBuilderTrue_AndExporterEnabled_Regis builder.Services.AddSingleton(sp => new Agent365ExporterOptions { UseS2SEndpoint = false, - TokenResolver = (_, _) => Task.FromResult("test-token") + TokenResolver = (_, _, _) => Task.FromResult("test-token") }); }); webHostBuilder.UseStartup(); @@ -160,7 +160,7 @@ public void AddA365Tracing_IHostBuilder_WithOpenTelemetryBuilderTrue_AndExporter builder.Services.AddSingleton(_ => new Agent365ExporterOptions { UseS2SEndpoint = false, - TokenResolver = (_, _) => Task.FromResult("test-token") + TokenResolver = (_, _, _) => Task.FromResult("test-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 4ff51b23..07b2685a 100644 --- a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterAsyncE2ETests.cs +++ b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterAsyncE2ETests.cs @@ -1,416 +1,425 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -using FluentAssertions; -using Microsoft.Agents.A365.Observability.Runtime.Tracing.Contracts; -using Microsoft.Agents.A365.Observability.Runtime.Tracing.Exporters; -using Microsoft.Agents.A365.Observability.Runtime.Tracing.Scopes; -using Microsoft.Extensions.DependencyInjection; -using Microsoft.Extensions.Hosting; -using System.Net; -using System.Text.Json; - -namespace Microsoft.Agents.A365.Observability.Runtime.Tests.IntegrationTests -{ - [TestClass] - public class Agent365ExporterAsyncE2ETests - { - private TestHttpMessageHandler? _handler; - private ServiceProvider? _provider; - private bool _receivedRequest; - private string? _receivedContent; - - [TestMethod] - public async Task AddTracing_And_InvokeAgentScope_ExporterMakesExpectedRequest() - { - // Arrange - this.SetupExporterTest(); - this._receivedRequest = false; - this._receivedContent = null; - var expectedAgentType = AgentType.EntraEmbodied; - var expectedAgentDetails = new AgentDetails( - agentId: Guid.NewGuid().ToString(), - agentName: "Test Agent", - agentDescription: "Agent for testing.", - agentAUID: Guid.NewGuid().ToString(), - agentUPN: "testagent@ztaitest12.onmicrosoft.com", - agentBlueprintId: Guid.NewGuid().ToString(), - tenantId: Guid.NewGuid().ToString(), - agentType: expectedAgentType); - var endpoint = new Uri("https://test-agent-endpoint"); - var invokeAgentDetails = new InvokeAgentDetails(endpoint: endpoint, details: expectedAgentDetails); - var tenantDetails = new TenantDetails(Guid.NewGuid()); - - var expectedRequest = new Request( - content: "Test request content", - executionType: ExecutionType.HumanToAgent, - channel: new Channel( - name: "msteams", - link: "https://testchannel.link")); - - var expectedCallerDetails = new CallerDetails( - callerId: "caller-123", - callerName: "Test Caller", - callerUpn: "caller-123@ztaitest12.onmicrosoft.com", - callerClientIP: IPAddress.Parse("203.0.113.42"), - tenantId: expectedAgentDetails.TenantId); - - // Act - using (var scope = InvokeAgentScope.Start( - invokeAgentDetails: invokeAgentDetails, - tenantDetails: tenantDetails, - request: expectedRequest, - callerDetails: expectedCallerDetails)) - { - scope.RecordInputMessages(new[] { "Input message 1", "Input message 2" }); - scope.RecordOutputMessages(new[] { "Output message 1" }); - } - - var timeout = TimeSpan.FromSeconds(30); - var start = DateTime.UtcNow; - while (!this._receivedRequest && DateTime.UtcNow - start < timeout) - { - await Task.Delay(1000).ConfigureAwait(false); - } - - this._receivedRequest.Should().BeTrue("Exporter should make the expected HTTP request."); - this._receivedContent.Should().NotBeNull("Exporter should send a request body."); - - using var doc = JsonDocument.Parse(this._receivedContent!); - var root = doc.RootElement; - - var attributes = root - .GetProperty("resourceSpans")[0] - .GetProperty("scopeSpans")[0] - .GetProperty("spans")[0] - .GetProperty("attributes"); - this.GetAttribute(attributes, "server.address").Should().Be(invokeAgentDetails.Endpoint?.Host); - this.GetAttribute(attributes, "microsoft.channel.name").Should().Be(expectedRequest.Channel?.Name); - this.GetAttribute(attributes, "microsoft.channel.link").Should().Be(expectedRequest.Channel?.Link); - this.GetAttribute(attributes, "microsoft.tenant.id").Should().Be(tenantDetails.TenantId.ToString()); - this.GetAttribute(attributes, "user.id").Should().Be(expectedCallerDetails.CallerId); - this.GetAttribute(attributes, "user.email").Should().Be(expectedCallerDetails.CallerUpn); - this.GetAttribute(attributes, "user.name").Should().Be(expectedCallerDetails.CallerName); - this.GetAttribute(attributes, "gen_ai.input.messages").Should().Be("Input message 1,Input message 2"); - this.GetAttribute(attributes, "gen_ai.output.messages").Should().Be("Output message 1"); - this.GetAttribute(attributes, "gen_ai.agent.id").Should().Be(expectedAgentDetails.AgentId); - this.GetAttribute(attributes, "gen_ai.agent.name").Should().Be(expectedAgentDetails.AgentName); - this.GetAttribute(attributes, "gen_ai.agent.description").Should().Be(expectedAgentDetails.AgentDescription); - this.GetAttribute(attributes, "microsoft.agent.user.id").Should().Be(expectedAgentDetails.AgentAUID); - this.GetAttribute(attributes, "microsoft.agent.user.email").Should().Be(expectedAgentDetails.AgentUPN); - this.GetAttribute(attributes, "microsoft.a365.agent.blueprint.id").Should().Be(expectedAgentDetails.AgentBlueprintId); - this.GetAttribute(attributes, "microsoft.tenant.id").Should().Be(tenantDetails.TenantId.ToString()); - this.GetAttribute(attributes, "gen_ai.operation.name").Should().Be("invoke_agent"); - } - - [TestMethod] - public async Task AddTracing_And_ExecuteToolScope_ExporterMakesExpectedRequest() - { - // Arrange - this.SetupExporterTest(); - this._receivedRequest = false; - this._receivedContent = null; - var expectedAgentType = AgentType.EntraEmbodied; - var expectedAgentDetails = new AgentDetails( - agentId: Guid.NewGuid().ToString(), - agentName: "Tool Agent", - agentDescription: "Agent for tool execution.", - agentAUID: Guid.NewGuid().ToString(), - agentUPN: "toolagent@ztaitest12.onmicrosoft.com", - agentBlueprintId: Guid.NewGuid().ToString(), - tenantId: Guid.NewGuid().ToString(), - agentType: expectedAgentType); - var tenantDetails = new TenantDetails(Guid.NewGuid()); - var endpoint = new Uri("https://tool-endpoint:8443"); - var toolCallDetails = new ToolCallDetails( - toolName: "TestTool", - arguments: "{\"param\":\"value\"}", - toolCallId: "call-456", - description: "Test tool call description", - toolType: "custom-type", - endpoint: endpoint); - - // Act - using (var scope = ExecuteToolScope.Start(toolCallDetails, expectedAgentDetails, tenantDetails)) - { - scope.RecordResponse("Tool response content"); - } - - var timeout = TimeSpan.FromSeconds(30); - var start = DateTime.UtcNow; - while (!this._receivedRequest && DateTime.UtcNow - start < timeout) - { - await Task.Delay(1000).ConfigureAwait(false); - } - - this._receivedRequest.Should().BeTrue("Exporter should make the expected HTTP request."); - this._receivedContent.Should().NotBeNull("Exporter should send a request body."); - - using var doc = JsonDocument.Parse(this._receivedContent!); - var root = doc.RootElement; - - var attributes = root - .GetProperty("resourceSpans")[0] - .GetProperty("scopeSpans")[0] - .GetProperty("spans")[0] - .GetProperty("attributes"); - - this.GetAttribute(attributes, "gen_ai.operation.name").Should().Be("execute_tool"); - this.GetAttribute(attributes, "gen_ai.agent.id").Should().Be(expectedAgentDetails.AgentId); - this.GetAttribute(attributes, "gen_ai.agent.name").Should().Be(expectedAgentDetails.AgentName); - this.GetAttribute(attributes, "gen_ai.agent.description").Should().Be(expectedAgentDetails.AgentDescription); - this.GetAttribute(attributes, "microsoft.agent.user.id").Should().Be(expectedAgentDetails.AgentAUID); - this.GetAttribute(attributes, "microsoft.agent.user.email").Should().Be(expectedAgentDetails.AgentUPN); - this.GetAttribute(attributes, "microsoft.a365.agent.blueprint.id").Should().Be(expectedAgentDetails.AgentBlueprintId); - this.GetAttribute(attributes, "microsoft.tenant.id").Should().Be(tenantDetails.TenantId.ToString()); - this.GetAttribute(attributes, "gen_ai.tool.name").Should().Be(toolCallDetails.ToolName); - this.GetAttribute(attributes, "gen_ai.tool.arguments").Should().Be(toolCallDetails.Arguments); - this.GetAttribute(attributes, "gen_ai.tool.call.id").Should().Be(toolCallDetails.ToolCallId); - this.GetAttribute(attributes, "gen_ai.tool.description").Should().Be(toolCallDetails.Description); - this.GetAttribute(attributes, "gen_ai.tool.type").Should().Be(toolCallDetails.ToolType); - this.GetAttribute(attributes, "server.address").Should().Be(endpoint.Host); - this.GetAttribute(attributes, "server.port").Should().Be(endpoint.Port.ToString()); - this.GetAttribute(attributes, "gen_ai.tool.call.result").Should().Be("Tool response content"); - } - - [TestMethod] - public async Task AddTracing_And_InferenceScope_ExporterMakesExpectedRequest() - { - // Arrange - this.SetupExporterTest(); - this._receivedRequest = false; - this._receivedContent = null; - var expectedAgentType = AgentType.EntraEmbodied; - var expectedAgentDetails = new AgentDetails( - agentId: Guid.NewGuid().ToString(), - agentName: "Inference Agent", - agentDescription: "Agent for inference testing.", - agentAUID: Guid.NewGuid().ToString(), - agentUPN: "inferenceagent@ztaitest12.onmicrosoft.com", - agentBlueprintId: Guid.NewGuid().ToString(), - tenantId: Guid.NewGuid().ToString(), - agentType: expectedAgentType); - var tenantDetails = new TenantDetails(Guid.NewGuid()); - - var inferenceDetails = new InferenceCallDetails( - operationName: InferenceOperationType.Chat, - model: "gpt-4", - providerName: "OpenAI", - inputTokens: 42, - outputTokens: 84, - finishReasons: new[] { "stop", "length" }, - responseId: "response-xyz"); - - // Act - using (var scope = InferenceScope.Start(inferenceDetails, expectedAgentDetails, tenantDetails)) - { - scope.RecordInputMessages(new[] { "Hello", "World" }); - scope.RecordOutputMessages(new[] { "Hi there!" }); - scope.RecordInputTokens(42); - scope.RecordOutputTokens(84); - scope.RecordFinishReasons(new[] { "stop", "length" }); - } - - var timeout = TimeSpan.FromSeconds(30); - var start = DateTime.UtcNow; - while (!this._receivedRequest && DateTime.UtcNow - start < timeout) - { - await Task.Delay(1000).ConfigureAwait(false); - } - - this._receivedRequest.Should().BeTrue("Exporter should make the expected HTTP request."); - this._receivedContent.Should().NotBeNull("Exporter should send a request body."); - - using var doc = JsonDocument.Parse(this._receivedContent!); - var root = doc.RootElement; - var attributes = root - .GetProperty("resourceSpans")[0] - .GetProperty("scopeSpans")[0] - .GetProperty("spans")[0] - .GetProperty("attributes"); - - this.GetAttribute(attributes, "gen_ai.operation.name").Should().Be(inferenceDetails.OperationName.ToString()); - this.GetAttribute(attributes, "gen_ai.agent.id").Should().Be(expectedAgentDetails.AgentId); - this.GetAttribute(attributes, "gen_ai.agent.name").Should().Be(expectedAgentDetails.AgentName); - this.GetAttribute(attributes, "gen_ai.agent.description").Should().Be(expectedAgentDetails.AgentDescription); - this.GetAttribute(attributes, "microsoft.agent.user.id").Should().Be(expectedAgentDetails.AgentAUID); - this.GetAttribute(attributes, "microsoft.agent.user.email").Should().Be(expectedAgentDetails.AgentUPN); - this.GetAttribute(attributes, "microsoft.a365.agent.blueprint.id").Should().Be(expectedAgentDetails.AgentBlueprintId); - this.GetAttribute(attributes, "microsoft.tenant.id").Should().Be(tenantDetails.TenantId.ToString()); - this.GetAttribute(attributes, "gen_ai.request.model").Should().Be(inferenceDetails.Model); - this.GetAttribute(attributes, "gen_ai.provider.name").Should().Be(inferenceDetails.ProviderName); - this.GetAttribute(attributes, "gen_ai.usage.input_tokens").Should().Be("42"); - this.GetAttribute(attributes, "gen_ai.usage.output_tokens").Should().Be("84"); - this.GetAttribute(attributes, "gen_ai.response.finish_reasons").Should().Be("stop,length"); - this.GetAttribute(attributes, "gen_ai.input.messages").Should().Be("Hello,World"); - this.GetAttribute(attributes, "gen_ai.output.messages").Should().Be("Hi there!"); - } - - [TestMethod] - public async Task AddTracing_NestedScopes_AllExporterRequestsReceived() - { - // Arrange - List receivedContents = new(); - - var agentType = AgentType.EntraEmbodied; - var agentDetails = new AgentDetails( - agentId: Guid.NewGuid().ToString(), - agentName: "Nested Agent", - agentDescription: "Agent for nested scope testing.", - agentAUID: Guid.NewGuid().ToString(), - agentUPN: "nestedagent@ztaitest12.onmicrosoft.com", - agentBlueprintId: Guid.NewGuid().ToString(), - tenantId: Guid.NewGuid().ToString(), - agentType: agentType); - - var tenantDetails = new TenantDetails(Guid.NewGuid()); - var endpoint = new Uri("https://nested-endpoint"); - - var handler = new TestHttpMessageHandler(req => - { - receivedContents.Add(req.Content?.ReadAsStringAsync().GetAwaiter().GetResult() ?? ""); - return new HttpResponseMessage(System.Net.HttpStatusCode.OK); - }); - var httpClient = new HttpClient(handler); - - this.CreateTestServiceProvider(httpClient); - - var invokeAgentDetails = new InvokeAgentDetails(endpoint: endpoint, details: agentDetails); - var request = new Request( - content: "Nested request", - executionType: ExecutionType.HumanToAgent, - channel: new Channel(name: "nested", link: "https://nestedchannel.link")); - - var toolCallDetails = new ToolCallDetails( - toolName: "NestedTool", - arguments: "{\"param\":\"nested\"}", - toolCallId: "call-nested", - description: "Nested tool call", - toolType: "nested-type", - endpoint: endpoint); - - var inferenceDetails = new InferenceCallDetails( - operationName: InferenceOperationType.Chat, - model: "gpt-nested", - providerName: "OpenAI", - inputTokens: 10, - outputTokens: 20, - finishReasons: new[] { "stop" }, - responseId: "response-nested"); - - // Act - using (var agentScope = InvokeAgentScope.Start(invokeAgentDetails, tenantDetails, request)) - { - agentScope.RecordInputMessages(new[] { "Agent input" }); - agentScope.RecordOutputMessages(new[] { "Agent output" }); - - using (var toolScope = ExecuteToolScope.Start(toolCallDetails, agentDetails, tenantDetails)) - { - toolScope.RecordResponse("Tool response"); - - using (var inferenceScope = InferenceScope.Start(inferenceDetails, agentDetails, tenantDetails)) - { - inferenceScope.RecordInputMessages(new[] { "Inference input" }); - inferenceScope.RecordOutputMessages(new[] { "Inference output" }); - inferenceScope.RecordInputTokens(10); - inferenceScope.RecordOutputTokens(20); - inferenceScope.RecordFinishReasons(new[] { "stop" }); - } - } - } - - // Wait for up to 5 seconds for all spans to be exported - await Task.Delay(5000).ConfigureAwait(false); - - // Assert - var allOperationNames = new List(); - foreach (var content in receivedContents) - { - using var doc = JsonDocument.Parse(content); - var root = doc.RootElement; - var spans = root - .GetProperty("resourceSpans")[0] - .GetProperty("scopeSpans")[0] - .GetProperty("spans") - .EnumerateArray(); - - foreach (var span in spans) - { - var opName = this.GetAttribute(span.GetProperty("attributes"), "gen_ai.operation.name"); - if (opName != null) - allOperationNames.Add(opName); - } - } - allOperationNames.Should().Contain(new[] { "invoke_agent", "execute_tool", InferenceOperationType.Chat.ToString() }, "All three nested scopes should be exported, even if batched in fewer requests."); - } - - private class TestHttpMessageHandler : HttpMessageHandler - { - private Func _handler; - public TestHttpMessageHandler(Func handler) - { - this._handler = handler; - } - public void SetHandler(Func handler) - { - this._handler = handler; - } - protected override Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) - { - return Task.FromResult(this._handler(request)); - } - } - - private string? GetAttribute(JsonElement attributes, string key) - { - if (attributes.TryGetProperty(key, out var value)) - { - if (value.ValueKind == JsonValueKind.String) - { - return value.GetString(); - } - if (value.ValueKind == JsonValueKind.Number) - { - return value.GetRawText(); - } - if (value.ValueKind == JsonValueKind.Object && value.TryGetProperty("stringValue", out var sv)) - { - return sv.GetString(); - } - } - return null; - } - - private ServiceProvider CreateTestServiceProvider(HttpClient httpClient) - { - HostApplicationBuilder builder = new HostApplicationBuilder(); - - builder.Configuration["EnableAgent365Exporter"] = "true"; - builder.Services.AddSingleton(httpClient); - builder.Services.AddSingleton(sp => - { - return new Agent365ExporterOptions - { - UseS2SEndpoint = false, - TokenResolver = (_, _) => Task.FromResult("test-token") - }; - }); - - builder.AddA365Tracing(useOpenTelemetryBuilder: false, agent365ExporterType: Agent365ExporterType.Agent365ExporterAsync); - return builder.Services.BuildServiceProvider(); - } - private void SetupExporterTest() - { - this._handler = new TestHttpMessageHandler(req => - { - this._receivedRequest = true; - this._receivedContent = req.Content?.ReadAsStringAsync().GetAwaiter().GetResult(); - req.RequestUri.Should().NotBeNull(); - req.Headers.Authorization.Should().NotBeNull(); - return new HttpResponseMessage(System.Net.HttpStatusCode.OK); - }); - var httpClient = new HttpClient(this._handler); - this._provider = this.CreateTestServiceProvider(httpClient); - } - } -} +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using FluentAssertions; +using Microsoft.Agents.A365.Observability.Runtime.Tracing.Contracts; +using Microsoft.Agents.A365.Observability.Runtime.Tracing.Exporters; +using Microsoft.Agents.A365.Observability.Runtime.Tracing.Scopes; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; +using System.Net; +using System.Text.Json; + +namespace Microsoft.Agents.A365.Observability.Runtime.Tests.IntegrationTests +{ + [TestClass] + public class Agent365ExporterAsyncE2ETests + { + private HttpClient? _httpClient; + private TestHttpMessageHandler? _handler; + private ServiceProvider? _provider; + private bool _receivedRequest; + private string? _receivedContent; + + [TestMethod] + public async Task AddTracing_And_InvokeAgentScope_ExporterMakesExpectedRequest() + { + // Arrange + this.SetupExporterTest(); + this._receivedRequest = false; + this._receivedContent = null; + var expectedAgentType = AgentType.EntraEmbodied; + var expectedAgentDetails = new AgentDetails( + agentId: Guid.NewGuid().ToString(), + agentName: "Test Agent", + agentDescription: "Agent for testing.", + agentAUID: Guid.NewGuid().ToString(), + agentUPN: "testagent@ztaitest12.onmicrosoft.com", + agentBlueprintId: Guid.NewGuid().ToString(), + tenantId: Guid.NewGuid().ToString(), + agentType: expectedAgentType); + var endpoint = new Uri("https://test-agent-endpoint"); + var invokeAgentDetails = new InvokeAgentDetails(endpoint: endpoint, details: expectedAgentDetails); + var tenantDetails = new TenantDetails(Guid.NewGuid()); + + var expectedRequest = new Request( + content: "Test request content", + executionType: ExecutionType.HumanToAgent, + channel: new Channel( + name: "msteams", + link: "https://testchannel.link")); + + var expectedCallerDetails = new CallerDetails( + callerId: "caller-123", + callerName: "Test Caller", + callerUpn: "caller-123@ztaitest12.onmicrosoft.com", + callerClientIP: IPAddress.Parse("203.0.113.42"), + tenantId: expectedAgentDetails.TenantId); + + // Act + using (var scope = InvokeAgentScope.Start( + invokeAgentDetails: invokeAgentDetails, + tenantDetails: tenantDetails, + request: expectedRequest, + callerDetails: expectedCallerDetails)) + { + scope.RecordInputMessages(new[] { "Input message 1", "Input message 2" }); + scope.RecordOutputMessages(new[] { "Output message 1" }); + } + + var timeout = TimeSpan.FromSeconds(30); + var start = DateTime.UtcNow; + while (!this._receivedRequest && DateTime.UtcNow - start < timeout) + { + await Task.Delay(1000).ConfigureAwait(false); + } + + this._receivedRequest.Should().BeTrue("Exporter should make the expected HTTP request."); + this._receivedContent.Should().NotBeNull("Exporter should send a request body."); + + var content = this._receivedContent!; + using var doc = JsonDocument.Parse(content); + var root = doc.RootElement; + + var attributes = root + .GetProperty("resourceSpans")[0] + .GetProperty("scopeSpans")[0] + .GetProperty("spans")[0] + .GetProperty("attributes"); + this.GetAttribute(attributes, "server.address").Should().Be(invokeAgentDetails.Endpoint?.Host); + this.GetAttribute(attributes, "microsoft.channel.name").Should().Be(expectedRequest.Channel?.Name); + this.GetAttribute(attributes, "microsoft.channel.link").Should().Be(expectedRequest.Channel?.Link); + this.GetAttribute(attributes, "microsoft.tenant.id").Should().Be(tenantDetails.TenantId.ToString()); + this.GetAttribute(attributes, "user.id").Should().Be(expectedCallerDetails.CallerId); + this.GetAttribute(attributes, "user.email").Should().Be(expectedCallerDetails.CallerUpn); + this.GetAttribute(attributes, "user.name").Should().Be(expectedCallerDetails.CallerName); + this.GetAttribute(attributes, "gen_ai.input.messages").Should().Be("Input message 1,Input message 2"); + this.GetAttribute(attributes, "gen_ai.output.messages").Should().Be("Output message 1"); + this.GetAttribute(attributes, "gen_ai.agent.id").Should().Be(expectedAgentDetails.AgentId); + this.GetAttribute(attributes, "gen_ai.agent.name").Should().Be(expectedAgentDetails.AgentName); + this.GetAttribute(attributes, "gen_ai.agent.description").Should().Be(expectedAgentDetails.AgentDescription); + this.GetAttribute(attributes, "microsoft.agent.user.id").Should().Be(expectedAgentDetails.AgentAUID); + this.GetAttribute(attributes, "microsoft.agent.user.email").Should().Be(expectedAgentDetails.AgentUPN); + this.GetAttribute(attributes, "microsoft.a365.agent.blueprint.id").Should().Be(expectedAgentDetails.AgentBlueprintId); + this.GetAttribute(attributes, "microsoft.tenant.id").Should().Be(tenantDetails.TenantId.ToString()); + this.GetAttribute(attributes, "gen_ai.operation.name").Should().Be("invoke_agent"); + } + + [TestMethod] + public async Task AddTracing_And_ExecuteToolScope_ExporterMakesExpectedRequest() + { + // Arrange + this.SetupExporterTest(); + this._receivedRequest = false; + this._receivedContent = null; + var expectedAgentType = AgentType.EntraEmbodied; + var expectedAgentDetails = new AgentDetails( + agentId: Guid.NewGuid().ToString(), + agentName: "Tool Agent", + agentDescription: "Agent for tool execution.", + agentAUID: Guid.NewGuid().ToString(), + agentUPN: "toolagent@ztaitest12.onmicrosoft.com", + agentBlueprintId: Guid.NewGuid().ToString(), + tenantId: Guid.NewGuid().ToString(), + agentType: expectedAgentType); + var tenantDetails = new TenantDetails(Guid.NewGuid()); + var endpoint = new Uri("https://tool-endpoint:8443"); + var toolCallDetails = new ToolCallDetails( + toolName: "TestTool", + arguments: "{\"param\":\"value\"}", + toolCallId: "call-456", + description: "Test tool call description", + toolType: "custom-type", + endpoint: endpoint); + + // Act + using (var scope = ExecuteToolScope.Start(toolCallDetails, expectedAgentDetails, tenantDetails)) + { + scope.RecordResponse("Tool response content"); + } + + var timeout = TimeSpan.FromSeconds(30); + var start = DateTime.UtcNow; + while (!this._receivedRequest && DateTime.UtcNow - start < timeout) + { + await Task.Delay(1000).ConfigureAwait(false); + } + + this._receivedRequest.Should().BeTrue("Exporter should make the expected HTTP request."); + this._receivedContent.Should().NotBeNull("Exporter should send a request body."); + + using var doc = JsonDocument.Parse(this._receivedContent!); + var root = doc.RootElement; + + var attributes = root + .GetProperty("resourceSpans")[0] + .GetProperty("scopeSpans")[0] + .GetProperty("spans")[0] + .GetProperty("attributes"); + + this.GetAttribute(attributes, "gen_ai.operation.name").Should().Be("execute_tool"); + this.GetAttribute(attributes, "gen_ai.agent.id").Should().Be(expectedAgentDetails.AgentId); + this.GetAttribute(attributes, "gen_ai.agent.name").Should().Be(expectedAgentDetails.AgentName); + this.GetAttribute(attributes, "gen_ai.agent.description").Should().Be(expectedAgentDetails.AgentDescription); + this.GetAttribute(attributes, "microsoft.agent.user.id").Should().Be(expectedAgentDetails.AgentAUID); + this.GetAttribute(attributes, "microsoft.agent.user.email").Should().Be(expectedAgentDetails.AgentUPN); + this.GetAttribute(attributes, "microsoft.a365.agent.blueprint.id").Should().Be(expectedAgentDetails.AgentBlueprintId); + this.GetAttribute(attributes, "microsoft.tenant.id").Should().Be(tenantDetails.TenantId.ToString()); + this.GetAttribute(attributes, "gen_ai.tool.name").Should().Be(toolCallDetails.ToolName); + this.GetAttribute(attributes, "gen_ai.tool.arguments").Should().Be(toolCallDetails.Arguments); + this.GetAttribute(attributes, "gen_ai.tool.call.id").Should().Be(toolCallDetails.ToolCallId); + this.GetAttribute(attributes, "gen_ai.tool.description").Should().Be(toolCallDetails.Description); + this.GetAttribute(attributes, "gen_ai.tool.type").Should().Be(toolCallDetails.ToolType); + this.GetAttribute(attributes, "server.address").Should().Be(endpoint.Host); + this.GetAttribute(attributes, "server.port").Should().Be(endpoint.Port.ToString()); + this.GetAttribute(attributes, "gen_ai.tool.call.result").Should().Be("Tool response content"); + } + + [TestMethod] + public async Task AddTracing_And_InferenceScope_ExporterMakesExpectedRequest() + { + // Arrange + this.SetupExporterTest(); + this._receivedRequest = false; + this._receivedContent = null; + var expectedAgentType = AgentType.EntraEmbodied; + var expectedAgentDetails = new AgentDetails( + agentId: Guid.NewGuid().ToString(), + agentName: "Inference Agent", + agentDescription: "Agent for inference testing.", + agentAUID: Guid.NewGuid().ToString(), + agentUPN: "inferenceagent@ztaitest12.onmicrosoft.com", + agentBlueprintId: Guid.NewGuid().ToString(), + tenantId: Guid.NewGuid().ToString(), + agentType: expectedAgentType); + var tenantDetails = new TenantDetails(Guid.NewGuid()); + + var inferenceDetails = new InferenceCallDetails( + operationName: InferenceOperationType.Chat, + model: "gpt-4", + providerName: "OpenAI", + inputTokens: 42, + outputTokens: 84, + finishReasons: new[] { "stop", "length" }, + responseId: "response-xyz"); + + // Act + using (var scope = InferenceScope.Start(inferenceDetails, expectedAgentDetails, tenantDetails)) + { + scope.RecordInputMessages(new[] { "Hello", "World" }); + scope.RecordOutputMessages(new[] { "Hi there!" }); + scope.RecordInputTokens(42); + scope.RecordOutputTokens(84); + scope.RecordFinishReasons(new[] { "stop", "length" }); + } + + var timeout = TimeSpan.FromSeconds(30); + var start = DateTime.UtcNow; + while (!this._receivedRequest && DateTime.UtcNow - start < timeout) + { + await Task.Delay(1000).ConfigureAwait(false); + } + + this._receivedRequest.Should().BeTrue("Exporter should make the expected HTTP request."); + this._receivedContent.Should().NotBeNull("Exporter should send a request body."); + + using var doc = JsonDocument.Parse(this._receivedContent!); + var root = doc.RootElement; + var attributes = root + .GetProperty("resourceSpans")[0] + .GetProperty("scopeSpans")[0] + .GetProperty("spans")[0] + .GetProperty("attributes"); + + this.GetAttribute(attributes, "gen_ai.operation.name").Should().Be(inferenceDetails.OperationName.ToString()); + this.GetAttribute(attributes, "gen_ai.agent.id").Should().Be(expectedAgentDetails.AgentId); + this.GetAttribute(attributes, "gen_ai.agent.name").Should().Be(expectedAgentDetails.AgentName); + this.GetAttribute(attributes, "gen_ai.agent.description").Should().Be(expectedAgentDetails.AgentDescription); + this.GetAttribute(attributes, "microsoft.agent.user.id").Should().Be(expectedAgentDetails.AgentAUID); + this.GetAttribute(attributes, "microsoft.agent.user.email").Should().Be(expectedAgentDetails.AgentUPN); + this.GetAttribute(attributes, "microsoft.a365.agent.blueprint.id").Should().Be(expectedAgentDetails.AgentBlueprintId); + this.GetAttribute(attributes, "microsoft.tenant.id").Should().Be(tenantDetails.TenantId.ToString()); + this.GetAttribute(attributes, "gen_ai.request.model").Should().Be(inferenceDetails.Model); + this.GetAttribute(attributes, "gen_ai.provider.name").Should().Be(inferenceDetails.ProviderName); + this.GetAttribute(attributes, "gen_ai.usage.input_tokens").Should().Be("42"); + this.GetAttribute(attributes, "gen_ai.usage.output_tokens").Should().Be("84"); + this.GetAttribute(attributes, "gen_ai.response.finish_reasons").Should().Be("stop,length"); + this.GetAttribute(attributes, "gen_ai.input.messages").Should().Be("Hello,World"); + this.GetAttribute(attributes, "gen_ai.output.messages").Should().Be("Hi there!"); + } + + [TestMethod] + public async Task AddTracing_NestedScopes_AllExporterRequestsReceived() + { + // Arrange + List receivedContents = new(); + + var agentType = AgentType.EntraEmbodied; + var agentDetails = new AgentDetails( + agentId: Guid.NewGuid().ToString(), + agentName: "Nested Agent", + agentDescription: "Agent for nested scope testing.", + agentAUID: Guid.NewGuid().ToString(), + agentUPN: "nestedagent@ztaitest12.onmicrosoft.com", + agentBlueprintId: Guid.NewGuid().ToString(), + tenantId: Guid.NewGuid().ToString(), + agentType: agentType); + + var tenantDetails = new TenantDetails(Guid.NewGuid()); + var endpoint = new Uri("https://nested-endpoint"); + + var handler = new TestHttpMessageHandler(req => + { + receivedContents.Add(req.Content?.ReadAsStringAsync().GetAwaiter().GetResult() ?? ""); + return new HttpResponseMessage(System.Net.HttpStatusCode.OK); + }); + var httpClient = new HttpClient(handler); + + this.CreateTestServiceProvider(httpClient); + + var invokeAgentDetails = new InvokeAgentDetails(endpoint: endpoint, details: agentDetails); + var request = new Request( + content: "Nested request", + executionType: ExecutionType.HumanToAgent, + channel: new Channel(name: "nested", link: "https://nestedchannel.link")); + + var toolCallDetails = new ToolCallDetails( + toolName: "NestedTool", + arguments: "{\"param\":\"nested\"}", + toolCallId: "call-nested", + description: "Nested tool call", + toolType: "nested-type", + endpoint: endpoint); + + var inferenceDetails = new InferenceCallDetails( + operationName: InferenceOperationType.Chat, + model: "gpt-nested", + providerName: "OpenAI", + inputTokens: 10, + outputTokens: 20, + finishReasons: new[] { "stop" }, + responseId: "response-nested"); + + // Act + using (var agentScope = InvokeAgentScope.Start(invokeAgentDetails, tenantDetails, request)) + { + agentScope.RecordInputMessages(new[] { "Agent input" }); + agentScope.RecordOutputMessages(new[] { "Agent output" }); + + using (var toolScope = ExecuteToolScope.Start(toolCallDetails, agentDetails, tenantDetails)) + { + toolScope.RecordResponse("Tool response"); + + using (var inferenceScope = InferenceScope.Start(inferenceDetails, agentDetails, tenantDetails)) + { + inferenceScope.RecordInputMessages(new[] { "Inference input" }); + inferenceScope.RecordOutputMessages(new[] { "Inference output" }); + inferenceScope.RecordInputTokens(10); + inferenceScope.RecordOutputTokens(20); + inferenceScope.RecordFinishReasons(new[] { "stop" }); + } + } + } + + // Wait for up to 5 seconds for all spans to be exported + await Task.Delay(5000).ConfigureAwait(false); + + // Assert + var allOperationNames = new List(); + foreach (var content in receivedContents) + { + using var doc = JsonDocument.Parse(content); + var root = doc.RootElement; + var spans = root + .GetProperty("resourceSpans")[0] + .GetProperty("scopeSpans")[0] + .GetProperty("spans") + .EnumerateArray(); + + allOperationNames.AddRange( + spans + .Select(span => this.GetAttribute(span.GetProperty("attributes"), "gen_ai.operation.name")) + .Where(opName => opName != null)!); + } + allOperationNames.Should().Contain(new[] { "invoke_agent", "execute_tool", InferenceOperationType.Chat.ToString() }, "All three nested scopes should be exported, even if batched in fewer requests."); + } + + private class TestHttpMessageHandler : HttpMessageHandler + { + private Func _handler; + public TestHttpMessageHandler(Func handler) + { + this._handler = handler; + } + public void SetHandler(Func handler) + { + this._handler = handler; + } + protected override Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + return Task.FromResult(this._handler(request)); + } + } + + private string? GetAttribute(JsonElement attributes, string key) + { + if (attributes.TryGetProperty(key, out var value)) + { + if (value.ValueKind == JsonValueKind.String) + { + return value.GetString(); + } + if (value.ValueKind == JsonValueKind.Number) + { + return value.GetRawText(); + } + if (value.ValueKind == JsonValueKind.Object && value.TryGetProperty("stringValue", out var sv)) + { + return sv.GetString(); + } + } + return null; + } + + private ServiceProvider CreateTestServiceProvider(HttpClient httpClient) + { + HostApplicationBuilder builder = new HostApplicationBuilder(); + + builder.Configuration["EnableAgent365Exporter"] = "true"; + builder.Services.AddSingleton(httpClient); + builder.Services.AddSingleton(sp => + { + return new Agent365ExporterOptions + { + UseS2SEndpoint = false, + TokenResolver = (_, _, _) => Task.FromResult("test-token") + }; + }); + + builder.AddA365Tracing(useOpenTelemetryBuilder: false, agent365ExporterType: Agent365ExporterType.Agent365ExporterAsync); + return builder.Services.BuildServiceProvider(); + } + private void SetupExporterTest() + { + this._handler = new TestHttpMessageHandler(req => + { + this._receivedRequest = true; + this._receivedContent = req.Content?.ReadAsStringAsync().GetAwaiter().GetResult(); + req.RequestUri.Should().NotBeNull(); + req.Headers.Authorization.Should().NotBeNull(); + return new HttpResponseMessage(System.Net.HttpStatusCode.OK); + }); + this._httpClient = new HttpClient(this._handler); + this._provider = this.CreateTestServiceProvider(this._httpClient); + } + + [TestCleanup] + public void Cleanup() + { + (this._provider as IDisposable)?.Dispose(); + this._provider = null; + this._httpClient?.Dispose(); + this._httpClient = null; + } + } +} diff --git a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterE2ETests.cs b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterE2ETests.cs index 05f152e7..668b0376 100644 --- a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterE2ETests.cs +++ b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.IntegrationTests/Agent365ExporterE2ETests.cs @@ -543,7 +543,7 @@ public async Task AddTracing_MultipleInvocations_NoDuplicateExports() builder.Services.AddSingleton(_ => new Agent365ExporterOptions { UseS2SEndpoint = false, - TokenResolver = (_, _) => Task.FromResult("test-token") + TokenResolver = (_, _, _) => Task.FromResult("test-token") }); // AddA365Tracing call @@ -636,7 +636,7 @@ private ServiceProvider CreateTestServiceProvider(HttpClient httpClient) return new Agent365ExporterOptions { UseS2SEndpoint = false, - TokenResolver = (_, _) => Task.FromResult("test-token") + TokenResolver = (_, _, _) => Task.FromResult("test-token") }; }); builder.AddA365Tracing(useOpenTelemetryBuilder: false, agent365ExporterType: Agent365ExporterType.Agent365Exporter); diff --git a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/BuilderTests.cs b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/BuilderTests.cs index 8c60b5a3..4f0b5e9e 100644 --- a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/BuilderTests.cs +++ b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/BuilderTests.cs @@ -71,7 +71,7 @@ public void Builder_WithExporterType_Sync_RegistersTracerProvider() // Provide required dependencies for exporter services.AddSingleton(_ => new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("unit-test-token"), + TokenResolver = (_, _, _) => Task.FromResult("unit-test-token"), UseS2SEndpoint = false }); @@ -106,7 +106,7 @@ public void Builder_WithExporterType_Async_RegistersTracerProvider() // Provide required dependencies for exporter services.AddSingleton(_ => new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("unit-test-token"), + TokenResolver = (_, _, _) => Task.FromResult("unit-test-token"), UseS2SEndpoint = false }); @@ -140,7 +140,7 @@ public void Builder_UseOpenTelemetryBuilder_False_WithExporterType_Sync_CreatesT services.AddSingleton(_ => new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("unit-test-token"), + TokenResolver = (_, _, _) => Task.FromResult("unit-test-token"), UseS2SEndpoint = false }); @@ -177,7 +177,7 @@ public void Builder_UseOpenTelemetryBuilder_False_WithExporterType_Async_Exports services.AddSingleton(_ => new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("unit-test-token"), + TokenResolver = (_, _, _) => Task.FromResult("unit-test-token"), UseS2SEndpoint = false }); diff --git a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/Tracing/Exporters/Agent365ExporterTests.cs b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/Tracing/Exporters/Agent365ExporterTests.cs index 27af125b..a5b6349d 100644 --- a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/Tracing/Exporters/Agent365ExporterTests.cs +++ b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/Tracing/Exporters/Agent365ExporterTests.cs @@ -99,7 +99,7 @@ private static Agent365Exporter CreateExporter(Func? to { var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("token") + TokenResolver = (_, _, _) => Task.FromResult("token") }; var resource = ResourceBuilder.CreateEmpty() @@ -118,7 +118,7 @@ public void Constructor_NullLogger_Throws() { var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("token") + TokenResolver = (_, _, _) => Task.FromResult("token") }; Action act = () => _ = new Agent365Exporter(Agent365ExporterTests._agent365ExporterCore, null!, options, resource: null); act.Should().Throw().WithParameterName("logger"); @@ -196,7 +196,7 @@ public void Agent365ExporterOptions_DefaultBatchingParameters_AreSet() // Arrange & Act var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("token") + TokenResolver = (_, _, _) => Task.FromResult("token") }; // Assert @@ -212,7 +212,7 @@ public void Agent365ExporterOptions_CustomBatchingParameters_CanBeSet() // Arrange & Act var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("token"), + TokenResolver = (_, _, _) => Task.FromResult("token"), MaxQueueSize = 4096, ScheduledDelayMilliseconds = 10000, ExporterTimeoutMilliseconds = 60000, @@ -232,7 +232,7 @@ public void Agent365ExporterOptions_UseS2SEndpoint_DefaultsToFalse() // Arrange & Act var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("token") + TokenResolver = (_, _, _) => Task.FromResult("token") }; // Assert @@ -245,7 +245,7 @@ public void Agent365ExporterOptions_UseS2SEndpoint_CanBeSetToTrue() // Arrange & Act var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("token"), + TokenResolver = (_, _, _) => Task.FromResult("token"), UseS2SEndpoint = true }; @@ -259,7 +259,7 @@ public void Agent365ExporterOptions_UseS2SEndpoint_CanBeSetToFalse() // Arrange & Act var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("token"), + TokenResolver = (_, _, _) => Task.FromResult("token"), UseS2SEndpoint = false }; @@ -275,7 +275,7 @@ public void UseS2SEndpoint_WhenFalse_UsesStandardEndpoint() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -297,7 +297,7 @@ public void UseS2SEndpoint_WhenTrue_UsesS2SEndpoint() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -328,7 +328,7 @@ public void UseS2SEndpoint_CanBeToggled_FromFalseToTrue() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -345,7 +345,7 @@ public void UseS2SEndpoint_CanBeToggled_FromTrueToFalse() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -362,7 +362,7 @@ public void Export_WithMultipleActivities_StandardEndpoint_GroupsByIdentity() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -396,7 +396,7 @@ public void Export_WithMultipleActivities_S2SEndpoint_GroupsByIdentity() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -433,7 +433,7 @@ public void Export_S2SEndpoint_TokenResolverCalled_WithCorrectParameters() var options = new Agent365ExporterOptions { - TokenResolver = (agentId, tenantId) => + TokenResolver = (agentId, tenantId, _) => { capturedAgentId = agentId; capturedTenantId = tenantId; @@ -472,7 +472,7 @@ public void Export_StandardEndpoint_TokenResolverCalled_WithCorrectParameters() var options = new Agent365ExporterOptions { - TokenResolver = (agentId, tenantId) => + TokenResolver = (agentId, tenantId, _) => { capturedAgentId = agentId; capturedTenantId = tenantId; @@ -508,7 +508,7 @@ public void Export_S2SEndpoint_NullToken_StillSendsRequest() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult(null), // Return null token + TokenResolver = (_, _, _) => Task.FromResult(null), // Return null token UseS2SEndpoint = true }; @@ -538,7 +538,7 @@ public void Export_S2SEndpoint_EmptyToken_StillSendsRequest() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult(string.Empty), // Return empty token + TokenResolver = (_, _, _) => Task.FromResult(string.Empty), // Return empty token UseS2SEndpoint = true }; @@ -568,7 +568,7 @@ public void Export_S2SEndpoint_TokenResolverThrows_ReturnsFailure() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromException(new InvalidOperationException("Token resolver failed")), + TokenResolver = (_, _, _) => Task.FromException(new InvalidOperationException("Token resolver failed")), UseS2SEndpoint = true }; @@ -598,7 +598,7 @@ public void Export_StandardEndpoint_TokenResolverThrows_ReturnsFailure() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromException(new InvalidOperationException("Token resolver failed")), + TokenResolver = (_, _, _) => Task.FromException(new InvalidOperationException("Token resolver failed")), UseS2SEndpoint = false }; @@ -628,7 +628,7 @@ public void Export_S2SEndpoint_ActivitiesWithoutIdentity_ReturnsSuccess() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -660,7 +660,7 @@ public void Export_StandardEndpoint_ActivitiesWithoutIdentity_ReturnsSuccess() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -692,7 +692,7 @@ public void Export_S2SEndpoint_ActivityWithOnlyTenantId_IsSkipped() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -722,7 +722,7 @@ public void Export_S2SEndpoint_ActivityWithOnlyAgentId_IsSkipped() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -752,7 +752,7 @@ public void Export_StandardEndpoint_ActivityWithOnlyTenantId_IsSkipped() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -782,7 +782,7 @@ public void Export_StandardEndpoint_ActivityWithOnlyAgentId_IsSkipped() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -812,7 +812,7 @@ public void Export_S2SEndpoint_MixedBatch_ProcessesOnlyValidActivities() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -846,7 +846,7 @@ public void Export_StandardEndpoint_MixedBatch_ProcessesOnlyValidActivities() // Arrange var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -889,7 +889,7 @@ public void Export_S2SEndpoint_WithCustomResource_ProcessesCorrectly() var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -924,7 +924,7 @@ public void Export_StandardEndpoint_WithCustomResource_ProcessesCorrectly() var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -954,7 +954,7 @@ public void Export_S2SEndpoint_WithDifferentClusterCategories_ProcessesCorrectly var options = new Agent365ExporterOptions { ClusterCategory = category, - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = true }; @@ -989,7 +989,7 @@ public void Export_StandardEndpoint_WithDifferentClusterCategories_ProcessesCorr var options = new Agent365ExporterOptions { ClusterCategory = category, - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; @@ -1036,7 +1036,7 @@ public void Export_RequestUri_EnvVar_Overrides_CustomDomainResolver_WhenBothSet( var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false, DomainResolver = tenantId => resolverDomain }; @@ -1086,7 +1086,7 @@ public void Export_RequestUri_UsesEnvVar_WhenNoResolverSet() var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false, }; @@ -1136,7 +1136,7 @@ public void Export_RequestUri_UsesCustomDomainResolver_WhenProvided() var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false, DomainResolver = tenantId => resolverDomain }; @@ -1184,7 +1184,7 @@ public void Export_RequestUri_UsesDefaultEndpoint_WhenNoResolverAndNoEnvVarSet() var options = new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), UseS2SEndpoint = false }; diff --git a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/Tracing/Exporters/ObservabilityBuilderExtensionsTests.cs b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/Tracing/Exporters/ObservabilityBuilderExtensionsTests.cs index 67ca662d..19c26a67 100644 --- a/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/Tracing/Exporters/ObservabilityBuilderExtensionsTests.cs +++ b/src/Tests/Microsoft.Agents.A365.Observability.Runtime.Tests/Tracing/Exporters/ObservabilityBuilderExtensionsTests.cs @@ -19,7 +19,7 @@ public void AddAgent365Exporter_WithCustomBatchingParameters_ShouldConfigureProc // Configure custom batching parameters services.AddSingleton(new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token"), + TokenResolver = (_, _, _) => Task.FromResult("test-token"), MaxQueueSize = 4096, ScheduledDelayMilliseconds = 10000, ExporterTimeoutMilliseconds = 60000, @@ -55,7 +55,7 @@ public void AddAgent365Exporter_WithDefaultBatchingParameters_ShouldUseDefaults( // Configure with default batching parameters services.AddSingleton(new Agent365ExporterOptions { - TokenResolver = (_, _) => Task.FromResult("test-token") + TokenResolver = (_, _, _) => Task.FromResult("test-token") }); var tracerProviderBuilder = services.AddOpenTelemetry().WithTracing(builder => 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 fcffb5c8..1b7ddddb 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.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/AddToolServersToAgent_Tests.cs b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Tests/Services/McpToolRegistrationServiceTests/AddToolServersToAgent_Tests.cs index 9dac3ae8..8646776f 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 @@ -65,7 +65,8 @@ await service.AddToolServersToAgent( TestAgentUserId, TestAuthToken, 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 874131c1..9f0e5c53 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 @@ -37,7 +37,8 @@ await service.GetMcpToolsAsync( TestAgentUserId, TestAuthToken, 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 80c7c975..9d84c58f 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 @@ -78,11 +78,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>())); @@ -102,11 +103,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()); 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(), + It.IsAny())); if (captureToolOptions != null) { - setup.Callback((_, _, _, options) => captureToolOptions(options)); + setup.Callback((_, _, _, options, _) => captureToolOptions(options)); } setup.ReturnsAsync((new List(), new Dictionary>())); @@ -114,7 +115,8 @@ private void SetupMocksForToolEnumeration(List servers, Diction It.IsAny(), It.IsAny(), It.IsAny(), - It.IsAny())) + It.IsAny(), + It.IsAny())) .ReturnsAsync((servers, toolsByServer)); } @@ -197,7 +199,8 @@ await _service.GetMcpToolDefinitionsAndResourcesAsync( TestAgentInstanceId, TestAuthToken, _mockTurnContext.Object, - It.Is(o => o.UserAgentConfiguration == Agent365AzureAIFoundrySdkUserAgentConfiguration.Instance)), + It.Is(o => o.UserAgentConfiguration == Agent365AzureAIFoundrySdkUserAgentConfiguration.Instance), + It.IsAny()), Times.Once); } diff --git a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.SemanticKernel.Tests/Services/McpToolRegistrationServiceTests.cs b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.SemanticKernel.Tests/Services/McpToolRegistrationServiceTests.cs index 5a8e29ea..610cbfcb 100644 --- a/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.SemanticKernel.Tests/Services/McpToolRegistrationServiceTests.cs +++ b/src/Tests/Microsoft.Agents.A365.Tooling.Extensions.SemanticKernel.Tests/Services/McpToolRegistrationServiceTests.cs @@ -111,11 +111,12 @@ private void SetupMocksForEmptyToolEnumeration(Action? captureToolO 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>())); @@ -168,7 +169,8 @@ await service.AddToolServersToAgentAsync( It.IsAny(), _testJwtToken, mockTurnContext.Object, - It.Is(o => o.UserAgentConfiguration == Agent365SemanticKernelSdkUserAgentConfiguration.Instance)), + It.Is(o => o.UserAgentConfiguration == Agent365SemanticKernelSdkUserAgentConfiguration.Instance), + It.IsAny()), Times.Once); } diff --git a/src/Tooling/Core/Services/IMcpToolServerConfigurationService.cs b/src/Tooling/Core/Services/IMcpToolServerConfigurationService.cs index fe5a4ff9..99f662e7 100644 --- a/src/Tooling/Core/Services/IMcpToolServerConfigurationService.cs +++ b/src/Tooling/Core/Services/IMcpToolServerConfigurationService.cs @@ -5,6 +5,7 @@ using Microsoft.Agents.A365.Tooling.Models; using Microsoft.Agents.Builder; using ModelContextProtocol.Client; +using System.Threading; namespace Microsoft.Agents.A365.Tooling.Services { @@ -21,6 +22,15 @@ public interface IMcpToolServerConfigurationService /// Returns the list of MCP Servers that are configured. Task> ListToolServersAsync(string agentInstanceId, string authToken); + /// + /// Gets the list of MCP Servers that are configured for the agent. + /// + /// Agent instance Id for the agent. + /// Auth token to access the MCP servers + /// A cancellation token to cancel the operation. + /// Returns the list of MCP Servers that are configured. + Task> ListToolServersAsync(string agentInstanceId, string authToken, CancellationToken cancellationToken); + /// /// Gets the list of MCP Servers that are configured for the agent. /// @@ -30,6 +40,16 @@ public interface IMcpToolServerConfigurationService /// Returns the list of MCP Servers that are configured. Task> ListToolServersAsync(string agentInstanceId, string authToken, ToolOptions toolOptions); + /// + /// Gets the list of MCP Servers that are configured for the agent. + /// + /// Agent instance Id for the agent. + /// Auth token to access the MCP servers + /// Tool options for listing servers. + /// A cancellation token to cancel the operation. + /// Returns the list of MCP Servers that are configured. + Task> ListToolServersAsync(string agentInstanceId, string authToken, ToolOptions toolOptions, CancellationToken cancellationToken); + /// /// Gets the MCP Client Tools from the specified MCP server. /// @@ -41,6 +61,18 @@ public interface IMcpToolServerConfigurationService /// Task> GetMcpClientToolsAsync(ITurnContext turnContext, MCPServerConfig mCPServerConfig, string authToken, ToolOptions toolOptions); + /// + /// Gets the MCP Client Tools from the specified MCP server. + /// + /// The turn context. + /// The MCP server configuration. + /// The authentication token. + /// Tool options for listing servers. + /// A cancellation token to cancel the operation. + /// MCP Client Tools + /// + Task> GetMcpClientToolsAsync(ITurnContext turnContext, MCPServerConfig mCPServerConfig, string authToken, ToolOptions toolOptions, CancellationToken cancellationToken); + /// /// Sends chat history to the MCP platform for real-time threat protection. /// @@ -84,6 +116,17 @@ public interface IMcpToolServerConfigurationService /// A tuple containing server configurations and a dictionary mapping server names to their available tools. Task<(List Servers, Dictionary> ToolsByServer)> EnumerateToolsFromServersAsync(string agentInstanceId, string authToken, ITurnContext turnContext, ToolOptions toolOptions); + /// + /// Enumerates all MCP tools from configured servers for a given agent. + /// + /// The agent instance ID. + /// Authentication token for MCP server access. + /// Turn context for the current request. + /// Tool options including user agent configuration. + /// A cancellation token to cancel the operation. + /// A tuple containing server configurations and a dictionary mapping server names to their available tools. + Task<(List Servers, Dictionary> ToolsByServer)> EnumerateToolsFromServersAsync(string agentInstanceId, string authToken, ITurnContext turnContext, ToolOptions toolOptions, CancellationToken cancellationToken); + /// /// Enumerates all MCP tools from configured servers, returning a flat list of all tools. /// @@ -93,5 +136,16 @@ public interface IMcpToolServerConfigurationService /// Tool options including user agent configuration. /// A flat list of all MCP tools from all configured servers. Task> EnumerateAllToolsAsync(string agentInstanceId, string authToken, ITurnContext turnContext, ToolOptions toolOptions); + + /// + /// Enumerates all MCP tools from configured servers, returning a flat list of all tools. + /// + /// The agent instance ID. + /// Authentication token for MCP server access. + /// Turn context for the current request. + /// Tool options including user agent configuration. + /// A cancellation token to cancel the operation. + /// A flat list of all MCP tools from all configured servers. + Task> EnumerateAllToolsAsync(string agentInstanceId, string authToken, ITurnContext turnContext, ToolOptions toolOptions, CancellationToken cancellationToken); } } diff --git a/src/Tooling/Core/Services/McpToolServerConfigurationService.ToolEnumeration.cs b/src/Tooling/Core/Services/McpToolServerConfigurationService.ToolEnumeration.cs index 57c94d95..deee4b7b 100644 --- a/src/Tooling/Core/Services/McpToolServerConfigurationService.ToolEnumeration.cs +++ b/src/Tooling/Core/Services/McpToolServerConfigurationService.ToolEnumeration.cs @@ -9,6 +9,7 @@ namespace Microsoft.Agents.A365.Tooling.Services using ModelContextProtocol.Client; using System; using System.Collections.Generic; + using System.Threading; using System.Threading.Tasks; /// @@ -17,11 +18,20 @@ namespace Microsoft.Agents.A365.Tooling.Services public partial class McpToolServerConfigurationService { /// - public virtual async Task<(List Servers, Dictionary> ToolsByServer)> EnumerateToolsFromServersAsync( + public virtual Task<(List Servers, Dictionary> ToolsByServer)> EnumerateToolsFromServersAsync( string agentInstanceId, string authToken, ITurnContext turnContext, ToolOptions toolOptions) + => EnumerateToolsFromServersAsync(agentInstanceId, authToken, turnContext, toolOptions, CancellationToken.None); + + /// + public virtual async Task<(List Servers, Dictionary> ToolsByServer)> EnumerateToolsFromServersAsync( + string agentInstanceId, + string authToken, + ITurnContext turnContext, + ToolOptions toolOptions, + CancellationToken cancellationToken) { var toolsByServer = new Dictionary>(StringComparer.OrdinalIgnoreCase); @@ -31,7 +41,8 @@ public partial class McpToolServerConfigurationService servers = await ListToolServersAsync( agentInstanceId, authToken, - toolOptions).ConfigureAwait(false); + toolOptions, + cancellationToken).ConfigureAwait(false); } catch (Exception ex) { @@ -68,7 +79,8 @@ public partial class McpToolServerConfigurationService turnContext, server, authToken, - toolOptions).ConfigureAwait(false); + toolOptions, + cancellationToken).ConfigureAwait(false); _logger.LogInformation( "Successfully loaded {ToolCount} tools from MCP server '{ServerName}'", @@ -107,17 +119,27 @@ public partial class McpToolServerConfigurationService } /// - public virtual async Task> EnumerateAllToolsAsync( + public virtual Task> EnumerateAllToolsAsync( string agentInstanceId, string authToken, ITurnContext turnContext, ToolOptions toolOptions) + => EnumerateAllToolsAsync(agentInstanceId, authToken, turnContext, toolOptions, CancellationToken.None); + + /// + public virtual async Task> EnumerateAllToolsAsync( + string agentInstanceId, + string authToken, + ITurnContext turnContext, + ToolOptions toolOptions, + CancellationToken cancellationToken) { var (_, toolsByServer) = await EnumerateToolsFromServersAsync( agentInstanceId, authToken, turnContext, - toolOptions).ConfigureAwait(false); + toolOptions, + cancellationToken).ConfigureAwait(false); var allTools = new List(); foreach (var tools in toolsByServer.Values) diff --git a/src/Tooling/Core/Services/McpToolServerConfigurationService.cs b/src/Tooling/Core/Services/McpToolServerConfigurationService.cs index db26decb..0ff7d7a3 100644 --- a/src/Tooling/Core/Services/McpToolServerConfigurationService.cs +++ b/src/Tooling/Core/Services/McpToolServerConfigurationService.cs @@ -11,6 +11,7 @@ namespace Microsoft.Agents.A365.Tooling.Services using System.Reflection; using System.Text; using System.Text.Json; + using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.A365.Runtime; using Microsoft.Agents.A365.Tooling.Handlers; @@ -49,23 +50,40 @@ public McpToolServerConfigurationService(ILogger - public virtual async Task> ListToolServersAsync(string agentInstanceId, string authToken) + public virtual Task> ListToolServersAsync(string agentInstanceId, string authToken) + => ListToolServersAsync(agentInstanceId, authToken, CancellationToken.None); + + /// + public virtual async Task> ListToolServersAsync(string agentInstanceId, string authToken, CancellationToken cancellationToken) { - return await ListToolServersAsync(agentInstanceId, authToken, new ToolOptions()); + return await ListToolServersAsync(agentInstanceId, authToken, new ToolOptions(), cancellationToken); } /// - public virtual async Task> ListToolServersAsync(string agentInstanceId, string authToken, ToolOptions toolOptions) + public virtual Task> ListToolServersAsync(string agentInstanceId, string authToken, ToolOptions toolOptions) + => ListToolServersAsync(agentInstanceId, authToken, toolOptions, CancellationToken.None); + + /// + public virtual async Task> ListToolServersAsync(string agentInstanceId, string authToken, ToolOptions toolOptions, CancellationToken cancellationToken) { - return IsDevScenario() ? GetMCPServersFromManifest() : await GetMCPServerFromToolingGatewayAsync(agentInstanceId, authToken, toolOptions); + return IsDevScenario() ? GetMCPServersFromManifest() : await GetMCPServerFromToolingGatewayAsync(agentInstanceId, authToken, toolOptions, cancellationToken); } /// - public virtual async Task> GetMcpClientToolsAsync( + public virtual Task> GetMcpClientToolsAsync( ITurnContext turnContext, MCPServerConfig mCPServerConfig, string authToken, ToolOptions toolOptions) + => GetMcpClientToolsAsync(turnContext, mCPServerConfig, authToken, toolOptions, CancellationToken.None); + + /// + public virtual async Task> GetMcpClientToolsAsync( + ITurnContext turnContext, + MCPServerConfig mCPServerConfig, + string authToken, + ToolOptions toolOptions, + CancellationToken cancellationToken) { try { @@ -78,8 +96,8 @@ public virtual async Task> GetMcpClientToolsAsync( this._logger.LogInformation($"Creating custom MCP client for: {mCPServerConfig.mcpServerName} at {mCPServerConfig.url}"); // Use custom HTTP-based implementation since MCP client library doesn't work - var mcpClient = await CreateMcpClientWithAuthHandlers(turnContext, new Uri(mCPServerConfig.url), authToken, toolOptions); - var tools = await mcpClient.ListToolsAsync(); + var mcpClient = await CreateMcpClientWithAuthHandlers(turnContext, new Uri(mCPServerConfig.url), authToken, toolOptions, cancellationToken); + var tools = await mcpClient.ListToolsAsync(cancellationToken: cancellationToken); this._logger.LogInformation($"Successfully retrieved {tools.Count} tools from {mCPServerConfig.mcpServerName}"); @@ -157,7 +175,7 @@ public async Task SendChatHistoryAsync(ITurnContext turnContext } private async Task> GetMCPServerFromToolingGatewayAsync( - string agentInstanceId, string authToken, ToolOptions toolOptions) + string agentInstanceId, string authToken, ToolOptions toolOptions, CancellationToken cancellationToken = default) { string configEndpoint = Utility.GetToolingGatewayForDigitalWorker(agentInstanceId, this._configuration); @@ -173,7 +191,7 @@ private async Task> GetMCPServerFromToolingGatewayAsync( httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", authToken); - var response = await httpClient.GetStringAsync(configEndpoint); + var response = await httpClient.GetStringAsync(configEndpoint, cancellationToken); var options = new JsonSerializerOptions { @@ -459,7 +477,7 @@ private List GetMCPServersFromManifest() /// /// Creates an MCP client with authentication handlers similar to your reference implementation /// - private async Task CreateMcpClientWithAuthHandlers(ITurnContext turnContext, Uri endpoint, string authToken, ToolOptions toolOptions) + private async Task CreateMcpClientWithAuthHandlers(ITurnContext turnContext, Uri endpoint, string authToken, ToolOptions toolOptions, CancellationToken cancellationToken = default) { // Create HTTP client handler chain for MCP service authentication var httpClientHandler = new HttpClientHandler(); @@ -507,7 +525,7 @@ private async Task CreateMcpClientWithAuthHandlers(ITurnContext turn try { - return await McpClientFactory.CreateAsync(clientTransport, loggerFactory: this._loggerFactory); + return await McpClientFactory.CreateAsync(clientTransport, loggerFactory: this._loggerFactory, cancellationToken: cancellationToken); } catch (Exception ex) { diff --git a/src/Tooling/Extensions/AgentFramework/Services/IMcpToolRegistrationService.cs b/src/Tooling/Extensions/AgentFramework/Services/IMcpToolRegistrationService.cs index aaf2f087..cb8c6d33 100644 --- a/src/Tooling/Extensions/AgentFramework/Services/IMcpToolRegistrationService.cs +++ b/src/Tooling/Extensions/AgentFramework/Services/IMcpToolRegistrationService.cs @@ -10,6 +10,7 @@ namespace Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Services; using Microsoft.Agents.Builder.App.UserAuth; using Microsoft.Extensions.AI; using System.Collections.Generic; +using System.Threading; using System.Threading.Tasks; /// @@ -43,6 +44,34 @@ Task AddToolServersToAgent( ITurnContext turnContext, string? authToken = null); + /// + /// Add new MCP servers to the agent by creating a new Agent instance. + /// + /// Note: Due to Microsoft.Extensions.AI framework limitations, MCP tools must be set during + /// Agent creation. If new tools are found, this method creates a new Agent + /// instance with all tools (existing + new) properly initialized. + /// + /// The configured IChatClient to use for creating the agent. + /// The agent instructions. + /// The existing tools to keep and add MCP tools to. + /// Agent User Id for the agent. + /// Turn context for the current request + /// User authorization information + /// Authentication Handler Name for use with the UserAuthorization System + /// Optional auth token to access the MCP servers. + /// A cancellation token to cancel the operation. + /// New Agent instance with all MCP tools, or agent with original tools if no new servers + Task AddToolServersToAgent( + IChatClient chatClient, + string agentInstructions, + IList initialTools, + string agentUserId, + UserAuthorization userAuthorization, + string authHandlerName, + ITurnContext turnContext, + string? authToken, + CancellationToken cancellationToken); + /// /// Returns a List of MCP tools to be added to the agent. /// @@ -59,6 +88,24 @@ Task> GetMcpToolsAsync( ITurnContext turnContext, string? authToken = null); + /// + /// Returns a List of MCP tools to be added to the agent. + /// + /// Agent User Id for the agent. + /// Turn context for the current request + /// User authorization information + /// Authentication Handler Name for use with the UserAuthorization System + /// Optional auth token to access the MCP servers. + /// A cancellation token to cancel the operation. + /// List of AI Tools be added to an agent. + Task> GetMcpToolsAsync( + string agentUserId, + UserAuthorization userAuthorization, + string authHandlerName, + ITurnContext turnContext, + string? authToken, + CancellationToken cancellationToken); + /// /// Sends chat history to the MCP platform. /// diff --git a/src/Tooling/Extensions/AgentFramework/Services/McpToolRegistrationService.cs b/src/Tooling/Extensions/AgentFramework/Services/McpToolRegistrationService.cs index 10927ab1..0390ad35 100644 --- a/src/Tooling/Extensions/AgentFramework/Services/McpToolRegistrationService.cs +++ b/src/Tooling/Extensions/AgentFramework/Services/McpToolRegistrationService.cs @@ -45,7 +45,7 @@ public McpToolRegistrationService( } /// - public async Task AddToolServersToAgent( + public Task AddToolServersToAgent( IChatClient chatClient, string agentInstructions, IList initialTools, @@ -54,6 +54,19 @@ public async Task AddToolServersToAgent( string authHandlerName, ITurnContext turnContext, string? authToken = null) + => AddToolServersToAgent(chatClient, agentInstructions, initialTools, agentUserId, userAuthorization, authHandlerName, turnContext, authToken, CancellationToken.None); + + /// + public async Task AddToolServersToAgent( + IChatClient chatClient, + string agentInstructions, + IList initialTools, + string agentUserId, + UserAuthorization userAuthorization, + string authHandlerName, + ITurnContext turnContext, + string? authToken, + CancellationToken cancellationToken) { if (chatClient == null) { @@ -83,7 +96,8 @@ public async Task AddToolServersToAgent( agentUserId, authToken, turnContext, - toolOptions).ConfigureAwait(false); + toolOptions, + cancellationToken).ConfigureAwait(false); // Add all MCP tools from all servers foreach (var serverEntry in toolsByServer) @@ -111,12 +125,22 @@ public async Task AddToolServersToAgent( } /// - public async Task> GetMcpToolsAsync( + public Task> GetMcpToolsAsync( string agentUserId, UserAuthorization userAuthorization, string authHandlerName, ITurnContext turnContext, string? authToken = null) + => GetMcpToolsAsync(agentUserId, userAuthorization, authHandlerName, turnContext, authToken, CancellationToken.None); + + /// + public async Task> GetMcpToolsAsync( + string agentUserId, + UserAuthorization userAuthorization, + string authHandlerName, + ITurnContext turnContext, + string? authToken, + CancellationToken cancellationToken) { try { @@ -135,7 +159,8 @@ public async Task> GetMcpToolsAsync( agentUserId, authToken, turnContext, - toolOptions).ConfigureAwait(false); + toolOptions, + cancellationToken).ConfigureAwait(false); // Convert to AITool list var a365ToolList = mcpTools.Cast().ToList(); diff --git a/src/Tooling/Extensions/AzureAIFoundry/Services/IMcpToolRegistrationService.cs b/src/Tooling/Extensions/AzureAIFoundry/Services/IMcpToolRegistrationService.cs index 2623c9a1..b26f77b6 100644 --- a/src/Tooling/Extensions/AzureAIFoundry/Services/IMcpToolRegistrationService.cs +++ b/src/Tooling/Extensions/AzureAIFoundry/Services/IMcpToolRegistrationService.cs @@ -33,6 +33,24 @@ Task AddToolServersToAgentAsync( ITurnContext turnContext, string? authToken = null); + /// + /// Loads/initializes configured MCP tool servers for the specified agent with full context. + /// This is the primary method that customers should use in orchestrators with full authentication context. + /// + /// The PersistentAgentsClient instance. + /// User authorization context. + /// Authentication Handler Name for use with the UserAuthorization System + /// Turn context for the conversation. + /// Optional auth token to access the MCP servers. + /// A cancellation token to cancel the operation. + Task AddToolServersToAgentAsync( + PersistentAgentsClient agentClient, + UserAuthorization userAuthorization, + string authHandlerName, + ITurnContext turnContext, + string? authToken, + CancellationToken cancellationToken); + /// /// Get MCP tool definitions and resources asynchronously. /// @@ -45,6 +63,20 @@ Task AddToolServersToAgentAsync( string authToken, ITurnContext turnContext); + /// + /// Get MCP tool definitions and resources asynchronously. + /// + /// Agent Instance Id for the agent. + /// Auth token to access the MCP servers. + /// Turn context for the conversation. + /// A cancellation token to cancel the operation. + /// A tuple containing the list of MCP tool definitions and tool resources. + Task<(IList ToolDefinitions, ToolResources? ToolResources)> GetMcpToolDefinitionsAndResourcesAsync( + string agentInstanceId, + string authToken, + ITurnContext turnContext, + CancellationToken cancellationToken); + /// /// Sends chat history to the MCP platform for real-time threat protection. /// Messages are provided directly by the caller as Azure AI Foundry messages. diff --git a/src/Tooling/Extensions/AzureAIFoundry/Services/McpToolRegistrationService.cs b/src/Tooling/Extensions/AzureAIFoundry/Services/McpToolRegistrationService.cs index 406e160d..d0bbacd3 100644 --- a/src/Tooling/Extensions/AzureAIFoundry/Services/McpToolRegistrationService.cs +++ b/src/Tooling/Extensions/AzureAIFoundry/Services/McpToolRegistrationService.cs @@ -92,12 +92,22 @@ public void AddToolServersToAgent( } /// - public async Task AddToolServersToAgentAsync( + public Task AddToolServersToAgentAsync( PersistentAgentsClient agentClient, UserAuthorization userAuthorization, string authHandlerName, ITurnContext turnContext, string? authToken = null) + => AddToolServersToAgentAsync(agentClient, userAuthorization, authHandlerName, turnContext, authToken, CancellationToken.None); + + /// + public async Task AddToolServersToAgentAsync( + PersistentAgentsClient agentClient, + UserAuthorization userAuthorization, + string authHandlerName, + ITurnContext turnContext, + string? authToken, + CancellationToken cancellationToken) { if (agentClient == null) { @@ -113,8 +123,7 @@ public async Task AddToolServersToAgentAsync( try { - // Perform the (potentially async) work in a dedicated task to keep this synchronous signature. - var (toolDefinitions, toolResources) = GetMcpToolDefinitionsAndResourcesAsync(agenticAppId, authToken ?? string.Empty, turnContext).GetAwaiter().GetResult(); + var (toolDefinitions, toolResources) = await GetMcpToolDefinitionsAndResourcesAsync(agenticAppId, authToken ?? string.Empty, turnContext, cancellationToken).ConfigureAwait(false); agentClient.Administration.UpdateAgent( agenticAppId, @@ -133,10 +142,18 @@ public async Task AddToolServersToAgentAsync( /// /// Get MCP tool definitions and resources. /// - public async Task<(IList ToolDefinitions, ToolResources? ToolResources)> GetMcpToolDefinitionsAndResourcesAsync( + public Task<(IList ToolDefinitions, ToolResources? ToolResources)> GetMcpToolDefinitionsAndResourcesAsync( string agentInstanceId, string authToken, ITurnContext turnContext) + => GetMcpToolDefinitionsAndResourcesAsync(agentInstanceId, authToken, turnContext, CancellationToken.None); + + /// + public async Task<(IList ToolDefinitions, ToolResources? ToolResources)> GetMcpToolDefinitionsAndResourcesAsync( + string agentInstanceId, + string authToken, + ITurnContext turnContext, + CancellationToken cancellationToken) { // TODO: Make this method private // Tool resources should ideally be accessible via agentClient after AddToolServersToAgent. @@ -153,7 +170,8 @@ public async Task AddToolServersToAgentAsync( agentInstanceId, authToken, turnContext, - toolOptions).ConfigureAwait(false); + toolOptions, + cancellationToken).ConfigureAwait(false); if (servers.Count == 0) { diff --git a/src/Tooling/Extensions/SemanticKernel/Services/IMcpToolRegistrationService.cs b/src/Tooling/Extensions/SemanticKernel/Services/IMcpToolRegistrationService.cs index d31aa8ac..0f969e45 100644 --- a/src/Tooling/Extensions/SemanticKernel/Services/IMcpToolRegistrationService.cs +++ b/src/Tooling/Extensions/SemanticKernel/Services/IMcpToolRegistrationService.cs @@ -25,10 +25,23 @@ public interface IMcpToolRegistrationService /// Authentication Handler Name for use with the UserAuthorization System /// /// Auth token to access the MCP servers - /// Returns a new object of the kernel + /// A task that completes when the MCP tool servers have been added to the provided kernel. /// Task AddToolServersToAgentAsync(Kernel kernel, UserAuthorization userAuthorization, string authHandlerName, ITurnContext turnContext, string? authToken = null); + /// + /// Adds the A365 MCP Tool Servers + /// + /// The kernel to which the tools will be added. + /// Agents SDK UserAuthorization System + /// Authentication Handler Name for use with the UserAuthorization System + /// + /// Auth token to access the MCP servers + /// A cancellation token to cancel the operation. + /// A task that completes when the MCP tool servers have been added to the provided kernel. + /// + Task AddToolServersToAgentAsync(Kernel kernel, UserAuthorization userAuthorization, string authHandlerName, ITurnContext turnContext, string? authToken, CancellationToken cancellationToken); + /// /// Sends chat history to the MCP platform for real-time threat protection. /// diff --git a/src/Tooling/Extensions/SemanticKernel/Services/McpToolRegistrationService.cs b/src/Tooling/Extensions/SemanticKernel/Services/McpToolRegistrationService.cs index c59575fb..9d91e243 100644 --- a/src/Tooling/Extensions/SemanticKernel/Services/McpToolRegistrationService.cs +++ b/src/Tooling/Extensions/SemanticKernel/Services/McpToolRegistrationService.cs @@ -48,7 +48,11 @@ public McpToolRegistrationService( } /// - public async Task AddToolServersToAgentAsync(Kernel kernel, UserAuthorization userAuthorization, string authHandlerName, ITurnContext turnContext, string? authToken = null) + public Task AddToolServersToAgentAsync(Kernel kernel, UserAuthorization userAuthorization, string authHandlerName, ITurnContext turnContext, string? authToken = null) + => AddToolServersToAgentAsync(kernel, userAuthorization, authHandlerName, turnContext, authToken, CancellationToken.None); + + /// + public async Task AddToolServersToAgentAsync(Kernel kernel, UserAuthorization userAuthorization, string authHandlerName, ITurnContext turnContext, string? authToken, CancellationToken cancellationToken) { if (kernel == null) { @@ -68,7 +72,7 @@ public async Task AddToolServersToAgentAsync(Kernel kernel, UserAuthorization us UserAgentConfiguration = Agent365SemanticKernelSdkUserAgentConfiguration.Instance }; - var (_, toolsByServer) = await _mcpServerConfigurationService.EnumerateToolsFromServersAsync(agenticAppId, authToken, turnContext, toolOptions).ConfigureAwait(false); + var (_, toolsByServer) = await _mcpServerConfigurationService.EnumerateToolsFromServersAsync(agenticAppId, authToken, turnContext, toolOptions, cancellationToken).ConfigureAwait(false); foreach (var serverEntry in toolsByServer) {