From 3ad91fc27e1769e1a69c40f01e0d07d0d3a590e0 Mon Sep 17 00:00:00 2001
From: Yujong Lee <yujong@berri.ai>
Date: Thu, 17 Sep 2026 06:53:19 -0700
Subject: [PATCH] fix(ocr): run hooks on completed Azure poll

---
 .../document_intelligence/transformation.rs   | 26 ++++++++++++-------
 1 file changed, 16 insertions(+), 10 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 e20ec29132d..1016ed02783 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
@@ -343,7 +343,7 @@ async fn read_operation_response(
     let bytes =
         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).await
+    poll_operation(http_client, operation, headers, connection, native, hooks).await
 }
 
 async fn poll_operation(
@@ -352,6 +352,7 @@ async fn poll_operation(
     headers: &[(String, String)],
     connection: &OcrConnection,
     native: bool,
+    hooks: &Arc<dyn OcrHooks>,
 ) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, crate::ocr::Error> {
     let deadline = Instant::now()
         .checked_add(connection.poll_timeout)
@@ -392,7 +393,10 @@ async fn poll_operation(
         .await
         .map_err(|_| crate::ocr::Error::PollTimeout)??;
         match &decoded.data.status {
-            Some(OperationStatus::Succeeded) => return Ok(decoded),
+            Some(OperationStatus::Succeeded) => {
+                crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?;
+                return Ok(decoded);
+            }
             Some(OperationStatus::Running | OperationStatus::NotStarted) => {
                 tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry)))
                     .await
@@ -985,7 +989,7 @@ mod tests {
 
     struct SubmissionBoundary {
         request_count: Arc<Mutex<Vec<String>>>,
-        post_calls: Arc<Mutex<Vec<Value>>>,
+        post_calls: Arc<Mutex<Vec<(usize, Value)>>>,
     }
 
     impl crate::ocr::hooks::OcrHooks for SubmissionBoundary {
@@ -994,18 +998,17 @@ mod tests {
             request: crate::ocr::hooks::OcrPostCallRequest,
         ) -> crate::ocr::hooks::OcrHookFuture<'_, crate::ocr::hooks::OcrPostCallRequest> {
             Box::pin(async move {
-                assert_eq!(self.request_count.lock().unwrap().len(), 1);
-                self.post_calls
-                    .lock()
-                    .unwrap()
-                    .push(request.original_response.clone());
+                self.post_calls.lock().unwrap().push((
+                    self.request_count.lock().unwrap().len(),
+                    request.original_response.clone(),
+                ));
                 Ok(request)
             })
         }
     }
 
     #[tokio::test]
-    async fn accepted_response_runs_post_call_once_before_polling() {
+    async fn accepted_response_runs_post_call_for_submission_and_completed_poll() {
         let (base, seen, server) = mock_server(vec![
             MockResponse {
                 status: 202,
@@ -1029,7 +1032,10 @@ mod tests {
         assert_eq!(seen.lock().unwrap().len(), 2);
         assert_eq!(
             *post_calls.lock().unwrap(),
-            [json!(r#"{"submitted":true}"#)]
+            [
+                (1, json!(r#"{"submitted":true}"#)),
+                (2, json!(r#"{"status":"succeeded"}"#)),
+            ]
         );
     }