From c635c35b3d2968dbf76ebec98eee833e2b5c8a0f Mon Sep 17 00:00:00 2001
From: Yujong Lee <yujong@berri.ai>
Date: Fri, 18 Sep 2026 18:47:36 -0700
Subject: [PATCH] fix(rust): keep native OCR on the proxy by declining only a
 supplied client

The proxy attaches its shared aiohttp session to every request as shared_session, so declining on it sent every proxy OCR call to Python, which never uses that session for OCR. aclient_session is a litellm global and never a call argument, so that check could not match. The proxy-shaped lifecycle test now asserts the call was served by Rust
---
 litellm-rust/crates/python-bridge/src/http.rs | 46 +++++++++----------
 tests/test_litellm_rust/ocr/test_lifecycle.py |  1 +
 2 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs
index 77542855fec..c5952c53132 100644
--- a/litellm-rust/crates/python-bridge/src/http.rs
+++ b/litellm-rust/crates/python-bridge/src/http.rs
@@ -12,8 +12,6 @@ use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings};
 static POOL: LazyLock<HttpClientPool> =
     LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver)));
 
-const LIVE_CLIENT_ARGUMENTS: [&str; 3] = ["client", "shared_session", "aclient_session"];
-
 pub(crate) fn pool() -> &'static HttpClientPool {
     &POOL
 }
@@ -23,7 +21,7 @@ pub(crate) fn call_config(
     kwargs: &Bound<'_, PyDict>,
     asynchronous: bool,
 ) -> PyResult<HttpClientConfig> {
-    decline_live_clients(kwargs)?;
+    decline_live_client(kwargs)?;
     decline_custom_url_policy(&PythonSettings::UrlPolicy.read(py)?)?;
     let configured = settings(&PythonSettings::Http.read(py)?)?
         .with_environment(&|name| std::env::var(name).ok());
@@ -53,13 +51,14 @@ fn for_call(
     }
 }
 
-pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> {
-    for name in LIVE_CLIENT_ARGUMENTS {
-        if kwargs.get_item(name)?.is_some_and(|value| !value.is_none()) {
-            return Err(RustBridgeDeclined::new_err(format!(
-                "{name} is a live Python HTTP client and cannot be used by the Rust route"
-            )));
-        }
+fn decline_live_client(kwargs: &Bound<'_, PyDict>) -> PyResult<()> {
+    if kwargs
+        .get_item("client")?
+        .is_some_and(|value| !value.is_none())
+    {
+        return Err(RustBridgeDeclined::new_err(
+            "client is a live Python HTTP client and cannot be used by the Rust route",
+        ));
     }
     Ok(())
 }
@@ -362,32 +361,29 @@ user_agent='litellm/9.9.9',
         assert_eq!(config.trust_proxy_env, expected);
     }
 
-    #[rstest]
-    #[case::client("client")]
-    #[case::shared_session("shared_session")]
-    #[case::aclient_session("aclient_session")]
-    fn live_python_clients_decline_before_dispatch(#[case] name: &str) {
+    #[test]
+    fn live_python_client_declines_before_dispatch() {
         Python::initialize();
         Python::attach(|py| {
             let kwargs = PyDict::new(py);
             kwargs
-                .set_item(name, py.eval(c"object()", None, None).unwrap())
+                .set_item("client", py.eval(c"object()", None, None).unwrap())
                 .unwrap();
-            let error = decline_live_clients(&kwargs).unwrap_err();
+            let error = decline_live_client(&kwargs).unwrap_err();
             assert!(error.is_instance_of::<RustBridgeDeclined>(py));
-            assert!(error.value(py).to_string().contains(name));
         });
     }
 
-    #[test]
-    fn none_valued_client_arguments_are_not_live_clients() {
+    #[rstest]
+    #[case::absent_client("{}")]
+    #[case::none_client("{'client': None}")]
+    #[case::proxy_shared_session("{'shared_session': object()}")]
+    fn calls_without_a_python_client_stay_on_the_rust_route(#[case] kwargs: &str) {
         Python::initialize();
         Python::attach(|py| {
-            let kwargs = PyDict::new(py);
-            for name in LIVE_CLIENT_ARGUMENTS {
-                kwargs.set_item(name, py.None()).unwrap();
-            }
-            decline_live_clients(&kwargs).unwrap();
+            let source = std::ffi::CString::new(kwargs).unwrap();
+            let kwargs = py.eval(&source, None, None).unwrap();
+            decline_live_client(kwargs.cast::<PyDict>().unwrap()).unwrap();
         });
     }
 }
diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py
index 5fca927bea3..264a666c685 100644
--- a/tests/test_litellm_rust/ocr/test_lifecycle.py
+++ b/tests/test_litellm_rust/ocr/test_lifecycle.py
@@ -39,6 +39,7 @@ async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer)
     )
     events: Final = await recorder.wait_for_async("async_log_success_event")
     assert response.pages[0].markdown == "native OCR response"
+    assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
     assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user"
     assert "metadata" not in ocr_server.requests[0].body