From ff3185f38568c0e7aee015044982928e2f1b8514 Mon Sep 17 00:00:00 2001
From: Yujong Lee <yujong@berri.ai>
Date: Tue, 15 Sep 2026 09:14:48 -0700
Subject: [PATCH] fix(ocr): avoid duplicate post-call callback after Azure
 polling

---
 .../document_intelligence/transformation.rs   |  8 ++----
 .../tests/azure_document_intelligence_ocr.rs  | 25 ++++++++++---------
 2 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs
index a987de3f842..ec65bac3774 100644
--- a/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs
+++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs
@@ -530,7 +530,7 @@ mod polling {
             crate::ocr::client::read_response_bytes(response, connection.max_response_bytes)
                 .await?;
         crate::ocr::handler::post_call(hooks, &bytes).await?;
-        poll_operation(http_client, operation, headers, connection, native, hooks).await
+        poll_operation(http_client, operation, headers, connection, native).await
     }
 
     async fn poll_operation(
@@ -539,7 +539,6 @@ mod polling {
         headers: &[(String, String)],
         connection: &OcrConnection,
         native: bool,
-        hooks: &Arc<dyn OcrHooks>,
     ) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
         let deadline = Instant::now()
             .checked_add(connection.poll_timeout)
@@ -583,10 +582,7 @@ mod polling {
             .await
             .map_err(|_| OcrPollingError::PollTimeout)??;
             match &decoded.data.status {
-                Some(OperationStatus::Succeeded) => {
-                    crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?;
-                    return Ok(decoded);
-                }
+                Some(OperationStatus::Succeeded) => return Ok(decoded),
                 Some(OperationStatus::Running | OperationStatus::NotStarted) => {
                     tokio::time::timeout_at(
                         deadline,
diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
index d20c1e8e83b..62d0c657a1a 100644
--- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
+++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
@@ -180,6 +180,7 @@ async fn accepted_response_polls_to_success_with_only_credentials() {
 
 struct SubmissionBoundary {
     request_count: Arc<Mutex<Vec<String>>>,
+    post_calls: Arc<Mutex<Vec<Value>>>,
 }
 
 impl super::hooks::OcrHooks for SubmissionBoundary {
@@ -188,24 +189,18 @@ impl super::hooks::OcrHooks for SubmissionBoundary {
         request: super::hooks::OcrPostCallRequest,
     ) -> super::hooks::OcrHookFuture<'_, super::hooks::OcrPostCallRequest> {
         Box::pin(async move {
-            match self.request_count.lock().unwrap().len() {
-                1 => assert_eq!(request.original_response, json!(r#"{"submitted":true}"#)),
-                2 => assert!(
-                    request
-                        .original_response
-                        .as_str()
-                        .unwrap()
-                        .contains("succeeded")
-                ),
-                count => panic!("unexpected callback after {count} requests"),
-            }
+            assert_eq!(self.request_count.lock().unwrap().len(), 1);
+            self.post_calls
+                .lock()
+                .unwrap()
+                .push(request.original_response.clone());
             Ok(request)
         })
     }
 }
 
 #[tokio::test]
-async fn accepted_response_runs_post_call_before_polling() {
+async fn accepted_response_runs_post_call_once_before_polling() {
     let (base, seen, server) = mock_server(vec![
         MockResponse {
             status: 202,
@@ -215,9 +210,11 @@ async fn accepted_response_runs_post_call_before_polling() {
         MockResponse::json(json!({"status":"succeeded"})),
     ])
     .await;
+    let post_calls = Arc::new(Mutex::new(Vec::new()));
     let request = super::LiteLLMOcrRequest {
         hooks: Arc::new(SubmissionBoundary {
             request_count: seen.clone(),
+            post_calls: post_calls.clone(),
         }),
         ..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}))
     };
@@ -225,6 +222,10 @@ async fn accepted_response_runs_post_call_before_polling() {
     perform_ocr(request).await.unwrap();
     server.await.unwrap();
     assert_eq!(seen.lock().unwrap().len(), 2);
+    assert_eq!(
+        *post_calls.lock().unwrap(),
+        [json!(r#"{"submitted":true}"#)]
+    );
 }
 
 #[tokio::test]