diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 14e1bfe91ca..dbcbd9bda92 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -59,6 +59,34 @@ pub(super) struct FacadeGuard { } impl ObjectGuard { + fn class_behaviors(class: &Bound<'_, PyType>) -> PyResult)>> { + let py = class.py(); + let builtins = py.import("builtins")?; + let property_type = builtins.getattr("property")?; + let staticmethod_type = builtins.getattr("staticmethod")?; + let classmethod_type = builtins.getattr("classmethod")?; + class + .getattr("__dict__")? + .call_method0("items")? + .try_iter()? + .map(|item| { + let item = item?; + let (name, value): (String, Py) = item.extract()?; + let value_bound = value.bind(py); + let is_behavior = value_bound.is_callable() + || value_bound.is_instance(&property_type)? + || value_bound.is_instance(&staticmethod_type)? + || value_bound.is_instance(&classmethod_type)?; + Ok(is_behavior.then_some((name, value))) + }) + .filter_map(|result| match result { + Ok(Some(attribute)) => Some(Ok(attribute)), + Ok(None) => None, + Err(error) => Some(Err(error)), + }) + .collect() + } + fn capture( py: Python<'_>, object: &Bound<'_, PyAny>, @@ -71,12 +99,7 @@ impl ObjectGuard { .iter() .map(|class| { let class = class.cast_into::()?; - let attributes = class - .getattr("__dict__")? - .call_method0("items")? - .try_iter()? - .map(|item| item?.extract::<(String, Py)>()) - .collect::>>()?; + let attributes = Self::class_behaviors(&class)?; Ok(ClassGuard { class: class.unbind(), attributes, @@ -129,15 +152,21 @@ impl ObjectGuard { } let instance = object.getattr("__dict__")?.cast_into::()?; for (class, expected) in mro.iter().zip(&self.classes) { + let class = class.cast_into::()?; if !class.is(expected.class.bind(py)) { return Ok(false); } - let attributes = class.getattr("__dict__")?; - if attributes.len()? != expected.attributes.len() { + let attributes = Self::class_behaviors(&class)?; + if attributes.len() != expected.attributes.len() { return Ok(false); } - for (name, value) in &expected.attributes { - if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + for ((name, value), (expected_name, expected_value)) in + attributes.iter().zip(&expected.attributes) + { + if name != expected_name + || instance.contains(name)? + || !value.bind(py).is(expected_value.bind(py)) + { return Ok(false); } } @@ -342,3 +371,51 @@ pub(super) fn resolve( } handle.service().map(Some) } + +#[cfg(test)] +mod tests { + use super::ObjectGuard; + use pyo3::{prelude::*, types::PyDict}; + + #[test] + fn class_data_shadowing_is_ignored_but_method_mutations_are_rejected() { + Python::initialize(); + Python::attach(|py| { + let namespace = PyDict::new(py); + py.run( + c"class Example:\n data = 1\n def method(self):\n return 1\nobject = Example()\nobject.data = 2", + None, + Some(&namespace), + ) + .unwrap(); + let object = namespace.get_item("object").unwrap().unwrap(); + let guard = ObjectGuard::capture(py, &object, &[]).unwrap(); + + assert!(guard.matches(py, &object).unwrap()); + + py.run(c"object.method = lambda: 2", None, Some(&namespace)) + .unwrap(); + assert!(!guard.matches(py, &object).unwrap()); + }); + } + + #[test] + fn class_method_replacement_is_rejected() { + Python::initialize(); + Python::attach(|py| { + let namespace = PyDict::new(py); + py.run( + c"class Example:\n def method(self):\n return 1\nobject = Example()", + None, + Some(&namespace), + ) + .unwrap(); + let object = namespace.get_item("object").unwrap().unwrap(); + let guard = ObjectGuard::capture(py, &object, &[]).unwrap(); + + py.run(c"Example.method = lambda self: 2", None, Some(&namespace)) + .unwrap(); + assert!(!guard.matches(py, &object).unwrap()); + }); + } +} diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index d33eaec0c27..02c159825bc 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -573,7 +573,9 @@ def test_qdrant_semantic_facade_binds_native_and_shares_entries( assert binding.kind == "native" assert binding.lookup(semantic_request("python-key", messages)) == {"id": "py"} binding.store(semantic_request("native-key", messages), {"id": "native"}) - assert facade.cache.get_cache("native-key", messages=messages) == {"id": "native"} + python_value: Final = facade.cache.get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] assert binding.lookup(semantic_request("native-key", unrelated)) is None assert facade.cache.get_cache("native-key", messages=unrelated) is None @@ -602,9 +604,20 @@ async def test_qdrant_semantic_async_parity( {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, messages=messages, ) - assert await binding.async_lookup(semantic_request("python-key", messages)) == {"id": "py"} + + async def lookup_after_commit() -> object: + for _ in range(20): + value: Final = await binding.async_lookup(semantic_request("python-key", messages)) + if value is not None: + return value + await asyncio.sleep(0.1) + return None + + assert await lookup_after_commit() == {"id": "py"} await binding.async_store(semantic_request("native-key", messages), {"id": "native"}) - assert await facade.cache.async_get_cache("native-key", messages=messages) == {"id": "native"} + python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( @@ -673,7 +686,9 @@ def test_qdrant_semantic_ignores_request_expiry( ) time.sleep(1.2) assert binding.lookup(semantic_request("persistent-key", messages)) == {"id": "persistent"} - assert facade.cache.get_cache("persistent-key", messages=messages) == {"id": "persistent"} + python_value: Final = facade.cache.get_cache("persistent-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "persistent"} def test_qdrant_semantic_mutation_and_projection_fallback(