From bd114df1716e1526c620d96199e4d1bddbdaddc6 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Mon, 27 Jul 2026 15:24:36 -0700 Subject: [PATCH 01/15] Add extensible chat client routing Introduce one-shot and failover routing primitives, ordered and semantic routing implementations, lifecycle reporting, streaming safeguards, and focused routing tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 113 ++ .../ChatRouting/RoutingContext.cs | 49 + .../Microsoft.Extensions.AI.Abstractions.json | 58 + .../ChatRouting/FailoverChatClient.cs | 337 ++++ .../ChatRouting/FailoverChatClientAttempt.cs | 90 + .../ChatRouting/OrderedFailoverChatClient.cs | 184 ++ .../ChatRouting/SemanticRoutingChatClient.cs | 370 ++++ .../Microsoft.Extensions.AI.json | 120 ++ src/Shared/DiagnosticIds/DiagnosticIds.cs | 1 + .../ChatRouting/RoutingChatClientTests.cs | 1758 +++++++++++++++++ 10 files changed, 3080 insertions(+) create mode 100644 src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs create mode 100644 src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs create mode 100644 src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs create mode 100644 src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs create mode 100644 src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs create mode 100644 src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs create mode 100644 test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs new file mode 100644 index 00000000000..5b2acfc830e --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -0,0 +1,113 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// +/// Provides a template for an that selects and invokes another chat client. +/// +/// +/// Derived classes implement to supply one client for each request. The selected +/// client is invoked once, and its response or failure is propagated to the caller. +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public abstract class RoutingChatClient : IChatClient +{ + /// Creates a routing client that selects one client for each request. + /// The callback that selects the client to invoke. + /// A routing client that uses . + /// The selected clients are caller-owned and are not disposed by the returned routing client. + /// is . + public static RoutingChatClient Create( + Func> clientSelector) + { + _ = Throw.IfNull(clientSelector); + return new CallbackRoutingChatClient(clientSelector); + } + + /// Selects the client to invoke for the request. + /// The request-specific inputs. + /// The cancellation token supplied for the request. + /// The client to invoke. + /// Exceptions from this method propagate to the caller. + protected abstract ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken); + + /// + public virtual async Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + var context = new RoutingContext(messages, options); + IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false); + return await client.GetResponseAsync( + context.Messages, + context.ChatOptions, + cancellationToken).ConfigureAwait(false); + } + + /// + public virtual async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + var context = new RoutingContext(messages, options); + IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false); + await foreach (ChatResponseUpdate update in + client.GetStreamingResponseAsync(context.Messages, context.ChatOptions, cancellationToken) + .WithCancellation(cancellationToken) + .ConfigureAwait(false)) + { + yield return update; + } + } + + /// + public virtual object? GetService(Type serviceType, object? serviceKey = null) + { + _ = Throw.IfNull(serviceType); + return serviceKey is null && serviceType.IsInstanceOfType(this) ? this : null; + } + + /// + public void Dispose() + { + Dispose(disposing: true); + GC.SuppressFinalize(this); + } + + /// Provides a mechanism for releasing resources owned by the derived instance. + /// when called from . + /// The default implementation performs no operation. + protected virtual void Dispose(bool disposing) + { + } + + private sealed class CallbackRoutingChatClient : RoutingChatClient + { + private readonly Func> _clientSelector; + + public CallbackRoutingChatClient( + Func> clientSelector) + { + _clientSelector = clientSelector; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + _clientSelector(context, cancellationToken); + } +} diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs new file mode 100644 index 00000000000..7a0c9b14011 --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs @@ -0,0 +1,49 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Provides request-specific inputs to a . +/// +/// +/// One context is created for each call to and for each enumeration +/// started from the sequence returned by . +/// +/// +/// A derived client may replace or during selection. The selected +/// client receives the updated values. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public class RoutingContext +{ + /// Initializes a new instance of the class. + /// The messages to route. + /// The options supplied for the request. + public RoutingContext( + IEnumerable messages, + ChatOptions? chatOptions) + { + Messages = Throw.IfNull(messages); + ChatOptions = chatOptions; + } + + /// Gets or sets the messages supplied to client selection and the selected client. + /// + /// The sequence may be enumerated during selection and again by one or more selected clients. Callers should + /// supply a repeatable sequence, or a selector may replace it with a materialized sequence when required. + /// + public IEnumerable Messages + { + get; + set => field = Throw.IfNull(value); + } + + /// Gets or sets the options supplied to client selection and the selected client. + public ChatOptions? ChatOptions { get; set; } +} 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 dc4542d4ae7..f7fa67c49bd 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 @@ -3966,6 +3966,64 @@ } ] }, + { + "Type": "abstract class Microsoft.Extensions.AI.RoutingChatClient : Microsoft.Extensions.AI.IChatClient, System.IDisposable", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.RoutingChatClient.RoutingChatClient();", + "Stage": "Experimental" + }, + { + "Member": "static Microsoft.Extensions.AI.RoutingChatClient Microsoft.Extensions.AI.RoutingChatClient.Create(System.Func> clientSelector);", + "Stage": "Experimental" + }, + { + "Member": "void Microsoft.Extensions.AI.RoutingChatClient.Dispose();", + "Stage": "Experimental" + }, + { + "Member": "virtual void Microsoft.Extensions.AI.RoutingChatClient.Dispose(bool disposing);", + "Stage": "Experimental" + }, + { + "Member": "virtual System.Threading.Tasks.Task Microsoft.Extensions.AI.RoutingChatClient.GetResponseAsync(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));", + "Stage": "Experimental" + }, + { + "Member": "virtual object? Microsoft.Extensions.AI.RoutingChatClient.GetService(System.Type serviceType, object? serviceKey = null);", + "Stage": "Experimental" + }, + { + "Member": "virtual System.Collections.Generic.IAsyncEnumerable Microsoft.Extensions.AI.RoutingChatClient.GetStreamingResponseAsync(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));", + "Stage": "Experimental" + }, + { + "Member": "abstract System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.RoutingChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + } + ] + }, + { + "Type": "class Microsoft.Extensions.AI.RoutingContext", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.RoutingContext.RoutingContext(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? chatOptions);", + "Stage": "Experimental" + } + ], + "Properties": [ + { + "Member": "Microsoft.Extensions.AI.ChatOptions? Microsoft.Extensions.AI.RoutingContext.ChatOptions { get; set; }", + "Stage": "Experimental" + }, + { + "Member": "System.Collections.Generic.IEnumerable Microsoft.Extensions.AI.RoutingContext.Messages { get; set; }", + "Stage": "Experimental" + } + ] + }, { "Type": "class Microsoft.Extensions.AI.SessionUpdateRealtimeClientMessage : Microsoft.Extensions.AI.RealtimeClientMessage", "Stage": "Experimental", diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs new file mode 100644 index 00000000000..4552c89826f --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -0,0 +1,337 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Diagnostics.CodeAnalysis; +using System.Runtime.CompilerServices; +using System.Runtime.ExceptionServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// +/// Provides a template for a that can select another client after an invocation fails. +/// +/// +/// +/// The client for each attempt is supplied by . After an invocation, +/// reports its outcome and whether routing is terminal. An uncanceled failure causes +/// another selection only when it happened before any streaming output was exposed and the attempt limit permits it. +/// +/// +/// The base class owns invocation, streaming commitment, attempt limits, and terminal reporting. Derived classes own +/// client selection, policy state, and the lifetime of clients they retain. +/// +/// +/// 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. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public abstract class FailoverChatClient : RoutingChatClient +{ + /// Gets or sets the maximum number of client invocations permitted for one request. + /// + /// A positive attempt limit, or to leave termination to client selection and request + /// cancellation. The default is . + /// + /// + /// The value is captured when a non-streaming request begins or when a streaming response begins enumeration. + /// Changing it does not affect requests or enumerations already in progress. + /// + /// The value is not or positive. + public int? MaximumAttemptsPerRequest + { + get; + set + { + if (value is <= 0) + { + Throw.ArgumentOutOfRangeException(nameof(value)); + } + + field = value; + } + } + + /// Invoked after a client invocation or when client selection terminates routing. + /// The request-specific inputs. + /// + /// The completed client invocation, or when + /// terminated routing before invoking a client. + /// + /// + /// if the base will not select another client after this callback completes successfully; + /// otherwise, . + /// + /// The cancellation token supplied for the request. + /// A task representing the update operation. + /// + /// + /// The default implementation performs no operation. A nonterminal update always contains an uncanceled, + /// pre-output failed attempt. State changes made by the override are visible to the next call to + /// . + /// + /// + /// A attempt is always terminal. If every callback completes successfully, every client + /// invocation produces one update and exactly one update per request is terminal. + /// + /// + /// Exceptions from this method propagate to the caller. A terminal update exception replaces the response or + /// exception already produced by the request. A nonterminal update exception stops routing without another update. + /// An override that retains per-request state must release that state before throwing because no later update is + /// made after an update exception. + /// + /// + protected virtual ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt? attempt, + bool isTerminal, + CancellationToken cancellationToken) => default; + + /// + public sealed override async Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + var context = new RoutingContext(messages, options); + int? maximumAttempts = MaximumAttemptsPerRequest; + int attemptCount = 0; + + while (true) + { + IChatClient selectedClient; + try + { + selectedClient = await SelectClientAsync(context, cancellationToken) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + } + catch (Exception) + { + await OnRoutingUpdateAsync(context, attempt: null, isTerminal: true, cancellationToken); + throw; + } + + attemptCount++; + ChatResponse? response = null; + Exception? exception = null; + long start = Stopwatch.GetTimestamp(); + + try + { + response = await selectedClient.GetResponseAsync( + context.Messages, + context.ChatOptions, + cancellationToken); + } + catch (Exception ex) + { + exception = ex; + } + + var attempt = new FailoverChatClientAttempt( + selectedClient, + exception, + GetElapsedTime(start), + timeToFirstUpdate: null, + responseCompleted: exception is null, + outputCommitted: false); + bool isTerminal = + exception is null || + cancellationToken.IsCancellationRequested || + (maximumAttempts is int limit && attemptCount >= limit); + + await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken); + + if (exception is null) + { + return response!; + } + + if (isTerminal) + { + Rethrow(exception); + } + } + } + + /// + public sealed override async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + var context = new RoutingContext(messages, options); + int? maximumAttempts = MaximumAttemptsPerRequest; + int attemptCount = 0; + + while (true) + { + IChatClient selectedClient; + try + { + selectedClient = await SelectClientAsync(context, cancellationToken) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + } + catch (Exception) + { + await OnRoutingUpdateAsync(context, attempt: null, isTerminal: true, cancellationToken); + throw; + } + + attemptCount++; + bool reachedAttemptLimit = maximumAttempts is int limit && attemptCount >= limit; + IAsyncEnumerator? enumerator = null; + TimeSpan? timeToFirstUpdate = null; + TimeSpan activeDuration = TimeSpan.Zero; + bool hasCurrent; + + long operationStart = Stopwatch.GetTimestamp(); + try + { + enumerator = selectedClient + .GetStreamingResponseAsync( + context.Messages, + context.ChatOptions, + cancellationToken) + .GetAsyncEnumerator(cancellationToken); + + hasCurrent = await enumerator.MoveNextAsync(); + } + catch (Exception ex) + { + activeDuration += GetElapsedTime(operationStart); + Exception exception = (await DisposeAsync(enumerator, ex, cancellationToken))!; + var attempt = new FailoverChatClientAttempt( + selectedClient, + exception, + activeDuration, + timeToFirstUpdate: null, + responseCompleted: false, + outputCommitted: false); + bool isTerminal = cancellationToken.IsCancellationRequested || reachedAttemptLimit; + + await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken); + + if (isTerminal) + { + Rethrow(exception); + } + + continue; + } + + activeDuration += GetElapsedTime(operationStart); + if (hasCurrent) + { + timeToFirstUpdate = activeDuration; + } + + bool responseCompleted = false; + bool outputCommitted = false; + bool isTerminalAttempt = false; + Exception? terminalException = null; + + try + { + while (hasCurrent) + { + outputCommitted = true; + yield return enumerator.Current; + + operationStart = Stopwatch.GetTimestamp(); + try + { + hasCurrent = await enumerator.MoveNextAsync(); + } + catch (Exception ex) + { + terminalException = ex; + break; + } + finally + { + activeDuration += GetElapsedTime(operationStart); + } + } + + responseCompleted = terminalException is null; + } + finally + { + terminalException = await DisposeAsync(enumerator, terminalException, cancellationToken); + + var attempt = new FailoverChatClientAttempt( + selectedClient, + terminalException, + activeDuration, + timeToFirstUpdate, + responseCompleted: responseCompleted && terminalException is null, + outputCommitted: outputCommitted); + isTerminalAttempt = + attempt.ResponseCompleted || + outputCommitted || + cancellationToken.IsCancellationRequested || + reachedAttemptLimit; + + await OnRoutingUpdateAsync(context, attempt, isTerminalAttempt, cancellationToken); + + if (terminalException is not null && isTerminalAttempt) + { + Rethrow(terminalException); + } + } + + if (!isTerminalAttempt) + { + continue; + } + + yield break; + } + } + + private static async ValueTask DisposeAsync( + IAsyncDisposable? disposable, + Exception? exception, + CancellationToken cancellationToken) + { + _ = cancellationToken; + + if (disposable is not null) + { + try + { + await disposable.DisposeAsync(); + } + catch (Exception ex) + { + return ex; + } + } + + return exception; + } + + private static TimeSpan GetElapsedTime(long startingTimestamp) => +#if NET + Stopwatch.GetElapsedTime(startingTimestamp); +#else + new((long)((Stopwatch.GetTimestamp() - startingTimestamp) * + ((double)TimeSpan.TicksPerSecond / Stopwatch.Frequency))); +#endif + + [DoesNotReturn] + private static void Rethrow(Exception exception) + { + ExceptionDispatchInfo.Capture(exception).Throw(); + throw exception; + } + +} diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs new file mode 100644 index 00000000000..04c32086a0a --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs @@ -0,0 +1,90 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Diagnostics; +using System.Diagnostics.CodeAnalysis; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Represents one client invocation performed by a . +/// +/// +/// A completed response has set to and +/// set to . +/// +/// +/// A failed invocation has set to and +/// set to the observed exception. If a streaming caller stops enumerating before the +/// response ends and disposal completes successfully, both and +/// are unset. +/// +/// +/// is independent of the outcome and indicates whether any streaming update was +/// exposed to the caller. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public sealed class FailoverChatClientAttempt +{ + internal FailoverChatClientAttempt( + IChatClient client, + Exception? exception, + TimeSpan duration, + TimeSpan? timeToFirstUpdate, + bool responseCompleted, + bool outputCommitted) + { + Debug.Assert(client is not null, "Expected a non-null invoked client."); + Debug.Assert(duration >= TimeSpan.Zero, "Expected a non-negative duration."); + Debug.Assert( + timeToFirstUpdate is not { } ttfu || (ttfu >= TimeSpan.Zero && ttfu <= duration), + "Expected time to first update to be within the active duration."); + Debug.Assert( + !responseCompleted || exception is null, + "A completed response should not have an exception."); + Debug.Assert( + outputCommitted == timeToFirstUpdate.HasValue, + "Output commitment and time to first update should agree."); + + Client = Throw.IfNull(client); + Exception = exception; + Duration = duration; + TimeToFirstUpdate = timeToFirstUpdate; + ResponseCompleted = responseCompleted; + OutputCommitted = outputCommitted; + } + + /// Gets the client that was invoked. + public IChatClient Client { get; } + + /// Gets the time spent actively invoking the client. + /// For streaming, time spent by the caller processing yielded updates is excluded. + public TimeSpan Duration { get; } + + /// Gets the exception observed while invoking or disposing the client, if any. + /// + /// If invocation and disposal both throw, this contains the disposal exception. + /// + /// A value does not necessarily indicate success; inspect + /// to distinguish a completed response from a streaming response that the caller stopped consuming. + /// + /// + public Exception? Exception { get; } + + /// Gets a value indicating whether any streaming update was exposed to the caller. + /// This is always for non-streaming invocations. + public bool OutputCommitted { get; } + + /// Gets a value indicating whether the response completed successfully. + /// + /// For streaming responses, this is when the caller stops enumeration before the + /// response stream ends. + /// + public bool ResponseCompleted { get; } + + /// Gets the time until the first streaming update, if this was a non-empty streaming invocation. + public TimeSpan? TimeToFirstUpdate { get; } +} diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs new file mode 100644 index 00000000000..e87130c8b17 --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -0,0 +1,184 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Runtime.ExceptionServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Provides ordered failover across a sequence of chat clients. +/// +/// +/// The clients are tried in order. An invocation failure before streaming output is exposed advances to the next +/// client. Cancellation and failures after streaming output is exposed are propagated without failover. +/// +/// +/// The configured clients are snapshotted by the constructor and must contain unique object references. When every +/// client has failed, the final failure is rethrown. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public sealed class OrderedFailoverChatClient : FailoverChatClient +{ + private readonly bool _leaveOpen; + private readonly IChatClient[] _clients; + private readonly Dictionary _requestStates = []; + private readonly object _sync = new(); + private bool _disposed; + + /// Initializes a new instance of the class. + /// The clients to invoke, in fallback order. + /// + /// to leave inner clients open when this instance is disposed; + /// otherwise, . + /// + /// is . + /// + /// is empty, contains , or contains the same client instance more than once. + /// + public OrderedFailoverChatClient(IReadOnlyList clients, bool leaveOpen = false) + { + _ = Throw.IfNull(clients); + _leaveOpen = leaveOpen; + if (clients.Count == 0) + { + Throw.ArgumentException(nameof(clients), "At least one client must be provided."); + } + + _clients = [.. clients]; + for (int i = 0; i < _clients.Length; i++) + { + if (_clients[i] is null) + { + Throw.ArgumentException(nameof(clients), "Clients must not contain null."); + } + + for (int j = 0; j < i; j++) + { + if (ReferenceEquals(_clients[j], _clients[i])) + { + Throw.ArgumentException(nameof(clients), "Each client instance must be unique."); + } + } + } + } + + /// + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) + { + _ = cancellationToken; + RequestState? state; + + lock (_sync) + { + if (!_requestStates.TryGetValue(context, out state)) + { + return new(_clients[0]); + } + + _ = _requestStates.Remove(context); + } + + if (state.ClientIndex < _clients.Length) + { + return new(_clients[state.ClientIndex]); + } + + ExceptionDispatchInfo.Capture(state.LastException).Throw(); + throw state.LastException; + } + + /// + protected override ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt? attempt, + bool isTerminal, + CancellationToken cancellationToken) + { + _ = cancellationToken; + + lock (_sync) + { + if (isTerminal) + { + _ = _requestStates.Remove(context); + return default; + } + + if (attempt?.Exception is null) + { + _ = _requestStates.Remove(context); + throw new InvalidOperationException("A nonterminal routing update requires a failed client invocation."); + } + + int clientIndex = IndexOfClient(attempt.Client); + if (clientIndex < 0) + { + _ = _requestStates.Remove(context); + throw new InvalidOperationException("The invocation did not use a configured ordered failover client."); + } + + _requestStates[context] = new(clientIndex + 1, attempt.Exception); + } + + return default; + } + + /// + protected override void Dispose(bool disposing) + { + if (_disposed) + { + return; + } + + _disposed = true; + try + { + lock (_sync) + { + _requestStates.Clear(); + } + + if (disposing && !_leaveOpen) + { + foreach (IChatClient client in _clients) + { + client.Dispose(); + } + } + } + finally + { + base.Dispose(disposing); + } + } + + private int IndexOfClient(IChatClient client) + { + for (int i = 0; i < _clients.Length; i++) + { + if (ReferenceEquals(_clients[i], client)) + { + return i; + } + } + + return -1; + } + + private sealed class RequestState(int clientIndex, Exception lastException) + { + public int ClientIndex { get; } = clientIndex; + + public Exception LastException { get; } = lastException; + } +} diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs new file mode 100644 index 00000000000..d6ae61fad65 --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -0,0 +1,370 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Linq; +using System.Numerics.Tensors; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Routes requests by semantic similarity to app-provided example utterances. +/// +/// +/// Profile embeddings are generated lazily and cached. Each request embeds the last user message and selects the +/// client with the highest score after aggregating the cosine similarities of the best-matching profile utterances. +/// The configured default client is selected when no user message is available or when the highest score is below the +/// configured threshold. +/// +/// +/// The configured clients are used as stable routing identities. By default this instance owns the clients and +/// embedding generator and disposes them when it is disposed. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public sealed class SemanticRoutingChatClient : RoutingChatClient +{ + /// Specifies how profile similarity scores are aggregated for each client. + public enum ScoreAggregation + { + /// Average the matching profile scores for each client. + Mean, + + /// Sum the matching profile scores for each client. + Sum, + } + + private readonly IChatClient[] _clients; + private readonly IEmbeddingGenerator> _embeddingGenerator; + private readonly SemaphoreSlim _indexGate = new(1, 1); + private readonly bool _leaveOpen; + private readonly (int ClientIndex, string Text)[] _profiles; + private readonly ScoreAggregation _scoreAggregation; + private readonly float _scoreThreshold; + private readonly int _topK; + + private bool _disposed; + private EmbeddedProfile[]? _index; + + /// Initializes a new instance of the class. + /// The generator used to embed profile utterances and request text. + /// The example utterances associated with each client. + /// The client selected when no profile satisfies . + /// The minimum aggregated score required to select a profiled client. + /// + /// to leave the configured clients and embedding generator open when this instance is + /// disposed; otherwise, . The default is . + /// + /// + /// The number of highest-scoring profile utterances, across all clients, whose scores are aggregated. + /// The default is 1. + /// + /// The method used to aggregate matching profile scores for each client. + /// + /// , , or + /// is . + /// + /// + /// is empty or contains a null client, an empty utterance list, or a blank + /// utterance. + /// + /// + /// is not positive, is invalid, or + /// is outside the possible range for the configured aggregation. + /// + public SemanticRoutingChatClient( + IEmbeddingGenerator> embeddingGenerator, + IReadOnlyDictionary> clientProfiles, + IChatClient defaultClient, + float scoreThreshold = 0.3f, + bool leaveOpen = false, + int topK = 1, + ScoreAggregation scoreAggregation = ScoreAggregation.Mean) + { + _embeddingGenerator = Throw.IfNull(embeddingGenerator); + _ = Throw.IfNull(clientProfiles); + _ = Throw.IfNull(defaultClient); + _leaveOpen = leaveOpen; + + if (clientProfiles.Count == 0) + { + Throw.ArgumentException(nameof(clientProfiles), "At least one client profile must be provided."); + } + + if (topK <= 0) + { + Throw.ArgumentOutOfRangeException(nameof(topK)); + } + + if (scoreAggregation is not ScoreAggregation.Mean and not ScoreAggregation.Sum) + { + Throw.ArgumentOutOfRangeException(nameof(scoreAggregation)); + } + + float scoreLimit = scoreAggregation == ScoreAggregation.Sum ? topK : 1; + if (float.IsNaN(scoreThreshold) || + float.IsInfinity(scoreThreshold) || + scoreThreshold < -scoreLimit || + scoreThreshold > scoreLimit) + { + Throw.ArgumentOutOfRangeException(nameof(scoreThreshold)); + } + + _topK = topK; + _scoreAggregation = scoreAggregation; + _scoreThreshold = scoreThreshold; + + var profiles = new List<(int ClientIndex, string Text)>(); + var clients = new List { defaultClient }; + foreach (KeyValuePair> profile in clientProfiles) + { + IChatClient client = profile.Key; + IReadOnlyList utterances = profile.Value; + if (client is null) + { + Throw.ArgumentException(nameof(clientProfiles), "Profile clients must not be null."); + } + + if (utterances is null || utterances.Count == 0) + { + Throw.ArgumentException( + nameof(clientProfiles), + "Every profile client must have at least one example utterance."); + } + + int clientIndex = clients.FindIndex(candidate => ReferenceEquals(candidate, client)); + if (clientIndex < 0) + { + clientIndex = clients.Count; + clients.Add(client); + } + + foreach (string utterance in utterances) + { + if (string.IsNullOrWhiteSpace(utterance)) + { + Throw.ArgumentException(nameof(clientProfiles), "Profile utterances must not be blank."); + } + + profiles.Add((clientIndex, utterance)); + } + } + + _clients = [.. clients]; + _profiles = [.. profiles]; + } + + /// + protected override async ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) + { + _ = Throw.IfNull(context); + string? query = LastUserText(context.Messages); + if (string.IsNullOrWhiteSpace(query)) + { + return _clients[0]; + } + + EmbeddedProfile[] index = await EnsureIndexAsync(cancellationToken).ConfigureAwait(false); + GeneratedEmbeddings> generated = + await _embeddingGenerator.GenerateAsync( + [query!], + cancellationToken: cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException("The embedding generator returned null."); + if (generated.Count != 1) + { + throw new InvalidOperationException("The embedding generator did not return one query embedding."); + } + + ReadOnlySpan queryVector = generated[0].Vector.Span; + if (queryVector.Length != index[0].Vector.Length) + { + throw new InvalidOperationException( + "The query embedding dimension does not match the profile embedding dimension."); + } + + if (_topK == 1) + { + int bestClientIndex = -1; + float bestScore = float.NegativeInfinity; + foreach (EmbeddedProfile profile in index) + { + float score = TensorPrimitives.CosineSimilarity(queryVector, profile.Vector); + if (score > bestScore) + { + bestClientIndex = profile.ClientIndex; + bestScore = score; + } + } + + return bestClientIndex >= 0 && bestScore >= _scoreThreshold + ? _clients[bestClientIndex] + : _clients[0]; + } + + var matches = new ScoredProfile[index.Length]; + for (int i = 0; i < index.Length; i++) + { + matches[i] = new( + i, + TensorPrimitives.CosineSimilarity(queryVector, index[i].Vector)); + } + + Array.Sort(matches, static (left, right) => + { + int scoreComparison = right.Score.CompareTo(left.Score); + return scoreComparison != 0 + ? scoreComparison + : left.ProfileIndex.CompareTo(right.ProfileIndex); + }); + + int matchCount = Math.Min(_topK, matches.Length); + var scoreSums = new float[_clients.Length]; + var scoreCounts = new int[_clients.Length]; + var clientOrder = new int[_clients.Length]; + int clientCount = 0; + for (int i = 0; i < matchCount; i++) + { + EmbeddedProfile profile = index[matches[i].ProfileIndex]; + int clientIndex = profile.ClientIndex; + if (scoreCounts[clientIndex] == 0) + { + clientOrder[clientCount++] = clientIndex; + } + + scoreSums[clientIndex] += matches[i].Score; + scoreCounts[clientIndex]++; + } + + int bestAggregatedClientIndex = -1; + float bestAggregatedScore = float.NegativeInfinity; + for (int i = 0; i < clientCount; i++) + { + int clientIndex = clientOrder[i]; + float score = _scoreAggregation == ScoreAggregation.Mean + ? scoreSums[clientIndex] / scoreCounts[clientIndex] + : scoreSums[clientIndex]; + if (score > bestAggregatedScore) + { + bestAggregatedClientIndex = clientIndex; + bestAggregatedScore = score; + } + } + + return bestAggregatedClientIndex >= 0 && bestAggregatedScore >= _scoreThreshold + ? _clients[bestAggregatedClientIndex] + : _clients[0]; + } + + /// + protected override void Dispose(bool disposing) + { + if (_disposed) + { + return; + } + + _disposed = true; + try + { + if (disposing) + { + _indexGate.Dispose(); + if (!_leaveOpen) + { + foreach (IChatClient client in _clients) + { + client.Dispose(); + } + + if (!Array.Exists( + _clients, + client => ReferenceEquals(client, _embeddingGenerator))) + { + _embeddingGenerator.Dispose(); + } + } + } + } + finally + { + base.Dispose(disposing); + } + } + + private static string? LastUserText(IEnumerable messages) + { + string? last = null; + foreach (ChatMessage message in messages) + { + if (message.Role == ChatRole.User) + { + last = message.Text; + } + } + + return last; + } + + private async Task EnsureIndexAsync(CancellationToken cancellationToken) + { + if (_index is { } cached) + { + return cached; + } + + await _indexGate.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + if (_index is { } existing) + { + return existing; + } + + GeneratedEmbeddings> embeddings = + await _embeddingGenerator.GenerateAsync( + _profiles.Select(profile => profile.Text), + cancellationToken: cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException("The embedding generator returned null."); + if (embeddings.Count != _profiles.Length) + { + throw new InvalidOperationException( + "The embedding generator did not return one embedding per profile utterance."); + } + + int dimensions = embeddings[0].Vector.Length; + if (dimensions == 0) + { + throw new InvalidOperationException("Profile embeddings must not be empty."); + } + + var index = new EmbeddedProfile[_profiles.Length]; + for (int i = 0; i < index.Length; i++) + { + if (embeddings[i].Vector.Length != dimensions) + { + throw new InvalidOperationException( + "All profile embeddings must have the same dimension."); + } + + index[i] = new(_profiles[i].ClientIndex, embeddings[i].Vector.ToArray()); + } + + return _index = index; + } + finally + { + _ = _indexGate.Release(); + } + } + + private sealed record EmbeddedProfile(int ClientIndex, float[] Vector); + + private readonly record struct ScoredProfile(int ProfileIndex, float Score); +} diff --git a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json index 9efb90882b9..c4dd509a8cf 100644 --- a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json +++ b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json @@ -489,6 +489,64 @@ } ] }, + { + "Type": "abstract class Microsoft.Extensions.AI.FailoverChatClient : Microsoft.Extensions.AI.RoutingChatClient", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.FailoverChatClient.FailoverChatClient();", + "Stage": "Experimental" + }, + { + "Member": "sealed override System.Threading.Tasks.Task Microsoft.Extensions.AI.FailoverChatClient.GetResponseAsync(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));", + "Stage": "Experimental" + }, + { + "Member": "sealed override System.Collections.Generic.IAsyncEnumerable Microsoft.Extensions.AI.FailoverChatClient.GetStreamingResponseAsync(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));", + "Stage": "Experimental" + }, + { + "Member": "virtual System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.FailoverChatClient.OnRoutingUpdateAsync(Microsoft.Extensions.AI.RoutingContext context, Microsoft.Extensions.AI.FailoverChatClientAttempt? attempt, bool isTerminal, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + } + ], + "Properties": [ + { + "Member": "int? Microsoft.Extensions.AI.FailoverChatClient.MaximumAttemptsPerRequest { get; set; }", + "Stage": "Experimental" + } + ] + }, + { + "Type": "sealed class Microsoft.Extensions.AI.FailoverChatClientAttempt", + "Stage": "Experimental", + "Properties": [ + { + "Member": "Microsoft.Extensions.AI.IChatClient Microsoft.Extensions.AI.FailoverChatClientAttempt.Client { get; }", + "Stage": "Experimental" + }, + { + "Member": "System.TimeSpan Microsoft.Extensions.AI.FailoverChatClientAttempt.Duration { get; }", + "Stage": "Experimental" + }, + { + "Member": "System.Exception? Microsoft.Extensions.AI.FailoverChatClientAttempt.Exception { get; }", + "Stage": "Experimental" + }, + { + "Member": "bool Microsoft.Extensions.AI.FailoverChatClientAttempt.OutputCommitted { get; }", + "Stage": "Experimental" + }, + { + "Member": "bool Microsoft.Extensions.AI.FailoverChatClientAttempt.ResponseCompleted { get; }", + "Stage": "Experimental" + }, + { + "Member": "System.TimeSpan? Microsoft.Extensions.AI.FailoverChatClientAttempt.TimeToFirstUpdate { get; }", + "Stage": "Experimental" + } + ] + }, { "Type": "class Microsoft.Extensions.AI.FunctionInvocationContext", "Stage": "Stable", @@ -1425,6 +1483,28 @@ } ] }, + { + "Type": "sealed class Microsoft.Extensions.AI.OrderedFailoverChatClient : Microsoft.Extensions.AI.FailoverChatClient", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.OrderedFailoverChatClient.OrderedFailoverChatClient(System.Collections.Generic.IReadOnlyList clients, bool leaveOpen = false);", + "Stage": "Experimental" + }, + { + "Member": "override void Microsoft.Extensions.AI.OrderedFailoverChatClient.Dispose(bool disposing);", + "Stage": "Experimental" + }, + { + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.OrderedFailoverChatClient.OnRoutingUpdateAsync(Microsoft.Extensions.AI.RoutingContext context, Microsoft.Extensions.AI.FailoverChatClientAttempt? attempt, bool isTerminal, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + }, + { + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.OrderedFailoverChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + } + ] + }, { "Type": "sealed class Microsoft.Extensions.AI.RealtimeClientBuilder", "Stage": "Experimental", @@ -1525,6 +1605,46 @@ } ] }, + { + "Type": "sealed class Microsoft.Extensions.AI.SemanticRoutingChatClient : Microsoft.Extensions.AI.RoutingChatClient", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.SemanticRoutingChatClient.SemanticRoutingChatClient(Microsoft.Extensions.AI.IEmbeddingGenerator> embeddingGenerator, System.Collections.Generic.IReadOnlyDictionary> clientProfiles, Microsoft.Extensions.AI.IChatClient defaultClient, float scoreThreshold = 0.3f, bool leaveOpen = false, int topK = 1, Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation scoreAggregation = Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.Mean);", + "Stage": "Experimental" + }, + { + "Member": "override void Microsoft.Extensions.AI.SemanticRoutingChatClient.Dispose(bool disposing);", + "Stage": "Experimental" + }, + { + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.SemanticRoutingChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + } + ] + }, + { + "Type": "enum Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.ScoreAggregation();", + "Stage": "Experimental" + } + ], + "Fields": [ + { + "Member": "const Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.Mean", + "Stage": "Experimental", + "Value": "0" + }, + { + "Member": "const Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.Sum", + "Stage": "Experimental", + "Value": "1" + } + ] + }, { "Type": "sealed class Microsoft.Extensions.AI.SpeechToTextClientBuilder", "Stage": "Experimental", diff --git a/src/Shared/DiagnosticIds/DiagnosticIds.cs b/src/Shared/DiagnosticIds/DiagnosticIds.cs index abddbf61c5a..29a3d276b48 100644 --- a/src/Shared/DiagnosticIds/DiagnosticIds.cs +++ b/src/Shared/DiagnosticIds/DiagnosticIds.cs @@ -61,6 +61,7 @@ internal static class Experiments internal const string AIRealTime = AIExperiments; internal const string AIFiles = AIExperiments; internal const string AIOpenAIRequestPolicies = AIExperiments; + internal const string AIRoutingChat = AIExperiments; // These diagnostic IDs are defined by the OpenAI package for its experimental APIs. // We use the same IDs so consumers do not need to suppress additional diagnostics diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs new file mode 100644 index 00000000000..6f551ee6876 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -0,0 +1,1758 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Extensions.AI; + +public class RoutingChatClientTests +{ + [Fact] + public void RoutingContext_CarriesMutableRequestInputs() + { + var messages = new List { new(ChatRole.User, "initial") }; + var options = new ChatOptions { ModelId = "initial" }; + var context = new RoutingContext(messages, options); + var replacementMessages = new List { new(ChatRole.User, "replacement") }; + var replacementOptions = new ChatOptions { ModelId = "replacement" }; + context.Messages = replacementMessages; + context.ChatOptions = replacementOptions; + + Assert.Same(replacementMessages, context.Messages); + Assert.Same(replacementOptions, context.ChatOptions); + Assert.Throws(() => new RoutingContext(null!, options)); + Assert.Throws(() => context.Messages = null!); + } + + [Fact] + public void Create_RejectsNullSelector() + { + Assert.Throws(() => RoutingChatClient.Create(null!)); + } + + [Fact] + public async Task Create_SelectsClientForRequest() + { + var messages = new ChatMessage[] { new(ChatRole.User, "hi") }; + var options = new ChatOptions(); + using var cancellationSource = new CancellationTokenSource(); + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var selected = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + RoutingContext? observedContext = null; + CancellationToken observedToken = default; + int selectionCount = 0; + using RoutingChatClient router = RoutingChatClient.Create((context, cancellationToken) => + { + observedContext = context; + observedToken = cancellationToken; + selectionCount++; + return new(selected); + }); + + ChatResponse response = await router.GetResponseAsync(messages, options, cancellationSource.Token); + + Assert.Same(expected, response); + Assert.Same(messages, observedContext!.Messages); + Assert.Same(options, observedContext.ChatOptions); + Assert.Equal(cancellationSource.Token, observedToken); + Assert.Equal(1, selectionCount); + } + + [Fact] + public async Task Create_DoesNotReselectAfterFailure() + { + var expected = new InvalidOperationException("failed"); + using var selected = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw expected, + }; + int selectionCount = 0; + using RoutingChatClient router = RoutingChatClient.Create((_, _) => + { + selectionCount++; + return new(selected); + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(expected, actual); + Assert.Equal(1, selectionCount); + } + + [Fact] + public async Task Create_DoesNotDisposeSelectedClient() + { + using var selected = new CountingDisposeClient(); + using (RoutingChatClient router = RoutingChatClient.Create((_, _) => new(selected))) + { + _ = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + } + + Assert.Equal(0, selected.DisposeCount); + } + + [Fact] + public void GetService_ReturnsSelfAndNullForUnknownOrKeyed() + { + using var client = new DelegatingTestRouter(_ => throw new NotSupportedException()); + + Assert.Same(client, client.GetService(typeof(DelegatingTestRouter))); + Assert.Same(client, client.GetService(typeof(RoutingChatClient))); + Assert.Same(client, client.GetService(typeof(IChatClient))); + Assert.Null(client.GetService(typeof(DelegatingTestRouter), serviceKey: "key")); + Assert.Null(client.GetService(typeof(string))); + } + + [Fact] + public void MaximumAttemptsPerRequest_ValidatesValue() + { + using var client = new DelegatingFailoverTestRouter( + _ => throw new NotSupportedException()); + + Assert.Null(client.MaximumAttemptsPerRequest); + client.MaximumAttemptsPerRequest = 2; + Assert.Equal(2, client.MaximumAttemptsPerRequest); + client.MaximumAttemptsPerRequest = null; + Assert.Null(client.MaximumAttemptsPerRequest); + Assert.Throws(() => client.MaximumAttemptsPerRequest = 0); + Assert.Throws(() => client.MaximumAttemptsPerRequest = -1); + } + + [Fact] + public async Task InitialSelectionFailureReportsTerminalUpdateWithoutAttempt() + { + var expected = new InvalidOperationException("selection failed"); + RoutingContext? selectedContext = null; + RoutingContext? completedContext = null; + int completionCount = 0; + using var router = new DelegatingFailoverTestRouter( + context => + { + selectedContext = context; + throw expected; + }, + (context, attempt, isTerminal) => + { + completionCount++; + completedContext = context; + Assert.Null(attempt); + Assert.True(isTerminal); + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(expected, actual); + Assert.Same(selectedContext, completedContext); + Assert.Equal(1, completionCount); + } + + [Fact] + public async Task StreamingInitialSelectionFailureReportsTerminalUpdateWithoutAttempt() + { + var expected = new InvalidOperationException("selection failed"); + RoutingContext? selectedContext = null; + RoutingContext? completedContext = null; + int completionCount = 0; + using var router = new DelegatingFailoverTestRouter( + context => + { + selectedContext = context; + throw expected; + }, + (context, attempt, isTerminal) => + { + completionCount++; + completedContext = context; + Assert.Null(attempt); + Assert.True(isTerminal); + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Same(expected, actual); + Assert.Same(selectedContext, completedContext); + Assert.Equal(1, completionCount); + } + + [Fact] + public async Task NullSelectionReportsTerminalUpdateWithoutAttempt() + { + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => null!, + (_, attempt, isTerminal) => + { + updateCount++; + Assert.Null(attempt); + Assert.True(isTerminal); + }); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Contains("SelectClientAsync", exception.Message, StringComparison.Ordinal); + Assert.Equal(1, updateCount); + } + + [Fact] + public async Task StreamingNullSelectionReportsTerminalUpdateWithoutAttempt() + { + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => null!, + (_, attempt, isTerminal) => + { + updateCount++; + Assert.Null(attempt); + Assert.True(isTerminal); + }); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Contains("SelectClientAsync", exception.Message, StringComparison.Ordinal); + Assert.Equal(1, updateCount); + } + + [Fact] + public async Task Failover_RetrySelectionFailureReportsTerminalUpdateWithoutAttempt() + { + var invocationException = new InvalidOperationException("invocation failed"); + var selectionException = new InvalidOperationException("selection failed"); + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw invocationException, + }; + int selections = 0; + var updates = new List<(FailoverChatClientAttempt? attempt, bool isTerminal)>(); + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? failing : throw selectionException, + (_, attempt, isTerminal) => + { + updates.Add((attempt, isTerminal)); + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(selectionException, actual); + Assert.Equal(2, selections); + Assert.Collection( + updates, + update => + { + Assert.Same(failing, update.attempt!.Client); + Assert.Same(invocationException, update.attempt.Exception); + Assert.False(update.isTerminal); + }, + update => + { + Assert.Null(update.attempt); + Assert.True(update.isTerminal); + }); + } + + [Fact] + public async Task TerminalRoutingUpdateFailureReplacesRequestFailure() + { + var invocationException = new InvalidOperationException("invocation failed"); + var completionException = new InvalidOperationException("completion failed"); + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw invocationException, + }; + using var router = new DelegatingFailoverTestRouter( + _ => failing, + (_, _, isTerminal) => + { + Assert.True(isTerminal); + throw completionException; + }) + { + MaximumAttemptsPerRequest = 1, + }; + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(completionException, actual); + } + + [Fact] + public async Task NonterminalRoutingUpdateFailureStopsWithoutAnotherUpdate() + { + var invocationException = new InvalidOperationException("invocation failed"); + var updateException = new InvalidOperationException("update failed"); + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw invocationException, + }; + int selections = 0; + int updates = 0; + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return failing; + }, + (_, attempt, isTerminal) => + { + updates++; + Assert.Same(invocationException, attempt!.Exception); + Assert.False(isTerminal); + throw updateException; + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(updateException, actual); + Assert.Equal(1, selections); + Assert.Equal(1, updates); + } + + [Fact] + public async Task TerminalSelectionUpdateFailureReplacesSelectionFailure() + { + var selectionException = new InvalidOperationException("selection failed"); + var updateException = new InvalidOperationException("update failed"); + using var router = new DelegatingFailoverTestRouter( + _ => throw selectionException, + (_, attempt, isTerminal) => + { + Assert.Null(attempt); + Assert.True(isTerminal); + throw updateException; + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(updateException, actual); + } + + [Fact] + public async Task Dispatch_ConfiguredClientPreservesAndOverridesRequestOptions() + { + ChatOptions? forwarded = null; + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + forwarded = options; + return Task.FromResult(new ChatResponse()); + }, + }; + using var configured = new ConfigureOptionsChatClient( + inner, + options => options.ModelId = "route"); + var requestOptions = new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }; + using var router = new DelegatingTestRouter(_ => configured); + + _ = await router.GetResponseAsync([new(ChatRole.User, "hi")], requestOptions); + + Assert.NotSame(requestOptions, forwarded); + Assert.Equal("route", forwarded!.ModelId); + Assert.Equal("caller", forwarded.Instructions); + Assert.Equal("request", requestOptions.ModelId); + } + + [Fact] + public async Task Policy_CanMutateContextBeforeDispatch() + { + IEnumerable? forwardedMessages = null; + ChatOptions? forwardedOptions = null; + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (messages, options, _) => + { + forwardedMessages = messages; + forwardedOptions = options; + return Task.FromResult(new ChatResponse()); + }, + }; + var replacementMessages = new List { new(ChatRole.User, "replacement") }; + var replacementOptions = new ChatOptions { ModelId = "replacement" }; + using var router = new DelegatingTestRouter(context => + { + context.Messages = replacementMessages; + context.ChatOptions = replacementOptions; + return inner; + }); + + _ = await router.GetResponseAsync([new(ChatRole.User, "original")], new ChatOptions()); + + Assert.Same(replacementMessages, forwardedMessages); + Assert.Same(replacementOptions, forwardedOptions); + } + + [Fact] + public async Task Failure_Propagates() + { + var expected = new InvalidOperationException("failed"); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw expected, + }; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observed = attempt; + }) + { + MaximumAttemptsPerRequest = 1, + }; + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(expected, actual); + Assert.Same(inner, observed!.Client); + Assert.Same(expected, observed.Exception); + Assert.False(observed.ResponseCompleted); + } + + [Fact] + public async Task Failover_UpdateChangesStateBeforeNextSelection() + { + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var working = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + RoutingContext? initialContext = null; + FailoverChatClientAttempt? failedAttempt = null; + FailoverChatClientAttempt? terminalAttempt = null; + using var router = new DelegatingFailoverTestRouter( + context => + { + initialContext ??= context; + Assert.Same(initialContext, context); + return failedAttempt is null ? failing : working; + }, + (context, attempt, isTerminal) => + { + Assert.Same(initialContext, context); + if (isTerminal) + { + terminalAttempt = attempt; + } + else + { + failedAttempt = attempt; + } + }); + + ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + Assert.Same(failing, failedAttempt!.Client); + Assert.Equal("failed", Assert.IsType(failedAttempt.Exception).Message); + Assert.True(failedAttempt.Duration >= TimeSpan.Zero); + Assert.Null(failedAttempt.TimeToFirstUpdate); + Assert.False(failedAttempt.OutputCommitted); + Assert.False(failedAttempt.ResponseCompleted); + Assert.Same(working, terminalAttempt!.Client); + Assert.Null(terminalAttempt.Exception); + Assert.True(terminalAttempt.ResponseCompleted); + } + + [Fact] + public async Task Failover_CanRetrySameClient() + { + int calls = 0; + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + ++calls == 1 + ? throw new InvalidOperationException("transient") + : Task.FromResult(new ChatResponse()), + }; + using var router = new DelegatingFailoverTestRouter( + _ => inner); + + _ = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Equal(2, calls); + } + + [Fact] + public async Task MaximumAttemptsPerRequest_AllowsSuccessAtLimit() + { + int calls = 0; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + ++calls == 1 + ? throw new InvalidOperationException("transient") + : Task.FromResult(expected), + }; + using var router = new DelegatingFailoverTestRouter( + _ => inner) + { + MaximumAttemptsPerRequest = 2, + }; + + ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + Assert.Equal(2, calls); + } + + [Fact] + public async Task MaximumAttemptsPerRequest_RethrowsLastFailure() + { + int calls = 0; + var terminalUpdates = new List(); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + throw new InvalidOperationException($"failure {++calls}"), + }; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.NotNull(attempt); + terminalUpdates.Add(isTerminal); + }) + { + MaximumAttemptsPerRequest = 2, + }; + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("failure 2", exception.Message); + Assert.Equal(2, calls); + Assert.Equal([false, true], terminalUpdates); + } + + [Fact] + public async Task Failover_NullResponseIsReturnedWithoutFailover() + { + using var nullClient = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(null!), + }; + int selections = 0; + FailoverChatClientAttempt? terminalAttempt = null; + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return nullClient; + }, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + terminalAttempt = attempt; + }); + + ChatResponse? response = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Null(response); + Assert.Equal(1, selections); + Assert.Null(terminalAttempt!.Exception); + Assert.True(terminalAttempt!.ResponseCompleted); + } + + [Fact] + public async Task Cancellation_DoesNotReselect() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + int selections = 0; + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, cancellationToken) => + throw new OperationCanceledException(cancellationToken), + }; + using var router = new DelegatingTestRouter(_ => + { + selections++; + return inner; + }); + + await Assert.ThrowsAnyAsync( + () => router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.Equal(1, selections); + } + + [Fact] + public async Task Failover_CancellationReportsTerminalUpdate() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, cancellationToken) => + throw new OperationCanceledException(cancellationToken), + }; + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + updateCount++; + Assert.IsAssignableFrom(attempt!.Exception); + Assert.True(isTerminal); + }); + + await Assert.ThrowsAnyAsync( + () => router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.Equal(1, updateCount); + } + + [Fact] + public async Task Streaming_PreOutputFailurePropagates() + { + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observed = attempt; + }) + { + MaximumAttemptsPerRequest = 1, + }; + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Equal("failed", exception.Message); + Assert.Same(exception, observed!.Exception); + Assert.False(observed.OutputCommitted); + Assert.False(observed.ResponseCompleted); + } + + [Fact] + public async Task Streaming_FallsBackBeforeFirstUpdate() + { + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? failing : working); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal(2, selections); + } + + [Fact] + public async Task Streaming_CancellationDoesNotSelectNext() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, cancellationToken) => + CanceledStream(cancellationToken), + }; + int selections = 0; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observed = attempt; + }); + + await Assert.ThrowsAnyAsync( + () => CollectAsync( + router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token))); + + Assert.Equal(1, selections); + Assert.IsAssignableFrom(observed!.Exception); + Assert.False(observed.OutputCommitted); + Assert.False(observed.ResponseCompleted); + } + + [Fact] + public async Task Streaming_MaximumAttemptsPerRequest_RethrowsLastFailure() + { + int calls = 0; + var terminalUpdates = new List(); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream($"failure {++calls}"), + }; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.NotNull(attempt); + terminalUpdates.Add(isTerminal); + }) + { + MaximumAttemptsPerRequest = 2, + }; + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Equal("failure 2", exception.Message); + Assert.Equal(2, calls); + Assert.Equal([false, true], terminalUpdates); + } + + [Fact] + public async Task Streaming_StreamCreationFailureFallsBack() + { + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? failing : working); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal(2, selections); + } + + [Fact] + public async Task Streaming_EnumeratorCreationFailureFallsBack() + { + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => + new ThrowingGetAsyncEnumeratorEnumerable(new InvalidOperationException("failed")), + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? failing : working); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal(2, selections); + } + + [Fact] + public async Task Streaming_EmptyStreamDisposalFailureFallsBack() + { + var disposalException = new InvalidOperationException("dispose failed"); + var failedStream = new TrackingAsyncEnumerable([], disposeException: disposalException); + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => failedStream, + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + FailoverChatClientAttempt? failedAttempt = null; + FailoverChatClientAttempt? terminalAttempt = null; + using var router = new DelegatingFailoverTestRouter( + _ => failedAttempt is null ? failing : working, + (_, attempt, isTerminal) => + { + if (isTerminal) + { + terminalAttempt = attempt; + } + else + { + failedAttempt = attempt; + } + }); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Same(disposalException, failedAttempt!.Exception); + Assert.False(failedAttempt.OutputCommitted); + Assert.False(failedAttempt.ResponseCompleted); + Assert.True(terminalAttempt!.ResponseCompleted); + Assert.Equal(1, failedStream.DisposeCount); + } + + [Fact] + public async Task Streaming_MidStreamFailureIsObservedAndDoesNotReselect() + { + var stream = new TrackingAsyncEnumerable( + [new ChatResponseUpdate(ChatRole.Assistant, "first")], + throwOnMove: 2, + exception: new InvalidOperationException("mid-stream")); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => stream, + }; + int selections = 0; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observed = attempt; + }); + var updates = new List(); + + async Task ConsumeAsync() + { + await foreach (ChatResponseUpdate update in router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])) + { + updates.Add(update); + } + } + + InvalidOperationException exception = await Assert.ThrowsAsync(ConsumeAsync); + + Assert.Equal("mid-stream", exception.Message); + Assert.Equal("first", Assert.Single(updates).Text); + Assert.Equal(1, selections); + Assert.Same(exception, observed!.Exception); + Assert.True(observed.OutputCommitted); + Assert.False(observed.ResponseCompleted); + Assert.NotNull(observed.TimeToFirstUpdate); + Assert.True(observed.Duration >= observed.TimeToFirstUpdate.Value); + Assert.Equal(1, stream.DisposeCount); + } + + [Fact] + public async Task Streaming_CompletionNotifiesHook() + { + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("a", "b"), + }; + int completions = 0; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + completions++; + observed = attempt; + }); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal(2, updates.Count); + Assert.Equal(1, completions); + Assert.Null(observed!.Exception); + Assert.True(observed.OutputCommitted); + Assert.True(observed.ResponseCompleted); + Assert.NotNull(observed.TimeToFirstUpdate); + } + + [Fact] + public async Task Streaming_EmptyCompletionNotifiesHook() + { + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates(), + }; + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + updateCount++; + Assert.NotNull(attempt); + Assert.Null(attempt.Exception); + Assert.False(attempt.OutputCommitted); + Assert.True(attempt.ResponseCompleted); + Assert.True(isTerminal); + }); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Empty(updates); + Assert.Equal(1, updateCount); + } + + [Fact] + public async Task Streaming_CallerStopsEarlyNotifiesHookAsIncomplete() + { + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("a", "b"), + }; + int completions = 0; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + completions++; + observed = attempt; + }); + + await ConsumeOneAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal(1, completions); + Assert.Null(observed!.Exception); + Assert.True(observed.OutputCommitted); + Assert.False(observed.ResponseCompleted); + } + + [Fact] + public void SemanticRouting_RejectsInvalidConfiguration() + { + using var client = new TestChatClient(); + using var generator = new TestEmbeddingGenerator(); + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [client] = ["profile"], + }; + + Assert.Throws(() => + new SemanticRoutingChatClient(null!, profiles, client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, null!, client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, null!)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary>(), + client)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary> + { + [client] = [], + }, + client)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary> + { + [client] = [" "], + }, + client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, client, scoreThreshold: 1.1f)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, client, topK: 0)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + profiles, + client, + scoreAggregation: (SemanticRoutingChatClient.ScoreAggregation)(-1))); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + profiles, + client, + scoreThreshold: 2.1f, + topK: 2, + scoreAggregation: SemanticRoutingChatClient.ScoreAggregation.Sum)); + } + + [Fact] + public async Task SemanticRouting_SelectsBestProfileAndCachesIndex() + { + var vectors = new Dictionary + { + ["code profile"] = [1, 0], + ["writing profile"] = [0, 1], + ["debug this code"] = [1, 0], + }; + int profileBatches = 0; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + { + string[] inputs = [.. values]; + if (inputs.Length > 1) + { + profileBatches++; + } + + return Task.FromResult(new GeneratedEmbeddings>( + [.. inputs.Select(input => new Embedding(vectors[input]))])); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "code")); + using var code = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + using var writing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "writing"))), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [code] = ["code profile"], + [writing] = ["writing profile"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient: writing, + leaveOpen: true); + + ChatResponse first = await router.GetResponseAsync([new(ChatRole.User, "debug this code")]); + ChatResponse second = await router.GetResponseAsync([new(ChatRole.User, "debug this code")]); + + Assert.Same(expected, first); + Assert.Same(expected, second); + Assert.Equal(1, profileBatches); + } + + [Theory] + [InlineData(SemanticRoutingChatClient.ScoreAggregation.Mean, "code")] + [InlineData(SemanticRoutingChatClient.ScoreAggregation.Sum, "writing")] + public async Task SemanticRouting_AggregatesGlobalTopKByClient( + SemanticRoutingChatClient.ScoreAggregation scoreAggregation, + string expectedResponse) + { + var vectors = new Dictionary + { + ["code"] = [1, 0], + ["writing one"] = [0.8f, 0.6f], + ["writing two"] = [0.8f, -0.6f], + ["query"] = [1, 0], + }; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + Task.FromResult(new GeneratedEmbeddings>( + [.. values.Select(input => new Embedding(vectors[input]))])), + }; + using var code = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "code"))), + }; + using var writing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "writing"))), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [code] = ["code"], + [writing] = ["writing one", "writing two"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient: code, + scoreThreshold: scoreAggregation == SemanticRoutingChatClient.ScoreAggregation.Sum ? 1.5f : 0.3f, + leaveOpen: true, + topK: 3, + scoreAggregation: scoreAggregation); + + ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "query")]); + + Assert.Equal(expectedResponse, response.Text); + } + + [Fact] + public async Task SemanticRouting_UsesDefaultBelowThreshold() + { + var vectors = new Dictionary + { + ["code profile"] = [1, 0], + ["unrelated query"] = [0, 1], + }; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + Task.FromResult(new GeneratedEmbeddings>( + [.. values.Select(input => new Embedding(vectors[input]))])), + }; + using var profiled = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "profiled"))), + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "default")); + using var defaultClient = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [profiled] = ["code profile"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient, + scoreThreshold: 0.5f, + leaveOpen: true); + + ChatResponse response = + await router.GetResponseAsync([new(ChatRole.User, "unrelated query")]); + + Assert.Same(expected, response); + } + + [Fact] + public void SemanticRouting_DisposesOwnedResourcesOnce() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var generator = new CountingEmbeddingGenerator(); + var profiled = new CountingDisposeClient(); + var defaultClient = new CountingDisposeClient(); + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [profiled] = ["profile"], + }; + var router = new SemanticRoutingChatClient(generator, profiles, defaultClient); +#pragma warning restore CA2000 + + router.Dispose(); + router.Dispose(); + + Assert.Equal(1, generator.DisposeCount); + Assert.Equal(1, profiled.DisposeCount); + Assert.Equal(1, defaultClient.DisposeCount); + } + + [Fact] + public void OrderedFailover_RejectsMissingClients() + { + using var inner = new TestChatClient(); + + Assert.Throws(() => new OrderedFailoverChatClient(null!)); + Assert.Throws(() => new OrderedFailoverChatClient([])); + Assert.Throws(() => new OrderedFailoverChatClient([inner, null!])); + Assert.Throws(() => new OrderedFailoverChatClient([inner, inner])); + } + + [Fact] + public async Task OrderedFailover_TriesClientsInOrderAndSkipsFailures() + { + var calls = new List(); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls.Add("first"); + throw new InvalidOperationException("first failed"); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls.Add("second"); + return Task.FromResult(expected); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + Assert.Equal(["first", "second"], calls); + } + + [Fact] + public async Task OrderedFailover_DistinguishesValueEqualClients() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var first = new ValueEqualChatClient( + () => throw new InvalidOperationException("failed")); + using var second = new ValueEqualChatClient( + () => Task.FromResult(expected)); + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + [Fact] + public async Task OrderedFailover_ExhaustionRethrowsLastFailure() + { + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("first"), + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("second"), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => client.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("second", exception.Message); + } + + [Fact] + public async Task OrderedFailover_StateIsScopedToOneRequest() + { + int firstCalls = 0; + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + firstCalls++; + throw new InvalidOperationException("failed"); + }, + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(new ChatResponse()), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + _ = await client.GetResponseAsync([new(ChatRole.User, "one")]); + _ = await client.GetResponseAsync([new(ChatRole.User, "two")]); + + Assert.Equal(2, firstCalls); + } + + [Fact] + public async Task OrderedFailover_ConcurrentRequestsHaveIndependentState() + { + const int RequestCount = 10; + int firstCalls = 0; + int secondCalls = 0; + var allStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var first = new TestChatClient + { + GetResponseAsyncCallback = async (_, _, _) => + { + if (Interlocked.Increment(ref firstCalls) == RequestCount) + { + allStarted.SetResult(true); + } + + await allStarted.Task; + throw new InvalidOperationException("failed"); + }, + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + Interlocked.Increment(ref secondCalls); + return Task.FromResult(new ChatResponse()); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + Task[] requests = + [ + .. Enumerable.Range(0, RequestCount).Select(index => + client.GetResponseAsync([new(ChatRole.User, index.ToString())])), + ]; + _ = await Task.WhenAll(requests); + + Assert.Equal(RequestCount, firstCalls); + Assert.Equal(RequestCount, secondCalls); + } + + [Fact] + public async Task OrderedFailover_DoesNotRetainStateWhileStreaming() + { + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("first", "second"), + }; + 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); + Assert.Equal(0, GetOrderedFailoverRequestStateCount(client)); + } + + [Fact] + public async Task OrderedFailover_SnapshotsConfiguredClients() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + var clients = new List { inner }; + using var failover = new OrderedFailoverChatClient(clients, leaveOpen: true); + clients.Clear(); + + ChatResponse response = await failover.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + [Fact] + public async Task OrderedFailover_UsesEachConfiguredClient() + { + var seenOptions = new List<(string? modelId, string? instructions)>(); + using var firstInner = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + seenOptions.Add((options?.ModelId, options?.Instructions)); + throw new InvalidOperationException("failed"); + }, + }; + using var first = new ConfigureOptionsChatClient( + firstInner, + options => options.ModelId = "first"); + using var secondInner = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + seenOptions.Add((options?.ModelId, options?.Instructions)); + return Task.FromResult(new ChatResponse()); + }, + }; + using var second = new ConfigureOptionsChatClient( + secondInner, + options => options.ModelId = "second"); + using var client = new OrderedFailoverChatClient( + [first, second], + leaveOpen: true); + + _ = await client.GetResponseAsync( + [new(ChatRole.User, "hi")], + new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }); + + Assert.Equal( + [("first", "caller"), ("second", "caller")], + seenOptions); + } + + [Fact] + public async Task OrderedFailover_StreamingFallsBackBeforeFirstUpdate() + { + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + List updates = + await CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("ok", Assert.Single(updates).Text); + } + + [Fact] + public async Task OrderedFailover_CancellationDoesNotFallback() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + bool secondCalled = false; + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, cancellationToken) => + throw new OperationCanceledException(cancellationToken), + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return Task.FromResult(new ChatResponse()); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + await Assert.ThrowsAnyAsync( + () => client.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.False(secondCalled); + } + + [Fact] + public void OrderedFailover_DisposesEachClientOnce() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var shared = new CountingDisposeClient(); + var other = new CountingDisposeClient(); + var client = new OrderedFailoverChatClient([shared, other]); +#pragma warning restore CA2000 + + client.Dispose(); + client.Dispose(); + + Assert.Equal(1, shared.DisposeCount); + Assert.Equal(1, other.DisposeCount); + } + + [Fact] + public void OrderedFailover_LeaveOpenDoesNotDisposeInnerClients() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var inner = new CountingDisposeClient(); + var client = new OrderedFailoverChatClient([inner], leaveOpen: true); +#pragma warning restore CA2000 + + client.Dispose(); + + Assert.Equal(0, inner.DisposeCount); + inner.Dispose(); + } + + [Fact] + public async Task NestedRouters_ReturnLeafResponse() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var leaf = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + using var inner = new OrderedFailoverChatClient([leaf]); + using var outer = new OrderedFailoverChatClient([inner]); + + ChatResponse response = await outer.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + private static async Task> CollectAsync( + IAsyncEnumerable updates) + { + var result = new List(); + await foreach (ChatResponseUpdate update in updates) + { + result.Add(update); + } + + return result; + } + + private static async Task ConsumeOneAsync(IAsyncEnumerable updates) + { + await using IAsyncEnumerator enumerator = updates.GetAsyncEnumerator(); + Assert.True(await enumerator.MoveNextAsync()); + } + + private static int GetOrderedFailoverRequestStateCount(OrderedFailoverChatClient client) + { + object requestStates = typeof(OrderedFailoverChatClient) + .GetField("_requestStates", System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic)! + .GetValue(client)!; + return ((System.Collections.IDictionary)requestStates).Count; + } + + 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 IAsyncEnumerable ThrowingStream(string message) + { + await Task.Yield(); + foreach (int _ in Array.Empty()) + { + yield return new ChatResponseUpdate(ChatRole.Assistant, "never"); + } + + throw new InvalidOperationException(message); + } + + private static async IAsyncEnumerable CanceledStream( + [EnumeratorCancellation] CancellationToken cancellationToken) + { + await Task.Yield(); + foreach (int _ in Array.Empty()) + { + yield return new ChatResponseUpdate(ChatRole.Assistant, "never"); + } + + throw new OperationCanceledException(cancellationToken); + } + + private sealed class DelegatingTestRouter : RoutingChatClient + { + private readonly Func _select; + + public DelegatingTestRouter(Func select) + { + _select = select; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + new(_select(context)); + } + + private sealed class DelegatingFailoverTestRouter : FailoverChatClient + { + private readonly Func _select; + private readonly Action? _onRoutingUpdate; + + public DelegatingFailoverTestRouter( + Func select, + Action? onRoutingUpdate = null) + { + _select = select; + _onRoutingUpdate = onRoutingUpdate; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + new(_select(context)); + + protected override ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt? attempt, + bool isTerminal, + CancellationToken cancellationToken) + { + _onRoutingUpdate?.Invoke(context, attempt, isTerminal); + return default; + } + } + + private sealed class ThrowingGetAsyncEnumeratorEnumerable : IAsyncEnumerable + { + private readonly Exception _exception; + + public ThrowingGetAsyncEnumeratorEnumerable(Exception exception) + { + _exception = exception; + } + + public IAsyncEnumerator GetAsyncEnumerator( + CancellationToken cancellationToken = default) => + throw _exception; + } + + private sealed class TrackingAsyncEnumerable : IAsyncEnumerable + { + private readonly IReadOnlyList _updates; + private readonly int? _throwOnMove; + private readonly Exception? _exception; + private readonly Exception? _disposeException; + + public TrackingAsyncEnumerable( + IReadOnlyList updates, + int? throwOnMove = null, + Exception? exception = null, + Exception? disposeException = null) + { + _updates = updates; + _throwOnMove = throwOnMove; + _exception = exception; + _disposeException = disposeException; + } + + public int DisposeCount { get; private set; } + + public IAsyncEnumerator GetAsyncEnumerator( + CancellationToken cancellationToken = default) => + new Enumerator(this); + + private sealed class Enumerator : IAsyncEnumerator + { + private readonly TrackingAsyncEnumerable _owner; + private int _moveCount; + + public Enumerator(TrackingAsyncEnumerable owner) + { + _owner = owner; + } + + public ChatResponseUpdate Current { get; private set; } = null!; + + public ValueTask MoveNextAsync() + { + _moveCount++; + if (_moveCount == _owner._throwOnMove) + { + throw _owner._exception!; + } + + int index = _moveCount - 1; + if (index >= _owner._updates.Count) + { + return new(false); + } + + Current = _owner._updates[index]; + return new(true); + } + + public ValueTask DisposeAsync() + { + _owner.DisposeCount++; + if (_owner._disposeException is { } exception) + { + throw exception; + } + + return default; + } + } + } + + private sealed class CountingDisposeClient : IChatClient + { + public int DisposeCount { get; private set; } + + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) => + Task.FromResult(new ChatResponse()); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + + public override bool Equals(object? obj) => obj is CountingDisposeClient; + + public override int GetHashCode() => 0; + } + + private sealed class CountingEmbeddingGenerator : + IEmbeddingGenerator> + { + public int DisposeCount { get; private set; } + + public Task>> GenerateAsync( + IEnumerable values, + EmbeddingGenerationOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + } + + private sealed class ChatClientReferenceComparer : IEqualityComparer + { + public static ChatClientReferenceComparer Instance { get; } = new(); + + public bool Equals(IChatClient? x, IChatClient? y) => ReferenceEquals(x, y); + + public int GetHashCode(IChatClient obj) => RuntimeHelpers.GetHashCode(obj); + } + + private sealed class ValueEqualChatClient(Func> getResponse) : IChatClient + { + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) => + getResponse(); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() + { + } + + public override bool Equals(object? obj) => obj is ValueEqualChatClient; + + public override int GetHashCode() => 0; + } +} From 665abe43eb163c102b0e29fd7759a9d345e678e7 Mon Sep 17 00:00:00 2001 From: Joshua Yue <109126602+joshuajyue@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:43:01 -0700 Subject: [PATCH 02/15] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../ChatRouting/RoutingChatClient.cs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index 5b2acfc830e..840681c0fac 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -50,7 +50,8 @@ public virtual async Task GetResponseAsync( CancellationToken cancellationToken = default) { var context = new RoutingContext(messages, options); - IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false); + IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); return await client.GetResponseAsync( context.Messages, context.ChatOptions, From f72f3d7f77eb10e655312f6ccdc1f2cc5787ffe4 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Mon, 27 Jul 2026 15:49:01 -0700 Subject: [PATCH 03/15] Address routing review feedback Treat provider-thrown cancellation as terminal, materialize semantic routing messages before inspection, and enforce non-null selection consistently for streaming. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 3 +- .../ChatRouting/FailoverChatClient.cs | 9 +- .../ChatRouting/SemanticRoutingChatClient.cs | 8 +- .../ChatRouting/RoutingChatClientTests.cs | 147 ++++++++++++++++++ 4 files changed, 162 insertions(+), 5 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index 840681c0fac..990a84095f5 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -65,7 +65,8 @@ public virtual async IAsyncEnumerable GetStreamingResponseAs [EnumeratorCancellation] CancellationToken cancellationToken = default) { var context = new RoutingContext(messages, options); - IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false); + IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); await foreach (ChatResponseUpdate update in client.GetStreamingResponseAsync(context.Messages, context.ChatOptions, cancellationToken) .WithCancellation(cancellationToken) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index 4552c89826f..9e000dc57f9 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -144,7 +144,7 @@ public sealed override async Task GetResponseAsync( outputCommitted: false); bool isTerminal = exception is null || - cancellationToken.IsCancellationRequested || + IsCancellation(exception, cancellationToken) || (maximumAttempts is int limit && attemptCount >= limit); await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken); @@ -215,7 +215,7 @@ public sealed override async IAsyncEnumerable GetStreamingRe timeToFirstUpdate: null, responseCompleted: false, outputCommitted: false); - bool isTerminal = cancellationToken.IsCancellationRequested || reachedAttemptLimit; + bool isTerminal = IsCancellation(exception, cancellationToken) || reachedAttemptLimit; await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken); @@ -277,7 +277,7 @@ public sealed override async IAsyncEnumerable GetStreamingRe isTerminalAttempt = attempt.ResponseCompleted || outputCommitted || - cancellationToken.IsCancellationRequested || + IsCancellation(terminalException, cancellationToken) || reachedAttemptLimit; await OnRoutingUpdateAsync(context, attempt, isTerminalAttempt, cancellationToken); @@ -327,6 +327,9 @@ private static TimeSpan GetElapsedTime(long startingTimestamp) => ((double)TimeSpan.TicksPerSecond / Stopwatch.Frequency))); #endif + private static bool IsCancellation(Exception? exception, CancellationToken cancellationToken) => + exception is OperationCanceledException || cancellationToken.IsCancellationRequested; + [DoesNotReturn] private static void Rethrow(Exception exception) { diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs index d6ae61fad65..8abae44540c 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -165,7 +165,13 @@ protected override async ValueTask SelectClientAsync( CancellationToken cancellationToken) { _ = Throw.IfNull(context); - string? query = LastUserText(context.Messages); + IEnumerable messages = context.Messages; + if (messages is not IReadOnlyList) + { + context.Messages = messages = messages.ToArray(); + } + + string? query = LastUserText(messages); if (string.IsNullOrWhiteSpace(query)) { return _clients[0]; diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs index 6f551ee6876..f78c6c710e4 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -89,6 +89,20 @@ public async Task Create_DoesNotReselectAfterFailure() Assert.Equal(1, selectionCount); } + [Fact] + public async Task Create_NullSelectionThrowsForNonStreamingAndStreaming() + { + using RoutingChatClient router = RoutingChatClient.Create((_, _) => new((IChatClient)null!)); + + InvalidOperationException nonStreaming = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + InvalidOperationException streaming = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Contains("SelectClientAsync", nonStreaming.Message, StringComparison.Ordinal); + Assert.Contains("SelectClientAsync", streaming.Message, StringComparison.Ordinal); + } + [Fact] public async Task Create_DoesNotDisposeSelectedClient() { @@ -720,6 +734,82 @@ await Assert.ThrowsAnyAsync( Assert.False(observed.ResponseCompleted); } + [Fact] + public async Task OperationCanceledException_DoesNotTriggerFailover() + { + bool secondCalled = false; + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new OperationCanceledException(), + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return Task.FromResult(new ChatResponse()); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + await Assert.ThrowsAnyAsync( + () => client.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.False(secondCalled); + } + + [Fact] + public async Task Streaming_OperationCanceledExceptionDoesNotTriggerFailover() + { + bool secondCalled = false; + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => CanceledStream(CancellationToken.None), + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return YieldUpdates("ok"); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + await Assert.ThrowsAnyAsync( + () => CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.False(secondCalled); + } + + [Fact] + public async Task Streaming_DisposalCancellationDoesNotTriggerFailover() + { + bool secondCalled = false; + var canceledStream = new TrackingAsyncEnumerable( + [], + disposeException: new OperationCanceledException()); + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => canceledStream, + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return YieldUpdates("ok"); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + await Assert.ThrowsAnyAsync( + () => CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.False(secondCalled); + Assert.Equal(1, canceledStream.DisposeCount); + } + [Fact] public async Task Streaming_MaximumAttemptsPerRequest_RethrowsLastFailure() { @@ -1072,6 +1162,46 @@ public async Task SemanticRouting_SelectsBestProfileAndCachesIndex() Assert.Equal(1, profileBatches); } + [Fact] + public async Task SemanticRouting_MaterializesMessagesBeforeSelection() + { + var vectors = new Dictionary + { + ["code profile"] = [1, 0], + ["debug this code"] = [1, 0], + }; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + Task.FromResult(new GeneratedEmbeddings>( + [.. values.Select(input => new Embedding(vectors[input]))])), + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "code")); + using var selected = new TestChatClient + { + GetResponseAsyncCallback = (messages, _, _) => + { + Assert.Equal("debug this code", Assert.Single(messages).Text); + return Task.FromResult(expected); + }, + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [selected] = ["code profile"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient: selected, + leaveOpen: true); + var messages = new SingleUseMessageEnumerable([new(ChatRole.User, "debug this code")]); + + ChatResponse response = await router.GetResponseAsync(messages); + + Assert.Same(expected, response); + Assert.Equal(1, messages.EnumerationCount); + } + [Theory] [InlineData(SemanticRoutingChatClient.ScoreAggregation.Mean, "code")] [InlineData(SemanticRoutingChatClient.ScoreAggregation.Sum, "writing")] @@ -1570,6 +1700,23 @@ protected override ValueTask SelectClientAsync( new(_select(context)); } + private sealed class SingleUseMessageEnumerable(IEnumerable messages) : IEnumerable + { + public int EnumerationCount { get; private set; } + + public IEnumerator GetEnumerator() + { + if (++EnumerationCount > 1) + { + throw new InvalidOperationException("The messages can only be enumerated once."); + } + + return messages.GetEnumerator(); + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() => GetEnumerator(); + } + private sealed class DelegatingFailoverTestRouter : FailoverChatClient { private readonly Func _select; From 26d7a5c9a4bd47c6357d84d7c2f7bdff1fe71221 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Mon, 27 Jul 2026 15:54:57 -0700 Subject: [PATCH 04/15] Credit semantic-router inspiration Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/SemanticRoutingChatClient.cs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs index 8abae44540c..bd3fcf00b69 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -25,6 +25,10 @@ namespace Microsoft.Extensions.AI; /// The configured clients are used as stable routing identities. By default this instance owns the clients and /// embedding generator and disposes them when it is disposed. /// +/// +/// The example-utterance routing approach is inspired by +/// Aurelio Labs' semantic-router project. +/// /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] public sealed class SemanticRoutingChatClient : RoutingChatClient From bcff6d39995e888f82b1c29e6f1a124c63656bae Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Tue, 28 Jul 2026 08:51:51 -0700 Subject: [PATCH 05/15] Handle streaming Current getter failures Capture Current access failures before committing output so pre-output failover and attempt reporting remain accurate. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/FailoverChatClient.cs | 38 ++++++++++-- .../ChatRouting/RoutingChatClientTests.cs | 59 +++++++++++++++++++ 2 files changed, 91 insertions(+), 6 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index 9e000dc57f9..b12fdc082c3 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -228,11 +228,6 @@ public sealed override async IAsyncEnumerable GetStreamingRe } activeDuration += GetElapsedTime(operationStart); - if (hasCurrent) - { - timeToFirstUpdate = activeDuration; - } - bool responseCompleted = false; bool outputCommitted = false; bool isTerminalAttempt = false; @@ -242,8 +237,20 @@ public sealed override async IAsyncEnumerable GetStreamingRe { while (hasCurrent) { + operationStart = Stopwatch.GetTimestamp(); + bool hasCurrentValue = TryGetCurrent( + enumerator, + out ChatResponseUpdate current, + out terminalException); + activeDuration += GetElapsedTime(operationStart); + if (!hasCurrentValue) + { + break; + } + + timeToFirstUpdate ??= activeDuration; outputCommitted = true; - yield return enumerator.Current; + yield return current; operationStart = Stopwatch.GetTimestamp(); try @@ -330,6 +337,25 @@ private static TimeSpan GetElapsedTime(long startingTimestamp) => private static bool IsCancellation(Exception? exception, CancellationToken cancellationToken) => exception is OperationCanceledException || cancellationToken.IsCancellationRequested; + private static bool TryGetCurrent( + IAsyncEnumerator enumerator, + out ChatResponseUpdate current, + out Exception? exception) + { + try + { + current = enumerator.Current; + exception = null; + return true; + } + catch (Exception ex) + { + current = null!; + exception = ex; + return false; + } + } + [DoesNotReturn] private static void Rethrow(Exception exception) { diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs index f78c6c710e4..4bd4751f344 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -883,6 +883,41 @@ public async Task Streaming_EnumeratorCreationFailureFallsBack() Assert.Equal(2, selections); } + [Fact] + public async Task Streaming_CurrentFailureFallsBackBeforeOutput() + { + var currentException = new InvalidOperationException("current failed"); + var failedStream = new ThrowingCurrentAsyncEnumerable(currentException); + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => failedStream, + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + FailoverChatClientAttempt? failedAttempt = null; + using var router = new DelegatingFailoverTestRouter( + _ => failedAttempt is null ? failing : working, + (_, attempt, isTerminal) => + { + if (!isTerminal) + { + failedAttempt = attempt; + } + }); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Same(currentException, failedAttempt!.Exception); + Assert.False(failedAttempt.OutputCommitted); + Assert.False(failedAttempt.ResponseCompleted); + Assert.Null(failedAttempt.TimeToFirstUpdate); + Assert.Equal(1, failedStream.DisposeCount); + } + [Fact] public async Task Streaming_EmptyStreamDisposalFailureFallsBack() { @@ -1760,6 +1795,30 @@ public IAsyncEnumerator GetAsyncEnumerator( throw _exception; } + private sealed class ThrowingCurrentAsyncEnumerable(Exception exception) : IAsyncEnumerable + { + public int DisposeCount { get; private set; } + + public IAsyncEnumerator GetAsyncEnumerator( + CancellationToken cancellationToken = default) => + new Enumerator(this, exception); + + private sealed class Enumerator( + ThrowingCurrentAsyncEnumerable owner, + Exception exception) : IAsyncEnumerator + { + public ChatResponseUpdate Current => throw exception; + + public ValueTask DisposeAsync() + { + owner.DisposeCount++; + return default; + } + + public ValueTask MoveNextAsync() => new(true); + } + } + private sealed class TrackingAsyncEnumerable : IAsyncEnumerable { private readonly IReadOnlyList _updates; From acc450e550e56d0ecdba054c2c256451b6e1d1f1 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Tue, 28 Jul 2026 14:22:48 -0700 Subject: [PATCH 06/15] Avoid context capture in failover routing Apply ConfigureAwait(false) consistently across selection, invocation, streaming, disposal, and routing update awaits. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/FailoverChatClient.cs | 42 +++++++++++++------ 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index b12fdc082c3..285a9ec8c9f 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -109,12 +109,16 @@ public sealed override async Task GetResponseAsync( IChatClient selectedClient; try { - selectedClient = await SelectClientAsync(context, cancellationToken) ?? + selectedClient = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); } catch (Exception) { - await OnRoutingUpdateAsync(context, attempt: null, isTerminal: true, cancellationToken); + await OnRoutingUpdateAsync( + context, + attempt: null, + isTerminal: true, + cancellationToken).ConfigureAwait(false); throw; } @@ -128,7 +132,7 @@ public sealed override async Task GetResponseAsync( response = await selectedClient.GetResponseAsync( context.Messages, context.ChatOptions, - cancellationToken); + cancellationToken).ConfigureAwait(false); } catch (Exception ex) { @@ -147,7 +151,7 @@ exception is null || IsCancellation(exception, cancellationToken) || (maximumAttempts is int limit && attemptCount >= limit); - await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken); + await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken).ConfigureAwait(false); if (exception is null) { @@ -176,12 +180,16 @@ public sealed override async IAsyncEnumerable GetStreamingRe IChatClient selectedClient; try { - selectedClient = await SelectClientAsync(context, cancellationToken) ?? + selectedClient = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); } catch (Exception) { - await OnRoutingUpdateAsync(context, attempt: null, isTerminal: true, cancellationToken); + await OnRoutingUpdateAsync( + context, + attempt: null, + isTerminal: true, + cancellationToken).ConfigureAwait(false); throw; } @@ -202,12 +210,13 @@ public sealed override async IAsyncEnumerable GetStreamingRe cancellationToken) .GetAsyncEnumerator(cancellationToken); - hasCurrent = await enumerator.MoveNextAsync(); + hasCurrent = await enumerator.MoveNextAsync().ConfigureAwait(false); } catch (Exception ex) { activeDuration += GetElapsedTime(operationStart); - Exception exception = (await DisposeAsync(enumerator, ex, cancellationToken))!; + Exception exception = + (await DisposeAsync(enumerator, ex, cancellationToken).ConfigureAwait(false))!; var attempt = new FailoverChatClientAttempt( selectedClient, exception, @@ -217,7 +226,7 @@ public sealed override async IAsyncEnumerable GetStreamingRe outputCommitted: false); bool isTerminal = IsCancellation(exception, cancellationToken) || reachedAttemptLimit; - await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken); + await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken).ConfigureAwait(false); if (isTerminal) { @@ -255,7 +264,7 @@ public sealed override async IAsyncEnumerable GetStreamingRe operationStart = Stopwatch.GetTimestamp(); try { - hasCurrent = await enumerator.MoveNextAsync(); + hasCurrent = await enumerator.MoveNextAsync().ConfigureAwait(false); } catch (Exception ex) { @@ -272,7 +281,10 @@ public sealed override async IAsyncEnumerable GetStreamingRe } finally { - terminalException = await DisposeAsync(enumerator, terminalException, cancellationToken); + terminalException = await DisposeAsync( + enumerator, + terminalException, + cancellationToken).ConfigureAwait(false); var attempt = new FailoverChatClientAttempt( selectedClient, @@ -287,7 +299,11 @@ public sealed override async IAsyncEnumerable GetStreamingRe IsCancellation(terminalException, cancellationToken) || reachedAttemptLimit; - await OnRoutingUpdateAsync(context, attempt, isTerminalAttempt, cancellationToken); + await OnRoutingUpdateAsync( + context, + attempt, + isTerminalAttempt, + cancellationToken).ConfigureAwait(false); if (terminalException is not null && isTerminalAttempt) { @@ -315,7 +331,7 @@ public sealed override async IAsyncEnumerable GetStreamingRe { try { - await disposable.DisposeAsync(); + await disposable.DisposeAsync().ConfigureAwait(false); } catch (Exception ex) { From 0abbfd135502cb5db4cc9d248cddd647752bc620 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Wed, 29 Jul 2026 03:01:33 -0700 Subject: [PATCH 07/15] Address routing API review feedback Clarify RoutingContext input semantics, add optional message buffering, centralize validation, simplify timing and disposal, and separate caller cancellation from provider failures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 7 +- .../ChatRouting/RoutingContext.cs | 41 ++- .../Microsoft.Extensions.AI.Abstractions.json | 8 +- .../ChatRouting/FailoverChatClient.cs | 113 ++++--- .../ChatRouting/SemanticRoutingChatClient.cs | 7 +- .../ChatRouting/RoutingChatClientTests.cs | 304 +++++++++++++++--- 6 files changed, 378 insertions(+), 102 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index 990a84095f5..986f61e1641 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -38,7 +38,10 @@ public static RoutingChatClient Create( /// The request-specific inputs. /// The cancellation token supplied for the request. /// The client to invoke. - /// Exceptions from this method propagate to the caller. + /// + /// Implementations should generally inspect the request inputs and return a client already configured for + /// route-specific invocation behavior. Exceptions from this method propagate to the caller. + /// protected abstract ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken); @@ -49,6 +52,7 @@ public virtual async Task GetResponseAsync( ChatOptions? options = null, CancellationToken cancellationToken = default) { + _ = Throw.IfNull(messages); var context = new RoutingContext(messages, options); IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); @@ -64,6 +68,7 @@ public virtual async IAsyncEnumerable GetStreamingResponseAs ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) { + _ = Throw.IfNull(messages); var context = new RoutingContext(messages, options); IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs index 7a0c9b14011..68a5b3b327b 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs @@ -15,8 +15,9 @@ namespace Microsoft.Extensions.AI; /// started from the sequence returned by . /// /// -/// A derived client may replace or during selection. The selected -/// client receives the updated values. +/// Selectors should generally treat the request inputs as read-only and return a client already configured for +/// route-specific behavior. provides explicit request-local buffering when repeatable +/// message enumeration is required. /// /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] @@ -33,17 +34,35 @@ public RoutingContext( ChatOptions = chatOptions; } - /// Gets or sets the messages supplied to client selection and the selected client. + /// Gets the messages supplied to client selection and the selected client. /// - /// The sequence may be enumerated during selection and again by one or more selected clients. Callers should - /// supply a repeatable sequence, or a selector may replace it with a materialized sequence when required. + /// Selectors should generally treat this sequence as input. A selector that must enumerate the sequence should use + /// when repeatable enumeration is required. /// - public IEnumerable Messages + public IEnumerable Messages { get; private set; } + + /// Gets the options supplied to client selection and the selected client. + /// + /// Selectors should generally treat this instance as input and return a client already configured for + /// route-specific behavior. Because is mutable, changes to the instance are observed by + /// the selected client and subsequent failover attempts. + /// + public ChatOptions? ChatOptions { get; } + + /// Returns the messages as a repeatable list, buffering the sequence when necessary. + /// The existing message list, or a list created and cached by enumerating once. + /// + /// The cached list is used by subsequent selectors and selected clients for this request. Existing + /// instances and individual messages are not cloned. + /// + public IReadOnlyList BufferMessages() { - get; - set => field = Throw.IfNull(value); - } + if (Messages is not IReadOnlyList buffered) + { + buffered = [.. Messages]; + Messages = buffered; + } - /// Gets or sets the options supplied to client selection and the selected client. - public ChatOptions? ChatOptions { get; set; } + return buffered; + } } 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 f7fa67c49bd..b893b6a4a6d 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 @@ -4011,15 +4011,19 @@ { "Member": "Microsoft.Extensions.AI.RoutingContext.RoutingContext(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? chatOptions);", "Stage": "Experimental" + }, + { + "Member": "System.Collections.Generic.IReadOnlyList Microsoft.Extensions.AI.RoutingContext.BufferMessages();", + "Stage": "Experimental" } ], "Properties": [ { - "Member": "Microsoft.Extensions.AI.ChatOptions? Microsoft.Extensions.AI.RoutingContext.ChatOptions { get; set; }", + "Member": "Microsoft.Extensions.AI.ChatOptions? Microsoft.Extensions.AI.RoutingContext.ChatOptions { get; }", "Stage": "Experimental" }, { - "Member": "System.Collections.Generic.IEnumerable Microsoft.Extensions.AI.RoutingContext.Messages { get; set; }", + "Member": "System.Collections.Generic.IEnumerable Microsoft.Extensions.AI.RoutingContext.Messages { get; private set; }", "Stage": "Experimental" } ] diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index 285a9ec8c9f..81dce14f1b8 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -78,8 +78,9 @@ public int? MaximumAttemptsPerRequest /// . /// /// - /// A attempt is always terminal. If every callback completes successfully, every client - /// invocation produces one update and exactly one update per request is terminal. + /// A attempt is always terminal and indicates that client selection terminated routing + /// before invoking a client. If every callback completes successfully, every client invocation produces one update + /// and exactly one update per request is terminal. /// /// /// Exceptions from this method propagate to the caller. A terminal update exception replaces the response or @@ -100,7 +101,9 @@ public sealed override async Task GetResponseAsync( ChatOptions? options = null, CancellationToken cancellationToken = default) { + _ = Throw.IfNull(messages); var context = new RoutingContext(messages, options); + _ = context.BufferMessages(); int? maximumAttempts = MaximumAttemptsPerRequest; int attemptCount = 0; @@ -114,18 +117,24 @@ public sealed override async Task GetResponseAsync( } catch (Exception) { + bool selectionCancellationRequested = cancellationToken.IsCancellationRequested; await OnRoutingUpdateAsync( context, attempt: null, isTerminal: true, cancellationToken).ConfigureAwait(false); + if (selectionCancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + } + throw; } attemptCount++; ChatResponse? response = null; Exception? exception = null; - long start = Stopwatch.GetTimestamp(); + Stopwatch stopwatch = Stopwatch.StartNew(); try { @@ -139,16 +148,20 @@ await OnRoutingUpdateAsync( exception = ex; } + stopwatch.Stop(); var attempt = new FailoverChatClientAttempt( selectedClient, exception, - GetElapsedTime(start), + stopwatch.Elapsed, timeToFirstUpdate: null, responseCompleted: exception is null, outputCommitted: false); + bool cancellationRequested = + exception is not null && + cancellationToken.IsCancellationRequested; bool isTerminal = exception is null || - IsCancellation(exception, cancellationToken) || + cancellationRequested || (maximumAttempts is int limit && attemptCount >= limit); await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken).ConfigureAwait(false); @@ -158,6 +171,11 @@ exception is null || return response!; } + if (cancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + } + if (isTerminal) { Rethrow(exception); @@ -171,7 +189,9 @@ public sealed override async IAsyncEnumerable GetStreamingRe ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) { + _ = Throw.IfNull(messages); var context = new RoutingContext(messages, options); + _ = context.BufferMessages(); int? maximumAttempts = MaximumAttemptsPerRequest; int attemptCount = 0; @@ -185,11 +205,17 @@ public sealed override async IAsyncEnumerable GetStreamingRe } catch (Exception) { + bool selectionCancellationRequested = cancellationToken.IsCancellationRequested; await OnRoutingUpdateAsync( context, attempt: null, isTerminal: true, cancellationToken).ConfigureAwait(false); + if (selectionCancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + } + throw; } @@ -197,10 +223,9 @@ await OnRoutingUpdateAsync( bool reachedAttemptLimit = maximumAttempts is int limit && attemptCount >= limit; IAsyncEnumerator? enumerator = null; TimeSpan? timeToFirstUpdate = null; - TimeSpan activeDuration = TimeSpan.Zero; bool hasCurrent; - long operationStart = Stopwatch.GetTimestamp(); + Stopwatch stopwatch = Stopwatch.StartNew(); try { enumerator = selectedClient @@ -214,20 +239,25 @@ await OnRoutingUpdateAsync( } catch (Exception ex) { - activeDuration += GetElapsedTime(operationStart); - Exception exception = - (await DisposeAsync(enumerator, ex, cancellationToken).ConfigureAwait(false))!; + stopwatch.Stop(); + Exception exception = (await DisposeAsync(enumerator, ex).ConfigureAwait(false))!; var attempt = new FailoverChatClientAttempt( selectedClient, exception, - activeDuration, + stopwatch.Elapsed, timeToFirstUpdate: null, responseCompleted: false, outputCommitted: false); - bool isTerminal = IsCancellation(exception, cancellationToken) || reachedAttemptLimit; + bool cancellationRequested = cancellationToken.IsCancellationRequested; + bool isTerminal = cancellationRequested || reachedAttemptLimit; await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken).ConfigureAwait(false); + if (cancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + } + if (isTerminal) { Rethrow(exception); @@ -236,7 +266,7 @@ await OnRoutingUpdateAsync( continue; } - activeDuration += GetElapsedTime(operationStart); + stopwatch.Stop(); bool responseCompleted = false; bool outputCommitted = false; bool isTerminalAttempt = false; @@ -246,22 +276,22 @@ await OnRoutingUpdateAsync( { while (hasCurrent) { - operationStart = Stopwatch.GetTimestamp(); + stopwatch.Start(); bool hasCurrentValue = TryGetCurrent( enumerator, out ChatResponseUpdate current, out terminalException); - activeDuration += GetElapsedTime(operationStart); + stopwatch.Stop(); if (!hasCurrentValue) { break; } - timeToFirstUpdate ??= activeDuration; + timeToFirstUpdate ??= stopwatch.Elapsed; outputCommitted = true; yield return current; - operationStart = Stopwatch.GetTimestamp(); + stopwatch.Start(); try { hasCurrent = await enumerator.MoveNextAsync().ConfigureAwait(false); @@ -273,7 +303,7 @@ await OnRoutingUpdateAsync( } finally { - activeDuration += GetElapsedTime(operationStart); + stopwatch.Stop(); } } @@ -281,22 +311,22 @@ await OnRoutingUpdateAsync( } finally { - terminalException = await DisposeAsync( - enumerator, - terminalException, - cancellationToken).ConfigureAwait(false); + terminalException = await DisposeAsync(enumerator, terminalException).ConfigureAwait(false); var attempt = new FailoverChatClientAttempt( selectedClient, terminalException, - activeDuration, + stopwatch.Elapsed, timeToFirstUpdate, responseCompleted: responseCompleted && terminalException is null, outputCommitted: outputCommitted); + bool cancellationRequested = + terminalException is not null && + cancellationToken.IsCancellationRequested; isTerminalAttempt = attempt.ResponseCompleted || outputCommitted || - IsCancellation(terminalException, cancellationToken) || + cancellationRequested || reachedAttemptLimit; await OnRoutingUpdateAsync( @@ -305,9 +335,17 @@ await OnRoutingUpdateAsync( isTerminalAttempt, cancellationToken).ConfigureAwait(false); - if (terminalException is not null && isTerminalAttempt) + if (terminalException is not null) { - Rethrow(terminalException); + if (cancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + } + + if (isTerminalAttempt) + { + Rethrow(terminalException); + } } } @@ -320,13 +358,14 @@ await OnRoutingUpdateAsync( } } + [SuppressMessage( + "Resilience", + "EA0014:The async method doesn't support cancellation", + Justification = "IAsyncDisposable.DisposeAsync does not support cancellation.")] private static async ValueTask DisposeAsync( IAsyncDisposable? disposable, - Exception? exception, - CancellationToken cancellationToken) + Exception? exception) { - _ = cancellationToken; - if (disposable is not null) { try @@ -342,17 +381,6 @@ await OnRoutingUpdateAsync( return exception; } - private static TimeSpan GetElapsedTime(long startingTimestamp) => -#if NET - Stopwatch.GetElapsedTime(startingTimestamp); -#else - new((long)((Stopwatch.GetTimestamp() - startingTimestamp) * - ((double)TimeSpan.TicksPerSecond / Stopwatch.Frequency))); -#endif - - private static bool IsCancellation(Exception? exception, CancellationToken cancellationToken) => - exception is OperationCanceledException || cancellationToken.IsCancellationRequested; - private static bool TryGetCurrent( IAsyncEnumerator enumerator, out ChatResponseUpdate current, @@ -372,11 +400,12 @@ private static bool TryGetCurrent( } } +#if NET [DoesNotReturn] +#endif private static void Rethrow(Exception exception) { ExceptionDispatchInfo.Capture(exception).Throw(); - throw exception; } } diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs index bd3fcf00b69..cbcccc78189 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -169,12 +169,7 @@ protected override async ValueTask SelectClientAsync( CancellationToken cancellationToken) { _ = Throw.IfNull(context); - IEnumerable messages = context.Messages; - if (messages is not IReadOnlyList) - { - context.Messages = messages = messages.ToArray(); - } - + IReadOnlyList messages = context.BufferMessages(); string? query = LastUserText(messages); if (string.IsNullOrWhiteSpace(query)) { diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs index 4bd4751f344..afa3fe86537 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -14,20 +14,32 @@ namespace Microsoft.Extensions.AI; public class RoutingChatClientTests { [Fact] - public void RoutingContext_CarriesMutableRequestInputs() + public void RoutingContext_CarriesRequestInputs() { var messages = new List { new(ChatRole.User, "initial") }; var options = new ChatOptions { ModelId = "initial" }; var context = new RoutingContext(messages, options); - var replacementMessages = new List { new(ChatRole.User, "replacement") }; - var replacementOptions = new ChatOptions { ModelId = "replacement" }; - context.Messages = replacementMessages; - context.ChatOptions = replacementOptions; - Assert.Same(replacementMessages, context.Messages); - Assert.Same(replacementOptions, context.ChatOptions); + Assert.Same(messages, context.Messages); + Assert.Same(messages, context.BufferMessages()); + Assert.Same(options, context.ChatOptions); Assert.Throws(() => new RoutingContext(null!, options)); - Assert.Throws(() => context.Messages = null!); + } + + [Fact] + public void RoutingContext_BufferMessagesEnumeratesOnceAndCaches() + { + var messages = new SingleUseMessageEnumerable( + [new ChatMessage(ChatRole.User, "one"), new ChatMessage(ChatRole.Assistant, "two")]); + var context = new RoutingContext(messages, chatOptions: null); + + IReadOnlyList first = context.BufferMessages(); + IReadOnlyList second = context.BufferMessages(); + + Assert.Same(first, second); + Assert.Same(first, context.Messages); + Assert.Equal(2, first.Count); + Assert.Equal(1, messages.EnumerationCount); } [Fact] @@ -103,6 +115,17 @@ public async Task Create_NullSelectionThrowsForNonStreamingAndStreaming() Assert.Contains("SelectClientAsync", streaming.Message, StringComparison.Ordinal); } + [Fact] + public async Task Create_RejectsNullMessagesForNonStreamingAndStreaming() + { + using var selected = new TestChatClient(); + using RoutingChatClient router = RoutingChatClient.Create((_, _) => new(selected)); + + await Assert.ThrowsAsync(() => router.GetResponseAsync(null!)); + await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync(null!))); + } + [Fact] public async Task Create_DoesNotDisposeSelectedClient() { @@ -142,6 +165,17 @@ public void MaximumAttemptsPerRequest_ValidatesValue() Assert.Throws(() => client.MaximumAttemptsPerRequest = -1); } + [Fact] + public async Task Failover_RejectsNullMessagesForNonStreamingAndStreaming() + { + using var selected = new TestChatClient(); + using var router = new DelegatingFailoverTestRouter(_ => selected); + + await Assert.ThrowsAsync(() => router.GetResponseAsync(null!)); + await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync(null!))); + } + [Fact] public async Task InitialSelectionFailureReportsTerminalUpdateWithoutAttempt() { @@ -171,6 +205,30 @@ public async Task InitialSelectionFailureReportsTerminalUpdateWithoutAttempt() Assert.Equal(1, completionCount); } + [Fact] + public async Task InitialSelectionFailureWithCallerCancellationThrowsCancellation() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + var selectionException = new InvalidOperationException("selection failed"); + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => throw selectionException, + (_, attempt, isTerminal) => + { + updateCount++; + Assert.Null(attempt); + Assert.True(isTerminal); + }); + + await Assert.ThrowsAnyAsync( + () => router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.Equal(1, updateCount); + } + [Fact] public async Task StreamingInitialSelectionFailureReportsTerminalUpdateWithoutAttempt() { @@ -388,32 +446,29 @@ public async Task Dispatch_ConfiguredClientPreservesAndOverridesRequestOptions() } [Fact] - public async Task Policy_CanMutateContextBeforeDispatch() + public async Task Policy_CanBufferMessagesBeforeDispatch() { IEnumerable? forwardedMessages = null; - ChatOptions? forwardedOptions = null; + IReadOnlyList? bufferedMessages = null; using var inner = new TestChatClient { - GetResponseAsyncCallback = (messages, options, _) => + GetResponseAsyncCallback = (messages, _, _) => { forwardedMessages = messages; - forwardedOptions = options; return Task.FromResult(new ChatResponse()); }, }; - var replacementMessages = new List { new(ChatRole.User, "replacement") }; - var replacementOptions = new ChatOptions { ModelId = "replacement" }; + var messages = new SingleUseMessageEnumerable([new(ChatRole.User, "original")]); using var router = new DelegatingTestRouter(context => { - context.Messages = replacementMessages; - context.ChatOptions = replacementOptions; + bufferedMessages = context.BufferMessages(); return inner; }); - _ = await router.GetResponseAsync([new(ChatRole.User, "original")], new ChatOptions()); + _ = await router.GetResponseAsync(messages); - Assert.Same(replacementMessages, forwardedMessages); - Assert.Same(replacementOptions, forwardedOptions); + Assert.Same(bufferedMessages, forwardedMessages); + Assert.Equal(1, messages.EnumerationCount); } [Fact] @@ -625,10 +680,10 @@ public async Task Failover_CancellationReportsTerminalUpdate() { using var cancellationSource = new CancellationTokenSource(); cancellationSource.Cancel(); + var failure = new InvalidOperationException("failed"); using var inner = new TestChatClient { - GetResponseAsyncCallback = (_, _, cancellationToken) => - throw new OperationCanceledException(cancellationToken), + GetResponseAsyncCallback = (_, _, _) => throw failure, }; int updateCount = 0; using var router = new DelegatingFailoverTestRouter( @@ -636,7 +691,7 @@ public async Task Failover_CancellationReportsTerminalUpdate() (_, attempt, isTerminal) => { updateCount++; - Assert.IsAssignableFrom(attempt!.Exception); + Assert.Same(failure, attempt!.Exception); Assert.True(isTerminal); }); @@ -648,6 +703,51 @@ await Assert.ThrowsAnyAsync( Assert.Equal(1, updateCount); } + [Fact] + public async Task Failover_CancellationDuringUpdateIsObservedByNextAttempt() + { + using var cancellationSource = new CancellationTokenSource(); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + }; + int selections = 0; + var updates = new List<(FailoverChatClientAttempt? attempt, bool isTerminal)>(); + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + updates.Add((attempt, isTerminal)); + if (!isTerminal) + { + cancellationSource.Cancel(); + } + }); + + await Assert.ThrowsAnyAsync( + () => router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.Equal(2, selections); + Assert.Collection( + updates, + update => + { + Assert.NotNull(update.attempt); + Assert.False(update.isTerminal); + }, + update => + { + Assert.NotNull(update.attempt); + Assert.True(update.isTerminal); + }); + } + [Fact] public async Task Streaming_PreOutputFailurePropagates() { @@ -698,6 +798,52 @@ public async Task Streaming_FallsBackBeforeFirstUpdate() Assert.Equal(2, selections); } + [Fact] + public async Task Streaming_CancellationDuringUpdateIsObservedByNextAttempt() + { + using var cancellationSource = new CancellationTokenSource(); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + int selections = 0; + var updates = new List<(FailoverChatClientAttempt? attempt, bool isTerminal)>(); + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + updates.Add((attempt, isTerminal)); + if (!isTerminal) + { + cancellationSource.Cancel(); + } + }); + + await Assert.ThrowsAnyAsync( + () => CollectAsync( + router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token))); + + Assert.Equal(2, selections); + Assert.Collection( + updates, + update => + { + Assert.NotNull(update.attempt); + Assert.False(update.isTerminal); + }, + update => + { + Assert.NotNull(update.attempt); + Assert.True(update.isTerminal); + }); + } + [Fact] public async Task Streaming_CancellationDoesNotSelectNext() { @@ -705,8 +851,7 @@ public async Task Streaming_CancellationDoesNotSelectNext() cancellationSource.Cancel(); using var inner = new TestChatClient { - GetStreamingResponseAsyncCallback = (_, _, cancellationToken) => - CanceledStream(cancellationToken), + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), }; int selections = 0; FailoverChatClientAttempt? observed = null; @@ -729,15 +874,16 @@ await Assert.ThrowsAnyAsync( cancellationToken: cancellationSource.Token))); Assert.Equal(1, selections); - Assert.IsAssignableFrom(observed!.Exception); + Assert.IsType(observed!.Exception); Assert.False(observed.OutputCommitted); Assert.False(observed.ResponseCompleted); } [Fact] - public async Task OperationCanceledException_DoesNotTriggerFailover() + public async Task OperationCanceledException_CanTriggerFailover() { bool secondCalled = false; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); using var first = new TestChatClient { GetResponseAsyncCallback = (_, _, _) => throw new OperationCanceledException(), @@ -747,19 +893,19 @@ public async Task OperationCanceledException_DoesNotTriggerFailover() GetResponseAsyncCallback = (_, _, _) => { secondCalled = true; - return Task.FromResult(new ChatResponse()); + return Task.FromResult(expected); }, }; using var client = new OrderedFailoverChatClient([first, second]); - await Assert.ThrowsAnyAsync( - () => client.GetResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); - Assert.False(secondCalled); + Assert.Same(expected, response); + Assert.True(secondCalled); } [Fact] - public async Task Streaming_OperationCanceledExceptionDoesNotTriggerFailover() + public async Task Streaming_OperationCanceledExceptionCanTriggerFailover() { bool secondCalled = false; using var first = new TestChatClient @@ -776,14 +922,15 @@ public async Task Streaming_OperationCanceledExceptionDoesNotTriggerFailover() }; using var client = new OrderedFailoverChatClient([first, second]); - await Assert.ThrowsAnyAsync( - () => CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + List updates = + await CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); - Assert.False(secondCalled); + Assert.Equal("ok", Assert.Single(updates).Text); + Assert.True(secondCalled); } [Fact] - public async Task Streaming_DisposalCancellationDoesNotTriggerFailover() + public async Task Streaming_DisposalCancellationCanTriggerFailover() { bool secondCalled = false; var canceledStream = new TrackingAsyncEnumerable( @@ -803,10 +950,11 @@ public async Task Streaming_DisposalCancellationDoesNotTriggerFailover() }; using var client = new OrderedFailoverChatClient([first, second]); - await Assert.ThrowsAnyAsync( - () => CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + List updates = + await CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); - Assert.False(secondCalled); + Assert.Equal("ok", Assert.Single(updates).Text); + Assert.True(secondCalled); Assert.Equal(1, canceledStream.DisposeCount); } @@ -918,6 +1066,54 @@ public async Task Streaming_CurrentFailureFallsBackBeforeOutput() Assert.Equal(1, failedStream.DisposeCount); } + [Fact] + public async Task Streaming_CurrentFailureCancellationIsObservedByNextAttempt() + { + using var cancellationSource = new CancellationTokenSource(); + var failedStream = new ThrowingCurrentAsyncEnumerable(new InvalidOperationException("current failed")); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => failedStream, + }; + int selections = 0; + var updates = new List<(FailoverChatClientAttempt? attempt, bool isTerminal)>(); + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + updates.Add((attempt, isTerminal)); + if (!isTerminal) + { + cancellationSource.Cancel(); + } + }); + + await Assert.ThrowsAnyAsync( + () => CollectAsync( + router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token))); + + Assert.Equal(2, selections); + Assert.Collection( + updates, + update => + { + Assert.NotNull(update.attempt); + Assert.False(update.isTerminal); + }, + update => + { + Assert.NotNull(update.attempt); + Assert.True(update.isTerminal); + }); + Assert.Equal(2, failedStream.DisposeCount); + } + [Fact] public async Task Streaming_EmptyStreamDisposalFailureFallsBack() { @@ -1447,6 +1643,35 @@ public async Task OrderedFailover_StateIsScopedToOneRequest() Assert.Equal(2, firstCalls); } + [Fact] + public async Task OrderedFailover_BuffersMessagesAcrossAttempts() + { + var messages = new SingleUseMessageEnumerable([new(ChatRole.User, "hi")]); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (forwarded, _, _) => + { + Assert.Equal("hi", Assert.Single(forwarded).Text); + throw new InvalidOperationException("failed"); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var second = new TestChatClient + { + GetResponseAsyncCallback = (forwarded, _, _) => + { + Assert.Equal("hi", Assert.Single(forwarded).Text); + return Task.FromResult(expected); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = await client.GetResponseAsync(messages); + + Assert.Same(expected, response); + Assert.Equal(1, messages.EnumerationCount); + } + [Fact] public async Task OrderedFailover_ConcurrentRequestsHaveIndependentState() { @@ -1595,8 +1820,7 @@ public async Task OrderedFailover_CancellationDoesNotFallback() bool secondCalled = false; using var first = new TestChatClient { - GetResponseAsyncCallback = (_, _, cancellationToken) => - throw new OperationCanceledException(cancellationToken), + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), }; using var second = new TestChatClient { From ea929cc43abed923cdca72fe088e9755690084d8 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Wed, 29 Jul 2026 17:38:32 -0700 Subject: [PATCH 08/15] Clarify routing context and attempt reporting Clone request options for selector shaping and restrict routing updates to concrete client attempts, leaving selection-failure cleanup to selectors. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 5 +- .../ChatRouting/RoutingContext.cs | 18 +-- .../ChatRouting/FailoverChatClient.cs | 34 ++--- .../ChatRouting/OrderedFailoverChatClient.cs | 4 +- .../Microsoft.Extensions.AI.json | 4 +- .../ChatRouting/RoutingChatClientTests.cs | 132 +++++++----------- 6 files changed, 78 insertions(+), 119 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index 986f61e1641..f4430515294 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -39,8 +39,9 @@ public static RoutingChatClient Create( /// The cancellation token supplied for the request. /// The client to invoke. /// - /// Implementations should generally inspect the request inputs and return a client already configured for - /// route-specific invocation behavior. Exceptions from this method propagate to the caller. + /// Implementations may adjust the request-local for dynamic request + /// shaping. Stable route-specific behavior should generally be attached to the returned client. Exceptions from + /// this method propagate to the caller. /// protected abstract ValueTask SelectClientAsync( RoutingContext context, diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs index 68a5b3b327b..ce9887817ea 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs @@ -15,9 +15,9 @@ namespace Microsoft.Extensions.AI; /// started from the sequence returned by . /// /// -/// Selectors should generally treat the request inputs as read-only and return a client already configured for -/// route-specific behavior. provides explicit request-local buffering when repeatable -/// message enumeration is required. +/// is cloned when the context is created, so request-specific changes do not mutate the +/// caller's instance. provides explicit request-local buffering when repeatable message +/// enumeration is required. /// /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] @@ -25,13 +25,13 @@ public class RoutingContext { /// Initializes a new instance of the class. /// The messages to route. - /// The options supplied for the request. + /// The options to clone for this request, or . public RoutingContext( IEnumerable messages, ChatOptions? chatOptions) { Messages = Throw.IfNull(messages); - ChatOptions = chatOptions; + ChatOptions = chatOptions?.Clone(); } /// Gets the messages supplied to client selection and the selected client. @@ -41,11 +41,11 @@ public RoutingContext( /// public IEnumerable Messages { get; private set; } - /// Gets the options supplied to client selection and the selected client. + /// Gets the request-local options supplied to client selection and the selected client. /// - /// Selectors should generally treat this instance as input and return a client already configured for - /// route-specific behavior. Because is mutable, changes to the instance are observed by - /// the selected client and subsequent failover attempts. + /// When the caller supplies options, this is a clone that selectors may adjust without mutating the caller's + /// instance. Changes are observed by the selected client and subsequent failover attempts. Stable route-specific + /// defaults should generally be attached to the returned client rather than reapplied during selection. /// public ChatOptions? ChatOptions { get; } diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index 81dce14f1b8..382bbd54be6 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -20,12 +20,13 @@ namespace Microsoft.Extensions.AI; /// /// /// The client for each attempt is supplied by . After an invocation, -/// reports its outcome and whether routing is terminal. An uncanceled failure causes -/// another selection only when it happened before any streaming output was exposed and the attempt limit permits it. +/// reports the concrete attempt and whether another selection will follow. An +/// uncanceled failure causes another selection only when it happened before any streaming output was exposed and the +/// attempt limit permits it. /// /// -/// The base class owns invocation, streaming commitment, attempt limits, and terminal reporting. Derived classes own -/// client selection, policy state, and the lifetime of clients they retain. +/// The base class owns invocation, streaming commitment, attempt limits, and attempt reporting. Derived classes own +/// client selection, policy state, selection-failure cleanup, and the lifetime of clients they retain. /// /// /// Once streaming enumeration begins, callers must dispose the enumerator. Abandoning an active enumerator without @@ -59,12 +60,9 @@ public int? MaximumAttemptsPerRequest } } - /// Invoked after a client invocation or when client selection terminates routing. + /// Invoked after a client invocation completes, fails, or is abandoned. /// The request-specific inputs. - /// - /// The completed client invocation, or when - /// terminated routing before invoking a client. - /// + /// The completed client invocation. /// /// if the base will not select another client after this callback completes successfully; /// otherwise, . @@ -78,9 +76,9 @@ public int? MaximumAttemptsPerRequest /// . /// /// - /// A attempt is always terminal and indicates that client selection terminated routing - /// before invoking a client. If every callback completes successfully, every client invocation produces one update - /// and exactly one update per request is terminal. + /// Every client invocation produces one update if the callback completes successfully. Selection failures are not + /// reported. A selector that retains request-scoped state must release it before throwing; a request may therefore + /// end without a terminal update when selection fails. /// /// /// Exceptions from this method propagate to the caller. A terminal update exception replaces the response or @@ -91,7 +89,7 @@ public int? MaximumAttemptsPerRequest /// protected virtual ValueTask OnRoutingUpdateAsync( RoutingContext context, - FailoverChatClientAttempt? attempt, + FailoverChatClientAttempt attempt, bool isTerminal, CancellationToken cancellationToken) => default; @@ -118,11 +116,6 @@ public sealed override async Task GetResponseAsync( catch (Exception) { bool selectionCancellationRequested = cancellationToken.IsCancellationRequested; - await OnRoutingUpdateAsync( - context, - attempt: null, - isTerminal: true, - cancellationToken).ConfigureAwait(false); if (selectionCancellationRequested) { cancellationToken.ThrowIfCancellationRequested(); @@ -206,11 +199,6 @@ public sealed override async IAsyncEnumerable GetStreamingRe catch (Exception) { bool selectionCancellationRequested = cancellationToken.IsCancellationRequested; - await OnRoutingUpdateAsync( - context, - attempt: null, - isTerminal: true, - cancellationToken).ConfigureAwait(false); if (selectionCancellationRequested) { cancellationToken.ThrowIfCancellationRequested(); diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs index e87130c8b17..f5ee88eed1c 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -99,7 +99,7 @@ protected override ValueTask SelectClientAsync( /// protected override ValueTask OnRoutingUpdateAsync( RoutingContext context, - FailoverChatClientAttempt? attempt, + FailoverChatClientAttempt attempt, bool isTerminal, CancellationToken cancellationToken) { @@ -113,7 +113,7 @@ protected override ValueTask OnRoutingUpdateAsync( return default; } - if (attempt?.Exception is null) + if (attempt.Exception is null) { _ = _requestStates.Remove(context); throw new InvalidOperationException("A nonterminal routing update requires a failed client invocation."); diff --git a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json index c4dd509a8cf..e55cd6b9619 100644 --- a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json +++ b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json @@ -506,7 +506,7 @@ "Stage": "Experimental" }, { - "Member": "virtual System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.FailoverChatClient.OnRoutingUpdateAsync(Microsoft.Extensions.AI.RoutingContext context, Microsoft.Extensions.AI.FailoverChatClientAttempt? attempt, bool isTerminal, System.Threading.CancellationToken cancellationToken);", + "Member": "virtual System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.FailoverChatClient.OnRoutingUpdateAsync(Microsoft.Extensions.AI.RoutingContext context, Microsoft.Extensions.AI.FailoverChatClientAttempt attempt, bool isTerminal, System.Threading.CancellationToken cancellationToken);", "Stage": "Experimental" } ], @@ -1496,7 +1496,7 @@ "Stage": "Experimental" }, { - "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.OrderedFailoverChatClient.OnRoutingUpdateAsync(Microsoft.Extensions.AI.RoutingContext context, Microsoft.Extensions.AI.FailoverChatClientAttempt? attempt, bool isTerminal, System.Threading.CancellationToken cancellationToken);", + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.OrderedFailoverChatClient.OnRoutingUpdateAsync(Microsoft.Extensions.AI.RoutingContext context, Microsoft.Extensions.AI.FailoverChatClientAttempt attempt, bool isTerminal, System.Threading.CancellationToken cancellationToken);", "Stage": "Experimental" }, { diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs index afa3fe86537..cda4dc79a5d 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -22,7 +22,10 @@ public void RoutingContext_CarriesRequestInputs() Assert.Same(messages, context.Messages); Assert.Same(messages, context.BufferMessages()); - Assert.Same(options, context.ChatOptions); + Assert.NotSame(options, context.ChatOptions); + Assert.Equal("initial", context.ChatOptions!.ModelId); + context.ChatOptions.ModelId = "changed"; + Assert.Equal("initial", options.ModelId); Assert.Throws(() => new RoutingContext(null!, options)); } @@ -52,12 +55,17 @@ public void Create_RejectsNullSelector() public async Task Create_SelectsClientForRequest() { var messages = new ChatMessage[] { new(ChatRole.User, "hi") }; - var options = new ChatOptions(); + var options = new ChatOptions { ModelId = "request" }; using var cancellationSource = new CancellationTokenSource(); ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + ChatOptions? forwardedOptions = null; using var selected = new TestChatClient { - GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + GetResponseAsyncCallback = (_, selectedOptions, _) => + { + forwardedOptions = selectedOptions; + return Task.FromResult(expected); + }, }; RoutingContext? observedContext = null; CancellationToken observedToken = default; @@ -67,6 +75,7 @@ public async Task Create_SelectsClientForRequest() observedContext = context; observedToken = cancellationToken; selectionCount++; + context.ChatOptions!.ModelId = "selected"; return new(selected); }); @@ -74,7 +83,10 @@ public async Task Create_SelectsClientForRequest() Assert.Same(expected, response); Assert.Same(messages, observedContext!.Messages); - Assert.Same(options, observedContext.ChatOptions); + Assert.NotSame(options, observedContext.ChatOptions); + Assert.Same(observedContext.ChatOptions, forwardedOptions); + Assert.Equal("selected", forwardedOptions!.ModelId); + Assert.Equal("request", options.ModelId); Assert.Equal(cancellationSource.Token, observedToken); Assert.Equal(1, selectionCount); } @@ -177,32 +189,25 @@ await Assert.ThrowsAsync( } [Fact] - public async Task InitialSelectionFailureReportsTerminalUpdateWithoutAttempt() + public async Task InitialSelectionFailureDoesNotReportAttempt() { var expected = new InvalidOperationException("selection failed"); RoutingContext? selectedContext = null; - RoutingContext? completedContext = null; - int completionCount = 0; + int updateCount = 0; using var router = new DelegatingFailoverTestRouter( context => { selectedContext = context; throw expected; }, - (context, attempt, isTerminal) => - { - completionCount++; - completedContext = context; - Assert.Null(attempt); - Assert.True(isTerminal); - }); + (_, _, _) => updateCount++); InvalidOperationException actual = await Assert.ThrowsAsync( () => router.GetResponseAsync([new(ChatRole.User, "hi")])); Assert.Same(expected, actual); - Assert.Same(selectedContext, completedContext); - Assert.Equal(1, completionCount); + Assert.NotNull(selectedContext); + Assert.Equal(0, updateCount); } [Fact] @@ -214,92 +219,70 @@ public async Task InitialSelectionFailureWithCallerCancellationThrowsCancellatio int updateCount = 0; using var router = new DelegatingFailoverTestRouter( _ => throw selectionException, - (_, attempt, isTerminal) => - { - updateCount++; - Assert.Null(attempt); - Assert.True(isTerminal); - }); + (_, _, _) => updateCount++); await Assert.ThrowsAnyAsync( () => router.GetResponseAsync( [new(ChatRole.User, "hi")], cancellationToken: cancellationSource.Token)); - Assert.Equal(1, updateCount); + Assert.Equal(0, updateCount); } [Fact] - public async Task StreamingInitialSelectionFailureReportsTerminalUpdateWithoutAttempt() + public async Task StreamingInitialSelectionFailureDoesNotReportAttempt() { var expected = new InvalidOperationException("selection failed"); RoutingContext? selectedContext = null; - RoutingContext? completedContext = null; - int completionCount = 0; + int updateCount = 0; using var router = new DelegatingFailoverTestRouter( context => { selectedContext = context; throw expected; }, - (context, attempt, isTerminal) => - { - completionCount++; - completedContext = context; - Assert.Null(attempt); - Assert.True(isTerminal); - }); + (_, _, _) => updateCount++); InvalidOperationException actual = await Assert.ThrowsAsync( () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); Assert.Same(expected, actual); - Assert.Same(selectedContext, completedContext); - Assert.Equal(1, completionCount); + Assert.NotNull(selectedContext); + Assert.Equal(0, updateCount); } [Fact] - public async Task NullSelectionReportsTerminalUpdateWithoutAttempt() + public async Task NullSelectionDoesNotReportAttempt() { int updateCount = 0; using var router = new DelegatingFailoverTestRouter( _ => null!, - (_, attempt, isTerminal) => - { - updateCount++; - Assert.Null(attempt); - Assert.True(isTerminal); - }); + (_, _, _) => updateCount++); InvalidOperationException exception = await Assert.ThrowsAsync( () => router.GetResponseAsync([new(ChatRole.User, "hi")])); Assert.Contains("SelectClientAsync", exception.Message, StringComparison.Ordinal); - Assert.Equal(1, updateCount); + Assert.Equal(0, updateCount); } [Fact] - public async Task StreamingNullSelectionReportsTerminalUpdateWithoutAttempt() + public async Task StreamingNullSelectionDoesNotReportAttempt() { int updateCount = 0; using var router = new DelegatingFailoverTestRouter( _ => null!, - (_, attempt, isTerminal) => - { - updateCount++; - Assert.Null(attempt); - Assert.True(isTerminal); - }); + (_, _, _) => updateCount++); InvalidOperationException exception = await Assert.ThrowsAsync( () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); Assert.Contains("SelectClientAsync", exception.Message, StringComparison.Ordinal); - Assert.Equal(1, updateCount); + Assert.Equal(0, updateCount); } [Fact] - public async Task Failover_RetrySelectionFailureReportsTerminalUpdateWithoutAttempt() + public async Task Failover_RetrySelectionFailureReportsOnlyCompletedAttempt() { var invocationException = new InvalidOperationException("invocation failed"); var selectionException = new InvalidOperationException("selection failed"); @@ -308,7 +291,7 @@ public async Task Failover_RetrySelectionFailureReportsTerminalUpdateWithoutAtte GetResponseAsyncCallback = (_, _, _) => throw invocationException, }; int selections = 0; - var updates = new List<(FailoverChatClientAttempt? attempt, bool isTerminal)>(); + var updates = new List<(FailoverChatClientAttempt attempt, bool isTerminal)>(); using var router = new DelegatingFailoverTestRouter( _ => ++selections == 1 ? failing : throw selectionException, (_, attempt, isTerminal) => @@ -321,19 +304,10 @@ public async Task Failover_RetrySelectionFailureReportsTerminalUpdateWithoutAtte Assert.Same(selectionException, actual); Assert.Equal(2, selections); - Assert.Collection( - updates, - update => - { - Assert.Same(failing, update.attempt!.Client); - Assert.Same(invocationException, update.attempt.Exception); - Assert.False(update.isTerminal); - }, - update => - { - Assert.Null(update.attempt); - Assert.True(update.isTerminal); - }); + (FailoverChatClientAttempt attempt, bool isTerminal) update = Assert.Single(updates); + Assert.Same(failing, update.attempt.Client); + Assert.Same(invocationException, update.attempt.Exception); + Assert.False(update.isTerminal); } [Fact] @@ -396,23 +370,19 @@ public async Task NonterminalRoutingUpdateFailureStopsWithoutAnotherUpdate() } [Fact] - public async Task TerminalSelectionUpdateFailureReplacesSelectionFailure() + public async Task SelectionFailureDoesNotInvokeRoutingUpdate() { var selectionException = new InvalidOperationException("selection failed"); - var updateException = new InvalidOperationException("update failed"); + int updateCount = 0; using var router = new DelegatingFailoverTestRouter( _ => throw selectionException, - (_, attempt, isTerminal) => - { - Assert.Null(attempt); - Assert.True(isTerminal); - throw updateException; - }); + (_, _, _) => updateCount++); InvalidOperationException actual = await Assert.ThrowsAsync( () => router.GetResponseAsync([new(ChatRole.User, "hi")])); - Assert.Same(updateException, actual); + Assert.Same(selectionException, actual); + Assert.Equal(0, updateCount); } [Fact] @@ -712,7 +682,7 @@ public async Task Failover_CancellationDuringUpdateIsObservedByNextAttempt() GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), }; int selections = 0; - var updates = new List<(FailoverChatClientAttempt? attempt, bool isTerminal)>(); + var updates = new List<(FailoverChatClientAttempt attempt, bool isTerminal)>(); using var router = new DelegatingFailoverTestRouter( _ => { @@ -807,7 +777,7 @@ public async Task Streaming_CancellationDuringUpdateIsObservedByNextAttempt() GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), }; int selections = 0; - var updates = new List<(FailoverChatClientAttempt? attempt, bool isTerminal)>(); + var updates = new List<(FailoverChatClientAttempt attempt, bool isTerminal)>(); using var router = new DelegatingFailoverTestRouter( _ => { @@ -1076,7 +1046,7 @@ public async Task Streaming_CurrentFailureCancellationIsObservedByNextAttempt() GetStreamingResponseAsyncCallback = (_, _, _) => failedStream, }; int selections = 0; - var updates = new List<(FailoverChatClientAttempt? attempt, bool isTerminal)>(); + var updates = new List<(FailoverChatClientAttempt attempt, bool isTerminal)>(); using var router = new DelegatingFailoverTestRouter( _ => { @@ -1979,11 +1949,11 @@ public IEnumerator GetEnumerator() private sealed class DelegatingFailoverTestRouter : FailoverChatClient { private readonly Func _select; - private readonly Action? _onRoutingUpdate; + private readonly Action? _onRoutingUpdate; public DelegatingFailoverTestRouter( Func select, - Action? onRoutingUpdate = null) + Action? onRoutingUpdate = null) { _select = select; _onRoutingUpdate = onRoutingUpdate; @@ -1996,7 +1966,7 @@ protected override ValueTask SelectClientAsync( protected override ValueTask OnRoutingUpdateAsync( RoutingContext context, - FailoverChatClientAttempt? attempt, + FailoverChatClientAttempt attempt, bool isTerminal, CancellationToken cancellationToken) { From 53f663960705881a480bd26d78be1e7577469f81 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Thu, 30 Jul 2026 08:46:03 -0700 Subject: [PATCH 09/15] Align routing message and selection semantics Follow MEAI's repeatable-enumerable convention and propagate selector failures without cancellation rewriting. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingContext.cs | 24 +-- .../Microsoft.Extensions.AI.Abstractions.json | 6 +- .../ChatRouting/FailoverChatClient.cs | 40 +---- .../ChatRouting/SemanticRoutingChatClient.cs | 18 +- .../ChatRouting/RoutingChatClientTests.cs | 155 +++--------------- 5 files changed, 36 insertions(+), 207 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs index ce9887817ea..10865674bd9 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs @@ -16,8 +16,7 @@ namespace Microsoft.Extensions.AI; /// /// /// is cloned when the context is created, so request-specific changes do not mutate the -/// caller's instance. provides explicit request-local buffering when repeatable message -/// enumeration is required. +/// caller's instance. /// /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] @@ -36,10 +35,9 @@ public RoutingContext( /// Gets the messages supplied to client selection and the selected client. /// - /// Selectors should generally treat this sequence as input. A selector that must enumerate the sequence should use - /// when repeatable enumeration is required. + /// Selection and failover may enumerate this sequence multiple times. Callers should supply a repeatable sequence. /// - public IEnumerable Messages { get; private set; } + public IEnumerable Messages { get; } /// Gets the request-local options supplied to client selection and the selected client. /// @@ -49,20 +47,4 @@ public RoutingContext( /// public ChatOptions? ChatOptions { get; } - /// Returns the messages as a repeatable list, buffering the sequence when necessary. - /// The existing message list, or a list created and cached by enumerating once. - /// - /// The cached list is used by subsequent selectors and selected clients for this request. Existing - /// instances and individual messages are not cloned. - /// - public IReadOnlyList BufferMessages() - { - if (Messages is not IReadOnlyList buffered) - { - buffered = [.. Messages]; - Messages = buffered; - } - - return buffered; - } } 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 b893b6a4a6d..8b5f337da10 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 @@ -4011,10 +4011,6 @@ { "Member": "Microsoft.Extensions.AI.RoutingContext.RoutingContext(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? chatOptions);", "Stage": "Experimental" - }, - { - "Member": "System.Collections.Generic.IReadOnlyList Microsoft.Extensions.AI.RoutingContext.BufferMessages();", - "Stage": "Experimental" } ], "Properties": [ @@ -4023,7 +4019,7 @@ "Stage": "Experimental" }, { - "Member": "System.Collections.Generic.IEnumerable Microsoft.Extensions.AI.RoutingContext.Messages { get; private set; }", + "Member": "System.Collections.Generic.IEnumerable Microsoft.Extensions.AI.RoutingContext.Messages { get; }", "Stage": "Experimental" } ] diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index 382bbd54be6..eab0b51eebc 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -101,28 +101,14 @@ public sealed override async Task GetResponseAsync( { _ = Throw.IfNull(messages); var context = new RoutingContext(messages, options); - _ = context.BufferMessages(); int? maximumAttempts = MaximumAttemptsPerRequest; int attemptCount = 0; while (true) { - IChatClient selectedClient; - try - { - selectedClient = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? - throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); - } - catch (Exception) - { - bool selectionCancellationRequested = cancellationToken.IsCancellationRequested; - if (selectionCancellationRequested) - { - cancellationToken.ThrowIfCancellationRequested(); - } - - throw; - } + IChatClient selectedClient = + await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); attemptCount++; ChatResponse? response = null; @@ -184,28 +170,14 @@ public sealed override async IAsyncEnumerable GetStreamingRe { _ = Throw.IfNull(messages); var context = new RoutingContext(messages, options); - _ = context.BufferMessages(); int? maximumAttempts = MaximumAttemptsPerRequest; int attemptCount = 0; while (true) { - IChatClient selectedClient; - try - { - selectedClient = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? - throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); - } - catch (Exception) - { - bool selectionCancellationRequested = cancellationToken.IsCancellationRequested; - if (selectionCancellationRequested) - { - cancellationToken.ThrowIfCancellationRequested(); - } - - throw; - } + IChatClient selectedClient = + await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); attemptCount++; bool reachedAttemptLimit = maximumAttempts is int limit && attemptCount >= limit; diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs index cbcccc78189..4d12214f97c 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -169,8 +169,8 @@ protected override async ValueTask SelectClientAsync( CancellationToken cancellationToken) { _ = Throw.IfNull(context); - IReadOnlyList messages = context.BufferMessages(); - string? query = LastUserText(messages); + string? query = context.Messages.LastOrDefault( + static message => message.Role == ChatRole.User)?.Text; if (string.IsNullOrWhiteSpace(query)) { return _clients[0]; @@ -303,20 +303,6 @@ protected override void Dispose(bool disposing) } } - private static string? LastUserText(IEnumerable messages) - { - string? last = null; - foreach (ChatMessage message in messages) - { - if (message.Role == ChatRole.User) - { - last = message.Text; - } - } - - return last; - } - private async Task EnsureIndexAsync(CancellationToken cancellationToken) { if (_index is { } cached) diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs index cda4dc79a5d..e2547deb3de 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -21,7 +21,6 @@ public void RoutingContext_CarriesRequestInputs() var context = new RoutingContext(messages, options); Assert.Same(messages, context.Messages); - Assert.Same(messages, context.BufferMessages()); Assert.NotSame(options, context.ChatOptions); Assert.Equal("initial", context.ChatOptions!.ModelId); context.ChatOptions.ModelId = "changed"; @@ -29,22 +28,6 @@ public void RoutingContext_CarriesRequestInputs() Assert.Throws(() => new RoutingContext(null!, options)); } - [Fact] - public void RoutingContext_BufferMessagesEnumeratesOnceAndCaches() - { - var messages = new SingleUseMessageEnumerable( - [new ChatMessage(ChatRole.User, "one"), new ChatMessage(ChatRole.Assistant, "two")]); - var context = new RoutingContext(messages, chatOptions: null); - - IReadOnlyList first = context.BufferMessages(); - IReadOnlyList second = context.BufferMessages(); - - Assert.Same(first, second); - Assert.Same(first, context.Messages); - Assert.Equal(2, first.Count); - Assert.Equal(1, messages.EnumerationCount); - } - [Fact] public void Create_RejectsNullSelector() { @@ -211,7 +194,7 @@ public async Task InitialSelectionFailureDoesNotReportAttempt() } [Fact] - public async Task InitialSelectionFailureWithCallerCancellationThrowsCancellation() + public async Task InitialSelectionFailureWithCallerCancellationPropagatesSelectionFailure() { using var cancellationSource = new CancellationTokenSource(); cancellationSource.Cancel(); @@ -221,11 +204,12 @@ public async Task InitialSelectionFailureWithCallerCancellationThrowsCancellatio _ => throw selectionException, (_, _, _) => updateCount++); - await Assert.ThrowsAnyAsync( + InvalidOperationException actual = await Assert.ThrowsAsync( () => router.GetResponseAsync( [new(ChatRole.User, "hi")], cancellationToken: cancellationSource.Token)); + Assert.Same(selectionException, actual); Assert.Equal(0, updateCount); } @@ -251,6 +235,27 @@ public async Task StreamingInitialSelectionFailureDoesNotReportAttempt() Assert.Equal(0, updateCount); } + [Fact] + public async Task StreamingInitialSelectionFailureWithCallerCancellationPropagatesSelectionFailure() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + var selectionException = new InvalidOperationException("selection failed"); + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => throw selectionException, + (_, _, _) => updateCount++); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => CollectAsync( + router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token))); + + Assert.Same(selectionException, actual); + Assert.Equal(0, updateCount); + } + [Fact] public async Task NullSelectionDoesNotReportAttempt() { @@ -415,32 +420,6 @@ public async Task Dispatch_ConfiguredClientPreservesAndOverridesRequestOptions() Assert.Equal("request", requestOptions.ModelId); } - [Fact] - public async Task Policy_CanBufferMessagesBeforeDispatch() - { - IEnumerable? forwardedMessages = null; - IReadOnlyList? bufferedMessages = null; - using var inner = new TestChatClient - { - GetResponseAsyncCallback = (messages, _, _) => - { - forwardedMessages = messages; - return Task.FromResult(new ChatResponse()); - }, - }; - var messages = new SingleUseMessageEnumerable([new(ChatRole.User, "original")]); - using var router = new DelegatingTestRouter(context => - { - bufferedMessages = context.BufferMessages(); - return inner; - }); - - _ = await router.GetResponseAsync(messages); - - Assert.Same(bufferedMessages, forwardedMessages); - Assert.Equal(1, messages.EnumerationCount); - } - [Fact] public async Task Failure_Propagates() { @@ -1363,46 +1342,6 @@ public async Task SemanticRouting_SelectsBestProfileAndCachesIndex() Assert.Equal(1, profileBatches); } - [Fact] - public async Task SemanticRouting_MaterializesMessagesBeforeSelection() - { - var vectors = new Dictionary - { - ["code profile"] = [1, 0], - ["debug this code"] = [1, 0], - }; - using var generator = new TestEmbeddingGenerator - { - GenerateAsyncCallback = (values, _, _) => - Task.FromResult(new GeneratedEmbeddings>( - [.. values.Select(input => new Embedding(vectors[input]))])), - }; - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "code")); - using var selected = new TestChatClient - { - GetResponseAsyncCallback = (messages, _, _) => - { - Assert.Equal("debug this code", Assert.Single(messages).Text); - return Task.FromResult(expected); - }, - }; - var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) - { - [selected] = ["code profile"], - }; - using var router = new SemanticRoutingChatClient( - generator, - profiles, - defaultClient: selected, - leaveOpen: true); - var messages = new SingleUseMessageEnumerable([new(ChatRole.User, "debug this code")]); - - ChatResponse response = await router.GetResponseAsync(messages); - - Assert.Same(expected, response); - Assert.Equal(1, messages.EnumerationCount); - } - [Theory] [InlineData(SemanticRoutingChatClient.ScoreAggregation.Mean, "code")] [InlineData(SemanticRoutingChatClient.ScoreAggregation.Sum, "writing")] @@ -1613,35 +1552,6 @@ public async Task OrderedFailover_StateIsScopedToOneRequest() Assert.Equal(2, firstCalls); } - [Fact] - public async Task OrderedFailover_BuffersMessagesAcrossAttempts() - { - var messages = new SingleUseMessageEnumerable([new(ChatRole.User, "hi")]); - using var first = new TestChatClient - { - GetResponseAsyncCallback = (forwarded, _, _) => - { - Assert.Equal("hi", Assert.Single(forwarded).Text); - throw new InvalidOperationException("failed"); - }, - }; - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); - using var second = new TestChatClient - { - GetResponseAsyncCallback = (forwarded, _, _) => - { - Assert.Equal("hi", Assert.Single(forwarded).Text); - return Task.FromResult(expected); - }, - }; - using var client = new OrderedFailoverChatClient([first, second]); - - ChatResponse response = await client.GetResponseAsync(messages); - - Assert.Same(expected, response); - Assert.Equal(1, messages.EnumerationCount); - } - [Fact] public async Task OrderedFailover_ConcurrentRequestsHaveIndependentState() { @@ -1929,23 +1839,6 @@ protected override ValueTask SelectClientAsync( new(_select(context)); } - private sealed class SingleUseMessageEnumerable(IEnumerable messages) : IEnumerable - { - public int EnumerationCount { get; private set; } - - public IEnumerator GetEnumerator() - { - if (++EnumerationCount > 1) - { - throw new InvalidOperationException("The messages can only be enumerated once."); - } - - return messages.GetEnumerator(); - } - - System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() => GetEnumerator(); - } - private sealed class DelegatingFailoverTestRouter : FailoverChatClient { private readonly Func _select; From a5767730d9574ac6208fe69d56270bcb0c32bd53 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Thu, 30 Jul 2026 11:08:39 -0700 Subject: [PATCH 10/15] Simplify ordered failover state Use ConcurrentDictionary atomic operations instead of explicit lock blocks while preserving the narrow request-state lifetime. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/OrderedFailoverChatClient.cs | 56 ++++++++----------- .../ChatRouting/RoutingChatClientTests.cs | 2 +- 2 files changed, 23 insertions(+), 35 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs index f5ee88eed1c..8794c22f4c9 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -2,6 +2,7 @@ // 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.CodeAnalysis; using System.Runtime.ExceptionServices; @@ -28,8 +29,7 @@ public sealed class OrderedFailoverChatClient : FailoverChatClient { private readonly bool _leaveOpen; private readonly IChatClient[] _clients; - private readonly Dictionary _requestStates = []; - private readonly object _sync = new(); + private readonly ConcurrentDictionary _requestStates = new(); private bool _disposed; /// Initializes a new instance of the class. @@ -75,16 +75,9 @@ protected override ValueTask SelectClientAsync( CancellationToken cancellationToken) { _ = cancellationToken; - RequestState? state; - - lock (_sync) + if (!_requestStates.TryRemove(context, out RequestState? state)) { - if (!_requestStates.TryGetValue(context, out state)) - { - return new(_clients[0]); - } - - _ = _requestStates.Remove(context); + return new(_clients[0]); } if (state.ClientIndex < _clients.Length) @@ -105,30 +98,28 @@ protected override ValueTask OnRoutingUpdateAsync( { _ = cancellationToken; - lock (_sync) + if (isTerminal) { - if (isTerminal) - { - _ = _requestStates.Remove(context); - return default; - } - - if (attempt.Exception is null) - { - _ = _requestStates.Remove(context); - throw new InvalidOperationException("A nonterminal routing update requires a failed client invocation."); - } + _ = _requestStates.TryRemove(context, out _); + return default; + } - int clientIndex = IndexOfClient(attempt.Client); - if (clientIndex < 0) - { - _ = _requestStates.Remove(context); - throw new InvalidOperationException("The invocation did not use a configured ordered failover client."); - } + if (attempt.Exception is null) + { + _ = _requestStates.TryRemove(context, out _); + throw new InvalidOperationException("A nonterminal routing update requires a failed client invocation."); + } - _requestStates[context] = new(clientIndex + 1, attempt.Exception); + int clientIndex = IndexOfClient(attempt.Client); + if (clientIndex < 0) + { + _ = _requestStates.TryRemove(context, out _); + throw new InvalidOperationException("The invocation did not use a configured ordered failover client."); } + // Selection immediately removes this state before invoking the next client. + _requestStates[context] = new(clientIndex + 1, attempt.Exception); + return default; } @@ -143,10 +134,7 @@ protected override void Dispose(bool disposing) _disposed = true; try { - lock (_sync) - { - _requestStates.Clear(); - } + _requestStates.Clear(); if (disposing && !_leaveOpen) { diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs index e2547deb3de..52811f8de07 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -1789,7 +1789,7 @@ private static int GetOrderedFailoverRequestStateCount(OrderedFailoverChatClient object requestStates = typeof(OrderedFailoverChatClient) .GetField("_requestStates", System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic)! .GetValue(client)!; - return ((System.Collections.IDictionary)requestStates).Count; + return (int)requestStates.GetType().GetProperty("Count")!.GetValue(requestStates)!; } private static async IAsyncEnumerable YieldUpdates(params string[] texts) From 87d382d6339f84527dbe6239cd18f6d5c2cba83a Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Thu, 30 Jul 2026 19:42:43 -0700 Subject: [PATCH 11/15] Isolate routing attempt options Keep client selection simple while preventing selector and failed-client option mutations from affecting downstream attempts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 10 +- .../ChatRouting/RoutingContext.cs | 11 +-- .../ChatRouting/FailoverChatClient.cs | 6 +- .../ChatRouting/RoutingChatClientTests.cs | 94 ++++++++++++++++++- 4 files changed, 106 insertions(+), 15 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index f4430515294..91625f66d85 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -39,9 +39,9 @@ public static RoutingChatClient Create( /// The cancellation token supplied for the request. /// The client to invoke. /// - /// Implementations may adjust the request-local for dynamic request - /// shaping. Stable route-specific behavior should generally be attached to the returned client. Exceptions from - /// this method propagate to the caller. + /// The context contains a snapshot of the caller's options for use during selection. Changes to that snapshot do + /// not affect the options passed to the selected client. Client-specific behavior should generally be attached to + /// the returned client. Exceptions from this method propagate to the caller. /// protected abstract ValueTask SelectClientAsync( RoutingContext context, @@ -59,7 +59,7 @@ public virtual async Task GetResponseAsync( throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); return await client.GetResponseAsync( context.Messages, - context.ChatOptions, + options?.Clone(), cancellationToken).ConfigureAwait(false); } @@ -74,7 +74,7 @@ public virtual async IAsyncEnumerable GetStreamingResponseAs IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); await foreach (ChatResponseUpdate update in - client.GetStreamingResponseAsync(context.Messages, context.ChatOptions, cancellationToken) + client.GetStreamingResponseAsync(context.Messages, options?.Clone(), cancellationToken) .WithCancellation(cancellationToken) .ConfigureAwait(false)) { diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs index 10865674bd9..6de0b9e25d1 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs @@ -15,8 +15,8 @@ namespace Microsoft.Extensions.AI; /// started from the sequence returned by . /// /// -/// is cloned when the context is created, so request-specific changes do not mutate the -/// caller's instance. +/// is cloned when the context is created. It is provided for client selection and is +/// independent of both the caller's instance and the options passed to the selected client. /// /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] @@ -39,11 +39,10 @@ public RoutingContext( /// public IEnumerable Messages { get; } - /// Gets the request-local options supplied to client selection and the selected client. + /// Gets a snapshot of the request options supplied to client selection. /// - /// When the caller supplies options, this is a clone that selectors may adjust without mutating the caller's - /// instance. Changes are observed by the selected client and subsequent failover attempts. Stable route-specific - /// defaults should generally be attached to the returned client rather than reapplied during selection. + /// Changes do not affect the caller's instance or the options passed to the selected client. Client-specific + /// behavior should generally be attached to the returned client. /// public ChatOptions? ChatOptions { get; } diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index eab0b51eebc..78d6cabd606 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -109,6 +109,7 @@ public sealed override async Task GetResponseAsync( IChatClient selectedClient = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + ChatOptions? attemptOptions = options?.Clone(); attemptCount++; ChatResponse? response = null; @@ -119,7 +120,7 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? { response = await selectedClient.GetResponseAsync( context.Messages, - context.ChatOptions, + attemptOptions, cancellationToken).ConfigureAwait(false); } catch (Exception ex) @@ -178,6 +179,7 @@ public sealed override async IAsyncEnumerable GetStreamingRe IChatClient selectedClient = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + ChatOptions? attemptOptions = options?.Clone(); attemptCount++; bool reachedAttemptLimit = maximumAttempts is int limit && attemptCount >= limit; @@ -191,7 +193,7 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? enumerator = selectedClient .GetStreamingResponseAsync( context.Messages, - context.ChatOptions, + attemptOptions, cancellationToken) .GetAsyncEnumerator(cancellationToken); diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs index 52811f8de07..7c671eb5d29 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -67,8 +67,9 @@ public async Task Create_SelectsClientForRequest() Assert.Same(expected, response); Assert.Same(messages, observedContext!.Messages); Assert.NotSame(options, observedContext.ChatOptions); - Assert.Same(observedContext.ChatOptions, forwardedOptions); - Assert.Equal("selected", forwardedOptions!.ModelId); + Assert.NotSame(observedContext.ChatOptions, forwardedOptions); + Assert.Equal("selected", observedContext.ChatOptions!.ModelId); + Assert.Equal("request", forwardedOptions!.ModelId); Assert.Equal("request", options.ModelId); Assert.Equal(cancellationSource.Token, observedToken); Assert.Equal(1, selectionCount); @@ -420,6 +421,51 @@ public async Task Dispatch_ConfiguredClientPreservesAndOverridesRequestOptions() Assert.Equal("request", requestOptions.ModelId); } + [Fact] + public async Task Failover_UsesFreshRequestOptionsForEachAttempt() + { + var requestOptions = new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }; + var invokedOptions = new List(); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + options!.ModelId = "changed by first client"; + throw new InvalidOperationException("failed"); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + return Task.FromResult(expected); + }, + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? first : second); + + ChatResponse response = await router.GetResponseAsync( + [new(ChatRole.User, "hi")], + requestOptions); + + Assert.Same(expected, response); + Assert.Equal(2, selections); + Assert.Equal(2, invokedOptions.Count); + Assert.NotSame(invokedOptions[0], invokedOptions[1]); + Assert.Equal("changed by first client", invokedOptions[0].ModelId); + Assert.Equal("request", invokedOptions[1].ModelId); + Assert.Equal("caller", invokedOptions[1].Instructions); + Assert.Equal("request", requestOptions.ModelId); + } + [Fact] public async Task Failure_Propagates() { @@ -747,6 +793,50 @@ public async Task Streaming_FallsBackBeforeFirstUpdate() Assert.Equal(2, selections); } + [Fact] + public async Task Streaming_FailoverUsesFreshRequestOptionsForEachAttempt() + { + var requestOptions = new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }; + var invokedOptions = new List(); + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + options!.ModelId = "changed by first client"; + return ThrowingStream("failed"); + }, + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + return YieldUpdates("ok"); + }, + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? first : second); + + List updates = await CollectAsync( + router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + requestOptions)); + + Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal(2, selections); + Assert.NotSame(invokedOptions[0], invokedOptions[1]); + Assert.Equal("changed by first client", invokedOptions[0].ModelId); + Assert.Equal("request", invokedOptions[1].ModelId); + Assert.Equal("caller", invokedOptions[1].Instructions); + Assert.Equal("request", requestOptions.ModelId); + } + [Fact] public async Task Streaming_CancellationDuringUpdateIsObservedByNextAttempt() { From 72598e1864413f29663ce616e5e883f3a85d6ea0 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Fri, 31 Jul 2026 09:39:48 -0700 Subject: [PATCH 12/15] Clarify routing identity semantics Document client-reference identity and keep semantic scoring parameters together in the constructor signature. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 4 ++++ .../ChatRouting/SemanticRoutingChatClient.cs | 17 +++++++++-------- .../Microsoft.Extensions.AI.json | 2 +- .../ChatRouting/RoutingChatClientTests.cs | 4 ++-- 4 files changed, 16 insertions(+), 11 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index 91625f66d85..493b20a4cec 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -18,6 +18,10 @@ namespace Microsoft.Extensions.AI; /// /// Derived classes implement to supply one client for each request. The selected /// client is invoked once, and its response or failure is propagated to the caller. +/// The exact client instance represents the selected routing identity. Caller-supplied options may vary per +/// invocation, but they are ephemeral and do not participate in that identity. Applications can use distinct +/// configured client wrappers when configurations require distinct routing identities, while custom policies may +/// maintain separate application-specific grouping keys. /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] public abstract class RoutingChatClient : IChatClient diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs index 4d12214f97c..c3c47be4af9 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -22,8 +22,9 @@ namespace Microsoft.Extensions.AI; /// configured threshold. /// /// -/// The configured clients are used as stable routing identities. By default this instance owns the clients and -/// embedding generator and disposes them when it is disposed. +/// The configured client instances are used as stable routing identities and are distinguished by reference. +/// Per-call options do not participate in that identity. By default this instance owns the clients and embedding +/// generator and disposes them when it is disposed. /// /// /// The example-utterance routing approach is inspired by @@ -60,15 +61,15 @@ public enum ScoreAggregation /// The example utterances associated with each client. /// The client selected when no profile satisfies . /// The minimum aggregated score required to select a profiled client. - /// - /// to leave the configured clients and embedding generator open when this instance is - /// disposed; otherwise, . The default is . - /// /// /// The number of highest-scoring profile utterances, across all clients, whose scores are aggregated. /// The default is 1. /// /// The method used to aggregate matching profile scores for each client. + /// + /// to leave the configured clients and embedding generator open when this instance is + /// disposed; otherwise, . The default is . + /// /// /// , , or /// is . @@ -86,9 +87,9 @@ public SemanticRoutingChatClient( IReadOnlyDictionary> clientProfiles, IChatClient defaultClient, float scoreThreshold = 0.3f, - bool leaveOpen = false, int topK = 1, - ScoreAggregation scoreAggregation = ScoreAggregation.Mean) + ScoreAggregation scoreAggregation = ScoreAggregation.Mean, + bool leaveOpen = false) { _embeddingGenerator = Throw.IfNull(embeddingGenerator); _ = Throw.IfNull(clientProfiles); diff --git a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json index e55cd6b9619..564b3369241 100644 --- a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json +++ b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json @@ -1610,7 +1610,7 @@ "Stage": "Experimental", "Methods": [ { - "Member": "Microsoft.Extensions.AI.SemanticRoutingChatClient.SemanticRoutingChatClient(Microsoft.Extensions.AI.IEmbeddingGenerator> embeddingGenerator, System.Collections.Generic.IReadOnlyDictionary> clientProfiles, Microsoft.Extensions.AI.IChatClient defaultClient, float scoreThreshold = 0.3f, bool leaveOpen = false, int topK = 1, Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation scoreAggregation = Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.Mean);", + "Member": "Microsoft.Extensions.AI.SemanticRoutingChatClient.SemanticRoutingChatClient(Microsoft.Extensions.AI.IEmbeddingGenerator> embeddingGenerator, System.Collections.Generic.IReadOnlyDictionary> clientProfiles, Microsoft.Extensions.AI.IChatClient defaultClient, float scoreThreshold = 0.3f, int topK = 1, Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation scoreAggregation = Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.Mean, bool leaveOpen = false);", "Stage": "Experimental" }, { diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs index 7c671eb5d29..68a56324306 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs @@ -1472,9 +1472,9 @@ public async Task SemanticRouting_AggregatesGlobalTopKByClient( profiles, defaultClient: code, scoreThreshold: scoreAggregation == SemanticRoutingChatClient.ScoreAggregation.Sum ? 1.5f : 0.3f, - leaveOpen: true, topK: 3, - scoreAggregation: scoreAggregation); + scoreAggregation: scoreAggregation, + leaveOpen: true); ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "query")]); From 8b79ae488c0ab850b7b0dee39ba5746341cc6bb6 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Fri, 31 Jul 2026 12:13:11 -0700 Subject: [PATCH 13/15] Address routing review feedback Align routing code with MEAI formatting and documentation conventions, simplify disposal, and split routing tests by component and assembly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 18 +- .../ChatRouting/RoutingContext.cs | 7 +- .../ChatRouting/FailoverChatClient.cs | 36 +- .../ChatRouting/FailoverChatClientAttempt.cs | 3 +- .../ChatRouting/OrderedFailoverChatClient.cs | 38 +- .../ChatRouting/SemanticRoutingChatClient.cs | 43 +- .../ChatRouting/RoutingChatClientTests.cs | 181 +++++ .../ChatRouting/RoutingContextTests.cs | 26 + ...entTests.cs => FailoverChatClientTests.cs} | 753 +----------------- .../OrderedFailoverChatClientTests.cs | 414 ++++++++++ .../SemanticRoutingChatClientTests.cs | 285 +++++++ 11 files changed, 978 insertions(+), 826 deletions(-) create mode 100644 test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs create mode 100644 test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingContextTests.cs rename test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/{RoutingChatClientTests.cs => FailoverChatClientTests.cs} (63%) create mode 100644 test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs create mode 100644 test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/SemanticRoutingChatClientTests.cs diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index 493b20a4cec..58a1f050a2f 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -35,6 +35,7 @@ public static RoutingChatClient Create( Func> clientSelector) { _ = Throw.IfNull(clientSelector); + return new CallbackRoutingChatClient(clientSelector); } @@ -43,9 +44,8 @@ public static RoutingChatClient Create( /// The cancellation token supplied for the request. /// The client to invoke. /// - /// The context contains a snapshot of the caller's options for use during selection. Changes to that snapshot do - /// not affect the options passed to the selected client. Client-specific behavior should generally be attached to - /// the returned client. Exceptions from this method propagate to the caller. + /// Client-specific behavior should generally be attached to the returned client. Exceptions from this method + /// propagate to the caller. /// protected abstract ValueTask SelectClientAsync( RoutingContext context, @@ -53,14 +53,14 @@ protected abstract ValueTask SelectClientAsync( /// public virtual async Task GetResponseAsync( - IEnumerable messages, - ChatOptions? options = null, - CancellationToken cancellationToken = default) + IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) { _ = Throw.IfNull(messages); + var context = new RoutingContext(messages, options); IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + return await client.GetResponseAsync( context.Messages, options?.Clone(), @@ -69,14 +69,15 @@ public virtual async Task GetResponseAsync( /// public virtual async IAsyncEnumerable GetStreamingResponseAsync( - IEnumerable messages, - ChatOptions? options = null, + IEnumerable messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) { _ = Throw.IfNull(messages); + var context = new RoutingContext(messages, options); IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + await foreach (ChatResponseUpdate update in client.GetStreamingResponseAsync(context.Messages, options?.Clone(), cancellationToken) .WithCancellation(cancellationToken) @@ -90,6 +91,7 @@ public virtual async IAsyncEnumerable GetStreamingResponseAs public virtual object? GetService(Type serviceType, object? serviceKey = null) { _ = Throw.IfNull(serviceType); + return serviceKey is null && serviceType.IsInstanceOfType(this) ? this : null; } diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs index 6de0b9e25d1..edd5bb7760f 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs @@ -24,12 +24,14 @@ public class RoutingContext { /// Initializes a new instance of the class. /// The messages to route. - /// The options to clone for this request, or . + /// The options associated with this context. public RoutingContext( IEnumerable messages, ChatOptions? chatOptions) { - Messages = Throw.IfNull(messages); + _ = Throw.IfNull(messages); + + Messages = messages; ChatOptions = chatOptions?.Clone(); } @@ -45,5 +47,4 @@ public RoutingContext( /// behavior should generally be attached to the returned client. /// public ChatOptions? ChatOptions { get; } - } diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index 78d6cabd606..d825a8a054d 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -62,9 +62,9 @@ public int? MaximumAttemptsPerRequest /// Invoked after a client invocation completes, fails, or is abandoned. /// The request-specific inputs. - /// The completed client invocation. + /// The attempted client invocation. /// - /// if the base will not select another client after this callback completes successfully; + /// if the base will not select another client after this method returns successfully; /// otherwise, . /// /// The cancellation token supplied for the request. @@ -76,9 +76,9 @@ public int? MaximumAttemptsPerRequest /// . /// /// - /// Every client invocation produces one update if the callback completes successfully. Selection failures are not - /// reported. A selector that retains request-scoped state must release it before throwing; a request may therefore - /// end without a terminal update when selection fails. + /// This method is invoked once after each selected-client invocation, whether it completes, fails, or is abandoned. + /// Selection failures are not reported. A selector that retains request-scoped state must release it before + /// throwing; a request may therefore end without a terminal update when selection fails. /// /// /// Exceptions from this method propagate to the caller. A terminal update exception replaces the response or @@ -95,11 +95,10 @@ protected virtual ValueTask OnRoutingUpdateAsync( /// public sealed override async Task GetResponseAsync( - IEnumerable messages, - ChatOptions? options = null, - CancellationToken cancellationToken = default) + IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) { _ = Throw.IfNull(messages); + var context = new RoutingContext(messages, options); int? maximumAttempts = MaximumAttemptsPerRequest; int attemptCount = 0; @@ -165,11 +164,11 @@ exception is null || /// public sealed override async IAsyncEnumerable GetStreamingResponseAsync( - IEnumerable messages, - ChatOptions? options = null, + IEnumerable messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) { _ = Throw.IfNull(messages); + var context = new RoutingContext(messages, options); int? maximumAttempts = MaximumAttemptsPerRequest; int attemptCount = 0; @@ -202,7 +201,7 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? catch (Exception ex) { stopwatch.Stop(); - Exception exception = (await DisposeAsync(enumerator, ex).ConfigureAwait(false))!; + Exception exception = (await DisposeEnumeratorAsync(enumerator, ex).ConfigureAwait(false))!; var attempt = new FailoverChatClientAttempt( selectedClient, exception, @@ -273,7 +272,8 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? } finally { - terminalException = await DisposeAsync(enumerator, terminalException).ConfigureAwait(false); + terminalException = + await DisposeEnumeratorAsync(enumerator, terminalException).ConfigureAwait(false); var attempt = new FailoverChatClientAttempt( selectedClient, @@ -320,13 +320,10 @@ await OnRoutingUpdateAsync( } } - [SuppressMessage( - "Resilience", - "EA0014:The async method doesn't support cancellation", - Justification = "IAsyncDisposable.DisposeAsync does not support cancellation.")] - private static async ValueTask DisposeAsync( - IAsyncDisposable? disposable, - Exception? exception) +#pragma warning disable EA0014 // IAsyncDisposable.DisposeAsync doesn't support cancellation. + private static async ValueTask DisposeEnumeratorAsync( + IAsyncDisposable? disposable, Exception? exception) +#pragma warning restore EA0014 { if (disposable is not null) { @@ -369,5 +366,4 @@ private static void Rethrow(Exception exception) { ExceptionDispatchInfo.Capture(exception).Throw(); } - } diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs index 04c32086a0a..0c4038ed8e0 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs @@ -9,7 +9,7 @@ namespace Microsoft.Extensions.AI; -/// Represents one client invocation performed by a . +/// Represents one client invocation attempt performed by a . /// /// /// A completed response has set to and @@ -37,7 +37,6 @@ internal FailoverChatClientAttempt( bool responseCompleted, bool outputCommitted) { - Debug.Assert(client is not null, "Expected a non-null invoked client."); Debug.Assert(duration >= TimeSpan.Zero, "Expected a non-negative duration."); Debug.Assert( timeToFirstUpdate is not { } ttfu || (ttfu >= TimeSpan.Zero && ttfu <= duration), diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs index 8794c22f4c9..7fcec7779bd 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -40,33 +40,37 @@ public sealed class OrderedFailoverChatClient : FailoverChatClient /// /// is . /// - /// is empty, contains , or contains the same client instance more than once. + /// is empty, contains , or contains the same client instance more + /// than once. /// public OrderedFailoverChatClient(IReadOnlyList clients, bool leaveOpen = false) { _ = Throw.IfNull(clients); - _leaveOpen = leaveOpen; - if (clients.Count == 0) + + IChatClient[] clientsSnapshot = [.. clients]; + if (clientsSnapshot.Length == 0) { Throw.ArgumentException(nameof(clients), "At least one client must be provided."); } - _clients = [.. clients]; - for (int i = 0; i < _clients.Length; i++) + for (int i = 0; i < clientsSnapshot.Length; i++) { - if (_clients[i] is null) + if (clientsSnapshot[i] is null) { Throw.ArgumentException(nameof(clients), "Clients must not contain null."); } for (int j = 0; j < i; j++) { - if (ReferenceEquals(_clients[j], _clients[i])) + if (ReferenceEquals(clientsSnapshot[j], clientsSnapshot[i])) { Throw.ArgumentException(nameof(clients), "Each client instance must be unique."); } } } + + _leaveOpen = leaveOpen; + _clients = clientsSnapshot; } /// @@ -75,6 +79,7 @@ protected override ValueTask SelectClientAsync( CancellationToken cancellationToken) { _ = cancellationToken; + if (!_requestStates.TryRemove(context, out RequestState? state)) { return new(_clients[0]); @@ -132,22 +137,17 @@ protected override void Dispose(bool disposing) } _disposed = true; - try - { - _requestStates.Clear(); + _requestStates.Clear(); - if (disposing && !_leaveOpen) + if (disposing && !_leaveOpen) + { + foreach (IChatClient client in _clients) { - foreach (IChatClient client in _clients) - { - client.Dispose(); - } + client.Dispose(); } } - finally - { - base.Dispose(disposing); - } + + base.Dispose(disposing); } private int IndexOfClient(IChatClient client) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs index c3c47be4af9..ff22411fff2 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -59,7 +59,9 @@ public enum ScoreAggregation /// Initializes a new instance of the class. /// The generator used to embed profile utterances and request text. /// The example utterances associated with each client. - /// The client selected when no profile satisfies . + /// + /// The client selected when no profile satisfies . + /// /// The minimum aggregated score required to select a profiled client. /// /// The number of highest-scoring profile utterances, across all clients, whose scores are aggregated. @@ -91,10 +93,9 @@ public SemanticRoutingChatClient( ScoreAggregation scoreAggregation = ScoreAggregation.Mean, bool leaveOpen = false) { - _embeddingGenerator = Throw.IfNull(embeddingGenerator); + _ = Throw.IfNull(embeddingGenerator); _ = Throw.IfNull(clientProfiles); _ = Throw.IfNull(defaultClient); - _leaveOpen = leaveOpen; if (clientProfiles.Count == 0) { @@ -120,6 +121,8 @@ public SemanticRoutingChatClient( Throw.ArgumentOutOfRangeException(nameof(scoreThreshold)); } + _embeddingGenerator = embeddingGenerator; + _leaveOpen = leaveOpen; _topK = topK; _scoreAggregation = scoreAggregation; _scoreThreshold = scoreThreshold; @@ -170,6 +173,7 @@ protected override async ValueTask SelectClientAsync( CancellationToken cancellationToken) { _ = Throw.IfNull(context); + string? query = context.Messages.LastOrDefault( static message => message.Role == ChatRole.User)?.Text; if (string.IsNullOrWhiteSpace(query)) @@ -277,31 +281,26 @@ protected override void Dispose(bool disposing) } _disposed = true; - try + if (disposing) { - if (disposing) + _indexGate.Dispose(); + if (!_leaveOpen) { - _indexGate.Dispose(); - if (!_leaveOpen) + foreach (IChatClient client in _clients) + { + client.Dispose(); + } + + if (!Array.Exists( + _clients, + client => ReferenceEquals(client, _embeddingGenerator))) { - foreach (IChatClient client in _clients) - { - client.Dispose(); - } - - if (!Array.Exists( - _clients, - client => ReferenceEquals(client, _embeddingGenerator))) - { - _embeddingGenerator.Dispose(); - } + _embeddingGenerator.Dispose(); } } } - finally - { - base.Dispose(disposing); - } + + base.Dispose(disposing); } private async Task EnsureIndexAsync(CancellationToken cancellationToken) diff --git a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs new file mode 100644 index 00000000000..b8b7d8439c3 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs @@ -0,0 +1,181 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Extensions.AI; + +public class RoutingChatClientTests +{ + [Fact] + public void Create_RejectsNullSelector() + { + Assert.Throws(() => RoutingChatClient.Create(null!)); + } + + [Fact] + public async Task Create_SelectsClientForRequest() + { + var messages = new ChatMessage[] { new(ChatRole.User, "hi") }; + var options = new ChatOptions { ModelId = "request" }; + using var cancellationSource = new CancellationTokenSource(); + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + ChatOptions? forwardedOptions = null; + using var selected = new TestChatClient + { + GetResponseAsyncCallback = (_, selectedOptions, _) => + { + forwardedOptions = selectedOptions; + return Task.FromResult(expected); + }, + }; + RoutingContext? observedContext = null; + CancellationToken observedToken = default; + int selectionCount = 0; + using RoutingChatClient router = RoutingChatClient.Create((context, cancellationToken) => + { + observedContext = context; + observedToken = cancellationToken; + selectionCount++; + context.ChatOptions!.ModelId = "selected"; + return new(selected); + }); + + ChatResponse response = await router.GetResponseAsync(messages, options, cancellationSource.Token); + + Assert.Same(expected, response); + Assert.Same(messages, observedContext!.Messages); + Assert.NotSame(options, observedContext.ChatOptions); + Assert.NotSame(observedContext.ChatOptions, forwardedOptions); + Assert.Equal("selected", observedContext.ChatOptions!.ModelId); + Assert.Equal("request", forwardedOptions!.ModelId); + Assert.Equal("request", options.ModelId); + Assert.Equal(cancellationSource.Token, observedToken); + Assert.Equal(1, selectionCount); + } + + [Fact] + public async Task Create_DoesNotReselectAfterFailure() + { + var expected = new InvalidOperationException("failed"); + using var selected = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw expected, + }; + int selectionCount = 0; + using RoutingChatClient router = RoutingChatClient.Create((_, _) => + { + selectionCount++; + return new(selected); + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(expected, actual); + Assert.Equal(1, selectionCount); + } + + [Fact] + public async Task Create_NullSelectionThrowsForNonStreamingAndStreaming() + { + using RoutingChatClient router = RoutingChatClient.Create((_, _) => new((IChatClient)null!)); + + InvalidOperationException nonStreaming = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + InvalidOperationException streaming = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Contains("SelectClientAsync", nonStreaming.Message, StringComparison.Ordinal); + Assert.Contains("SelectClientAsync", streaming.Message, StringComparison.Ordinal); + } + + [Fact] + public async Task Create_RejectsNullMessagesForNonStreamingAndStreaming() + { + using var selected = new TestChatClient(); + using RoutingChatClient router = RoutingChatClient.Create((_, _) => new(selected)); + + await Assert.ThrowsAsync(() => router.GetResponseAsync(null!)); + await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync(null!))); + } + + [Fact] + public async Task Create_DoesNotDisposeSelectedClient() + { + using var selected = new CountingDisposeClient(); + using (RoutingChatClient router = RoutingChatClient.Create((_, _) => new(selected))) + { + _ = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + } + + Assert.Equal(0, selected.DisposeCount); + } + + [Fact] + public void GetService_ReturnsSelfAndNullForUnknownOrKeyed() + { + using var client = new DelegatingTestRouter(_ => throw new NotSupportedException()); + + Assert.Same(client, client.GetService(typeof(DelegatingTestRouter))); + Assert.Same(client, client.GetService(typeof(RoutingChatClient))); + Assert.Same(client, client.GetService(typeof(IChatClient))); + Assert.Null(client.GetService(typeof(DelegatingTestRouter), serviceKey: "key")); + Assert.Null(client.GetService(typeof(string))); + } + + private static async Task> CollectAsync( + IAsyncEnumerable updates) + { + var result = new List(); + await foreach (ChatResponseUpdate update in updates) + { + result.Add(update); + } + + return result; + } + + private sealed class DelegatingTestRouter : RoutingChatClient + { + private readonly Func _select; + + public DelegatingTestRouter(Func select) + { + _select = select; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + new(_select(context)); + } + + private sealed class CountingDisposeClient : IChatClient + { + public int DisposeCount { get; private set; } + + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + Task.FromResult(new ChatResponse()); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + + public override bool Equals(object? obj) => obj is CountingDisposeClient; + + public override int GetHashCode() => 0; + } +} diff --git a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingContextTests.cs b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingContextTests.cs new file mode 100644 index 00000000000..d97195262a7 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingContextTests.cs @@ -0,0 +1,26 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using Xunit; + +namespace Microsoft.Extensions.AI; + +public class RoutingContextTests +{ + [Fact] + public void RoutingContext_CarriesRequestInputs() + { + var messages = new List { new(ChatRole.User, "initial") }; + var options = new ChatOptions { ModelId = "initial" }; + var context = new RoutingContext(messages, options); + + Assert.Same(messages, context.Messages); + Assert.NotSame(options, context.ChatOptions); + Assert.Equal("initial", context.ChatOptions!.ModelId); + context.ChatOptions.ModelId = "changed"; + Assert.Equal("initial", options.ModelId); + Assert.Throws(() => new RoutingContext(null!, options)); + } +} diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs similarity index 63% rename from test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs rename to test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs index 68a56324306..3c03108b665 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs @@ -3,7 +3,6 @@ using System; using System.Collections.Generic; -using System.Linq; using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; @@ -11,141 +10,8 @@ namespace Microsoft.Extensions.AI; -public class RoutingChatClientTests +public class FailoverChatClientTests { - [Fact] - public void RoutingContext_CarriesRequestInputs() - { - var messages = new List { new(ChatRole.User, "initial") }; - var options = new ChatOptions { ModelId = "initial" }; - var context = new RoutingContext(messages, options); - - Assert.Same(messages, context.Messages); - Assert.NotSame(options, context.ChatOptions); - Assert.Equal("initial", context.ChatOptions!.ModelId); - context.ChatOptions.ModelId = "changed"; - Assert.Equal("initial", options.ModelId); - Assert.Throws(() => new RoutingContext(null!, options)); - } - - [Fact] - public void Create_RejectsNullSelector() - { - Assert.Throws(() => RoutingChatClient.Create(null!)); - } - - [Fact] - public async Task Create_SelectsClientForRequest() - { - var messages = new ChatMessage[] { new(ChatRole.User, "hi") }; - var options = new ChatOptions { ModelId = "request" }; - using var cancellationSource = new CancellationTokenSource(); - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); - ChatOptions? forwardedOptions = null; - using var selected = new TestChatClient - { - GetResponseAsyncCallback = (_, selectedOptions, _) => - { - forwardedOptions = selectedOptions; - return Task.FromResult(expected); - }, - }; - RoutingContext? observedContext = null; - CancellationToken observedToken = default; - int selectionCount = 0; - using RoutingChatClient router = RoutingChatClient.Create((context, cancellationToken) => - { - observedContext = context; - observedToken = cancellationToken; - selectionCount++; - context.ChatOptions!.ModelId = "selected"; - return new(selected); - }); - - ChatResponse response = await router.GetResponseAsync(messages, options, cancellationSource.Token); - - Assert.Same(expected, response); - Assert.Same(messages, observedContext!.Messages); - Assert.NotSame(options, observedContext.ChatOptions); - Assert.NotSame(observedContext.ChatOptions, forwardedOptions); - Assert.Equal("selected", observedContext.ChatOptions!.ModelId); - Assert.Equal("request", forwardedOptions!.ModelId); - Assert.Equal("request", options.ModelId); - Assert.Equal(cancellationSource.Token, observedToken); - Assert.Equal(1, selectionCount); - } - - [Fact] - public async Task Create_DoesNotReselectAfterFailure() - { - var expected = new InvalidOperationException("failed"); - using var selected = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => throw expected, - }; - int selectionCount = 0; - using RoutingChatClient router = RoutingChatClient.Create((_, _) => - { - selectionCount++; - return new(selected); - }); - - InvalidOperationException actual = await Assert.ThrowsAsync( - () => router.GetResponseAsync([new(ChatRole.User, "hi")])); - - Assert.Same(expected, actual); - Assert.Equal(1, selectionCount); - } - - [Fact] - public async Task Create_NullSelectionThrowsForNonStreamingAndStreaming() - { - using RoutingChatClient router = RoutingChatClient.Create((_, _) => new((IChatClient)null!)); - - InvalidOperationException nonStreaming = await Assert.ThrowsAsync( - () => router.GetResponseAsync([new(ChatRole.User, "hi")])); - InvalidOperationException streaming = await Assert.ThrowsAsync( - () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); - - Assert.Contains("SelectClientAsync", nonStreaming.Message, StringComparison.Ordinal); - Assert.Contains("SelectClientAsync", streaming.Message, StringComparison.Ordinal); - } - - [Fact] - public async Task Create_RejectsNullMessagesForNonStreamingAndStreaming() - { - using var selected = new TestChatClient(); - using RoutingChatClient router = RoutingChatClient.Create((_, _) => new(selected)); - - await Assert.ThrowsAsync(() => router.GetResponseAsync(null!)); - await Assert.ThrowsAsync( - () => CollectAsync(router.GetStreamingResponseAsync(null!))); - } - - [Fact] - public async Task Create_DoesNotDisposeSelectedClient() - { - using var selected = new CountingDisposeClient(); - using (RoutingChatClient router = RoutingChatClient.Create((_, _) => new(selected))) - { - _ = await router.GetResponseAsync([new(ChatRole.User, "hi")]); - } - - Assert.Equal(0, selected.DisposeCount); - } - - [Fact] - public void GetService_ReturnsSelfAndNullForUnknownOrKeyed() - { - using var client = new DelegatingTestRouter(_ => throw new NotSupportedException()); - - Assert.Same(client, client.GetService(typeof(DelegatingTestRouter))); - Assert.Same(client, client.GetService(typeof(RoutingChatClient))); - Assert.Same(client, client.GetService(typeof(IChatClient))); - Assert.Null(client.GetService(typeof(DelegatingTestRouter), serviceKey: "key")); - Assert.Null(client.GetService(typeof(string))); - } - [Fact] public void MaximumAttemptsPerRequest_ValidatesValue() { @@ -1322,540 +1188,6 @@ public async Task Streaming_CallerStopsEarlyNotifiesHookAsIncomplete() Assert.False(observed.ResponseCompleted); } - [Fact] - public void SemanticRouting_RejectsInvalidConfiguration() - { - using var client = new TestChatClient(); - using var generator = new TestEmbeddingGenerator(); - var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) - { - [client] = ["profile"], - }; - - Assert.Throws(() => - new SemanticRoutingChatClient(null!, profiles, client)); - Assert.Throws(() => - new SemanticRoutingChatClient(generator, null!, client)); - Assert.Throws(() => - new SemanticRoutingChatClient(generator, profiles, null!)); - Assert.Throws(() => - new SemanticRoutingChatClient( - generator, - new Dictionary>(), - client)); - Assert.Throws(() => - new SemanticRoutingChatClient( - generator, - new Dictionary> - { - [client] = [], - }, - client)); - Assert.Throws(() => - new SemanticRoutingChatClient( - generator, - new Dictionary> - { - [client] = [" "], - }, - client)); - Assert.Throws(() => - new SemanticRoutingChatClient(generator, profiles, client, scoreThreshold: 1.1f)); - Assert.Throws(() => - new SemanticRoutingChatClient(generator, profiles, client, topK: 0)); - Assert.Throws(() => - new SemanticRoutingChatClient( - generator, - profiles, - client, - scoreAggregation: (SemanticRoutingChatClient.ScoreAggregation)(-1))); - Assert.Throws(() => - new SemanticRoutingChatClient( - generator, - profiles, - client, - scoreThreshold: 2.1f, - topK: 2, - scoreAggregation: SemanticRoutingChatClient.ScoreAggregation.Sum)); - } - - [Fact] - public async Task SemanticRouting_SelectsBestProfileAndCachesIndex() - { - var vectors = new Dictionary - { - ["code profile"] = [1, 0], - ["writing profile"] = [0, 1], - ["debug this code"] = [1, 0], - }; - int profileBatches = 0; - using var generator = new TestEmbeddingGenerator - { - GenerateAsyncCallback = (values, _, _) => - { - string[] inputs = [.. values]; - if (inputs.Length > 1) - { - profileBatches++; - } - - return Task.FromResult(new GeneratedEmbeddings>( - [.. inputs.Select(input => new Embedding(vectors[input]))])); - }, - }; - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "code")); - using var code = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), - }; - using var writing = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "writing"))), - }; - var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) - { - [code] = ["code profile"], - [writing] = ["writing profile"], - }; - using var router = new SemanticRoutingChatClient( - generator, - profiles, - defaultClient: writing, - leaveOpen: true); - - ChatResponse first = await router.GetResponseAsync([new(ChatRole.User, "debug this code")]); - ChatResponse second = await router.GetResponseAsync([new(ChatRole.User, "debug this code")]); - - Assert.Same(expected, first); - Assert.Same(expected, second); - Assert.Equal(1, profileBatches); - } - - [Theory] - [InlineData(SemanticRoutingChatClient.ScoreAggregation.Mean, "code")] - [InlineData(SemanticRoutingChatClient.ScoreAggregation.Sum, "writing")] - public async Task SemanticRouting_AggregatesGlobalTopKByClient( - SemanticRoutingChatClient.ScoreAggregation scoreAggregation, - string expectedResponse) - { - var vectors = new Dictionary - { - ["code"] = [1, 0], - ["writing one"] = [0.8f, 0.6f], - ["writing two"] = [0.8f, -0.6f], - ["query"] = [1, 0], - }; - using var generator = new TestEmbeddingGenerator - { - GenerateAsyncCallback = (values, _, _) => - Task.FromResult(new GeneratedEmbeddings>( - [.. values.Select(input => new Embedding(vectors[input]))])), - }; - using var code = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "code"))), - }; - using var writing = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "writing"))), - }; - var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) - { - [code] = ["code"], - [writing] = ["writing one", "writing two"], - }; - using var router = new SemanticRoutingChatClient( - generator, - profiles, - defaultClient: code, - scoreThreshold: scoreAggregation == SemanticRoutingChatClient.ScoreAggregation.Sum ? 1.5f : 0.3f, - topK: 3, - scoreAggregation: scoreAggregation, - leaveOpen: true); - - ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "query")]); - - Assert.Equal(expectedResponse, response.Text); - } - - [Fact] - public async Task SemanticRouting_UsesDefaultBelowThreshold() - { - var vectors = new Dictionary - { - ["code profile"] = [1, 0], - ["unrelated query"] = [0, 1], - }; - using var generator = new TestEmbeddingGenerator - { - GenerateAsyncCallback = (values, _, _) => - Task.FromResult(new GeneratedEmbeddings>( - [.. values.Select(input => new Embedding(vectors[input]))])), - }; - using var profiled = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "profiled"))), - }; - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "default")); - using var defaultClient = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), - }; - var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) - { - [profiled] = ["code profile"], - }; - using var router = new SemanticRoutingChatClient( - generator, - profiles, - defaultClient, - scoreThreshold: 0.5f, - leaveOpen: true); - - ChatResponse response = - await router.GetResponseAsync([new(ChatRole.User, "unrelated query")]); - - Assert.Same(expected, response); - } - - [Fact] - public void SemanticRouting_DisposesOwnedResourcesOnce() - { -#pragma warning disable CA2000 // Dispose objects before losing scope - var generator = new CountingEmbeddingGenerator(); - var profiled = new CountingDisposeClient(); - var defaultClient = new CountingDisposeClient(); - var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) - { - [profiled] = ["profile"], - }; - var router = new SemanticRoutingChatClient(generator, profiles, defaultClient); -#pragma warning restore CA2000 - - router.Dispose(); - router.Dispose(); - - Assert.Equal(1, generator.DisposeCount); - Assert.Equal(1, profiled.DisposeCount); - Assert.Equal(1, defaultClient.DisposeCount); - } - - [Fact] - public void OrderedFailover_RejectsMissingClients() - { - using var inner = new TestChatClient(); - - Assert.Throws(() => new OrderedFailoverChatClient(null!)); - Assert.Throws(() => new OrderedFailoverChatClient([])); - Assert.Throws(() => new OrderedFailoverChatClient([inner, null!])); - Assert.Throws(() => new OrderedFailoverChatClient([inner, inner])); - } - - [Fact] - public async Task OrderedFailover_TriesClientsInOrderAndSkipsFailures() - { - var calls = new List(); - using var first = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - { - calls.Add("first"); - throw new InvalidOperationException("first failed"); - }, - }; - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); - using var second = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - { - calls.Add("second"); - return Task.FromResult(expected); - }, - }; - using var client = new OrderedFailoverChatClient([first, second]); - - ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); - - Assert.Same(expected, response); - Assert.Equal(["first", "second"], calls); - } - - [Fact] - public async Task OrderedFailover_DistinguishesValueEqualClients() - { - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); - using var first = new ValueEqualChatClient( - () => throw new InvalidOperationException("failed")); - using var second = new ValueEqualChatClient( - () => Task.FromResult(expected)); - using var client = new OrderedFailoverChatClient([first, second]); - - ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); - - Assert.Same(expected, response); - } - - [Fact] - public async Task OrderedFailover_ExhaustionRethrowsLastFailure() - { - using var first = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("first"), - }; - using var second = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("second"), - }; - using var client = new OrderedFailoverChatClient([first, second]); - - InvalidOperationException exception = await Assert.ThrowsAsync( - () => client.GetResponseAsync([new(ChatRole.User, "hi")])); - - Assert.Equal("second", exception.Message); - } - - [Fact] - public async Task OrderedFailover_StateIsScopedToOneRequest() - { - int firstCalls = 0; - using var first = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - { - firstCalls++; - throw new InvalidOperationException("failed"); - }, - }; - using var second = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => Task.FromResult(new ChatResponse()), - }; - using var client = new OrderedFailoverChatClient([first, second]); - - _ = await client.GetResponseAsync([new(ChatRole.User, "one")]); - _ = await client.GetResponseAsync([new(ChatRole.User, "two")]); - - Assert.Equal(2, firstCalls); - } - - [Fact] - public async Task OrderedFailover_ConcurrentRequestsHaveIndependentState() - { - const int RequestCount = 10; - int firstCalls = 0; - int secondCalls = 0; - var allStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - using var first = new TestChatClient - { - GetResponseAsyncCallback = async (_, _, _) => - { - if (Interlocked.Increment(ref firstCalls) == RequestCount) - { - allStarted.SetResult(true); - } - - await allStarted.Task; - throw new InvalidOperationException("failed"); - }, - }; - using var second = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - { - Interlocked.Increment(ref secondCalls); - return Task.FromResult(new ChatResponse()); - }, - }; - using var client = new OrderedFailoverChatClient([first, second]); - - Task[] requests = - [ - .. Enumerable.Range(0, RequestCount).Select(index => - client.GetResponseAsync([new(ChatRole.User, index.ToString())])), - ]; - _ = await Task.WhenAll(requests); - - Assert.Equal(RequestCount, firstCalls); - Assert.Equal(RequestCount, secondCalls); - } - - [Fact] - public async Task OrderedFailover_DoesNotRetainStateWhileStreaming() - { - using var first = new TestChatClient - { - GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), - }; - using var second = new TestChatClient - { - GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("first", "second"), - }; - 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); - Assert.Equal(0, GetOrderedFailoverRequestStateCount(client)); - } - - [Fact] - public async Task OrderedFailover_SnapshotsConfiguredClients() - { - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); - using var inner = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), - }; - var clients = new List { inner }; - using var failover = new OrderedFailoverChatClient(clients, leaveOpen: true); - clients.Clear(); - - ChatResponse response = await failover.GetResponseAsync([new(ChatRole.User, "hi")]); - - Assert.Same(expected, response); - } - - [Fact] - public async Task OrderedFailover_UsesEachConfiguredClient() - { - var seenOptions = new List<(string? modelId, string? instructions)>(); - using var firstInner = new TestChatClient - { - GetResponseAsyncCallback = (_, options, _) => - { - seenOptions.Add((options?.ModelId, options?.Instructions)); - throw new InvalidOperationException("failed"); - }, - }; - using var first = new ConfigureOptionsChatClient( - firstInner, - options => options.ModelId = "first"); - using var secondInner = new TestChatClient - { - GetResponseAsyncCallback = (_, options, _) => - { - seenOptions.Add((options?.ModelId, options?.Instructions)); - return Task.FromResult(new ChatResponse()); - }, - }; - using var second = new ConfigureOptionsChatClient( - secondInner, - options => options.ModelId = "second"); - using var client = new OrderedFailoverChatClient( - [first, second], - leaveOpen: true); - - _ = await client.GetResponseAsync( - [new(ChatRole.User, "hi")], - new ChatOptions - { - Instructions = "caller", - ModelId = "request", - }); - - Assert.Equal( - [("first", "caller"), ("second", "caller")], - seenOptions); - } - - [Fact] - public async Task OrderedFailover_StreamingFallsBackBeforeFirstUpdate() - { - using var first = new TestChatClient - { - GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), - }; - using var second = new TestChatClient - { - GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), - }; - using var client = new OrderedFailoverChatClient([first, second]); - - List updates = - await CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); - - Assert.Equal("ok", Assert.Single(updates).Text); - } - - [Fact] - public async Task OrderedFailover_CancellationDoesNotFallback() - { - using var cancellationSource = new CancellationTokenSource(); - cancellationSource.Cancel(); - bool secondCalled = false; - using var first = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), - }; - using var second = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => - { - secondCalled = true; - return Task.FromResult(new ChatResponse()); - }, - }; - using var client = new OrderedFailoverChatClient([first, second]); - - await Assert.ThrowsAnyAsync( - () => client.GetResponseAsync( - [new(ChatRole.User, "hi")], - cancellationToken: cancellationSource.Token)); - - Assert.False(secondCalled); - } - - [Fact] - public void OrderedFailover_DisposesEachClientOnce() - { -#pragma warning disable CA2000 // Dispose objects before losing scope - var shared = new CountingDisposeClient(); - var other = new CountingDisposeClient(); - var client = new OrderedFailoverChatClient([shared, other]); -#pragma warning restore CA2000 - - client.Dispose(); - client.Dispose(); - - Assert.Equal(1, shared.DisposeCount); - Assert.Equal(1, other.DisposeCount); - } - - [Fact] - public void OrderedFailover_LeaveOpenDoesNotDisposeInnerClients() - { -#pragma warning disable CA2000 // Dispose objects before losing scope - var inner = new CountingDisposeClient(); - var client = new OrderedFailoverChatClient([inner], leaveOpen: true); -#pragma warning restore CA2000 - - client.Dispose(); - - Assert.Equal(0, inner.DisposeCount); - inner.Dispose(); - } - - [Fact] - public async Task NestedRouters_ReturnLeafResponse() - { - ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); - using var leaf = new TestChatClient - { - GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), - }; - using var inner = new OrderedFailoverChatClient([leaf]); - using var outer = new OrderedFailoverChatClient([inner]); - - ChatResponse response = await outer.GetResponseAsync([new(ChatRole.User, "hi")]); - - Assert.Same(expected, response); - } - private static async Task> CollectAsync( IAsyncEnumerable updates) { @@ -1874,14 +1206,6 @@ private static async Task ConsumeOneAsync(IAsyncEnumerable u Assert.True(await enumerator.MoveNextAsync()); } - 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 static async IAsyncEnumerable YieldUpdates(params string[] texts) { foreach (string text in texts) @@ -2063,79 +1387,4 @@ public ValueTask DisposeAsync() } } } - - private sealed class CountingDisposeClient : IChatClient - { - public int DisposeCount { get; private set; } - - public Task GetResponseAsync( - IEnumerable messages, - ChatOptions? options = null, - CancellationToken cancellationToken = default) => - Task.FromResult(new ChatResponse()); - - public IAsyncEnumerable GetStreamingResponseAsync( - IEnumerable messages, - ChatOptions? options = null, - CancellationToken cancellationToken = default) => - throw new NotSupportedException(); - - public object? GetService(Type serviceType, object? serviceKey = null) => null; - - public void Dispose() => DisposeCount++; - - public override bool Equals(object? obj) => obj is CountingDisposeClient; - - public override int GetHashCode() => 0; - } - - private sealed class CountingEmbeddingGenerator : - IEmbeddingGenerator> - { - public int DisposeCount { get; private set; } - - public Task>> GenerateAsync( - IEnumerable values, - EmbeddingGenerationOptions? options = null, - CancellationToken cancellationToken = default) => - throw new NotSupportedException(); - - public object? GetService(Type serviceType, object? serviceKey = null) => null; - - public void Dispose() => DisposeCount++; - } - - private sealed class ChatClientReferenceComparer : IEqualityComparer - { - public static ChatClientReferenceComparer Instance { get; } = new(); - - public bool Equals(IChatClient? x, IChatClient? y) => ReferenceEquals(x, y); - - public int GetHashCode(IChatClient obj) => RuntimeHelpers.GetHashCode(obj); - } - - private sealed class ValueEqualChatClient(Func> getResponse) : IChatClient - { - public Task GetResponseAsync( - IEnumerable messages, - ChatOptions? options = null, - CancellationToken cancellationToken = default) => - getResponse(); - - public IAsyncEnumerable GetStreamingResponseAsync( - IEnumerable messages, - ChatOptions? options = null, - CancellationToken cancellationToken = default) => - throw new NotSupportedException(); - - public object? GetService(Type serviceType, object? serviceKey = null) => null; - - public void Dispose() - { - } - - public override bool Equals(object? obj) => obj is ValueEqualChatClient; - - public override int GetHashCode() => 0; - } } diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs new file mode 100644 index 00000000000..ea9de58c785 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs @@ -0,0 +1,414 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Extensions.AI; + +public class OrderedFailoverChatClientTests +{ + [Fact] + public void OrderedFailover_RejectsMissingClients() + { + using var inner = new TestChatClient(); + + Assert.Throws(() => new OrderedFailoverChatClient(null!)); + Assert.Throws(() => new OrderedFailoverChatClient([])); + Assert.Throws(() => new OrderedFailoverChatClient([inner, null!])); + Assert.Throws(() => new OrderedFailoverChatClient([inner, inner])); + } + + [Fact] + public async Task OrderedFailover_TriesClientsInOrderAndSkipsFailures() + { + var calls = new List(); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls.Add("first"); + throw new InvalidOperationException("first failed"); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls.Add("second"); + return Task.FromResult(expected); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + Assert.Equal(["first", "second"], calls); + } + + [Fact] + public async Task OrderedFailover_DistinguishesValueEqualClients() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var first = new ValueEqualChatClient( + () => throw new InvalidOperationException("failed")); + using var second = new ValueEqualChatClient( + () => Task.FromResult(expected)); + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + [Fact] + public async Task OrderedFailover_ExhaustionRethrowsLastFailure() + { + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("first"), + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("second"), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => client.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("second", exception.Message); + } + + [Fact] + public async Task OrderedFailover_StateIsScopedToOneRequest() + { + int firstCalls = 0; + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + firstCalls++; + throw new InvalidOperationException("failed"); + }, + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(new ChatResponse()), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + _ = await client.GetResponseAsync([new(ChatRole.User, "one")]); + _ = await client.GetResponseAsync([new(ChatRole.User, "two")]); + + Assert.Equal(2, firstCalls); + } + + [Fact] + public async Task OrderedFailover_ConcurrentRequestsHaveIndependentState() + { + const int RequestCount = 10; + int firstCalls = 0; + int secondCalls = 0; + var allStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var first = new TestChatClient + { + GetResponseAsyncCallback = async (_, _, _) => + { + if (Interlocked.Increment(ref firstCalls) == RequestCount) + { + allStarted.SetResult(true); + } + + await allStarted.Task; + throw new InvalidOperationException("failed"); + }, + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + Interlocked.Increment(ref secondCalls); + return Task.FromResult(new ChatResponse()); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + Task[] requests = + [ + .. Enumerable.Range(0, RequestCount).Select(index => + client.GetResponseAsync([new(ChatRole.User, index.ToString())])), + ]; + _ = await Task.WhenAll(requests); + + Assert.Equal(RequestCount, firstCalls); + Assert.Equal(RequestCount, secondCalls); + } + + [Fact] + public async Task OrderedFailover_DoesNotRetainStateWhileStreaming() + { + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("first", "second"), + }; + 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); + Assert.Equal(0, GetOrderedFailoverRequestStateCount(client)); + } + + [Fact] + public async Task OrderedFailover_SnapshotsConfiguredClients() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + var clients = new List { inner }; + using var failover = new OrderedFailoverChatClient(clients, leaveOpen: true); + clients.Clear(); + + ChatResponse response = await failover.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + [Fact] + public async Task OrderedFailover_UsesEachConfiguredClient() + { + var seenOptions = new List<(string? modelId, string? instructions)>(); + using var firstInner = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + seenOptions.Add((options?.ModelId, options?.Instructions)); + throw new InvalidOperationException("failed"); + }, + }; + using var first = new ConfigureOptionsChatClient( + firstInner, + options => options.ModelId = "first"); + using var secondInner = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + seenOptions.Add((options?.ModelId, options?.Instructions)); + return Task.FromResult(new ChatResponse()); + }, + }; + using var second = new ConfigureOptionsChatClient( + secondInner, + options => options.ModelId = "second"); + using var client = new OrderedFailoverChatClient( + [first, second], + leaveOpen: true); + + _ = await client.GetResponseAsync( + [new(ChatRole.User, "hi")], + new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }); + + Assert.Equal( + [("first", "caller"), ("second", "caller")], + seenOptions); + } + + [Fact] + public async Task OrderedFailover_StreamingFallsBackBeforeFirstUpdate() + { + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + List updates = + await CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("ok", Assert.Single(updates).Text); + } + + [Fact] + public async Task OrderedFailover_CancellationDoesNotFallback() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + bool secondCalled = false; + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return Task.FromResult(new ChatResponse()); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + await Assert.ThrowsAnyAsync( + () => client.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.False(secondCalled); + } + + [Fact] + public void OrderedFailover_DisposesEachClientOnce() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var shared = new CountingDisposeClient(); + var other = new CountingDisposeClient(); + var client = new OrderedFailoverChatClient([shared, other]); +#pragma warning restore CA2000 + + client.Dispose(); + client.Dispose(); + + Assert.Equal(1, shared.DisposeCount); + Assert.Equal(1, other.DisposeCount); + } + + [Fact] + public void OrderedFailover_LeaveOpenDoesNotDisposeInnerClients() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var inner = new CountingDisposeClient(); + var client = new OrderedFailoverChatClient([inner], leaveOpen: true); +#pragma warning restore CA2000 + + client.Dispose(); + + Assert.Equal(0, inner.DisposeCount); + inner.Dispose(); + } + + [Fact] + public async Task NestedRouters_ReturnLeafResponse() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var leaf = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + using var inner = new OrderedFailoverChatClient([leaf]); + using var outer = new OrderedFailoverChatClient([inner]); + + ChatResponse response = await outer.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + private static async Task> CollectAsync( + IAsyncEnumerable updates) + { + var result = new List(); + await foreach (ChatResponseUpdate update in updates) + { + result.Add(update); + } + + return result; + } + + 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 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 IAsyncEnumerable ThrowingStream(string message) + { + await Task.Yield(); + foreach (int _ in Array.Empty()) + { + yield return new ChatResponseUpdate(ChatRole.Assistant, "never"); + } + + throw new InvalidOperationException(message); + } + + private sealed class CountingDisposeClient : IChatClient + { + public int DisposeCount { get; private set; } + + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + Task.FromResult(new ChatResponse()); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + + public override bool Equals(object? obj) => obj is CountingDisposeClient; + + public override int GetHashCode() => 0; + } + + private sealed class ValueEqualChatClient(Func> getResponse) : IChatClient + { + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + getResponse(); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() + { + } + + public override bool Equals(object? obj) => obj is ValueEqualChatClient; + + public override int GetHashCode() => 0; + } +} diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/SemanticRoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/SemanticRoutingChatClientTests.cs new file mode 100644 index 00000000000..c41627ae698 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/SemanticRoutingChatClientTests.cs @@ -0,0 +1,285 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Extensions.AI; + +public class SemanticRoutingChatClientTests +{ + [Fact] + public void SemanticRouting_RejectsInvalidConfiguration() + { + using var client = new TestChatClient(); + using var generator = new TestEmbeddingGenerator(); + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [client] = ["profile"], + }; + + Assert.Throws(() => + new SemanticRoutingChatClient(null!, profiles, client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, null!, client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, null!)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary>(), + client)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary> + { + [client] = [], + }, + client)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary> + { + [client] = [" "], + }, + client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, client, scoreThreshold: 1.1f)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, client, topK: 0)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + profiles, + client, + scoreAggregation: (SemanticRoutingChatClient.ScoreAggregation)(-1))); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + profiles, + client, + scoreThreshold: 2.1f, + topK: 2, + scoreAggregation: SemanticRoutingChatClient.ScoreAggregation.Sum)); + } + + [Fact] + public async Task SemanticRouting_SelectsBestProfileAndCachesIndex() + { + var vectors = new Dictionary + { + ["code profile"] = [1, 0], + ["writing profile"] = [0, 1], + ["debug this code"] = [1, 0], + }; + int profileBatches = 0; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + { + string[] inputs = [.. values]; + if (inputs.Length > 1) + { + profileBatches++; + } + + return Task.FromResult(new GeneratedEmbeddings>( + [.. inputs.Select(input => new Embedding(vectors[input]))])); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "code")); + using var code = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + using var writing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "writing"))), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [code] = ["code profile"], + [writing] = ["writing profile"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient: writing, + leaveOpen: true); + + ChatResponse first = await router.GetResponseAsync([new(ChatRole.User, "debug this code")]); + ChatResponse second = await router.GetResponseAsync([new(ChatRole.User, "debug this code")]); + + Assert.Same(expected, first); + Assert.Same(expected, second); + Assert.Equal(1, profileBatches); + } + + [Theory] + [InlineData(SemanticRoutingChatClient.ScoreAggregation.Mean, "code")] + [InlineData(SemanticRoutingChatClient.ScoreAggregation.Sum, "writing")] + public async Task SemanticRouting_AggregatesGlobalTopKByClient( + SemanticRoutingChatClient.ScoreAggregation scoreAggregation, + string expectedResponse) + { + var vectors = new Dictionary + { + ["code"] = [1, 0], + ["writing one"] = [0.8f, 0.6f], + ["writing two"] = [0.8f, -0.6f], + ["query"] = [1, 0], + }; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + Task.FromResult(new GeneratedEmbeddings>( + [.. values.Select(input => new Embedding(vectors[input]))])), + }; + using var code = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "code"))), + }; + using var writing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "writing"))), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [code] = ["code"], + [writing] = ["writing one", "writing two"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient: code, + scoreThreshold: scoreAggregation == SemanticRoutingChatClient.ScoreAggregation.Sum ? 1.5f : 0.3f, + topK: 3, + scoreAggregation: scoreAggregation, + leaveOpen: true); + + ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "query")]); + + Assert.Equal(expectedResponse, response.Text); + } + + [Fact] + public async Task SemanticRouting_UsesDefaultBelowThreshold() + { + var vectors = new Dictionary + { + ["code profile"] = [1, 0], + ["unrelated query"] = [0, 1], + }; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + Task.FromResult(new GeneratedEmbeddings>( + [.. values.Select(input => new Embedding(vectors[input]))])), + }; + using var profiled = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "profiled"))), + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "default")); + using var defaultClient = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [profiled] = ["code profile"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient, + scoreThreshold: 0.5f, + leaveOpen: true); + + ChatResponse response = + await router.GetResponseAsync([new(ChatRole.User, "unrelated query")]); + + Assert.Same(expected, response); + } + + [Fact] + public void SemanticRouting_DisposesOwnedResourcesOnce() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var generator = new CountingEmbeddingGenerator(); + var profiled = new CountingDisposeClient(); + var defaultClient = new CountingDisposeClient(); + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [profiled] = ["profile"], + }; + var router = new SemanticRoutingChatClient(generator, profiles, defaultClient); +#pragma warning restore CA2000 + + router.Dispose(); + router.Dispose(); + + Assert.Equal(1, generator.DisposeCount); + Assert.Equal(1, profiled.DisposeCount); + Assert.Equal(1, defaultClient.DisposeCount); + } + + private sealed class CountingDisposeClient : IChatClient + { + public int DisposeCount { get; private set; } + + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + Task.FromResult(new ChatResponse()); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + + public override bool Equals(object? obj) => obj is CountingDisposeClient; + + public override int GetHashCode() => 0; + } + + private sealed class CountingEmbeddingGenerator : + IEmbeddingGenerator> + { + public int DisposeCount { get; private set; } + + public Task>> GenerateAsync( + IEnumerable values, + EmbeddingGenerationOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + } + + private sealed class ChatClientReferenceComparer : IEqualityComparer + { + public static ChatClientReferenceComparer Instance { get; } = new(); + + public bool Equals(IChatClient? x, IChatClient? y) => ReferenceEquals(x, y); + + public int GetHashCode(IChatClient obj) => RuntimeHelpers.GetHashCode(obj); + } +} From 938e0b97b6c35bff6353d71d73b83b3f2aea7626 Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Fri, 31 Jul 2026 16:36:40 -0700 Subject: [PATCH 14/15] Refine routing tests Consolidate shared streaming and non-streaming cases, reuse routing test helpers, and simplify semantic router disposal. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/SemanticRoutingChatClient.cs | 7 +- .../ChatRouting/FailoverChatClientTests.cs | 172 ++++++------------ .../OrderedFailoverChatClientTests.cs | 39 +--- 3 files changed, 62 insertions(+), 156 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs index ff22411fff2..2cbce7be87e 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -291,12 +291,7 @@ protected override void Dispose(bool disposing) client.Dispose(); } - if (!Array.Exists( - _clients, - client => ReferenceEquals(client, _embeddingGenerator))) - { - _embeddingGenerator.Dispose(); - } + _embeddingGenerator.Dispose(); } } diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs index 3c03108b665..6af1df345aa 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs @@ -38,8 +38,10 @@ await Assert.ThrowsAsync( () => CollectAsync(router.GetStreamingResponseAsync(null!))); } - [Fact] - public async Task InitialSelectionFailureDoesNotReportAttempt() + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task InitialSelectionFailureDoesNotReportAttempt(bool streaming) { var expected = new InvalidOperationException("selection failed"); RoutingContext? selectedContext = null; @@ -52,16 +54,21 @@ public async Task InitialSelectionFailureDoesNotReportAttempt() }, (_, _, _) => updateCount++); - InvalidOperationException actual = await Assert.ThrowsAsync( - () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + Task operation = streaming + ? router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync() + : router.GetResponseAsync([new(ChatRole.User, "hi")]); + InvalidOperationException actual = + await Assert.ThrowsAsync(() => operation); Assert.Same(expected, actual); Assert.NotNull(selectedContext); Assert.Equal(0, updateCount); } - [Fact] - public async Task InitialSelectionFailureWithCallerCancellationPropagatesSelectionFailure() + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task InitialSelectionFailureWithCallerCancellationPropagatesSelectionFailure(bool streaming) { using var cancellationSource = new CancellationTokenSource(); cancellationSource.Cancel(); @@ -71,83 +78,35 @@ public async Task InitialSelectionFailureWithCallerCancellationPropagatesSelecti _ => throw selectionException, (_, _, _) => updateCount++); - InvalidOperationException actual = await Assert.ThrowsAsync( - () => router.GetResponseAsync( + Task operation = streaming + ? router.GetStreamingResponseAsync( [new(ChatRole.User, "hi")], - cancellationToken: cancellationSource.Token)); - - Assert.Same(selectionException, actual); - Assert.Equal(0, updateCount); - } - - [Fact] - public async Task StreamingInitialSelectionFailureDoesNotReportAttempt() - { - var expected = new InvalidOperationException("selection failed"); - RoutingContext? selectedContext = null; - int updateCount = 0; - using var router = new DelegatingFailoverTestRouter( - context => - { - selectedContext = context; - throw expected; - }, - (_, _, _) => updateCount++); - - InvalidOperationException actual = await Assert.ThrowsAsync( - () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); - - Assert.Same(expected, actual); - Assert.NotNull(selectedContext); - Assert.Equal(0, updateCount); - } - - [Fact] - public async Task StreamingInitialSelectionFailureWithCallerCancellationPropagatesSelectionFailure() - { - using var cancellationSource = new CancellationTokenSource(); - cancellationSource.Cancel(); - var selectionException = new InvalidOperationException("selection failed"); - int updateCount = 0; - using var router = new DelegatingFailoverTestRouter( - _ => throw selectionException, - (_, _, _) => updateCount++); - - InvalidOperationException actual = await Assert.ThrowsAsync( - () => CollectAsync( - router.GetStreamingResponseAsync( - [new(ChatRole.User, "hi")], - cancellationToken: cancellationSource.Token))); + cancellationToken: cancellationSource.Token).ToChatResponseAsync() + : router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token); + InvalidOperationException actual = + await Assert.ThrowsAsync(() => operation); Assert.Same(selectionException, actual); Assert.Equal(0, updateCount); } - [Fact] - public async Task NullSelectionDoesNotReportAttempt() - { - int updateCount = 0; - using var router = new DelegatingFailoverTestRouter( - _ => null!, - (_, _, _) => updateCount++); - - InvalidOperationException exception = await Assert.ThrowsAsync( - () => router.GetResponseAsync([new(ChatRole.User, "hi")])); - - Assert.Contains("SelectClientAsync", exception.Message, StringComparison.Ordinal); - Assert.Equal(0, updateCount); - } - - [Fact] - public async Task StreamingNullSelectionDoesNotReportAttempt() + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task NullSelectionDoesNotReportAttempt(bool streaming) { int updateCount = 0; using var router = new DelegatingFailoverTestRouter( _ => null!, (_, _, _) => updateCount++); - InvalidOperationException exception = await Assert.ThrowsAsync( - () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + Task operation = streaming + ? router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync() + : router.GetResponseAsync([new(ChatRole.User, "hi")]); + InvalidOperationException exception = + await Assert.ThrowsAsync(() => operation); Assert.Contains("SelectClientAsync", exception.Message, StringComparison.Ordinal); Assert.Equal(0, updateCount); @@ -241,22 +200,6 @@ public async Task NonterminalRoutingUpdateFailureStopsWithoutAnotherUpdate() Assert.Equal(1, updates); } - [Fact] - public async Task SelectionFailureDoesNotInvokeRoutingUpdate() - { - var selectionException = new InvalidOperationException("selection failed"); - int updateCount = 0; - using var router = new DelegatingFailoverTestRouter( - _ => throw selectionException, - (_, _, _) => updateCount++); - - InvalidOperationException actual = await Assert.ThrowsAsync( - () => router.GetResponseAsync([new(ChatRole.User, "hi")])); - - Assert.Same(selectionException, actual); - Assert.Equal(0, updateCount); - } - [Fact] public async Task Dispatch_ConfiguredClientPreservesAndOverridesRequestOptions() { @@ -652,10 +595,10 @@ public async Task Streaming_FallsBackBeforeFirstUpdate() using var router = new DelegatingFailoverTestRouter( _ => ++selections == 1 ? failing : working); - List updates = - await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); Assert.Equal(2, selections); } @@ -689,12 +632,11 @@ public async Task Streaming_FailoverUsesFreshRequestOptionsForEachAttempt() using var router = new DelegatingFailoverTestRouter( _ => ++selections == 1 ? first : second); - List updates = await CollectAsync( - router.GetStreamingResponseAsync( - [new(ChatRole.User, "hi")], - requestOptions)); + ChatResponse response = await router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + requestOptions).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); Assert.Equal(2, selections); Assert.NotSame(invokedOptions[0], invokedOptions[1]); Assert.Equal("changed by first client", invokedOptions[0].ModelId); @@ -827,10 +769,10 @@ public async Task Streaming_OperationCanceledExceptionCanTriggerFailover() }; using var client = new OrderedFailoverChatClient([first, second]); - List updates = - await CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = + await client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); Assert.True(secondCalled); } @@ -855,10 +797,10 @@ public async Task Streaming_DisposalCancellationCanTriggerFailover() }; using var client = new OrderedFailoverChatClient([first, second]); - List updates = - await CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = + await client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); Assert.True(secondCalled); Assert.Equal(1, canceledStream.DisposeCount); } @@ -906,10 +848,10 @@ public async Task Streaming_StreamCreationFailureFallsBack() using var router = new DelegatingFailoverTestRouter( _ => ++selections == 1 ? failing : working); - List updates = - await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); Assert.Equal(2, selections); } @@ -929,10 +871,10 @@ public async Task Streaming_EnumeratorCreationFailureFallsBack() using var router = new DelegatingFailoverTestRouter( _ => ++selections == 1 ? failing : working); - List updates = - await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); Assert.Equal(2, selections); } @@ -960,10 +902,10 @@ public async Task Streaming_CurrentFailureFallsBackBeforeOutput() } }); - List updates = - await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); Assert.Same(currentException, failedAttempt!.Exception); Assert.False(failedAttempt.OutputCommitted); Assert.False(failedAttempt.ResponseCompleted); @@ -1048,10 +990,10 @@ public async Task Streaming_EmptyStreamDisposalFailureFallsBack() } }); - List updates = - await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); Assert.Same(disposalException, failedAttempt!.Exception); Assert.False(failedAttempt.OutputCommitted); Assert.False(failedAttempt.ResponseCompleted); @@ -1206,7 +1148,7 @@ private static async Task ConsumeOneAsync(IAsyncEnumerable u Assert.True(await enumerator.MoveNextAsync()); } - private static async IAsyncEnumerable YieldUpdates(params string[] texts) + internal static async IAsyncEnumerable YieldUpdates(params string[] texts) { foreach (string text in texts) { @@ -1215,7 +1157,7 @@ private static async IAsyncEnumerable YieldUpdates(params st } } - private static async IAsyncEnumerable ThrowingStream(string message) + internal static async IAsyncEnumerable ThrowingStream(string message) { await Task.Yield(); foreach (int _ in Array.Empty()) diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs index ea9de58c785..5cb2cf34ed8 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs @@ -7,6 +7,7 @@ using System.Threading; using System.Threading.Tasks; using Xunit; +using static Microsoft.Extensions.AI.FailoverChatClientTests; namespace Microsoft.Extensions.AI; @@ -244,10 +245,10 @@ public async Task OrderedFailover_StreamingFallsBackBeforeFirstUpdate() }; using var client = new OrderedFailoverChatClient([first, second]); - List updates = - await CollectAsync(client.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + ChatResponse response = + await client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); - Assert.Equal("ok", Assert.Single(updates).Text); + Assert.Equal("ok", response.Text); } [Fact] @@ -324,18 +325,6 @@ public async Task NestedRouters_ReturnLeafResponse() Assert.Same(expected, response); } - private static async Task> CollectAsync( - IAsyncEnumerable updates) - { - var result = new List(); - await foreach (ChatResponseUpdate update in updates) - { - result.Add(update); - } - - return result; - } - private static int GetOrderedFailoverRequestStateCount(OrderedFailoverChatClient client) { object requestStates = typeof(OrderedFailoverChatClient) @@ -346,26 +335,6 @@ private static int GetOrderedFailoverRequestStateCount(OrderedFailoverChatClient return (int)requestStates.GetType().GetProperty("Count")!.GetValue(requestStates)!; } - 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 IAsyncEnumerable ThrowingStream(string message) - { - await Task.Yield(); - foreach (int _ in Array.Empty()) - { - yield return new ChatResponseUpdate(ChatRole.Assistant, "never"); - } - - throw new InvalidOperationException(message); - } - private sealed class CountingDisposeClient : IChatClient { public int DisposeCount { get; private set; } From 78ae0232e535e837595b603423dece9cf3668bee Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Fri, 31 Jul 2026 18:04:50 -0700 Subject: [PATCH 15/15] Add per-attempt routing selections Return the selected client with complete ephemeral options while preserving client-reference identity and recording routing-boundary attempt options. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/RoutingChatClient.cs | 56 ++++---- .../ChatRouting/RoutingContext.cs | 9 +- .../ChatRouting/RoutingSelection.cs | 35 +++++ .../Microsoft.Extensions.AI.Abstractions.json | 24 +++- .../ChatRouting/FailoverChatClient.cs | 19 ++- .../ChatRouting/FailoverChatClientAttempt.cs | 9 ++ .../ChatRouting/OrderedFailoverChatClient.cs | 7 +- .../ChatRouting/SemanticRoutingChatClient.cs | 12 +- .../Microsoft.Extensions.AI.json | 8 +- .../ChatRouting/RoutingChatClientTests.cs | 40 ++++-- .../ChatRouting/FailoverChatClientTests.cs | 126 +++++++++++++++++- 11 files changed, 280 insertions(+), 65 deletions(-) create mode 100644 src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingSelection.cs diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs index 58a1f050a2f..0f28c518c3b 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs @@ -16,38 +16,35 @@ namespace Microsoft.Extensions.AI; /// Provides a template for an that selects and invokes another chat client. /// /// -/// Derived classes implement to supply one client for each request. The selected -/// client is invoked once, and its response or failure is propagated to the caller. -/// The exact client instance represents the selected routing identity. Caller-supplied options may vary per -/// invocation, but they are ephemeral and do not participate in that identity. Applications can use distinct -/// configured client wrappers when configurations require distinct routing identities, while custom policies may -/// maintain separate application-specific grouping keys. +/// Derived classes implement to supply one selection for each request. The selected +/// client is invoked once with the selected options, and its response or failure is propagated to the caller. /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] public abstract class RoutingChatClient : IChatClient { - /// Creates a routing client that selects one client for each request. - /// The callback that selects the client to invoke. - /// A routing client that uses . + /// Creates a routing client that selects one client and its options for each request. + /// The callback that produces the selection to invoke. + /// A routing client that uses . /// The selected clients are caller-owned and are not disposed by the returned routing client. - /// is . + /// is . public static RoutingChatClient Create( - Func> clientSelector) + Func> selector) { - _ = Throw.IfNull(clientSelector); + _ = Throw.IfNull(selector); - return new CallbackRoutingChatClient(clientSelector); + return new CallbackRoutingChatClient(selector); } - /// Selects the client to invoke for the request. + /// Selects the client and complete options to use for the request. /// The request-specific inputs. /// The cancellation token supplied for the request. - /// The client to invoke. + /// The selection to invoke. /// - /// Client-specific behavior should generally be attached to the returned client. Exceptions from this method - /// propagate to the caller. + /// The returned are complete attempt options, not an overlay. + /// Implementations may clone and adjust the clone for the selected + /// attempt. Exceptions from this method propagate to the caller. /// - protected abstract ValueTask SelectClientAsync( + protected abstract ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken); @@ -58,12 +55,12 @@ public virtual async Task GetResponseAsync( _ = Throw.IfNull(messages); var context = new RoutingContext(messages, options); - IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + RoutingSelection selection = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); - return await client.GetResponseAsync( + return await selection.Client.GetResponseAsync( context.Messages, - options?.Clone(), + selection.ChatOptions, cancellationToken).ConfigureAwait(false); } @@ -75,11 +72,14 @@ public virtual async IAsyncEnumerable GetStreamingResponseAs _ = Throw.IfNull(messages); var context = new RoutingContext(messages, options); - IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + RoutingSelection selection = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); await foreach (ChatResponseUpdate update in - client.GetStreamingResponseAsync(context.Messages, options?.Clone(), cancellationToken) + selection.Client.GetStreamingResponseAsync( + context.Messages, + selection.ChatOptions, + cancellationToken) .WithCancellation(cancellationToken) .ConfigureAwait(false)) { @@ -111,17 +111,17 @@ protected virtual void Dispose(bool disposing) private sealed class CallbackRoutingChatClient : RoutingChatClient { - private readonly Func> _clientSelector; + private readonly Func> _selector; public CallbackRoutingChatClient( - Func> clientSelector) + Func> selector) { - _clientSelector = clientSelector; + _selector = selector; } - protected override ValueTask SelectClientAsync( + protected override ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) => - _clientSelector(context, cancellationToken); + _selector(context, cancellationToken); } } diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs index edd5bb7760f..1c72a4d18fb 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs @@ -15,8 +15,8 @@ namespace Microsoft.Extensions.AI; /// started from the sequence returned by . /// /// -/// is cloned when the context is created. It is provided for client selection and is -/// independent of both the caller's instance and the options passed to the selected client. +/// is cloned when the context is created and represents the request baseline. Selectors can +/// clone and modify it to produce complete attempt options through . /// /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] @@ -41,10 +41,9 @@ public RoutingContext( /// public IEnumerable Messages { get; } - /// Gets a snapshot of the request options supplied to client selection. + /// Gets the request options supplied to client selection. /// - /// Changes do not affect the caller's instance or the options passed to the selected client. Client-specific - /// behavior should generally be attached to the returned client. + /// This is a clone of the caller's options. Selectors should clone it before applying attempt-specific changes. /// public ChatOptions? ChatOptions { get; } } diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingSelection.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingSelection.cs new file mode 100644 index 00000000000..5f8947b83cd --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingSelection.cs @@ -0,0 +1,35 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Diagnostics.CodeAnalysis; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Represents the client and complete options selected for one routing attempt. +/// +/// The exact instance represents the routing identity. are ephemeral +/// attempt data and do not participate in that identity. The selected client may further transform the options. +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public sealed class RoutingSelection +{ + /// Initializes a new instance of the class. + /// The client selected for the attempt. + /// The complete options selected for the attempt, or . + /// is . + public RoutingSelection(IChatClient client, ChatOptions? chatOptions) + { + Client = Throw.IfNull(client); + ChatOptions = chatOptions?.Clone(); + } + + /// Gets the client selected for the attempt. + public IChatClient Client { get; } + + /// Gets the complete options selected for the attempt. + /// This is a clone of the options supplied to the constructor. + public ChatOptions? ChatOptions { get; } +} 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..0518ee08b60 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 @@ -3975,7 +3975,7 @@ "Stage": "Experimental" }, { - "Member": "static Microsoft.Extensions.AI.RoutingChatClient Microsoft.Extensions.AI.RoutingChatClient.Create(System.Func> clientSelector);", + "Member": "static Microsoft.Extensions.AI.RoutingChatClient Microsoft.Extensions.AI.RoutingChatClient.Create(System.Func> selector);", "Stage": "Experimental" }, { @@ -3999,7 +3999,7 @@ "Stage": "Experimental" }, { - "Member": "abstract System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.RoutingChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", + "Member": "abstract System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.RoutingChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", "Stage": "Experimental" } ] @@ -4024,6 +4024,26 @@ } ] }, + { + "Type": "sealed class Microsoft.Extensions.AI.RoutingSelection", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.RoutingSelection.RoutingSelection(Microsoft.Extensions.AI.IChatClient client, Microsoft.Extensions.AI.ChatOptions? chatOptions);", + "Stage": "Experimental" + } + ], + "Properties": [ + { + "Member": "Microsoft.Extensions.AI.ChatOptions? Microsoft.Extensions.AI.RoutingSelection.ChatOptions { get; }", + "Stage": "Experimental" + }, + { + "Member": "Microsoft.Extensions.AI.IChatClient Microsoft.Extensions.AI.RoutingSelection.Client { get; }", + "Stage": "Experimental" + } + ] + }, { "Type": "class Microsoft.Extensions.AI.SessionUpdateRealtimeClientMessage : Microsoft.Extensions.AI.RealtimeClientMessage", "Stage": "Experimental", diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs index d825a8a054d..b85268e13f0 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -105,10 +105,12 @@ public sealed override async Task GetResponseAsync( while (true) { - IChatClient selectedClient = + RoutingSelection selection = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); - ChatOptions? attemptOptions = options?.Clone(); + IChatClient selectedClient = selection.Client; + ChatOptions? invocationOptions = selection.ChatOptions; + ChatOptions? attemptOptions = invocationOptions?.Clone(); attemptCount++; ChatResponse? response = null; @@ -119,7 +121,7 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? { response = await selectedClient.GetResponseAsync( context.Messages, - attemptOptions, + invocationOptions, cancellationToken).ConfigureAwait(false); } catch (Exception ex) @@ -130,6 +132,7 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? stopwatch.Stop(); var attempt = new FailoverChatClientAttempt( selectedClient, + attemptOptions, exception, stopwatch.Elapsed, timeToFirstUpdate: null, @@ -175,10 +178,12 @@ public sealed override async IAsyncEnumerable GetStreamingRe while (true) { - IChatClient selectedClient = + RoutingSelection selection = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); - ChatOptions? attemptOptions = options?.Clone(); + IChatClient selectedClient = selection.Client; + ChatOptions? invocationOptions = selection.ChatOptions; + ChatOptions? attemptOptions = invocationOptions?.Clone(); attemptCount++; bool reachedAttemptLimit = maximumAttempts is int limit && attemptCount >= limit; @@ -192,7 +197,7 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? enumerator = selectedClient .GetStreamingResponseAsync( context.Messages, - attemptOptions, + invocationOptions, cancellationToken) .GetAsyncEnumerator(cancellationToken); @@ -204,6 +209,7 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? Exception exception = (await DisposeEnumeratorAsync(enumerator, ex).ConfigureAwait(false))!; var attempt = new FailoverChatClientAttempt( selectedClient, + attemptOptions, exception, stopwatch.Elapsed, timeToFirstUpdate: null, @@ -277,6 +283,7 @@ await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? var attempt = new FailoverChatClientAttempt( selectedClient, + attemptOptions, terminalException, stopwatch.Elapsed, timeToFirstUpdate, diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs index 0c4038ed8e0..21cb0e32585 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs @@ -31,6 +31,7 @@ public sealed class FailoverChatClientAttempt { internal FailoverChatClientAttempt( IChatClient client, + ChatOptions? chatOptions, Exception? exception, TimeSpan duration, TimeSpan? timeToFirstUpdate, @@ -49,6 +50,7 @@ internal FailoverChatClientAttempt( "Output commitment and time to first update should agree."); Client = Throw.IfNull(client); + ChatOptions = chatOptions; Exception = exception; Duration = duration; TimeToFirstUpdate = timeToFirstUpdate; @@ -59,6 +61,13 @@ internal FailoverChatClientAttempt( /// Gets the client that was invoked. public IChatClient Client { get; } + /// Gets a snapshot of the options passed by the router to the selected client. + /// + /// The selected client or middleware in its pipeline may subsequently clone or transform these options before + /// invoking an underlying provider. + /// + public ChatOptions? ChatOptions { get; } + /// Gets the time spent actively invoking the client. /// For streaming, time spent by the caller processing yielded updates is excluded. public TimeSpan Duration { get; } diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs index 7fcec7779bd..f90e3694d7a 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -74,20 +74,21 @@ public OrderedFailoverChatClient(IReadOnlyList clients, bool leaveO } /// - protected override ValueTask SelectClientAsync( + protected override ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) { + _ = Throw.IfNull(context); _ = cancellationToken; if (!_requestStates.TryRemove(context, out RequestState? state)) { - return new(_clients[0]); + return new(new RoutingSelection(_clients[0], context.ChatOptions)); } if (state.ClientIndex < _clients.Length) { - return new(_clients[state.ClientIndex]); + return new(new RoutingSelection(_clients[state.ClientIndex], context.ChatOptions)); } ExceptionDispatchInfo.Capture(state.LastException).Throw(); diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs index 2cbce7be87e..f51aa5c33d7 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -168,7 +168,7 @@ public SemanticRoutingChatClient( } /// - protected override async ValueTask SelectClientAsync( + protected override async ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) { @@ -178,7 +178,7 @@ protected override async ValueTask SelectClientAsync( static message => message.Role == ChatRole.User)?.Text; if (string.IsNullOrWhiteSpace(query)) { - return _clients[0]; + return new(_clients[0], context.ChatOptions); } EmbeddedProfile[] index = await EnsureIndexAsync(cancellationToken).ConfigureAwait(false); @@ -213,9 +213,11 @@ await _embeddingGenerator.GenerateAsync( } } - return bestClientIndex >= 0 && bestScore >= _scoreThreshold + IChatClient client = bestClientIndex >= 0 && bestScore >= _scoreThreshold ? _clients[bestClientIndex] : _clients[0]; + + return new(client, context.ChatOptions); } var matches = new ScoredProfile[index.Length]; @@ -267,9 +269,11 @@ await _embeddingGenerator.GenerateAsync( } } - return bestAggregatedClientIndex >= 0 && bestAggregatedScore >= _scoreThreshold + IChatClient selectedClient = bestAggregatedClientIndex >= 0 && bestAggregatedScore >= _scoreThreshold ? _clients[bestAggregatedClientIndex] : _clients[0]; + + return new(selectedClient, context.ChatOptions); } /// diff --git a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json index 564b3369241..1dbff0278a5 100644 --- a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json +++ b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json @@ -521,6 +521,10 @@ "Type": "sealed class Microsoft.Extensions.AI.FailoverChatClientAttempt", "Stage": "Experimental", "Properties": [ + { + "Member": "Microsoft.Extensions.AI.ChatOptions? Microsoft.Extensions.AI.FailoverChatClientAttempt.ChatOptions { get; }", + "Stage": "Experimental" + }, { "Member": "Microsoft.Extensions.AI.IChatClient Microsoft.Extensions.AI.FailoverChatClientAttempt.Client { get; }", "Stage": "Experimental" @@ -1500,7 +1504,7 @@ "Stage": "Experimental" }, { - "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.OrderedFailoverChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.OrderedFailoverChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", "Stage": "Experimental" } ] @@ -1618,7 +1622,7 @@ "Stage": "Experimental" }, { - "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.SemanticRoutingChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.SemanticRoutingChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", "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..f15f3896df8 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs @@ -17,6 +17,21 @@ public void Create_RejectsNullSelector() Assert.Throws(() => RoutingChatClient.Create(null!)); } + [Fact] + public void RoutingSelection_ValidatesClientAndClonesOptions() + { + using var client = new TestChatClient(); + var options = new ChatOptions { ModelId = "original" }; + + var selection = new RoutingSelection(client, options); + options.ModelId = "changed"; + + Assert.Same(client, selection.Client); + Assert.NotSame(options, selection.ChatOptions); + Assert.Equal("original", selection.ChatOptions!.ModelId); + Assert.Throws(() => new RoutingSelection(null!, options)); + } + [Fact] public async Task Create_SelectsClientForRequest() { @@ -41,8 +56,9 @@ public async Task Create_SelectsClientForRequest() observedContext = context; observedToken = cancellationToken; selectionCount++; - context.ChatOptions!.ModelId = "selected"; - return new(selected); + ChatOptions selectedOptions = context.ChatOptions!.Clone(); + selectedOptions.ModelId = "selected"; + return new(new RoutingSelection(selected, selectedOptions)); }); ChatResponse response = await router.GetResponseAsync(messages, options, cancellationSource.Token); @@ -51,8 +67,8 @@ public async Task Create_SelectsClientForRequest() Assert.Same(messages, observedContext!.Messages); Assert.NotSame(options, observedContext.ChatOptions); Assert.NotSame(observedContext.ChatOptions, forwardedOptions); - Assert.Equal("selected", observedContext.ChatOptions!.ModelId); - Assert.Equal("request", forwardedOptions!.ModelId); + Assert.Equal("request", observedContext.ChatOptions!.ModelId); + Assert.Equal("selected", forwardedOptions!.ModelId); Assert.Equal("request", options.ModelId); Assert.Equal(cancellationSource.Token, observedToken); Assert.Equal(1, selectionCount); @@ -67,10 +83,10 @@ public async Task Create_DoesNotReselectAfterFailure() GetResponseAsyncCallback = (_, _, _) => throw expected, }; int selectionCount = 0; - using RoutingChatClient router = RoutingChatClient.Create((_, _) => + using RoutingChatClient router = RoutingChatClient.Create((context, _) => { selectionCount++; - return new(selected); + return new(new RoutingSelection(selected, context.ChatOptions)); }); InvalidOperationException actual = await Assert.ThrowsAsync( @@ -83,7 +99,7 @@ public async Task Create_DoesNotReselectAfterFailure() [Fact] public async Task Create_NullSelectionThrowsForNonStreamingAndStreaming() { - using RoutingChatClient router = RoutingChatClient.Create((_, _) => new((IChatClient)null!)); + using RoutingChatClient router = RoutingChatClient.Create((_, _) => new((RoutingSelection)null!)); InvalidOperationException nonStreaming = await Assert.ThrowsAsync( () => router.GetResponseAsync([new(ChatRole.User, "hi")])); @@ -98,7 +114,8 @@ public async Task Create_NullSelectionThrowsForNonStreamingAndStreaming() public async Task Create_RejectsNullMessagesForNonStreamingAndStreaming() { using var selected = new TestChatClient(); - using RoutingChatClient router = RoutingChatClient.Create((_, _) => new(selected)); + using RoutingChatClient router = RoutingChatClient.Create( + (context, _) => new(new RoutingSelection(selected, context.ChatOptions))); await Assert.ThrowsAsync(() => router.GetResponseAsync(null!)); await Assert.ThrowsAsync( @@ -109,7 +126,8 @@ await Assert.ThrowsAsync( public async Task Create_DoesNotDisposeSelectedClient() { using var selected = new CountingDisposeClient(); - using (RoutingChatClient router = RoutingChatClient.Create((_, _) => new(selected))) + using (RoutingChatClient router = RoutingChatClient.Create( + (context, _) => new(new RoutingSelection(selected, context.ChatOptions)))) { _ = await router.GetResponseAsync([new(ChatRole.User, "hi")]); } @@ -150,10 +168,10 @@ public DelegatingTestRouter(Func select) _select = select; } - protected override ValueTask SelectClientAsync( + protected override ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) => - new(_select(context)); + new(new RoutingSelection(_select(context), context.ChatOptions)); } private sealed class CountingDisposeClient : IChatClient diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs index 6af1df345aa..b2c0289c345 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs @@ -220,13 +220,21 @@ public async Task Dispatch_ConfiguredClientPreservesAndOverridesRequestOptions() Instructions = "caller", ModelId = "request", }; - using var router = new DelegatingTestRouter(_ => configured); + FailoverChatClientAttempt? observedAttempt = null; + using var router = new DelegatingFailoverTestRouter( + _ => configured, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observedAttempt = attempt; + }); _ = await router.GetResponseAsync([new(ChatRole.User, "hi")], requestOptions); Assert.NotSame(requestOptions, forwarded); Assert.Equal("route", forwarded!.ModelId); Assert.Equal("caller", forwarded.Instructions); + Assert.Equal("request", observedAttempt!.ChatOptions!.ModelId); Assert.Equal("request", requestOptions.ModelId); } @@ -275,6 +283,84 @@ public async Task Failover_UsesFreshRequestOptionsForEachAttempt() Assert.Equal("request", requestOptions.ModelId); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task SelectedOptionsAreScopedToAttempt(bool streaming) + { + var requestOptions = new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }; + var invokedOptions = new List(); + var attempts = new List(); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + options!.ModelId = "mutated"; + throw new InvalidOperationException("failed"); + }, + GetStreamingResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + options!.ModelId = "mutated"; + return ThrowingStream("failed"); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + return Task.FromResult(expected); + }, + GetStreamingResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + return YieldUpdates("ok"); + }, + }; + int selections = 0; + using var router = new DelegatingSelectionFailoverTestRouter( + context => + { + Assert.Equal("request", context.ChatOptions!.ModelId); + ChatOptions selectedOptions = context.ChatOptions.Clone(); + bool firstAttempt = selections++ == 0; + selectedOptions.ModelId = firstAttempt ? "first" : "second"; + return new(firstAttempt ? first : second, selectedOptions); + }, + (_, attempt, _) => attempts.Add(attempt)); + + ChatResponse response = streaming + ? await router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + requestOptions).ToChatResponseAsync() + : await router.GetResponseAsync([new(ChatRole.User, "hi")], requestOptions); + + Assert.Equal("ok", response.Text); + Assert.Equal(2, selections); + Assert.Equal("mutated", invokedOptions[0].ModelId); + Assert.Equal("second", invokedOptions[1].ModelId); + Assert.Collection( + attempts, + attempt => + { + Assert.Equal("first", attempt.ChatOptions!.ModelId); + Assert.NotSame(invokedOptions[0], attempt.ChatOptions); + }, + attempt => + { + Assert.Equal("second", attempt.ChatOptions!.ModelId); + Assert.NotSame(invokedOptions[1], attempt.ChatOptions); + }); + Assert.Equal("request", requestOptions.ModelId); + } + [Fact] public async Task Failure_Propagates() { @@ -1189,10 +1275,10 @@ public DelegatingTestRouter(Func select) _select = select; } - protected override ValueTask SelectClientAsync( + protected override ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) => - new(_select(context)); + new(new RoutingSelection(_select(context), context.ChatOptions)); } private sealed class DelegatingFailoverTestRouter : FailoverChatClient @@ -1208,7 +1294,39 @@ public DelegatingFailoverTestRouter( _onRoutingUpdate = onRoutingUpdate; } - protected override ValueTask SelectClientAsync( + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) + { + IChatClient client = _select(context); + return new(client is null ? null! : new RoutingSelection(client, context.ChatOptions)); + } + + protected override ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt attempt, + bool isTerminal, + CancellationToken cancellationToken) + { + _onRoutingUpdate?.Invoke(context, attempt, isTerminal); + return default; + } + } + + private sealed class DelegatingSelectionFailoverTestRouter : FailoverChatClient + { + private readonly Func _select; + private readonly Action? _onRoutingUpdate; + + public DelegatingSelectionFailoverTestRouter( + Func select, + Action? onRoutingUpdate = null) + { + _select = select; + _onRoutingUpdate = onRoutingUpdate; + } + + protected override ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) => new(_select(context));