From 40a2dd5cd0978896313b0025f8da13240212c230 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+rogerbarreto@users.noreply.github.com> Date: Mon, 15 Jun 2026 14:21:53 +0100 Subject: [PATCH] .NET: Restore ambient client-header scope between non-streaming ClientHeadersAgent runs (#6517) * Restore ambient client-header scope between non-streaming runs (#6516) Make ClientHeadersAgent.RunCoreAsync async + await so the per-run ClientHeadersScope is unwound on return, matching the streaming path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Assert per-run on wire instead of brittle exact request count Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../ClientHeadersAgent.cs | 12 +-- .../ClientHeadersExtensionsTests.cs | 77 +++++++++++++++++++ 2 files changed, 84 insertions(+), 5 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry/ClientHeadersAgent.cs b/dotnet/src/Microsoft.Agents.AI.Foundry/ClientHeadersAgent.cs index 4366a84506..697d9120b8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry/ClientHeadersAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry/ClientHeadersAgent.cs @@ -35,7 +35,7 @@ internal sealed class ClientHeadersAgent : DelegatingAIAgent } /// - protected override Task RunCoreAsync( + protected override async Task RunCoreAsync( IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, @@ -44,13 +44,15 @@ internal sealed class ClientHeadersAgent : DelegatingAIAgent var snapshot = TrySnapshot(options); if (snapshot is not null) { - // AsyncLocal mutations made inside an awaited async method do not leak back to the - // caller after the method returns, so we do not need an explicit restore step here. - // See ClientHeadersScope remarks. + // This method is async, so the runtime restores the caller's ExecutionContext (and + // therefore the previous ClientHeadersScope.Current value) when the returned task + // completes. Awaiting the inner call is what establishes that async-method boundary, + // so the per-run scope set here cannot carry into a later run on the same async flow. + // See ClientHeadersScope remarks. The streaming path relies on the same behavior. ClientHeadersScope.Current = snapshot; } - return this.InnerAgent.RunAsync(messages, session, options, cancellationToken); + return await this.InnerAgent.RunAsync(messages, session, options, cancellationToken).ConfigureAwait(false); } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/ClientHeadersExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/ClientHeadersExtensionsTests.cs index cdc4ef343a..bdb6dec41f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/ClientHeadersExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/ClientHeadersExtensionsTests.cs @@ -4,6 +4,7 @@ using System; using System.ClientModel; using System.ClientModel.Primitives; using System.Collections.Generic; +using System.Linq; using System.Net; using System.Net.Http; using System.Reflection; @@ -561,6 +562,82 @@ public sealed class ClientHeadersExtensionsTests Assert.Equal(1, EntriesCount(policies!)); } + // ------------------------------------------------------------------------------------------- + // 21. Non-streaming hardening: a non-streaming run must restore the ambient ClientHeadersScope + // on return so a previous run's x-client-* headers do not carry into a later headerless run + // on the same async flow. (Streaming already restores naturally via its async iterator.) + // ------------------------------------------------------------------------------------------- + + [Fact] + public async Task NonStreaming_DoesNotCarryClientHeadersToSubsequentRunAsync() + { + // Arrange: a probe inner agent records ClientHeadersScope.Current observed at each run. + var observed = new List?>(); + var inner = new ProbeAgent(_ => + { + observed.Add(ClientHeadersScope.Current); + return Task.CompletedTask; + }); + var agent = new ClientHeadersAgent(inner); + + // Act: run 1 supplies a client header; run 2 supplies fresh, empty ChatOptions (no headers). + var run1Options = new ChatOptions(); + run1Options.WithClientHeader("x-client-end-user-id", "alice"); + await agent.RunAsync(messages: [], options: new ChatClientAgentRunOptions(run1Options)); + + // The scope must not carry back into the caller's flow after run 1 returns. + Assert.Null(ClientHeadersScope.Current); + + var run2Options = new ChatOptions(); + await agent.RunAsync(messages: [], options: new ChatClientAgentRunOptions(run2Options)); + + // Assert: run 1 observed "alice"; run 2 observed no headers (did not inherit run 1's value). + Assert.Equal(2, observed.Count); + Assert.NotNull(observed[0]); + Assert.Equal("alice", observed[0]!["x-client-end-user-id"]); + Assert.Null(observed[1]); + Assert.Null(ClientHeadersScope.Current); + } + + // ------------------------------------------------------------------------------------------- + // 22. End-to-end non-streaming: a second headerless run on the same async flow must not carry + // the first run's x-client-end-user-id onto the wire. + // ------------------------------------------------------------------------------------------- + + [Fact] + public async Task EndToEnd_NonStreaming_SecondRunDoesNotInheritHeaderOnWireAsync() + { + // Arrange: a real OpenAI ResponsesClient pointed at a recording handler. + using var handler = new RecordingHandler(MinimalResponseJson()); +#pragma warning disable CA5399 + using var http = new HttpClient(handler); +#pragma warning restore CA5399 + var openAIOptions = new OpenAIClientOptions { Transport = new HttpClientPipelineTransport(http) }; + var openAIClient = new OpenAIClient(new ApiKeyCredential("fake"), openAIOptions); + IChatClient chatClient = openAIClient.GetResponsesClient().AsIChatClient(); + + AIAgent agent = new ChatClientAgent(chatClient).AsBuilder().UseClientHeaders().Build(); + + // Act: run 1 carries x-client-end-user-id; run 2 supplies fresh options with no client headers. + var run1 = new ChatClientAgentRunOptions(new ChatOptions()); + run1.ChatOptions!.WithClientHeader("x-client-end-user-id", "alice"); + await agent.RunAsync("hi", options: run1); + var afterRun1 = handler.Requests.Count; + + var run2 = new ChatClientAgentRunOptions(new ChatOptions()); + await agent.RunAsync("hi", options: run2); + + // Assert: run 1 stamped the header; none of the requests issued by run 2 carry it. + // (Assert per-run rather than on an exact total count, which would be brittle to + // any extra/internal SDK requests.) + Assert.True(afterRun1 > 0); + Assert.Equal("alice", handler.Requests[0].Headers["x-client-end-user-id"]); + + var run2Requests = handler.Requests.Skip(afterRun1).ToList(); + Assert.NotEmpty(run2Requests); + Assert.All(run2Requests, r => Assert.False(r.Headers.ContainsKey("x-client-end-user-id"))); + } + // ------------------------------------------------------------------------------------------- // Helpers // -------------------------------------------------------------------------------------------