scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

refactor(rust-bridge): require the native extension (#48852)

Browse source 
* refactor(rust-bridge): require the native extension

_native ships in every wheel and the sdist will not install without it, yet the bridge loaded it optionally and carried None paths that could fail open for fork safety, exception mapping and OTLP errors. Delete loader.py and import _native directly. Tests inject a fake module through NativeBinding(source=...) or with_source() instead of patching the loader

The install smoke checks now verify the compiled module, and tests fail instead of skipping when _native is missing. The interactions spec canary installs the project so it has the extension

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(ci): build Rust bridge in Windows smoke tests

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yujonglee 2026-10-28 13:16:56 -07:00 • committed by GitHub
parent 99ac0d3e39
commit 1f68cc5649
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
34 changed files with 144 additions and 345 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
.github/workflows/interactions-spec-canary.yml vendored
View file

@ -19,7 +19,7 @@ jobs:
name: interactions-spec-canary
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 15
timeout-minutes: 30
permissions:
contents: read
outputs:
@ -39,8 +39,11 @@ jobs:
with:
version: "0.10.9"
- name: Set up the Rust bridge
uses: ./.github/actions/rust-bridge
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
- name: Run the compliance tests against Google's live spec
id: compliance

7
.github/workflows/test-merge-smoke.yml vendored
View file

@ -132,7 +132,7 @@ jobs:
windows-sdk:
name: windows-sdk
runs-on: windows-2025
timeout-minutes: 10
timeout-minutes: 30
env:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
defaults:
@ -154,8 +154,11 @@ jobs:
with:
version: "0.10.9"
- name: Set up the Rust bridge
uses: ./.github/actions/rust-bridge
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --no-install-workspace --group dev --python 3.11
run: .github/scripts/uv_sync_with_retries.sh --frozen --group dev --python 3.11
- name: Run the Windows tests
run: uv run --no-sync python -m pytest --tb=short -v tests/windows_tests

10
.github/workflows/test-unit.yml vendored
View file

@ -1306,6 +1306,8 @@ jobs:
done
openai-sdk-compat:
name: OpenAI ${{ matrix.openai-version }} compatibility
needs: rust-bridge
if: ${{ !cancelled() }}
permissions:
contents: read
pull-requests: read
@ -1342,10 +1344,18 @@ jobs:
with:
version: "0.10.9"
- name: Set up the Rust bridge
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 5
uses: ./.github/actions/rust-bridge
with:
artifact: ${{ needs.rust-bridge.outputs.artifact }}
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace
cp rust-bridge-dist/litellm/rust_bridge/_native.* litellm/rust_bridge/
uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z
uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__'

2
gateway/Dockerfile
View file

@ -153,7 +153,7 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
python -c "import litellm, litellm.rust_bridge._native; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__"
USER nonroot

2
litellm-rust/crates/python-bridge/src/logger/tests.rs
View file

@ -187,6 +187,7 @@ fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication
locals
.set_item("http_warning", wrap_pyfunction!(http_warning, py).unwrap())
.unwrap();
locals.set_item("native", crate::native_module(py)).unwrap();
let importable = py
.eval(
c"__import__('importlib.util', fromlist=['util']).find_spec('dotenv') is not None",
@ -205,6 +206,7 @@ import asyncio
import logging
import sys
sys.path.insert(0, repo_root)
sys.modules['litellm.rust_bridge._native'] = native
import litellm
from litellm._logging import verbose_logger, session_id_var, trace_id_var

2
litellm/litellm_core_utils/README.md
View file

@ -11,4 +11,4 @@ Core files:
- `duration_parser.py`: code for parsing durations - e.g. "1d", "1mo", "10s"
- `api_route_to_call_types.py`: mapping of API routes to their corresponding CallTypes (e.g., `/chat/completions` -> [acompletion, completion])
Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Missing native bindings or unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer.
Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer.

6
litellm/rust_bridge/AGENTS.md
View file

@ -46,6 +46,12 @@
- A failure after execution starts is terminal and never triggers fallback
- Creating an unstarted native coroutine acquires no clients or credentials and captures no execution context
# Native extension
- `_native` ships in every wheel and the sdist builds it, so it is always importable
- Import it with `from litellm.rust_bridge import _native` and use the attributes `_native.pyi` declares
- Never add a `None` branch, `getattr` default or stand-in class for a missing extension
# Verification
- Check cross-language behavior against a freshly installed extension, covering unstarted-call release, setup ordering, exception identity, sync and async streams, cancellation and GC

7
litellm/rust_bridge/__init__.py
View file

@ -1,10 +1,5 @@
"""LiteLLM Rust bridge package."""
from litellm.rust_bridge.configuration import rust
from litellm.rust_bridge.loader import (
get_native_bridge,
native_bridge_available,
reset_native_bridge_cache,
)
__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "rust"]
__all__ = ["rust"]

26
litellm/rust_bridge/bindings.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
from collections.abc import Callable
from typing import Final, Generic, TypeVar
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge import _native
BindingT = TypeVar("BindingT")
@ -18,32 +18,24 @@ _UNSET: Final = _Unset()
class NativeBinding(Generic[BindingT]):
"""Resolve one native attribute with an explicit, resettable test override."""
def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None:
def __init__(
self, attribute: str, *, validate: Callable[[object], BindingT | None], source: object = _native
) -> None:
self._attribute: Final = attribute
self._validate: Final = validate
self._source: Final = source
self._override: BindingT | None | _Unset = _UNSET
def load(self) -> BindingT | None:
if not isinstance(self._override, _Unset):
return self._override
native: Final = get_native_bridge()
if native is None:
return None
return self._validate(getattr(native, self._attribute, None))
return self._validate(getattr(self._source, self._attribute, None))
def with_source(self, source: object) -> NativeBinding[BindingT]:
return NativeBinding(self._attribute, validate=self._validate, source=source)
def override(self, value: BindingT | None) -> None:
self._override = value
def reset(self) -> None:
self._override = _UNSET
def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None:
native: Final = get_native_bridge()
if native is None:
return None
declined: Final = getattr(native, "RustBridgeDeclined", None)
upstream: Final = getattr(native, "RustUpstreamError", None)
if not isinstance(declined, type) or not isinstance(upstream, type):
return None
return declined, upstream

37
litellm/rust_bridge/clickhouse.py
View file

@ -1,12 +1,10 @@
import os
from collections.abc import Awaitable, Mapping, Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, runtime_checkable
from pydantic import ConfigDict, TypeAdapter
from typing import Final
from litellm.constants import DEFAULT_AGENT_TRACING_RETENTION_DAYS, DEFAULT_CLICKHOUSE_DATABASE
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge import _native
@dataclass(frozen=True, slots=True, repr=False)
@ -16,27 +14,6 @@ class SpendStorageConfig:
retention_days: int = DEFAULT_AGENT_TRACING_RETENTION_DAYS
class NativeSpendConfig(Protocol):
def __init__(self, database: str, url: str, retention_days: int) -> None: ...
class NativeSpendStorage(Protocol):
def __init__(self, config: NativeSpendConfig) -> None: ...
def ensure_schema(self) -> Awaitable[None]: ...
def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Awaitable[None]: ...
@runtime_checkable
class NativeSpendBridge(Protocol):
NativeClickHouseSpendConfig: type[NativeSpendConfig]
NativeClickHouseSpendStorage: type[NativeSpendStorage]
_NATIVE: Final[TypeAdapter[NativeSpendBridge]] = TypeAdapter(
NativeSpendBridge, config=ConfigDict(arbitrary_types_allowed=True)
)
def spend_storage_config(environ: Mapping[str, str] = os.environ) -> SpendStorageConfig:
url: Final = environ.get("CLICKHOUSE_URL")
if not url:
@ -53,12 +30,8 @@ def spend_storage_config(environ: Mapping[str, str] = os.environ) -> SpendStorag
class ClickHouseSpendStorage:
def __init__(self, config: SpendStorageConfig) -> None:
bridge: Final = get_native_bridge()
if bridge is None:
raise RuntimeError("ClickHouse spend logging requires the Rust extension")
native: Final = _NATIVE.validate_python(bridge)
self._native: Final = native.NativeClickHouseSpendStorage(
native.NativeClickHouseSpendConfig(config.database, config.url, config.retention_days)
self._native: Final = _native.NativeClickHouseSpendStorage(
_native.NativeClickHouseSpendConfig(config.database, config.url, config.retention_days)
)
async def ensure_schema(self) -> None:

20
litellm/rust_bridge/fork_guard.py
View file

@ -11,32 +11,22 @@ site, so the runtime can never start there. Workers forked from it are unaffecte
from __future__ import annotations
from collections.abc import Callable
from typing import Final
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge import _native
class NativeStateStartedBeforeFork(RuntimeError):
pass
class _NeverRaised(RuntimeError):
"""Stands in for a native exception when the extension is unavailable or predates it."""
ForkedAfterNativeRuntimeStarted: Final = _native.ForkedAfterNativeRuntimeStarted
ProcessReservedForForking: Final = _native.ProcessReservedForForking
_native: Final = get_native_bridge()
ForkedAfterNativeRuntimeStarted: Final[type[RuntimeError]] = getattr(
_native, "ForkedAfterNativeRuntimeStarted", _NeverRaised
)
ProcessReservedForForking: Final[type[RuntimeError]] = getattr(_native, "ProcessReservedForForking", _NeverRaised)
def reserve_process_for_forking(where: str) -> None:
def reserve_process_for_forking(where: str, reserve: Callable[[], None] = _native.reserve_process_for_forking) -> None:
"""Forbid native routes in this process. Raises if one already ran here."""
native: Final = get_native_bridge()
reserve: Final = getattr(native, "reserve_process_for_forking", None)
if not callable(reserve):
return
try:
reserve()
except RuntimeError as error:

35
litellm/rust_bridge/loader.py
View file

@ -1,35 +0,0 @@
"""Loader for the packaged LiteLLM Rust extension."""
from __future__ import annotations
from types import ModuleType
from typing import Final
_BRIDGE_SENTINEL: Final = object()
_cached_bridge: ModuleType | None | object = _BRIDGE_SENTINEL
def get_native_bridge() -> ModuleType | None:
"""Return the packaged Rust extension, or ``None`` when unavailable."""
global _cached_bridge
if _cached_bridge is not _BRIDGE_SENTINEL:
return _cached_bridge if isinstance(_cached_bridge, ModuleType) else None
try:
from litellm.rust_bridge import _native
except ImportError:
_cached_bridge = None
return None
_cached_bridge = _native
return _native
def reset_native_bridge_cache() -> None:
"""Forget the cached extension so the next lookup reimports it from disk."""
global _cached_bridge
_cached_bridge = _BRIDGE_SENTINEL
def native_bridge_available() -> bool:
"""Whether the packaged Rust extension is importable."""
return get_native_bridge() is not None

25
litellm/rust_bridge/responses/websocket.py
View file

@ -2,32 +2,33 @@
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, Protocol
import httpx
from websockets.exceptions import ConnectionClosedOK
from litellm.rust_bridge import _native
from litellm.rust_bridge.host.timeouts import timeout_to_seconds
from litellm.rust_bridge.loader import get_native_bridge
class RustResponsesWebSocket(Protocol):
async def send_text(self, text: str) -> None: ...
def send_text(self, text: str) -> Awaitable[None]: ...
async def recv_text(self) -> str | None: ...
def recv_text(self) -> Awaitable[str | None]: ...
async def close(self) -> None: ...
def close(self) -> Awaitable[None]: ...
class RustResponsesWebSocketConnection(Protocol):
@classmethod
async def connect(
def connect(
cls,
url: str,
headers: dict[str, str],
timeout_seconds: float | None,
) -> RustResponsesWebSocket: ...
) -> Awaitable[RustResponsesWebSocket]: ...
class _Unset:
@ -53,16 +54,10 @@ def set_rust_responses_websocket(
_STATE.connection = connection
def load_rust_responses_websocket() -> RustResponsesWebSocketConnection | None:
def load_rust_responses_websocket() -> RustResponsesWebSocketConnection:
if _STATE.connection is not None:
return _STATE.connection
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
connection_type: Final[RustResponsesWebSocketConnection | None] = getattr(
native_bridge, "ResponsesWebSocketConnection", None
)
return connection_type
return _native.ResponsesWebSocketConnection
class _ConnectionAdapter:
@ -89,8 +84,6 @@ async def connect(
timeout: float | httpx.Timeout | None,
) -> _ConnectionAdapter | None:
connection_type: Final = load_rust_responses_websocket()
if connection_type is None:
return None
try:
connection: Final = await connection_type.connect(
url=url,

19
litellm/rust_bridge/runtime.py
View file

@ -7,7 +7,8 @@ from typing import Final, Generic, NoReturn, TypeAlias, TypeVar
from typing_extensions import assert_never
from litellm.exceptions import APIError
from litellm.rust_bridge.bindings import NativeBinding, native_exception_types
from litellm.rust_bridge import _native
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import RouteContext, Rules, decision
from litellm.rust_bridge.configuration import Decision
from litellm.rust_bridge.response_metadata import mark_rust_response
@ -153,15 +154,11 @@ def attempt(
) -> RustAttempt[ResultT]:
if native_call is None:
return RustUnavailable()
exceptions: Final = native_exception_types()
if exceptions is None:
return RustHandled(adapt(native_call()))
declined, upstream = exceptions
try:
value: Final = native_call()
except declined as error:
except _native.RustBridgeDeclined as error:
return RustDeclined(reason=_decline_reason(error))
except upstream as error:
except _native.RustUpstreamError as error:
_raise_upstream(error, context)
return RustHandled(adapt(value))
@ -174,15 +171,11 @@ async def aattempt(
) -> RustAttempt[ResultT]:
if native_call is None:
return RustUnavailable()
exceptions: Final = native_exception_types()
if exceptions is None:
return RustHandled(adapt(await native_call()))
declined, upstream = exceptions
try:
value: Final = await native_call()
except declined as error:
except _native.RustBridgeDeclined as error:
return RustDeclined(reason=_decline_reason(error))
except upstream as error:
except _native.RustUpstreamError as error:
_raise_upstream(error, context)
return RustHandled(adapt(value))

18
litellm/tracing/otlp_http.py
View file

@ -1,20 +1,9 @@
import json
from typing import Final, Protocol, runtime_checkable
from typing import Final
from pydantic import ConfigDict, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm.rust_bridge.loader import get_native_bridge
@runtime_checkable
class NativeOtlpError(Protocol):
def trace_encode_error(self, message: str) -> bytes: ...
_NATIVE: Final[TypeAdapter[NativeOtlpError]] = TypeAdapter(
NativeOtlpError, config=ConfigDict(arbitrary_types_allowed=True)
)
from litellm.rust_bridge import _native
class OTLPError(TypedDict):
@ -22,8 +11,7 @@ class OTLPError(TypedDict):
def encode_error(message: str) -> bytes:
native: Final = get_native_bridge()
return b"" if native is None else _NATIVE.validate_python(native).trace_encode_error(message)
return _native.trace_encode_error(message)
def encode_otlp_response(content_type: str | None, error: str | None = None) -> tuple[bytes, str]:

2
pyproject.toml
View file

@ -401,7 +401,7 @@ markers = [
"asyncio: mark test as an asyncio test",
"limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')",
"no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests",
"requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension",
"requires_rust_extension: public Python contract tests that run against the release-built wheel",
]
filterwarnings = [
# Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests)

20
tests/base_sdk_tests/check_base_sdk_install.py
View file

@ -8,6 +8,8 @@ the very class of undeclared-dependency bug this guards against.
import argparse
import importlib.util
from importlib.machinery import EXTENSION_SUFFIXES
from pathlib import Path
import re
import sys
import traceback
@ -89,6 +91,18 @@ def check_import() -> str:
return f"imported {distribution} {sdk_version}"
def check_native_extension() -> str:
import litellm
from litellm.rust_bridge import _native
native_path: Final = Path(_native.__file__).resolve()
package_path: Final = Path(litellm.__file__).resolve().parent
_require(native_path.parent == package_path / "rust_bridge", "native extension is outside the SDK package")
_require(any(str(native_path).endswith(suffix) for suffix in EXTENSION_SUFFIXES), "native module is not compiled")
_require(_native.trace_encode_error("ping") == b"\x12\x04ping", "native OTLP error encoding failed")
return f"loaded and exercised {native_path}"
def check_completion() -> str:
import litellm
@ -230,6 +244,7 @@ def check_bedrock_credential_resolution() -> str:
CHECKS: tuple[tuple[str, Callable[[], str]], ...] = (
("environment is base-only", check_environment_is_base_only),
("import litellm", check_import),
("native extension", check_native_extension),
("optional MCP installation guidance", check_mcp_install_guidance),
("chat completion", check_completion),
("chat completion with a function tool", check_completion_with_function_tool),
@ -250,8 +265,11 @@ def _run(check: Callable[[], str]) -> tuple[bool, str]:
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--profile", choices=("legacy", "core", "dependencies"), default="legacy")
parser.add_argument("--profile", choices=("legacy", "core", "dependencies", "native"), default="legacy")
profile = parser.parse_args().profile
if profile == "native":
print(check_native_extension())
return 0
for module in ("boto3", "botocore", "tokenizers", "huggingface_hub"):
present = importlib.util.find_spec(module) is not None
_require(present == (profile != "core"), f"{profile}: unexpected presence of {module}: {present}")

12
tests/test_litellm_rust/conftest.py
View file

@ -1,6 +1,4 @@
from collections.abc import AsyncIterator, Generator
from importlib import import_module
from typing import Final
import pytest
import pytest_asyncio
@ -23,16 +21,6 @@ def recording_server() -> Generator[RecordingServer]:
yield server
def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
try:
import_module("litellm.rust_bridge._native")
except ImportError:
skip: Final = pytest.mark.skip(reason="requires a compiled litellm.rust_bridge._native extension")
for item in items:
if "test_litellm_rust" in item.path.parts:
item.add_marker(skip)
@pytest.fixture
def isolated_azure_auth(monkeypatch: pytest.MonkeyPatch) -> None:
for name in (

8
tests/unit/litellm_core_utils/test_tokenizer.py
View file

@ -241,20 +241,14 @@ print("compatible")
@pytest.mark.parametrize("rust", (None, "0", "1"))
def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None:
def test_tokenization_stays_offline(tmp_path: Path, rust: str | None) -> None:
script: Final = """
import importlib.abc
import sys
sys.path.insert(0, sys.argv[1])
def reject_network(event, args):
if event == "socket.connect":
raise AssertionError("tokenizer attempted a network connection")
sys.addaudithook(reject_network)
class Block(importlib.abc.MetaPathFinder):
def find_spec(self, fullname, path=None, target=None):
if fullname == "litellm.rust_bridge._native":
raise ImportError("native extension is unavailable")
sys.meta_path.insert(0, Block())
import litellm
from litellm.rust_bridge.tokenizer import get_encoding
import tiktoken

11
tests/unit/proxy/proxy_server/test_exception_handlers.py
View file

@ -488,22 +488,17 @@ async def test_otel_unhandled_exception_handler_reraises_http_exception_invalid(
@pytest.mark.asyncio
@pytest.mark.parametrize("media_type", ["application/json", "application/x-protobuf"])
@pytest.mark.parametrize("root_path", ["", "/tenant-a"])
@pytest.mark.parametrize("native_available", [True, False])
@pytest.mark.parametrize("error", [
ProxyException("database credentials: secret", "auth_error", None, 401),
HTTPException(403, "database credentials: secret"),
])
async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native(
media_type: str, root_path: str, native_available: bool,
error: ProxyException | HTTPException, monkeypatch: pytest.MonkeyPatch,
async def test_otlp_auth_errors_hide_internal_details(
media_type: str, root_path: str, error: ProxyException | HTTPException,
) -> None:
from google.rpc.status_pb2 import Status
from litellm.proxy.proxy_server import otlp_http_exception_handler
from litellm.rust_bridge import loader
if not native_available:
monkeypatch.setattr(loader, "_cached_bridge", None)
request: Final = Request({
"type": "http", "method": "POST", "path": root_path + "/v1/traces", "root_path": root_path,
"headers": [(b"content-type", media_type.encode())],
@ -522,7 +517,7 @@ async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native
else Status.FromString(response.body).message
)
expected: Final = "Unauthorized" if isinstance(error, ProxyException) else "Forbidden"
assert message == (expected if native_available or media_type == "application/json" else "")
assert message == expected
def _websocket(sent: list[Message]) -> WebSocket:

1
tests/unit/proxy/test_tracing_endpoints.py
View file

@ -30,7 +30,6 @@ from litellm.proxy.auth.authorization import OwnedRows, ReadScope
from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.tracing_runtime import manage_tracing, provide_storage
from litellm.rust_bridge import loader
from litellm.tracing.errors import TraceChanged
from litellm.tracing.generated.models import TraceQueryHelp
from litellm.tracing.generated.responses import TraceSQLResponse

15
tests/unit/responses/test_rust_bridge_websocket.py
View file

@ -55,21 +55,6 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None:
await adapter.recv()
@pytest.mark.asyncio
async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(responses_websocket, "_STATE", responses_websocket._RustResponsesWebSocketState())
monkeypatch.setattr(responses_websocket, "get_native_bridge", lambda: None)
assert (
await responses_websocket.connect(
url="wss://example.test/responses",
headers={},
timeout=None,
)
is None
)
@pytest.mark.asyncio
async def test_enabled_bridge_connects_and_adapts_socket(
monkeypatch: pytest.MonkeyPatch,

46
tests/unit/rust_bridge/host/test_secret_manager.py
View file

@ -17,6 +17,7 @@ from botocore.credentials import Credentials
from pydantic import JsonValue
import litellm
from litellm.rust_bridge import _native as native
from litellm.rust_bridge.host.secret_manager import capture_secret_manager, native_secret_manager_config
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager
@ -67,7 +68,6 @@ def _vault_body(value: str) -> dict[str, object]:
def _native_handle(client: object) -> "_SecretManagerRuntime":
native: Final = pytest.importorskip("litellm.rust_bridge._native")
handle: Final = native._SecretManagerRuntime.from_client(client)
assert handle is not None
return handle
@ -75,7 +75,6 @@ def _native_handle(client: object) -> "_SecretManagerRuntime":
@pytest.mark.parametrize("system", ("aws_secret_manager", "hashicorp_vault", "cyberark"))
async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(system: str) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
environment: Final = {
"AWS_REGION_NAME": "us-east-1",
@ -109,7 +108,6 @@ async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(s
def test_shared_initializer_captures_credentials_and_tracks_instance_settings(monkeypatch: pytest.MonkeyPatch) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body={"SecretString": "native-value"})
server.expected_requests = 2
@ -141,7 +139,6 @@ def test_shared_initializer_captures_credentials_and_tracks_instance_settings(mo
def test_configuration_replacement_rebuilds_without_invalidating_existing_handles(
monkeypatch: pytest.MonkeyPatch,
) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as first_server, recording_service() as second_server:
first_server.default_response = ResponseSpec(body=_vault_body("first"))
second_server.default_response = ResponseSpec(body=_vault_body("second"))
@ -158,7 +155,6 @@ def test_configuration_replacement_rebuilds_without_invalidating_existing_handle
def test_custom_subclass_is_not_captured_for_a_native_handle() -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
class CustomManager(AWSSecretsManagerV2):
def sync_read_secret(self, secret_name: str, primary_secret_name: str | None = None) -> str:
@ -171,19 +167,16 @@ def test_custom_subclass_is_not_captured_for_a_native_handle() -> None:
@pytest.mark.parametrize("system", ("google_secret_manager", "hashicorp_vault", "cyberark"))
def test_enterprise_backends_cannot_initialize_without_entitlement(system: str) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with pytest.raises(ValueError, match=r"[Ee]nterprise|[Pp]remium"):
native._SecretManagerRuntime.from_config(system, {"CYBERARK_API_KEY": "key"})
def test_azure_factory_rejects_unencrypted_vault_endpoints() -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with pytest.raises(ValueError, match="https"):
native._SecretManagerRuntime.from_config("azure_key_vault", {"AZURE_KEY_VAULT_URI": "http://127.0.0.1:1"})
def test_direct_builtin_constructor_can_use_native_without_registration(monkeypatch: pytest.MonkeyPatch) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body=_vault_body("native-value"))
manager: Final = _vault(monkeypatch, server.base_url)
@ -193,12 +186,10 @@ def test_direct_builtin_constructor_can_use_native_without_registration(monkeypa
def test_uncaptured_clients_have_no_native_handle() -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
assert native._SecretManagerRuntime.from_client(SimpleNamespace()) is None
def test_native_failure_is_not_replayed_in_python(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(status=403, body={"errors": ["denied"]})
server.expected_requests = 1
@ -267,7 +258,6 @@ def test_same_named_custom_client_is_not_captured() -> None:
@pytest.mark.parametrize("value", ("text", "", True, False, 42, 2**100, [1, "two"], {"nested": True}, None))
async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.MonkeyPatch, value: JsonValue) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body={"SecretString": json.dumps({"KEY": value})})
server.expected_requests = 3
@ -290,7 +280,6 @@ async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.Monke
@pytest.mark.parametrize("primary", (None, "primary"))
async def test_aws_primary_values_match_python_reads(monkeypatch: pytest.MonkeyPatch, primary: str | None) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body={"SecretString": '{"KEY":"value"}'})
server.expected_requests = 3
@ -343,7 +332,6 @@ def test_aws_absence_and_failed_reads_match_python_without_environment_fallback(
async def test_aws_primary_json_errors_preserve_python_exception_details(
monkeypatch: pytest.MonkeyPatch, document: str
) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body={"SecretString": document})
server.expected_requests = 3
@ -370,7 +358,6 @@ async def test_aws_primary_json_errors_preserve_python_exception_details(
async def test_public_aws_reads_preserve_coroutines_and_per_call_credentials(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as initial, recording_service() as selected:
initial.expected_requests = 0
selected.expected_requests = 2
@ -421,7 +408,6 @@ async def test_public_aws_reads_preserve_coroutines_and_per_call_credentials(
async def test_public_aws_primary_reads_ignore_operation_overrides(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server, recording_service() as unused:
server.default_response = ResponseSpec(body={"SecretString": '{"KEY":true}'})
server.expected_requests = 2
@ -440,7 +426,6 @@ async def test_public_aws_primary_reads_ignore_operation_overrides(monkeypatch:
async def test_public_aws_bootstrap_names_only_bypass_sync_reads(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body={"SecretString": "remote-access"})
server.expected_requests = 1
@ -457,7 +442,6 @@ async def test_public_aws_bootstrap_names_only_bypass_sync_reads(monkeypatch: py
async def test_public_aws_read_timeouts_follow_the_python_http_handler(
monkeypatch: pytest.MonkeyPatch, timeout: float | httpx.Timeout
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body={"SecretString": "too-late"}, delay=0.25)
server.expected_requests = 2
@ -470,7 +454,6 @@ async def test_public_aws_read_timeouts_follow_the_python_http_handler(
async def test_public_vault_reads_keep_overrides_cache_and_coroutines(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body=_vault_body("value"))
server.expected_requests = 2
@ -490,7 +473,6 @@ async def test_public_vault_reads_keep_overrides_cache_and_coroutines(monkeypatc
async def test_public_vault_failed_reads_return_none_without_replay(
monkeypatch: pytest.MonkeyPatch, status: int
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(status=status, body={"errors": ["unavailable"]})
server.expected_requests = 2
@ -502,7 +484,6 @@ async def test_public_vault_failed_reads_return_none_without_replay(
async def test_public_cyberark_reads_reuse_authentication_and_cached_values(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
from litellm.proxy import proxy_server
from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager
@ -536,7 +517,6 @@ def test_public_aws_bootstrap_read_does_not_hit_the_provider(monkeypatch: pytest
def test_public_vault_prefix_overrides_match_python_string_conversion(
monkeypatch: pytest.MonkeyPatch, override: str | int
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("HCP_VAULT_PATH_PREFIX", "default-prefix")
with recording_service() as server:
server.default_response = ResponseSpec(body=_vault_body("value"))
@ -562,7 +542,6 @@ def _cyberark(monkeypatch: pytest.MonkeyPatch, address: str) -> CyberArkSecretMa
async def test_public_cyberark_writes_and_deletes_share_the_read_cache(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
for body in (b"token", b"old", {}, {}, b"provider-after-delete"):
server.enqueue(ResponseSpec(body=body))
@ -609,7 +588,6 @@ async def test_public_cyberark_write_errors_match_python_without_http_retries(
status: int,
authentication: bool,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
responses: Final = (
(ResponseSpec(status=status, body={}),) * 2
@ -635,7 +613,6 @@ async def test_public_cyberark_write_errors_match_python_without_http_retries(
async def test_public_cyberark_write_recovers_from_initial_policy_authentication_failure(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.enqueue(ResponseSpec(status=401, body={}))
server.enqueue(ResponseSpec(body=b"token"))
@ -658,7 +635,6 @@ async def test_public_cyberark_write_rejects_unsafe_names_before_authentication(
monkeypatch: pytest.MonkeyPatch,
name: str,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.expected_requests = 0
manager: Final = _cyberark(monkeypatch, server.base_url)
@ -673,7 +649,6 @@ async def test_public_cyberark_rotation_returns_the_write_response_and_retains_o
monkeypatch: pytest.MonkeyPatch,
same_name: bool,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
for body in (b"token", b"old-value", {}, {}):
server.enqueue(ResponseSpec(body=body))
@ -696,7 +671,6 @@ async def test_public_native_cyberark_rotation_requires_a_fresh_matching_replace
monkeypatch: pytest.MonkeyPatch,
replacement: bytes | None,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
for body in (b"token", b"old-value", {}, {}):
server.enqueue(ResponseSpec(body=body))
@ -720,7 +694,6 @@ async def test_public_native_cyberark_rotation_requires_a_fresh_matching_replace
async def test_public_cyberark_cached_authentication_does_not_retry_denied_reads(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.enqueue(ResponseSpec(body=b"token"))
server.enqueue(ResponseSpec(body=b"value"))
@ -734,7 +707,6 @@ async def test_public_cyberark_cached_authentication_does_not_retry_denied_reads
async def test_cyberark_handler_errors_match_python_after_cached_authentication_is_denied(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
for response in (
ResponseSpec(body=b"token"),
@ -757,7 +729,6 @@ async def test_cyberark_handler_errors_match_python_after_cached_authentication_
async def test_public_cyberark_connection_errors_match_python(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.expected_requests = 0
address: Final = server.base_url
@ -769,7 +740,6 @@ async def test_public_cyberark_connection_errors_match_python(
async def test_public_native_cyberark_rotation_stops_after_a_failed_write(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
for body in (b"token", b"old-value", {}):
server.enqueue(ResponseSpec(body=body))
@ -788,7 +758,6 @@ async def test_public_vault_writes_preserve_complete_responses_and_request_field
monkeypatch: pytest.MonkeyPatch,
description: str | None,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
response: Final = {
"request_id": "test-request",
@ -829,7 +798,6 @@ async def test_public_vault_mutation_http_errors_match_python_without_retry(
operation: str,
status: int,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(status=status, body={"errors": ["denied"]})
server.expected_requests = 2
@ -856,7 +824,6 @@ async def test_public_vault_write_response_conversion_matches_python(
monkeypatch: pytest.MonkeyPatch,
body: bytes,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body=body)
server.expected_requests = 2
@ -870,7 +837,6 @@ async def test_public_vault_write_response_conversion_matches_python(
async def test_public_vault_deletion_invalidates_cached_fields(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.enqueue(ResponseSpec(body=_vault_body("old")))
server.enqueue(ResponseSpec(status=204, body=b""))
@ -893,7 +859,6 @@ async def test_public_vault_rotation_preserves_response_and_best_effort_deletion
same_name: bool,
delete_status: int,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
response: Final = {"request_id": "write-id", "data": {"version": 3}, "extra": [1, 2]}
server.enqueue(ResponseSpec(body=b"current-existence-is-status-only"))
@ -924,7 +889,6 @@ async def test_public_vault_rotation_failure_messages_and_request_counts_match_p
stage: str,
status: int,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
before: Final = (
()
@ -960,7 +924,6 @@ async def test_public_vault_rotation_mismatches_do_not_delete_the_old_alias(
monkeypatch: pytest.MonkeyPatch,
value: JsonValue,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
responses: Final = (
ResponseSpec(body=_vault_body("old")),
@ -983,7 +946,6 @@ async def test_public_vault_mutation_timeouts_match_python(
monkeypatch: pytest.MonkeyPatch,
operation: str,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.default_response = ResponseSpec(body=_vault_body("value"), delay=0.25)
server.expected_requests = 2
@ -1021,7 +983,6 @@ async def test_public_vault_unsafe_names_fail_before_authentication(
monkeypatch: pytest.MonkeyPatch,
operation: str,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.expected_requests = 0
manager: Final = _vault(monkeypatch, server.base_url)
@ -1038,7 +999,6 @@ async def test_public_vault_authentication_errors_match_python(
monkeypatch: pytest.MonkeyPatch,
operation: str,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("HCP_VAULT_APPROLE_ROLE_ID", "role")
monkeypatch.setenv("HCP_VAULT_APPROLE_SECRET_ID", "secret-id")
with recording_service() as server:
@ -1064,7 +1024,6 @@ async def test_public_vault_authentication_errors_match_python(
async def test_public_vault_rotation_stops_on_a_success_response_containing_an_error(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
result: Final = {"status": "error", "message": "write rejected", "extra": 2**100}
for response in (ResponseSpec(body=_vault_body("old")), ResponseSpec(body=result)) * 2:
@ -1078,7 +1037,6 @@ async def test_public_vault_rotation_stops_on_a_success_response_containing_an_e
async def test_public_native_vault_write_invalidates_stale_cached_values(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.enqueue(ResponseSpec(body=_vault_body("old")))
server.enqueue(ResponseSpec(body={"data": {"version": 2}}))
@ -1094,7 +1052,6 @@ async def test_public_native_vault_write_invalidates_stale_cached_values(monkeyp
async def test_public_native_vault_write_rejects_description_overwriting_the_secret(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
server.expected_requests = 0
handle: Final = _native_handle(_vault(monkeypatch, server.base_url))
@ -1111,7 +1068,6 @@ async def test_public_vault_rotation_preserves_malformed_verification_errors(
monkeypatch: pytest.MonkeyPatch,
body: JsonValue,
) -> None:
pytest.importorskip("litellm.rust_bridge._native")
with recording_service() as server:
responses: Final = (
ResponseSpec(body=_vault_body("old")),

2
tests/unit/rust_bridge/messages/test_secrets.py
View file

@ -17,8 +17,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service
pytest.importorskip("litellm.rust_bridge._native")
pytestmark = pytest.mark.usefixtures("local_model_cost_map")

3
tests/unit/rust_bridge/ocr/test_secrets.py
View file

@ -12,6 +12,7 @@ import pytest
import litellm
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge import _native as native
from litellm.rust_bridge.host import settings
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
from litellm.rust_bridge.public_call import NativeCall
@ -19,8 +20,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem
from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service
native: Final = pytest.importorskip("litellm.rust_bridge._native")
AccessMode: TypeAlias = Literal["read_only", "write_only", "read_and_write"]

57
tests/unit/rust_bridge/test_bindings.py
View file

@ -6,7 +6,7 @@ from typing import Final, cast # noqa: TID251 # narrows the async binding's un
import pytest
from litellm.rust_bridge import bindings
from litellm.rust_bridge import _native, bindings
from litellm.rust_bridge.chat_completions import entrypoints as chat_completions
from litellm.rust_bridge.embeddings import entrypoints as embeddings
from litellm.rust_bridge.execution import Complete, Step
@ -17,12 +17,15 @@ from litellm.rust_bridge.responses import entrypoints as responses
from litellm.rust_bridge.transcription import native as transcription
def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None:
native = SimpleNamespace(route=lambda: "native")
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
binding: bindings.NativeBinding[object] = bindings.NativeBinding("route", validate=lambda value: value)
def test_binding_distinguishes_disable_from_reset() -> None:
def route() -> str:
return "native"
assert binding.load() is native.route
binding: bindings.NativeBinding[object] = bindings.NativeBinding(
"route", validate=lambda value: value, source=SimpleNamespace(route=route)
)
assert binding.load() is route
binding.override(None)
assert binding.load() is None
@ -32,16 +35,15 @@ def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None:
assert binding.load() is replacement
binding.reset()
assert binding.load() is native.route
assert binding.load() is route
@pytest.mark.parametrize(("value", "expected"), ((3, 3), ("invalid", None), (None, None)))
def test_binding_validates_native_attribute(
monkeypatch: pytest.MonkeyPatch, value: object, expected: int | None
) -> None:
def test_binding_validates_native_attribute(value: object, expected: int | None) -> None:
native: Final = SimpleNamespace(route=value)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None)
binding: Final = bindings.NativeBinding(
"route", validate=lambda item: item if isinstance(item, int) else None, source=native
)
assert binding.load() == expected
@ -72,28 +74,33 @@ ROUTE_BINDINGS: Final = tuple((attribute, binding) for attribute, binding, _ in
@pytest.mark.parametrize(
("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS]
)
def test_route_bindings_reject_non_callable_native_attributes(
monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object]
def test_route_bindings_resolve_the_packaged_native_route(
attribute: str, route_binding: bindings.NativeBinding[object]
) -> None:
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: "not callable"}))
route_binding.reset()
assert route_binding.load() is None
route_binding.reset()
assert attribute in dir(_native)
assert callable(route_binding.load())
@pytest.mark.parametrize(
("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS]
)
def test_route_bindings_reject_non_callable_native_attributes(
attribute: str, route_binding: bindings.NativeBinding[object]
) -> None:
assert route_binding.with_source(SimpleNamespace(**{attribute: "not callable"})).load() is None
@pytest.mark.parametrize(
("attribute", "route_binding"), TRANSCRIPTION_BINDINGS, ids=[attribute for attribute, _ in TRANSCRIPTION_BINDINGS]
)
def test_transcription_bindings_expose_the_native_callable_itself(
monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object]
attribute: str, route_binding: bindings.NativeBinding[object]
) -> None:
def native_route() -> None:
pass
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route}))
route_binding.reset()
assert route_binding.load() is native_route
route_binding.reset()
assert route_binding.with_source(SimpleNamespace(**{attribute: native_route})).load() is native_route
@dataclass(frozen=True, slots=True)
@ -116,7 +123,6 @@ class _EchoingExecution:
ids=[attribute for attribute, _, _ in EXECUTION_BINDINGS],
)
def test_execution_bindings_start_the_native_execution_and_return_its_result(
monkeypatch: pytest.MonkeyPatch,
attribute: str,
route_binding: bindings.NativeBinding[object],
asynchronous: bool,
@ -128,9 +134,7 @@ def test_execution_bindings_start_the_native_execution_and_return_its_result(
received.append(incoming)
return _EchoingExecution()
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route}))
route_binding.reset()
bound: Final = route_binding.load()
bound: Final = route_binding.with_source(SimpleNamespace(**{attribute: native_route})).load()
assert callable(bound)
outcome: Final = bound(call)
@ -139,7 +143,6 @@ def test_execution_bindings_start_the_native_execution_and_return_its_result(
if asynchronous
else outcome
)
route_binding.reset()
assert received == [call]
assert result == ("completed with", None)

23
tests/unit/rust_bridge/test_fork_guard.py
View file

@ -1,36 +1,21 @@
from types import SimpleNamespace
import pytest
from litellm.rust_bridge import fork_guard
def _reserve_with(monkeypatch: pytest.MonkeyPatch, native: object) -> None:
monkeypatch.setattr(fork_guard, "get_native_bridge", lambda: native)
fork_guard.reserve_process_for_forking("the gunicorn master")
def test_missing_extension_has_nothing_to_reserve(monkeypatch: pytest.MonkeyPatch) -> None:
assert _reserve_with(monkeypatch, None) is None
def test_extension_built_before_reservation_existed_passes(monkeypatch: pytest.MonkeyPatch) -> None:
assert _reserve_with(monkeypatch, SimpleNamespace()) is None
def test_unused_extension_is_reserved(monkeypatch: pytest.MonkeyPatch) -> None:
def test_unused_extension_is_reserved() -> None:
calls: list[None] = []
_reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=lambda: calls.append(None)))
fork_guard.reserve_process_for_forking("the gunicorn master", reserve=lambda: calls.append(None))
assert calls == [None]
def test_used_extension_refuses_and_names_the_place(monkeypatch: pytest.MonkeyPatch) -> None:
def test_used_extension_refuses_and_names_the_place() -> None:
def reserve() -> None:
raise RuntimeError("the native runtime already started in this process")
with pytest.raises(fork_guard.NativeStateStartedBeforeFork, match="the gunicorn master") as raised:
_reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=reserve))
fork_guard.reserve_process_for_forking("the gunicorn master", reserve=reserve)
assert isinstance(raised.value.__cause__, RuntimeError)

3
tests/unit/rust_bridge/test_logger.py
View file

@ -80,7 +80,6 @@ def test_native_logging_observes_level_changes(caplog: pytest.LogCaptureFixture)
),
)
def test_native_credential_patterns_match_python(text: str) -> None:
pytest.importorskip("litellm.rust_bridge._native")
from litellm.rust_bridge._native import NativeDiagnosticProcessor
processor: Final = NativeDiagnosticProcessor(MINIMUM_CUSTOM_KEY_LENGTH)
@ -91,7 +90,6 @@ def test_native_credential_patterns_match_python(text: str) -> None:
def test_native_client_redaction_matches_python() -> None:
pytest.importorskip("litellm.rust_bridge._native")
from litellm.rust_bridge._native import NativeDiagnosticProcessor
text: Final = "error at /etc/secrets/config on db.internal\nTraceback (most recent call last):\nsecret"
@ -100,7 +98,6 @@ def test_native_client_redaction_matches_python() -> None:
def test_native_diagnostic_batch_matches_python() -> None:
pytest.importorskip("litellm.rust_bridge._native")
from litellm.rust_bridge._native import NativeDiagnosticProcessor
message: Final = "é" * 110 + "sk-" + "q" * 48 + "界" * 1000

17
tests/unit/rust_bridge/test_runtime.py
View file

@ -1,7 +1,6 @@
from __future__ import annotations
from collections.abc import Callable, Generator
from types import SimpleNamespace
from typing import Final, Protocol
import pytest
@ -10,26 +9,14 @@ from litellm.exceptions import APIError
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.rust_bridge import bindings, configuration, runtime
from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.execution import AsyncStream, Complete, StreamChunk, SyncStream
class RustBridgeDeclined(Exception):
pass
class RustUpstreamError(Exception):
pass
@pytest.fixture(autouse=True)
def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
native: Final = SimpleNamespace(
RustBridgeDeclined=RustBridgeDeclined,
RustUpstreamError=RustUpstreamError,
)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
def rust_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield

4
tests/unit/rust_bridge/test_token_counter.py
View file

@ -16,6 +16,7 @@ import litellm
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding
from litellm.proxy.spend_tracking.input_tokens import count_input_tokens_for_model
from litellm.rust_bridge import _native as native
from litellm.rust_bridge import token_counter as bridge
from litellm.rust_bridge import tokenizer as tokenizer_dispatch
from litellm.rust_bridge._native import Tokenizer
@ -313,7 +314,6 @@ PARITY_MODELS: Final[tuple[tuple[str, bridge.RustTokenizer], ...]] = (
async def test_native_count_matches_python_budget_counter(
request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer
) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
body: Final = json.dumps(request_body).replace(MODEL, model)
parsed: Final = json.loads(body)
@ -328,7 +328,6 @@ async def test_tiktoken_counts_long_text_exactly_where_python_chunks(
model: str, tokenizer: bridge.RustTokenizer
) -> None:
"""Python encodes tiktoken text in fixed-size chunks (drift of up to one token per chunk boundary); Rust does not."""
native: Final = pytest.importorskip("litellm.rust_bridge._native")
text: Final = "x " * 20_000
body: Final = {"model": model, "messages": [{"role": "user", "content": text}]}
encoding: Final = Tokenizer.from_tiktoken(tokenizer)
@ -362,7 +361,6 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
async def test_native_declines_shapes_python_prices_differently(
request_body: dict[str, object], tokenizer: bridge.RustTokenizer
) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
raw, _ = _counted(request_body, MODEL_BY_TOKENIZER[tokenizer])
with pytest.raises(native.RustBridgeDeclined):

14
tests/unit/test_audio_transcription_rust_bridge.py
View file

@ -1,14 +1,14 @@
from __future__ import annotations
from collections.abc import Generator
from types import SimpleNamespace
from typing import Final
import pytest
import litellm
from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge import configuration
from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError
from litellm.rust_bridge.public_call import NativeCall
from litellm.rust_bridge.transcription.native import NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION
@ -28,18 +28,8 @@ TRANSCRIPTION_FIELDS: Final = frozenset(
)
class RustBridgeDeclined(Exception):
pass
class RustUpstreamError(Exception):
pass
@pytest.fixture(autouse=True)
def isolated_bridge(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
native: Final = SimpleNamespace(RustBridgeDeclined=RustBridgeDeclined, RustUpstreamError=RustUpstreamError)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield

14
tests/unit/test_logging.py
View file

@ -862,8 +862,6 @@ def test_secret_filter_keeps_truncated_traceback(monkeypatch):
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
def test_diagnostic_redaction_precedes_a_credential_cut(monkeypatch, native):
if native:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
monkeypatch.setenv("MAX_STRING_LENGTH_STDOUT_LOG", "500")
monkeypatch.setenv("MAX_BASE64_LENGTH_STDOUT_LOG", "0")
@ -890,8 +888,6 @@ def test_correlation_id_redacts_before_its_length_bound(monkeypatch):
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native):
if native:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
record = _make_record(logging.WARNING, "bad % api_key=secret123", ("value",))
@ -905,8 +901,6 @@ def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native):
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
def test_key_pattern_template_keeps_the_rendered_redacted_line(monkeypatch, native):
if native:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
record = _make_record(logging.INFO, "password=%s ok", ("hunter2",))
@ -1733,8 +1727,6 @@ def test_access_log_path_filter_keeps_a_record_without_a_string_path_arg(monkeyp
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, native):
if native:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
secret = "sk-" + "q" * 48
@ -1761,8 +1753,6 @@ def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, nativ
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch, native):
if native:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
record = _make_record(logging.WARNING, "api_key=secret123")
@ -1775,8 +1765,6 @@ def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch,
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native):
if native:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
secret = "sk-" + "q" * 48
@ -1791,8 +1779,6 @@ def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native):
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
def test_diagnostic_filter_redacts_a_non_string_message_object(monkeypatch, native):
if native:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
secret = "sk-" + "q" * 48

2
tests/unit/test_secret_redaction.py
View file

@ -76,8 +76,6 @@ def test_redact_string_catches_secret_patterns():
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
def test_diagnostic_redaction_policy_matches_across_backends(monkeypatch: pytest.MonkeyPatch, native: bool) -> None:
if native:
pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
assert redact_string("GET /v1?api_key=abcdefgh12345&page=2") == "GET /v1?REDACTED&page=2"

12
tests/windows_tests/check_windows_wheel_install.py
View file

@ -87,10 +87,20 @@ def main(argv):
f"::error::installing {os.path.basename(wheel)} into a deep prefix failed"
)
return 1
if _run([python, "-c", "import litellm; import litellm.types.utils"]) != 0:
if _run([python, "-I", "-c", "import litellm; import litellm.types.utils"]) != 0:
print("::error::litellm did not import after install (half-unpacked package)")
return 1
checker = os.path.abspath(os.path.join("tests", "base_sdk_tests", "check_base_sdk_install.py"))
for rust in ("0", "1"):
result = subprocess.run(
[python, "-I", checker, "--profile", "native"],
env={**os.environ, "LITELLM_RUST": rust, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
check=False,
)
if result.returncode != 0:
return result.returncode
print(
f"ok: {os.path.basename(wheel)} installs into a worst-case prefix and imports"
)