From 2e5355d38d5d719e08d3865ac30eee300b9dc08f Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 16:18:07 -0700
Subject: [PATCH 01/10] Refresh Fabric auth during SynapseML request retries

AB#3582121

## Summary
Add Fabric authentication provenance to implicit OpenAI requests, reconstruct requests for retry-safe replay, refresh authorization once after a trusted 401, and reacquire current authorization across 429 retries. Add bounded endpoint validation, runtime cache invalidation compatibility, and targeted Scala coverage.

## Prompting Intent
The engineer asked to close the Fabric authentication refresh gap for long-running PySpark AI Functions calls, preserve request behavior across retries, review the implementation, and prepare it for an upstream pull request.

## Linked Sources
- Work item: https://dev.azure.com/msdata/A365/_workitems/edit/3582121
- Requirements: engineer request captured in Copilot session 98752775-8e28-43eb-9609-2249dd586f80
- Related source change: companion SynapseML-Internal pandas pull request, linked from the PR description

## Rationale
Use internal provenance rather than inferring authentication from user headers, strip that marker before transmission, and refresh only requests constrained to the trusted Fabric OpenAI endpoint. Reconstruct requests from buffered bodies and original headers so a single 401 replay and existing 429 retries preserve payloads and caller intent without extending Fabric refresh behavior to custom endpoints.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../ml/services/CognitiveServiceBase.scala    |  77 +++++--
 .../synapse/ml/services/openai/OpenAI.scala   |  10 +
 .../openai/OpenAIFabricHeadersSuite.scala     |  51 +++++
 .../synapse/ml/fabric/FabricClient.scala      |  56 ++++-
 .../synapse/ml/fabric/TokenLibrary.scala      |  99 +++++++-
 .../synapse/ml/io/http/HTTPClients.scala      | 148 +++++++++++-
 .../azure/synapse/ml/io/http/HTTPSchema.scala |  13 +-
 .../ml/fabric/VerifyTokenInvalidation.scala   |  42 ++++
 .../ml/io/split1/VerifySendWithRetries.scala  | 211 +++++++++++++++++-
 9 files changed, 682 insertions(+), 25 deletions(-)
 create mode 100644 core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala

diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala
index 2b8f90549b0..37fe70c56ab 100644
--- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala
+++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala
@@ -329,6 +329,8 @@ object URLEncodingUtils {
 }
 
 private[ml] object ServiceAuthHeaders {
+  private[ml] case class Resolution(headers: Map[String, String], usesFabricFallback: Boolean)
+
   private[ml] def nonBlank(value: String): Boolean = value != null && value.trim.nonEmpty
 
   // Normalize a header map from any caller (a writer-supplied java.util.HashMap included) into a
@@ -360,15 +362,15 @@ private[ml] object ServiceAuthHeaders {
       .headOption
       .map { case (_, value) => canonicalName -> value }
 
-  def build(subscriptionKey: Option[String],
-            subscriptionKeyHeaderName: String,
-            aadHeaderName: String,
-            aadToken: Option[String],
-            customAuthHeader: Option[String],
-            customHeaders: Option[Map[String, String]],
-            fabricFallbackAuthHeader: => Option[String],
-            telemHeaders: Option[Map[String, String]],
-            contentType: Option[String]): Map[String, String] = {
+  def resolve(subscriptionKey: Option[String],
+              subscriptionKeyHeaderName: String,
+              aadHeaderName: String,
+              aadToken: Option[String],
+              customAuthHeader: Option[String],
+              customHeaders: Option[Map[String, String]],
+              fabricFallbackAuthHeader: => Option[String],
+              telemHeaders: Option[Map[String, String]],
+              contentType: Option[String]): Resolution = {
     val providedCustomHeaders = sanitizeHeaderMap(customHeaders)
 
     // Header names that carry credentials in this context, compared case-insensitively.
@@ -383,13 +385,18 @@ private[ml] object ServiceAuthHeaders {
     // higher-priority source above is absent. A fallback that acquires a token (and may throw) is
     // therefore never run while a subscription key, AAD token, explicit custom-auth header, or an
     // embedded customHeaders credential is present.
-    val authHeader: Option[(String, String)] = subscriptionKey.filter(nonBlank)
+    val explicitAuthHeader: Option[(String, String)] = subscriptionKey.filter(nonBlank)
       .map(value => subscriptionKeyHeaderName -> value)
       .orElse(aadToken.filter(nonBlank).map(value => aadHeaderName -> ("Bearer " + value)))
       .orElse(customAuthHeader.filter(nonBlank).map(value => aadHeaderName -> value))
       .orElse(embeddedCredential(providedCustomHeaders, subscriptionKeyHeaderName))
       .orElse(embeddedCredential(providedCustomHeaders, aadHeaderName))
-      .orElse(fabricFallbackAuthHeader.filter(nonBlank).map(value => aadHeaderName -> value))
+    val fabricAuthHeader = if (explicitAuthHeader.isDefined) {
+      None
+    } else {
+      fabricFallbackAuthHeader.filter(nonBlank).map(value => aadHeaderName -> value)
+    }
+    val authHeader = explicitAuthHeader.orElse(fabricAuthHeader)
 
     // Generic headers never carry auth: strip api-key/Authorization entries (any casing) so they
     // can neither override the resolved credential nor duplicate it under a different case.
@@ -412,7 +419,30 @@ private[ml] object ServiceAuthHeaders {
     }
     contentType.filterNot(StringUtils.isEmpty).foreach(value => headers += ("Content-Type" -> value))
 
-    new scala.collection.immutable.TreeMap[String, String]() ++ headers
+    Resolution(
+      new scala.collection.immutable.TreeMap[String, String]() ++ headers,
+      usesFabricFallback = fabricAuthHeader.isDefined)
+  }
+
+  def build(subscriptionKey: Option[String],
+            subscriptionKeyHeaderName: String,
+            aadHeaderName: String,
+            aadToken: Option[String],
+            customAuthHeader: Option[String],
+            customHeaders: Option[Map[String, String]],
+            fabricFallbackAuthHeader: => Option[String],
+            telemHeaders: Option[Map[String, String]],
+            contentType: Option[String]): Map[String, String] = {
+    resolve(
+      subscriptionKey,
+      subscriptionKeyHeaderName,
+      aadHeaderName,
+      aadToken,
+      customAuthHeader,
+      customHeaders,
+      fabricFallbackAuthHeader,
+      telemHeaders,
+      contentType).headers
   }
 }
 
@@ -503,12 +533,22 @@ trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAA
     getValueOpt(row, customHeaders)
   }
 
+  protected def supportsImplicitFabricAuthRetry: Boolean = false
+
   protected def addHeaders(req: HttpRequestBase,
                            row: Row,
                            addContentType: Boolean = true): Unit = {
 
-    val headers = getHeaders(row, addContentType)
-    headers.foreach { case (headerName, headerValue) => req.addHeader(headerName, headerValue) }
+    val resolution = resolveServiceAuthHeaders(row, addContentType, getFabricFallbackAuthHeader(row))
+    req.removeHeaders(HTTPRequestData.FabricAuthMarkerHeader)
+    resolution.headers.foreach { case (headerName, headerValue) =>
+      if (!HTTPRequestData.isFabricAuthMarker(headerName)) {
+        req.addHeader(headerName, headerValue)
+      }
+    }
+    if (resolution.usesFabricFallback && supportsImplicitFabricAuthRetry) {
+      req.addHeader(HTTPRequestData.FabricAuthMarkerHeader, "true")
+    }
   }
 
   // Returns a list of key-value pairs representing the headers
@@ -523,7 +563,14 @@ trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAA
   private[ml] def buildServiceAuthHeaders(row: Row,
                                           addContentType: Boolean,
                                           fabricFallbackAuthHeader: => Option[String]): Map[String, String] = {
-    ServiceAuthHeaders.build(
+    resolveServiceAuthHeaders(row, addContentType, fabricFallbackAuthHeader).headers
+  }
+
+  private def resolveServiceAuthHeaders(
+      row: Row,
+      addContentType: Boolean,
+      fabricFallbackAuthHeader: => Option[String]): ServiceAuthHeaders.Resolution = {
+    ServiceAuthHeaders.resolve(
       getValueOpt(row, subscriptionKey),
       subscriptionKeyHeaderName,
       aadHeaderName,
diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala
index 8549412f88a..1b1bb55e902 100644
--- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala
+++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala
@@ -577,6 +577,16 @@ trait HasOpenAIFabricHeaders extends HasCognitiveServiceInput {
     runningOnFabric && usingDefaultOpenAIEndpoint && !hasCustomUrlRoot
   }
 
+  override protected def supportsImplicitFabricAuthRetry: Boolean = usingImplicitFabricEndpoint
+
+  abstract override protected def getFabricFallbackAuthHeader(row: Row): Option[String] = {
+    if (usingImplicitFabricEndpoint) {
+      super.getFabricFallbackAuthHeader(row)
+    } else {
+      None
+    }
+  }
+
   abstract override protected def getCustomHeaders(row: Row): Option[Map[String, String]] = {
     val headers = super.getCustomHeaders(row)
     if (usingImplicitFabricEndpoint) {
diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala
index 35f60b2c2dc..8cd94c1c8c2 100644
--- a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala
+++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala
@@ -4,8 +4,10 @@
 package com.microsoft.azure.synapse.ml.services.openai
 
 import com.microsoft.azure.synapse.ml.core.test.base.TestBase
+import com.microsoft.azure.synapse.ml.io.http.HTTPRequestData
 import com.microsoft.azure.synapse.ml.logging.common.PlatformDetails
 import com.microsoft.azure.synapse.ml.services.HasCognitiveServiceInput
+import org.apache.http.client.methods.HttpPost
 import org.apache.spark.sql.Row
 import spray.json._
 
@@ -58,6 +60,25 @@ class OpenAIFabricHeadersSuite extends TestBase {
       override protected val usesDefaultEndpoint: Boolean)
     extends OpenAIResponses with InspectableFabricHeaders
 
+  private class AuthProvenanceProbe(fallbackAuthHeader: Option[String]) extends OpenAIChatCompletion {
+    var fallbackCalls = 0
+
+    override protected[openai] def runningOnFabric: Boolean = true
+
+    override protected[openai] def usingDefaultOpenAIEndpoint: Boolean = true
+
+    override protected def getFabricFallbackAuthHeader(row: Row): Option[String] = {
+      fallbackCalls += 1
+      fallbackAuthHeader
+    }
+
+    def requestData: HTTPRequestData = {
+      val request = new HttpPost("https://example.test/openai")
+      addHeaders(request, Row.empty, addContentType = false)
+      new HTTPRequestData(request)
+    }
+  }
+
   private def transformers(
       isFabric: Boolean,
       usesDefaultEndpoint: Boolean): Seq[InspectableFabricHeaders] = {
@@ -160,4 +181,34 @@ class OpenAIFabricHeadersSuite extends TestBase {
     assert(headers.values.forall(_ != null))
     assert(!headers.contains("Other"))
   }
+
+  test("only implicit Fabric authentication marks a request for refresh") {
+    val implicitProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+    val implicitRequestData = implicitProbe.requestData
+    val implicitHttpRequest = implicitRequestData.toHTTPCore
+
+    assert(implicitProbe.fallbackCalls === 1)
+    assert(implicitRequestData.usesFabricAuth)
+    assert(implicitHttpRequest.getFirstHeader("Authorization").getValue === "MwcToken implicit")
+    assert(Option(implicitHttpRequest.getFirstHeader(HTTPRequestData.FabricAuthMarkerHeader)).isEmpty)
+
+    val explicitProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+      .setCustomAuthHeader("MwcToken explicit")
+      .setCustomHeaders(Map(HTTPRequestData.FabricAuthMarkerHeader -> "true"))
+      .asInstanceOf[AuthProvenanceProbe]
+    val explicitRequestData = explicitProbe.requestData
+
+    assert(explicitProbe.fallbackCalls === 0)
+    assert(!explicitRequestData.usesFabricAuth)
+    assert(explicitRequestData.toHTTPCore.getFirstHeader("Authorization").getValue === "MwcToken explicit")
+
+    val apiKeyProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+      .setSubscriptionKey("explicit-key")
+      .asInstanceOf[AuthProvenanceProbe]
+    val apiKeyRequestData = apiKeyProbe.requestData
+
+    assert(apiKeyProbe.fallbackCalls === 0)
+    assert(!apiKeyRequestData.usesFabricAuth)
+    assert(apiKeyRequestData.toHTTPCore.getFirstHeader("api-key").getValue === "explicit-key")
+  }
 }
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala
index 4a6d92b02ac..c68582d455e 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala
@@ -6,9 +6,11 @@ package com.microsoft.azure.synapse.ml.fabric
 import spray.json.DefaultJsonProtocol.{StringJsonFormat, mapFormat}
 import spray.json._
 
-import java.net.{MalformedURLException, URL}
+import java.net.{MalformedURLException, URI, URL}
 import java.util.UUID
+import scala.collection.concurrent.TrieMap
 import scala.io.Source
+import scala.util.control.NonFatal
 
 object FabricClient extends RESTUtils {
   private val WorkloadEndpointTypeML = "ML";
@@ -19,6 +21,7 @@ object FabricClient extends RESTUtils {
   private val ContextFilePath = "/home/trusted-service-user/.trident-context";
   private val SparkConfPath = "/opt/spark/conf/spark-defaults.conf";
   private val ClusterInfoPath = "/opt/health-agent/conf/cluster-info.json";
+  private val CognitiveMwcRefreshLocks = TrieMap.empty[(String, String), AnyRef]
 
   lazy val CapacityID: Option[String] = getCapacityID;
   lazy val WorkspaceID: Option[String] = getWorkspaceID;
@@ -195,4 +198,55 @@ object FabricClient extends RESTUtils {
   def getCognitiveMWCTokenAuthHeader: String = {
     TokenLibrary.getCognitiveMwcTokenAuthHeader(WorkspaceID.getOrElse(""), ArtifactID.getOrElse(""))
   }
+
+  private[ml] def isEndpointUnder(requestUrl: String, endpointRoot: String): Boolean = {
+    try {
+      val request = new URL(requestUrl)
+      val root = new URL(endpointRoot)
+      val requestPort = if (request.getPort >= 0) request.getPort else request.getDefaultPort
+      val rootPort = if (root.getPort >= 0) root.getPort else root.getDefaultPort
+      val requestPath = trustedPath(new URI(requestUrl).getRawPath)
+      val rootPath = trustedPath(new URI(endpointRoot).getRawPath)
+      requestPath.exists(path =>
+        rootPath.exists(rootPrefix =>
+          request.getProtocol.equalsIgnoreCase("https") &&
+            root.getProtocol.equalsIgnoreCase("https") &&
+            request.getHost.equalsIgnoreCase(root.getHost) &&
+            requestPort == rootPort &&
+            path.startsWith(rootPrefix)))
+    } catch {
+      case _: MalformedURLException | _: java.net.URISyntaxException => false
+    }
+  }
+
+  private def trustedPath(rawPath: String): Option[String] = {
+    val lowerPath = rawPath.toLowerCase
+    val hasAmbiguousEncoding = Seq("%2e", "%2f", "%5c").exists(lowerPath.contains)
+    val normalizedPath = rawPath.replaceAll("/+", "/")
+    val hasDotSegment = normalizedPath.split("/", -1).exists(segment => segment == "." || segment == "..")
+    if (hasAmbiguousEncoding || rawPath.contains("\\") || hasDotSegment) None else Some(normalizedPath)
+  }
+
+  private[ml] def isOpenAIEndpoint(requestUrl: String): Boolean = {
+    try {
+      isEndpointUnder(requestUrl, MLWorkloadEndpointOpenAI)
+    } catch {
+      case NonFatal(_) => false
+    }
+  }
+
+  def refreshCognitiveMWCTokenAuthHeader(rejectedAuthHeader: String): String = {
+    val workspaceId = WorkspaceID.getOrElse("")
+    val artifactId = ArtifactID.getOrElse("")
+    val refreshLock = CognitiveMwcRefreshLocks.getOrElseUpdate((workspaceId, artifactId), new Object())
+    refreshLock.synchronized {
+      val currentAuthHeader = TokenLibrary.getCognitiveMwcTokenAuthHeader(workspaceId, artifactId)
+      if (currentAuthHeader != rejectedAuthHeader) {
+        currentAuthHeader
+      } else {
+        TokenLibrary.invalidateSparkMwcToken(workspaceId, artifactId)
+        TokenLibrary.getCognitiveMwcTokenAuthHeader(workspaceId, artifactId)
+      }
+    }
+  }
 }
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
index 1ced3b3d7a8..a036d0eb8b8 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
@@ -3,10 +3,26 @@
 
 package com.microsoft.azure.synapse.ml.fabric
 
+import java.io.File
+import java.nio.file.{Files, Path, Paths}
 import scala.reflect.runtime.currentMirror
 import scala.reflect.runtime.universe._
 
 object TokenLibrary {
+  private val TokenLibraryClass = "com.microsoft.azure.trident.tokenlibrary.TokenLibrary$"
+  private val InMemoryCacheClasses = Seq(
+    "com.microsoft.azure.trident.tokenlibrary.InMemoryCacheClient$",
+    "com.microsoft.azure.trident.tokenlibrary.cache.InMemoryCacheClient$",
+    "com.microsoft.fabric.tokenlibrary.InMemoryCacheClient$",
+    "com.microsoft.fabric.tokenlibrary.cache.InMemoryCacheClient$")
+  private val NfsCacheClasses = Seq(
+    "com.microsoft.azure.trident.tokenlibrary.NFSCache$",
+    "com.microsoft.azure.trident.tokenlibrary.cache.NFSCache$",
+    "com.microsoft.fabric.tokenlibrary.NFSCache$",
+    "com.microsoft.fabric.tokenlibrary.cache.NFSCache$")
+  private val SparkTokenVersion = 2
+  private val SparkWorkloadType = "SparkCore"
+
   def getAccessToken: String = {
     val objectName = "com.microsoft.azure.trident.tokenlibrary.TokenLibrary"
     val mirror = currentMirror
@@ -39,10 +55,91 @@ object TokenLibrary {
       m.asMethod.paramLists.flatten.map(_.typeSignature).zip(argTypes).forall { case (a, b) => a =:= b }
     }.getOrElse(throw new NoSuchMethodException(s"Method $methodName with argument type not found"))
     val methodMirror = mirror.reflect(obj).reflectMethod(selectedMethodSymbol.asMethod)
-    methodMirror(workspaceId, artifactId, 2, "SparkCore")
+    methodMirror(workspaceId, artifactId, SparkTokenVersion, SparkWorkloadType)
       .asInstanceOf[String]
   }
 
+  private def objectMethod(classNames: Seq[String],
+                           methodName: String,
+                           parameterCount: Int): Option[(AnyRef, java.lang.reflect.Method)] = {
+    classNames.iterator.flatMap { className =>
+      try {
+        val cls = Class.forName(className)
+        val module = cls.getField("MODULE$").get(null).asInstanceOf[AnyRef] //scalastyle:ignore null
+        (cls.getMethods ++ cls.getDeclaredMethods)
+          .find(method => method.getName == methodName && method.getParameterCount == parameterCount)
+          .map { method =>
+            method.setAccessible(true)
+            module -> method
+          }
+      } catch {
+        case _: ClassNotFoundException | _: NoSuchFieldException => None
+      }
+    }.toSeq.headOption
+  }
+
+  private def invalidateWithRuntimeApi(workspaceId: String, artifactId: String): Boolean = {
+    objectMethod(Seq(TokenLibraryClass), "invalidateMwcToken", 4).exists { case (module, method) =>
+      method.invoke(
+        module,
+        workspaceId,
+        artifactId,
+        Int.box(SparkTokenVersion),
+        SparkWorkloadType)
+      true
+    }
+  }
+
+  private def nfsCacheKey(cacheKey: String): String = {
+    objectMethod(NfsCacheClasses, "getNFSCacheKey", 1)
+      .map { case (module, method) => method.invoke(module, cacheKey).asInstanceOf[String] }
+      .getOrElse(cacheKey)
+  }
+
+  private def deleteNfsToken(resolvedCacheKey: String): Boolean = {
+    objectMethod(NfsCacheClasses, "getNFSTokenFilePath", 1).exists { case (module, method) =>
+      val tokenPath = method.invoke(module, resolvedCacheKey) match {
+        case path: Path => path
+        case file: File => file.toPath
+        case path: String => Paths.get(path)
+        case other =>
+          throw new IllegalStateException(
+            s"Unsupported Fabric token cache path type: ${Option(other).map(_.getClass.getName).orNull}")
+      }
+      Files.deleteIfExists(tokenPath)
+      true
+    }
+  }
+
+  private def clearInMemoryTokenCache(): Boolean = {
+    objectMethod(InMemoryCacheClasses, "clear", 0).exists { case (module, method) =>
+      method.invoke(module)
+      true
+    }
+  }
+
+  private[ml] def invalidateSparkMwcTokenCaches(
+      cacheKey: String,
+      encodeNfsCacheKey: String => String,
+      deleteNfsCacheEntry: String => Boolean,
+      clearInMemoryCache: () => Boolean): Unit = {
+    val deletedNfsToken = deleteNfsCacheEntry(encodeNfsCacheKey(cacheKey))
+    val clearedInMemoryToken = clearInMemoryCache()
+    if (!deletedNfsToken && !clearedInMemoryToken) {
+      throw new NoSuchMethodException("Fabric runtime does not expose MWC token cache invalidation.")
+    }
+  }
+
+  def invalidateSparkMwcToken(workspaceId: String, artifactId: String): Unit = {
+    if (!invalidateWithRuntimeApi(workspaceId, artifactId)) {
+      val cacheKey = workspaceId + artifactId + SparkTokenVersion + SparkWorkloadType
+      invalidateSparkMwcTokenCaches(
+        cacheKey,
+        nfsCacheKey,
+        deleteNfsToken,
+        () => clearInMemoryTokenCache())
+    }
+  }
 
   def getMLWorkloadAADAuthHeader: String = "Bearer " + getAccessToken
 
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
index cfb8e53f1c6..ed103465341 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
@@ -3,6 +3,7 @@
 
 package com.microsoft.azure.synapse.ml.io.http
 
+import com.microsoft.azure.synapse.ml.fabric.FabricClient
 import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
 import org.apache.commons.io.IOUtils
 import org.apache.http.client.config.RequestConfig
@@ -193,17 +194,158 @@ object HandlingUtils extends SparkLogging {
   //scalastyle:on method.length
   //scalastyle:on cyclomatic.complexity
 
+  //scalastyle:off cyclomatic.complexity
+  //scalastyle:off method.length
+  //scalastyle:off magic.number
+  private[ml] def sendWithFabricAuthRetries(
+      client: CloseableHttpClient,
+      requestData: HTTPRequestData,
+      retriesLeft: Array[Int],
+      extraCodesToRetry: Set[Int] = Set(),
+      getAuthHeader: () => String = () => FabricClient.getCognitiveMWCTokenAuthHeader,
+      refreshAuthHeader: String => String = FabricClient.refreshCognitiveMWCTokenAuthHeader,
+      backoff429Ms: Long = 0,
+      authRetryUsed: Boolean = false,
+      authOverride: Option[String] = None): (CloseableHttpResponse, HttpRequestBase) = {
+    val request = requestData.toHTTPCore
+    val authHeader = authOverride.getOrElse(getAuthHeader())
+    request.setHeader("Authorization", authHeader)
+    var executingRequest = true
+    try {
+      val response = client.execute(request)
+      executingRequest = false
+      val code = response.getStatusLine.getStatusCode
+      val capacityLimitExceeded = if (code == 429) {
+        val maxInspectionBytes = 1024 * 1024L
+        Option(response.getEntity).exists { entity =>
+          if (entity.getContentLength > maxInspectionBytes) {
+            false
+          } else {
+            response.setEntity(new BufferedHttpEntity(entity))
+            Option(response.getEntity)
+              .flatMap(e => Try(IOUtils.toString(e.getContent, "UTF-8")).toOption)
+              .exists(_.contains("CapacityLimitExceeded"))
+          }
+        }
+      } else {
+        false
+      }
+
+      val successful = Set(200, 201, 202)(code)
+      val retryable = if (code == 429) {
+        !capacityLimitExceeded
+      } else if (code == 401) {
+        false
+      } else if (extraCodesToRetry(code)) {
+        true
+      } else {
+        !code.toString.startsWith("4")
+      }
+
+      if (code == 401 && !authRetryUsed) {
+        response.close()
+        request.releaseConnection()
+        val refreshedAuthHeader = refreshAuthHeader(authHeader)
+        sendWithFabricAuthRetries(
+          client,
+          requestData,
+          retriesLeft,
+          extraCodesToRetry,
+          getAuthHeader,
+          refreshAuthHeader,
+          backoff429Ms,
+          authRetryUsed = true,
+          authOverride = Some(refreshedAuthHeader))
+      } else if (successful || !retryable || retriesLeft.isEmpty) {
+        if (capacityLimitExceeded) {
+          logWarning(s"Capacity limit exceeded (non-retryable 429) on ${request.getURI}")
+        }
+        response -> request
+      } else {
+        val retryAfterMs = if (code == 429) {
+          Option(response.getFirstHeader("Retry-After"))
+            .flatMap(h => Try(h.getValue.toLong * 1000).toOption)
+            .filter(_ >= 0)
+            .map(math.min(_, MaxBackoffMs))
+        } else {
+          None
+        }
+        response.close()
+        request.releaseConnection()
+        if (code == 429) {
+          val baseBackoff = retryAfterMs.getOrElse {
+            val current = math.max(backoff429Ms, retriesLeft.head.toLong)
+            math.min(current * 2, MaxBackoffMs)
+          }
+          val jitter = Random.nextInt(math.max((baseBackoff / 10).toInt, 1))
+          Thread.sleep(math.min(baseBackoff + jitter, MaxBackoffMs))
+          sendWithFabricAuthRetries(
+            client,
+            requestData,
+            retriesLeft,
+            extraCodesToRetry,
+            getAuthHeader,
+            refreshAuthHeader,
+            baseBackoff,
+            authRetryUsed)
+        } else {
+          Thread.sleep(retriesLeft.head.toLong)
+          sendWithFabricAuthRetries(
+            client,
+            requestData,
+            retriesLeft.tail,
+            extraCodesToRetry,
+            getAuthHeader,
+            refreshAuthHeader,
+            authRetryUsed = authRetryUsed)
+        }
+      }
+    } catch {
+      case e: java.io.IOException if executingRequest =>
+        request.releaseConnection()
+        if (retriesLeft.isEmpty) {
+          throw e
+        }
+        logError("Encountering a connection error", e)
+        Thread.sleep(retriesLeft.head.toLong)
+        sendWithFabricAuthRetries(
+          client,
+          requestData,
+          retriesLeft.tail,
+          extraCodesToRetry,
+          getAuthHeader,
+          refreshAuthHeader,
+          backoff429Ms,
+          authRetryUsed)
+    }
+  }
+  //scalastyle:on magic.number
+  //scalastyle:on method.length
+  //scalastyle:on cyclomatic.complexity
+
   def advanced(retryTimes: Int*)(client: CloseableHttpClient,
                                  request: HTTPRequestData): HTTPResponseData = {
     try {
-      val req = request.toHTTPCore
-      val message = req match {
+      val previewRequest = request.toHTTPCore
+      val message = previewRequest match {
         case r: HttpPost => Try(IOUtils.toString(r.getEntity.getContent, "UTF-8")).getOrElse("")
         case r => r.getURI
       }
+      previewRequest.releaseConnection()
       SynapseMLLogging.logDebug(s"sending $message")
       val start = System.currentTimeMillis()
-      val resp = sendWithRetries(client, req, retryTimes.toArray)
+      val usesTrustedFabricAuth = request.usesFabricAuth &&
+        FabricClient.isOpenAIEndpoint(request.requestLine.uri)
+      val (resp, req) = if (usesTrustedFabricAuth) {
+        sendWithFabricAuthRetries(
+          client,
+          request,
+          retryTimes.toArray,
+          authOverride = request.authorizationHeader)
+      } else {
+        val httpRequest = request.toHTTPCore
+        sendWithRetries(client, httpRequest, retryTimes.toArray) -> httpRequest
+      }
       SynapseMLLogging.logMessage(
         s"finished sending to ${req.getURI} took (${System.currentTimeMillis() - start}ms)")
       val respData = convertAndClose(resp)
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
index 46ffbfc237c..bfff172bb37 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
@@ -166,6 +166,12 @@ case class HTTPRequestData(requestLine: RequestLineData,
                            headers: Array[HeaderData],
                            entity: Option[EntityData]) {
 
+  private[ml] def usesFabricAuth: Boolean =
+    headers.exists(h => HTTPRequestData.isFabricAuthMarker(h.name))
+
+  private[ml] def authorizationHeader: Option[String] =
+    headers.find(_.name.equalsIgnoreCase("Authorization")).map(_.value)
+
   def this(r: HttpRequestBase) = {
     this(new RequestLineData(r.getRequestLine),
       r.getAllHeaders.map(new HeaderData(_)),
@@ -197,7 +203,7 @@ case class HTTPRequestData(requestLine: RequestLineData,
     request.setURI(new URI(requestLine.uri))
     requestLine.protocolVersion.foreach(pv =>
       request.setProtocolVersion(pv.toHTTPCore))
-    request.setHeaders(headers.map(_.toHTTPCore) ++
+    request.setHeaders(headers.filterNot(h => HTTPRequestData.isFabricAuthMarker(h.name)).map(_.toHTTPCore) ++
       Array(new BasicHeader(
         "User-Agent", s"synapseml/${BuildInfo.version}${HeaderValues.PlatformInfo}")))
     request
@@ -206,6 +212,11 @@ case class HTTPRequestData(requestLine: RequestLineData,
 }
 
 object HTTPRequestData extends SparkBindings[HTTPRequestData] {
+  private[ml] val FabricAuthMarkerHeader = "X-SynapseML-Implicit-Fabric-Auth"
+
+  private[ml] def isFabricAuthMarker(headerName: String): Boolean =
+    FabricAuthMarkerHeader.equalsIgnoreCase(headerName)
+
   def fromHTTPExchange(httpEx: HttpExchange): HTTPRequestData = {
     val requestHeaders = httpEx.getRequestHeaders
     val isChunked = Option(requestHeaders.getFirst("Transfer-Encoding") == "chunked").getOrElse(false)
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
new file mode 100644
index 00000000000..1d8b8703ef7
--- /dev/null
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
@@ -0,0 +1,42 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.fabric
+
+import com.microsoft.azure.synapse.ml.core.test.base.TestBase
+
+import java.nio.file.Files
+
+class VerifyTokenInvalidation extends TestBase {
+
+  test("Spark MWC invalidation resolves the encoded NFS cache key") {
+    val tokenPath = Files.createTempFile("synapseml-mwc-token", ".cache")
+    val events = scala.collection.mutable.ArrayBuffer.empty[(String, String)]
+    val logicalCacheKey = "WorkspaceArtifact2SparkCore"
+
+    try {
+      TokenLibrary.invalidateSparkMwcTokenCaches(
+        logicalCacheKey,
+        cacheKey => {
+          events += ("encode" -> cacheKey)
+          "encoded-cache-key"
+        },
+        cacheKey => {
+          events += ("delete" -> cacheKey)
+          Files.deleteIfExists(tokenPath)
+        },
+        () => {
+          events += ("clear" -> "")
+          true
+        })
+
+      assert(events === Seq(
+        "encode" -> logicalCacheKey,
+        "delete" -> "encoded-cache-key",
+        "clear" -> ""))
+      assert(!Files.exists(tokenPath))
+    } finally {
+      Files.deleteIfExists(tokenPath)
+    }
+  }
+}
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
index af1860b9c3e..b868c1dff5c 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
@@ -3,16 +3,20 @@
 
 package com.microsoft.azure.synapse.ml.io.split1
 
-import com.microsoft.azure.synapse.ml.io.http.HandlingUtils
+import com.microsoft.azure.synapse.ml.fabric.FabricClient
+import com.microsoft.azure.synapse.ml.io.http.{HTTPRequestData, HandlingUtils}
 
 import com.microsoft.azure.synapse.ml.core.test.base.TestBase
 import com.sun.net.httpserver.{HttpExchange, HttpServer}
-import org.apache.http.client.methods.HttpGet
+import org.apache.http.client.methods.{HttpGet, HttpPost}
+import org.apache.http.entity.StringEntity
 import org.apache.http.impl.client.HttpClients
 
 import java.net.{InetSocketAddress, ServerSocket}
-import java.util.concurrent.Executors
-import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.{ConcurrentLinkedQueue, Executors}
+import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
+import scala.collection.JavaConverters._
+import scala.io.Source
 
 class VerifySendWithRetries extends TestBase {
 
@@ -44,6 +48,27 @@ class VerifySendWithRetries extends TestBase {
     exchange.close()
   }
 
+  private def readRequestBody(exchange: HttpExchange): String = {
+    val source = Source.fromInputStream(exchange.getRequestBody, "UTF-8")
+    try {
+      source.mkString
+    } finally {
+      source.close()
+    }
+  }
+
+  private def fabricPost(port: Int): HTTPRequestData = {
+    val request = new HttpPost(s"http://localhost:$port/test")
+    request.setHeader(HTTPRequestData.FabricAuthMarkerHeader, "true")
+    request.setHeader("X-Custom", "preserved")
+    request.setHeader("X-Taxonomy-TrafficType", "Background")
+    request.setHeader("X-Llm-Service-Tier", "flex")
+    request.setHeader("X-Taxonomy-ExtendedProperties", """{"feature":"synapseml"}""")
+    request.setHeader("x-ms-llm-feature-name", "SparkCodeFirst")
+    request.setEntity(new StringEntity("""{"prompt":"hello"}""", "UTF-8"))
+    new HTTPRequestData(request)
+  }
+
   test("429 without Retry-After uses exponential backoff") {
     val port = getFreePort
     val requestCount = new AtomicInteger(0)
@@ -427,4 +452,182 @@ class VerifySendWithRetries extends TestBase {
       server.stop(0)
     }
   }
+
+  test("implicit Fabric auth refreshes and replays a request once after 401") {
+    val port = getFreePort
+    val requestCount = new AtomicInteger(0)
+    val refreshCount = new AtomicInteger(0)
+    val currentAuth = new AtomicReference("MwcToken stale")
+    val authHeaders = new ConcurrentLinkedQueue[String]()
+    val requestBodies = new ConcurrentLinkedQueue[String]()
+    val customHeaders = new ConcurrentLinkedQueue[String]()
+    val taxonomyHeaders = new ConcurrentLinkedQueue[String]()
+    val server = startServer(port) { exchange =>
+      authHeaders.add(exchange.getRequestHeaders.getFirst("Authorization"))
+      customHeaders.add(exchange.getRequestHeaders.getFirst("X-Custom"))
+      taxonomyHeaders.add(exchange.getRequestHeaders.getFirst("X-Taxonomy-TrafficType"))
+      requestBodies.add(readRequestBody(exchange))
+      if (requestCount.incrementAndGet() == 1) {
+        respond(exchange, 401, "expired")
+      } else {
+        respond(exchange, 200, """{"ok":true}""")
+      }
+    }
+    try {
+      val client = HttpClients.createDefault()
+      val requestData = fabricPost(port)
+      val (response, request) = HandlingUtils.sendWithFabricAuthRetries(
+        client,
+        requestData,
+        Array(10),
+        getAuthHeader = () => currentAuth.get(),
+        refreshAuthHeader = rejectedAuthHeader => {
+          assert(rejectedAuthHeader == "MwcToken stale")
+          refreshCount.incrementAndGet()
+          currentAuth.set("MwcToken fresh")
+          currentAuth.get()
+        })
+      val code = response.getStatusLine.getStatusCode
+      response.close()
+      request.releaseConnection()
+      client.close()
+
+      assert(code === 200)
+      assert(requestCount.get() === 2)
+      assert(refreshCount.get() === 1)
+      assert(authHeaders.asScala.toSeq === Seq("MwcToken stale", "MwcToken fresh"))
+      assert(requestBodies.asScala.toSeq === Seq("""{"prompt":"hello"}""", """{"prompt":"hello"}"""))
+      assert(customHeaders.asScala.toSeq === Seq("preserved", "preserved"))
+      assert(taxonomyHeaders.asScala.toSeq === Seq("Background", "Background"))
+      Seq(
+        "X-Taxonomy-TrafficType",
+        "X-Llm-Service-Tier",
+        "X-Taxonomy-ExtendedProperties",
+        "x-ms-llm-feature-name"
+      ).foreach { headerName =>
+        assert(requestData.headers.exists(_.name.equalsIgnoreCase(headerName)))
+      }
+    } finally {
+      server.stop(0)
+    }
+  }
+
+  test("implicit Fabric auth returns the second 401 without retrying again") {
+    val port = getFreePort
+    val requestCount = new AtomicInteger(0)
+    val refreshCount = new AtomicInteger(0)
+    val server = startServer(port) { exchange =>
+      requestCount.incrementAndGet()
+      readRequestBody(exchange)
+      respond(exchange, 401, "unauthorized")
+    }
+    try {
+      val client = HttpClients.createDefault()
+      val (response, request) = HandlingUtils.sendWithFabricAuthRetries(
+        client,
+        fabricPost(port),
+        Array(10, 10),
+        extraCodesToRetry = Set(401),
+        getAuthHeader = () => "MwcToken stale",
+        refreshAuthHeader = _ => {
+          refreshCount.incrementAndGet()
+          "MwcToken refreshed"
+        })
+      val code = response.getStatusLine.getStatusCode
+      response.close()
+      request.releaseConnection()
+      client.close()
+
+      assert(code === 401)
+      assert(requestCount.get() === 2)
+      assert(refreshCount.get() === 1)
+    } finally {
+      server.stop(0)
+    }
+  }
+
+  test("implicit Fabric auth is reacquired for a 429 retry") {
+    val port = getFreePort
+    val requestCount = new AtomicInteger(0)
+    val authCallCount = new AtomicInteger(0)
+    val authHeaders = new ConcurrentLinkedQueue[String]()
+    val server = startServer(port) { exchange =>
+      authHeaders.add(exchange.getRequestHeaders.getFirst("Authorization"))
+      readRequestBody(exchange)
+      if (requestCount.incrementAndGet() == 1) {
+        respond(exchange, 429, headers = Map("Retry-After" -> "0"))
+      } else {
+        respond(exchange, 200, """{"ok":true}""")
+      }
+    }
+    try {
+      val client = HttpClients.createDefault()
+      val (response, request) = HandlingUtils.sendWithFabricAuthRetries(
+        client,
+        fabricPost(port),
+        Array(10),
+        getAuthHeader = () => s"MwcToken token-${authCallCount.incrementAndGet()}",
+        refreshAuthHeader = _ => fail("401 refresh should not run for a 429"))
+      val code = response.getStatusLine.getStatusCode
+      response.close()
+      request.releaseConnection()
+      client.close()
+
+      assert(code === 200)
+      assert(requestCount.get() === 2)
+      assert(authCallCount.get() === 2)
+      assert(authHeaders.asScala.toSeq === Seq("MwcToken token-1", "MwcToken token-2"))
+    } finally {
+      server.stop(0)
+    }
+  }
+
+  test("untrusted marker does not replace explicit auth on a non-Fabric endpoint") {
+    val port = getFreePort
+    val authorization = new AtomicReference[String]()
+    val server = startServer(port) { exchange =>
+      authorization.set(exchange.getRequestHeaders.getFirst("Authorization"))
+      respond(exchange, 200, """{"ok":true}""")
+    }
+    try {
+      val client = HttpClients.createDefault()
+      val request = new HttpGet(s"http://localhost:$port/test")
+      request.setHeader(HTTPRequestData.FabricAuthMarkerHeader, "true")
+      request.setHeader("Authorization", "Bearer explicit")
+
+      val response = HandlingUtils.advanced(10)(client, new HTTPRequestData(request))
+
+      client.close()
+      assert(response.statusLine.statusCode === 200)
+      assert(authorization.get() === "Bearer explicit")
+    } finally {
+      server.stop(0)
+    }
+  }
+
+  test("Fabric endpoint validation requires HTTPS host and path containment") {
+    val endpointRoot = "https://workspace.fabric.microsoft.com/cognitive/openai/"
+
+    assert(FabricClient.isEndpointUnder(
+      "https://workspace.fabric.microsoft.com//cognitive/openai/chat",
+      endpointRoot))
+    assert(FabricClient.isEndpointUnder(
+      "https://workspace.fabric.microsoft.com:443/cognitive/openai/chat",
+      endpointRoot))
+    assert(!FabricClient.isEndpointUnder(
+      "http://workspace.fabric.microsoft.com/cognitive/openai/chat",
+      endpointRoot))
+    assert(!FabricClient.isEndpointUnder(
+      "https://attacker.example/cognitive/openai/chat",
+      endpointRoot))
+    assert(!FabricClient.isEndpointUnder(
+      "https://workspace.fabric.microsoft.com/other",
+      endpointRoot))
+    assert(!FabricClient.isEndpointUnder(
+      "https://workspace.fabric.microsoft.com/cognitive/openai/../other",
+      endpointRoot))
+    assert(!FabricClient.isEndpointUnder(
+      "https://workspace.fabric.microsoft.com/cognitive/openai/%2e%2e/other",
+      endpointRoot))
+  }
 }

From 15c9eb72b23cf037bfb0737298c5a1329601e351 Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 17:48:01 -0700
Subject: [PATCH 02/10] fix: harden Fabric auth retry trust boundaries

AB#3380998

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../synapse/ml/fabric/FabricClient.scala      |  33 ++++--
 .../synapse/ml/io/http/HTTPClients.scala      |  83 ++++++++------
 .../azure/synapse/ml/io/http/HTTPSchema.scala |  12 +-
 .../ml/fabric/VerifyTokenInvalidation.scala   |  47 ++++++++
 .../ml/io/split1/VerifySendWithRetries.scala  | 107 ++++++++++++++++++
 5 files changed, 239 insertions(+), 43 deletions(-)

diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala
index c68582d455e..6be1d4d56ea 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala
@@ -7,7 +7,7 @@ import spray.json.DefaultJsonProtocol.{StringJsonFormat, mapFormat}
 import spray.json._
 
 import java.net.{MalformedURLException, URI, URL}
-import java.util.UUID
+import java.util.{Locale, UUID}
 import scala.collection.concurrent.TrieMap
 import scala.io.Source
 import scala.util.control.NonFatal
@@ -22,6 +22,7 @@ object FabricClient extends RESTUtils {
   private val SparkConfPath = "/opt/spark/conf/spark-defaults.conf";
   private val ClusterInfoPath = "/opt/health-agent/conf/cluster-info.json";
   private val CognitiveMwcRefreshLocks = TrieMap.empty[(String, String), AnyRef]
+  private val AmbiguousPathEncoding = "(?i)%(?:25)*(?:2e|2f|5c)".r
 
   lazy val CapacityID: Option[String] = getCapacityID;
   lazy val WorkspaceID: Option[String] = getWorkspaceID;
@@ -220,11 +221,13 @@ object FabricClient extends RESTUtils {
   }
 
   private def trustedPath(rawPath: String): Option[String] = {
-    val lowerPath = rawPath.toLowerCase
-    val hasAmbiguousEncoding = Seq("%2e", "%2f", "%5c").exists(lowerPath.contains)
-    val normalizedPath = rawPath.replaceAll("/+", "/")
-    val hasDotSegment = normalizedPath.split("/", -1).exists(segment => segment == "." || segment == "..")
-    if (hasAmbiguousEncoding || rawPath.contains("\\") || hasDotSegment) None else Some(normalizedPath)
+    Option(rawPath).flatMap { path =>
+      val lowerPath = path.toLowerCase(Locale.ROOT)
+      val hasAmbiguousEncoding = AmbiguousPathEncoding.findFirstIn(lowerPath).nonEmpty
+      val normalizedPath = path.replaceAll("/+", "/")
+      val hasDotSegment = normalizedPath.split("/", -1).exists(segment => segment == "." || segment == "..")
+      if (hasAmbiguousEncoding || path.contains("\\") || hasDotSegment) None else Some(normalizedPath)
+    }
   }
 
   private[ml] def isOpenAIEndpoint(requestUrl: String): Boolean = {
@@ -239,13 +242,25 @@ object FabricClient extends RESTUtils {
     val workspaceId = WorkspaceID.getOrElse("")
     val artifactId = ArtifactID.getOrElse("")
     val refreshLock = CognitiveMwcRefreshLocks.getOrElseUpdate((workspaceId, artifactId), new Object())
+    refreshAuthHeader(
+      rejectedAuthHeader,
+      refreshLock,
+      () => TokenLibrary.getCognitiveMwcTokenAuthHeader(workspaceId, artifactId),
+      () => TokenLibrary.invalidateSparkMwcToken(workspaceId, artifactId))
+  }
+
+  private[ml] def refreshAuthHeader(
+      rejectedAuthHeader: String,
+      refreshLock: AnyRef,
+      getCurrentAuthHeader: () => String,
+      invalidateAuthHeader: () => Unit): String = {
     refreshLock.synchronized {
-      val currentAuthHeader = TokenLibrary.getCognitiveMwcTokenAuthHeader(workspaceId, artifactId)
+      val currentAuthHeader = getCurrentAuthHeader()
       if (currentAuthHeader != rejectedAuthHeader) {
         currentAuthHeader
       } else {
-        TokenLibrary.invalidateSparkMwcToken(workspaceId, artifactId)
-        TokenLibrary.getCognitiveMwcTokenAuthHeader(workspaceId, artifactId)
+        invalidateAuthHeader()
+        getCurrentAuthHeader()
       }
     }
   }
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
index ed103465341..a5f44e0708d 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
@@ -8,7 +8,7 @@ import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
 import org.apache.commons.io.IOUtils
 import org.apache.http.client.config.RequestConfig
 import org.apache.http.client.methods.{CloseableHttpResponse, HttpPost, HttpRequestBase}
-import org.apache.http.entity.BufferedHttpEntity
+import org.apache.http.entity.{BasicHttpEntity, ByteArrayEntity}
 import org.apache.http.impl.client.{CloseableHttpClient, HttpClientBuilder}
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager
 import org.apache.spark.injections.UDFUtils
@@ -16,6 +16,8 @@ import org.apache.spark.internal.{Logging => SparkLogging}
 import org.apache.spark.sql.expressions.UserDefinedFunction
 import org.apache.spark.sql.types.StringType
 
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, SequenceInputStream}
+import java.nio.charset.StandardCharsets
 import scala.concurrent.duration.Duration
 import scala.concurrent.{ExecutionContext, blocking}
 import scala.util.{Random, Try}
@@ -88,6 +90,48 @@ object HandlingUtils extends SparkLogging {
   }
 
   private val MaxBackoffMs: Long = 60000L // 1 minute cap for 429 backoff
+  private val MaxResponseInspectionBytes = 1024 * 1024L
+
+  private def responseBodyForInspection(response: CloseableHttpResponse): Option[String] = {
+    Option(response.getEntity).flatMap { entity =>
+      if (entity.getContentLength > MaxResponseInspectionBytes) {
+        None
+      } else {
+        val output = new ByteArrayOutputStream()
+        val input = entity.getContent
+        var keepInputOpen = false
+        try {
+          IOUtils.copyLarge(input, output, 0, MaxResponseInspectionBytes + 1)
+          val bytes = output.toByteArray
+          if (bytes.length > MaxResponseInspectionBytes) {
+            val replayEntity = new BasicHttpEntity()
+            replayEntity.setContent(new SequenceInputStream(new ByteArrayInputStream(bytes), input))
+            replayEntity.setContentLength(entity.getContentLength)
+            Option(entity.getContentEncoding).foreach(replayEntity.setContentEncoding)
+            Option(entity.getContentType).foreach(replayEntity.setContentType)
+            replayEntity.setChunked(entity.isChunked)
+            response.setEntity(replayEntity)
+            keepInputOpen = true
+            None
+          } else {
+            val bufferedEntity = new ByteArrayEntity(bytes)
+            Option(entity.getContentEncoding).foreach(bufferedEntity.setContentEncoding)
+            Option(entity.getContentType).foreach(bufferedEntity.setContentType)
+            bufferedEntity.setChunked(entity.isChunked)
+            response.setEntity(bufferedEntity)
+            Some(new String(bytes, StandardCharsets.UTF_8))
+          }
+        } finally {
+          if (!keepInputOpen) {
+            input.close()
+          }
+        }
+      }
+    }
+  }
+
+  private def capacityLimitExceeded(response: CloseableHttpResponse): Boolean =
+    responseBodyForInspection(response).exists(_.contains("CapacityLimitExceeded"))
 
   //scalastyle:off cyclomatic.complexity
   //scalastyle:off method.length
@@ -107,19 +151,7 @@ object HandlingUtils extends SparkLogging {
         case 202 => true
         case 429 =>
           // Inspect body to distinguish capacity errors from transient rate limits.
-          // Guard with Content-Length cap to avoid buffering unexpectedly large payloads.
-          val MaxInspectionBytes = 1024 * 1024L
-          val bodyStr = Option(response.getEntity).flatMap { entity =>
-            val contentLength = entity.getContentLength
-            if (contentLength > MaxInspectionBytes) {
-              None
-            } else {
-              response.setEntity(new BufferedHttpEntity(entity))
-              Option(response.getEntity)
-                .flatMap(e => Try(IOUtils.toString(e.getContent, "UTF-8")).toOption)
-            }
-          }.getOrElse("")
-          if (bodyStr.contains("CapacityLimitExceeded")) {
+          if (capacityLimitExceeded(response)) {
             // Fabric capacity-exceeded 429s are NOT transient rate limits —
             // retrying will not help and causes hangs
             logWarning(s"Capacity limit exceeded (non-retryable 429) on ${request.getURI}")
@@ -209,31 +241,18 @@ object HandlingUtils extends SparkLogging {
       authOverride: Option[String] = None): (CloseableHttpResponse, HttpRequestBase) = {
     val request = requestData.toHTTPCore
     val authHeader = authOverride.getOrElse(getAuthHeader())
+    request.removeHeaders("Authorization")
     request.setHeader("Authorization", authHeader)
     var executingRequest = true
     try {
       val response = client.execute(request)
       executingRequest = false
       val code = response.getStatusLine.getStatusCode
-      val capacityLimitExceeded = if (code == 429) {
-        val maxInspectionBytes = 1024 * 1024L
-        Option(response.getEntity).exists { entity =>
-          if (entity.getContentLength > maxInspectionBytes) {
-            false
-          } else {
-            response.setEntity(new BufferedHttpEntity(entity))
-            Option(response.getEntity)
-              .flatMap(e => Try(IOUtils.toString(e.getContent, "UTF-8")).toOption)
-              .exists(_.contains("CapacityLimitExceeded"))
-          }
-        }
-      } else {
-        false
-      }
+      val capacityLimitExceededResponse = code == 429 && capacityLimitExceeded(response)
 
       val successful = Set(200, 201, 202)(code)
       val retryable = if (code == 429) {
-        !capacityLimitExceeded
+        !capacityLimitExceededResponse
       } else if (code == 401) {
         false
       } else if (extraCodesToRetry(code)) {
@@ -257,7 +276,7 @@ object HandlingUtils extends SparkLogging {
           authRetryUsed = true,
           authOverride = Some(refreshedAuthHeader))
       } else if (successful || !retryable || retriesLeft.isEmpty) {
-        if (capacityLimitExceeded) {
+        if (capacityLimitExceededResponse) {
           logWarning(s"Capacity limit exceeded (non-retryable 429) on ${request.getURI}")
         }
         response -> request
@@ -282,7 +301,7 @@ object HandlingUtils extends SparkLogging {
           sendWithFabricAuthRetries(
             client,
             requestData,
-            retriesLeft,
+            retriesLeft.tail,
             extraCodesToRetry,
             getAuthHeader,
             refreshAuthHeader,
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
index bfff172bb37..f6cb652e502 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
@@ -167,10 +167,13 @@ case class HTTPRequestData(requestLine: RequestLineData,
                            entity: Option[EntityData]) {
 
   private[ml] def usesFabricAuth: Boolean =
-    headers.exists(h => HTTPRequestData.isFabricAuthMarker(h.name))
+    headers.exists(h =>
+      HTTPRequestData.isFabricAuthMarker(h.name) &&
+        Option(h.value).exists(_.equalsIgnoreCase("true"))) &&
+      authorizationHeader.exists(HTTPRequestData.isCognitiveMwcAuthHeader)
 
   private[ml] def authorizationHeader: Option[String] =
-    headers.find(_.name.equalsIgnoreCase("Authorization")).map(_.value)
+    headers.find(h => "Authorization".equalsIgnoreCase(h.name)).map(_.value)
 
   def this(r: HttpRequestBase) = {
     this(new RequestLineData(r.getRequestLine),
@@ -217,6 +220,11 @@ object HTTPRequestData extends SparkBindings[HTTPRequestData] {
   private[ml] def isFabricAuthMarker(headerName: String): Boolean =
     FabricAuthMarkerHeader.equalsIgnoreCase(headerName)
 
+  private def isCognitiveMwcAuthHeader(value: String): Boolean = {
+    val parts = Option(value).map(_.trim.split("\\s+", 2)).getOrElse(Array.empty)
+    parts.length == 2 && parts.head.equalsIgnoreCase("MwcToken") && parts.last.nonEmpty
+  }
+
   def fromHTTPExchange(httpEx: HttpExchange): HTTPRequestData = {
     val requestHeaders = httpEx.getRequestHeaders
     val isChunked = Option(requestHeaders.getFirst("Transfer-Encoding") == "chunked").getOrElse(false)
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
index 1d8b8703ef7..23c48089853 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
@@ -6,6 +6,8 @@ package com.microsoft.azure.synapse.ml.fabric
 import com.microsoft.azure.synapse.ml.core.test.base.TestBase
 
 import java.nio.file.Files
+import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
+import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit}
 
 class VerifyTokenInvalidation extends TestBase {
 
@@ -39,4 +41,49 @@ class VerifyTokenInvalidation extends TestBase {
       Files.deleteIfExists(tokenPath)
     }
   }
+
+  test("Spark MWC invalidation fails when the runtime exposes no supported cache API") {
+    val error = intercept[NoSuchMethodException] {
+      TokenLibrary.invalidateSparkMwcTokenCaches(
+        "cache-key",
+        identity,
+        _ => false,
+        () => false)
+    }
+
+    assert(error.getMessage.contains("does not expose MWC token cache invalidation"))
+  }
+
+  test("concurrent refreshes invalidate a rejected MWC token once") {
+    val workers = 8
+    val executor = Executors.newFixedThreadPool(workers)
+    val start = new CountDownLatch(1)
+    val currentAuth = new AtomicReference("MwcToken stale")
+    val invalidationCount = new AtomicInteger(0)
+    val refreshLock = new Object()
+
+    try {
+      val futures = (1 to workers).map { _ =>
+        executor.submit(new Callable[String] {
+          override def call(): String = {
+            start.await()
+            FabricClient.refreshAuthHeader(
+              "MwcToken stale",
+              refreshLock,
+              () => currentAuth.get(),
+              () => {
+                invalidationCount.incrementAndGet()
+                currentAuth.set("MwcToken fresh")
+              })
+          }
+        })
+      }
+      start.countDown()
+
+      assert(futures.map(_.get(10, TimeUnit.SECONDS)) === Seq.fill(workers)("MwcToken fresh"))
+      assert(invalidationCount.get() === 1)
+    } finally {
+      executor.shutdownNow()
+    }
+  }
 }
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
index b868c1dff5c..011d2233fd7 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
@@ -582,6 +582,107 @@ class VerifySendWithRetries extends TestBase {
     }
   }
 
+  test("implicit Fabric auth bounds 429 retries") {
+    val port = getFreePort
+    val requestCount = new AtomicInteger(0)
+    val server = startServer(port) { exchange =>
+      requestCount.incrementAndGet()
+      readRequestBody(exchange)
+      respond(exchange, 429, """{"error":{"code":"RateLimitExceeded"}}""",
+        headers = Map("Retry-After" -> "0"))
+    }
+    try {
+      val client = HttpClients.createDefault()
+      val (response, request) = HandlingUtils.sendWithFabricAuthRetries(
+        client,
+        fabricPost(port),
+        Array(0),
+        getAuthHeader = () => "MwcToken current",
+        refreshAuthHeader = _ => fail("401 refresh should not run for a 429"))
+      val code = response.getStatusLine.getStatusCode
+      response.close()
+      request.releaseConnection()
+      client.close()
+
+      assert(code === 429)
+      assert(requestCount.get() === 2, "Initial request plus one configured retry should be sent")
+    } finally {
+      server.stop(0)
+    }
+  }
+
+  test("unknown-length 429 response bodies remain readable") {
+    val port = getFreePort
+    val responseBody = """{"error":{"code":"RateLimitExceeded"}}"""
+    val server = startServer(port) { exchange =>
+      exchange.sendResponseHeaders(429, 0)
+      val output = exchange.getResponseBody
+      output.write(responseBody.getBytes("UTF-8"))
+      output.close()
+      exchange.close()
+    }
+    try {
+      val client = HttpClients.createDefault()
+      val request = new HttpGet(s"http://localhost:$port/test")
+      val response = HandlingUtils.sendWithRetries(client, request, Array.empty)
+      val body = Source.fromInputStream(response.getEntity.getContent, "UTF-8")
+      val actualBody = try {
+        body.mkString
+      } finally {
+        body.close()
+      }
+      response.close()
+      client.close()
+
+      assert(actualBody === responseBody)
+    } finally {
+      server.stop(0)
+    }
+  }
+
+  test("large unknown-length 429 response bodies remain readable") {
+    val port = getFreePort
+    val responseBody = "x" * (1024 * 1024 + 128)
+    val server = startServer(port) { exchange =>
+      exchange.sendResponseHeaders(429, 0)
+      val output = exchange.getResponseBody
+      output.write(responseBody.getBytes("UTF-8"))
+      output.close()
+      exchange.close()
+    }
+    try {
+      val client = HttpClients.createDefault()
+      val request = new HttpGet(s"http://localhost:$port/test")
+      val response = HandlingUtils.sendWithRetries(client, request, Array.empty)
+      val body = Source.fromInputStream(response.getEntity.getContent, "UTF-8")
+      val actualBody = try {
+        body.mkString
+      } finally {
+        body.close()
+      }
+      response.close()
+      client.close()
+
+      assert(actualBody === responseBody)
+    } finally {
+      server.stop(0)
+    }
+  }
+
+  test("Fabric auth marker requires an MWC authorization header") {
+    def requestData(markerValue: String, authHeader: String): HTTPRequestData = {
+      val request = new HttpGet("https://workspace.fabric.microsoft.com/cognitive/openai/chat")
+      request.setHeader(HTTPRequestData.FabricAuthMarkerHeader, markerValue)
+      request.setHeader("Authorization", authHeader)
+      new HTTPRequestData(request)
+    }
+
+    assert(requestData("true", "MwcToken token").usesFabricAuth)
+    assert(!requestData("false", "MwcToken token").usesFabricAuth)
+    assert(!requestData("true", "Bearer explicit").usesFabricAuth)
+    assert(!requestData("true", "MwcToken ").usesFabricAuth)
+  }
+
   test("untrusted marker does not replace explicit auth on a non-Fabric endpoint") {
     val port = getFreePort
     val authorization = new AtomicReference[String]()
@@ -629,5 +730,11 @@ class VerifySendWithRetries extends TestBase {
     assert(!FabricClient.isEndpointUnder(
       "https://workspace.fabric.microsoft.com/cognitive/openai/%2e%2e/other",
       endpointRoot))
+    assert(!FabricClient.isEndpointUnder(
+      "https://workspace.fabric.microsoft.com/cognitive/openai/%252e%252e/other",
+      endpointRoot))
+    assert(!FabricClient.isEndpointUnder(
+      "https://workspace.fabric.microsoft.com",
+      endpointRoot))
   }
 }

From bea3088873e108ebdbd3062af416351fc0cf571d Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 17:52:47 -0700
Subject: [PATCH 03/10] test: cover Fabric retry authorization deduplication

AB#3582121

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../ml/io/split1/VerifySendWithRetries.scala  | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
index 011d2233fd7..1f209065cff 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
@@ -683,6 +683,39 @@ class VerifySendWithRetries extends TestBase {
     assert(!requestData("true", "MwcToken ").usesFabricAuth)
   }
 
+  test("Fabric auth retries replace duplicate authorization headers") {
+    val port = getFreePort
+    val authorizationHeaders = new AtomicReference[Seq[String]]()
+    val server = startServer(port) { exchange =>
+      authorizationHeaders.set(exchange.getRequestHeaders.get("Authorization").asScala.toSeq)
+      readRequestBody(exchange)
+      respond(exchange, 200, """{"ok":true}""")
+    }
+    try {
+      val request = new HttpPost(s"http://localhost:$port/test")
+      request.setHeader(HTTPRequestData.FabricAuthMarkerHeader, "true")
+      request.addHeader("Authorization", "MwcToken stale")
+      request.addHeader("authorization", "Bearer duplicate")
+      request.setEntity(new StringEntity("""{"prompt":"hello"}""", "UTF-8"))
+
+      val client = HttpClients.createDefault()
+      val (response, replayedRequest) = HandlingUtils.sendWithFabricAuthRetries(
+        client,
+        new HTTPRequestData(request),
+        Array.empty,
+        getAuthHeader = () => "MwcToken current")
+      val code = response.getStatusLine.getStatusCode
+      response.close()
+      replayedRequest.releaseConnection()
+      client.close()
+
+      assert(code === 200)
+      assert(authorizationHeaders.get() === Seq("MwcToken current"))
+    } finally {
+      server.stop(0)
+    }
+  }
+
   test("untrusted marker does not replace explicit auth on a non-Fabric endpoint") {
     val port = getFreePort
     val authorization = new AtomicReference[String]()

From d2e84892fb54ba836ca0a3c2e034810ad96f364c Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 18:07:42 -0700
Subject: [PATCH 04/10] fix: preserve retry responses after inspection failures

Short-circuit Fabric runtime reflection after the first compatible class and make bounded response inspection best-effort while replaying partially read bytes.

AB#3582121

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../synapse/ml/fabric/TokenLibrary.scala      |  2 +-
 .../synapse/ml/io/http/HTTPClients.scala      | 68 +++++++++++++----
 .../split1/VerifyResponseBodyInspection.scala | 73 +++++++++++++++++++
 3 files changed, 126 insertions(+), 17 deletions(-)
 create mode 100644 core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala

diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
index a036d0eb8b8..5dbc884d723 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
@@ -75,7 +75,7 @@ object TokenLibrary {
       } catch {
         case _: ClassNotFoundException | _: NoSuchFieldException => None
       }
-    }.toSeq.headOption
+    }.take(1).toSeq.headOption
   }
 
   private def invalidateWithRuntimeApi(workspaceId: String, artifactId: String): Boolean = {
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
index a5f44e0708d..85b2cadcd9c 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
@@ -6,9 +6,10 @@ package com.microsoft.azure.synapse.ml.io.http
 import com.microsoft.azure.synapse.ml.fabric.FabricClient
 import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
 import org.apache.commons.io.IOUtils
+import org.apache.http.HttpEntity
 import org.apache.http.client.config.RequestConfig
 import org.apache.http.client.methods.{CloseableHttpResponse, HttpPost, HttpRequestBase}
-import org.apache.http.entity.{BasicHttpEntity, ByteArrayEntity}
+import org.apache.http.entity.{AbstractHttpEntity, BasicHttpEntity, ByteArrayEntity}
 import org.apache.http.impl.client.{CloseableHttpClient, HttpClientBuilder}
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager
 import org.apache.spark.injections.UDFUtils
@@ -16,11 +17,12 @@ import org.apache.spark.internal.{Logging => SparkLogging}
 import org.apache.spark.sql.expressions.UserDefinedFunction
 import org.apache.spark.sql.types.StringType
 
-import java.io.{ByteArrayInputStream, ByteArrayOutputStream, SequenceInputStream}
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, InputStream, SequenceInputStream}
 import java.nio.charset.StandardCharsets
 import scala.concurrent.duration.Duration
 import scala.concurrent.{ExecutionContext, blocking}
 import scala.util.{Random, Try}
+import scala.util.control.NonFatal
 
 trait Handler {
 
@@ -92,38 +94,72 @@ object HandlingUtils extends SparkLogging {
   private val MaxBackoffMs: Long = 60000L // 1 minute cap for 429 backoff
   private val MaxResponseInspectionBytes = 1024 * 1024L
 
-  private def responseBodyForInspection(response: CloseableHttpResponse): Option[String] = {
+  private def copyResponseEntityMetadata(source: HttpEntity, target: AbstractHttpEntity): Unit = {
+    Option(source.getContentEncoding).foreach(target.setContentEncoding)
+    Option(source.getContentType).foreach(target.setContentType)
+    target.setChunked(source.isChunked)
+  }
+
+  private def replayResponseEntity(response: CloseableHttpResponse,
+                                   source: HttpEntity,
+                                   bytes: Array[Byte],
+                                   input: InputStream): Unit = {
+    val replay = new BasicHttpEntity()
+    replay.setContent(new SequenceInputStream(new ByteArrayInputStream(bytes), input))
+    replay.setContentLength(source.getContentLength)
+    copyResponseEntityMetadata(source, replay)
+    response.setEntity(replay)
+  }
+
+  private def closeInspectionInput(input: InputStream): Unit = {
+    try {
+      input.close()
+    } catch {
+      case NonFatal(error) =>
+        logWarning("Could not close the HTTP response inspection stream.", error)
+    }
+  }
+
+  private[ml] def responseBodyForInspection(response: CloseableHttpResponse): Option[String] = {
     Option(response.getEntity).flatMap { entity =>
       if (entity.getContentLength > MaxResponseInspectionBytes) {
         None
       } else {
         val output = new ByteArrayOutputStream()
-        val input = entity.getContent
+        var input = Option.empty[InputStream]
         var keepInputOpen = false
         try {
-          IOUtils.copyLarge(input, output, 0, MaxResponseInspectionBytes + 1)
+          val responseInput = entity.getContent
+          input = Some(responseInput)
+          IOUtils.copyLarge(responseInput, output, 0, MaxResponseInspectionBytes + 1)
           val bytes = output.toByteArray
           if (bytes.length > MaxResponseInspectionBytes) {
-            val replayEntity = new BasicHttpEntity()
-            replayEntity.setContent(new SequenceInputStream(new ByteArrayInputStream(bytes), input))
-            replayEntity.setContentLength(entity.getContentLength)
-            Option(entity.getContentEncoding).foreach(replayEntity.setContentEncoding)
-            Option(entity.getContentType).foreach(replayEntity.setContentType)
-            replayEntity.setChunked(entity.isChunked)
-            response.setEntity(replayEntity)
+            replayResponseEntity(response, entity, bytes, responseInput)
             keepInputOpen = true
             None
           } else {
             val bufferedEntity = new ByteArrayEntity(bytes)
-            Option(entity.getContentEncoding).foreach(bufferedEntity.setContentEncoding)
-            Option(entity.getContentType).foreach(bufferedEntity.setContentType)
-            bufferedEntity.setChunked(entity.isChunked)
+            copyResponseEntityMetadata(entity, bufferedEntity)
             response.setEntity(bufferedEntity)
             Some(new String(bytes, StandardCharsets.UTF_8))
           }
+        } catch {
+          case NonFatal(error) =>
+            logWarning("Could not inspect the HTTP response body; preserving it for retry handling.", error)
+            input.foreach { responseInput =>
+              try {
+                replayResponseEntity(response, entity, output.toByteArray, responseInput)
+                keepInputOpen = true
+              } catch {
+                case NonFatal(replayError) =>
+                  error.addSuppressed(replayError)
+                  logWarning("Could not reconstruct the partially inspected HTTP response body.", replayError)
+              }
+            }
+            None
         } finally {
           if (!keepInputOpen) {
-            input.close()
+            input.foreach(closeInspectionInput)
           }
         }
       }
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala
new file mode 100644
index 00000000000..9211c965e0c
--- /dev/null
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala
@@ -0,0 +1,73 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.io.split1
+
+import com.microsoft.azure.synapse.ml.core.test.base.TestBase
+import com.microsoft.azure.synapse.ml.io.http.HandlingUtils
+import org.apache.http.HttpVersion
+import org.apache.http.client.methods.CloseableHttpResponse
+import org.apache.http.entity.BasicHttpEntity
+import org.apache.http.message.{BasicHttpResponse, BasicStatusLine}
+
+import java.io.{IOException, InputStream}
+import scala.io.Source
+
+class VerifyResponseBodyInspection extends TestBase {
+
+  test("response inspection replays bytes after an input failure") {
+    val content = """{"error":{"code":"RateLimitExceeded"}}""".getBytes("UTF-8")
+    val input = new InputStream {
+      private var index = 0
+      private var failed = false
+
+      override def read(): Int = {
+        if (index < content.length) {
+          val value = content(index) & 0xff
+          index += 1
+          value
+        } else {
+          failOrFinish()
+        }
+      }
+
+      override def read(buffer: Array[Byte], offset: Int, length: Int): Int = {
+        if (index < content.length) {
+          val count = math.min(length, content.length - index)
+          System.arraycopy(content, index, buffer, offset, count)
+          index += count
+          count
+        } else {
+          failOrFinish()
+        }
+      }
+
+      private def failOrFinish(): Int = {
+        if (!failed) {
+          failed = true
+          throw new IOException("simulated response read failure")
+        }
+        -1
+      }
+    }
+    val entity = new BasicHttpEntity()
+    entity.setContent(input)
+    entity.setContentLength(-1)
+    val response = new BasicHttpResponse(
+      new BasicStatusLine(HttpVersion.HTTP_1_1, 429, "Too Many Requests"))
+      with CloseableHttpResponse {
+      override def close(): Unit = ()
+    }
+    response.setEntity(entity)
+
+    assert(HandlingUtils.responseBodyForInspection(response).isEmpty)
+    val replayed = Source.fromInputStream(response.getEntity.getContent, "UTF-8")
+    val replayedBody = try {
+      replayed.mkString
+    } finally {
+      replayed.close()
+    }
+
+    assert(replayedBody === new String(content, "UTF-8"))
+  }
+}

From c530008a6020fc1081f3f3f14f2f302a8d460010 Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 18:13:32 -0700
Subject: [PATCH 05/10] fix: close retry preview resources

Close preview entity streams and always release preview requests, with focused cleanup coverage.

AB#3582121

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../synapse/ml/io/http/HTTPClients.scala      | 29 +++++++++++++----
 .../split1/VerifyResponseBodyInspection.scala | 31 +++++++++++++++++--
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
index 85b2cadcd9c..bc1286ce92e 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
@@ -169,6 +169,28 @@ object HandlingUtils extends SparkLogging {
   private def capacityLimitExceeded(response: CloseableHttpResponse): Boolean =
     responseBodyForInspection(response).exists(_.contains("CapacityLimitExceeded"))
 
+  private[ml] def previewMessage(previewRequest: HttpRequestBase): String = {
+    try {
+      previewRequest match {
+        case request: HttpPost =>
+          Option(request.getEntity).map { entity =>
+            Try {
+              val input = entity.getContent
+              try {
+                IOUtils.toString(input, "UTF-8")
+              } finally {
+                input.close()
+              }
+            }.getOrElse("")
+          }.getOrElse("")
+        case request =>
+          request.getURI.toString
+      }
+    } finally {
+      previewRequest.releaseConnection()
+    }
+  }
+
   //scalastyle:off cyclomatic.complexity
   //scalastyle:off method.length
   private[ml] def sendWithRetries(client: CloseableHttpClient,
@@ -381,12 +403,7 @@ object HandlingUtils extends SparkLogging {
   def advanced(retryTimes: Int*)(client: CloseableHttpClient,
                                  request: HTTPRequestData): HTTPResponseData = {
     try {
-      val previewRequest = request.toHTTPCore
-      val message = previewRequest match {
-        case r: HttpPost => Try(IOUtils.toString(r.getEntity.getContent, "UTF-8")).getOrElse("")
-        case r => r.getURI
-      }
-      previewRequest.releaseConnection()
+      val message = previewMessage(request.toHTTPCore)
       SynapseMLLogging.logDebug(s"sending $message")
       val start = System.currentTimeMillis()
       val usesTrustedFabricAuth = request.usesFabricAuth &&
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala
index 9211c965e0c..6b6dc08c13d 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala
@@ -6,11 +6,12 @@ package com.microsoft.azure.synapse.ml.io.split1
 import com.microsoft.azure.synapse.ml.core.test.base.TestBase
 import com.microsoft.azure.synapse.ml.io.http.HandlingUtils
 import org.apache.http.HttpVersion
-import org.apache.http.client.methods.CloseableHttpResponse
+import org.apache.http.client.methods.{CloseableHttpResponse, HttpPost}
 import org.apache.http.entity.BasicHttpEntity
 import org.apache.http.message.{BasicHttpResponse, BasicStatusLine}
 
-import java.io.{IOException, InputStream}
+import java.io.{ByteArrayInputStream, IOException, InputStream}
+import java.util.concurrent.atomic.AtomicBoolean
 import scala.io.Source
 
 class VerifyResponseBodyInspection extends TestBase {
@@ -70,4 +71,30 @@ class VerifyResponseBodyInspection extends TestBase {
 
     assert(replayedBody === new String(content, "UTF-8"))
   }
+
+  test("request preview closes its entity stream and releases the request") {
+    val content = """{"prompt":"hello"}""".getBytes("UTF-8")
+    val streamClosed = new AtomicBoolean(false)
+    val requestReleased = new AtomicBoolean(false)
+    val input = new ByteArrayInputStream(content) {
+      override def close(): Unit = {
+        streamClosed.set(true)
+        super.close()
+      }
+    }
+    val entity = new BasicHttpEntity()
+    entity.setContent(input)
+    entity.setContentLength(content.length)
+    val request = new HttpPost("https://example.test/openai") {
+      override def releaseConnection(): Unit = {
+        requestReleased.set(true)
+        super.releaseConnection()
+      }
+    }
+    request.setEntity(entity)
+
+    assert(HandlingUtils.previewMessage(request) === new String(content, "UTF-8"))
+    assert(streamClosed.get())
+    assert(requestReleased.get())
+  }
 }

From 538755a31c320b7d8b3077d4e55b53647cae41ba Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 18:22:15 -0700
Subject: [PATCH 06/10] fix: fall back across Fabric token cache APIs

Continue through broken reflection candidates and unsupported NFS path types so compatible cache invalidation mechanisms can still run.

AB#3582121

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../synapse/ml/fabric/TokenLibrary.scala      | 71 ++++++++++++-------
 .../ml/fabric/VerifyTokenInvalidation.scala   | 21 ++++++
 2 files changed, 67 insertions(+), 25 deletions(-)

diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
index 5dbc884d723..8d83e37fdbc 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
@@ -59,11 +59,21 @@ object TokenLibrary {
       .asInstanceOf[String]
   }
 
-  private def objectMethod(classNames: Seq[String],
-                           methodName: String,
-                           parameterCount: Int): Option[(AnyRef, java.lang.reflect.Method)] = {
+  private def reflectionOrElse[T](fallback: => T)(operation: => T): T = {
+    try {
+      operation
+    } catch {
+      case _: ReflectiveOperationException | _: SecurityException |
+           _: LinkageError | _: IllegalArgumentException =>
+        fallback
+    }
+  }
+
+  private[ml] def objectMethod(classNames: Seq[String],
+                               methodName: String,
+                               parameterCount: Int): Option[(AnyRef, java.lang.reflect.Method)] = {
     classNames.iterator.flatMap { className =>
-      try {
+      reflectionOrElse(Option.empty[(AnyRef, java.lang.reflect.Method)]) {
         val cls = Class.forName(className)
         val module = cls.getField("MODULE$").get(null).asInstanceOf[AnyRef] //scalastyle:ignore null
         (cls.getMethods ++ cls.getDeclaredMethods)
@@ -72,49 +82,60 @@ object TokenLibrary {
             method.setAccessible(true)
             module -> method
           }
-      } catch {
-        case _: ClassNotFoundException | _: NoSuchFieldException => None
       }
     }.take(1).toSeq.headOption
   }
 
   private def invalidateWithRuntimeApi(workspaceId: String, artifactId: String): Boolean = {
     objectMethod(Seq(TokenLibraryClass), "invalidateMwcToken", 4).exists { case (module, method) =>
-      method.invoke(
-        module,
-        workspaceId,
-        artifactId,
-        Int.box(SparkTokenVersion),
-        SparkWorkloadType)
-      true
+      reflectionOrElse(false) {
+        method.invoke(
+          module,
+          workspaceId,
+          artifactId,
+          Int.box(SparkTokenVersion),
+          SparkWorkloadType)
+        true
+      }
     }
   }
 
   private def nfsCacheKey(cacheKey: String): String = {
     objectMethod(NfsCacheClasses, "getNFSCacheKey", 1)
-      .map { case (module, method) => method.invoke(module, cacheKey).asInstanceOf[String] }
+      .flatMap { case (module, method) =>
+        reflectionOrElse(Option.empty[String]) {
+          method.invoke(module, cacheKey) match {
+            case resolved: String => Some(resolved)
+            case _ => None
+          }
+        }
+      }
       .getOrElse(cacheKey)
   }
 
   private def deleteNfsToken(resolvedCacheKey: String): Boolean = {
     objectMethod(NfsCacheClasses, "getNFSTokenFilePath", 1).exists { case (module, method) =>
-      val tokenPath = method.invoke(module, resolvedCacheKey) match {
-        case path: Path => path
-        case file: File => file.toPath
-        case path: String => Paths.get(path)
-        case other =>
-          throw new IllegalStateException(
-            s"Unsupported Fabric token cache path type: ${Option(other).map(_.getClass.getName).orNull}")
+      reflectionOrElse(false) {
+        val tokenPath = method.invoke(module, resolvedCacheKey) match {
+          case path: Path => Some(path)
+          case file: File => Some(file.toPath)
+          case path: String => Some(Paths.get(path))
+          case _ => None
+        }
+        tokenPath.exists { path =>
+          Files.deleteIfExists(path)
+          true
+        }
       }
-      Files.deleteIfExists(tokenPath)
-      true
     }
   }
 
   private def clearInMemoryTokenCache(): Boolean = {
     objectMethod(InMemoryCacheClasses, "clear", 0).exists { case (module, method) =>
-      method.invoke(module)
-      true
+      reflectionOrElse(false) {
+        method.invoke(module)
+        true
+      }
     }
   }
 
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
index 23c48089853..70ff4ac5a8f 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
@@ -9,6 +9,16 @@ import java.nio.file.Files
 import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
 import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit}
 
+object BrokenReflectionCandidate {
+  val Initialize: Unit = throw new RuntimeException("simulated initialization failure")
+
+  def clear(): Unit = ()
+}
+
+object WorkingReflectionCandidate {
+  def clear(): Unit = ()
+}
+
 class VerifyTokenInvalidation extends TestBase {
 
   test("Spark MWC invalidation resolves the encoded NFS cache key") {
@@ -54,6 +64,17 @@ class VerifyTokenInvalidation extends TestBase {
     assert(error.getMessage.contains("does not expose MWC token cache invalidation"))
   }
 
+  test("runtime reflection skips a broken candidate and uses the next compatible class") {
+    val method = TokenLibrary.objectMethod(
+      Seq(
+        "com.microsoft.azure.synapse.ml.fabric.BrokenReflectionCandidate$",
+        "com.microsoft.azure.synapse.ml.fabric.WorkingReflectionCandidate$"),
+      "clear",
+      0)
+
+    assert(method.exists(_._1.getClass.getName.endsWith("WorkingReflectionCandidate$")))
+  }
+
   test("concurrent refreshes invalidate a rejected MWC token once") {
     val workers = 8
     val executor = Executors.newFixedThreadPool(workers)

From b11ea80051a191d83095dff9c4d231ce8e275670 Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 18:42:46 -0700
Subject: [PATCH 07/10] test: isolate Fabric auth retry coverage

Move auth provenance coverage into a focused test file so the master patch replays cleanly onto the spark4.1 release branch.

AB#3582121

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../openai/OpenAIFabricAuthRetrySuite.scala   | 61 +++++++++++++++++++
 .../openai/OpenAIFabricHeadersSuite.scala     | 51 ----------------
 2 files changed, 61 insertions(+), 51 deletions(-)
 create mode 100644 cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricAuthRetrySuite.scala

diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricAuthRetrySuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricAuthRetrySuite.scala
new file mode 100644
index 00000000000..5c56ef71d6a
--- /dev/null
+++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricAuthRetrySuite.scala
@@ -0,0 +1,61 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.services.openai
+
+import com.microsoft.azure.synapse.ml.core.test.base.TestBase
+import com.microsoft.azure.synapse.ml.io.http.HTTPRequestData
+import org.apache.http.client.methods.HttpPost
+import org.apache.spark.sql.Row
+
+class OpenAIFabricAuthRetrySuite extends TestBase {
+
+  private class AuthProvenanceProbe(fallbackAuthHeader: Option[String]) extends OpenAIChatCompletion {
+    var fallbackCalls = 0
+
+    override protected[openai] def runningOnFabric: Boolean = true
+
+    override protected[openai] def usingDefaultOpenAIEndpoint: Boolean = true
+
+    override protected def getFabricFallbackAuthHeader(row: Row): Option[String] = {
+      fallbackCalls += 1
+      fallbackAuthHeader
+    }
+
+    def requestData: HTTPRequestData = {
+      val request = new HttpPost("https://example.test/openai")
+      addHeaders(request, Row.empty, addContentType = false)
+      new HTTPRequestData(request)
+    }
+  }
+
+  test("only implicit Fabric authentication marks a request for refresh") {
+    val implicitProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+    val implicitRequestData = implicitProbe.requestData
+    val implicitHttpRequest = implicitRequestData.toHTTPCore
+
+    assert(implicitProbe.fallbackCalls === 1)
+    assert(implicitRequestData.usesFabricAuth)
+    assert(implicitHttpRequest.getFirstHeader("Authorization").getValue === "MwcToken implicit")
+    assert(Option(implicitHttpRequest.getFirstHeader(HTTPRequestData.FabricAuthMarkerHeader)).isEmpty)
+
+    val explicitProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+      .setCustomAuthHeader("MwcToken explicit")
+      .setCustomHeaders(Map(HTTPRequestData.FabricAuthMarkerHeader -> "true"))
+      .asInstanceOf[AuthProvenanceProbe]
+    val explicitRequestData = explicitProbe.requestData
+
+    assert(explicitProbe.fallbackCalls === 0)
+    assert(!explicitRequestData.usesFabricAuth)
+    assert(explicitRequestData.toHTTPCore.getFirstHeader("Authorization").getValue === "MwcToken explicit")
+
+    val apiKeyProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+      .setSubscriptionKey("explicit-key")
+      .asInstanceOf[AuthProvenanceProbe]
+    val apiKeyRequestData = apiKeyProbe.requestData
+
+    assert(apiKeyProbe.fallbackCalls === 0)
+    assert(!apiKeyRequestData.usesFabricAuth)
+    assert(apiKeyRequestData.toHTTPCore.getFirstHeader("api-key").getValue === "explicit-key")
+  }
+}
diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala
index 8cd94c1c8c2..35f60b2c2dc 100644
--- a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala
+++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala
@@ -4,10 +4,8 @@
 package com.microsoft.azure.synapse.ml.services.openai
 
 import com.microsoft.azure.synapse.ml.core.test.base.TestBase
-import com.microsoft.azure.synapse.ml.io.http.HTTPRequestData
 import com.microsoft.azure.synapse.ml.logging.common.PlatformDetails
 import com.microsoft.azure.synapse.ml.services.HasCognitiveServiceInput
-import org.apache.http.client.methods.HttpPost
 import org.apache.spark.sql.Row
 import spray.json._
 
@@ -60,25 +58,6 @@ class OpenAIFabricHeadersSuite extends TestBase {
       override protected val usesDefaultEndpoint: Boolean)
     extends OpenAIResponses with InspectableFabricHeaders
 
-  private class AuthProvenanceProbe(fallbackAuthHeader: Option[String]) extends OpenAIChatCompletion {
-    var fallbackCalls = 0
-
-    override protected[openai] def runningOnFabric: Boolean = true
-
-    override protected[openai] def usingDefaultOpenAIEndpoint: Boolean = true
-
-    override protected def getFabricFallbackAuthHeader(row: Row): Option[String] = {
-      fallbackCalls += 1
-      fallbackAuthHeader
-    }
-
-    def requestData: HTTPRequestData = {
-      val request = new HttpPost("https://example.test/openai")
-      addHeaders(request, Row.empty, addContentType = false)
-      new HTTPRequestData(request)
-    }
-  }
-
   private def transformers(
       isFabric: Boolean,
       usesDefaultEndpoint: Boolean): Seq[InspectableFabricHeaders] = {
@@ -181,34 +160,4 @@ class OpenAIFabricHeadersSuite extends TestBase {
     assert(headers.values.forall(_ != null))
     assert(!headers.contains("Other"))
   }
-
-  test("only implicit Fabric authentication marks a request for refresh") {
-    val implicitProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
-    val implicitRequestData = implicitProbe.requestData
-    val implicitHttpRequest = implicitRequestData.toHTTPCore
-
-    assert(implicitProbe.fallbackCalls === 1)
-    assert(implicitRequestData.usesFabricAuth)
-    assert(implicitHttpRequest.getFirstHeader("Authorization").getValue === "MwcToken implicit")
-    assert(Option(implicitHttpRequest.getFirstHeader(HTTPRequestData.FabricAuthMarkerHeader)).isEmpty)
-
-    val explicitProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
-      .setCustomAuthHeader("MwcToken explicit")
-      .setCustomHeaders(Map(HTTPRequestData.FabricAuthMarkerHeader -> "true"))
-      .asInstanceOf[AuthProvenanceProbe]
-    val explicitRequestData = explicitProbe.requestData
-
-    assert(explicitProbe.fallbackCalls === 0)
-    assert(!explicitRequestData.usesFabricAuth)
-    assert(explicitRequestData.toHTTPCore.getFirstHeader("Authorization").getValue === "MwcToken explicit")
-
-    val apiKeyProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
-      .setSubscriptionKey("explicit-key")
-      .asInstanceOf[AuthProvenanceProbe]
-    val apiKeyRequestData = apiKeyProbe.requestData
-
-    assert(apiKeyProbe.fallbackCalls === 0)
-    assert(!apiKeyRequestData.usesFabricAuth)
-    assert(apiKeyRequestData.toHTTPCore.getFirstHeader("api-key").getValue === "explicit-key")
-  }
 }

From dde47fe1c67a4def98e21dfd934c75462e7f7ef8 Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 18:50:57 -0700
Subject: [PATCH 08/10] fix: make Fabric cache invalidation best effort

Avoid inaccessible reflective methods, tolerate supported reflection fallback failures, and continue past NFS deletion errors so in-memory invalidation can still run.

AB#3582121

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../azure/synapse/ml/fabric/TokenLibrary.scala | 18 +++++++++++++-----
 .../ml/fabric/VerifyTokenInvalidation.scala    | 14 ++++++++++++++
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
index 8d83e37fdbc..4b60bb74b7c 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala
@@ -66,6 +66,9 @@ object TokenLibrary {
       case _: ReflectiveOperationException | _: SecurityException |
            _: LinkageError | _: IllegalArgumentException =>
         fallback
+      case error: RuntimeException
+          if error.getClass.getName == "java.lang.reflect.InaccessibleObjectException" =>
+        fallback
     }
   }
 
@@ -79,7 +82,6 @@ object TokenLibrary {
         (cls.getMethods ++ cls.getDeclaredMethods)
           .find(method => method.getName == methodName && method.getParameterCount == parameterCount)
           .map { method =>
-            method.setAccessible(true)
             module -> method
           }
       }
@@ -122,14 +124,20 @@ object TokenLibrary {
           case path: String => Some(Paths.get(path))
           case _ => None
         }
-        tokenPath.exists { path =>
-          Files.deleteIfExists(path)
-          true
-        }
+        tokenPath.exists(deleteTokenPath)
       }
     }
   }
 
+  private[ml] def deleteTokenPath(path: Path): Boolean = {
+    try {
+      Files.deleteIfExists(path)
+      true
+    } catch {
+      case _: java.io.IOException | _: SecurityException => false
+    }
+  }
+
   private def clearInMemoryTokenCache(): Boolean = {
     objectMethod(InMemoryCacheClasses, "clear", 0).exists { case (module, method) =>
       reflectionOrElse(false) {
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
index 70ff4ac5a8f..46f93251f26 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala
@@ -75,6 +75,20 @@ class VerifyTokenInvalidation extends TestBase {
     assert(method.exists(_._1.getClass.getName.endsWith("WorkingReflectionCandidate$")))
   }
 
+  test("NFS token deletion reports filesystem failures without throwing") {
+    val directory = Files.createTempDirectory("synapseml-mwc-cache")
+    val child = Files.createFile(directory.resolve("token"))
+
+    try {
+      assert(!TokenLibrary.deleteTokenPath(directory))
+      assert(Files.exists(directory))
+      assert(Files.exists(child))
+    } finally {
+      Files.deleteIfExists(child)
+      Files.deleteIfExists(directory)
+    }
+  }
+
   test("concurrent refreshes invalidate a rejected MWC token once") {
     val workers = 8
     val executor = Executors.newFixedThreadPool(workers)

From 2ea0d9ffd53f5ccab8172224271c63af5fe05172 Mon Sep 17 00:00:00 2001
From: Ranadeep Singh <ranadeep.dtu@gmail.com>
Date: Tue, 1 Sep 2026 18:59:13 -0700
Subject: [PATCH 09/10] perf: defer HTTP request previews to debug logging

Keep preview body reconstruction lazy so normal request execution does not read and allocate a duplicate POST body solely for disabled debug logging.

AB#3582121

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
---
 .../com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala   | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
index bc1286ce92e..6cfd2b4245b 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
@@ -403,8 +403,7 @@ object HandlingUtils extends SparkLogging {
   def advanced(retryTimes: Int*)(client: CloseableHttpClient,
                                  request: HTTPRequestData): HTTPResponseData = {
     try {
-      val message = previewMessage(request.toHTTPCore)
-      SynapseMLLogging.logDebug(s"sending $message")
+      SynapseMLLogging.logDebug(s"sending ${previewMessage(request.toHTTPCore)}")
       val start = System.currentTimeMillis()
       val usesTrustedFabricAuth = request.usesFabricAuth &&
         FabricClient.isOpenAIEndpoint(request.requestLine.uri)

From 504c399c5a10b2aea0db69d6c7961b46c5a4ed54 Mon Sep 17 00:00:00 2001
From: Rana Singh <ranadeep.dtu@gmail.com>
Date: Thu, 3 Sep 2026 20:23:08 -0700
Subject: [PATCH 10/10] fix: preserve bounded inspection for large retry
 responses

## Summary

Inspect a bounded prefix of large HTTP response bodies while replaying every consumed byte, so capacity-limit errors remain detectable without consuming caller-visible content. Treat null Authorization header values as absent.

## Prompting Intent

The engineer asked to rebase and safely validate every open pull request before authorizing Azure Pipelines. Copilot review of the rebased head found that oversized 429 bodies could bypass CapacityLimitExceeded detection and that a null header could escape through Option[String]; this change resolves both findings with regression coverage.

## Linked Sources

- Large-response finding: https://github.com/microsoft/SynapseML/pull/2685#discussion_r3930569863
- Null-authorization finding: https://github.com/microsoft/SynapseML/pull/2685#discussion_r3930569931
- Pull request: https://github.com/microsoft/SynapseML/pull/2685

## Rationale

Reading at most one MiB plus one byte keeps inspection memory bounded. Replaying the buffered prefix before the original stream preserves the complete response, while returning the decoded prefix is sufficient to detect the ASCII capacity error code. Null filtering keeps Option semantics honest and fails Fabric-auth detection closed.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 476d113f-dd35-40c6-bc79-005dcccd7b79
---
 .../synapse/ml/io/http/HTTPClients.scala      | 70 +++++++++----------
 .../azure/synapse/ml/io/http/HTTPSchema.scala |  2 +-
 .../split1/VerifyResponseBodyInspection.scala | 28 ++++++++
 .../ml/io/split1/VerifySendWithRetries.scala  | 14 +++-
 4 files changed, 75 insertions(+), 39 deletions(-)

diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
index 6cfd2b4245b..c0e0cacf56a 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala
@@ -122,45 +122,41 @@ object HandlingUtils extends SparkLogging {
 
   private[ml] def responseBodyForInspection(response: CloseableHttpResponse): Option[String] = {
     Option(response.getEntity).flatMap { entity =>
-      if (entity.getContentLength > MaxResponseInspectionBytes) {
-        None
-      } else {
-        val output = new ByteArrayOutputStream()
-        var input = Option.empty[InputStream]
-        var keepInputOpen = false
-        try {
-          val responseInput = entity.getContent
-          input = Some(responseInput)
-          IOUtils.copyLarge(responseInput, output, 0, MaxResponseInspectionBytes + 1)
-          val bytes = output.toByteArray
-          if (bytes.length > MaxResponseInspectionBytes) {
-            replayResponseEntity(response, entity, bytes, responseInput)
-            keepInputOpen = true
-            None
-          } else {
-            val bufferedEntity = new ByteArrayEntity(bytes)
-            copyResponseEntityMetadata(entity, bufferedEntity)
-            response.setEntity(bufferedEntity)
-            Some(new String(bytes, StandardCharsets.UTF_8))
-          }
-        } catch {
-          case NonFatal(error) =>
-            logWarning("Could not inspect the HTTP response body; preserving it for retry handling.", error)
-            input.foreach { responseInput =>
-              try {
-                replayResponseEntity(response, entity, output.toByteArray, responseInput)
-                keepInputOpen = true
-              } catch {
-                case NonFatal(replayError) =>
-                  error.addSuppressed(replayError)
-                  logWarning("Could not reconstruct the partially inspected HTTP response body.", replayError)
-              }
+      val output = new ByteArrayOutputStream()
+      var input = Option.empty[InputStream]
+      var keepInputOpen = false
+      try {
+        val responseInput = entity.getContent
+        input = Some(responseInput)
+        IOUtils.copyLarge(responseInput, output, 0, MaxResponseInspectionBytes + 1)
+        val bytes = output.toByteArray
+        if (bytes.length > MaxResponseInspectionBytes) {
+          replayResponseEntity(response, entity, bytes, responseInput)
+          keepInputOpen = true
+          Some(new String(bytes, 0, MaxResponseInspectionBytes.toInt, StandardCharsets.UTF_8))
+        } else {
+          val bufferedEntity = new ByteArrayEntity(bytes)
+          copyResponseEntityMetadata(entity, bufferedEntity)
+          response.setEntity(bufferedEntity)
+          Some(new String(bytes, StandardCharsets.UTF_8))
+        }
+      } catch {
+        case NonFatal(error) =>
+          logWarning("Could not inspect the HTTP response body; preserving it for retry handling.", error)
+          input.foreach { responseInput =>
+            try {
+              replayResponseEntity(response, entity, output.toByteArray, responseInput)
+              keepInputOpen = true
+            } catch {
+              case NonFatal(replayError) =>
+                error.addSuppressed(replayError)
+                logWarning("Could not reconstruct the partially inspected HTTP response body.", replayError)
             }
-            None
-        } finally {
-          if (!keepInputOpen) {
-            input.foreach(closeInspectionInput)
           }
+          None
+      } finally {
+        if (!keepInputOpen) {
+          input.foreach(closeInspectionInput)
         }
       }
     }
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
index f6cb652e502..2dfa4a03d47 100644
--- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
+++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
@@ -173,7 +173,7 @@ case class HTTPRequestData(requestLine: RequestLineData,
       authorizationHeader.exists(HTTPRequestData.isCognitiveMwcAuthHeader)
 
   private[ml] def authorizationHeader: Option[String] =
-    headers.find(h => "Authorization".equalsIgnoreCase(h.name)).map(_.value)
+    headers.find(h => "Authorization".equalsIgnoreCase(h.name)).flatMap(h => Option(h.value))
 
   def this(r: HttpRequestBase) = {
     this(new RequestLineData(r.getRequestLine),
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala
index 6b6dc08c13d..7626e6a5fad 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala
@@ -72,6 +72,34 @@ class VerifyResponseBodyInspection extends TestBase {
     assert(replayedBody === new String(content, "UTF-8"))
   }
 
+  test("response inspection detects capacity errors in large bodies without consuming them") {
+    val content = ("""{"error":{"code":"CapacityLimitExceeded"}}""" +
+      ("x" * (1024 * 1024 + 128))).getBytes("UTF-8")
+
+    Seq(-1L, content.length.toLong).foreach { contentLength =>
+      val entity = new BasicHttpEntity()
+      entity.setContent(new ByteArrayInputStream(content))
+      entity.setContentLength(contentLength)
+      val response = new BasicHttpResponse(
+        new BasicStatusLine(HttpVersion.HTTP_1_1, 429, "Too Many Requests"))
+        with CloseableHttpResponse {
+        override def close(): Unit = ()
+      }
+      response.setEntity(entity)
+
+      val inspected = HandlingUtils.responseBodyForInspection(response)
+      assert(inspected.exists(_.contains("CapacityLimitExceeded")))
+      val replayed = Source.fromInputStream(response.getEntity.getContent, "UTF-8")
+      val replayedBody = try {
+        replayed.mkString
+      } finally {
+        replayed.close()
+      }
+
+      assert(replayedBody === new String(content, "UTF-8"))
+    }
+  }
+
   test("request preview closes its entity stream and releases the request") {
     val content = """{"prompt":"hello"}""".getBytes("UTF-8")
     val streamClosed = new AtomicBoolean(false)
diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
index 1f209065cff..6105ac59de8 100644
--- a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
+++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala
@@ -4,7 +4,7 @@
 package com.microsoft.azure.synapse.ml.io.split1
 
 import com.microsoft.azure.synapse.ml.fabric.FabricClient
-import com.microsoft.azure.synapse.ml.io.http.{HTTPRequestData, HandlingUtils}
+import com.microsoft.azure.synapse.ml.io.http.{HTTPRequestData, HandlingUtils, HeaderData, RequestLineData}
 
 import com.microsoft.azure.synapse.ml.core.test.base.TestBase
 import com.sun.net.httpserver.{HttpExchange, HttpServer}
@@ -683,6 +683,18 @@ class VerifySendWithRetries extends TestBase {
     assert(!requestData("true", "MwcToken ").usesFabricAuth)
   }
 
+  test("null authorization values are treated as absent") {
+    val requestData = HTTPRequestData(
+      RequestLineData("GET", "https://workspace.fabric.microsoft.com/cognitive/openai/chat", None),
+      Array(
+        HeaderData(HTTPRequestData.FabricAuthMarkerHeader, "true"),
+        HeaderData("Authorization", null)), //scalastyle:ignore null
+      None)
+
+    assert(requestData.authorizationHeader.isEmpty)
+    assert(!requestData.usesFabricAuth)
+  }
+
   test("Fabric auth retries replace duplicate authorization headers") {
     val port = getFreePort
     val authorizationHeaders = new AtomicReference[Seq[String]]()