.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>

Author: rogerbarreto

dotnet/src/Microsoft.Agents.AI.Foundry/ClientHeadersAgent.cs

@@ -35,7 +35,7 @@ public ClientHeadersAgent(AIAgent innerAgent)
     }
 
     /// <inheritdoc/>
-    protected override Task<AgentResponse> RunCoreAsync(
+    protected override async Task<AgentResponse> RunCoreAsync(
         IEnumerable<ChatMessage> messages,
         AgentSession? session = null,
         AgentRunOptions? options = null,
@@ -44,13 +44,15 @@ protected override Task<AgentResponse> RunCoreAsync(
         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);
     }
 
     /// <inheritdoc/>

dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/ClientHeadersExtensionsTests.cs

@@ -4,6 +4,7 @@
 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 void UseClientHeaders_RepeatedRegistrations_OnSameChatClient_OnlyRegister
         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<IReadOnlyDictionary<string, string>?>();
+        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
     // -------------------------------------------------------------------------------------------