Add a routing context factory for request-scoped policy state

src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs

@@ -39,6 +39,25 @@ public static RoutingChatClient Create(
         return new CallbackRoutingChatClient(clientSelector);
     }
 
+    /// <summary>Creates the context supplied to client selection for one request.</summary>
+    /// <param name="messages">The messages to route.</param>
+    /// <param name="options">The options supplied by the caller, or <see langword="null"/>.</param>
+    /// <returns>The context for the request.</returns>
+    /// <remarks>
+    /// <para>
+    /// The default implementation returns a new <see cref="RoutingContext"/>. A derived class can return its own
+    /// <see cref="RoutingContext"/> subclass to supply additional request inputs or to carry request-scoped policy
+    /// state, which <see cref="SelectClientAsync"/> then reads by casting the supplied context. State stored on the
+    /// context is released with the request.
+    /// </para>
+    /// <para>
+    /// This method is called once per request by <see cref="GetResponseAsync"/> and
+    /// <see cref="GetStreamingResponseAsync"/>. Exceptions from this method propagate to the caller.
+    /// </para>
+    /// </remarks>
+    protected virtual RoutingContext CreateContext(IEnumerable<ChatMessage> messages, ChatOptions? options) =>
+        new(messages, options);
+
     /// <summary>Selects the client to invoke for the request.</summary>
     /// <param name="context">The request-specific inputs.</param>
     /// <param name="cancellationToken">The cancellation token supplied for the request.</param>
@@ -57,7 +76,8 @@ public virtual async Task<ChatResponse> 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<ChatResponseUpdate> 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.");
 

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<Microsoft.Extensions.AI.RoutingContext, System.Threading.CancellationToken, System.Threading.Tasks.ValueTask<Microsoft.Extensions.AI.IChatClient>> clientSelector);",
           "Stage": "Experimental"
         },
+        {
+          "Member": "virtual Microsoft.Extensions.AI.RoutingContext Microsoft.Extensions.AI.RoutingChatClient.CreateContext(System.Collections.Generic.IEnumerable<Microsoft.Extensions.AI.ChatMessage> messages, Microsoft.Extensions.AI.ChatOptions? options);",
+          "Stage": "Experimental"
+        },
         {
           "Member": "void Microsoft.Extensions.AI.RoutingChatClient.Dispose();",
           "Stage": "Experimental"

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.
 /// </para>
 /// <para>
+/// One <see cref="RoutingContext"/> 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
+/// <see cref="RoutingChatClient.CreateContext"/> to return its own <see cref="RoutingContext"/> subclass and store
+/// that state on the context, which releases it with the request instead of requiring explicit cleanup.
+/// </para>
+/// <para>
 /// 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.
 /// </para>
@@ -99,7 +105,8 @@ public sealed override async Task<ChatResponse> 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<ChatResponseUpdate> 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;
 

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<RoutingContext, int> _requestStates = new();
     private bool _disposed;
 
     /// <summary>Initializes a new instance of the <see cref="OrderedFailoverChatClient"/> class.</summary>
@@ -66,6 +61,10 @@ public OrderedFailoverChatClient(IReadOnlyList<IChatClient> clients, bool leaveO
         _clients = clientsSnapshot;
     }
 
+    /// <inheritdoc/>
+    protected override RoutingContext CreateContext(IEnumerable<ChatMessage> messages, ChatOptions? options) =>
+        new OrderedRoutingContext(messages, options);
+
     /// <inheritdoc/>
     protected override ValueTask<IChatClient> SelectClientAsync(
         RoutingContext context,
@@ -74,9 +73,7 @@ protected override ValueTask<IChatClient> SelectClientAsync(
         _ = Throw.IfNull(context);
         _ = cancellationToken;
 
-        int clientIndex = _requestStates.TryGetValue(context, out int nextClientIndex) ? nextClientIndex : 0;
-
-        return new(_clients[clientIndex]);
+        return new(_clients[GetState(context).NextClientIndex]);
     }
 
     /// <inheritdoc/>
@@ -90,22 +87,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 +114,6 @@ protected override void Dispose(bool disposing)
         }
 
         _disposed = true;
-        _requestStates.Clear();
 
         if (disposing && !_leaveOpen)
         {
@@ -131,4 +125,16 @@ protected override void Dispose(bool disposing)
 
         base.Dispose(disposing);
     }
+
+    private static OrderedRoutingContext GetState(RoutingContext context)
+    {
+        Debug.Assert(context is OrderedRoutingContext, "The context was created by CreateContext.");
+        return (OrderedRoutingContext)context;
+    }
+
+    private sealed class OrderedRoutingContext(IEnumerable<ChatMessage> messages, ChatOptions? chatOptions)
+        : RoutingContext(messages, chatOptions)
+    {
+        public int NextClientIndex { get; set; }
+    }
 }

src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json

@@ -1491,6 +1491,10 @@
           "Member": "Microsoft.Extensions.AI.OrderedFailoverChatClient.OrderedFailoverChatClient(System.Collections.Generic.IReadOnlyList<Microsoft.Extensions.AI.IChatClient> clients, bool leaveOpen = false);",
           "Stage": "Experimental"
         },
+        {
+          "Member": "override Microsoft.Extensions.AI.RoutingContext Microsoft.Extensions.AI.OrderedFailoverChatClient.CreateContext(System.Collections.Generic.IEnumerable<Microsoft.Extensions.AI.ChatMessage> messages, Microsoft.Extensions.AI.ChatOptions? options);",
+          "Stage": "Experimental"
+        },
         {
           "Member": "override void Microsoft.Extensions.AI.OrderedFailoverChatClient.Dispose(bool disposing);",
           "Stage": "Experimental"

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<InvalidOperationException>(
+            () => 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<ChatResponseUpdate> YieldUpdates(params string[] texts)
+    {
+        foreach (string text in texts)
+        {
+            await Task.Yield();
+            yield return new ChatResponseUpdate(ChatRole.Assistant, text);
+        }
+    }
+
     private static async Task<List<ChatResponseUpdate>> CollectAsync(
         IAsyncEnumerable<ChatResponseUpdate> updates)
     {
@@ -156,6 +199,61 @@ protected override ValueTask<IChatClient> 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<ChatMessage> messages, ChatOptions? options)
+        {
+            ContextsCreated++;
+            return new TaggedRoutingContext(messages, options);
+        }
+
+        protected override ValueTask<IChatClient> SelectClientAsync(
+            RoutingContext context,
+            CancellationToken cancellationToken)
+        {
+            if (context is TaggedRoutingContext)
+            {
+                CustomContextsObserved++;
+            }
+
+            return new(_client);
+        }
+
+        private sealed class TaggedRoutingContext(IEnumerable<ChatMessage> 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<ChatMessage> messages, ChatOptions? options) =>
+            null!;
+
+        protected override ValueTask<IChatClient> SelectClientAsync(
+            RoutingContext context,
+            CancellationToken cancellationToken) =>
+            new(_client);
+    }
+
     private sealed class CountingDisposeClient : IChatClient
     {
         public int DisposeCount { get; private set; }