From bd935ff5655b117a4ca996e5e7b66a78ea379834 Mon Sep 17 00:00:00 2001
From: Jacob Viau <javia@microsoft.com>
Date: Thu, 11 Jan 2024 13:19:41 -0800
Subject: [PATCH] Add CreateCheckStatusResponseAsync APIs (#2722)

* Add CreateCheckStatusResponseAsync APIs

* Undo changes to DirectInput.cs

* update release_notes.md
---
 release_notes.md                              |  2 +
 .../DurableTaskClientExtensions.cs            | 66 +++++++++++++++++--
 2 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/release_notes.md b/release_notes.md
index 730f39293..28441f566 100644
--- a/release_notes.md
+++ b/release_notes.md
@@ -4,6 +4,8 @@
 
 ### New Features
 
+- Add `CreateCheckStatusResponseAsync` APIs. (https://github.com/Azure/azure-functions-durable-extension/pull/2722)
+
 ### Bug Fixes
 
 - Fix issue with isolated entities: custom deserialization was not working because IServices was not passed along (https://github.com/Azure/azure-functions-durable-extension/pull/2686)
diff --git a/src/Worker.Extensions.DurableTask/DurableTaskClientExtensions.cs b/src/Worker.Extensions.DurableTask/DurableTaskClientExtensions.cs
index ec9311528..1a487f388 100644
--- a/src/Worker.Extensions.DurableTask/DurableTaskClientExtensions.cs
+++ b/src/Worker.Extensions.DurableTask/DurableTaskClientExtensions.cs
@@ -4,6 +4,7 @@
 using System;
 using System.Net;
 using System.Threading;
+using System.Threading.Tasks;
 using Azure.Core.Serialization;
 using Microsoft.Azure.Functions.Worker.Http;
 using Microsoft.DurableTask.Client;
@@ -17,6 +18,52 @@ namespace Microsoft.Azure.Functions.Worker;
 /// </summary>
 public static class DurableTaskClientExtensions
 {
+    /// <summary>
+    /// Creates an HTTP response that is useful for checking the status of the specified instance.
+    /// </summary>
+    /// <param name="client">The <see cref="DurableTaskClient"/>.</param>
+    /// <param name="request">The HTTP request that this response is for.</param>
+    /// <param name="instanceId">The ID of the orchestration instance to check.</param>
+    /// <param name="cancellation">The cancellation token.</param>
+    /// <returns>An HTTP 202 response with a Location header and a payload containing instance control URLs.</returns>
+    public static Task<HttpResponseData> CreateCheckStatusResponseAsync(
+        this DurableTaskClient client,
+        HttpRequestData request,
+        string instanceId,
+        CancellationToken cancellation = default)
+    {
+        return client.CreateCheckStatusResponseAsync(request, instanceId, HttpStatusCode.Accepted, cancellation);
+    }
+
+    /// <summary>
+    /// Creates an HTTP response that is useful for checking the status of the specified instance.
+    /// </summary>
+    /// <param name="client">The <see cref="DurableTaskClient"/>.</param>
+    /// <param name="request">The HTTP request that this response is for.</param>
+    /// <param name="instanceId">The ID of the orchestration instance to check.</param>
+    /// <param name="statusCode">The status code.</param>
+    /// <param name="cancellation">The cancellation token.</param>
+    /// <returns>An HTTP response with a Location header and a payload containing instance control URLs.</returns>
+    public static async Task<HttpResponseData> CreateCheckStatusResponseAsync(
+        this DurableTaskClient client,
+        HttpRequestData request,
+        string instanceId,
+        HttpStatusCode statusCode,
+        CancellationToken cancellation = default)
+    {
+        if (client is null)
+        {
+            throw new ArgumentNullException(nameof(client));
+        }
+
+        HttpResponseData response = request.CreateResponse(statusCode);
+        object payload = SetHeadersAndGetPayload(client, request, response, instanceId);
+
+        ObjectSerializer serializer = GetObjectSerializer(response);
+        await serializer.SerializeAsync(response.Body, payload, payload.GetType(), cancellation);
+        return response;
+    }
+
     /// <summary>
     /// Creates an HTTP response that is useful for checking the status of the specified instance.
     /// </summary>
@@ -55,6 +102,17 @@ public static HttpResponseData CreateCheckStatusResponse(
             throw new ArgumentNullException(nameof(client));
         }
 
+        HttpResponseData response = request.CreateResponse(statusCode);
+        object payload = SetHeadersAndGetPayload(client, request, response, instanceId);
+
+        ObjectSerializer serializer = GetObjectSerializer(response);
+        serializer.Serialize(response.Body, payload, payload.GetType(), cancellation);
+        return response;
+    }
+
+    private static object SetHeadersAndGetPayload(
+        DurableTaskClient client, HttpRequestData request, HttpResponseData response, string instanceId)
+    {
         static string BuildUrl(string url, params string?[] queryValues)
         {
             bool appended = false;
@@ -79,13 +137,10 @@ static string BuildUrl(string url, params string?[] queryValues)
         string formattedInstanceId = Uri.EscapeDataString(instanceId);
         string instanceUrl = $"{baseUrl}/runtime/webhooks/durabletask/instances/{formattedInstanceId}";
         string? commonQueryParameters = GetQueryParams(client);
-
-        HttpResponseData response = request.CreateResponse(statusCode);
         response.Headers.Add("Location", BuildUrl(instanceUrl, commonQueryParameters));
         response.Headers.Add("Content-Type", "application/json");
 
-        ObjectSerializer serializer = GetObjectSerializer(response);
-        var payload = new
+        return new
         {
             id = instanceId,
             purgeHistoryDeleteUri = BuildUrl(instanceUrl, commonQueryParameters),
@@ -93,9 +148,6 @@ static string BuildUrl(string url, params string?[] queryValues)
             statusQueryGetUri = BuildUrl(instanceUrl, commonQueryParameters),
             terminatePostUri = BuildUrl($"{instanceUrl}/terminate", "reason={{text}}}", commonQueryParameters),
         };
-
-        serializer.Serialize(response.Body, payload, payload.GetType(), cancellation);
-        return response;
     }
 
     private static ObjectSerializer GetObjectSerializer(HttpResponseData response)