@azure-typespec/http-client-csharp-provisioning 1.0.0-alpha.20260313.4

package/dist/generator/Shared/Core/NextLinkOperationImplementation.cs

@@ -0,0 +1,701 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#nullable enable
+
+using System;
+using System.ClientModel.Primitives;
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using System.Net;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core.Pipeline;
+
+namespace Azure.Core
+{
+    internal class NextLinkOperationImplementation : IOperation
+    {
+        internal const string NotSet = "NOT_SET";
+        internal const string RehydrationTokenVersion = "1.0.0";
+        private const string ApiVersionParam = "api-version";
+        private static readonly string[] FailureStates = { "failed", "canceled" };
+        private static readonly string[] SuccessStates = { "succeeded" };
+
+        private readonly HeaderSource _headerSource;
+        private readonly Uri _startRequestUri;
+        private readonly OperationFinalStateVia _finalStateVia;
+        private readonly HttpPipeline _pipeline;
+        private readonly string? _apiVersion;
+
+        private string? _lastKnownLocation;
+        private string _nextRequestUri;
+
+        // We can only get OperationId when
+        // - The operation is still in progress and nextRequestUri contains it
+        // - During rehydration, rehydrationToken.Id is the operation id
+        public string OperationId { get; private set; } = NotSet;
+        public RequestMethod RequestMethod { get; }
+
+        public static IOperation Create(
+            HttpPipeline pipeline,
+            RequestMethod requestMethod,
+            Uri startRequestUri,
+            Response response,
+            OperationFinalStateVia finalStateVia,
+            bool skipApiVersionOverride = false,
+            string? apiVersionOverrideValue = null)
+        {
+            string? apiVersionStr = null;
+            if (apiVersionOverrideValue is not null)
+            {
+                apiVersionStr = apiVersionOverrideValue;
+            }
+            else
+            {
+                apiVersionStr = !skipApiVersionOverride && TryGetApiVersion(startRequestUri, out ReadOnlySpan<char> apiVersion) ? apiVersion.ToString() : null;
+            }
+            var headerSource = GetHeaderSource(requestMethod, startRequestUri, response, apiVersionStr, out string nextRequestUri, out bool isNextRequestPolling);
+
+            string? lastKnownLocation;
+            if (!response.Headers.TryGetValue("Location", out lastKnownLocation))
+            {
+                lastKnownLocation = null;
+            }
+
+            NextLinkOperationImplementation operation = new(pipeline, requestMethod, startRequestUri, nextRequestUri, headerSource, lastKnownLocation, finalStateVia, apiVersionStr, isNextRequestPolling: isNextRequestPolling);
+
+            if (headerSource == HeaderSource.None && IsFinalState(response, headerSource, out var failureState, out _))
+            {
+                return new CompletedOperation(failureState ?? GetOperationStateFromFinalResponse(requestMethod, response), operation);
+            }
+
+            return operation;
+        }
+
+        public static IOperation<T> Create<T>(
+            IOperationSource<T> operationSource,
+            HttpPipeline pipeline,
+            RequestMethod requestMethod,
+            Uri startRequestUri,
+            Response response,
+            OperationFinalStateVia finalStateVia,
+            bool skipApiVersionOverride = false,
+            string? apiVersionOverrideValue = null)
+        {
+            var operation = Create(pipeline, requestMethod, startRequestUri, response, finalStateVia, skipApiVersionOverride, apiVersionOverrideValue);
+            return new OperationToOperationOfT<T>(operationSource, operation);
+        }
+
+        public static IOperation<T> Create<T>(
+            IOperationSource<T> operationSource,
+            IOperation operation)
+            => new OperationToOperationOfT<T>(operationSource, operation);
+
+        public static IOperation Create(
+            HttpPipeline pipeline,
+            RehydrationToken rehydrationToken)
+        {
+            AssertNotNull(rehydrationToken, nameof(rehydrationToken));
+            AssertNotNull(pipeline, nameof(pipeline));
+
+            // TODO: Once we remove NextLinkOperationImplementation from internal shared and make it internal to Azure.Core only in https://github.com/Azure/azure-sdk-for-net/issues/43260
+            // We can access the internal members from RehydrationToken directly
+            var data = ModelReaderWriter.Write(rehydrationToken!, ModelReaderWriterOptions.Json, AzureCoreContext.Default);
+            using var document = JsonDocument.Parse(data);
+            var lroDetails = document.RootElement;
+
+            // We are sure that the following properties exists in the serialized rehydrationToken
+            var initialUri = lroDetails.GetProperty("initialUri").GetString();
+            if (!Uri.TryCreate(initialUri, UriKind.Absolute, out var startRequestUri))
+            {
+                throw new ArgumentException($"\"initialUri\" property on \"rehydrationToken\" is an invalid Uri", nameof(rehydrationToken));
+            }
+
+            // We are sure that the following properties(apart from nullable lastKnownLocation) are not null as they are required in the rehydrationToken
+            string nextRequestUri = lroDetails.GetProperty("nextRequestUri").GetString()!;
+            string requestMethodStr = lroDetails.GetProperty("requestMethod").GetString()!;
+            RequestMethod requestMethod = new RequestMethod(requestMethodStr)!;
+            string? lastKnownLocation = lroDetails.GetProperty("lastKnownLocation").GetString();
+
+            string finalStateViaStr = lroDetails.GetProperty("finalStateVia").GetString()!;
+            OperationFinalStateVia finalStateVia;
+            if (Enum.IsDefined(typeof(OperationFinalStateVia), finalStateViaStr))
+            {
+                finalStateVia = (OperationFinalStateVia)Enum.Parse(typeof(OperationFinalStateVia), finalStateViaStr);
+            }
+            else
+            {
+                finalStateVia = OperationFinalStateVia.Location;
+            }
+
+            string headerSourceStr = lroDetails.GetProperty("headerSource").GetString()!;
+            HeaderSource headerSource;
+            if (Enum.IsDefined(typeof(HeaderSource), headerSourceStr))
+            {
+                headerSource = (HeaderSource)Enum.Parse(typeof(HeaderSource), headerSourceStr);
+            }
+            else
+            {
+                headerSource = HeaderSource.None;
+            }
+
+            return new NextLinkOperationImplementation(pipeline, requestMethod, startRequestUri, nextRequestUri, headerSource, lastKnownLocation, finalStateVia, null, rehydrationToken.Id);
+        }
+
+        private NextLinkOperationImplementation(
+            HttpPipeline pipeline,
+            RequestMethod requestMethod,
+            Uri startRequestUri,
+            string nextRequestUri,
+            HeaderSource headerSource,
+            string? lastKnownLocation,
+            OperationFinalStateVia finalStateVia,
+            string? apiVersion,
+            string? operationId = null,
+            bool isNextRequestPolling = false)
+        {
+            AssertNotNull(pipeline, nameof(pipeline));
+            AssertNotNull(requestMethod, nameof(requestMethod));
+            AssertNotNull(startRequestUri, nameof(startRequestUri));
+            AssertNotNull(nextRequestUri, nameof(nextRequestUri));
+            AssertNotNull(headerSource, nameof(headerSource));
+            AssertNotNull(finalStateVia, nameof(finalStateVia));
+
+            RequestMethod = requestMethod;
+            _headerSource = headerSource;
+            _startRequestUri = startRequestUri;
+            _nextRequestUri = nextRequestUri;
+            _lastKnownLocation = lastKnownLocation;
+            _finalStateVia = finalStateVia;
+            _pipeline = pipeline;
+            _apiVersion = apiVersion;
+            if (operationId is not null)
+            {
+                OperationId = operationId;
+            }
+            else if (isNextRequestPolling)
+            {
+                OperationId = ParseOperationId(startRequestUri, nextRequestUri);
+            }
+        }
+
+        private static string ParseOperationId(Uri startRequestUri, string nextRequestUri)
+        {
+            if (Uri.TryCreate(nextRequestUri, UriKind.Absolute, out var nextLink) && nextLink.Scheme != "file")
+            {
+                return nextLink.Segments.Last();
+            }
+            else
+            {
+                return new Uri(startRequestUri, nextRequestUri).Segments.Last();
+            }
+        }
+
+        public RehydrationToken GetRehydrationToken()
+            => GetRehydrationToken(RequestMethod, _startRequestUri, _nextRequestUri, _headerSource.ToString(), _lastKnownLocation, _finalStateVia.ToString(), OperationId);
+
+        public static RehydrationToken GetRehydrationToken(
+            RequestMethod requestMethod,
+            Uri startRequestUri,
+            Response response,
+            OperationFinalStateVia finalStateVia)
+        {
+            AssertNotNull(requestMethod, nameof(requestMethod));
+            AssertNotNull(startRequestUri, nameof(startRequestUri));
+            AssertNotNull(response, nameof(response));
+            AssertNotNull(finalStateVia, nameof(finalStateVia));
+
+            var headerSource = GetHeaderSource(requestMethod, startRequestUri, response, null, out string nextRequestUri, out bool isNextRequestPolling);
+            string? lastKnownLocation;
+            if (!response.Headers.TryGetValue("Location", out lastKnownLocation))
+            {
+                lastKnownLocation = null;
+            }
+            return GetRehydrationToken(requestMethod, startRequestUri, nextRequestUri, headerSource.ToString(), lastKnownLocation, finalStateVia.ToString(), isNextRequestPolling ? ParseOperationId(startRequestUri, nextRequestUri) : null);
+        }
+
+        public static RehydrationToken GetRehydrationToken(
+            RequestMethod requestMethod,
+            Uri startRequestUri,
+            string nextRequestUri,
+            string headerSource,
+            string? lastKnownLocation,
+            string finalStateVia,
+            string? operationId = null)
+        {
+            // TODO: Once we remove NextLinkOperationImplementation from internal shared and make it internal to Azure.Core only in https://github.com/Azure/azure-sdk-for-net/issues/43260
+            // We can access the internal members from RehydrationToken directly
+            var json = $$"""
+            {"version":"{{RehydrationTokenVersion}}","id":{{ConstructStringValue(operationId)}},"requestMethod":"{{requestMethod}}","initialUri":"{{startRequestUri.AbsoluteUri}}","nextRequestUri":"{{nextRequestUri}}","headerSource":"{{headerSource}}","finalStateVia":"{{finalStateVia}}","lastKnownLocation":{{ConstructStringValue(lastKnownLocation)}}}
+            """;
+            var data = new BinaryData(json);
+            return ModelReaderWriter.Read<RehydrationToken>(data, ModelReaderWriterOptions.Json, AzureCoreContext.Default);
+        }
+
+        private static string? ConstructStringValue(string? value) => value is null ? "null" : $"\"{value}\"";
+
+        public async ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken)
+        {
+            Response response = async
+                ? await GetResponseAsync(_nextRequestUri, cancellationToken).ConfigureAwait(false)
+                : GetResponse(_nextRequestUri, cancellationToken);
+
+            var hasCompleted = IsFinalState(response, _headerSource, out var failureState, out var resourceLocation);
+            if (failureState != null)
+            {
+                return failureState.Value;
+            }
+
+            if (hasCompleted)
+            {
+                string? finalUri = GetFinalUri(resourceLocation);
+                Response finalResponse;
+                if (finalUri != null)
+                {
+                    finalResponse = async
+                        ? await GetResponseAsync(finalUri, cancellationToken).ConfigureAwait(false)
+                        : GetResponse(finalUri, cancellationToken);
+                }
+                else
+                {
+                    finalResponse = response;
+                }
+                return GetOperationStateFromFinalResponse(RequestMethod, finalResponse);
+            }
+
+            UpdateNextRequestUri(response.Headers);
+            return OperationState.Pending(response);
+        }
+
+        private static OperationState GetOperationStateFromFinalResponse(RequestMethod requestMethod, Response response)
+        {
+            switch (response.Status)
+            {
+                case 200:
+                case 201 when requestMethod == RequestMethod.Put:
+                case 204 when requestMethod != RequestMethod.Put && requestMethod != RequestMethod.Patch:
+                    return OperationState.Success(response);
+                default:
+                    return OperationState.Failure(response);
+            }
+        }
+
+        private void UpdateNextRequestUri(ResponseHeaders headers)
+        {
+            var hasLocation = headers.TryGetValue("Location", out string? location);
+            if (hasLocation)
+            {
+                _lastKnownLocation = location;
+            }
+
+            switch (_headerSource)
+            {
+                case HeaderSource.OperationLocation when headers.TryGetValue("Operation-Location", out string? operationLocation):
+                    _nextRequestUri = AppendOrReplaceApiVersion(operationLocation, _apiVersion);
+                    OperationId = ParseOperationId(_startRequestUri, _nextRequestUri);
+                    return;
+                case HeaderSource.AzureAsyncOperation when headers.TryGetValue("Azure-AsyncOperation", out string? azureAsyncOperation):
+                    _nextRequestUri = AppendOrReplaceApiVersion(azureAsyncOperation, _apiVersion);
+                    OperationId = ParseOperationId(_startRequestUri, _nextRequestUri);
+                    return;
+                case HeaderSource.Location when hasLocation:
+                    _nextRequestUri = AppendOrReplaceApiVersion(location!, _apiVersion);
+                    OperationId = ParseOperationId(_startRequestUri, _nextRequestUri);
+                    return;
+            }
+        }
+
+        internal static string AppendOrReplaceApiVersion(string uri, string? apiVersion)
+        {
+            if (!string.IsNullOrEmpty(apiVersion))
+            {
+                var uriSpan = uri.AsSpan();
+                var apiVersionParamSpan = ApiVersionParam.AsSpan();
+                var apiVersionIndex = uriSpan.IndexOf(apiVersionParamSpan);
+                if (apiVersionIndex == -1)
+                {
+                    var concatSymbol = uriSpan.IndexOf('?') > -1 ? "&" : "?";
+                    return $"{uri}{concatSymbol}api-version={apiVersion}";
+                }
+                else
+                {
+                    var lengthToEndOfApiVersionParam = apiVersionIndex + ApiVersionParam.Length;
+                    ReadOnlySpan<char> remaining = uriSpan.Slice(lengthToEndOfApiVersionParam);
+                    bool apiVersionHasEqualSign = false;
+                    if (remaining.IndexOf('=') == 0)
+                    {
+                        remaining = remaining.Slice(1);
+                        lengthToEndOfApiVersionParam += 1;
+                        apiVersionHasEqualSign = true;
+                    }
+                    var indexOfFirstSignAfterApiVersion = remaining.IndexOf('&');
+                    ReadOnlySpan<char> uriBeforeApiVersion = uriSpan.Slice(0, lengthToEndOfApiVersionParam);
+                    if (indexOfFirstSignAfterApiVersion == -1)
+                    {
+                        return string.Concat(uriBeforeApiVersion.ToString(), apiVersionHasEqualSign ? string.Empty : "=", apiVersion);
+                    }
+                    else
+                    {
+                        ReadOnlySpan<char> uriAfterApiVersion = uriSpan.Slice(indexOfFirstSignAfterApiVersion + lengthToEndOfApiVersionParam);
+                        return string.Concat(uriBeforeApiVersion.ToString(), apiVersionHasEqualSign ? string.Empty : "=", apiVersion, uriAfterApiVersion.ToString());
+                    }
+                }
+            }
+            return uri;
+        }
+
+        internal static bool TryGetApiVersion(Uri startRequestUri, out ReadOnlySpan<char> apiVersion)
+        {
+            apiVersion = default;
+            ReadOnlySpan<char> uriSpan = startRequestUri.Query.AsSpan();
+            int startIndex = uriSpan.IndexOf(ApiVersionParam.AsSpan());
+            if (startIndex == -1)
+            {
+                return false;
+            }
+            startIndex += ApiVersionParam.Length;
+            ReadOnlySpan<char> remaining = uriSpan.Slice(startIndex);
+            if (remaining.IndexOf('=') == 0)
+            {
+                remaining = remaining.Slice(1);
+                startIndex += 1;
+            }
+            else
+            {
+                return false;
+            }
+            int endIndex = remaining.IndexOf('&');
+            int length = endIndex == -1 ? uriSpan.Length - startIndex : endIndex;
+            apiVersion = uriSpan.Slice(startIndex, length);
+            return true;
+        }
+
+        /// <summary>
+        /// This function is used to get the final request uri after the lro has completed.
+        /// </summary>
+        private string? GetFinalUri(string? resourceLocation)
+        {
+            // Set final uri as null if the response for initial request doesn't contain header "Operation-Location" or "Azure-AsyncOperation".
+            if (_headerSource is not (HeaderSource.OperationLocation or HeaderSource.AzureAsyncOperation))
+            {
+                return null;
+            }
+
+            // Set final uri as null if initial request is a delete method.
+            if (RequestMethod == RequestMethod.Delete)
+            {
+                return null;
+            }
+
+            // Handle final-state-via options: https://github.com/Azure/autorest/blob/main/docs/extensions/readme.md#x-ms-long-running-operation-options
+            switch (_finalStateVia)
+            {
+                case OperationFinalStateVia.LocationOverride when !string.IsNullOrEmpty(_lastKnownLocation):
+                    return _lastKnownLocation;
+                case OperationFinalStateVia.OperationLocation or OperationFinalStateVia.AzureAsyncOperation when RequestMethod == RequestMethod.Post:
+                    return null;
+                case OperationFinalStateVia.OriginalUri:
+                    return _startRequestUri.AbsoluteUri;
+            }
+
+            // If response body contains resourceLocation, use it: https://github.com/microsoft/api-guidelines/blob/vNext/Guidelines.md#target-resource-location
+            if (resourceLocation != null)
+            {
+                return resourceLocation;
+            }
+
+            // If initial request is PUT or PATCH, return initial request Uri
+            if (RequestMethod == RequestMethod.Put || RequestMethod == RequestMethod.Patch)
+            {
+                return _startRequestUri.AbsoluteUri;
+            }
+
+            // If response for initial request contains header "Location", return last known location
+            if (!string.IsNullOrEmpty(_lastKnownLocation))
+            {
+                return _lastKnownLocation;
+            }
+
+            return null;
+        }
+
+        private Response GetResponse(string uri, CancellationToken cancellationToken)
+        {
+            using HttpMessage message = CreateRequest(uri);
+            _pipeline.Send(message, cancellationToken);
+
+            // If we are doing final get for a delete LRO with 404, just return empty response with 204
+            if (message.Response.Status == 404 && RequestMethod == RequestMethod.Delete)
+            {
+                return new EmptyResponse(HttpStatusCode.NoContent, message.Response.ClientRequestId);
+            }
+            return message.Response;
+        }
+
+        private async ValueTask<Response> GetResponseAsync(string uri, CancellationToken cancellationToken)
+        {
+            using HttpMessage message = CreateRequest(uri);
+            await _pipeline.SendAsync(message, cancellationToken).ConfigureAwait(false);
+
+            // If we are doing final get for a delete LRO with 404, just return empty response with 204
+            if (message.Response.Status == 404 && RequestMethod == RequestMethod.Delete)
+            {
+                return new EmptyResponse(HttpStatusCode.NoContent, message.Response.ClientRequestId);
+            }
+            return message.Response;
+        }
+
+        /// <summary>
+        /// This is only used for final get of the delete LRO, we just want to return an empty response with 204 to the user for this case.
+        /// </summary>
+        private sealed class EmptyResponse : Response
+        {
+            public EmptyResponse(HttpStatusCode status, string clientRequestId)
+            {
+                Status = (int)status;
+                ReasonPhrase = status.ToString();
+                ClientRequestId = clientRequestId;
+            }
+
+            public override int Status { get; }
+
+            public override string ReasonPhrase { get; }
+
+            public override Stream? ContentStream { get => null; set => throw new InvalidOperationException("Should not set ContentStream for an empty response."); }
+            public override string ClientRequestId { get; set; }
+
+            public override void Dispose()
+            {
+            }
+
+            /// <inheritdoc />
+#if HAS_INTERNALS_VISIBLE_CORE
+            internal
+#endif
+            protected override bool ContainsHeader(string name) => false;
+
+            /// <inheritdoc />
+#if HAS_INTERNALS_VISIBLE_CORE
+            internal
+#endif
+            protected override IEnumerable<HttpHeader> EnumerateHeaders() => Array.Empty<HttpHeader>();
+
+            /// <inheritdoc />
+#if HAS_INTERNALS_VISIBLE_CORE
+            internal
+#endif
+            protected override bool TryGetHeader(string name, out string value)
+            {
+                value = string.Empty;
+                return false;
+            }
+
+            /// <inheritdoc />
+#if HAS_INTERNALS_VISIBLE_CORE
+            internal
+#endif
+            protected override bool TryGetHeaderValues(string name, out IEnumerable<string> values)
+            {
+                values = Array.Empty<string>();
+                return false;
+            }
+        }
+
+        private HttpMessage CreateRequest(string uri)
+        {
+            HttpMessage message = _pipeline.CreateMessage();
+            Request request = message.Request;
+            request.Method = RequestMethod.Get;
+
+            if (Uri.TryCreate(uri, UriKind.Absolute, out var nextLink) && nextLink.Scheme != "file")
+            {
+                request.Uri.Reset(nextLink);
+            }
+            else
+            {
+                request.Uri.Reset(new Uri(_startRequestUri, uri));
+            }
+
+            return message;
+        }
+
+        private static bool IsFinalState(Response response, HeaderSource headerSource, out OperationState? failureState, out string? resourceLocation)
+        {
+            failureState = null;
+            resourceLocation = null;
+
+            if (headerSource == HeaderSource.Location)
+            {
+                return response.Status != 202;
+            }
+
+            if (response.Status is >= 200 and <= 204)
+            {
+                if (response.ContentStream is { Length: > 0 })
+                {
+                    try
+                    {
+                        using JsonDocument document = JsonDocument.Parse(response.ContentStream);
+                        var root = document.RootElement;
+                        switch (headerSource)
+                        {
+                            case HeaderSource.None when root.TryGetProperty("properties", out var properties) && properties.TryGetProperty("provisioningState", out JsonElement property):
+                            case HeaderSource.OperationLocation when root.TryGetProperty("status", out property):
+                            case HeaderSource.AzureAsyncOperation when root.TryGetProperty("status", out property):
+                                var state = GetRequiredString(property).ToLowerInvariant();
+                                if (FailureStates.Contains(state))
+                                {
+                                    failureState = OperationState.Failure(response);
+                                    return true;
+                                }
+                                else if (!SuccessStates.Contains(state))
+                                {
+                                    return false;
+                                }
+                                else
+                                {
+                                    if (headerSource is HeaderSource.OperationLocation or HeaderSource.AzureAsyncOperation && root.TryGetProperty("resourceLocation", out var resourceLocationProperty))
+                                    {
+                                        resourceLocation = resourceLocationProperty.GetString();
+                                    }
+                                    return true;
+                                }
+                        }
+                    }
+                    finally
+                    {
+                        // It is required to reset the position of the content after reading as this response may be used for deserialization.
+                        response.ContentStream.Position = 0;
+                    }
+                }
+
+                // If headerSource is None and provisioningState was not found, it defaults to Succeeded.
+                if (headerSource == HeaderSource.None)
+                {
+                    return true;
+                }
+            }
+
+            failureState = OperationState.Failure(response);
+            return true;
+        }
+
+        private static string GetRequiredString(in JsonElement element)
+        {
+            var value = element.GetString();
+            if (value == null)
+                throw new InvalidOperationException($"The requested operation requires an element of type 'String', but the target element has type '{element.ValueKind}'.");
+
+            return value;
+        }
+
+        private static bool ShouldIgnoreHeader(RequestMethod method, Response response)
+            => method.Method == RequestMethod.Patch.Method && response.Status == 200;
+
+        // Since this method is static, we can't manipulate the instance property OperationId of the class. We need to return isRequestPolling to update the OperationId after creaing the instance.
+        private static HeaderSource GetHeaderSource(RequestMethod requestMethod, Uri requestUri, Response response, string? apiVersion, out string nextRequestUri, out bool isNextRequestPolling)
+        {
+            isNextRequestPolling = false;
+            if (ShouldIgnoreHeader(requestMethod, response))
+            {
+                nextRequestUri = requestUri.AbsoluteUri;
+                return HeaderSource.None;
+            }
+
+            var headers = response.Headers;
+            if (headers.TryGetValue("Operation-Location", out var operationLocationUri))
+            {
+                nextRequestUri = AppendOrReplaceApiVersion(operationLocationUri, apiVersion);
+                isNextRequestPolling = true;
+                return HeaderSource.OperationLocation;
+            }
+
+            if (headers.TryGetValue("Azure-AsyncOperation", out var azureAsyncOperationUri))
+            {
+                nextRequestUri = AppendOrReplaceApiVersion(azureAsyncOperationUri, apiVersion);
+                isNextRequestPolling = true;
+                return HeaderSource.AzureAsyncOperation;
+            }
+
+            if (headers.TryGetValue("Location", out var locationUri))
+            {
+                nextRequestUri = AppendOrReplaceApiVersion(locationUri, apiVersion);
+                isNextRequestPolling = true;
+                return HeaderSource.Location;
+            }
+
+            nextRequestUri = requestUri.AbsoluteUri;
+            return HeaderSource.None;
+        }
+
+        private static void AssertNotNull<T>(T value, string name)
+        {
+            if (value is null)
+            {
+                throw new ArgumentNullException(name);
+            }
+        }
+
+        private enum HeaderSource
+        {
+            None,
+            OperationLocation,
+            AzureAsyncOperation,
+            Location
+        }
+
+        private class CompletedOperation : IOperation
+        {
+            private readonly OperationState _operationState;
+
+            private readonly NextLinkOperationImplementation _operation;
+
+            public CompletedOperation(OperationState operationState, NextLinkOperationImplementation operation)
+            {
+                _operationState = operationState;
+                _operation = operation;
+            }
+
+            public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => new(_operationState);
+
+            public RehydrationToken GetRehydrationToken() => _operation.GetRehydrationToken();
+        }
+
+        private sealed class OperationToOperationOfT<T> : IOperation<T>
+        {
+            private readonly IOperationSource<T> _operationSource;
+            private readonly IOperation _operation;
+
+            public OperationToOperationOfT(IOperationSource<T> operationSource, IOperation operation)
+            {
+                _operationSource = operationSource;
+                _operation = operation;
+            }
+
+            public async ValueTask<OperationState<T>> UpdateStateAsync(bool async, CancellationToken cancellationToken)
+            {
+                var state = await _operation.UpdateStateAsync(async, cancellationToken).ConfigureAwait(false);
+                if (state.HasSucceeded)
+                {
+                    var result = async
+                        ? await _operationSource.CreateResultAsync(state.RawResponse, cancellationToken).ConfigureAwait(false)
+                        : _operationSource.CreateResult(state.RawResponse, cancellationToken);
+
+                    return OperationState<T>.Success(state.RawResponse, result);
+                }
+
+                if (state.HasCompleted)
+                {
+                    return OperationState<T>.Failure(state.RawResponse, state.OperationFailedException);
+                }
+
+                return OperationState<T>.Pending(state.RawResponse);
+            }
+
+            public RehydrationToken GetRehydrationToken() => _operation.GetRehydrationToken();
+        }
+    }
+}