diff --git a/src/Functions.WorkerProxy/ExtensionRpc/ExtensionCall.cs b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionCall.cs
new file mode 100644
index 0000000000..cd1ac9fa25
--- /dev/null
+++ b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionCall.cs
@@ -0,0 +1,130 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Channels;
+using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+///
+/// Represents one worker-originated gRPC call multiplexed over the shared extension RPC stream.
+///
+/// The physical extension RPC stream carrying the call.
+/// The identifier used to correlate lifecycle messages for the call.
+/// The negotiated transport settings for the stream.
+internal sealed class ExtensionCall(ExtensionRpcStream stream, string callId, ExtensionRpcReady ready)
+ : IAsyncDisposable
+{
+ private const int InboundQueueCapacity = 32;
+
+ private readonly Channel _inbound = Channel.CreateBounded(
+ new BoundedChannelOptions(InboundQueueCapacity)
+ {
+ SingleReader = true,
+ SingleWriter = true,
+ AllowSynchronousContinuations = false,
+ FullMode = BoundedChannelFullMode.Wait,
+ });
+
+ private readonly ExtensionRpcCreditWindow _requestCredits = new(ready.InitialReceiveWindowBytes);
+ private int _disposed;
+
+ ///
+ /// Gets the identifier used to correlate lifecycle messages for the call.
+ ///
+ public string CallId { get; } = callId;
+
+ ///
+ /// Gets the identifier of the physical extension RPC stream carrying this call.
+ ///
+ public string StreamId => stream.StreamId;
+
+ ///
+ /// Gets the transport settings negotiated for the stream.
+ ///
+ public ExtensionRpcReady Ready { get; } = ready;
+
+ ///
+ /// Gets the token that is cancelled when the physical stream closes.
+ ///
+ public CancellationToken CancellationToken => stream.CancellationToken;
+
+ ///
+ /// Gets a snapshot of the logical calls currently registered with the physical stream.
+ ///
+ public int ActiveCallCount => stream.ActiveCallCount;
+
+ ///
+ /// Reads ordered response lifecycle messages received from the host.
+ ///
+ /// A token that cancels response enumeration.
+ /// The ordered response messages for this call.
+ public IAsyncEnumerable ReadAllAsync(CancellationToken cancellationToken)
+ {
+ return _inbound.Reader.ReadAllAsync(cancellationToken);
+ }
+
+ ///
+ /// Writes a request lifecycle message to the shared extension RPC stream.
+ ///
+ /// The message to write.
+ /// A token that cancels the write.
+ /// A task that represents the asynchronous write.
+ public async ValueTask WriteAsync(ExtensionRpcMessage message, CancellationToken cancellationToken)
+ {
+ using CancellationTokenSource cancellationTokenSource =
+ CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, CancellationToken);
+ cancellationToken = cancellationTokenSource.Token;
+ if (message.ContentCase is ExtensionRpcMessage.ContentOneofCase.Data)
+ {
+ await _requestCredits.ReserveAsync((ulong)message.Data.Payload.Length, cancellationToken);
+ }
+
+ await stream.WriteExtensionMessageAsync(CallId, message, cancellationToken);
+ }
+
+ ///
+ /// Processes a response lifecycle message routed from the shared extension RPC stream.
+ ///
+ /// The inbound message to process.
+ /// A token that cancels queueing the message.
+ /// A task that represents the asynchronous operation.
+ public async ValueTask HandleInboundAsync(ExtensionRpcMessage message, CancellationToken cancellationToken)
+ {
+ if (message.ContentCase is ExtensionRpcMessage.ContentOneofCase.WindowUpdate)
+ {
+ _requestCredits.Add(message.WindowUpdate.ByteCount);
+ return;
+ }
+
+ await _inbound.Writer.WriteAsync(message, cancellationToken);
+ if (message.ContentCase is ExtensionRpcMessage.ContentOneofCase.Complete)
+ {
+ await DisposeAsync();
+ }
+ }
+
+ ///
+ /// Completes the inbound response queue because its physical stream has ended.
+ ///
+ public void Complete()
+ {
+ _inbound.Writer.TryComplete();
+ }
+
+ ///
+ public ValueTask DisposeAsync()
+ {
+ if (Interlocked.Exchange(ref _disposed, 1) is 0)
+ {
+ stream.RemoveCall(CallId);
+ _inbound.Writer.TryComplete();
+ }
+
+ return ValueTask.CompletedTask;
+ }
+}
diff --git a/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcCreditWindow.cs b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcCreditWindow.cs
new file mode 100644
index 0000000000..8234bb8506
--- /dev/null
+++ b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcCreditWindow.cs
@@ -0,0 +1,76 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+///
+/// Coordinates byte-based flow-control credits for one direction of an extension RPC call.
+///
+/// The number of bytes initially available to the sender.
+internal sealed class ExtensionRpcCreditWindow(ulong initialCredits)
+{
+ private readonly Lock _syncLock = new();
+ private ulong _available = initialCredits;
+ private TaskCompletionSource _changed = CreateChangedSource();
+
+ ///
+ /// Adds byte credits granted by the receiver.
+ ///
+ /// The number of credits to add.
+ public void Add(ulong credits)
+ {
+ if (credits is 0)
+ {
+ return;
+ }
+
+ TaskCompletionSource changed;
+ lock (_syncLock)
+ {
+ _available = ulong.MaxValue - _available < credits ? ulong.MaxValue : _available + credits;
+ changed = _changed;
+ _changed = CreateChangedSource();
+ }
+
+ changed.TrySetResult();
+ }
+
+ ///
+ /// Waits for and reserves the requested number of byte credits.
+ ///
+ /// The number of credits to reserve.
+ /// A token that cancels the wait.
+ /// A task that completes when the credits have been reserved.
+ public async ValueTask ReserveAsync(ulong credits, CancellationToken cancellationToken)
+ {
+ if (credits is 0)
+ {
+ return;
+ }
+
+ while (true)
+ {
+ Task waitTask;
+ lock (_syncLock)
+ {
+ if (_available >= credits)
+ {
+ _available -= credits;
+ return;
+ }
+
+ waitTask = _changed.Task;
+ }
+
+ await waitTask.WaitAsync(cancellationToken);
+ }
+ }
+
+ private static TaskCompletionSource CreateChangedSource()
+ {
+ return new(TaskCreationOptions.RunContinuationsAsynchronously);
+ }
+}
diff --git a/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcRelay.cs b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcRelay.cs
new file mode 100644
index 0000000000..dbae50144c
--- /dev/null
+++ b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcRelay.cs
@@ -0,0 +1,119 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Grpc.Core;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+using ExtensionRpcService = Microsoft.Azure.WebJobs.Script.Grpc.Messages.ExtensionRpc;
+using GrpcRpcException = Grpc.Core.RpcException;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+///
+/// Implements the runtime-facing extension RPC service and attaches its stream to the coordinator.
+///
+/// The WorkerProxy listener configuration.
+/// The coordinator that owns the active extension RPC stream.
+internal sealed class ExtensionRpcRelay(
+ WorkerProxyEndpointConfiguration endpoints, ExtensionRpcStreamCoordinator streamCoordinator)
+ : ExtensionRpcService.ExtensionRpcBase
+{
+ ///
+ public override Task EventStream(
+ IAsyncStreamReader requestStream,
+ IServerStreamWriter responseStream,
+ ServerCallContext context)
+ {
+ HttpContext httpContext = context.GetHttpContext();
+ if (!endpoints.TryGetRelaySide(httpContext.Connection.LocalPort, out FunctionRpcRelaySide side))
+ {
+ throw new GrpcRpcException(
+ new Status(StatusCode.Unimplemented, "ExtensionRpc is unavailable on this listener."));
+ }
+
+ if (side is not FunctionRpcRelaySide.Runtime)
+ {
+ throw new GrpcRpcException(
+ new Status(StatusCode.PermissionDenied, "ExtensionRpc is only available on the runtime gRPC port."));
+ }
+
+ return RelayAsync(requestStream, responseStream, context.CancellationToken);
+ }
+
+ ///
+ /// Relays inbound and outbound lifecycle messages for one physical extension RPC stream.
+ ///
+ /// Messages received from the host runtime.
+ /// Messages sent to the host runtime.
+ /// A token that is cancelled when the stream ends.
+ /// A task that represents the stream lifetime.
+ internal async Task RelayAsync(
+ IAsyncStreamReader requestStream,
+ IServerStreamWriter responseStream,
+ CancellationToken cancellationToken)
+ {
+ using CancellationTokenSource cancellationTokenSource =
+ CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+ await using ExtensionRpcStreamLease lease = OpenStream(cancellationTokenSource.Token);
+ Task readTask = ReadInboundAsync(lease.Stream, requestStream, cancellationTokenSource.Token);
+ Task writeTask = WriteOutboundAsync(lease.Stream, responseStream, cancellationTokenSource.Token);
+
+ try
+ {
+ Task completedTask = await Task.WhenAny(readTask, writeTask);
+ await completedTask;
+ }
+ catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+ {
+ }
+ finally
+ {
+ cancellationTokenSource.Cancel();
+ await Task.WhenAll(
+ ObserveCompletionAsync(readTask),
+ ObserveCompletionAsync(writeTask));
+ }
+ }
+
+ private ExtensionRpcStreamLease OpenStream(CancellationToken cancellationToken)
+ {
+ try
+ {
+ return streamCoordinator.Open(cancellationToken);
+ }
+ catch (InvalidOperationException exception)
+ {
+ throw new GrpcRpcException(new Status(StatusCode.AlreadyExists, exception.Message));
+ }
+ }
+
+ private static async Task ReadInboundAsync(
+ ExtensionRpcStream stream,
+ IAsyncStreamReader requestStream,
+ CancellationToken cancellationToken)
+ {
+ while (await requestStream.MoveNext(cancellationToken))
+ {
+ await stream.HandleInboundAsync(requestStream.Current, cancellationToken);
+ }
+ }
+
+ private static async Task WriteOutboundAsync(
+ ExtensionRpcStream stream,
+ IServerStreamWriter responseStream,
+ CancellationToken cancellationToken)
+ {
+ await foreach (ExtensionRpcMessage message in stream.Outbound.ReadAllAsync(cancellationToken))
+ {
+ await responseStream.WriteAsync(message, cancellationToken);
+ }
+ }
+
+ private static async Task ObserveCompletionAsync(Task task)
+ {
+ await task.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
+ }
+}
diff --git a/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStream.cs b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStream.cs
new file mode 100644
index 0000000000..fbda2bf45a
--- /dev/null
+++ b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStream.cs
@@ -0,0 +1,243 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Concurrent;
+using System.Threading;
+using System.Threading.Channels;
+using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+///
+/// Represents the physical host extension RPC stream that multiplexes logical extension calls.
+///
+internal sealed class ExtensionRpcStream
+{
+ private readonly ExtensionRpcStreamCoordinator _owner;
+ private readonly CancellationTokenSource _cancellationTokenSource;
+ private readonly ConcurrentDictionary _calls = new();
+ private readonly Channel _outbound = Channel.CreateBounded(
+ new BoundedChannelOptions(256)
+ {
+ SingleReader = true,
+ SingleWriter = false,
+ AllowSynchronousContinuations = false,
+ FullMode = BoundedChannelFullMode.Wait,
+ });
+
+ private volatile ExtensionRpcReady? _ready;
+ private int _closed;
+
+ ///
+ /// Initializes a new instance of the class.
+ ///
+ /// The coordinator that owns this stream.
+ /// The identifier shared by reconnects in the current session.
+ /// The identifier for this physical stream instance.
+ /// A token that is cancelled when the transport ends.
+ public ExtensionRpcStream(
+ ExtensionRpcStreamCoordinator owner, string sessionId, string streamId, CancellationToken cancellationToken)
+ {
+ _owner = owner;
+ SessionId = sessionId;
+ StreamId = streamId;
+ _cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+ CancellationToken = _cancellationTokenSource.Token;
+ }
+
+ ///
+ /// Gets the extension RPC session identifier.
+ ///
+ public string SessionId { get; }
+
+ ///
+ /// Gets the physical stream identifier.
+ ///
+ public string StreamId { get; }
+
+ ///
+ /// Gets the token that is cancelled when this stream closes.
+ ///
+ public CancellationToken CancellationToken { get; }
+
+ ///
+ /// Gets the ordered messages waiting to be written to the host.
+ ///
+ public ChannelReader Outbound => _outbound.Reader;
+
+ ///
+ /// Gets the number of logical calls registered with this stream.
+ ///
+ public int ActiveCallCount => _calls.Count;
+
+ ///
+ /// Gets a value indicating whether negotiation completed with the stream enabled.
+ ///
+ public bool IsReady => _ready is { Enabled: true };
+
+ ///
+ /// Gets a value indicating whether the host has completed stream negotiation.
+ ///
+ public bool IsNegotiated => _ready is not null;
+
+ ///
+ /// Processes a lifecycle message received from the host.
+ ///
+ /// The inbound message.
+ /// A token that cancels message processing.
+ /// A task that represents the asynchronous operation.
+ public async ValueTask HandleInboundAsync(ExtensionRpcMessage message, CancellationToken cancellationToken)
+ {
+ if (!string.Equals(message.SessionId, SessionId, StringComparison.Ordinal)
+ || !string.Equals(message.ShardId, StreamId, StringComparison.Ordinal))
+ {
+ return;
+ }
+
+ if (message.ContentCase is ExtensionRpcMessage.ContentOneofCase.SessionClosed)
+ {
+ _owner.CloseSession(this);
+ return;
+ }
+
+ if (message.ContentCase is ExtensionRpcMessage.ContentOneofCase.Ready)
+ {
+ _ready = IsValidReady(message.Ready)
+ ? message.Ready
+ : new ExtensionRpcReady
+ {
+ Enabled = false,
+ RejectionReason = "The host returned invalid extension RPC negotiation settings.",
+ };
+ _owner.SignalAvailabilityChanged();
+ return;
+ }
+
+ if (_calls.TryGetValue(message.CallId, out ExtensionCall? call))
+ {
+ await call.HandleInboundAsync(message, cancellationToken);
+ }
+ }
+
+ ///
+ /// Registers a logical call and writes its start message.
+ ///
+ /// The identifier assigned to the call.
+ /// The call start message.
+ /// A token that cancels opening the call.
+ /// The registered extension call.
+ public async Task OpenExtensionCallAsync(
+ string callId, ExtensionRpcStart start, CancellationToken cancellationToken)
+ {
+ ExtensionRpcReady ready = _ready
+ ?? throw new InvalidOperationException($"Extension RPC stream '{StreamId}' is not ready.");
+ if (!ready.Enabled)
+ {
+ throw new InvalidOperationException(
+ $"The host disabled extension RPC stream '{StreamId}': {ready.RejectionReason}");
+ }
+
+ ExtensionCall call = new(this, callId, ready);
+ if (!_calls.TryAdd(callId, call))
+ {
+ throw new InvalidOperationException($"Extension call '{callId}' is already registered.");
+ }
+
+ bool opened = false;
+ try
+ {
+ await call.WriteAsync(new ExtensionRpcMessage { Start = start }, cancellationToken);
+ opened = true;
+ return call;
+ }
+ finally
+ {
+ if (!opened)
+ {
+ _calls.TryRemove(callId, out _);
+ call.Complete();
+ }
+ }
+ }
+
+ ///
+ /// Closes the stream and completes all logical calls registered with it.
+ ///
+ public void Close()
+ {
+ if (Interlocked.Exchange(ref _closed, 1) is not 0)
+ {
+ return;
+ }
+
+ if (!_cancellationTokenSource.IsCancellationRequested)
+ {
+ _cancellationTokenSource.Cancel();
+ }
+
+ _outbound.Writer.TryComplete();
+ foreach (ExtensionCall call in _calls.Values)
+ {
+ call.Complete();
+ }
+
+ _calls.Clear();
+ _cancellationTokenSource.Dispose();
+ }
+
+ ///
+ /// Adds stream and call correlation identifiers and queues a message for the host.
+ ///
+ /// The logical call identifier.
+ /// The message to queue.
+ /// A token that cancels queueing the message.
+ /// A task that represents the asynchronous operation.
+ internal async ValueTask WriteExtensionMessageAsync(
+ string callId, ExtensionRpcMessage message, CancellationToken cancellationToken)
+ {
+ message.SessionId = SessionId;
+ message.ShardId = StreamId;
+ message.CallId = callId;
+ await _outbound.Writer.WriteAsync(message, cancellationToken);
+ }
+
+ ///
+ /// Removes a completed logical call from this stream.
+ ///
+ /// The identifier of the call to remove.
+ internal void RemoveCall(string callId)
+ {
+ _calls.TryRemove(callId, out _);
+ }
+
+ ///
+ /// Queues the initial protocol and transport-capability negotiation message.
+ ///
+ internal void QueueHello()
+ {
+ _outbound.Writer.TryWrite(
+ new ExtensionRpcMessage
+ {
+ SessionId = SessionId,
+ ShardId = StreamId,
+ Hello = new ExtensionRpcHello
+ {
+ InitialReceiveWindowBytes = ExtensionRpcStreamCoordinator.DefaultInitialWindowSize,
+ MaxDataChunkBytes = ExtensionRpcStreamCoordinator.DefaultMaxChunkSize,
+ MaxMessageBytes = ExtensionRpcStreamCoordinator.DefaultMaxMessageSize,
+ SupportedVersions = { ExtensionRpcStreamCoordinator.ProtocolVersion },
+ },
+ });
+ }
+
+ private static bool IsValidReady(ExtensionRpcReady ready)
+ {
+ return !ready.Enabled || (ready.SelectedVersion == ExtensionRpcStreamCoordinator.ProtocolVersion
+ && ready.InitialReceiveWindowBytes is > 0 and <= ExtensionRpcStreamCoordinator.DefaultInitialWindowSize
+ && ready.MaxDataChunkBytes is > 0 and <= ExtensionRpcStreamCoordinator.DefaultMaxChunkSize
+ && ready.MaxDataChunkBytes <= ready.InitialReceiveWindowBytes
+ && ready.MaxMessageBytes is > 0 and <= ExtensionRpcStreamCoordinator.DefaultMaxMessageSize);
+ }
+}
diff --git a/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStreamCoordinator.cs b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStreamCoordinator.cs
new file mode 100644
index 0000000000..e670466cff
--- /dev/null
+++ b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStreamCoordinator.cs
@@ -0,0 +1,253 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Diagnostics;
+using System.Linq;
+using System.Threading;
+using System.Threading.Channels;
+using System.Threading.Tasks;
+using Google.Protobuf;
+using Google.Protobuf.WellKnownTypes;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+///
+/// Coordinates the active host extension RPC stream and worker-originated logical calls.
+///
+internal sealed class ExtensionRpcStreamCoordinator
+{
+ internal const uint ProtocolVersion = 1;
+ internal const uint DefaultMaxChunkSize = 64 * 1024;
+ internal const ulong DefaultInitialWindowSize = 1024 * 1024;
+ internal const ulong DefaultMaxMessageSize = 16 * 1024 * 1024;
+
+ private readonly Lock _syncLock = new();
+ private TaskCompletionSource _availabilityChanged = CreateChangedSource();
+ private ExtensionRpcStream? _stream;
+ private string? _sessionId;
+ private long _nextCallId;
+ private long _nextStreamId;
+
+ ///
+ /// Gets a value indicating whether a host extension RPC stream is currently connected.
+ ///
+ public bool HasConnectedStream => Volatile.Read(ref _stream) is not null;
+
+ ///
+ /// Registers the physical extension RPC stream used to relay logical calls.
+ ///
+ /// A token that is cancelled when the physical stream ends.
+ /// A lease that unregisters and closes the stream when disposed.
+ public ExtensionRpcStreamLease Open(CancellationToken cancellationToken)
+ {
+ ExtensionRpcStream stream;
+ lock (_syncLock)
+ {
+ if (_stream is not null)
+ {
+ throw new InvalidOperationException("An extension RPC stream is already connected.");
+ }
+
+ _sessionId ??= Guid.NewGuid().ToString("N");
+ string streamId = Interlocked.Increment(ref _nextStreamId).ToStringInvariant();
+ stream = new ExtensionRpcStream(this, _sessionId, streamId, cancellationToken);
+ _stream = stream;
+ SignalAvailabilityChangedUnsynchronized();
+ }
+
+ stream.QueueHello();
+
+ return new ExtensionRpcStreamLease(this, stream);
+ }
+
+ ///
+ /// Opens a logical extension call on the connected and negotiated stream.
+ ///
+ /// The start message describing the worker-facing gRPC request.
+ /// A token that cancels opening the call.
+ /// The opened extension call.
+ public async Task OpenExtensionCallAsync(
+ ExtensionRpcStart start, CancellationToken cancellationToken)
+ {
+ TimeSpan? timeout = start.Timeout?.ToTimeSpan();
+ long waitStart = Stopwatch.GetTimestamp();
+ while (true)
+ {
+ ExtensionRpcStream? stream;
+ Task? availabilityTask = null;
+ lock (_syncLock)
+ {
+ stream = _stream ?? throw new InvalidOperationException("No extension RPC stream is connected.");
+ if (!stream.IsReady)
+ {
+ if (stream.IsNegotiated)
+ {
+ throw new InvalidOperationException("The host disabled the extension RPC stream.");
+ }
+
+ availabilityTask = _availabilityChanged.Task;
+ stream = null;
+ }
+ }
+
+ if (stream is not null)
+ {
+ if (timeout is not null)
+ {
+ TimeSpan remaining = GetRemainingTimeout(timeout.Value, waitStart, cancellationToken);
+ start.Timeout = Duration.FromTimeSpan(remaining);
+ UpdateTimeoutMetadata(start, remaining);
+ }
+
+ string callId = Interlocked.Increment(ref _nextCallId).ToStringInvariant();
+ try
+ {
+ return await stream.OpenExtensionCallAsync(callId, start, cancellationToken);
+ }
+ catch (OperationCanceledException) when (
+ stream.CancellationToken.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+ {
+ }
+ catch (ChannelClosedException) when (!cancellationToken.IsCancellationRequested)
+ {
+ }
+
+ continue;
+ }
+
+ if (timeout is null)
+ {
+ await availabilityTask!.WaitAsync(cancellationToken);
+ }
+ else
+ {
+ TimeSpan remaining = GetRemainingTimeout(timeout.Value, waitStart, cancellationToken);
+ await availabilityTask!.WaitAsync(remaining, cancellationToken);
+ }
+ }
+ }
+
+ ///
+ /// Unregisters and closes the specified physical stream.
+ ///
+ /// The stream to close.
+ internal void Close(ExtensionRpcStream stream)
+ {
+ Close(stream, endSession: false);
+ }
+
+ ///
+ /// Ends the session and closes its active physical stream.
+ ///
+ /// The stream whose session is ending.
+ internal void CloseSession(ExtensionRpcStream stream)
+ {
+ Close(stream, endSession: true);
+ }
+
+ private void Close(ExtensionRpcStream stream, bool endSession)
+ {
+ lock (_syncLock)
+ {
+ if (ReferenceEquals(_stream, stream))
+ {
+ _stream = null;
+ if (endSession)
+ {
+ _sessionId = null;
+ }
+ }
+
+ SignalAvailabilityChangedUnsynchronized();
+ }
+
+ stream.Close();
+ }
+
+ private static TimeSpan GetRemainingTimeout(
+ TimeSpan timeout, long waitStart, CancellationToken cancellationToken)
+ {
+ TimeSpan remaining = timeout - Stopwatch.GetElapsedTime(waitStart);
+ if (remaining <= TimeSpan.Zero)
+ {
+ cancellationToken.ThrowIfCancellationRequested();
+ throw new TimeoutException("The extension RPC call timed out before it could be opened.");
+ }
+
+ return remaining;
+ }
+
+ ///
+ /// Signals callers waiting for the stream negotiation state to change.
+ ///
+ internal void SignalAvailabilityChanged()
+ {
+ lock (_syncLock)
+ {
+ SignalAvailabilityChangedUnsynchronized();
+ }
+ }
+
+ private static void UpdateTimeoutMetadata(ExtensionRpcStart start, TimeSpan timeout)
+ {
+ string value = FormatTimeout(timeout);
+ ExtensionRpcMetadataEntry? timeoutEntry = start.Metadata.FirstOrDefault(
+ entry => string.Equals(entry.Key, "grpc-timeout", StringComparison.OrdinalIgnoreCase));
+ if (timeoutEntry is null)
+ {
+ start.Metadata.Add(
+ new ExtensionRpcMetadataEntry
+ {
+ Key = "grpc-timeout",
+ Value = ByteString.CopyFromUtf8(value),
+ });
+ }
+ else
+ {
+ timeoutEntry.Value = ByteString.CopyFromUtf8(value);
+ }
+ }
+
+ private static string FormatTimeout(TimeSpan timeout)
+ {
+ long ticks = Math.Max(1, timeout.Ticks);
+ if (ticks <= 999_999)
+ {
+ return $"{(ticks * 100).ToStringInvariant()}n";
+ }
+
+ (long Divisor, char Unit)[] units =
+ [
+ (TimeSpan.TicksPerMicrosecond, 'u'),
+ (TimeSpan.TicksPerMillisecond, 'm'),
+ (TimeSpan.TicksPerSecond, 'S'),
+ (TimeSpan.TicksPerMinute, 'M'),
+ (TimeSpan.TicksPerHour, 'H'),
+ ];
+
+ foreach ((long divisor, char unit) in units)
+ {
+ long value = (ticks + divisor - 1) / divisor;
+ if (value <= 99_999_999)
+ {
+ return $"{value.ToStringInvariant()}{unit}";
+ }
+ }
+
+ throw new InvalidOperationException("The extension gRPC timeout exceeds the protocol limit.");
+ }
+
+ private void SignalAvailabilityChangedUnsynchronized()
+ {
+ TaskCompletionSource changed = _availabilityChanged;
+ _availabilityChanged = CreateChangedSource();
+ changed.TrySetResult();
+ }
+
+ private static TaskCompletionSource CreateChangedSource()
+ {
+ return new(TaskCreationOptions.RunContinuationsAsynchronously);
+ }
+}
diff --git a/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStreamLease.cs b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStreamLease.cs
new file mode 100644
index 0000000000..0ae76fac48
--- /dev/null
+++ b/src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcStreamLease.cs
@@ -0,0 +1,29 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Threading.Tasks;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+///
+/// Owns the lifetime of the single physical extension RPC stream registered with the coordinator.
+///
+/// The coordinator that owns the stream registration.
+/// The registered stream.
+internal sealed class ExtensionRpcStreamLease(ExtensionRpcStreamCoordinator owner, ExtensionRpcStream stream)
+ : IAsyncDisposable
+{
+ ///
+ /// Gets the registered extension RPC stream.
+ ///
+ public ExtensionRpcStream Stream => stream;
+
+ ///
+ public ValueTask DisposeAsync()
+ {
+ owner.Close(Stream);
+
+ return ValueTask.CompletedTask;
+ }
+}
diff --git a/src/Functions.WorkerProxy/NumberExtensions.cs b/src/Functions.WorkerProxy/NumberExtensions.cs
new file mode 100644
index 0000000000..b5b68f57a7
--- /dev/null
+++ b/src/Functions.WorkerProxy/NumberExtensions.cs
@@ -0,0 +1,32 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System.Globalization;
+
+namespace Azure.Functions.WorkerProxy;
+
+///
+/// Provides culture-invariant numeric formatting helpers.
+///
+internal static class NumberExtensions
+{
+ ///
+ /// Formats a 32-bit integer using the invariant culture.
+ ///
+ /// The value to format.
+ /// The invariant string representation.
+ public static string ToStringInvariant(this int value)
+ {
+ return value.ToString(CultureInfo.InvariantCulture);
+ }
+
+ ///
+ /// Formats a 64-bit integer using the invariant culture.
+ ///
+ /// The value to format.
+ /// The invariant string representation.
+ public static string ToStringInvariant(this long value)
+ {
+ return value.ToString(CultureInfo.InvariantCulture);
+ }
+}
diff --git a/src/Functions.WorkerProxy/WorkerProxyApplication.cs b/src/Functions.WorkerProxy/WorkerProxyApplication.cs
index 4be19a627a..d7a625b6ed 100644
--- a/src/Functions.WorkerProxy/WorkerProxyApplication.cs
+++ b/src/Functions.WorkerProxy/WorkerProxyApplication.cs
@@ -5,6 +5,7 @@
using System.Net;
using System.Net.Http;
using System.Threading.Tasks;
+using Azure.Functions.WorkerProxy.ExtensionRpc;
using Azure.Functions.WorkerProxy.Http;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Hosting;
@@ -50,6 +51,7 @@ public static WebApplication Build(string[] args)
});
builder.Services.AddSingleton();
builder.Services.AddHostedService(static services => services.GetRequiredService());
+ builder.Services.AddSingleton();
ConfigureHttpForwarding(builder);
WebApplication app = builder.Build();
@@ -77,6 +79,7 @@ private static void ConfigureGrpcPipeline(IApplicationBuilder app)
app.UseEndpoints(static endpoints =>
{
endpoints.MapGrpcService();
+ endpoints.MapGrpcService();
});
}
diff --git a/test/Functions.WorkerProxy.Tests/ExtensionRpcTransportTests.cs b/test/Functions.WorkerProxy.Tests/ExtensionRpcTransportTests.cs
new file mode 100644
index 0000000000..f2dbf289c8
--- /dev/null
+++ b/test/Functions.WorkerProxy.Tests/ExtensionRpcTransportTests.cs
@@ -0,0 +1,506 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Channels;
+using System.Threading.Tasks;
+using Azure.Functions.WorkerProxy.ExtensionRpc;
+using Google.Protobuf;
+using Google.Protobuf.WellKnownTypes;
+using Grpc.Core;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+using Microsoft.Extensions.Options;
+using Xunit;
+using GrpcRpcException = Grpc.Core.RpcException;
+
+namespace Azure.Functions.WorkerProxy.Tests;
+
+public class ExtensionRpcTransportTests
+{
+ [Fact]
+ public async Task OpenExtensionCall_RoutesMessagesThroughSession()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+
+ ExtensionRpcMessage hello = await lease.Stream.Outbound.ReadAsync();
+ Assert.Contains(ExtensionRpcStreamCoordinator.ProtocolVersion, hello.Hello.SupportedVersions);
+ await lease.Stream.HandleInboundAsync(CreateReady(hello), CancellationToken.None);
+
+ await using ExtensionCall call = await streamCoordinator.OpenExtensionCallAsync(
+ new ExtensionRpcStart { Method = "/extensions.Echo/Unary" },
+ CancellationToken.None);
+
+ ExtensionRpcMessage start = await lease.Stream.Outbound.ReadAsync();
+ Assert.Equal(hello.SessionId, start.SessionId);
+ Assert.Equal(hello.ShardId, start.ShardId);
+ Assert.Equal(call.CallId, start.CallId);
+ Assert.Equal("/extensions.Echo/Unary", start.Start.Method);
+
+ await lease.Stream.HandleInboundAsync(
+ new ExtensionRpcMessage
+ {
+ SessionId = hello.SessionId,
+ ShardId = hello.ShardId,
+ CallId = call.CallId,
+ Headers = new ExtensionRpcHeaders(),
+ },
+ CancellationToken.None);
+
+ ExtensionRpcMessage received = await FirstAsync(call.ReadAllAsync(CancellationToken.None));
+ Assert.Equal(ExtensionRpcMessage.ContentOneofCase.Headers, received.ContentCase);
+ }
+
+ [Fact]
+ public async Task Open_RejectsSecondConcurrentPhysicalStream()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+
+ Assert.Throws(() => streamCoordinator.Open(CancellationToken.None));
+ }
+
+ [Fact]
+ public async Task OpenExtensionCall_AllowsMoreThanPreviousCapacity()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage hello = await lease.Stream.Outbound.ReadAsync();
+ await lease.Stream.HandleInboundAsync(CreateReady(hello), CancellationToken.None);
+ var calls = new List();
+
+ try
+ {
+ for (int i = 0; i < 129; i++)
+ {
+ calls.Add(await streamCoordinator.OpenExtensionCallAsync(
+ new ExtensionRpcStart { Method = "/extensions.Echo/Unary" },
+ CancellationToken.None));
+ await lease.Stream.Outbound.ReadAsync();
+ }
+
+ Assert.Equal(129, lease.Stream.ActiveCallCount);
+ }
+ finally
+ {
+ foreach (ExtensionCall call in calls)
+ {
+ await call.DisposeAsync();
+ }
+ }
+ }
+
+ [Fact]
+ public async Task OpenExtensionCall_PropagatesRemainingTimeout()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage hello = await lease.Stream.Outbound.ReadAsync();
+ await lease.Stream.HandleInboundAsync(CreateReady(hello), CancellationToken.None);
+ var start = new ExtensionRpcStart
+ {
+ Method = "/extensions.Echo/Unary",
+ Timeout = Duration.FromTimeSpan(TimeSpan.FromSeconds(1)),
+ };
+ start.Metadata.Add(
+ new ExtensionRpcMetadataEntry
+ {
+ Key = "grpc-timeout",
+ Value = ByteString.CopyFromUtf8("1S"),
+ });
+
+ await using ExtensionCall call =
+ await streamCoordinator.OpenExtensionCallAsync(start, CancellationToken.None);
+ ExtensionRpcMessage message = await lease.Stream.Outbound.ReadAsync();
+
+ Assert.InRange(message.Start.Timeout.ToTimeSpan(), TimeSpan.Zero, TimeSpan.FromSeconds(1));
+ Assert.NotEqual(
+ "1S",
+ Assert.Single(
+ message.Start.Metadata,
+ entry => string.Equals(entry.Key, "grpc-timeout", StringComparison.OrdinalIgnoreCase))
+ .Value.ToStringUtf8());
+ }
+
+ [Fact]
+ public async Task OpenExtensionCall_ThrowsTimeoutExceptionWhenTimeoutExpires()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage hello = await lease.Stream.Outbound.ReadAsync();
+ await lease.Stream.HandleInboundAsync(CreateReady(hello), CancellationToken.None);
+
+ await Assert.ThrowsAsync(
+ () => streamCoordinator.OpenExtensionCallAsync(
+ new ExtensionRpcStart
+ {
+ Method = "/extensions.Echo/Unary",
+ Timeout = Duration.FromTimeSpan(TimeSpan.Zero),
+ },
+ CancellationToken.None));
+ }
+
+ [Fact]
+ public async Task OpenExtensionCall_ThrowsTimeoutExceptionWhenNegotiationTimesOut()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+ await lease.Stream.Outbound.ReadAsync();
+ using var cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+
+ await Assert.ThrowsAsync(
+ () => streamCoordinator.OpenExtensionCallAsync(
+ new ExtensionRpcStart
+ {
+ Method = "/extensions.Echo/Unary",
+ Timeout = Duration.FromTimeSpan(TimeSpan.FromMilliseconds(50)),
+ },
+ cancellationTokenSource.Token));
+ }
+
+ [Fact]
+ public async Task SessionClosed_AllowsNewSession()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ string firstSessionId;
+ string firstStreamId;
+ await using (ExtensionRpcStreamLease firstLease = streamCoordinator.Open(CancellationToken.None))
+ {
+ ExtensionRpcMessage hello = await firstLease.Stream.Outbound.ReadAsync();
+ firstSessionId = hello.SessionId;
+ firstStreamId = hello.ShardId;
+ await firstLease.Stream.HandleInboundAsync(
+ new ExtensionRpcMessage
+ {
+ SessionId = hello.SessionId,
+ ShardId = hello.ShardId,
+ SessionClosed = new ExtensionRpcSessionClosed(),
+ },
+ CancellationToken.None);
+ }
+
+ Assert.False(streamCoordinator.HasConnectedStream);
+
+ await using ExtensionRpcStreamLease secondLease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage secondHello = await secondLease.Stream.Outbound.ReadAsync();
+ Assert.NotEqual(firstSessionId, secondHello.SessionId);
+ Assert.NotEqual(firstStreamId, secondHello.ShardId);
+ }
+
+ [Fact]
+ public async Task TransportDisconnect_PreservesSessionForReconnect()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ string firstSessionId;
+ string firstStreamId;
+ await using (ExtensionRpcStreamLease firstLease = streamCoordinator.Open(CancellationToken.None))
+ {
+ ExtensionRpcMessage hello = await firstLease.Stream.Outbound.ReadAsync();
+ firstSessionId = hello.SessionId;
+ firstStreamId = hello.ShardId;
+ }
+
+ await using ExtensionRpcStreamLease secondLease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage secondHello = await secondLease.Stream.Outbound.ReadAsync();
+
+ Assert.Equal(firstSessionId, secondHello.SessionId);
+ Assert.NotEqual(firstStreamId, secondHello.ShardId);
+ }
+
+ [Fact]
+ public async Task OpenExtensionCall_RejectsInvalidNegotiation()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage hello = await lease.Stream.Outbound.ReadAsync();
+ ExtensionRpcMessage ready = CreateReady(hello);
+ ready.Ready.SelectedVersion++;
+
+ await lease.Stream.HandleInboundAsync(ready, CancellationToken.None);
+
+ InvalidOperationException exception = await Assert.ThrowsAsync(
+ () => streamCoordinator.OpenExtensionCallAsync(
+ new ExtensionRpcStart { Method = "/extensions.Echo/Unary" },
+ CancellationToken.None));
+ Assert.Contains("disabled", exception.Message, StringComparison.OrdinalIgnoreCase);
+ }
+
+ [Fact]
+ public async Task WriteAsync_WaitsForReceiveWindowCredit()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage hello = await lease.Stream.Outbound.ReadAsync();
+ ExtensionRpcMessage ready = CreateReady(hello);
+ ready.Ready.InitialReceiveWindowBytes = 2;
+ ready.Ready.MaxDataChunkBytes = 2;
+ await lease.Stream.HandleInboundAsync(ready, CancellationToken.None);
+ await using ExtensionCall call = await streamCoordinator.OpenExtensionCallAsync(
+ new ExtensionRpcStart { Method = "/extensions.Echo/Unary" },
+ CancellationToken.None);
+ await lease.Stream.Outbound.ReadAsync();
+
+ await call.WriteAsync(
+ new ExtensionRpcMessage
+ {
+ Data = new ExtensionRpcData { Payload = ByteString.CopyFrom([1, 2]) },
+ },
+ CancellationToken.None);
+ await lease.Stream.Outbound.ReadAsync();
+
+ ValueTask blockedWrite = call.WriteAsync(
+ new ExtensionRpcMessage
+ {
+ Data = new ExtensionRpcData { Payload = ByteString.CopyFrom([3, 4]) },
+ },
+ CancellationToken.None);
+ Assert.False(blockedWrite.IsCompleted);
+
+ await lease.Stream.HandleInboundAsync(
+ new ExtensionRpcMessage
+ {
+ SessionId = hello.SessionId,
+ ShardId = hello.ShardId,
+ CallId = call.CallId,
+ WindowUpdate = new ExtensionRpcWindowUpdate { ByteCount = 2 },
+ },
+ CancellationToken.None);
+
+ await blockedWrite;
+ ExtensionRpcMessage data = await lease.Stream.Outbound.ReadAsync();
+ Assert.Equal(ByteString.CopyFrom([3, 4]), data.Data.Payload);
+ }
+
+ [Fact]
+ public async Task WriteAsync_StopsWaitingForReceiveWindowWhenStreamCloses()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage hello = await lease.Stream.Outbound.ReadAsync();
+ ExtensionRpcMessage ready = CreateReady(hello);
+ ready.Ready.InitialReceiveWindowBytes = 1;
+ ready.Ready.MaxDataChunkBytes = 1;
+ await lease.Stream.HandleInboundAsync(ready, CancellationToken.None);
+ await using ExtensionCall call = await streamCoordinator.OpenExtensionCallAsync(
+ new ExtensionRpcStart { Method = "/extensions.Echo/Unary" },
+ CancellationToken.None);
+ await lease.Stream.Outbound.ReadAsync();
+ await call.WriteAsync(
+ new ExtensionRpcMessage
+ {
+ Data = new ExtensionRpcData { Payload = ByteString.CopyFrom([1]) },
+ },
+ CancellationToken.None);
+ await lease.Stream.Outbound.ReadAsync();
+
+ ValueTask blockedWrite = call.WriteAsync(
+ new ExtensionRpcMessage
+ {
+ Data = new ExtensionRpcData { Payload = ByteString.CopyFrom([2]) },
+ },
+ CancellationToken.None);
+ Assert.False(blockedWrite.IsCompleted);
+
+ await lease.DisposeAsync();
+
+ await Assert.ThrowsAnyAsync(() => blockedWrite.AsTask());
+ }
+
+ [Fact]
+ public async Task Complete_RemovesCallFromStream()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ await using ExtensionRpcStreamLease lease = streamCoordinator.Open(CancellationToken.None);
+ ExtensionRpcMessage hello = await lease.Stream.Outbound.ReadAsync();
+ await lease.Stream.HandleInboundAsync(CreateReady(hello), CancellationToken.None);
+ await using ExtensionCall call = await streamCoordinator.OpenExtensionCallAsync(
+ new ExtensionRpcStart { Method = "/extensions.Echo/Unary" },
+ CancellationToken.None);
+ await lease.Stream.Outbound.ReadAsync();
+
+ await lease.Stream.HandleInboundAsync(
+ new ExtensionRpcMessage
+ {
+ SessionId = hello.SessionId,
+ ShardId = hello.ShardId,
+ CallId = call.CallId,
+ Complete = new ExtensionRpcComplete(),
+ },
+ CancellationToken.None);
+
+ Assert.Equal(0, lease.Stream.ActiveCallCount);
+ ExtensionRpcMessage complete = await FirstAsync(call.ReadAllAsync(CancellationToken.None));
+ Assert.Equal(ExtensionRpcMessage.ContentOneofCase.Complete, complete.ContentCase);
+ }
+
+ [Fact]
+ public async Task RelayAsync_ClosesCoordinatorWhenRuntimeStreamEnds()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ var endpoints = new WorkerProxyEndpointConfiguration(
+ Options.Create(new WorkerProxyOptions()));
+ var relay = new ExtensionRpcRelay(endpoints, streamCoordinator);
+ Channel inbound = Channel.CreateUnbounded();
+ var outbound = new TestServerStreamWriter();
+
+ Task relayTask = relay.RelayAsync(
+ new TestAsyncStreamReader(inbound.Reader),
+ outbound,
+ CancellationToken.None);
+
+ ExtensionRpcMessage hello = await outbound.Messages.ReadAsync();
+ Assert.Equal(ExtensionRpcMessage.ContentOneofCase.Hello, hello.ContentCase);
+ Assert.True(streamCoordinator.HasConnectedStream);
+
+ inbound.Writer.TryComplete();
+ await relayTask;
+
+ Assert.False(streamCoordinator.HasConnectedStream);
+ }
+
+ [Fact]
+ public async Task RelayAsync_RejectsSecondConcurrentPhysicalStreamWithAlreadyExists()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ var endpoints = new WorkerProxyEndpointConfiguration(
+ Options.Create(new WorkerProxyOptions()));
+ var relay = new ExtensionRpcRelay(endpoints, streamCoordinator);
+ Channel firstInbound = Channel.CreateUnbounded();
+ var firstOutbound = new TestServerStreamWriter();
+ Task firstRelayTask = relay.RelayAsync(
+ new TestAsyncStreamReader(firstInbound.Reader),
+ firstOutbound,
+ CancellationToken.None);
+ await firstOutbound.Messages.ReadAsync();
+
+ GrpcRpcException exception = await Assert.ThrowsAsync(
+ () => relay.RelayAsync(
+ new TestAsyncStreamReader(
+ Channel.CreateUnbounded().Reader),
+ new TestServerStreamWriter(),
+ CancellationToken.None));
+
+ Assert.Equal(StatusCode.AlreadyExists, exception.StatusCode);
+
+ firstInbound.Writer.TryComplete();
+ await firstRelayTask;
+ }
+
+ [Fact]
+ public async Task RelayAsync_PreservesFirstFailureDuringCleanup()
+ {
+ var streamCoordinator = new ExtensionRpcStreamCoordinator();
+ var endpoints = new WorkerProxyEndpointConfiguration(
+ Options.Create(new WorkerProxyOptions()));
+ var relay = new ExtensionRpcRelay(endpoints, streamCoordinator);
+ var expectedException = new InvalidOperationException("The inbound stream failed.");
+
+ InvalidOperationException exception = await Assert.ThrowsAsync(
+ () => relay.RelayAsync(
+ new ThrowingAsyncStreamReader(expectedException),
+ new FailAfterCancellationServerStreamWriter(),
+ CancellationToken.None));
+
+ Assert.Same(expectedException, exception);
+ }
+
+ private static ExtensionRpcMessage CreateReady(ExtensionRpcMessage hello)
+ {
+ return new ExtensionRpcMessage
+ {
+ SessionId = hello.SessionId,
+ ShardId = hello.ShardId,
+ Ready = new ExtensionRpcReady
+ {
+ SelectedVersion = ExtensionRpcStreamCoordinator.ProtocolVersion,
+ Enabled = true,
+ InitialReceiveWindowBytes = ExtensionRpcStreamCoordinator.DefaultInitialWindowSize,
+ MaxDataChunkBytes = ExtensionRpcStreamCoordinator.DefaultMaxChunkSize,
+ MaxMessageBytes = ExtensionRpcStreamCoordinator.DefaultMaxMessageSize,
+ },
+ };
+ }
+
+ private static async Task FirstAsync(IAsyncEnumerable source)
+ {
+ await foreach (T item in source)
+ {
+ return item;
+ }
+
+ throw new InvalidOperationException("The sequence contained no items.");
+ }
+
+ private sealed class TestAsyncStreamReader(ChannelReader reader) : IAsyncStreamReader
+ {
+ public T Current { get; private set; } = default!;
+
+ public async Task MoveNext(CancellationToken cancellationToken)
+ {
+ while (await reader.WaitToReadAsync(cancellationToken))
+ {
+ if (reader.TryRead(out T? item))
+ {
+ Current = item;
+ return true;
+ }
+ }
+
+ return false;
+ }
+ }
+
+ private sealed class ThrowingAsyncStreamReader(Exception exception) : IAsyncStreamReader
+ {
+ public T Current { get; } = default!;
+
+ public Task MoveNext(CancellationToken cancellationToken)
+ {
+ return Task.FromException(exception);
+ }
+ }
+
+ private sealed class FailAfterCancellationServerStreamWriter : IServerStreamWriter
+ {
+ public WriteOptions? WriteOptions { get; set; }
+
+ public Task WriteAsync(T message)
+ {
+ return Task.FromException(new InvalidOperationException("The outbound stream failed."));
+ }
+
+ public async Task WriteAsync(T message, CancellationToken cancellationToken)
+ {
+ try
+ {
+ await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken);
+ }
+ catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+ {
+ throw new InvalidOperationException("The outbound stream failed.");
+ }
+ }
+ }
+
+ private sealed class TestServerStreamWriter : IServerStreamWriter
+ {
+ private readonly Channel _messages = Channel.CreateUnbounded();
+
+ public ChannelReader Messages => _messages.Reader;
+
+ public WriteOptions? WriteOptions { get; set; }
+
+ public Task WriteAsync(T message)
+ {
+ return _messages.Writer.WriteAsync(message).AsTask();
+ }
+
+ public Task WriteAsync(T message, CancellationToken cancellationToken)
+ {
+ return _messages.Writer.WriteAsync(message, cancellationToken).AsTask();
+ }
+ }
+}