diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs index 05f7d9610d5..ac37440d682 100644 --- a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs +++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs @@ -2,10 +2,10 @@ //! API has to resolve its own crypto provider, in a test binary where nothing //! has installed a process-wide one, and has to leave it uninstalled. -use std::collections::HashMap; use std::time::Duration; -use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection; +use futures_util::{sink, stream}; +use litellm_ai_gateway::io::responses_ws::async_responses_websocket; use tokio::net::TcpListener; async fn dead_tls_server() -> u16 { @@ -30,10 +30,15 @@ async fn dead_tls_server() -> u16 { async fn dialing_wss_returns_an_error_instead_of_panicking() { let port = dead_tls_server().await; - let result = ResponsesWebSocketConnection::connect_url( - &format!("wss://127.0.0.1:{port}/"), - &HashMap::new(), + let result = async_responses_websocket( + "gpt-5", + Some("test-key"), + Some(&format!("wss://127.0.0.1:{port}/")), + None, Some(Duration::from_secs(10)), + |_| {}, + stream::empty(), + sink::drain(), ) .await; diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index e00a6e6d3a3..6dc6b34b73d 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -174,7 +174,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result value .as_u64() .and_then(|value| usize::try_from(value).ok()) - .filter(|value| *value > 0) + .filter(|value| *value > 0 && *value <= defaults.max_response_bytes) .ok_or_else(|| OcrRequestError::RequestField { path: "max_response_bytes".into(), }) diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index ac64681bb01..55f8713d76e 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -727,6 +727,7 @@ fn response_limit_is_validated_and_not_forwarded_to_the_provider() { json!(true), json!("123"), json!(1.5), + json!(crate::constants::OCR_RESPONSE_MAX_BYTES + 1), Value::Null, ] { let wire = serde_json::from_value(json!({ diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index 65e4f1c3fb5..014564ae89d 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -45,17 +45,13 @@ pub(crate) trait PythonRoute: Send + Sync { fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; } -type HostResumeStep = HostStep< - NativeCallStep< - <::Call as NativeCall>::Operation, - <::Call as NativeCall>::Complete, - >, - Py, ->; +type NativeStep = NativeCallStep<::Operation, ::Complete>; +type NativeResult = Result, litellm_core::Error>; +type HostResumeStep = HostStep::Call>, Py>; struct NativeCallState { call: C, - result: Option, litellm_core::Error>>, + result: Option>, } enum PendingOperation { @@ -141,11 +137,7 @@ impl PythonLifecycle { } } - fn take_native_result( - &self, - ) -> PyResult< - NativeCallStep<::Operation, ::Complete>, - > { + fn take_native_result(&self) -> PyResult> { self.call .as_ref() .ok_or_else(missing_state)? diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index 16b8ba52e71..12d902a3544 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -22,7 +22,7 @@ struct PythonOcrHost { enum OcrHostData { Unprojected { request: Py }, - Projected(ProjectedOcrHost), + Projected(Box), Released, } @@ -192,13 +192,13 @@ impl PythonRoute for PythonOcrHost { let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?; let has_token_provider = projected.fields.azure_ad_token_provider.is_some(); let request = projected.request; - self.data = OcrHostData::Projected(ProjectedOcrHost { + self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost { fields: projected.fields, pre_call: None, retained_fields: None, body: None, headers: None, - }); + })); OcrHostResult::Request(Ok((Box::new(request), has_token_provider))) } OcrHostOperation::AcquireAzureAdToken => {