feat: add opt-in recovery for AIPrompt results

codecov.yaml

@@ -3,9 +3,9 @@ codecov:
     require_ci_to_pass: no
     # Wait for expected number of coverage uploads before sending notifications
     # This prevents premature comments with incomplete data.
-    # 54 = UnitTests 40 + PythonTests 7 + RTests 6 + WebsiteSamplesTests 1
+    # 55 = UnitTests 41 + PythonTests 7 + RTests 6 + WebsiteSamplesTests 1
     # (every leg of those jobs runs templates/codecov.yml on succeededOrFailed).
-    after_n_builds: 54
+    after_n_builds: 55
 
 comment:
   layout: "reach,diff,flags,files"
@@ -14,7 +14,7 @@ comment:
   require_base: false
   require_head: true
   # Wait for all expected uploads before commenting (see note above)
-  after_n_builds: 54
+  after_n_builds: 55
 
 coverage:
   precision: 2

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala

@@ -329,7 +329,9 @@ object URLEncodingUtils {
 }
 
 private[ml] object ServiceAuthHeaders {
-  private[ml] case class Resolution(headers: Map[String, String], usesFabricFallback: Boolean)
+  private[ml] case class Resolution(headers: Map[String, String],
+                                    usesFabricFallback: Boolean,
+                                    usesSubscriptionKey: Boolean)
 
   private[ml] def nonBlank(value: String): Boolean = value != null && value.trim.nonEmpty
 
@@ -377,20 +379,15 @@ private[ml] object ServiceAuthHeaders {
     def isAuthHeaderName(name: String): Boolean =
       name.equalsIgnoreCase(subscriptionKeyHeaderName) || name.equalsIgnoreCase(aadHeaderName)
 
-    // Resolve exactly one auth header across every credential source in a single, case-insensitive
-    // precedence step. Blank values are skipped so they cannot suppress a valid lower-priority
-    // credential, and the automatic Fabric fallback ranks below a credential embedded in
-    // customHeaders (matching how management-index requests, which have no fallback, resolve auth).
-    // fabricFallbackAuthHeader is by-name and lowest priority, so it is evaluated only when every
-    // 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 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))
+    // Select one non-blank credential and retain its source even when the auth header names coincide.
+    // The by-name Fabric fallback is evaluated only when no explicit or embedded credential exists.
+    val explicitAuth: Option[((String, String), Boolean)] = subscriptionKey.filter(nonBlank)
+      .map(value => (subscriptionKeyHeaderName -> value) -> true)
+      .orElse(aadToken.filter(nonBlank).map(value => (aadHeaderName -> ("Bearer " + value)) -> false))
+      .orElse(customAuthHeader.filter(nonBlank).map(value => (aadHeaderName -> value) -> false))
+      .orElse(embeddedCredential(providedCustomHeaders, subscriptionKeyHeaderName).map(_ -> true))
+      .orElse(embeddedCredential(providedCustomHeaders, aadHeaderName).map(_ -> false))
+    val explicitAuthHeader = explicitAuth.map(_._1)
     val fabricAuthHeader = if (explicitAuthHeader.isDefined) {
       None
     } else {
@@ -421,7 +418,8 @@ private[ml] object ServiceAuthHeaders {
 
     Resolution(
       new scala.collection.immutable.TreeMap[String, String]() ++ headers,
-      usesFabricFallback = fabricAuthHeader.isDefined)
+      usesFabricFallback = fabricAuthHeader.isDefined,
+      usesSubscriptionKey = explicitAuth.exists(_._2))
   }
 
   def build(subscriptionKey: Option[String],

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIColumnPlan.scala

@@ -0,0 +1,83 @@
+// 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.ai
+
+import org.apache.spark.sql.{Column, Row, functions => F}
+import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructField, StructType}
+
+import java.util.{HashSet, Set => JavaSet, TreeSet}
+
+private[ai] final case class AIColumnPlan(input: StructType, outputs: Seq[StructField]) {
+  AIColumnPlan.distinct(input.fieldNames.toSeq, "input columns")
+  AIColumnPlan.distinct(outputs.map(_.name), "output roles")
+  require(outputs.forall(_.name.nonEmpty), "output roles must have non-empty names")
+  private val outputNames = outputs.map(_.name).toSet
+  private val replacements = input.fields.map(f =>
+    outputs.find(o => AIColumnPlan.matches(f.name, o.name)))
+  val schema: StructType = StructType(input.fields.zip(replacements).map {
+    case (original, replacement) => replacement.getOrElse(original)
+  } ++ outputs.filterNot(o => input.fields.exists(f => AIColumnPlan.matches(f.name, o.name))))
+
+  private val indices = schema.fields.map(f =>
+    if (outputNames(f.name)) -1 else input.fieldIndex(f.name))
+
+  def assemble(row: Row, values: Map[String, Any]): Row =
+    Row.fromSeq(schema.fields.zip(indices).map { case (field, index) =>
+      if (index < 0) values.getOrElse(field.name, None.orNull) else row.get(index)
+    })
+
+  def scratch(prefix: String, additional: String*): String =
+    AIColumnPlan.unused(prefix, schema.fieldNames.toSeq ++ input.fieldNames ++ additional)
+}
+
+private[ai] object AIColumnPlan {
+  def matches(left: String, right: String): Boolean = SQLConf.get.resolver(left, right)
+
+  def distinct(names: Seq[String], description: String): Unit = {
+    require(names.forall(_ != null), s"$description must not have null names")
+    // This comparator matches Spark's equalsIgnoreCase, including Unicode variants lost by lowercasing.
+    val seen: JavaSet[String] = if (SQLConf.get.caseSensitiveAnalysis) new HashSet[String]()
+    else new TreeSet[String](String.CASE_INSENSITIVE_ORDER)
+    require(!names.exists(name => !seen.add(name)),
+      s"Ambiguous $description under Spark's column resolver")
+  }
+
+  def resolve(schema: StructType, name: String): String =
+    schema.fieldNames.find(matches(_, name)).getOrElse(
+      throw new IllegalArgumentException(s"Input column '$name' does not exist"))
+
+  def unused(prefix: String, names: Seq[String]): String =
+    (Iterator(prefix) ++ Iterator.from(1).map(i => s"${prefix}_$i"))
+      .find(candidate => !names.exists(matches(_, candidate))).get
+
+  def column(name: String): Column = F.col("`" + name.replace("`", "``") + "`")
+
+  // Spark can eliminate serialization between Dataset maps and retain external JVM arrays.
+  def arrayValues(value: Any): scala.collection.Seq[Any] = value match {
+    case values: scala.collection.Seq[_] => values
+    case values: Array[_] => values.toSeq
+    case _ => throw new IllegalArgumentException("Expected a Spark array value")
+  }
+
+  // Fused Dataset maps can pass schema-less Rows directly to the protocol serializers.
+  def withSchema(row: Row, schema: StructType): Row =
+    new GenericRowWithSchema(schema.fields.zipWithIndex.map { case (field, index) =>
+      nestedSchema(row.get(index), field.dataType)
+    }, schema)
+
+  private def nestedSchema(value: Any, dataType: DataType): Any = Option(value).map { present =>
+    dataType match {
+      case struct: StructType => withSchema(present.asInstanceOf[Row], struct)
+      case ArrayType(element, _) =>
+        arrayValues(present).iterator.map(nestedSchema(_, element)).toVector
+      case MapType(key, element, _) =>
+        present.asInstanceOf[scala.collection.Map[Any, Any]].iterator.map { case (k, v) =>
+          nestedSchema(k, key) -> nestedSchema(v, element)
+        }.toMap
+      case _ => present
+    }
+  }.orNull
+}

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIDurableInput.scala

@@ -0,0 +1,194 @@
+// 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.ai
+
+import com.microsoft.azure.synapse.ml.services.openai.OpenAIAttachmentUtils
+import org.apache.spark.sql.{DataFrame, Row, functions => F}
+import org.apache.spark.sql.types.{StringType, StructType}
+import spray.json._
+
+private[ai] object AIDurableInput {
+  private val ExternalBindings = Set(
+    "previous_response_id", "conversation", "conversation_id", "file_id", "file_ids",
+    "tools", "tool_choice", "tool_resources", "functions", "function_call", "retrieve", "run_id",
+    "data_sources", "web_search_options", "prompt", "tool_call_id")
+  private val OutputSchemas = Map(
+    "chat_completions" -> Set(List("response_format", "json_schema", "schema")),
+    "responses" -> Set(List("text", "format", "schema")))
+
+  private def externalBindings(value: JsValue, path: List[String] = Nil,
+                               outputSchemas: Set[List[String]] = Set.empty): Unit = value match {
+    case JsObject(fields) if !outputSchemas(path) =>
+      val bindings = fields.iterator.collect {
+        case (name, entry) if ExternalBindings(name) && entry != JsNull => name
+      }.toSeq.sorted
+      require(bindings.isEmpty,
+        "Durable requests cannot depend on provider state, retrieval, tools, or external resources; " +
+          s"restricted fields: ${bindings.mkString(", ")}")
+      content(value)
+      fields.foreach { case (name, entry) => externalBindings(entry, path :+ name, outputSchemas) }
+    case JsArray(items) => items.foreach(externalBindings(_))
+    case _ =>
+  }
+
+  private def body(value: JsValue, api: String): Unit =
+    externalBindings(JsObject(AIJson.obj(value).fields - "metadata"),
+      outputSchemas = OutputSchemas.getOrElse(api, Set.empty))
+
+  private def responseFormatBody(value: Any, api: String): JsObject = {
+    val format = AIRequestAdapter.responseFormat(value, api)
+    if (api == "responses") JsObject("text" -> JsObject("format" -> format))
+    else JsObject("response_format" -> format)
+  }
+
+  private def questions(api: String, value: JsValue): Unit = (api, value) match {
+    case ("systemone", JsObject(named)) =>
+      named.values.foreach {
+        // Question names and criteria labels are user-defined text, not provider binding fields.
+        case JsObject(fields) =>
+          externalBindings(JsObject(fields - "criteria"))
+          fields.get("criteria").foreach {
+            case JsObject(criteria) => criteria.values.foreach(externalBindings(_))
+            case other => externalBindings(other)
+          }
+        case other => externalBindings(other)
+      }
+    case _ => externalBindings(value)
+  }
+
+  def validate(stage: AIPrompt, schema: StructType): Unit = {
+    val api = stage.getApiType
+    Seq(stage.subscriptionKey, stage.AADToken, stage.CustomAuthHeader,
+      stage.customHeaders, stage.telemHeaders).foreach { p =>
+      require(!stage.get(p).orElse(stage.getDefault(p)).exists(_.isRight),
+        s"Durable mode rejects ${p.name}Col; rebind scalar credentials on resume")
+    }
+    require(!stage.isDefined(stage.previousResponseId),
+      "Durable mode does not support previousResponseId or previousResponseIdCol")
+    stage.get(stage.extraBody).orElse(stage.getDefault(stage.extraBody)).foreach {
+      case Left(value) => body(AIJson.fromAny(value), api)
+      case _ =>
+    }
+    stage.get(stage.responseFormat).orElse(stage.getDefault(stage.responseFormat)).foreach {
+      case Left(value) => body(responseFormatBody(value, api), api)
+      case _ =>
+    }
+    stage.get(stage.questions).orElse(stage.getDefault(stage.questions)).foreach {
+      case Left(value) => questions(api, AIJson.fromAny(value))
+      case _ =>
+    }
+    stage.getColumnTypes.foreach { case (name, _) => AIColumnPlan.resolve(schema, name) }
+  }
+
+  private def inline(value: JsValue, literalContent: Boolean = false): Unit = value match {
+    case JsString(url) if !literalContent && url.trim.nonEmpty =>
+      // Request validation handles malformed inline media as row errors; only external dependencies abort here.
+      require(OpenAIAttachmentUtils.isDataUrl(url.trim),
+        "Durable media must be inline base64 data URLs, not remote content or paths")
+    case JsArray(items) => items.foreach(inline(_))
+    case JsObject(fields) => fields.get("url").foreach(inline(_))
+    case _ =>
+  }
+
+  private def content(part: JsValue): Unit = part match {
+    case JsObject(fields) =>
+      Seq("image_url", "file_data", "file_url").foreach { name =>
+        // A string file_data is literal content, not a location to resolve.
+        fields.get(name).foreach(inline(_, literalContent = name == "file_data"))
+      }
+    case _ =>
+  }
+
+  private def messages(value: JsValue): Unit = {
+    externalBindings(value)
+    value match {
+      case JsArray(items) => items.foreach {
+        case JsObject(fields) =>
+          require(!fields.get("id").exists(_ != JsNull) && !fields.get("type").exists {
+            case JsString("item_reference" | "function_call" | "function_call_output") => true
+            case _ => false
+          }, "Durable messages cannot reference provider items or tool results")
+          fields.get("content").foreach {
+            case JsArray(parts) => parts.foreach(content)
+            case _ =>
+          }
+        case _ =>
+      }
+      case _ =>
+    }
+  }
+
+  private def validateDeferredJson(kind: String, text: String): Unit = {
+    val parsed = try Some(AIJson.parse(text)) catch {
+      case _: IllegalArgumentException => None
+    }
+    // Inference parses these cells again and commits malformed values as terminal row errors.
+    parsed.foreach { value =>
+      if (kind.endsWith("_questions")) questions(kind.stripSuffix("_questions"), value)
+      else if (kind == "systemone") externalBindings(value)
+      else messages(value)
+    }
+  }
+
+  private def validateBodyCell(kind: String, text: String, api: String): Unit = {
+    // Malformed bodies and formats become terminal row errors; resource checks must not be deferred.
+    val parsed = try Some(if (kind == "body") AIJson.objectString(text)
+      else responseFormatBody(text, kind.stripSuffix("_format"))) catch {
+      case _: IllegalArgumentException => None
+    }
+    parsed.foreach(body(_, api))
+  }
+
+  private def validateCell(kind: String, text: String, api: String): Unit = kind match {
+    case "prompt" => () // Selecting the rendered value forces evaluation before any inference requests.
+    case "inline" =>
+      val normalized = text.trim
+      if (normalized.nonEmpty) inline(JsString(normalized))
+    case "body" | "chat_completions_format" | "responses_format" => validateBodyCell(kind, text, api)
+    // Inference encodes message Rows directly, so parser failures cannot be deferred to row errors.
+    case "messages" => messages(AIJson.parse(text))
+    case "systemone" | "decisions" | "systemone_questions" | "decisions_questions" =>
+      validateDeferredJson(kind, text)
+    case _ => throw new IllegalArgumentException("Unsupported durable input kind")
+  }
+
+  private def validateCells(kinds: Seq[String], row: Row, api: String): Unit = {
+    kinds.zipWithIndex.foreach { case (kind, index) =>
+      if (!row.isNullAt(index)) validateCell(kind, row.getString(index), api)
+    }
+  }
+
+  def validateRows(stage: AIPrompt, snapshot: DataFrame, preparedInput: Option[String] = None): Unit = {
+    val api = stage.getApiType
+    val kinds = Seq.newBuilder[String]
+    val columns = Seq.newBuilder[org.apache.spark.sql.Column]
+    def add(name: String, kind: String, json: Boolean = false): Unit = {
+      val resolved = AIColumnPlan.resolve(snapshot.schema, name)
+      kinds += kind
+      val column = AIColumnPlan.column(resolved)
+      columns += (if (json) F.to_json(column) else column)
+    }
+    stage.getColumnTypes.foreach { case (name, kind) => if (kind == "path") add(name, "inline") }
+    if (stage.isSet(stage.messagesCol)) add(stage.getMessagesCol, "messages", json = true)
+    stage.get(stage.extraBody).orElse(stage.getDefault(stage.extraBody)).flatMap(_.right.toOption)
+      .foreach(add(_, "body"))
+    stage.get(stage.responseFormat).orElse(stage.getDefault(stage.responseFormat)).flatMap(_.right.toOption)
+      .foreach { name =>
+        val dataType = snapshot.schema(AIColumnPlan.resolve(snapshot.schema, name)).dataType
+        add(name, api + "_format", json = dataType != StringType)
+      }
+    stage.get(stage.questions).orElse(stage.getDefault(stage.questions)).flatMap(_.right.toOption)
+      .foreach(add(_, api + "_questions"))
+    if (AIRequestAdapter.NativeApis(api) && stage.getInputMode == "json") {
+      preparedInput.orElse(stage.get(stage.stateCol)).orElse(stage.get(stage.inputCol))
+        .foreach(add(_, api))
+    } else preparedInput.foreach(add(_, "prompt"))
+    val selectedKinds = kinds.result()
+    if (selectedKinds.nonEmpty) {
+      snapshot.select(columns.result(): _*).foreachPartition { rows: Iterator[Row] =>
+        rows.foreach(validateCells(selectedKinds, _, api))
+      }
+    }
+  }
+}