Skip to content

.NET: Graduate message injection out of experimental #7044

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
westey-m wants to merge 2 commits into
base: main
Choose a base branch
from
Open

.NET: Graduate message injection out of experimental #7044

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
70fbb5c
Remove experiemental flags for MessageInjection component
westey-m Jul 10, 2026
67abdad
Improve locking on message injection.
westey-m Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions dotnet/samples/02-agents/Agents/Agent_Step19_InFunctionLoopCheckpointing/Program.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<string> CheckTravelAdvisory([Description("The city name.")] string city)
{
// Simulated travel advisory data.
var advisory = city.ToUpperInvariant() switch
Expand All @@ -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<MessageInjectingChatClient>()?.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<MessageInjectingChatClient>();
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;
}
Expand Down
26 changes: 14 additions & 12 deletions dotnet/samples/02-agents/Harness/Harness_Shared_Console/HarnessAgentRunner.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,16 +111,15 @@ internal async Task OnUserInputAsync(string text)
/// enqueued via the <see cref="MessageInjectingChatClient"/> so it can be picked up
/// by the agent on its next opportunity.
/// </summary>
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));
}

/// <summary>
Expand Down Expand Up @@ -151,7 +150,9 @@ internal async Task StartAgentTurnAsync(IList<ChatMessage> messages)
private async Task RunAgentLoopAsync(IList<ChatMessage> messages)
{
IList<ChatMessage>? nextMessages = messages;
IReadOnlyList<ChatMessage> lastPendingMessages = this._messageInjector?.GetPendingMessages(this._session) ?? [];
IReadOnlyList<ChatMessage> lastPendingMessages = this._messageInjector is not null
? await this._messageInjector.GetPendingMessagesAsync(this._session).ConfigureAwait(false)
: [];

while (nextMessages is not null)
{
Expand Down Expand Up @@ -199,7 +200,7 @@ private async Task RunAgentLoopAsync(IList<ChatMessage> messages)
}
}

this.SyncQueuedMessageDisplay(ref lastPendingMessages);
lastPendingMessages = await this.SyncQueuedMessageDisplayAsync(lastPendingMessages).ConfigureAwait(false);
}
}
catch (Exception ex)
Expand All @@ -208,7 +209,7 @@ private async Task RunAgentLoopAsync(IList<ChatMessage> 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);
Expand Down Expand Up @@ -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.
/// </summary>
private void SyncQueuedMessageDisplay(ref IReadOnlyList<ChatMessage> lastPendingMessages)
/// <returns>The updated snapshot of pending messages.</returns>
private async Task<IReadOnlyList<ChatMessage>> SyncQueuedMessageDisplayAsync(IReadOnlyList<ChatMessage> 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++)
Expand All @@ -291,7 +293,7 @@ private void SyncQueuedMessageDisplay(ref IReadOnlyList<ChatMessage> lastPending
this._ux.WriteUserInputEcho(text);
}

lastPendingMessages = pending;
this._ux.SetQueuedMessages(pending);
return pending;
}
}
3 changes: 0 additions & 3 deletions dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -177,7 +175,6 @@ public sealed class ChatClientAgentOptions
/// <value>
/// Default is <see langword="false"/>.
/// </value>
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public bool EnableMessageInjection { get; set; }

/// <summary>
Expand Down
3 changes: 0 additions & 3 deletions dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -143,7 +141,6 @@ public static ChatClientBuilder UsePerServiceCallChatHistoryPersistence(this Cha
/// </remarks>
/// <param name="builder">The <see cref="ChatClientBuilder"/> to add the decorator to.</param>
/// <returns>The <paramref name="builder"/> for chaining.</returns>
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public static ChatClientBuilder UseMessageInjection(this ChatClientBuilder builder)
{
return builder.Use(innerClient => new MessageInjectingChatClient(innerClient));
Expand Down
118 changes: 79 additions & 39 deletions dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -41,14 +39,21 @@ namespace Microsoft.Agents.AI;
/// method is called.
/// </para>
/// </remarks>
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public sealed class MessageInjectingChatClient : DelegatingChatClient
{
/// <summary>
/// The key used to store the pending injected messages queue in the session's <see cref="AgentSessionStateBag"/>.
/// </summary>
internal const string PendingMessagesStateKey = "MessageInjectingChatClient.PendingInjectedMessages";

/// <summary>
/// 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 <see cref="ConditionalWeakTable{TKey, TValue}"/> so it
/// is released automatically when the session is collected.
/// </summary>
private readonly ConditionalWeakTable<AgentSession, SemaphoreSlim> _sessionLocks = new();

/// <summary>
/// Initializes a new instance of the <see cref="MessageInjectingChatClient"/> class.
/// </summary>
Expand All @@ -65,9 +70,8 @@ public override async Task<ChatResponse> GetResponseAsync(
CancellationToken cancellationToken = default)
{
var session = GetRequiredSession();
var queue = GetOrCreateQueue(session);

var newMessages = DrainInjectedMessages(queue, messages as IList<ChatMessage> ?? messages.ToList());
var newMessages = await this.DrainInjectedMessagesAsync(session, messages as IList<ChatMessage> ?? 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
Expand All @@ -86,13 +90,7 @@ public override async Task<ChatResponse> 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;
}
Expand All @@ -101,7 +99,7 @@ public override async Task<ChatResponse> GetResponseAsync(
// continue within the same conversation.
UpdateOptionsForNextIteration(ref options, response.ConversationId);

newMessages = DrainInjectedMessages(queue, Array.Empty<ChatMessage>());
newMessages = await this.DrainInjectedMessagesAsync(session, Array.Empty<ChatMessage>(), cancellationToken).ConfigureAwait(false);
}
}

Expand All @@ -112,9 +110,8 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
[EnumeratorCancellation] CancellationToken cancellationToken = default)
{
var session = GetRequiredSession();
var queue = GetOrCreateQueue(session);

var newMessages = DrainInjectedMessages(queue, messages as IList<ChatMessage> ?? messages.ToList());
var newMessages = await this.DrainInjectedMessagesAsync(session, messages as IList<ChatMessage> ?? 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
Expand Down Expand Up @@ -161,13 +158,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> 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;
}
Expand All @@ -176,7 +167,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
// continue within the same conversation.
UpdateOptionsForNextIteration(ref options, lastConversationId);

newMessages = DrainInjectedMessages(queue, Array.Empty<ChatMessage>());
newMessages = await this.DrainInjectedMessagesAsync(session, Array.Empty<ChatMessage>(), cancellationToken).ConfigureAwait(false);
}
}

Expand All @@ -190,20 +181,27 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
/// </remarks>
/// <param name="session">The agent session to enqueue messages for.</param>
/// <param name="messages">The messages to enqueue.</param>
public void EnqueueMessages(AgentSession session, IEnumerable<ChatMessage> messages)
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
/// <returns>A <see cref="Task"/> that represents the asynchronous operation.</returns>
public async Task EnqueueMessagesAsync(AgentSession session, IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)

@rogerbarreto rogerbarreto Jul 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1

{
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();
}
}

/// <summary>
Expand All @@ -215,30 +213,47 @@ public void EnqueueMessages(AgentSession session, IEnumerable<ChatMessage> messa
/// list is a point-in-time snapshot; messages may be consumed between calls.
/// </remarks>
/// <param name="session">The agent session to check.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
/// <returns>A read-only list of pending messages, or an empty list if none are pending.</returns>
public IReadOnlyList<ChatMessage> GetPendingMessages(AgentSession session)
public async Task<IReadOnlyList<ChatMessage>> GetPendingMessagesAsync(AgentSession session, CancellationToken cancellationToken = default)
{
Throw.IfNull(session);

if (!session.StateBag.TryGetValue<List<ChatMessage>>(PendingMessagesStateKey, out var queue) || queue is null)
SemaphoreSlim sessionLock = this.GetSessionLock(session);
await sessionLock.WaitAsync(cancellationToken).ConfigureAwait(false);
try
{
return Array.Empty<ChatMessage>();
}
if (!session.StateBag.TryGetValue<List<ChatMessage>>(PendingMessagesStateKey, out var queue) || queue is null || queue.Count == 0)
{
return Array.Empty<ChatMessage>();
}

lock (queue)
return queue.ToList();
}
finally
{
return queue.Count == 0 ? Array.Empty<ChatMessage>() : queue.ToList();
sessionLock.Release();
}
}

/// <summary>
/// Returns the per-session semaphore used to serialize access to the session's pending messages queue.
/// </summary>
private SemaphoreSlim GetSessionLock(AgentSession session)
=> this._sessionLocks.GetValue(session, static _ => new SemaphoreSlim(1, 1));

/// <summary>
/// Gets or creates the pending injected messages queue from the session's <see cref="AgentSessionStateBag"/>.
/// </summary>
/// <remarks>
/// Callers must hold the session lock (see <see cref="GetSessionLock"/>) while calling this method and
/// while operating on the returned queue, since the get-or-create is a non-atomic check-then-act.
/// </remarks>
private static List<ChatMessage> GetOrCreateQueue(AgentSession session)
{
if (session.StateBag.TryGetValue<List<ChatMessage>>(PendingMessagesStateKey, out var queue))
if (session.StateBag.TryGetValue<List<ChatMessage>>(PendingMessagesStateKey, out var queue) && queue is not null)
{
return queue!;
return queue;
}

var newQueue = new List<ChatMessage>();
Expand All @@ -263,13 +278,34 @@ private static AgentSession GetRequiredSession()
}

/// <summary>
/// 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 <see langword="true"/> if the session's pending messages queue is empty or has not been created.
/// </summary>
private async Task<bool> IsQueueEmptyAsync(AgentSession session, CancellationToken cancellationToken)
{
SemaphoreSlim sessionLock = this.GetSessionLock(session);
await sessionLock.WaitAsync(cancellationToken).ConfigureAwait(false);
try
{
return !session.StateBag.TryGetValue<List<ChatMessage>>(PendingMessagesStateKey, out var queue) || queue is null || queue.Count == 0;
}
finally
{
sessionLock.Release();
}
}

/// <summary>
/// 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 <paramref name="newMessages"/> list
/// is never modified.
/// </summary>
private static IList<ChatMessage> DrainInjectedMessages(List<ChatMessage> queue, IList<ChatMessage> newMessages)
private async Task<IList<ChatMessage>> DrainInjectedMessagesAsync(AgentSession session, IList<ChatMessage> 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;
Expand All @@ -280,6 +316,10 @@ private static IList<ChatMessage> DrainInjectedMessages(List<ChatMessage> queue,
queue.Clear();
return combined;
}
finally
{
sessionLock.Release();
}
}

/// <summary>
Expand Down
Loading
Loading