From 36c2d91d85b900366d694ef1e45c555d0b5bb889 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 04:19:27 +0000 Subject: [PATCH] fix(azure-di): bound operation polling Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .../ai-gateway/src/io/ocr/common_utils.rs | 11 ++- .../crates/ai-gateway/src/io/ocr/tests.rs | 88 ++++++++++++++++--- 2 files changed, 86 insertions(+), 13 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs index e965209a0dd..dd9d11efc90 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs @@ -515,11 +515,12 @@ pub(super) async fn poll_document_intelligence( AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS, )); loop { - if start.elapsed() > timeout { + let remaining = timeout.saturating_sub(start.elapsed()); + if remaining.is_zero() { return Err(CoreError::Timeout); } - let mut request_builder = http_client().get(operation_url); + let mut request_builder = http_client().get(operation_url).timeout(remaining); for (key, value) in headers { if key.eq_ignore_ascii_case("ocp-apim-subscription-key") { request_builder = request_builder.header(key, value); @@ -544,7 +545,11 @@ pub(super) async fn poll_document_intelligence( if operation_status(&response_json)? == "succeeded" { return Ok(response_json); } - tokio::time::sleep(Duration::from_secs(retry_after)).await; + let remaining = timeout.saturating_sub(start.elapsed()); + if remaining.is_zero() { + return Err(CoreError::Timeout); + } + tokio::time::sleep(remaining.min(Duration::from_secs(retry_after))).await; } } diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs index 901e20275bc..f2a4ef8dce9 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs @@ -366,8 +366,21 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() { .await .expect("writes post response"); - let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); - let poll_request = read_http_headers(&mut poll_socket).await; + let (mut running_socket, _) = listener.accept().await.expect("accepts running poll"); + let running_request = read_http_headers(&mut running_socket).await; + let running_body = r#"{"status":"running"}"#; + let running_response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\nretry-after: 0\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + running_body.len(), + running_body + ); + running_socket + .write_all(running_response.as_bytes()) + .await + .expect("writes running response"); + + let (mut poll_socket, _) = listener.accept().await.expect("accepts completed poll"); + let completed_request = read_http_headers(&mut poll_socket).await; let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#; let poll_response = format!( "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", @@ -378,7 +391,7 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() { .write_all(poll_response.as_bytes()) .await .expect("writes poll response"); - (post_request, poll_request) + (post_request, running_request, completed_request) }); let response = ocr(OcrRequest { @@ -399,19 +412,74 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() { assert_eq!(response["pages"][0]["markdown"], "ok"); - let (post_request, poll_request) = server.await.expect("server task completes"); + let (post_request, running_request, completed_request) = + server.await.expect("server task completes"); assert!( post_request .to_ascii_lowercase() .contains("ocp-apim-subscription-key: di-key"), "{post_request}" ); - assert!( - poll_request - .to_ascii_lowercase() - .contains("ocp-apim-subscription-key: di-key"), - "{poll_request}" - ); + for poll_request in [running_request, completed_request] { + assert!( + poll_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{poll_request}" + ); + } +} + +#[tokio::test] +async fn document_intelligence_retry_after_respects_polling_timeout() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + let operation_url = format!("http://{addr}/operations/1"); + + tokio::spawn(async move { + let (mut post_socket, _) = listener.accept().await.expect("accepts post request"); + read_http_headers(&mut post_socket).await; + let post_response = format!( + "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" + ); + post_socket + .write_all(post_response.as_bytes()) + .await + .expect("writes post response"); + + let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); + read_http_headers(&mut poll_socket).await; + let body = r#"{"status":"running"}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\nretry-after: 60\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + poll_socket + .write_all(response.as_bytes()) + .await + .expect("writes poll response"); + }); + + let error = ocr(OcrRequest { + model: "prebuilt-read", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("di-key"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: "azure_ai/doc-intelligence", + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_millis(100)), + }) + .await + .expect_err("polling deadline should bound retry-after"); + + assert_eq!(error, CoreError::Timeout); } #[test]