fix: refresh Fabric authentication during AI Functions retries - #2685
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
|
Hey Rana Singh (@ranadeepsingh) 👋! We use semantic commit messages to streamline the release process. Examples of commit messages with semantic prefixes:
To test your commit locally, please follow our guild on building from source. |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The current marker-trust logic can allow a spoofed implicit-auth marker to trigger auth-switching/refresh behavior on trusted endpoints, and endpoint-path normalization has a null/locale edge case that can break validation.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review tier: Lite
Findings: 1
New issues introduced by this change (2)
| Severity | Finding |
|---|---|
 |
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala — usesTrustedFabricAuth currently trusts the presence of X-SynapseML-Implicit-Fabric-Auth alone.… |
 |
core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala — trustedPath calls rawPath.toLowerCase and later rawPath.replaceAll(...) without guarding… |
What changed in this PR
This PR improves reliability of Fabric-hosted OpenAI (AI Functions) calls by enabling a safe, replayable auth refresh on trusted 401s, while preserving explicit caller auth and maintaining existing 429 backoff behavior. It introduces an internal marker header to opt only implicit Fabric-default OpenAI requests into auth-aware retries, adds endpoint validation, and adds token-cache invalidation + refresh single-flight.
Changes:
- Add
sendWithFabricAuthRetriesand route only trusted implicit Fabric OpenAI requests through auth-refresh-aware retry logic. - Add an internal Fabric auth marker header that is stripped before transmission and used to decide eligibility for replay/refresh.
- Add Fabric token cache invalidation helpers + workspace/artifact single-flight refresh lock, plus targeted Scala tests.
| File | Description |
|---|---|
| core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala | Adds Fabric-auth-aware retry path and routes eligible requests through it. |
| core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala | Adds Fabric auth marker constant/helpers; strips marker from outgoing headers. |
| core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala | Adds trusted endpoint validation and synchronized token refresh logic. |
| core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala | Adds runtime/NFS/in-memory token cache invalidation plumbing. |
| cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala | Tracks whether Fabric fallback auth was used and conditionally sets marker header. |
| cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala | Limits Fabric fallback auth + retry eligibility to implicit default Fabric OpenAI endpoint. |
| cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricHeadersSuite.scala | Verifies marker is set only for implicit Fabric auth and is stripped before sending. |
| core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifySendWithRetries.scala | Adds tests for 401 replay/refresh, 429 reacquisition, and endpoint validation behavior. |
| core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/VerifyTokenInvalidation.scala | Adds unit test for encoded NFS cache key invalidation flow. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Rana Singh (ranadeepsingh)
changed the title
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
A small but concrete performance/reliability issue remains in TokenLibrary.objectMethod where .toSeq.headOption forces reflection over all candidates instead of short-circuiting on the first match.
Review tier: Lite
Findings: None
Issues resolved since last review (2)
| Severity | Finding |
|---|---|
|
core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala — trustedPath calls rawPath.toLowerCase and later rawPath.replaceAll(...) without guarding… View resolved comment |
|
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala — usesTrustedFabricAuth currently trusts the presence of X-SynapseML-Implicit-Fabric-Auth alone.… View resolved comment |
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala:79
objectMethodmaterializes the full iterator with.toSeq.headOption, so it will attempt to load and reflect on all candidate class names even after a match is found. This adds unnecessary classloading/reflection overhead and can also trigger avoidable failures if a later candidate class has initialization/linkage issues. Short-circuit after the first successful match.
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
responseBodyForInspection can throw during 429/401 inspection and break retry handling by propagating exceptions and/or partially consuming the response stream.
Review tier: Lite
Findings: None
Issues resolved since last review (2)
| Severity | Finding |
|---|---|
|
core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala — trustedPath calls rawPath.toLowerCase and later rawPath.replaceAll(...) without guarding… View resolved comment |
|
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala — usesTrustedFabricAuth currently trusts the presence of X-SynapseML-Implicit-Fabric-Auth alone.… View resolved comment |
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala:124
responseBodyForInspectioncurrently lets exceptions fromentity.getContent,IOUtils.copyLarge, and/or stream cleanup propagate. That can abort the retry logic on 429/401 inspection paths, and can also leave the response entity partially consumed. Make inspection best-effort by catchingNonFataland, on failure, replaying any bytes already read back into the response entity before continuing.
|
Addressed the current and previously suppressed automated-review findings in d029d5d: |
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
|
Addressed the final head-specific suppressed findings in d029d5d: runtime reflection now short-circuits after the first compatible class, and bounded response inspection is best-effort while replaying partially read bytes after failures. Validation: scalastyle/test:scalastyle; VerifyResponseBodyInspection 1 passed; VerifySendWithRetries 23 passed; VerifyTokenInvalidation 3 passed. |
There was a problem hiding this comment.
Copilot review overview
🟢 Approval recommended
The retry/auth changes are well-scoped, preserve explicit credentials, and are backed by targeted regression tests; only a minor resource-cleanup improvement remains.
Review tier: Lite
Findings: None
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala:390
advancedbuilds a preview message by readingr.getEntity.getContentbut never closes that InputStream. Even withByteArrayEntitythis is best-effort resource cleanup, and it becomes a leak if a different streaming entity type is ever used inHTTPRequestData.toHTTPCoreor downstream callers. Wrap the read in atry/finallythat closes the stream, and ensurepreviewRequest.releaseConnection()runs in afinallyas well.
|
Addressed the final current-head suppressed cleanup finding in 28ae704: request previewing now closes the entity stream in |
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
TokenLibrary’s new reflection-based cache invalidation path can throw on unexpected runtime behaviors (instead of safely degrading), which risks breaking auth refresh/retry in production.
Review tier: Lite
Findings: None
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala:77
objectMethodonly catchesClassNotFoundException/NoSuchFieldException. Other reflective failures (e.g.,IllegalAccessException,InvocationTargetExceptionfrom module lookup, orSecurityException) will currently escape and can break auth refresh/invalidation paths instead of cleanly falling back to the next strategy.
core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/TokenLibrary.scala:112deleteNfsTokenthrows anIllegalStateExceptionfor an unexpected return type fromgetNFSTokenFilePath. That exception will abort token invalidation entirely (and prevents the in-memory cache clear from running), which is risky in a best-effort refresh path. Prefer treating unknown types as "can't delete" and continuing so other invalidation mechanisms can still succeed.
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It modifies authentication, retry, and request/response replay behavior in core HTTP handling, which is high-impact and warrants final human review despite strong targeted test coverage.
Review tier: Lite
Findings: None
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #2685 +/- ##
==========================================
- Coverage 86.89% 86.85% -0.04%
==========================================
Files 341 341
Lines 20742 20871 +129
Branches 2166 2231 +65
==========================================
+ Hits 18023 18127 +104
- Misses 2719 2744 +25
🚀 New features to boost your workflow:
|
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
AB#3380998 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
AB#3582121 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 98752775-8e28-43eb-9609-2249dd586f80
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
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
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
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
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
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
dea4a52 to
2ea0d9f
Compare
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The new response-body inspection helper can miss CapacityLimitExceeded detection on large/unknown-length 429 payloads, risking incorrect retry behavior, and authorizationHeader should avoid returning Some(null) to prevent null propagation.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review tier: Lite
Findings: 2
New issues introduced by this change (2)
| Severity | Finding |
|---|---|
|
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala — responseBodyForInspection returns None whenever more than 1MiB is read, even though the first… |
|
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala — authorizationHeader can return Some(null) if an Authorization header is present with a null… |
## 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: microsoft#2685 (comment) - Null-authorization finding: microsoft#2685 (comment) - Pull request: microsoft#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
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It changes core authentication and HTTP retry behavior (including reflection-based token invalidation and replay logic), which is security- and reliability-critical and warrants final human review despite strong targeted test coverage.
Review tier: Lite
Findings: None
Issues resolved since last review (2)
| Severity | Finding |
|---|---|
|
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala — authorizationHeader can return Some(null) if an Authorization header is present with a null… View resolved comment |
|
core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala — responseBodyForInspection returns None whenever more than 1MiB is read, even though the first… View resolved comment |
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
Rana Singh (ranadeepsingh)
changed the title
Related Issues/PRs
What changes are proposed in this pull request?
Refresh Fabric authentication safely during long-running AI Functions requests:
How is this patch tested?
Current-head local validation:
VerifySendWithRetries: 24 passedVerifyResponseBodyInspection: 3 passedVerifyTokenInvalidation: 5 passedOpenAIFabricHeadersSuite: 7 passedOpenAIFabricAuthRetrySuite: 1 passedAzureSearchAuthSuite: 48 passedAddDocumentsHeaderPersistenceSuite: 1 passedsbt codegenandsbt packagePythonpassed; generated OpenAI wrappers compiled and were present in the wheelDoes this PR change any dependencies?
Does this PR add a new feature? If so, have you added samples on website?
Current integration evidence
25815741901ab8d9062c51eb3a7f06070550e607after the target CI repair merged; final head504c399c5a10b2aea0db69d6c7961b46c5a4ed54is 10 commits ahead and 0 behind.Authorizationheader values are treated as absent.Findings: None; all review threads are resolved.fbbcad43035627fc53e4799428ef191af1043ba9and passed all 65 jobs.