[core-lro] Support path rewriting

sdk/core/core-lro/CHANGELOG.md

@@ -1,9 +1,11 @@
 # Release History
 
-## 3.1.1 (Unreleased)
+## 3.2.0 (Unreleased)
 
 ### Features Added
 
+- Supports a `baseUrl` option that can be used to rewrite the polling URL to use that base URL instead. This makes sure that polling works for operations that live behind proxies or API gateways.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/core/core-lro/package.json

@@ -2,7 +2,7 @@
   "name": "@azure/core-lro",
   "author": "Microsoft Corporation",
   "sdk-type": "client",
-  "version": "3.1.1",
+  "version": "3.2.0",
   "type": "module",
   "description": "Isomorphic client library for supporting long-running operations in node.js and browser.",
   "exports": {
@@ -101,9 +101,11 @@
     "@azure/core-util": "^1.11.0",
     "@azure/dev-tool": "^1.0.0",
     "@azure/eslint-plugin-azure-sdk": "^3.0.0",
+    "@types/chai-as-promised": "^8.0.2",
     "@types/node": "^18.0.0",
     "@vitest/browser": "^3.0.3",
     "@vitest/coverage-istanbul": "^3.0.3",
+    "chai-as-promised": "^8.0.1",
     "eslint": "^9.9.0",
     "playwright": "^1.41.2",
     "typescript": "~5.7.2",

sdk/core/core-lro/review/core-lro.api.md

@@ -14,6 +14,7 @@ export function createHttpPoller<TResult, TState extends OperationState<TResult>
 
 // @public
 export interface CreateHttpPollerOptions<TResult, TState> {
+    baseUrl?: string;
     intervalInMs?: number;
     processResult?: (result: unknown, state: TState) => Promise<TResult>;
     resolveOnUnsuccessful?: boolean;

sdk/core/core-lro/src/http/models.ts

@@ -103,6 +103,10 @@ export interface CreateHttpPollerOptions<TResult, TState> {
    * The potential location of the result of the LRO if specified by the LRO extension in the swagger.
    */
   resourceLocationConfig?: ResourceLocationConfig;
+  /**
+   * The base URL to use when making requests.
+   */
+  baseUrl?: string;
   /**
    * A function to process the result of the LRO.
    */

sdk/core/core-lro/src/http/operation.ts

@@ -62,7 +62,7 @@ function findResourceLocation(inputs: {
     }
   }
 
-  function getDefault() {
+  function getDefault(): string | undefined {
     switch (resourceLocationConfig) {
       case "operation-location":
       case "azure-async-operation": {

sdk/core/core-lro/src/http/poller.ts

@@ -15,6 +15,7 @@ import {
 } from "./operation.js";
 import type { CreateHttpPollerOptions } from "./models.js";
 import { buildCreatePoller } from "../poller/poller.js";
+import { rewriteUrl } from "./utils.js";
 
 /**
  * Creates a poller that can be used to poll a long-running operation.
@@ -34,6 +35,7 @@ export function createHttpPoller<TResult, TState extends OperationState<TResult>
     updateState,
     withOperationLocation,
     resolveOnUnsuccessful = false,
+    baseUrl,
   } = options || {};
   return buildCreatePoller<OperationResponse, TResult, TState>({
     getStatusFromInitialResponse,
@@ -51,8 +53,8 @@ export function createHttpPoller<TResult, TState extends OperationState<TResult>
         const config = inferLroMode(response.rawResponse, resourceLocationConfig);
         return {
           response,
-          operationLocation: config?.operationLocation,
-          resourceLocation: config?.resourceLocation,
+          operationLocation: rewriteUrl({ url: config?.operationLocation, baseUrl }),
+          resourceLocation: rewriteUrl({ url: config?.resourceLocation, baseUrl }),
           initialRequestUrl: config?.initialRequestUrl,
           requestMethod: config?.requestMethod,
           ...(config?.mode ? { metadata: { mode: config.mode } } : {}),

sdk/core/core-lro/src/http/utils.ts

@@ -0,0 +1,55 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+/**
+ * Rewrites a given URL to use the specified base URL.
+ * It preserves the pathname, search parameters, and fragment.
+ * Handles relative URLs and proper encoding.
+ *
+ * @param params - An object containing the url and baseUrl.
+ *                    url - The original URL (absolute or relative).
+ *                    baseUrl - The new base URL to use.
+ * @returns The rewritten URL as a string.
+ */
+export function rewriteUrl({
+  url,
+  baseUrl,
+}: {
+  url?: string;
+  baseUrl?: string;
+}): string | undefined {
+  if (!url) {
+    return undefined;
+  }
+  if (!baseUrl) {
+    return url;
+  }
+
+  let originalUrl: URL;
+
+  try {
+    // Try to parse inputUrl as an absolute URL.
+    originalUrl = new URL(url);
+  } catch {
+    // If inputUrl is relative, resolve using the provided baseUrl.
+    try {
+      originalUrl = new URL(url, baseUrl);
+    } catch (e) {
+      throw new Error(`Invalid input URL provided: ${url}`);
+    }
+  }
+
+  let newBase: URL;
+  try {
+    newBase = new URL(baseUrl);
+  } catch (e) {
+    throw new Error(`Invalid base URL provided: ${baseUrl}`);
+  }
+
+  const rewrittenUrl = new URL(
+    `${originalUrl.pathname}${originalUrl.search}${originalUrl.hash}`,
+    newBase,
+  );
+
+  return rewrittenUrl.toString();
+}

sdk/core/core-lro/test/lro.spec.ts

@@ -3,11 +3,14 @@
 
 import type { ImplementationName, Result } from "./utils/utils.js";
 import { assertDivergentBehavior, assertError, createDoubleHeaders } from "./utils/utils.js";
-import { describe, it, assert, expect } from "vitest";
+import { describe, it, assert, expect, chai } from "vitest";
+import chaiAsPromised from "chai-as-promised";
 import { createRunLroWith, createTestPoller } from "./utils/router.js";
 import { delay } from "@azure/core-util";
 import { matrix } from "@azure-tools/test-utils-vitest";
 
+chai.use(chaiAsPromised);
+
 matrix(
   [["createPoller"], [true, false]] as const,
   async function (implName: ImplementationName, throwOnNon2xxResponse: boolean) {
@@ -2345,7 +2348,7 @@ matrix(
             implName,
             throwOnNon2xxResponse,
           });
-          assert.isUndefined(poller.operationState);
+          assert.isUndefined(poller.operationState as any);
           const serialized = await poller.serialize();
           const expectedSerialized = JSON.stringify({
             state: {
@@ -2369,7 +2372,7 @@ matrix(
             throwOnNon2xxResponse,
           });
           assert.equal(serialized, await restoredPoller.serialize());
-          assert.equal(poller.operationState!.status, "succeeded");
+          assert.equal(poller.operationState.status, "succeeded");
           assert.deepEqual(poller.result, retResult);
         });
 
@@ -2445,7 +2448,7 @@ matrix(
           assert.equal(pollCount, 0);
           assert.deepEqual(retResult, await restoredPoller);
           assert.equal(pollCount, 11);
-          assert.equal(restoredPoller.operationState.status, "succeeded");
+          assert.equal(restoredPoller.operationState?.status, "succeeded");
           assert.deepEqual(restoredPoller.result, retResult);
           assert.isUndefined(poller.result);
           // duplicate awaitting would not trigger extra pollings
@@ -2950,6 +2953,43 @@ matrix(
           });
           assert.equal(poller.result?.properties?.provisioningState, "Canceled");
         });
+        it("polling URL is rewritten if the baseUrl option is provided ", async function () {
+          const baseUrl = "https://example2.com";
+          const initialRequestUrl = `${baseUrl}/path`;
+          const poller = createTestPoller({
+            routes: [
+              {
+                method: "POST",
+                status: 200,
+                body: JSON.stringify({ status: "Running" }),
+                headers: {
+                  "operation-location": "https://example1.com/path/polling",
+                  location: "https://example1.com/location",
+                },
+                path: initialRequestUrl,
+              },
+              {
+                method: "GET",
+                status: 200,
+                body: JSON.stringify({ status: "Succeeded" }),
+                path: `${baseUrl}/path/polling`,
+              },
+              {
+                method: "GET",
+                status: 200,
+                body: "",
+                path: `${baseUrl}/location`,
+              },
+            ],
+            baseUrl,
+            throwOnNon2xxResponse,
+          });
+          const pollerState = JSON.parse(await poller.serialize());
+          assert.equal(pollerState.state.config.operationLocation, `${baseUrl}/path/polling`);
+          assert.equal(pollerState.state.config.resourceLocation, `${baseUrl}/location`);
+          assert.equal(pollerState.state.config.initialRequestUrl, initialRequestUrl);
+          await assert.isFulfilled(poller.pollUntilDone());
+        });
         it("prints an error message based on the error in the status monitor", async () => {
           const pollingPath = "/postlocation/retry/succeeded/operationResults/200/";
           const code = "InvalidRequest";
@@ -3172,7 +3212,7 @@ matrix(
               implName,
               throwOnNon2xxResponse: true,
             });
-            await expect(poller).rejects.toThrow(errMsg);
+            await expect(poller).to.be.rejectedWith(errMsg);
           });
           it("should work properly when mixing catch and await", async () => {
             const body = { status: "canceled", results: [1, 2] };
@@ -3207,9 +3247,9 @@ matrix(
               err = e;
             });
             assert.equal(err.message, errMsg);
-            await expect(poller).rejects.toThrow(errMsg);
+            await expect(poller).to.be.rejectedWith(errMsg);
             assert.equal(callbackCounts, 1);
-            await expect(poller.finally(() => callbackCounts++)).rejects.toThrow(errMsg);
+            await expect(poller.finally(() => callbackCounts++)).to.be.rejectedWith(errMsg);
             assert.equal(callbackCounts, 2);
           });
         });