fix(rust): initialize callback lifecycle runtime

This commit is contained in:
Yujong Lee 2026-09-05 21:45:18 -07:00
parent e4f8aa57af
commit 7780d877bb
4 changed files with 7 additions and 1 deletions

View file

@ -162,6 +162,8 @@ mod _native {
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
super::errors::register(module)?;
litellm_python_interop::callback_runtime::register(module)?;
super::callback_bindings::register(module)?;
super::routes::register(module)?;
module.add_class::<super::ResponsesWebSocketConnection>()?;
module.add_function(wrap_pyfunction!(

View file

@ -398,6 +398,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
fake_module = types.ModuleType("litellm.rust_bridge._native")
fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
fake_module.ready_endpoints = {"ocr": {"callbacks"}} # type: ignore[attr-defined]
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge.bindings"),
"get_native_bridge",

View file

@ -113,6 +113,7 @@ async def test_connection_forwards_session_callback_adapter() -> None:
cls,
request: NativeResponsesWebSocketRequest,
*,
options: object,
context: NativeRequestContext,
callback_adapter: object | None = None,
) -> _FakeNativeConnection:
@ -167,7 +168,7 @@ async def test_connection_dispatch_cleans_up_without_reconnecting(native, sessio
class Native:
@classmethod
async def connect(cls, request, *, options, context):
async def connect(cls, request, *, options, context, callback_adapter=None):
connections.append("native")
assert options.custom_llm_provider == "azure"
return native_socket

View file

@ -34,6 +34,7 @@ class SyncBridge:
*,
options: NativeRequestOptions,
context: NativeRequestContext,
callback_adapter: object | None = None,
) -> dict[str, object]:
self.calls.append(
{
@ -53,6 +54,7 @@ class AsyncBridge:
*,
options: NativeRequestOptions,
context: NativeRequestContext,
callback_adapter: object | None = None,
) -> dict[str, object]:
return {"text": "async"}