refactor(rust): move credential inheritance and the SDK limits out of the legacy callback crate into a driver preflight (#43259)

The legacy callback crate carried two rewrites that have nothing to do with the
Logging contract: litellm_credential_name inheritance and the max_budget and
num_retries_per_request checks. Any later callback host would need them
unchanged, which is the smell the crate's AGENTS.md now names. They are now a
Preflight the driver in litellm-host-python runs on the keyword view begin
returned, before the host projects from it, supplied by python-bridge and passed
through run_legacy_call. The call order is unchanged (setup, deployment hook,
credentials, limits) and a rejection still fails the call as a host failure, so
the failure callbacks run as before. The preflight rewrites the adapter's own
copy in place, so no extra dict copy and no new lifecycle method

Co-authored-by: Yujong Lee <yujong@berri.ai>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

litellm-rust/Cargo.lock

@@ -3311,6 +3311,7 @@ dependencies = [
  "serde_json",
  "serde_with",
  "sha2 0.10.9",
+ "strum",
  "thiserror 2.0.19",
  "tokio",
  "tokio-tungstenite",

litellm-rust/crates/callbacks-legacy-python/AGENTS.md

@@ -1,19 +1,19 @@
 - Target invariants, not completion claims
-- Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits)
+- This crate is the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on)
+  - Smell test: if a future callback host (`callbacks-v1-python`, WASM, in-process Rust) could share a piece of this crate, it does not belong here
+  - SDK request policy (credential inheritance, the budget and retry-count limits) is the driver's preflight, supplied by `python-bridge`; this crate only adopts the keyword view it produces
   - The driver in `litellm-host-python`, the routes and core see one `PythonLifecycle`; they never learn which Python objects consume a call
-- Rust drives the call; every litellm Python internal it still borrows is a variant of `LegacyPython`, grouped by subsystem (`Wrapper`, `Logging`, `DeploymentHooks`)
+- Every litellm Python internal Rust still borrows is a variant of `LegacyPython`, grouped by subsystem, with its signature pinned in `python_contract.json`
   - The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it
   - Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython`
-  - `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy
-- `setup` reuses a `Logging` the caller passed as `litellm_logging_obj` (the proxy and Router are the live cases) and otherwise builds one through `function_setup`, as `@client` does
-  - Either way every phase calls the same `Logging` method the Python path calls; which callbacks run is `Logging`'s decision, never this crate's
+- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the call rewrites it (setup, deployment hook, preflight) and the bound request object backing omitted keywords; routes hand it over through `run_legacy_call` and keep no copy
+- `setup` reuses a `Logging` passed as `litellm_logging_obj` (the proxy and Router) and otherwise builds one through `function_setup`; which callbacks run is `Logging`'s decision, never this crate's
 - Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation
-  - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view
-  - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup`
-  - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only
-  - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-host`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view
-- Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts
-  - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch
-  - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once
+  - Retain complete boundary arguments, opaque values, aliases, omitted/default distinctions and deliberate copies; preserve the deployment-hook kwargs view
+  - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object, resolved through `litellm_host_python::lookup`
+  - Retain body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only
+- Success and failure handlers receive the exact selected public response or exception
+  - A failure-handler error cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch
+  - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy releases deferred success at most once
   - Delivery follows the registry, not the callable's type: direct, awaited, executor-submitted, logging-worker and deferred paths stay distinct
 - Traverse every retained Python edge; `close` is idempotent and restores the correlation context once

litellm-rust/crates/callbacks-legacy-python/python_contract.json

@@ -6,9 +6,6 @@
     "start_time",
     "asynchronous"
   ],
-  "check_limits": [
-    "kwargs"
-  ],
   "finalize": [
     "response",
     "logger",
@@ -76,11 +73,6 @@
   ],
   "custom_pricing_fields": [],
   "is_internal_call": [],
-  "credential_list": [],
-  "warn_unknown_credential": [
-    "name",
-    "loaded"
-  ],
   "before_deployment_call": [
     "kwargs",
     "call_type"

litellm-rust/crates/callbacks-legacy-python/src/adapter.rs

@@ -19,7 +19,7 @@ use serde_json::Value;
 use crate::{
     DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
     deferred::{PendingLogging, PendingSuccess},
-    finalize, is_internal_call, prepare,
+    finalize, is_internal_call,
     python::Streaming,
     setup,
 };
@@ -117,9 +117,13 @@ impl LegacyLogging {
         })
     }
 
+    /// The keyword view the rest of the call reads: a copy, so the deployment hook's own
+    /// dict is left as the hook returned it, carrying the logger as `@client` injects it.
+    /// The driver's preflight rewrites this same dict before the host projects from it.
     fn prepare(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
-        let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind();
-        self.call.set_kwargs(prepared);
+        let prepared = self.call.kwargs().bind(py).copy()?;
+        prepared.set_item("litellm_logging_obj", self.logger()?.object(py))?;
+        self.call.set_kwargs(prepared.unbind());
         Ok(LifecycleStep::Arguments(self.call.kwargs().clone_ref(py)))
     }
 
@@ -580,8 +584,6 @@ assert prepared['document'] is replacement
 assert prepared['pages'] is replaced_kwargs['pages']
 assert prepared['litellm_logging_obj'] is logger
 assert 'litellm_logging_obj' not in replaced_kwargs
-[checked] = [value for name, value in logger.calls if name == 'check_limits']
-assert checked is prepared
 ",
             );
         });
@@ -616,8 +618,6 @@ kwargs = {'logger': logger, 'vendor_extension': opaque}
                 &locals,
                 c"
 assert prepared['vendor_extension'] is opaque
-[checked] = [value for name, value in logger.calls if name == 'check_limits']
-assert checked['vendor_extension'] is opaque
 assert hooked == ([opaque] if asynchronous else []), hooked
 ",
             );
@@ -733,45 +733,6 @@ assert all(value is failure for name, value in logger.calls if name.endswith('_h
             );
         });
     }
-
-    #[rstest]
-    #[case::synchronous(false)]
-    #[case::asynchronous(true)]
-    fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) {
-        Python::initialize();
-        Python::attach(|py| {
-            let locals = namespace(
-                py,
-                c"
-class BudgetExceeded(Exception):
-    pass
-
-rejection = BudgetExceeded('over budget')
-
-class LimitedLogger(StubLogger):
-    def check_limits(self, arguments):
-        raise rejection
-
-logger = LimitedLogger()
-logger.hooks = {'pre': lambda kwargs: kwargs}
-kwargs = {'logger': logger}
-",
-            );
-            let mut logging = legacy_call(py, &locals, asynchronous);
-            let kwargs = local(&locals, "kwargs")
-                .cast_into::<PyDict>()
-                .unwrap()
-                .unbind();
-            let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step {
-                LifecycleStep::Await(_) => {
-                    logging.resume(py, Ok(local(&locals, "kwargs").unbind()))
-                }
-                step => Ok(step),
-            });
-            let error = result.err().unwrap();
-            assert!(error.value(py).is(local(&locals, "rejection")));
-        });
-    }
 }
 
 #[cfg(test)]

litellm-rust/crates/callbacks-legacy-python/src/call.rs

@@ -4,7 +4,7 @@
 //! this crate holds them.
 
 use litellm_host::{machine::Machine, protocol::Protocol};
-use litellm_host_python::{ProtocolHost, lookup, run_call};
+use litellm_host_python::{Preflight, ProtocolHost, lookup, run_call};
 use pyo3::{
     gc::{PyTraverseError, PyVisit},
     prelude::*,
@@ -39,7 +39,8 @@ impl PublicCall {
     }
 
     /// The keyword view the legacy path currently reads: the caller's copy until
-    /// `function_setup`, then each rewrite (setup, deployment hook, prepare) in turn.
+    /// `function_setup`, then each rewrite (setup, deployment hook, the driver's preflight)
+    /// in turn.
     pub(crate) fn kwargs(&self) -> &Py<PyDict> {
         &self.kwargs
     }
@@ -64,13 +65,15 @@ impl PublicCall {
 }
 
 /// Runs one native call under the legacy `Logging` contract: the protocol host projects from
-/// the keyword view the contract prepares, and the contract observes the call.
+/// the keyword view the contract prepares and `preflight` rewrites, and the contract
+/// observes the call.
 pub fn run_legacy_call<H, M>(
     py: Python<'_>,
     surface: LegacySurface,
     call: PublicCall,
     machine: M,
     host: H,
+    preflight: Preflight,
     asynchronous: bool,
 ) -> PyResult<Py<PyAny>>
 where
@@ -83,6 +86,7 @@ where
         machine,
         host,
         Box::new(LegacyLogging::new(py, surface, call, asynchronous)),
+        preflight,
         arguments,
         asynchronous,
     )

litellm-rust/crates/callbacks-legacy-python/src/lib.rs

@@ -1,9 +1,10 @@
 //! The legacy `@client` wrapper as the native call sees it: litellm's `Logging` object, the
-//! sync and async callback registries it fans out to, the deployment hooks, the deferred
-//! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name
-//! inheritance, budget and retry-count limits). All of it sits behind one
+//! sync and async callback registries it fans out to, the deployment hooks and the deferred
+//! proxy release. All of it sits behind one
 //! [`PythonLifecycle`](litellm_host_python::PythonLifecycle), so the driver, the routes and
-//! core never learn which Python object is on the other end.
+//! core never learn which Python object is on the other end. The SDK's own request policy
+//! (credential inheritance, the budget and retry limits) is the driver's preflight, not this
+//! crate's.
 //!
 //! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`]
 //! is where those objects live, and [`run_legacy_call`] is how a route hands them over
@@ -14,14 +15,12 @@ mod call;
 mod callbacks;
 mod deferred;
 mod logger;
-mod preparation;
 mod python;
 pub(crate) use adapter::LegacyLogging;
 pub use adapter::{LegacySurface, PassThroughStream};
 pub use call::{PublicCall, run_legacy_call};
 pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
 pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};
-pub(crate) use preparation::prepare;
 
 #[cfg(test)]
 mod test_support {
@@ -77,7 +76,6 @@ FAKES = {
         logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'],
         kwargs=kwargs,
     ),
-    'check_limits': lambda arguments: arguments['logger'].check_limits(arguments),
     'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response),
     'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs(
         kwargs=kwargs,
@@ -104,8 +102,6 @@ FAKES = {
     'restore_context': lambda logger: logger.record('restore', None),
     'custom_pricing_fields': lambda: ('ocr_cost_per_page',),
     'is_internal_call': lambda: legacy.is_internal.get(),
-    'credential_list': lambda: [],
-    'warn_unknown_credential': lambda name, loaded: None,
     'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type),
     'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook(
         'success', response, call_type
@@ -162,9 +158,6 @@ class StubLogger:
         self.record(phase + '_hook', call_type)
         return self.hooks.get(phase, lambda value: 'awaitable')(value)
 
-    def check_limits(self, arguments):
-        self.record('check_limits', arguments)
-
     def failure_handler(self, error, trace, start, end):
         self.record('failure_handler', error)
 

litellm-rust/crates/callbacks-legacy-python/src/python.rs

@@ -19,18 +19,12 @@ pub(crate) enum LegacyPython {
     Streaming(Streaming),
 }
 
-/// The `@client` wrapper around the call: `function_setup`, limits, credentials,
-/// response metadata and the correlation context.
+/// The `@client` wrapper around the call: `function_setup`, response metadata and the
+/// correlation context.
 #[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)]
 pub(crate) enum Wrapper {
     #[strum(serialize = "setup")]
     Setup,
-    #[strum(serialize = "check_limits")]
-    CheckLimits,
-    #[strum(serialize = "credential_list")]
-    CredentialList,
-    #[strum(serialize = "warn_unknown_credential")]
-    WarnUnknownCredential,
     #[strum(serialize = "is_internal_call")]
     IsInternalCall,
     #[strum(serialize = "finalize")]

litellm-rust/crates/host-python/AGENTS.md

@@ -2,7 +2,7 @@
 - Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`ProtocolHost` traits
   - No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features
   - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business
-  - `ProtocolHost::project` receives the keyword view the adapter's `begin` returned, not the caller's dict; a protocol host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance)
+  - `ProtocolHost::project` receives the keyword view the adapter's `begin` returned, rewritten in place by the route's `Preflight`, not the caller's dict; a protocol host that projects from it inherits the adapter's rewrites (for the legacy adapter: setup, deployment hooks) and the preflight's (credential inheritance)
   - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is
   - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed
 - Use standard PyO3 ownership and conversion APIs

litellm-rust/crates/host-python/src/adapter.rs

@@ -9,6 +9,12 @@ pub fn missing_state() -> PyErr {
     PyRuntimeError::new_err("missing native call state")
 }
 
+/// The SDK's request policy, run by the driver on the keyword view `begin` returned and
+/// before the protocol host projects from it. It rewrites that view in place, so the
+/// lifecycle that returned it sees the rewrite too; a rejection fails the call as a host
+/// failure, so the lifecycle still observes it.
+pub type Preflight = fn(Python<'_>, &Bound<'_, PyDict>) -> PyResult<()>;
+
 /// What an adapter step produced: either the value the driver asked for, or a Python
 /// awaitable the driver hands back to the caller's task before asking again.
 pub enum LifecycleStep {

litellm-rust/crates/host-python/src/driver.rs

@@ -14,7 +14,8 @@ use pyo3::types::PyDict;
 use tokio::sync::Mutex;
 
 use crate::adapter::{
-    InvokeError, LifecycleEvent, LifecycleStep, ProtocolHost, PythonLifecycle, missing_state,
+    InvokeError, LifecycleEvent, LifecycleStep, Preflight, ProtocolHost, PythonLifecycle,
+    missing_state,
 };
 use crate::execution::{poll_async_value, run_async_value, run_sync_value};
 use crate::handle::{Execution, ExecutionBody, ExecutionStep};
@@ -83,6 +84,7 @@ where
 {
     host: H,
     adapter: Box<dyn PythonLifecycle>,
+    preflight: Preflight,
     machine: Option<Arc<Mutex<MachineState<M>>>>,
     arguments: Option<Py<PyDict>>,
     started_at: f64,
@@ -95,12 +97,14 @@ where
 }
 
 /// Runs one native call for Python: synchronously, or as a coroutine that awaits every
-/// host suspension inline in the caller's task.
+/// host suspension inline in the caller's task. `preflight` runs once, on the keyword view
+/// the adapter's `begin` returned, before the host projects from it.
 pub fn run_call<H, M>(
     py: Python<'_>,
     machine: M,
     host: H,
     adapter: Box<dyn PythonLifecycle>,
+    preflight: Preflight,
     arguments: Py<PyDict>,
     asynchronous: bool,
 ) -> PyResult<Py<PyAny>>
@@ -111,6 +115,7 @@ where
     let mut driver = PythonDriver {
         host,
         adapter,
+        preflight,
         machine: Some(Arc::new(Mutex::new(MachineState {
             machine,
             result: None,
@@ -213,6 +218,9 @@ where
         match (expect, step) {
             (Expect::Started, LifecycleStep::Done) => self.begin(py),
             (Expect::Arguments, LifecycleStep::Arguments(arguments)) => {
+                if let Err(error) = (self.preflight)(py, arguments.bind(py)) {
+                    return self.adapter_failed(py, error);
+                }
                 self.arguments = Some(arguments);
                 self.stage = Stage::Call;
                 self.resume_machine(py, None)
@@ -869,6 +877,21 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
         host: SyntheticHost,
         script: AdapterScript,
         asynchronous: bool,
+    ) -> (PyResult<Py<PyAny>>, Vec<String>) {
+        run_preflighted(py, machine, host, script, no_preflight, asynchronous)
+    }
+
+    fn no_preflight(_: Python<'_>, _: &Bound<'_, PyDict>) -> PyResult<()> {
+        Ok(())
+    }
+
+    fn run_preflighted(
+        py: Python<'_>,
+        machine: CallMachine<Synthetic>,
+        host: SyntheticHost,
+        script: AdapterScript,
+        preflight: Preflight,
+        asynchronous: bool,
     ) -> (PyResult<Py<PyAny>>, Vec<String>) {
         let log = Log(host.log.0.clone());
         let adapter = SyntheticAdapter {
@@ -882,6 +905,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
             machine,
             host,
             Box::new(adapter),
+            preflight,
             arguments.unbind(),
             asynchronous,
         );
@@ -1088,6 +1112,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
                     streaming_machine(),
                     StreamingHost,
                     Box::new(adapter),
+                    no_preflight,
                     PyDict::new(py).unbind(),
                     asynchronous,
                 )
@@ -1291,6 +1316,87 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
         });
     }
 
+    /// The rejection a preflight raised, kept so a test can check the caller receives that
+    /// exact object. A `Preflight` is a plain `fn`, so it cannot capture one itself.
+    static REJECTION: Mutex<Option<Py<PyBaseException>>> = Mutex::new(None);
+
+    fn rejecting_preflight(py: Python<'_>, _: &Bound<'_, PyDict>) -> PyResult<()> {
+        let error = PyValueError::new_err("over budget");
+        *REJECTION.lock().unwrap() = Some(error.value(py).clone().unbind());
+        Err(error)
+    }
+
+    fn inheriting_preflight(_: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<()> {
+        arguments.set_item("api_key", "inherited")
+    }
+
+    #[test]
+    fn a_preflight_rejection_is_the_callers_error_and_the_machine_never_starts() {
+        let _guard = PYTHON_GLOBALS
+            .lock()
+            .unwrap_or_else(|error| error.into_inner());
+        crate::initialize_python();
+        Python::attach(|py| {
+            install_lifecycle_module(py);
+            for asynchronous in [false, true] {
+                let (result, log) = run_preflighted(
+                    py,
+                    success_machine(),
+                    SyntheticHost {
+                        log: Log::default(),
+                        op: OpScript::Answer,
+                        classifier_fails: false,
+                    },
+                    AdapterScript::Plain,
+                    rejecting_preflight,
+                    asynchronous,
+                );
+                let error = result.unwrap_err();
+                let raised = REJECTION.lock().unwrap().take().unwrap();
+                assert!(error.value(py).is(&raised));
+                assert_eq!(
+                    log,
+                    [
+                        "started",
+                        "begin",
+                        "failed:Host:over budget",
+                        "adapter.close",
+                        "host.close"
+                    ]
+                );
+            }
+        });
+    }
+
+    #[test]
+    fn the_host_projects_from_the_keyword_view_the_preflight_rewrote() {
+        let _guard = PYTHON_GLOBALS
+            .lock()
+            .unwrap_or_else(|error| error.into_inner());
+        crate::initialize_python();
+        Python::attach(|py| {
+            install_lifecycle_module(py);
+            for asynchronous in [false, true] {
+                let (result, _) = run_preflighted(
+                    py,
+                    success_machine(),
+                    SyntheticHost {
+                        log: Log::default(),
+                        op: OpScript::Answer,
+                        classifier_fails: false,
+                    },
+                    AdapterScript::Plain,
+                    inheriting_preflight,
+                    asynchronous,
+                );
+                assert_eq!(
+                    result.unwrap().extract::<String>(py).unwrap(),
+                    "project:2|sign|rewritten"
+                );
+            }
+        });
+    }
+
     #[test]
     fn the_adapters_finalized_response_is_what_the_call_returns_and_reports() {
         let _guard = PYTHON_GLOBALS
@@ -1412,6 +1518,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
                 success_machine(),
                 host,
                 Box::new(adapter),
+                no_preflight,
                 PyDict::new(py).unbind(),
                 false,
             )

litellm-rust/crates/host-python/src/lib.rs

@@ -15,7 +15,8 @@ mod handle;
 mod marshal;
 
 pub use adapter::{
-    InvokeError, LifecycleEvent, LifecycleStep, ProtocolHost, PythonLifecycle, missing_state,
+    InvokeError, LifecycleEvent, LifecycleStep, Preflight, ProtocolHost, PythonLifecycle,
+    missing_state,
 };
 pub use argument::lookup;
 pub use callable::wrap_failure;

litellm-rust/crates/python-bridge/Cargo.toml

@@ -55,6 +55,7 @@ pyo3-async-runtimes.workspace = true
 reqwest.workspace = true
 redis = { version = "1.7.0", features = ["tls-rustls"] }
 serde_json.workspace = true
+strum.workspace = true
 veil.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = ["rt", "sync"] }

litellm-rust/crates/python-bridge/preflight_contract.json

@@ -0,0 +1,10 @@
+{
+  "credential_list": [],
+  "warn_unknown_credential": [
+    "name",
+    "loaded"
+  ],
+  "check_limits": [
+    "kwargs"
+  ]
+}

litellm-rust/crates/python-bridge/src/lib.rs

@@ -6,6 +6,7 @@ mod errors;
 mod http;
 mod logger;
 mod marshal;
+mod preflight;
 mod python_settings;
 mod routes;
 mod secrets;

litellm-rust/crates/python-bridge/src/preflight.rs

@@ -1,9 +1,51 @@
+//! The SDK's request policy the driver runs on every route's keyword view before the host
+//! projects from it: credential-name inheritance from `litellm.credential_list`, then the
+//! budget and retry-count limits. It is the `@client` prologue after `function_setup` and the
+//! deployment hook, and belongs to no callback contract.
+
 use pyo3::{
     prelude::*,
     types::{PyDict, PyList},
 };
+use strum::{IntoStaticStr, VariantArray};
 
-use crate::python::Wrapper;
+const MODULE: &str = "litellm.rust_bridge.preflight";
+
+/// The litellm globals the preflight still reads through Python. `preflight_contract.json`
+/// pins each function's parameters on both sides.
+#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)]
+pub(crate) enum PythonPreflight {
+    #[strum(serialize = "credential_list")]
+    CredentialList,
+    #[strum(serialize = "warn_unknown_credential")]
+    WarnUnknownCredential,
+    #[strum(serialize = "check_limits")]
+    CheckLimits,
+}
+
+impl PythonPreflight {
+    fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult<Bound<'py, PyAny>>
+    where
+        A: pyo3::call::PyCallArgs<'py>,
+    {
+        py.import(MODULE)?.getattr(<&str>::from(self))?.call1(args)
+    }
+}
+
+#[cfg(test)]
+pub(crate) const PYTHON_CONTRACT: &str = include_str!("../preflight_contract.json");
+
+/// Rewrites `arguments` in place, in the order the Python wrapper runs: credentials first,
+/// so the limits see the same view the provider request is built from.
+pub(crate) fn sdk_preflight(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<()> {
+    inherit_credentials(py, arguments, || {
+        Ok(PythonPreflight::CredentialList
+            .call(py, ())?
+            .cast_into::<PyList>()?)
+    })?;
+    PythonPreflight::CheckLimits.call(py, (arguments,))?;
+    Ok(())
+}
 
 struct CredentialEntry<'py>(Bound<'py, PyAny>);
 
@@ -17,22 +59,6 @@ impl<'py> CredentialEntry<'py> {
     }
 }
 
-pub fn prepare<'py>(
-    py: Python<'py>,
-    kwargs: &Bound<'py, PyDict>,
-    logger: &crate::PythonLogger,
-) -> PyResult<Bound<'py, PyDict>> {
-    let arguments = kwargs.copy()?;
-    arguments.set_item("litellm_logging_obj", logger.object(py))?;
-    inherit_credentials(py, &arguments, || {
-        Ok(Wrapper::CredentialList
-            .call(py, ())?
-            .cast_into::<PyList>()?)
-    })?;
-    Wrapper::CheckLimits.call(py, (&arguments,))?;
-    Ok(arguments)
-}
-
 fn inherit_credentials<'py>(
     py: Python<'py>,
     arguments: &Bound<'py, PyDict>,
@@ -54,7 +80,7 @@ fn inherit_credentials<'py>(
         .map(|credential| CredentialEntry(credential).name())
         .collect::<PyResult<Vec<_>>>()?;
     let Some(index) = names.iter().position(|name| *name == requested) else {
-        Wrapper::WarnUnknownCredential.call(py, (requested, names.len()))?;
+        PythonPreflight::WarnUnknownCredential.call(py, (requested, names.len()))?;
         return Ok(());
     };
     let selected = CredentialEntry(credentials.get_item(index)?);
@@ -71,7 +97,42 @@ fn inherit_credentials<'py>(
 
 #[cfg(test)]
 mod tests {
+    use std::collections::BTreeSet;
+    use std::sync::Mutex;
+
     use super::*;
+    use strum::VariantArray;
+
+    /// Tests share one interpreter, and the stub module below is global state, so the
+    /// tests that install it run one at a time.
+    static PREFLIGHT_MODULE: Mutex<()> = Mutex::new(());
+
+    /// A fresh stand-in for `litellm.rust_bridge.preflight` that records every call, then
+    /// `script` run against it with the module bound as `preflight`.
+    fn preflight_stubs<'py>(py: Python<'py>, script: &std::ffi::CStr) -> Bound<'py, PyDict> {
+        let locals = PyDict::new(py);
+        py.run(
+            c"
+import sys
+import types
+
+for name in ('litellm', 'litellm.rust_bridge'):
+    sys.modules.setdefault(name, types.ModuleType(name))
+preflight = types.ModuleType('litellm.rust_bridge.preflight')
+preflight.warnings = []
+preflight.checked = []
+preflight.credential_list = lambda: []
+preflight.warn_unknown_credential = lambda name, loaded: preflight.warnings.append((name, loaded))
+preflight.check_limits = lambda kwargs: preflight.checked.append(kwargs)
+sys.modules['litellm.rust_bridge.preflight'] = preflight
+",
+            Some(&locals),
+            Some(&locals),
+        )
+        .unwrap();
+        py.run(script, Some(&locals), Some(&locals)).unwrap();
+        locals
+    }
 
     fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
         let locals = PyDict::new(py);
@@ -312,4 +373,110 @@ arguments = {'litellm_credential_name': 'ocr-test'}
             }
         });
     }
+
+    #[test]
+    fn every_borrowed_function_is_in_the_python_contract() {
+        Python::initialize();
+        Python::attach(|py| {
+            let contract = litellm_host_python::json_loads(py, PYTHON_CONTRACT.as_bytes()).unwrap();
+            let declared: BTreeSet<String> = contract
+                .bind(py)
+                .cast::<PyDict>()
+                .unwrap()
+                .keys()
+                .extract()
+                .map(|names: Vec<String>| names.into_iter().collect())
+                .unwrap();
+            let called: BTreeSet<String> = PythonPreflight::VARIANTS
+                .iter()
+                .map(|&function| <&str>::from(function).to_owned())
+                .collect();
+            assert_eq!(
+                called.len(),
+                PythonPreflight::VARIANTS.len(),
+                "a function is borrowed twice"
+            );
+            assert_eq!(called, declared);
+        });
+    }
+
+    #[test]
+    fn an_unknown_name_is_reported_with_the_loaded_count_and_leaves_the_arguments_alone() {
+        let _guard = PREFLIGHT_MODULE
+            .lock()
+            .unwrap_or_else(|error| error.into_inner());
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = preflight_stubs(
+                py,
+                c"
+class Credential:
+    credential_name = 'listed'
+    credential_values = {'api_key': 'listed-key'}
+preflight.credential_list = lambda: [Credential(), Credential()]
+arguments = {'litellm_credential_name': 'missing'}
+",
+            );
+            sdk_preflight(py, &argument_dict(&locals)).unwrap();
+            py.run(
+                c"
+assert arguments == {'litellm_credential_name': 'missing'}, arguments
+assert preflight.warnings == [('missing', 2)], preflight.warnings
+assert preflight.checked == [arguments]
+",
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+        });
+    }
+
+    #[test]
+    fn limits_are_checked_on_the_arguments_after_credentials_are_inherited() {
+        let _guard = PREFLIGHT_MODULE
+            .lock()
+            .unwrap_or_else(|error| error.into_inner());
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = preflight_stubs(
+                py,
+                c"
+class Credential:
+    credential_name = 'ocr-test'
+    credential_values = {'api_key': 'inherited'}
+preflight.credential_list = lambda: [Credential()]
+rejection = RuntimeError('Max retries per request hit!')
+def check_limits(arguments):
+    preflight.checked.append(dict(arguments))
+    raise rejection
+preflight.check_limits = check_limits
+arguments = {'litellm_credential_name': 'ocr-test'}
+",
+            );
+            let error = sdk_preflight(py, &argument_dict(&locals)).unwrap_err();
+            assert!(
+                error
+                    .value(py)
+                    .is(locals.get_item("rejection").unwrap().unwrap())
+            );
+            py.run(
+                c"
+assert preflight.checked == [{'litellm_credential_name': 'ocr-test', 'api_key': 'inherited'}], preflight.checked
+assert arguments['api_key'] == 'inherited'
+",
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+        });
+    }
+
+    fn argument_dict<'py>(locals: &Bound<'py, PyDict>) -> Bound<'py, PyDict> {
+        locals
+            .get_item("arguments")
+            .unwrap()
+            .unwrap()
+            .cast_into::<PyDict>()
+            .unwrap()
+    }
 }

litellm-rust/crates/python-bridge/src/routes/messages/mod.rs

@@ -33,6 +33,7 @@ fn run_messages(
         PublicCall::capture(&request, &args, &kwargs)?,
         crate::logger::LoggedMachine::new(messages_machine(secrets)),
         MessagesPythonHost::new(request.unbind()),
+        crate::preflight::sdk_preflight,
         asynchronous,
     )
 }

litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs

@@ -70,6 +70,7 @@ fn run_ocr(
         PublicCall::capture(&request, &args, &kwargs)?,
         crate::logger::LoggedMachine::new(ocr_machine(client)),
         OcrPythonHost::new(request.unbind()),
+        crate::preflight::sdk_preflight,
         asynchronous,
     )
 }

litellm/rust_bridge/callbacks_legacy_python.py

@@ -22,7 +22,6 @@ from typing import (
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging
-    from litellm.types.utils import CredentialItem
 
 
 class MetadataUpdater(Protocol):
@@ -72,21 +71,6 @@ def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> Call
     return call_setup
 
 
-def check_limits(kwargs: Mapping[str, object]) -> None:
-    from litellm import (
-        BudgetExceededError,
-        _current_cost,  # pyright: ignore[reportPrivateUsage]  # shared SDK budget counter has no public accessor
-        max_budget,
-        num_retries_per_request,
-    )
-    from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
-
-    if max_budget and _current_cost > max_budget:
-        raise BudgetExceededError(current_cost=_current_cost, max_budget=max_budget)
-    if max_retries_per_request_hit(kwargs, num_retries_per_request):
-        raise RuntimeError("Max retries per request hit!")
-
-
 def finalize(
     response: object,
     logger: Logging,
@@ -299,22 +283,6 @@ def is_internal_call() -> bool:
     return internal.get()
 
 
-def credential_list() -> list[CredentialItem]:
-    from litellm import credential_list as credentials
-
-    return credentials
-
-
-def warn_unknown_credential(name: str, loaded: int) -> None:
-    from litellm._logging import verbose_logger
-
-    verbose_logger.warning(
-        "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it",
-        name,
-        loaded,
-    )
-
-
 def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitable[object]:
     from litellm import utils
 

litellm/rust_bridge/preflight.py

@@ -0,0 +1,45 @@
+"""The SDK request policy the native driver runs before a route's host projects.
+
+These are the `@client` prologue steps after `function_setup` and the deployment hook:
+credential-name inheritance and the budget and retry-count limits. Rust owns the
+inheritance itself; it borrows only the globals below.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from litellm.types.utils import CredentialItem
+
+
+def credential_list() -> list[CredentialItem]:
+    from litellm import credential_list as credentials
+
+    return credentials
+
+
+def warn_unknown_credential(name: str, loaded: int) -> None:
+    from litellm._logging import verbose_logger
+
+    verbose_logger.warning(
+        "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it",
+        name,
+        loaded,
+    )
+
+
+def check_limits(kwargs: Mapping[str, object]) -> None:
+    from litellm import (
+        BudgetExceededError,
+        _current_cost,  # pyright: ignore[reportPrivateUsage]  # shared SDK budget counter has no public accessor
+        max_budget,
+        num_retries_per_request,
+    )
+    from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
+
+    if max_budget and _current_cost > max_budget:
+        raise BudgetExceededError(current_cost=_current_cost, max_budget=max_budget)
+    if max_retries_per_request_hit(kwargs, num_retries_per_request):
+        raise RuntimeError("Max retries per request hit!")

tests/unit/rust_bridge/test_callbacks_legacy_python.py

@@ -8,11 +8,10 @@ from typing import Final
 import pytest
 from pydantic import TypeAdapter
 
-import litellm
 from litellm._internal_context import is_internal_call
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.rust_bridge import callbacks_legacy_python as legacy
-from litellm.rust_bridge.callbacks_legacy_python import check_limits, failure_handler, setup
+from litellm.rust_bridge.callbacks_legacy_python import failure_handler, setup
 
 _OCR_KWARGS: Final = MappingProxyType(
     {
@@ -22,33 +21,6 @@ _OCR_KWARGS: Final = MappingProxyType(
 )
 
 
-@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
-@pytest.mark.parametrize(
-    "cap, request_retry_count, refused",
-    [(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)],
-    ids=[
-        "cap-above-four-reached",
-        "cap-above-four-not-reached",
-        "first-attempt-passes-cap-of-zero",
-        "cap-of-zero-refuses-first-retry",
-    ],
-)
-def test_check_limits_reads_request_retry_count(
-    monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, request_retry_count: int, refused: bool
-) -> None:
-    monkeypatch.setattr(litellm, "num_retries_per_request", cap)
-    monkeypatch.setattr(litellm, "max_budget", None)
-    kwargs: Final = {
-        "model": "mistral/mistral-ocr-latest",
-        metadata_key: {"request_retry_count": request_retry_count},
-    }
-    if refused:
-        with pytest.raises(RuntimeError, match="Max retries per request hit!"):
-            check_limits(kwargs)
-    else:
-        check_limits(kwargs)
-
-
 def _supplied_logger() -> Logging:
     return Logging(
         model="mistral/mistral-ocr-latest",

tests/unit/rust_bridge/test_preflight.py

@@ -0,0 +1,63 @@
+import inspect
+from collections.abc import Callable
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.rust_bridge import preflight
+from litellm.rust_bridge.preflight import check_limits
+
+
+@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
+@pytest.mark.parametrize(
+    "cap, request_retry_count, refused",
+    [(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)],
+    ids=[
+        "cap-above-four-reached",
+        "cap-above-four-not-reached",
+        "first-attempt-passes-cap-of-zero",
+        "cap-of-zero-refuses-first-retry",
+    ],
+)
+def test_check_limits_reads_request_retry_count(
+    monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, request_retry_count: int, refused: bool
+) -> None:
+    monkeypatch.setattr(litellm, "num_retries_per_request", cap)
+    monkeypatch.setattr(litellm, "max_budget", None)
+    kwargs: Final = {
+        "model": "mistral/mistral-ocr-latest",
+        metadata_key: {"request_retry_count": request_retry_count},
+    }
+    if refused:
+        with pytest.raises(RuntimeError, match="Max retries per request hit!"):
+            check_limits(kwargs)
+    else:
+        check_limits(kwargs)
+
+
+def test_check_limits_refuses_a_call_over_the_budget(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "num_retries_per_request", None)
+    monkeypatch.setattr(litellm, "max_budget", 1.0)
+    monkeypatch.setattr(litellm, "_current_cost", 1.5)
+    with pytest.raises(litellm.BudgetExceededError):
+        check_limits({"model": "mistral/mistral-ocr-latest"})
+
+
+CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/preflight_contract.json"
+_SHIMS: Final[MappingProxyType[str, Callable[..., object]]] = MappingProxyType(
+    {
+        "credential_list": preflight.credential_list,
+        "warn_unknown_credential": preflight.warn_unknown_credential,
+        "check_limits": preflight.check_limits,
+    }
+)
+
+
+def test_the_rust_contract_matches_the_shim_signatures() -> None:
+    contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text())
+
+    assert contract == {name: list(inspect.signature(_SHIMS[name]).parameters) for name in contract}