From 8338e650e4767aa3d3ee3a75aab4b9231872a06f Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Thu, 6 Aug 2026 02:20:40 -0700 Subject: [PATCH 1/2] Add a routing context factory for request-scoped policy state Add a protected virtual CreateContext to FailoverChatClient so a derived class can return its own RoutingContext subclass. One context is already created per request and supplied to every selection and routing update, so state stored on it is scoped to the request and released with it. Previously a policy that needed state across attempts had to keep a side table keyed by the context and remove the entry on the terminal update. That state outlives the request whenever routing ends without a terminal update, such as when selection throws after a nonterminal update or when a streaming enumerator is abandoned without being disposed. Move OrderedFailoverChatClient to the new pattern. Its next-client index is now a field on its own context, which removes the ConcurrentDictionary, the lookups on every selection and update, and the explicit cleanup on termination. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 25 +++- .../Microsoft.Extensions.AI.Abstractions.json | 4 + .../ChatRouting/FailoverChatClient.cs | 12 +- .../ChatRouting/OrderedFailoverChatClient.cs | 44 ++++--- .../Microsoft.Extensions.AI.json | 4 + .../ChatRouting/RoutingChatClientTests.cs | 98 ++++++++++++++++ .../ChatRouting/FailoverChatClientTests.cs | 108 ++++++++++++++++++ .../OrderedFailoverChatClientTests.cs | 37 +++--- 8 files changed, 293 insertions(+), 39 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index 58a1f050a2f..f4615706c0b 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -39,6 +39,25 @@ public static RoutingChatClient Create( return new CallbackRoutingChatClient(clientSelector); } + /// Creates the context supplied to client selection for one request. + /// The messages to route. + /// The options supplied by the caller, or . + /// The context for the request. + /// + /// + /// The default implementation returns a new . A derived class can return its own + /// subclass to supply additional request inputs or to carry request-scoped policy + /// state, which then reads by casting the supplied context. State stored on the + /// context is released with the request. + /// + /// + /// This method is called once per request by and + /// . Exceptions from this method propagate to the caller. + /// + /// + protected virtual RoutingContext CreateContext(IEnumerable messages, ChatOptions? options) => + new(messages, options); + /// Selects the client to invoke for the request. /// The request-specific inputs. /// The cancellation token supplied for the request. @@ -57,7 +76,8 @@ public virtual async Task GetResponseAsync( { _ = Throw.IfNull(messages); - var context = new RoutingContext(messages, options); + RoutingContext context = CreateContext(messages, options) ?? + throw new InvalidOperationException($"{nameof(CreateContext)} returned null."); IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); @@ -74,7 +94,8 @@ public virtual async IAsyncEnumerable GetStreamingResponseAs { _ = Throw.IfNull(messages); - var context = new RoutingContext(messages, options); + RoutingContext context = CreateContext(messages, options) ?? + throw new InvalidOperationException($"{nameof(CreateContext)} returned null."); IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json b/src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json index 8b5f337da10..4fd0fa2ce5b 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json @@ -3978,6 +3978,10 @@ "Member": "static Microsoft.Extensions.AI.RoutingChatClient Microsoft.Extensions.AI.RoutingChatClient.Create(System.Func> clientSelector);", "Stage": "Experimental" }, + { + "Member": "virtual Microsoft.Extensions.AI.RoutingContext Microsoft.Extensions.AI.RoutingChatClient.CreateContext(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? options);", + "Stage": "Experimental" + }, { "Member": "void Microsoft.Extensions.AI.RoutingChatClient.Dispose();", "Stage": "Experimental" diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index d825a8a054d..5fabb1270bc 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -29,6 +29,12 @@ namespace Microsoft.Extensions.AI; /// client selection, policy state, selection-failure cleanup, and the lifetime of clients they retain. /// /// +/// One is created for each request and is supplied to every selection and routing +/// update for that request. A derived class that keeps request-scoped policy state can override +/// to return its own subclass and store +/// that state on the context, which releases it with the request instead of requiring explicit cleanup. +/// +/// /// Once streaming enumeration begins, callers must dispose the enumerator. Abandoning an active enumerator without /// disposing it prevents both inner enumerator disposal and the terminal routing update. /// @@ -99,7 +105,8 @@ public sealed override async Task GetResponseAsync( { _ = Throw.IfNull(messages); - var context = new RoutingContext(messages, options); + RoutingContext context = CreateContext(messages, options) ?? + throw new InvalidOperationException($"{nameof(CreateContext)} returned null."); int? maximumAttempts = MaximumAttemptsPerRequest; int attemptCount = 0; @@ -169,7 +176,8 @@ public sealed override async IAsyncEnumerable GetStreamingRe { _ = Throw.IfNull(messages); - var context = new RoutingContext(messages, options); + RoutingContext context = CreateContext(messages, options) ?? + throw new InvalidOperationException($"{nameof(CreateContext)} returned null."); int? maximumAttempts = MaximumAttemptsPerRequest; int attemptCount = 0; diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs index 925c0e9544e..639491c0dcc 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -2,7 +2,6 @@ // The .NET Foundation licenses this file to you under the MIT license. using System; -using System.Collections.Concurrent; using System.Collections.Generic; using System.Diagnostics; using System.Diagnostics.CodeAnalysis; @@ -30,10 +29,6 @@ public sealed class OrderedFailoverChatClient : FailoverChatClient { private readonly bool _leaveOpen; private readonly IChatClient[] _clients; - - // Holds the next client index for a request that has a failed attempt. A nonterminal update is always followed - // by another selection, so a stored index is always in range. - private readonly ConcurrentDictionary _requestStates = new(); private bool _disposed; /// Initializes a new instance of the class. @@ -66,17 +61,18 @@ public OrderedFailoverChatClient(IReadOnlyList clients, bool leaveO _clients = clientsSnapshot; } + /// + protected override RoutingContext CreateContext(IEnumerable messages, ChatOptions? options) => + new OrderedRoutingContext(messages, options); + /// protected override ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) { - _ = Throw.IfNull(context); _ = cancellationToken; - int clientIndex = _requestStates.TryGetValue(context, out int nextClientIndex) ? nextClientIndex : 0; - - return new(_clients[clientIndex]); + return new(_clients[GetState(context).NextClientIndex]); } /// @@ -90,22 +86,20 @@ protected override ValueTask OnRoutingUpdateAsync( if (isTerminal) { - _ = _requestStates.TryRemove(context, out _); return default; } Exception? exception = attempt.Exception; Debug.Assert(exception is not null, "A nonterminal update always reports a failed invocation."); - int nextClientIndex = (_requestStates.TryGetValue(context, out int attemptedIndex) ? attemptedIndex : 0) + 1; - if (nextClientIndex < _clients.Length) + OrderedRoutingContext state = GetState(context); + if (state.NextClientIndex + 1 < _clients.Length) { - _requestStates[context] = nextClientIndex; + state.NextClientIndex++; return default; } - // Every client has failed. Release the state before the final failure ends routing. - _ = _requestStates.TryRemove(context, out _); + // Every client has failed, so the final failure ends routing. ExceptionDispatchInfo.Capture(exception!).Throw(); throw exception!; } @@ -119,7 +113,6 @@ protected override void Dispose(bool disposing) } _disposed = true; - _requestStates.Clear(); if (disposing && !_leaveOpen) { @@ -131,4 +124,23 @@ protected override void Dispose(bool disposing) base.Dispose(disposing); } + + private static OrderedRoutingContext GetState(RoutingContext context) + { + if (context is not OrderedRoutingContext state) + { + Throw.ArgumentException( + nameof(context), + $"The context was not created by {nameof(CreateContext)}."); + return null!; + } + + return state; + } + + private sealed class OrderedRoutingContext(IEnumerable messages, ChatOptions? chatOptions) + : RoutingContext(messages, chatOptions) + { + public int NextClientIndex { get; set; } + } } diff --git a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json index 564b3369241..6b6e7a6a8b5 100644 --- a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json +++ b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json @@ -1491,6 +1491,10 @@ "Member": "Microsoft.Extensions.AI.OrderedFailoverChatClient.OrderedFailoverChatClient(System.Collections.Generic.IReadOnlyList clients, bool leaveOpen = false);", "Stage": "Experimental" }, + { + "Member": "override Microsoft.Extensions.AI.RoutingContext Microsoft.Extensions.AI.OrderedFailoverChatClient.CreateContext(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? options);", + "Stage": "Experimental" + }, { "Member": "override void Microsoft.Extensions.AI.OrderedFailoverChatClient.Dispose(bool disposing);", "Stage": "Experimental" diff --git a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs index b8b7d8439c3..aa5938f7e5e 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs @@ -117,6 +117,40 @@ public async Task Create_DoesNotDisposeSelectedClient() Assert.Equal(0, selected.DisposeCount); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task CreateContext_SuppliesCustomContextToSelection(bool streaming) + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var selected = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + using var router = new CustomContextTestRouter(selected); + + ChatResponse response = streaming + ? await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync() + : await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Equal("ok", response.Text); + Assert.Equal(1, router.ContextsCreated); + Assert.Equal(1, router.CustomContextsObserved); + } + + [Fact] + public async Task CreateContext_NullResultThrows() + { + using var selected = new TestChatClient(); + using var router = new NullContextTestRouter(selected); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Contains("CreateContext", exception.Message, StringComparison.Ordinal); + } + [Fact] public void GetService_ReturnsSelfAndNullForUnknownOrKeyed() { @@ -129,6 +163,15 @@ public void GetService_ReturnsSelfAndNullForUnknownOrKeyed() Assert.Null(client.GetService(typeof(string))); } + private static async IAsyncEnumerable YieldUpdates(params string[] texts) + { + foreach (string text in texts) + { + await Task.Yield(); + yield return new ChatResponseUpdate(ChatRole.Assistant, text); + } + } + private static async Task> CollectAsync( IAsyncEnumerable updates) { @@ -156,6 +199,61 @@ protected override ValueTask SelectClientAsync( new(_select(context)); } + private sealed class CustomContextTestRouter : RoutingChatClient + { + private readonly IChatClient _client; + + public CustomContextTestRouter(IChatClient client) + { + _client = client; + } + + public int ContextsCreated { get; private set; } + + public int CustomContextsObserved { get; private set; } + + protected override RoutingContext CreateContext( + IEnumerable messages, ChatOptions? options) + { + ContextsCreated++; + return new TaggedRoutingContext(messages, options); + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) + { + if (context is TaggedRoutingContext) + { + CustomContextsObserved++; + } + + return new(_client); + } + + private sealed class TaggedRoutingContext(IEnumerable messages, ChatOptions? chatOptions) + : RoutingContext(messages, chatOptions); + } + + private sealed class NullContextTestRouter : RoutingChatClient + { + private readonly IChatClient _client; + + public NullContextTestRouter(IChatClient client) + { + _client = client; + } + + protected override RoutingContext CreateContext( + IEnumerable messages, ChatOptions? options) => + null!; + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + new(_client); + } + private sealed class CountingDisposeClient : IChatClient { public int DisposeCount { get; private set; } diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs index 6af1df345aa..58fd8a78ea2 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs @@ -27,6 +27,46 @@ public void MaximumAttemptsPerRequest_ValidatesValue() Assert.Throws(() => client.MaximumAttemptsPerRequest = -1); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task CreateContext_SuppliesCustomContextToSelectionAndUpdates(bool streaming) + { + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var succeeding = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + using var router = new StatefulFailoverTestRouter(failing, succeeding); + + ChatResponse response = streaming + ? await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync() + : await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Equal("ok", response.Text); + Assert.Equal(2, router.ObservedContexts.Count); + Assert.Same(router.ObservedContexts[0], router.ObservedContexts[1]); + Assert.Equal(1, router.LastObservedAttemptNumber); + } + + [Fact] + public async Task CreateContext_NullResultThrows() + { + using var selected = new TestChatClient(); + using var router = new NullContextFailoverTestRouter(selected); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Contains("CreateContext", exception.Message, StringComparison.Ordinal); + } + [Fact] public async Task Failover_RejectsNullMessagesForNonStreamingAndStreaming() { @@ -1224,6 +1264,74 @@ protected override ValueTask OnRoutingUpdateAsync( } } + private sealed class StatefulFailoverTestRouter : FailoverChatClient + { + private readonly IChatClient[] _clients; + + public StatefulFailoverTestRouter(params IChatClient[] clients) + { + _clients = clients; + } + + public List ObservedContexts { get; } = []; + + public int LastObservedAttemptNumber { get; private set; } + + protected override RoutingContext CreateContext( + IEnumerable messages, ChatOptions? options) => + new CountingRoutingContext(messages, options); + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) + { + ObservedContexts.Add(context); + var state = (CountingRoutingContext)context; + return new(_clients[state.AttemptNumber]); + } + + protected override ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt attempt, + bool isTerminal, + CancellationToken cancellationToken) + { + var state = (CountingRoutingContext)context; + LastObservedAttemptNumber = state.AttemptNumber; + if (!isTerminal) + { + state.AttemptNumber++; + } + + return default; + } + + private sealed class CountingRoutingContext(IEnumerable messages, ChatOptions? chatOptions) + : RoutingContext(messages, chatOptions) + { + public int AttemptNumber { get; set; } + } + } + + private sealed class NullContextFailoverTestRouter : FailoverChatClient + { + private readonly IChatClient _client; + + public NullContextFailoverTestRouter(IChatClient client) + { + _client = client; + } + + protected override RoutingContext CreateContext( + IEnumerable messages, ChatOptions? options) => + null!; + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + new(_client); + } + private sealed class ThrowingGetAsyncEnumeratorEnumerable : IAsyncEnumerable { private readonly Exception _exception; diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs index dd2f326617b..af81af07285 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs @@ -189,11 +189,16 @@ .. Enumerable.Range(0, RequestCount).Select(index => } [Fact] - public async Task OrderedFailover_ReleasesStateWhenStreamingEnds() + public async Task OrderedFailover_AbandonedStreamDoesNotAffectLaterRequests() { + int firstCalls = 0; using var first = new TestChatClient { - GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + GetStreamingResponseAsyncCallback = (_, _, _) => + { + firstCalls++; + return ThrowingStream("failed"); + }, }; using var second = new TestChatClient { @@ -201,14 +206,18 @@ public async Task OrderedFailover_ReleasesStateWhenStreamingEnds() }; using var client = new OrderedFailoverChatClient([first, second]); - await using (IAsyncEnumerator enumerator = - client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).GetAsyncEnumerator()) - { - Assert.True(await enumerator.MoveNextAsync()); - Assert.Equal("first", enumerator.Current.Text); - } + // Abandon the enumerator mid-stream without disposing it, so no terminal routing update is made. + IAsyncEnumerator abandoned = + client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).GetAsyncEnumerator(); + Assert.True(await abandoned.MoveNextAsync()); + Assert.Equal("first", abandoned.Current.Text); - Assert.Equal(0, GetOrderedFailoverRequestStateCount(client)); + // A later request still starts from the first client because state is scoped to its own context. + ChatResponse response = + await client.GetStreamingResponseAsync([new(ChatRole.User, "again")]).ToChatResponseAsync(); + + Assert.Equal(2, firstCalls); + Assert.Equal("firstsecond", response.Text); } [Fact] @@ -364,16 +373,6 @@ public async Task NestedRouters_ReturnLeafResponse() Assert.Same(expected, response); } - private static int GetOrderedFailoverRequestStateCount(OrderedFailoverChatClient client) - { - object requestStates = typeof(OrderedFailoverChatClient) - .GetField( - "_requestStates", - System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic)! - .GetValue(client)!; - return (int)requestStates.GetType().GetProperty("Count")!.GetValue(requestStates)!; - } - private sealed class CountingDisposeClient : IChatClient { public int DisposeCount { get; private set; } From 42919bb1a653190ee8eacd3f3ccb7297c1af7acf Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Fri, 7 Aug 2026 10:01:52 -0700 Subject: [PATCH 2/2] Assert the ordered failover context type instead of throwing The cast cannot fail: OrderedFailoverChatClient is sealed, FailoverChatClient seals both invocation methods, and the selection and update methods are protected, so the only context they receive is the one CreateContext produced. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/OrderedFailoverChatClient.cs | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs index 639491c0dcc..5f177a53e76 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -70,6 +70,7 @@ protected override ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) { + _ = Throw.IfNull(context); _ = cancellationToken; return new(_clients[GetState(context).NextClientIndex]); @@ -127,15 +128,8 @@ protected override void Dispose(bool disposing) private static OrderedRoutingContext GetState(RoutingContext context) { - if (context is not OrderedRoutingContext state) - { - Throw.ArgumentException( - nameof(context), - $"The context was not created by {nameof(CreateContext)}."); - return null!; - } - - return state; + Debug.Assert(context is OrderedRoutingContext, "The context was created by CreateContext."); + return (OrderedRoutingContext)context; } private sealed class OrderedRoutingContext(IEnumerable messages, ChatOptions? chatOptions)