From 70fbb5cf96e059856db283dce27e33bce59db0d4 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 10 Jul 2026 09:45:17 +0000 Subject: [PATCH 1/2] Remove experiemental flags for MessageInjection component --- .../Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs | 3 --- .../ChatClient/ChatClientBuilderExtensions.cs | 3 --- .../ChatClient/MessageInjectingChatClient.cs | 3 --- 3 files changed, 9 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs index a9748d5cca7..ab90f735080 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs @@ -1,9 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System.Collections.Generic; -using System.Diagnostics.CodeAnalysis; using Microsoft.Extensions.AI; -using Microsoft.Shared.DiagnosticIds; namespace Microsoft.Agents.AI; @@ -177,7 +175,6 @@ public sealed class ChatClientAgentOptions /// /// Default is . /// - [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] public bool EnableMessageInjection { get; set; } /// diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs index ee3f4572e89..6da11e5eb6e 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs @@ -2,11 +2,9 @@ using System; using System.Collections.Generic; -using System.Diagnostics.CodeAnalysis; using Microsoft.Agents.AI; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Logging; -using Microsoft.Shared.DiagnosticIds; using Microsoft.Shared.Diagnostics; namespace Microsoft.Extensions.AI; @@ -143,7 +141,6 @@ public static ChatClientBuilder UsePerServiceCallChatHistoryPersistence(this Cha /// /// The to add the decorator to. /// The for chaining. - [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] public static ChatClientBuilder UseMessageInjection(this ChatClientBuilder builder) { return builder.Use(innerClient => new MessageInjectingChatClient(innerClient)); diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs index cbf5a18626d..d2c92c6d3d1 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs @@ -2,13 +2,11 @@ using System; using System.Collections.Generic; -using System.Diagnostics.CodeAnalysis; using System.Linq; using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; -using Microsoft.Shared.DiagnosticIds; using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -41,7 +39,6 @@ namespace Microsoft.Agents.AI; /// method is called. /// /// -[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] public sealed class MessageInjectingChatClient : DelegatingChatClient { /// From 67abdad2958a293942b51a82b15b80dd0bac51a7 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 10 Jul 2026 15:00:54 +0000 Subject: [PATCH 2/2] Improve locking on message injection. --- .../Program.cs | 12 +- .../HarnessAgentRunner.cs | 26 ++-- .../ChatClient/MessageInjectingChatClient.cs | 115 ++++++++++++------ .../MessageInjectingChatClientTests.cs | 93 ++++++++++---- 4 files changed, 171 insertions(+), 75 deletions(-) diff --git a/dotnet/samples/02-agents/Agents/Agent_Step19_InFunctionLoopCheckpointing/Program.cs b/dotnet/samples/02-agents/Agents/Agent_Step19_InFunctionLoopCheckpointing/Program.cs index b2b3f4b3270..e378ec70364 100644 --- a/dotnet/samples/02-agents/Agents/Agent_Step19_InFunctionLoopCheckpointing/Program.cs +++ b/dotnet/samples/02-agents/Agents/Agent_Step19_InFunctionLoopCheckpointing/Program.cs @@ -67,7 +67,7 @@ static string GetTime([Description("The city name.")] string city) => // asking for alternative destinations. The model will process this injected message on the next // service call — even though the parent FunctionInvokingChatClient loop would otherwise stop. [Description("Check current travel advisories for a city.")] -static string CheckTravelAdvisory([Description("The city name.")] string city) +static async Task CheckTravelAdvisory([Description("The city name.")] string city) { // Simulated travel advisory data. var advisory = city.ToUpperInvariant() switch @@ -85,9 +85,13 @@ static string CheckTravelAdvisory([Description("The city name.")] string city) // When an advisory is found, inject a follow-up question so the model automatically // suggests alternatives without the user needing to ask. var runContext = AIAgent.CurrentRunContext!; - runContext.Agent.GetService()?.EnqueueMessages( - runContext.Session!, - [new ChatMessage(ChatRole.User, $"Given the travel advisory for {city}, what alternative cities would you recommend instead?")]); + var injector = runContext.Agent.GetService(); + if (injector is not null) + { + await injector.EnqueueMessagesAsync( + runContext.Session!, + [new ChatMessage(ChatRole.User, $"Given the travel advisory for {city}, what alternative cities would you recommend instead?")]); + } return advisory; } diff --git a/dotnet/samples/02-agents/Harness/Harness_Shared_Console/HarnessAgentRunner.cs b/dotnet/samples/02-agents/Harness/Harness_Shared_Console/HarnessAgentRunner.cs index ced5444c5ad..142de2d481e 100644 --- a/dotnet/samples/02-agents/Harness/Harness_Shared_Console/HarnessAgentRunner.cs +++ b/dotnet/samples/02-agents/Harness/Harness_Shared_Console/HarnessAgentRunner.cs @@ -111,16 +111,15 @@ internal async Task OnUserInputAsync(string text) /// enqueued via the so it can be picked up /// by the agent on its next opportunity. /// - internal Task OnStreamingInputAsync(string text) + internal async Task OnStreamingInputAsync(string text) { if (this._messageInjector is null) { - return Task.CompletedTask; + return; } - this._messageInjector.EnqueueMessages(this._session, [new ChatMessage(ChatRole.User, text)]); - this._ux.SetQueuedMessages(this._messageInjector.GetPendingMessages(this._session)); - return Task.CompletedTask; + await this._messageInjector.EnqueueMessagesAsync(this._session, [new ChatMessage(ChatRole.User, text)]).ConfigureAwait(false); + this._ux.SetQueuedMessages(await this._messageInjector.GetPendingMessagesAsync(this._session).ConfigureAwait(false)); } /// @@ -151,7 +150,9 @@ internal async Task StartAgentTurnAsync(IList messages) private async Task RunAgentLoopAsync(IList messages) { IList? nextMessages = messages; - IReadOnlyList lastPendingMessages = this._messageInjector?.GetPendingMessages(this._session) ?? []; + IReadOnlyList lastPendingMessages = this._messageInjector is not null + ? await this._messageInjector.GetPendingMessagesAsync(this._session).ConfigureAwait(false) + : []; while (nextMessages is not null) { @@ -199,7 +200,7 @@ private async Task RunAgentLoopAsync(IList messages) } } - this.SyncQueuedMessageDisplay(ref lastPendingMessages); + lastPendingMessages = await this.SyncQueuedMessageDisplayAsync(lastPendingMessages).ConfigureAwait(false); } } catch (Exception ex) @@ -208,7 +209,7 @@ private async Task RunAgentLoopAsync(IList messages) } // Final sync after streaming. - this.SyncQueuedMessageDisplay(ref lastPendingMessages); + lastPendingMessages = await this.SyncQueuedMessageDisplayAsync(lastPendingMessages).ConfigureAwait(false); this._ux.StopSpinner(); await this._ux.EndStreamingOutputAsync().ConfigureAwait(false); @@ -275,14 +276,15 @@ private void CompleteTurn() /// Messages that have been consumed (drained by the service) are echoed to the output /// area as regular user-input entries. /// - private void SyncQueuedMessageDisplay(ref IReadOnlyList lastPendingMessages) + /// The updated snapshot of pending messages. + private async Task> SyncQueuedMessageDisplayAsync(IReadOnlyList lastPendingMessages) { if (this._messageInjector is null) { - return; + return lastPendingMessages; } - var pending = this._messageInjector.GetPendingMessages(this._session); + var pending = await this._messageInjector.GetPendingMessagesAsync(this._session).ConfigureAwait(false); int consumedCount = lastPendingMessages.Count - pending.Count; for (int i = 0; i < consumedCount && i < lastPendingMessages.Count; i++) @@ -291,7 +293,7 @@ private void SyncQueuedMessageDisplay(ref IReadOnlyList lastPending this._ux.WriteUserInputEcho(text); } - lastPendingMessages = pending; this._ux.SetQueuedMessages(pending); + return pending; } } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs index d2c92c6d3d1..5251b301447 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs @@ -46,6 +46,14 @@ public sealed class MessageInjectingChatClient : DelegatingChatClient /// internal const string PendingMessagesStateKey = "MessageInjectingChatClient.PendingInjectedMessages"; + /// + /// Per-session semaphore used to serialize all access to a session's pending messages queue, + /// including its creation. A single client instance is shared across sessions, so the lock is + /// keyed on the session and stored in a so it + /// is released automatically when the session is collected. + /// + private readonly ConditionalWeakTable _sessionLocks = new(); + /// /// Initializes a new instance of the class. /// @@ -62,9 +70,8 @@ public override async Task GetResponseAsync( CancellationToken cancellationToken = default) { var session = GetRequiredSession(); - var queue = GetOrCreateQueue(session); - var newMessages = DrainInjectedMessages(queue, messages as IList ?? messages.ToList()); + var newMessages = await this.DrainInjectedMessagesAsync(session, messages as IList ?? messages.ToList(), cancellationToken).ConfigureAwait(false); // Loop to process injected messages: after each service call, if no actionable function calls // are pending but new messages have been injected into the queue, we call the service again @@ -83,13 +90,7 @@ public override async Task GetResponseAsync( // No actionable function calls. If there are pending injected messages, loop again // to send them to the service. Otherwise, we're done. - bool queueEmpty; - lock (queue) - { - queueEmpty = queue.Count == 0; - } - - if (queueEmpty) + if (await this.IsQueueEmptyAsync(session, cancellationToken).ConfigureAwait(false)) { return response; } @@ -98,7 +99,7 @@ public override async Task GetResponseAsync( // continue within the same conversation. UpdateOptionsForNextIteration(ref options, response.ConversationId); - newMessages = DrainInjectedMessages(queue, Array.Empty()); + newMessages = await this.DrainInjectedMessagesAsync(session, Array.Empty(), cancellationToken).ConfigureAwait(false); } } @@ -109,9 +110,8 @@ public override async IAsyncEnumerable GetStreamingResponseA [EnumeratorCancellation] CancellationToken cancellationToken = default) { var session = GetRequiredSession(); - var queue = GetOrCreateQueue(session); - var newMessages = DrainInjectedMessages(queue, messages as IList ?? messages.ToList()); + var newMessages = await this.DrainInjectedMessagesAsync(session, messages as IList ?? messages.ToList(), cancellationToken).ConfigureAwait(false); // Loop to process injected messages: after each service call, if no actionable function calls // are pending but new messages have been injected into the queue, we call the service again @@ -158,13 +158,7 @@ public override async IAsyncEnumerable GetStreamingResponseA // No actionable function calls. If there are pending injected messages, loop again // to send them to the service. Otherwise, we're done. - bool queueEmpty; - lock (queue) - { - queueEmpty = queue.Count == 0; - } - - if (queueEmpty) + if (await this.IsQueueEmptyAsync(session, cancellationToken).ConfigureAwait(false)) { yield break; } @@ -173,7 +167,7 @@ public override async IAsyncEnumerable GetStreamingResponseA // continue within the same conversation. UpdateOptionsForNextIteration(ref options, lastConversationId); - newMessages = DrainInjectedMessages(queue, Array.Empty()); + newMessages = await this.DrainInjectedMessagesAsync(session, Array.Empty(), cancellationToken).ConfigureAwait(false); } } @@ -187,20 +181,27 @@ public override async IAsyncEnumerable GetStreamingResponseA /// /// The agent session to enqueue messages for. /// The messages to enqueue. - public void EnqueueMessages(AgentSession session, IEnumerable messages) + /// The to monitor for cancellation requests. + /// A that represents the asynchronous operation. + public async Task EnqueueMessagesAsync(AgentSession session, IEnumerable messages, CancellationToken cancellationToken = default) { Throw.IfNull(session); Throw.IfNull(messages); - var queue = GetOrCreateQueue(session); - - lock (queue) + SemaphoreSlim sessionLock = this.GetSessionLock(session); + await sessionLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try { + var queue = GetOrCreateQueue(session); foreach (var message in messages) { queue.Add(message); } } + finally + { + sessionLock.Release(); + } } /// @@ -212,30 +213,47 @@ public void EnqueueMessages(AgentSession session, IEnumerable messa /// list is a point-in-time snapshot; messages may be consumed between calls. /// /// The agent session to check. + /// The to monitor for cancellation requests. /// A read-only list of pending messages, or an empty list if none are pending. - public IReadOnlyList GetPendingMessages(AgentSession session) + public async Task> GetPendingMessagesAsync(AgentSession session, CancellationToken cancellationToken = default) { Throw.IfNull(session); - if (!session.StateBag.TryGetValue>(PendingMessagesStateKey, out var queue) || queue is null) + SemaphoreSlim sessionLock = this.GetSessionLock(session); + await sessionLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try { - return Array.Empty(); - } + if (!session.StateBag.TryGetValue>(PendingMessagesStateKey, out var queue) || queue is null || queue.Count == 0) + { + return Array.Empty(); + } - lock (queue) + return queue.ToList(); + } + finally { - return queue.Count == 0 ? Array.Empty() : queue.ToList(); + sessionLock.Release(); } } + /// + /// Returns the per-session semaphore used to serialize access to the session's pending messages queue. + /// + private SemaphoreSlim GetSessionLock(AgentSession session) + => this._sessionLocks.GetValue(session, static _ => new SemaphoreSlim(1, 1)); + /// /// Gets or creates the pending injected messages queue from the session's . /// + /// + /// Callers must hold the session lock (see ) while calling this method and + /// while operating on the returned queue, since the get-or-create is a non-atomic check-then-act. + /// private static List GetOrCreateQueue(AgentSession session) { - if (session.StateBag.TryGetValue>(PendingMessagesStateKey, out var queue)) + if (session.StateBag.TryGetValue>(PendingMessagesStateKey, out var queue) && queue is not null) { - return queue!; + return queue; } var newQueue = new List(); @@ -260,13 +278,34 @@ private static AgentSession GetRequiredSession() } /// - /// Drains all pending injected messages from the queue and returns a new list combining - /// the original messages with the drained messages. The original list is never modified. + /// Returns if the session's pending messages queue is empty or has not been created. + /// + private async Task IsQueueEmptyAsync(AgentSession session, CancellationToken cancellationToken) + { + SemaphoreSlim sessionLock = this.GetSessionLock(session); + await sessionLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + return !session.StateBag.TryGetValue>(PendingMessagesStateKey, out var queue) || queue is null || queue.Count == 0; + } + finally + { + sessionLock.Release(); + } + } + + /// + /// Drains all pending injected messages from the session's queue and returns a new list combining + /// the original messages with the drained messages. The original list + /// is never modified. /// - private static IList DrainInjectedMessages(List queue, IList newMessages) + private async Task> DrainInjectedMessagesAsync(AgentSession session, IList newMessages, CancellationToken cancellationToken) { - lock (queue) + SemaphoreSlim sessionLock = this.GetSessionLock(session); + await sessionLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try { + var queue = GetOrCreateQueue(session); if (queue.Count == 0) { return newMessages; @@ -277,6 +316,10 @@ private static IList DrainInjectedMessages(List queue, queue.Clear(); return combined; } + finally + { + sessionLock.Release(); + } } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs index 93f6b02b030..4bd69b089d2 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs @@ -175,17 +175,17 @@ public async Task RunAsync_LoopsInternally_WhenNoActionableFCCButPendingMessages It.IsAny>(), It.IsAny(), It.IsAny())) - .Returns((IEnumerable msgs, ChatOptions? _, CancellationToken _) => + .Returns(async (IEnumerable msgs, ChatOptions? _, CancellationToken ct) => { serviceCallCount++; if (serviceCallCount == 1) { // First call — simulate that something enqueues a message (e.g., a provider or background task) - injectorRef!.EnqueueMessages(sessionRef!, [new ChatMessage(ChatRole.User, "injected during first call")]); + await injectorRef!.EnqueueMessagesAsync(sessionRef!, [new ChatMessage(ChatRole.User, "injected during first call")], ct); } // Return a plain text response (no FunctionCallContent) to trigger the internal loop - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, $"response {serviceCallCount}")])); + return new ChatResponse([new(ChatRole.Assistant, $"response {serviceCallCount}")]); }); Mock mockChatHistoryProvider = new(null, null, null); @@ -236,20 +236,20 @@ public async Task RunAsync_DoesNotLoopInternally_WhenActionableFCCPresentAsync() It.IsAny>(), It.IsAny(), It.IsAny())) - .Returns((IEnumerable msgs, ChatOptions? _, CancellationToken _) => + .Returns(async (IEnumerable msgs, ChatOptions? _, CancellationToken ct) => { serviceCallCount++; if (serviceCallCount == 1) { // Enqueue a message during the first call - injectorRef!.EnqueueMessages(sessionRef!, [new ChatMessage(ChatRole.User, "injected")]); + await injectorRef!.EnqueueMessagesAsync(sessionRef!, [new ChatMessage(ChatRole.User, "injected")], ct); // Return a response with an actionable FunctionCallContent - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, - [new FunctionCallContent("call1", "myTool", new Dictionary())])])); + return new ChatResponse([new(ChatRole.Assistant, + [new FunctionCallContent("call1", "myTool", new Dictionary())])]); } // Subsequent calls return plain text (the FCC loop will call back after tool execution) - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, "final")])); + return new ChatResponse([new(ChatRole.Assistant, "final")]); }); Mock mockChatHistoryProvider = new(null, null, null); @@ -306,19 +306,19 @@ public async Task RunAsync_LoopsInternally_WhenOnlyInformationalOnlyFCCAndPendin It.IsAny>(), It.IsAny(), It.IsAny())) - .Returns((IEnumerable msgs, ChatOptions? _, CancellationToken _) => + .Returns(async (IEnumerable msgs, ChatOptions? _, CancellationToken ct) => { serviceCallCount++; if (serviceCallCount == 1) { // Enqueue a message during the first call - injectorRef!.EnqueueMessages(sessionRef!, [new ChatMessage(ChatRole.User, "injected")]); + await injectorRef!.EnqueueMessagesAsync(sessionRef!, [new ChatMessage(ChatRole.User, "injected")], ct); // Return a response with InformationalOnly FCC (not actionable) - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, - [new FunctionCallContent("call1", "myTool", new Dictionary()) { InformationalOnly = true }])])); + return new ChatResponse([new(ChatRole.Assistant, + [new FunctionCallContent("call1", "myTool", new Dictionary()) { InformationalOnly = true }])]); } - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, "final")])); + return new ChatResponse([new(ChatRole.Assistant, "final")]); }); Mock mockChatHistoryProvider = new(null, null, null); @@ -370,7 +370,7 @@ public async Task RunAsync_PropagatesConversationId_AcrossInternalLoopIterations It.IsAny>(), It.IsAny(), It.IsAny())) - .Returns((IEnumerable _, ChatOptions? opts, CancellationToken _) => + .Returns(async (IEnumerable _, ChatOptions? opts, CancellationToken ct) => { serviceCallCount++; capturedConversationIds.Add(opts?.ConversationId); @@ -378,15 +378,15 @@ public async Task RunAsync_PropagatesConversationId_AcrossInternalLoopIterations if (serviceCallCount == 1) { // First call: inject a message and return a ConversationId - injectorRef!.EnqueueMessages(sessionRef!, [new ChatMessage(ChatRole.User, "injected")]); - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, "first response")]) + await injectorRef!.EnqueueMessagesAsync(sessionRef!, [new ChatMessage(ChatRole.User, "injected")], ct); + return new ChatResponse([new(ChatRole.Assistant, "first response")]) { ConversationId = "conv-123", - }); + }; } // Second call (from loop): should have the propagated ConversationId - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, "second response")])); + return new ChatResponse([new(ChatRole.Assistant, "second response")]); }); ChatClientAgent agent = new(mockService.Object, options: new() @@ -427,7 +427,7 @@ public async Task RunAsync_DeliversInjectedMessages_AfterSessionSerializationRou It.IsAny>(), It.IsAny(), It.IsAny())) - .Returns((IEnumerable msgs, ChatOptions? _, CancellationToken _) => + .Returns(async (IEnumerable msgs, ChatOptions? _, CancellationToken ct) => { if (runCount == 1) { @@ -435,16 +435,16 @@ public async Task RunAsync_DeliversInjectedMessages_AfterSessionSerializationRou // Inject a message during the first run — this will remain pending (not drained) // because we return an actionable FCC that causes the parent loop to take over. - injectorRef!.EnqueueMessages(sessionRef!, [new ChatMessage(ChatRole.User, "injected before serialization")]); + await injectorRef!.EnqueueMessagesAsync(sessionRef!, [new ChatMessage(ChatRole.User, "injected before serialization")], ct); // Return actionable FCC so the injection loop does NOT drain the message - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, - [new FunctionCallContent("call1", "myTool", new Dictionary())])])); + return new ChatResponse([new(ChatRole.Assistant, + [new FunctionCallContent("call1", "myTool", new Dictionary())])]); } // Second run (after deserialization) — capture what messages come through capturedMessagesSecondRun.AddRange(msgs); - return Task.FromResult(new ChatResponse([new(ChatRole.Assistant, "final response")])); + return new ChatResponse([new(ChatRole.Assistant, "final response")]); }); Mock mockChatHistoryProvider = new(null, null, null); @@ -490,4 +490,51 @@ public async Task RunAsync_DeliversInjectedMessages_AfterSessionSerializationRou Assert.Contains(capturedMessagesSecondRun, m => m.Text == "injected before serialization"); Assert.Contains(capturedMessagesSecondRun, m => m.Text == "second run message"); } + + /// + /// Verifies that concurrent calls on the same + /// session do not lose messages. This guards against the queue-creation race where two callers could + /// each create a separate backing queue and one overwrites the other. + /// + [Fact] + public async Task EnqueueMessages_ConcurrentEnqueues_DoesNotLoseMessagesAsync() + { + // Arrange + Mock mockService = new(); + ChatClientAgent agent = new(mockService.Object, options: new() + { + EnableMessageInjection = true, + }); + + var injector = agent.ChatClient.GetService(); + Assert.NotNull(injector); + + var session = await agent.CreateSessionAsync(); + + const int ThreadCount = 16; + const int MessagesPerThread = 100; + + // Release all threads at once to maximize the chance of hitting the queue-creation race. + using var barrier = new Barrier(ThreadCount); + var tasks = new List(ThreadCount); + for (int t = 0; t < ThreadCount; t++) + { + int threadIndex = t; + tasks.Add(Task.Run(async () => + { + barrier.SignalAndWait(); + for (int i = 0; i < MessagesPerThread; i++) + { + await injector!.EnqueueMessagesAsync(session, [new ChatMessage(ChatRole.User, $"{threadIndex}-{i}")]); + } + })); + } + + // Act + await Task.WhenAll(tasks); + + // Assert — every enqueued message must be present (none lost to a creation race). + IReadOnlyList pending = await injector!.GetPendingMessagesAsync(session); + Assert.Equal(ThreadCount * MessagesPerThread, pending.Count); + } }