fix: Refresh Fabric authentication during AI Functions retries - #2685
fix: Refresh Fabric authentication during AI Functions retries#2685Rana Singh (ranadeepsingh) wants to merge 9 commits into
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.
|
Addressed the eager preview finding in d4e7293: request reconstruction and POST body reads now remain inside the by-name debug log expression, so normal execution does not allocate or consume a preview. Existing preview cleanup coverage remains green. |
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It changes core HTTP/auth retry behavior (including token invalidation and endpoint trust validation), which is high impact and merits final human review despite strong targeted tests.
Review tier: Lite
Findings: None
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
|
Current exact-head CI evidence (not a readiness claim): Azure build 234078853 validates merge |
d4e7293 to
d778d98
Compare
|
/azp run |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It changes security- and reliability-critical authentication and retry behavior (including reflection-based runtime invalidation), so it warrants final human review despite strong targeted test coverage.
Review tier: Lite
Findings: None
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
d778d98 to
dea4a52
Compare
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:
|
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: 23 passedVerifyResponseBodyInspection: 2 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
1db3c40b1a7397f4801ab6077dfa282da9781ac4; final headdea4a52d6037d0a1bdbb87277ade71a348d81642is 9 ahead and 0 behind.Findings: None; all review threads are resolved.5b75adab673d8de2d0b4af525c3aae2549ca06a7; all 65 jobs completed with 51 succeeded and 14 failed.1db3c40b1a7397f4801ab6077dfa282da9781ac4and failed the same 14 job names as PR build 234219169.OpenAIPromptSuite.scala:155receives the same AI Foundry HTTP 403 and reports0 did not equal 3.PythonTests lightgbmandDatabricks GPU E2Epassed in build 234219169, clearing the previous run's infrastructure-only LightGBM cleanup failure.