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 2a2e284f6f83b8530ccc94fc3496602d9c48879d Mon Sep 17 00:00:00 2001 From: Joshua Yue Date: Mon, 3 Aug 2026 15:20:31 -0700 Subject: [PATCH 15/15] Simplify ordered failover request state Allow repeated client instances in ordered failover. Selection now reads a stored index instead of reverse-resolving the attempted client by reference, so the constructor no longer requires unique instances and a client may appear more than once, being invoked once per position. The per-request state holds only the next index. Exhaustion is detected when the index advances past the last client, so the final failure is rethrown from the routing update rather than stored and rethrown from the next selection. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74d04840-2379-4615-93f7-84f2299ada74 --- .../ChatRouting/OrderedFailoverChatClient.cs | 84 +++++------------- .../OrderedFailoverChatClientTests.cs | 86 +++++++++++-------- 2 files changed, 74 insertions(+), 96 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs index 7fcec7779bd..925c0e9544e 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -4,6 +4,7 @@ using System; using System.Collections.Concurrent; using System.Collections.Generic; +using System.Diagnostics; using System.Diagnostics.CodeAnalysis; using System.Runtime.ExceptionServices; using System.Threading; @@ -20,8 +21,8 @@ namespace Microsoft.Extensions.AI; /// 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. +/// The configured clients are snapshotted by the constructor. The same client may appear more than once, in which +/// case it is invoked once per position. When every client has failed, the final failure is rethrown. /// /// [Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] @@ -29,7 +30,10 @@ public sealed class OrderedFailoverChatClient : FailoverChatClient { private readonly bool _leaveOpen; private readonly IChatClient[] _clients; - private readonly ConcurrentDictionary _requestStates = new(); + + // Holds the next client index for a request that has a failed attempt. A nonterminal update is always followed + // by another selection, so a stored index is always in range. + private readonly ConcurrentDictionary _requestStates = new(); private bool _disposed; /// Initializes a new instance of the class. @@ -39,10 +43,7 @@ public sealed class OrderedFailoverChatClient : FailoverChatClient /// otherwise, . /// /// is . - /// - /// is empty, contains , or contains the same client instance more - /// than once. - /// + /// is empty or contains . public OrderedFailoverChatClient(IReadOnlyList clients, bool leaveOpen = false) { _ = Throw.IfNull(clients); @@ -53,20 +54,12 @@ public OrderedFailoverChatClient(IReadOnlyList clients, bool leaveO Throw.ArgumentException(nameof(clients), "At least one client must be provided."); } - for (int i = 0; i < clientsSnapshot.Length; i++) + foreach (IChatClient client in clientsSnapshot) { - if (clientsSnapshot[i] is null) + if (client is null) { Throw.ArgumentException(nameof(clients), "Clients must not contain null."); } - - for (int j = 0; j < i; j++) - { - if (ReferenceEquals(clientsSnapshot[j], clientsSnapshot[i])) - { - Throw.ArgumentException(nameof(clients), "Each client instance must be unique."); - } - } } _leaveOpen = leaveOpen; @@ -78,20 +71,12 @@ protected override ValueTask SelectClientAsync( RoutingContext context, CancellationToken cancellationToken) { + _ = Throw.IfNull(context); _ = cancellationToken; - if (!_requestStates.TryRemove(context, out RequestState? state)) - { - return new(_clients[0]); - } - - if (state.ClientIndex < _clients.Length) - { - return new(_clients[state.ClientIndex]); - } + int clientIndex = _requestStates.TryGetValue(context, out int nextClientIndex) ? nextClientIndex : 0; - ExceptionDispatchInfo.Capture(state.LastException).Throw(); - throw state.LastException; + return new(_clients[clientIndex]); } /// @@ -109,23 +94,20 @@ protected override ValueTask OnRoutingUpdateAsync( return default; } - if (attempt.Exception is null) - { - _ = _requestStates.TryRemove(context, out _); - throw new InvalidOperationException("A nonterminal routing update requires a failed client invocation."); - } + Exception? exception = attempt.Exception; + Debug.Assert(exception is not null, "A nonterminal update always reports a failed invocation."); - int clientIndex = IndexOfClient(attempt.Client); - if (clientIndex < 0) + int nextClientIndex = (_requestStates.TryGetValue(context, out int attemptedIndex) ? attemptedIndex : 0) + 1; + if (nextClientIndex < _clients.Length) { - _ = _requestStates.TryRemove(context, out _); - throw new InvalidOperationException("The invocation did not use a configured ordered failover client."); + _requestStates[context] = nextClientIndex; + return default; } - // Selection immediately removes this state before invoking the next client. - _requestStates[context] = new(clientIndex + 1, attempt.Exception); - - return default; + // Every client has failed. Release the state before the final failure ends routing. + _ = _requestStates.TryRemove(context, out _); + ExceptionDispatchInfo.Capture(exception!).Throw(); + throw exception!; } /// @@ -149,24 +131,4 @@ protected override void Dispose(bool disposing) 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/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs index 5cb2cf34ed8..dd2f326617b 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs @@ -21,7 +21,6 @@ public void OrderedFailover_RejectsMissingClients() Assert.Throws(() => new OrderedFailoverChatClient(null!)); Assert.Throws(() => new OrderedFailoverChatClient([])); Assert.Throws(() => new OrderedFailoverChatClient([inner, null!])); - Assert.Throws(() => new OrderedFailoverChatClient([inner, inner])); } [Fact] @@ -54,18 +53,55 @@ public async Task OrderedFailover_TriesClientsInOrderAndSkipsFailures() } [Fact] - public async Task OrderedFailover_DistinguishesValueEqualClients() + public async Task OrderedFailover_InvokesRepeatedClientOncePerPosition() { + var calls = new List(); 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]); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls.Add("first"); + return calls.Count < 3 + ? throw new InvalidOperationException("first failed") + : Task.FromResult(expected); + }, + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls.Add("second"); + throw new InvalidOperationException("second failed"); + }, + }; + using var client = new OrderedFailoverChatClient([first, second, first]); ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); Assert.Same(expected, response); + Assert.Equal(["first", "second", "first"], calls); + } + + [Fact] + public async Task OrderedFailover_RepeatedClientExhaustionRethrowsLastFailure() + { + int calls = 0; + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls++; + throw new InvalidOperationException($"failure {calls}"); + }, + }; + using var client = new OrderedFailoverChatClient([inner, inner]); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => client.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal(2, calls); + Assert.Equal("failure 2", exception.Message); } [Fact] @@ -153,7 +189,7 @@ .. Enumerable.Range(0, RequestCount).Select(index => } [Fact] - public async Task OrderedFailover_DoesNotRetainStateWhileStreaming() + public async Task OrderedFailover_ReleasesStateWhenStreamingEnds() { using var first = new TestChatClient { @@ -164,11 +200,14 @@ public async Task OrderedFailover_DoesNotRetainStateWhileStreaming() 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); + 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)); } @@ -357,27 +396,4 @@ public IAsyncEnumerable GetStreamingResponseAsync( 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; - } }