diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index ca98fe2b5b2..ee90a734185 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_logger from litellm.rust_bridge import catalog, runtime from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import Route, RouteContext, Rules +from litellm.rust_bridge.observation import observe_early RequestT: Final = TypeVar("RequestT") NativeT: Final = TypeVar("NativeT") @@ -34,6 +35,7 @@ class PublicDispatch(Generic[RequestT]): def _native_request(self, args: tuple[object, ...], kwargs: Mapping[str, object]) -> RequestT: request: Final = self.request(args, kwargs) if request is None: + observe_early(self.route, "projection_failure") raise runtime.NoPythonImplementationError( f"{self.route.value} has no Python implementation, so every call must project to a native request" ) @@ -43,6 +45,7 @@ class PublicDispatch(Generic[RequestT]): try: return self.context(request) except Exception: # noqa: BLE001 # detectors call legacy handlers, so any error must fall back to Python + observe_early(self.route, "detector_failure") verbose_logger.warning( "%s route detection failed, serving the call on Python", self.route.value, exc_info=True ) @@ -69,9 +72,11 @@ class PublicDispatch(Generic[RequestT]): rules=selected_rules, ) if not catalog.rust_might_handle(self.route, selected_rules): + observe_early(self.route, "no_candidate") return python(*args, **kwargs) request: Final = self.request(args, kwargs) if request is None: + observe_early(self.route, "projection_failure") return python(*args, **kwargs) context: Final = self._guarded_context(request) if context is None: @@ -105,9 +110,11 @@ class PublicDispatch(Generic[RequestT]): rules=selected_rules, ) if not catalog.rust_might_handle(self.route, selected_rules): + observe_early(self.route, "no_candidate") return await python(*args, **kwargs) request: Final = self.request(args, kwargs) if request is None: + observe_early(self.route, "projection_failure") return await python(*args, **kwargs) context: Final = self._guarded_context(request) if context is None: diff --git a/litellm/rust_bridge/observation.py b/litellm/rust_bridge/observation.py new file mode 100644 index 00000000000..be98fc28ae5 --- /dev/null +++ b/litellm/rust_bridge/observation.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import ContextVar, Token +from dataclasses import dataclass +from threading import Lock +from typing import Final, Protocol +from uuid import uuid4 + +from litellm.rust_bridge.catalog import Route, RouteContext +from litellm.rust_bridge.configuration import Decision + + +@dataclass(frozen=True, slots=True) +class RouteObservation: + event: str + trace_id: str | None + attempt_id: str | None + parent_attempt_id: str | None + depth: int + route: Route | None + context: RouteContext | None = None + decision: Decision | None = None + outcome: str | None = None + reason: str | None = None + engine: str | None = None + + +class RouteObserver(Protocol): + def record(self, observation: RouteObservation) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _RequestScope: + trace_id: str + depth: int = 0 + parent_attempt_id: str | None = None + + +class _EventBuffer: + def __init__(self) -> None: + self._events: tuple[RouteObservation, ...] = () + self._lock: Final = Lock() + + def append(self, event: RouteObservation) -> None: + with self._lock: + self._events = (*self._events, event) + + def snapshot(self) -> tuple[RouteObservation, ...]: + with self._lock: + return self._events + + +_OBSERVER: Final[ContextVar[RouteObserver | None]] = ContextVar("litellm_rust_route_observer", default=None) +_REQUEST: Final[ContextVar[_RequestScope | None]] = ContextVar("litellm_rust_request_scope", default=None) +_CURRENT_ATTEMPT: Final[ContextVar[tuple[str, int, str | None] | None]] = ContextVar( + "litellm_rust_current_attempt", default=None +) +_EVENTS: Final[ContextVar[_EventBuffer | None]] = ContextVar("litellm_rust_route_events", default=None) + + +def set_observer(observer: RouteObserver | None) -> Token[RouteObserver | None]: + return _OBSERVER.set(observer) + + +def reset_observer(token: Token[RouteObserver | None]) -> None: + _OBSERVER.reset(token) + + +@contextmanager +def request_scope(trace_id: str) -> Generator[None, None, None]: + request_token: Final = _REQUEST.set(_RequestScope(trace_id)) + try: + yield + finally: + _REQUEST.reset(request_token) + + +@contextmanager +def trace_scope() -> Generator[None, None, None]: + token: Final = _EVENTS.set(_EventBuffer()) + try: + yield + finally: + _EVENTS.reset(token) + + +def current_events() -> tuple[RouteObservation, ...]: + events: Final = _EVENTS.get() + return () if events is None else events.snapshot() + + +def record_event(observation: RouteObservation) -> None: + events: Final = _EVENTS.get() + if events is not None: + events.append(observation) + + +def observe_early(route: Route, reason: str) -> None: + request: Final = _REQUEST.get() + current: Final = _CURRENT_ATTEMPT.get() + observer: Final = _OBSERVER.get() + if observer is None: + return + observer.record( + RouteObservation( + event="early", + trace_id=None if request is None else request.trace_id, + attempt_id=None, + parent_attempt_id=( + current[0] if current is not None else None if request is None else request.parent_attempt_id + ), + depth=(current[1] + 1 if current is not None else 0 if request is None else request.depth), + route=route, + reason=reason, + ) + ) + + +@contextmanager +def observed_attempt(context: RouteContext, selected: Decision) -> Generator[str, None, None]: + request: Final = _REQUEST.get() + observer: Final = _OBSERVER.get() + if observer is None: + yield "" + return + attempt_id: Final = uuid4().hex + trace_id: Final = None if request is None else request.trace_id + parent: Final = _CURRENT_ATTEMPT.get() + parent_attempt_id: Final = ( + parent[0] if parent is not None else None if request is None else request.parent_attempt_id + ) + depth: Final = parent[1] + 1 if parent is not None else 0 if request is None else request.depth + attempt_token: Final = _CURRENT_ATTEMPT.set((attempt_id, depth, parent_attempt_id)) + observer.record( + RouteObservation( + event="selected", + trace_id=trace_id, + attempt_id=attempt_id, + parent_attempt_id=parent_attempt_id, + depth=depth, + route=context.route, + context=context, + decision=selected, + ) + ) + try: + yield attempt_id + except BaseException: + observer.record( + RouteObservation( + event="terminal", + trace_id=trace_id, + attempt_id=attempt_id, + parent_attempt_id=parent_attempt_id, + depth=depth, + route=context.route, + context=context, + decision=selected, + outcome="failure", + reason="execution_error", + ) + ) + raise + finally: + _CURRENT_ATTEMPT.reset(attempt_token) + + +def observe_terminal( + attempt_id: str, + context: RouteContext, + selected: Decision, + *, + engine: str, +) -> None: + observer: Final = _OBSERVER.get() + if observer is None: + return + request: Final = _REQUEST.get() + current: Final = _CURRENT_ATTEMPT.get() + observer.record( + RouteObservation( + event="terminal", + trace_id=None if request is None else request.trace_id, + attempt_id=attempt_id, + parent_attempt_id=( + current[2] if current is not None else None if request is None else request.parent_attempt_id + ), + depth=current[1] if current is not None else 0 if request is None else request.depth, + route=context.route, + context=context, + decision=selected, + outcome="success", + engine=engine, + ) + ) + + +def observe_request_outcome(trace_id: str, outcome: str) -> None: + observer: Final = _OBSERVER.get() + if observer is None: + return + observer.record( + RouteObservation( + event="request", + trace_id=trace_id, + attempt_id=None, + parent_attempt_id=None, + depth=0, + route=None, + outcome=outcome, + ) + ) diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index cf632c09074..6b5f25eb8e3 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -10,6 +10,7 @@ from litellm.exceptions import APIError from litellm.rust_bridge.bindings import NativeBinding, native_exception_types from litellm.rust_bridge.catalog import RouteContext, Rules, decision from litellm.rust_bridge.configuration import Decision +from litellm.rust_bridge.observation import observe_terminal, observed_attempt from litellm.rust_bridge.response_metadata import mark_rust_response NativeT = TypeVar("NativeT") @@ -62,19 +63,28 @@ def run( rules: Rules | None = None, ) -> ResultT: selected: Final = decision(context, rules) - if isinstance(python, NoPythonImplementation): - _require_rust(context, selected) - return _required(_attempt_native(context, binding, native), context) - match selected: - case Decision.PYTHON: - return python() - case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: - result: Final = _attempt_native(context, binding, native) - if isinstance(result, RustHandled) or selected is Decision.RUST_REQUIRED: - return _required(result, context) - return python() - case _: - assert_never(selected) + with observed_attempt(context, selected) as attempt_id: + if isinstance(python, NoPythonImplementation): + _require_rust(context, selected) + return _observed_result( + _required(_attempt_native(context, binding, native), context), + attempt_id, + context, + selected, + "rust", + ) + match selected: + case Decision.PYTHON: + return _observed_result(python(), attempt_id, context, selected, "python") + case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: + native_result: Final = _attempt_native(context, binding, native) + if isinstance(native_result, RustHandled): + return _observed_result(_required(native_result, context), attempt_id, context, selected, "rust") + if selected is Decision.RUST_REQUIRED: + return _observed_result(_required(native_result, context), attempt_id, context, selected, "rust") + return _observed_result(python(), attempt_id, context, selected, "python") + case _: + assert_never(selected) async def arun( @@ -86,19 +96,39 @@ async def arun( rules: Rules | None = None, ) -> ResultT: selected: Final = decision(context, rules) - if isinstance(python, NoPythonImplementation): - _require_rust(context, selected) - return _required(await _aattempt_native(context, binding, native), context) - match selected: - case Decision.PYTHON: - return await python() - case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: - result: Final = await _aattempt_native(context, binding, native) - if isinstance(result, RustHandled) or selected is Decision.RUST_REQUIRED: - return _required(result, context) - return await python() - case _: - assert_never(selected) + with observed_attempt(context, selected) as attempt_id: + if isinstance(python, NoPythonImplementation): + _require_rust(context, selected) + return _observed_result( + _required(await _aattempt_native(context, binding, native), context), + attempt_id, + context, + selected, + "rust", + ) + match selected: + case Decision.PYTHON: + return _observed_result(await python(), attempt_id, context, selected, "python") + case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: + native_result: Final = await _aattempt_native(context, binding, native) + if isinstance(native_result, RustHandled): + return _observed_result(_required(native_result, context), attempt_id, context, selected, "rust") + if selected is Decision.RUST_REQUIRED: + return _observed_result(_required(native_result, context), attempt_id, context, selected, "rust") + return _observed_result(await python(), attempt_id, context, selected, "python") + case _: + assert_never(selected) + + +def _observed_result( + value: ResultT, + attempt_id: str, + context: RouteContext, + selected: Decision, + engine: str, +) -> ResultT: + observe_terminal(attempt_id, context, selected, engine=engine) + return value def _require_rust(context: RouteContext, selected: Decision) -> None: diff --git a/tests/integration/AGENTS.md b/tests/integration/AGENTS.md index a727b20596c..f5781c33452 100644 --- a/tests/integration/AGENTS.md +++ b/tests/integration/AGENTS.md @@ -37,3 +37,9 @@ together with its dedicated job, and list a new dedicated GitHub workflow with t `DEDICATED_INTEGRATION_WORKFLOWS` in `.github/scripts/assert_ci_coverage.py`. The `ci-coverage` check fails a file here that a GitHub workflow runs unless it is in `GITHUB_FILES` and that workflow owns its directory (#46780) + +Owned proxy tests that opt into rule-aware routing use +`integration._support.routing:install_rule_aware_harness` as a worker startup hook. The hook installs +worker-local trace middleware and resolves the shared catalog into explicit Python-only or native-required +rules. Assertions correlate the generated `x-integration-routing-id` response header with its trace and +use `assert_routing_matches_policy` rather than assuming an engine from a test row diff --git a/tests/integration/_support/routing.py b/tests/integration/_support/routing.py index 1f8f5626031..f53194f880a 100644 --- a/tests/integration/_support/routing.py +++ b/tests/integration/_support/routing.py @@ -6,8 +6,10 @@ import json import os import re import sys -from collections.abc import Iterator, Mapping -from dataclasses import dataclass +import threading +import uuid +from collections.abc import Collection, Iterator, Mapping +from dataclasses import dataclass, replace from pathlib import Path from types import MappingProxyType from typing import Final @@ -16,6 +18,11 @@ from urllib.parse import urlsplit, urlunsplit import psycopg import pytest from pydantic import TypeAdapter +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from litellm.rust_bridge import catalog, observation +from litellm.rust_bridge.catalog import Feature, RouteContext, RouteRule, Rules, Target +from litellm.rust_bridge.configuration import Decision, Rollout, decision WRITER_ROLE: Final = "litellm_writer" READER_ROLE: Final = "litellm_reader" @@ -39,6 +46,283 @@ def normalize(query: str) -> str: return _PLACEHOLDERS.sub("$n", " ".join(query.split())) +def strict_current_policy(rules: Rules | None = None) -> Rules: + source: Final = catalog.RULES if rules is None else rules + return tuple( + replace( + rule, + rollout=(Rollout.PYTHON_ONLY if decision(rule.rollout) is Decision.PYTHON else Rollout.RUST_REQUIRED), + ) + if rule.rollout in (Rollout.RUST_OPT_IN, Rollout.RUST_OPT_OUT) + else rule + for rule in source + ) + + +def _rule_document(rule: RouteRule) -> dict[str, object]: + def values(value: Collection[str | Feature | Target] | None) -> tuple[str, ...] | None: + return ( + None + if value is None + else tuple(sorted(item.value if isinstance(item, (Feature, Target)) else item for item in value)) + ) + + return { + "route": rule.route.value, + "rollout": rule.rollout.value, + "providers": values(rule.providers), + "models": values(rule.models), + "targets": values(rule.targets), + "features": values(rule.features), + } + + +def _policy_document(rules: Rules) -> tuple[dict[str, object], ...]: + return tuple(_rule_document(rule) for rule in rules) + + +class _TraceObserver: + def __init__(self, rules: Rules, trace_dir: Path) -> None: + self._rules: Final = rules + self._trace_dir: Final = trace_dir + self._lock: Final = threading.Lock() + + def record(self, event: observation.RouteObservation) -> None: + observation.record_event(event) + + def write(self, trace_id: str, events: tuple[observation.RouteObservation, ...]) -> Path: + self._trace_dir.mkdir(parents=True, exist_ok=True) + destination: Final = self._trace_dir / f"trace-{trace_id}.json" + document: Final = { + "trace_id": trace_id, + "worker_pid": os.getpid(), + "policy": _policy_document(self._rules), + "events": tuple(_event_document(event) for event in events), + } + with self._lock: + destination.write_text(json.dumps(document, indent=2, sort_keys=True) + "\n") + return destination + + +def _context_document(context: RouteContext | None) -> dict[str, object] | None: + if context is None: + return None + return { + "route": context.route.value, + "provider": context.provider, + "model": context.model, + "target": None if context.target is None else context.target.value, + "features": tuple(sorted(feature.value for feature in context.features)), + } + + +def _event_document(event: observation.RouteObservation) -> dict[str, object]: + return { + "event": event.event, + "trace_id": event.trace_id, + "attempt_id": event.attempt_id, + "parent_attempt_id": event.parent_attempt_id, + "depth": event.depth, + "route": None if event.route is None else event.route.value, + "context": _context_document(event.context), + "decision": None if event.decision is None else event.decision.value, + "outcome": event.outcome, + "reason": event.reason, + "engine": event.engine, + } + + +class _RoutingTraceMiddleware: + def __init__(self, app: ASGIApp, observer: _TraceObserver) -> None: + self._app: Final = app + self._observer: Final = observer + + async def __call__( + self, + scope: Scope, + receive: Receive, + send: Send, + ) -> None: + if scope.get("type") != "http": + await self._app(scope, receive, send) + return + trace_id: Final = uuid.uuid4().hex + + async def traced_send(message: Message) -> None: + if message.get("type") == "http.response.start" and not any( + event.event == "request" for event in observation.current_events() + ): + events: Final = observation.current_events() + outcome: Final = ( + "completed" + if any(event.event == "terminal" and event.outcome == "success" for event in events) + else "failure" + if any(event.event == "terminal" and event.outcome == "failure" for event in events) + else "no_dispatch" + ) + observation.observe_request_outcome(trace_id, outcome) + self._observer.write(trace_id, observation.current_events()) + headers: Final = [ + (name, value) + for name, value in message.get("headers", []) + if name.lower() != b"x-integration-routing-id" + ] + response: Final = { + **message, + "headers": [*headers, (b"x-integration-routing-id", trace_id.encode())], + } + await send(response) + return + await send(message) + + observer_token: Final = observation.set_observer(self._observer) + try: + with observation.request_scope(trace_id), observation.trace_scope(): + try: + await self._app(scope, receive, traced_send) + finally: + if not any(event.event == "request" for event in observation.current_events()): + events: Final = observation.current_events() + outcome: Final = ( + "failure" + if any(event.event == "terminal" and event.outcome == "failure" for event in events) + else "no_dispatch" + ) + observation.observe_request_outcome(trace_id, outcome) + self._observer.write(trace_id, observation.current_events()) + finally: + observation.reset_observer(observer_token) + + +def install_rule_aware_harness() -> None: + trace_dir: Final = Path(os.environ["LITELLM_INTEGRATION_ROUTING_TRACE_DIR"]) + rules: Final = strict_current_policy() + policy_path: Final = trace_dir / f"policy-{os.getpid()}.json" + trace_dir.mkdir(parents=True, exist_ok=True) + policy_path.write_text(json.dumps(_policy_document(rules), indent=2, sort_keys=True) + "\n") + setattr(catalog, "RULES", rules) + observer: Final = _TraceObserver(rules, trace_dir) + from litellm.proxy.proxy_server import app + + middleware_stack: Final = app.middleware_stack + if middleware_stack is None: + app.add_middleware(_RoutingTraceMiddleware, observer=observer) + else: + app.middleware_stack = _RoutingTraceMiddleware(middleware_stack, observer) + + +def trace_path(trace_dir: Path, trace_id: str) -> Path: + return trace_dir / f"trace-{trace_id}.json" + + +def rule_aware_overrides(directory: Path, *, rust: bool) -> Mapping[str, str]: + return MappingProxyType( + { + "LITELLM_RUST": "true" if rust else "false", + "LITELLM_INTEGRATION_ROUTING_TRACE_DIR": str(directory / "routing-traces"), + "LITELLM_WORKER_STARTUP_HOOKS": "integration._support.routing:install_rule_aware_harness", + } + ) + + +def assert_routing_matches_policy(headers: Mapping[str, str], trace_location: Path) -> None: + trace_id: Final = headers.get("x-integration-routing-id") + assert trace_id is not None, dict(headers) + assert trace_location.name == f"trace-{trace_id}.json", (trace_location, trace_id) + assert trace_location.exists(), f"routing trace missing: {trace_location}" + document: Final = TypeAdapter(dict[str, object]).validate_python(json.loads(trace_location.read_text())) + assert document.get("trace_id") == trace_id, document + events: Final = TypeAdapter(list[dict[str, object]]).validate_python(document.get("events", [])) + assert all(event.get("trace_id") == trace_id for event in events), events + rules: Final = _rules_from_document(document.get("policy", [])) + requests: Final = tuple(event for event in events if event.get("event") == "request") + assert len(requests) == 1, requests + assert requests[0].get("outcome") in ("completed", "no_dispatch"), requests + assert not any( + event.get("event") == "early" and event.get("reason") in ("projection_failure", "detector_failure") + for event in events + ), events + terminals: Final = tuple( + event + for event in events + if event.get("event") == "terminal" + and event.get("outcome") in ("success", "failure") + and event.get("parent_attempt_id") is None + ) + successes: Final = tuple(event for event in terminals if event.get("outcome") == "success") + if successes: + assert requests[0].get("outcome") == "completed", requests + assert len(successes) == 1, successes + terminal: Final = successes[0] + context_document: Final = TypeAdapter(dict[str, object]).validate_python(terminal["context"]) + context: Final = RouteContext( + route=catalog.Route(str(context_document["route"])), + provider=None if context_document.get("provider") is None else str(context_document["provider"]), + model=None if context_document.get("model") is None else str(context_document["model"]), + target=( + None if context_document.get("target") is None else catalog.Target(str(context_document["target"])) + ), + features=frozenset( + Feature(feature) + for feature in TypeAdapter(list[str]).validate_python(context_document.get("features", [])) + ), + ) + expected: Final = catalog.decision(context, rules) + assert expected in (Decision.PYTHON, Decision.RUST_REQUIRED), (terminal, expected) + assert terminal.get("decision") == expected.value, (terminal, expected) + marker: Final = headers.get("x-litellm-rust") == "true" + if expected is Decision.RUST_REQUIRED: + assert marker, headers + assert terminal.get("engine") == "rust", terminal + else: + assert not marker, headers + assert terminal.get("engine") == "python", terminal + return + + assert not terminals, terminals + no_candidates: Final = tuple( + event + for event in events + if event.get("event") == "early" + and event.get("reason") == "no_candidate" + and event.get("parent_attempt_id") is None + ) + assert len(no_candidates) == 1, events + assert requests[0].get("outcome") == "no_dispatch", requests + early: Final = no_candidates[0] + assert early.get("route") is not None, early + route: Final = catalog.Route(str(early["route"])) + assert not catalog.rust_might_handle(route, rules), (early, rules) + assert headers.get("x-litellm-rust") != "true", headers + return + + +def _rules_from_document(value: object) -> Rules: + documents: Final = TypeAdapter(list[dict[str, object]]).validate_python(value) + + return tuple( + RouteRule( + route=catalog.Route(str(document["route"])), + rollout=Rollout(TypeAdapter(int).validate_python(document["rollout"])), + providers=None + if document.get("providers") is None + else frozenset(TypeAdapter(list[str]).validate_python(document["providers"])), + models=None + if document.get("models") is None + else frozenset(TypeAdapter(list[str]).validate_python(document["models"])), + targets=None + if document.get("targets") is None + else frozenset(Target(str(value)) for value in TypeAdapter(list[str]).validate_python(document["targets"])), + features=None + if document.get("features") is None + else frozenset( + Feature(str(value)) for value in TypeAdapter(list[str]).validate_python(document["features"]) + ), + ) + for document in documents + ) + + @dataclass(frozen=True, slots=True) class Observation: queries: RoutingMap diff --git a/tests/integration/providers/_anthropic_safeguards_beta_support.py b/tests/integration/providers/_anthropic_safeguards_beta_support.py index 0307ea11724..f71819693aa 100644 --- a/tests/integration/providers/_anthropic_safeguards_beta_support.py +++ b/tests/integration/providers/_anthropic_safeguards_beta_support.py @@ -20,6 +20,7 @@ from integration._support.anthropic_sse import ( text_delta, ) from integration._support.client import Gateway, object_value +from integration._support.routing import assert_routing_matches_policy, rule_aware_overrides, trace_path from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire from pydantic import JsonValue, TypeAdapter @@ -219,6 +220,16 @@ def rust_routed(headers: Mapping[str, str]) -> bool: return headers.get(RUST_ROUTE_HEADER) == "true" +def assert_routing_policy(headers: Mapping[str, str], trace_dir: Path) -> None: + trace_id: Final = headers.get("x-integration-routing-id") + assert trace_id is not None, dict(headers) + assert_routing_matches_policy(headers, trace_path(trace_dir, trace_id)) + + +def routing_overrides(directory: Path, *, rust: bool) -> Mapping[str, str]: + return rule_aware_overrides(directory, rust=rust) + + def messages_body(model: str, marker: str, *, stream: bool = False) -> dict[str, JsonValue]: return { "model": model, diff --git a/tests/integration/providers/_idle_upstream_support.py b/tests/integration/providers/_idle_upstream_support.py index a10017425f1..81c43a679f6 100644 --- a/tests/integration/providers/_idle_upstream_support.py +++ b/tests/integration/providers/_idle_upstream_support.py @@ -153,6 +153,7 @@ class Answer: call_id: str response_id: str rust: bool = False + routing_id: str = "" @dataclass(frozen=True, slots=True) @@ -169,6 +170,10 @@ def _rust(headers: httpx.Headers) -> bool: return "true" in headers.get_list(RUST_SERVED) +def _routing_id(headers: httpx.Headers) -> str: + return next(iter(headers.get_list("x-integration-routing-id")), "") + + def _refused(status: int, response: httpx.Response) -> Answer: return Answer(status, response.text, _call_id(response.headers), "") @@ -272,7 +277,7 @@ def _post( if response.status_code != 200: return _refused(response.status_code, response) text, identity = read(_JSON.validate_json(response.content)) - return Answer(200, text, _call_id(response.headers), identity, _rust(response.headers)) + return Answer(200, text, _call_id(response.headers), identity, _rust(response.headers), _routing_id(response.headers)) def _chat_body(body: Mapping[str, JsonValue]) -> tuple[str, str]: @@ -345,6 +350,7 @@ def transcription(target: ProxyTarget, model: str, marker: str) -> Answer: _call_id(response.headers), "", _rust(response.headers), + _routing_id(response.headers), ) @@ -405,6 +411,7 @@ def messages_sdk(target: ProxyTarget, model: str, marker: str) -> Answer: _call_id(raw.headers), message.id, _rust(raw.headers), + _routing_id(raw.headers), ) @@ -432,6 +439,7 @@ async def _messages_stream(target: ProxyTarget, model: str, marker: str) -> Answ _call_id(raw.headers), next((event.message.id for event in events if isinstance(event, RawMessageStartEvent)), ""), _rust(raw.headers), + _routing_id(raw.headers), ) @@ -648,6 +656,7 @@ class UpstreamProxy: workers: tuple[int, int] rows: Mapping[str, Row] gateways: Mapping[str, IdleGateway] + trace_dir: Path @property def target(self) -> ProxyTarget: @@ -712,7 +721,13 @@ def proxy_for( healthcheck=WorkerHealthcheck.LENIENT, state=OwnedState.SHARED_WITH_WORKER, ) as owned: - yield UpstreamProxy(owned, two_workers(owned), MappingProxyType({row.name: row for row in rows}), gateways) + yield UpstreamProxy( + owned, + two_workers(owned), + MappingProxyType({row.name: row for row in rows}), + gateways, + directory / "routing-traces", + ) @contextmanager diff --git a/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py index 8d92fb609e9..079356959ee 100644 --- a/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py +++ b/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py @@ -5,6 +5,7 @@ import pytest from integration._support.client import Gateway from integration._support.idle_gateway import Drop from integration._support.process import graceful_stop_seconds +from integration._support.routing import assert_routing_matches_policy, rule_aware_overrides, trace_path from integration.providers._idle_upstream_support import ( IDLE_SECONDS, Row, @@ -16,7 +17,6 @@ from integration.providers._idle_upstream_support import ( landed_once_undropped, messages_sdk, messages_sdk_stream, - served_by_rust, spend_statuses, stale_pools, ) @@ -24,33 +24,38 @@ from integration.providers._idle_upstream_support import ( pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS)) ROWS: Final = ( - Row("messages-rust-anthropic-sdk-sync", anthropic_deployment, messages_sdk, rust=True), - Row("messages-rust-anthropic-sdk-async-stream", anthropic_deployment, messages_sdk_stream, rust=True), + Row("messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk), + Row("messages-anthropic-sdk-async-stream", anthropic_deployment, messages_sdk_stream), Row( - "messages-rust-upstream-resets-stale-connection", + "messages-upstream-resets-stale-connection", anthropic_deployment, messages_sdk, drop=Drop.RESET, - rust=True, ), ) @pytest.fixture(scope="module") def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]: - with stale_pools( - session_gateway, tmp_path_factory.mktemp("idle-upstream-rust"), ROWS, {"LITELLM_RUST": "1"} - ) as value: + directory: Final = tmp_path_factory.mktemp("idle-upstream-rust") + with stale_pools(session_gateway, directory, ROWS, rule_aware_overrides(directory, rust=True)) as value: yield value @pytest.mark.parametrize("name", [row.name for row in ROWS]) -def test_rust_messages_after_upstream_idle_past_keepalive_open_a_fresh_connection_on_each_worker( +def test_messages_after_upstream_idle_follow_route_policy_and_open_fresh_worker_connections( pools: StalePools, name: str ) -> None: row: Final = pools.rows[name] sents, arrivals = pools.probe(name) assert answered(sents) == expected_answers(row, sents), sents assert landed(arrivals) == landed_once_undropped(sents), arrivals - assert served_by_rust(sents) == (row.rust,) * len(sents), sents + for sent in sents: + assert_routing_matches_policy( + { + "x-integration-routing-id": sent.answer.routing_id, + "x-litellm-rust": "true" if sent.answer.rust else "false", + }, + trace_path(pools.proxy.trace_dir, sent.answer.routing_id), + ) assert spend_statuses(sents) == ("success",) * len(sents), sents diff --git a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py index c13d4da7dbb..93aa9e783e1 100644 --- a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py +++ b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py @@ -5,6 +5,7 @@ import threading from collections.abc import Iterator, Mapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass +from pathlib import Path from typing import Final import anthropic @@ -12,7 +13,6 @@ import httpx import pytest from integration._support.client import ( Gateway, - Scenario, eventually, gateway_from_environment, ) @@ -46,6 +46,7 @@ from integration.providers._anthropic_safeguards_beta_support import ( WEB_SEARCH, WEB_SEARCH_TOOL, anthropic_peer, + assert_routing_policy, chat_body, deliveries, json_object, @@ -55,7 +56,7 @@ from integration.providers._anthropic_safeguards_beta_support import ( outage_peer, owned_config, responses_body, - rust_routed, + routing_overrides, safeguarded, started_workers, ) @@ -79,6 +80,7 @@ class _Rig: owned: OwnedProxy wire: Wire outage: threading.Event + trace_dir: Path @property def gateway(self) -> Gateway: @@ -91,7 +93,6 @@ class _Answer: marker: str status: int text: str - rust: bool @pytest.fixture(scope="module") @@ -103,14 +104,16 @@ def shared_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: wire_server(outage_peer(outage)) as wire, ): config: Final = owned_config(directory, environment, wire.url) - with owned_proxy_process(environment, directory, {"LITELLM_RUST": "true"}, config=config, workers=2) as owned: + with owned_proxy_process( + environment, directory, routing_overrides(directory, rust=True), config=config, workers=2 + ) as owned: eventually( lambda: started_workers(owned.log), lambda pids: len(pids) == 2, seconds=60, ) wire.drain() - yield _Rig(owned, wire, outage) + yield _Rig(owned, wire, outage, directory / "routing-traces") @pytest.fixture @@ -155,7 +158,7 @@ def test_sdk_messages_with_safeguards_go_upstream_with_the_dangerous_tool_use_be extra_body={"safeguards": list(SAFEGUARDS), "tools": [dict(BASH_TOOL)]}, ) message: Final = raw.parse() - assert rust_routed(raw.headers), dict(raw.headers) + assert_routing_policy(raw.headers, rig.trace_dir) received: Final = json_object(message.model_dump_json()) assert message.id == f"msg_{marker}", received assert received.get("safeguard_results") == list(SAFEGUARD_RESULTS), received @@ -192,7 +195,7 @@ async def test_sdk_streamed_messages_with_safeguards_go_upstream_with_the_beta( ) as stream: text: Final = "".join([piece async for piece in stream.text_stream]) final: Final = await stream.get_final_message() - assert rust_routed(stream.response.headers), dict(stream.response.headers) + assert_routing_policy(stream.response.headers, rig.trace_dir) assert (final.id, text) == (f"msg_{marker}", REPLY_TEXT), (final, text) sent: Final = one_delivery(rig.wire, marker) assert (sent.target, sent.beta, sent.safeguards) == ( @@ -209,7 +212,7 @@ def test_raw_messages_merge_the_safeguards_beta_into_the_client_betas(rig: _Rig, body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker, stream=stream), list(SAFEGUARDS)) response: Final = _post(rig.gateway, body, beta=CLIENT_BETAS) assert response.status_code == 200, response.text - assert rust_routed(response.headers), dict(response.headers) + assert_routing_policy(response.headers, rig.trace_dir) assert f"msg_{marker}" in response.text and "safeguard_results" in response.text, response.text sent: Final = one_delivery(rig.wire, marker) assert (sent.target, sent.beta) == (ANTHROPIC_TARGET, MERGED_BETAS), sent @@ -221,7 +224,7 @@ def test_null_or_absent_safeguards_add_no_beta(rig: _Rig, extra: Mapping[str, Js marker: Final = new_marker() response: Final = _post(rig.gateway, {**messages_body(ANTHROPIC_DEPLOYMENT, marker), **extra}) assert response.status_code == 200, response.text - assert rust_routed(response.headers), dict(response.headers) + assert_routing_policy(response.headers, rig.trace_dir) assert f"msg_{marker}" in response.text and "safeguard_results" not in response.text, response.text sent: Final = one_delivery(rig.wire, marker) assert (sent.beta, sent.safeguards) == (None, None), sent @@ -253,7 +256,7 @@ def test_malformed_safeguards_still_add_the_beta_and_relay_the_vendor_verdict( control: Final = new_marker() after: Final = _post(rig.gateway, messages_body(ANTHROPIC_DEPLOYMENT, control)) assert after.status_code == 200, after.text - assert rust_routed(after.headers), dict(after.headers) + assert_routing_policy(after.headers, rig.trace_dir) assert one_delivery(rig.wire, control).beta is None assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 @@ -276,7 +279,7 @@ def test_duplicated_safeguards_keys_resolve_to_the_last_value(rig: _Rig, last: s }, ) assert response.status_code == 200, response.text - assert rust_routed(response.headers), dict(response.headers) + assert_routing_policy(response.headers, rig.trace_dir) assert ("safeguard_results" in response.text) == (last == "list"), response.text sent: Final = one_delivery(rig.wire, marker) assert sent.beta == (DANGEROUS_TOOL_USE if last == "list" else None), sent @@ -304,7 +307,7 @@ def test_vertex_messages_with_safeguards_keep_the_beta_from_the_shared_helper( body: Final = safeguarded(messages_body(VERTEX_DEPLOYMENT, marker, stream=stream), list(SAFEGUARDS)) response: Final = _post(rig.gateway, {**body, "tools": tools}) assert response.status_code == 200, response.text - assert rust_routed(response.headers), dict(response.headers) + assert_routing_policy(response.headers, rig.trace_dir) assert f"msg_{marker}" in response.text and "safeguard_results" in response.text, response.text sent: Final = one_delivery(rig.wire, marker) assert sent.target == (VERTEX_STREAM_TARGET if stream else VERTEX_TARGET), sent @@ -326,24 +329,28 @@ def test_deepseek_messages_forward_safeguards_without_adding_a_beta(rig: _Rig, b beta=beta, ) assert response.status_code == 200, response.text - assert rust_routed(response.headers), dict(response.headers) + assert_routing_policy(response.headers, rig.trace_dir) assert f"msg_{marker}" in response.text, response.text sent: Final = one_delivery(rig.wire, marker) assert (sent.target, sent.api_key, sent.safeguards) == (DEEPSEEK_TARGET, DEEPSEEK_KEY, list(SAFEGUARDS)), sent assert _beta_set(sent.beta) <= _beta_set(beta) - {DANGEROUS_TOOL_USE}, sent -def test_python_route_adds_the_beta_for_safeguards(gateway: Gateway, scenario: Scenario) -> None: - with wire_server(anthropic_peer) as wire: +def test_python_policy_route_adds_the_beta_for_safeguards(gateway: Gateway, tmp_path: Path) -> None: + with ( + wire_server(anthropic_peer) as wire, + owned_proxy_process(gateway, tmp_path, routing_overrides(tmp_path, rust=False)) as owned, + owned.gateway.scenario() as scenario, + ): deployment: Final = scenario.model(model=f"anthropic/{BACKEND}", api_base=wire.url, api_key=ANTHROPIC_KEY) marker: Final = new_marker() - response: Final = _post(gateway, safeguarded(messages_body(deployment, marker), list(SAFEGUARDS))) + response: Final = _post(owned.gateway, safeguarded(messages_body(deployment, marker), list(SAFEGUARDS))) assert response.status_code == 200, response.text assert "safeguard_results" in response.text, response.text (sent,) = deliveries(wire) assert (sent.marker, sent.safeguards) == (marker, list(SAFEGUARDS)), sent assert DANGEROUS_TOOL_USE in _beta_set(sent.beta), sent - assert not rust_routed(response.headers), dict(response.headers) + assert_routing_policy(response.headers, tmp_path / "routing-traces") def _burst_request(index: int, marker: str) -> tuple[str, dict[str, JsonValue]]: @@ -356,21 +363,25 @@ def _burst_request(index: int, marker: str) -> tuple[str, dict[str, JsonValue]]: return "/v1/responses", responses_body(ANTHROPIC_DEPLOYMENT, marker) -def _answer(gateway: Gateway, path: str, marker: str, body: Mapping[str, JsonValue]) -> _Answer: - response: Final = gateway.request("POST", path, body) - return _Answer(path, marker, response.status_code, response.text, rust_routed(response.headers)) +def _answer(rig: _Rig, path: str, marker: str, body: Mapping[str, JsonValue]) -> _Answer: + response: Final = rig.gateway.request("POST", path, body) + if response.status_code == 200: + if path == "/v1/messages": + assert_routing_policy(response.headers, rig.trace_dir) + else: + assert response.headers.get("x-litellm-rust") != "true", dict(response.headers) + return _Answer(path, marker, response.status_code, response.text) -def _wave(pool: ThreadPoolExecutor, gateway: Gateway, markers: tuple[str, ...], offset: int) -> tuple[_Answer, ...]: +def _wave(pool: ThreadPoolExecutor, rig: _Rig, markers: tuple[str, ...], offset: int) -> tuple[_Answer, ...]: planned: Final = tuple((marker, *_burst_request(offset + index, marker)) for index, marker in enumerate(markers)) - submitted: Final = tuple(pool.submit(_answer, gateway, path, marker, body) for marker, path, body in planned) + submitted: Final = tuple(pool.submit(_answer, rig, path, marker, body) for marker, path, body in planned) return tuple(future.result(timeout=120) for future in submitted) def _assert_served(answer: _Answer) -> None: assert answer.status == 200, answer assert REPLY_TEXT in answer.text, answer - assert answer.rust == (answer.path == "/v1/messages"), answer if answer.path == "/v1/messages": assert f"msg_{answer.marker}" in answer.text and "safeguard_results" in answer.text, answer @@ -383,12 +394,12 @@ def test_peer_outage_mid_safeguards_burst_relays_each_verdict_once_and_recovers( outage: Final = tuple(new_marker() for _ in range(_BURST // 2)) recovery: Final = tuple(new_marker() for _ in range(_RECOVERY)) with ThreadPoolExecutor(max_workers=8) as pool: - before: Final = _wave(pool, rig.gateway, healthy, 0) + before: Final = _wave(pool, rig, healthy, 0) rig.outage.set() assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 - during: Final = _wave(pool, rig.gateway, outage, len(healthy)) + during: Final = _wave(pool, rig, outage, len(healthy)) rig.outage.clear() - after: Final = _wave(pool, rig.gateway, recovery, len(healthy) + len(outage)) + after: Final = _wave(pool, rig, recovery, len(healthy) + len(outage)) for answer in (*before, *after): _assert_served(answer) for answer in during: diff --git a/tests/integration/routing/test_rule_aware_integration_harness.py b/tests/integration/routing/test_rule_aware_integration_harness.py new file mode 100644 index 00000000000..2a0dcae1479 --- /dev/null +++ b/tests/integration/routing/test_rule_aware_integration_harness.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.process import owned_proxy_process +from integration._support.routing import assert_routing_matches_policy, install_rule_aware_harness, trace_path +from integration._support.wire import wire_server +from integration.providers._anthropic_safeguards_beta_support import ( + ANTHROPIC_DEPLOYMENT, + ANTHROPIC_KEY, + anthropic_peer, + deliveries, + new_marker, + owned_config, +) + +from litellm.rust_bridge import catalog +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.configuration import Rollout + +_REQUIRED_NATIVE_RULES: Final = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, providers=frozenset({"anthropic"})),) + + +def install_required_native_test_harness() -> None: + setattr(catalog, "RULES", _REQUIRED_NATIVE_RULES) + install_rule_aware_harness() + + +def _native_overrides(directory: Path) -> Mapping[str, str]: + return MappingProxyType( + { + "LITELLM_RUST": "true", + "LITELLM_INTEGRATION_ROUTING_TRACE_DIR": str(directory / "routing-traces"), + "LITELLM_WORKER_STARTUP_HOOKS": ( + "integration.routing.test_rule_aware_integration_harness:install_required_native_test_harness" + ), + } + ) + + +@pytest.mark.timeout(120) +def test_required_native_profile_confirms_supported_messages_request( + gateway: Gateway, + tmp_path: Path, +) -> None: + with wire_server(anthropic_peer) as wire: + config: Final = owned_config(tmp_path, gateway, wire.url) + with owned_proxy_process(gateway, tmp_path, _native_overrides(tmp_path), config=config) as owned: + marker: Final = new_marker() + response: Final = owned.gateway.request( + "POST", + "/v1/messages", + { + "model": ANTHROPIC_DEPLOYMENT, + "max_tokens": 64, + "messages": [{"role": "user", "content": marker}], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert f"msg_{marker}" in response.text, response.text + assert response.headers.get("x-litellm-rust") == "true", dict(response.headers) + trace_id: Final = response.headers["x-integration-routing-id"] + assert_routing_matches_policy(response.headers, trace_path(tmp_path / "routing-traces", trace_id)) + sent: Final = deliveries(wire) + + assert len(sent) == 1, sent + assert sent[0].api_key == ANTHROPIC_KEY, sent + assert sent[0].target == "/v1/messages", sent + assert sent[0].marker == marker, sent diff --git a/tests/test_litellm_rust/test_dispatch.py b/tests/test_litellm_rust/test_dispatch.py index 193866409a2..83d1e801bfd 100644 --- a/tests/test_litellm_rust/test_dispatch.py +++ b/tests/test_litellm_rust/test_dispatch.py @@ -8,6 +8,7 @@ from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import Feature, Route, RouteContext, RouteRule, Rules, Target from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.dispatch import PublicDispatch +from litellm.rust_bridge import observation from litellm.rust_bridge.runtime import NO_PYTHON, NoPythonImplementationError @@ -272,6 +273,79 @@ def test_a_detector_error_keeps_the_call_on_python() -> None: assert result is expected +def test_dispatch_observes_detector_failure_without_projecting_a_python_call() -> None: + class Collector: + def __init__(self) -> None: + self.events: tuple[observation.RouteObservation, ...] = () + + def record(self, event: observation.RouteObservation) -> None: + self.events = (*self.events, event) + + collector: Final = Collector() + token: Final = observation.set_observer(collector) + try: + with observation.request_scope("trace"), observation.trace_scope(): + assert ( + failing_dispatch().run( + ("model",), + {}, + python=lambda *_: "python", + binding=binding(), + native=reject_native, + rules=(RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),), + ) + == "python" + ) + assert any(event.reason == "detector_failure" for event in collector.events) + finally: + observation.reset_observer(token) + + +def test_dispatch_observes_no_candidate_and_projection_failure() -> None: + class Collector: + def __init__(self) -> None: + self.events: tuple[observation.RouteObservation, ...] = () + + def record(self, event: observation.RouteObservation) -> None: + self.events = (*self.events, event) + + collector: Final = Collector() + token: Final = observation.set_observer(collector) + dispatch: Final = PublicDispatch( + route=Route.CHAT_COMPLETIONS, + request=lambda args, kwargs: None, + context=lambda _: RouteContext(Route.CHAT_COMPLETIONS), + ) + try: + with observation.request_scope("trace"), observation.trace_scope(): + assert ( + dispatch.run( + ("model",), + {}, + python=lambda *_: "python", + binding=binding(), + native=reject_native, + rules=(), + ) + == "python" + ) + with observation.request_scope("trace"), observation.trace_scope(): + assert ( + dispatch.run( + ("model",), + {}, + python=lambda *_: "python", + binding=binding(), + native=reject_native, + rules=(RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),), + ) + == "python" + ) + assert {event.reason for event in collector.events} == {"no_candidate", "projection_failure"} + finally: + observation.reset_observer(token) + + async def test_a_detector_error_keeps_the_async_call_on_python() -> None: expected: Final = object() diff --git a/tests/unit/rust_bridge/test_observation.py b/tests/unit/rust_bridge/test_observation.py new file mode 100644 index 00000000000..e071888e2ab --- /dev/null +++ b/tests/unit/rust_bridge/test_observation.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Final + +import pytest + +from litellm.rust_bridge import bindings, observation, runtime +from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule +from litellm.rust_bridge.configuration import Decision, Rollout + + +@dataclass +class Collector: + events: tuple[observation.RouteObservation, ...] = () + + def record(self, event: observation.RouteObservation) -> None: + observation.record_event(event) + self.events = (*self.events, event) + + +def binding() -> bindings.NativeBinding[object]: + bound: Final = bindings.NativeBinding("observation", validate=lambda _: None) + bound.override(lambda: "native") + return bound + + +def test_runtime_observes_context_decision_and_terminal_engine() -> None: + collector: Final = Collector() + token: Final = observation.set_observer(collector) + try: + with observation.request_scope("trace"), observation.trace_scope(): + assert ( + runtime.run( + RouteContext(Route.MESSAGES, provider="anthropic", model="model"), + binding=binding(), + native=lambda hook: hook(), + python=lambda: "python", + rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),), + ) + == "native" + ) + selected: Final = next(event for event in collector.events if event.event == "selected") + terminal: Final = next(event for event in collector.events if event.event == "terminal") + assert selected.context == RouteContext(Route.MESSAGES, provider="anthropic", model="model") + assert selected.decision is Decision.RUST_REQUIRED + assert terminal.engine == "rust" + assert terminal.outcome == "success" + finally: + observation.reset_observer(token) + + +def test_execution_error_is_observed_as_failure_and_re_raised() -> None: + collector: Final = Collector() + token: Final = observation.set_observer(collector) + failure: Final = ValueError("native failure") + try: + with ( + observation.request_scope("trace"), + observation.trace_scope(), + pytest.raises(ValueError, match="native failure") as raised, + ): + runtime.run( + RouteContext(Route.MESSAGES), + binding=binding(), + native=lambda hook: (_ for _ in ()).throw(failure), + python=lambda: "python", + rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),), + ) + assert raised.value is failure + terminal: Final = next(event for event in collector.events if event.event == "terminal") + assert (terminal.outcome, terminal.reason) == ("failure", "execution_error") + finally: + observation.reset_observer(token) + + +@pytest.mark.asyncio +async def test_async_request_scopes_are_isolated() -> None: + collector: Final = Collector() + token: Final = observation.set_observer(collector) + + async def call(trace_id: str) -> None: + with observation.request_scope(trace_id), observation.trace_scope(): + await runtime.arun( + RouteContext(Route.MESSAGES, model=trace_id), + binding=binding(), + native=lambda hook: asyncio.sleep(0, result=hook()), + python=lambda: asyncio.sleep(0, result="python"), + rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),), + ) + + try: + await asyncio.gather(call("one"), call("two")) + selected: Final = tuple(event for event in collector.events if event.event == "selected") + assert {event.trace_id for event in selected} == {"one", "two"} + assert {event.context.model for event in selected if event.context is not None} == {"one", "two"} + finally: + observation.reset_observer(token) + + +@pytest.mark.asyncio +async def test_copied_worker_context_records_dispatch_events_in_request_trace() -> None: + collector: Final = Collector() + token: Final = observation.set_observer(collector) + try: + with observation.request_scope("thread"), observation.trace_scope(): + await asyncio.to_thread( + lambda: runtime.run( + RouteContext(Route.MESSAGES, model="thread"), + binding=binding(), + native=lambda hook: hook(), + python=lambda: "python", + rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),), + ) + ) + events: Final = observation.current_events() + assert {event.event for event in events} == {"selected", "terminal"} + assert {event.trace_id for event in events} == {"thread"} + finally: + observation.reset_observer(token) + + +def test_nested_attempts_have_parent_depth_and_cleanup() -> None: + collector: Final = Collector() + token: Final = observation.set_observer(collector) + try: + with observation.request_scope("trace"), observation.trace_scope(): + with observation.observed_attempt(RouteContext(Route.MESSAGES), Decision.RUST_REQUIRED) as parent: + with observation.observed_attempt(RouteContext(Route.MESSAGES), Decision.RUST_REQUIRED) as child: + observation.observe_terminal( + child, RouteContext(Route.MESSAGES), Decision.RUST_REQUIRED, engine="rust" + ) + observation.observe_terminal(parent, RouteContext(Route.MESSAGES), Decision.RUST_REQUIRED, engine="rust") + with observation.observed_attempt(RouteContext(Route.MESSAGES), Decision.RUST_REQUIRED) as next_attempt: + observation.observe_terminal( + next_attempt, RouteContext(Route.MESSAGES), Decision.RUST_REQUIRED, engine="rust" + ) + nested: Final = tuple(event for event in collector.events if event.event == "terminal") + assert nested[0].depth == 1 + assert nested[1].depth == 0 + assert nested[0].parent_attempt_id == nested[1].attempt_id + assert nested[2].depth == 0 + assert nested[2].parent_attempt_id is None + finally: + observation.reset_observer(token) diff --git a/tests/unit/rust_bridge/test_routing_harness.py b/tests/unit/rust_bridge/test_routing_harness.py new file mode 100644 index 00000000000..ac4be0e71df --- /dev/null +++ b/tests/unit/rust_bridge/test_routing_harness.py @@ -0,0 +1,429 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.routing import ( + _RoutingTraceMiddleware, + _TraceObserver, + assert_routing_matches_policy, + strict_current_policy, + trace_path, +) +from pydantic import TypeAdapter +from starlette.types import Receive, Scope, Send + +from litellm.rust_bridge.catalog import Feature, Route, RouteRule, Target +from litellm.rust_bridge.configuration import Decision, Rollout, reset_rust_configuration + + +@pytest.mark.parametrize( + ("switch", "optional_rollout"), + ((True, Rollout.RUST_REQUIRED), (False, Rollout.PYTHON_ONLY)), +) +def test_strict_policy_resolves_synthetic_rollouts_in_order( + monkeypatch: pytest.MonkeyPatch, switch: bool, optional_rollout: Rollout +) -> None: + monkeypatch.setenv("LITELLM_RUST", str(switch)) + reset_rust_configuration() + rules: Final = ( + RouteRule( + Route.MESSAGES, + Rollout.RUST_OPT_IN, + providers=frozenset({"anthropic"}), + targets=frozenset({Target.MESSAGES}), + ), + RouteRule( + Route.MESSAGES, + Rollout.PYTHON_ONLY, + features=frozenset({Feature.SAFEGUARDS_CLASSIFIER}), + ), + RouteRule(Route.RESPONSES, Rollout.RUST_OPT_OUT, models=frozenset({"synthetic"})), + RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"openai"})), + ) + resolved: Final = strict_current_policy(rules) + reset_rust_configuration() + + assert tuple(rule.route for rule in resolved) == tuple(rule.route for rule in rules) + assert tuple(rule.rollout for rule in resolved) == ( + optional_rollout, + Rollout.PYTHON_ONLY, + optional_rollout, + Rollout.RUST_REQUIRED, + ) + assert tuple((rule.providers, rule.models, rule.targets, rule.features) for rule in resolved) == tuple( + (rule.providers, rule.models, rule.targets, rule.features) for rule in rules + ) + + +def _trace( + directory: Path, + *, + trace_id: str, + rollout: Rollout, + decision: Decision, + engine: str, +) -> Path: + destination: Final = directory / f"trace-{trace_id}.json" + destination.write_text( + json.dumps( + { + "trace_id": trace_id, + "policy": [ + { + "route": Route.MESSAGES.value, + "rollout": rollout.value, + "providers": None, + "models": None, + "targets": None, + "features": None, + } + ], + "events": [ + { + "event": "terminal", + "trace_id": trace_id, + "attempt_id": "attempt", + "parent_attempt_id": None, + "depth": 0, + "route": Route.MESSAGES.value, + "context": { + "route": Route.MESSAGES.value, + "provider": "anthropic", + "model": "synthetic", + "target": Target.MESSAGES.value, + "features": [], + }, + "decision": decision.value, + "outcome": "success", + "reason": None, + "engine": engine, + }, + { + "event": "request", + "trace_id": trace_id, + "attempt_id": None, + "parent_attempt_id": None, + "depth": 0, + "route": None, + "context": None, + "decision": None, + "outcome": "completed", + "reason": None, + "engine": None, + }, + ], + } + ) + ) + return destination + + +def test_policy_assertion_tracks_the_injected_python_or_native_rule(tmp_path: Path) -> None: + python_trace: Final = _trace( + tmp_path, + trace_id="python", + rollout=Rollout.PYTHON_ONLY, + decision=Decision.PYTHON, + engine="python", + ) + assert_routing_matches_policy({"x-integration-routing-id": "python"}, python_trace) + + native_trace: Final = _trace( + tmp_path, + trace_id="native", + rollout=Rollout.RUST_REQUIRED, + decision=Decision.RUST_REQUIRED, + engine="rust", + ) + assert_routing_matches_policy( + {"x-integration-routing-id": "native", "x-litellm-rust": "true"}, + native_trace, + ) + + +@pytest.mark.asyncio +async def test_middleware_generates_and_persists_a_request_trace(tmp_path: Path) -> None: + rules: Final = (RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY),) + observer: Final = _TraceObserver(rules, tmp_path) + + async def app(scope: Scope, _receive: Receive, send: Send) -> None: + assert scope["type"] == "http" + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"x-integration-routing-id", b"client-value")], + } + ) + await send( + { + "type": "http.response.body", + "body": b"ok", + "more_body": False, + } + ) + + transport: Final = httpx.ASGITransport(app=_RoutingTraceMiddleware(app, observer)) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response: Final = await client.get("/", headers={"x-integration-routing-id": "client-value"}) + + trace_id: Final = response.headers["x-integration-routing-id"] + assert trace_id != "client-value" + assert trace_path(tmp_path, trace_id).exists() + document: Final = TypeAdapter(dict[str, object]).validate_python( + json.loads(trace_path(tmp_path, trace_id).read_text()) + ) + events: Final = TypeAdapter(list[dict[str, object]]).validate_python(document["events"]) + assert events[-1]["outcome"] == "no_dispatch" + + +def test_policy_assertion_rejects_missing_trace_and_mismatched_engine(tmp_path: Path) -> None: + with pytest.raises(AssertionError, match="routing trace missing"): + assert_routing_matches_policy({"x-integration-routing-id": "missing"}, tmp_path / "trace-missing.json") + + trace: Final = _trace( + tmp_path, + trace_id="mismatch", + rollout=Rollout.RUST_REQUIRED, + decision=Decision.RUST_REQUIRED, + engine="python", + ) + with pytest.raises(AssertionError): + assert_routing_matches_policy( + {"x-integration-routing-id": "mismatch", "x-litellm-rust": "true"}, + trace, + ) + + +def test_no_dispatch_trace_validates_a_python_exclusion_without_inventing_a_terminal( + tmp_path: Path, +) -> None: + trace_id: Final = "early" + trace: Final = tmp_path / f"trace-{trace_id}.json" + trace.write_text( + json.dumps( + { + "trace_id": trace_id, + "policy": [ + { + "route": Route.MESSAGES.value, + "rollout": Rollout.PYTHON_ONLY.value, + "providers": None, + "models": None, + "targets": None, + "features": None, + } + ], + "events": [ + { + "event": "early", + "trace_id": trace_id, + "attempt_id": None, + "parent_attempt_id": None, + "depth": 0, + "route": Route.MESSAGES.value, + "context": None, + "decision": None, + "outcome": None, + "reason": "no_candidate", + "engine": None, + }, + { + "event": "request", + "trace_id": trace_id, + "attempt_id": None, + "parent_attempt_id": None, + "depth": 0, + "route": None, + "context": None, + "decision": None, + "outcome": "no_dispatch", + "reason": None, + "engine": None, + }, + ], + } + ) + ) + + assert_routing_matches_policy({"x-integration-routing-id": trace_id}, trace) + + +def _request_event(trace_id: str, outcome: str) -> dict[str, object]: + return { + "event": "request", + "trace_id": trace_id, + "attempt_id": None, + "parent_attempt_id": None, + "depth": 0, + "route": None, + "context": None, + "decision": None, + "outcome": outcome, + "reason": None, + "engine": None, + } + + +def _early_event(trace_id: str, reason: str) -> dict[str, object]: + return { + "event": "early", + "trace_id": trace_id, + "attempt_id": None, + "parent_attempt_id": None, + "depth": 0, + "route": Route.MESSAGES.value, + "context": None, + "decision": None, + "outcome": None, + "reason": reason, + "engine": None, + } + + +def _terminal_event(trace_id: str, outcome: str, decision: Decision, engine: str) -> dict[str, object]: + return { + "event": "terminal", + "trace_id": trace_id, + "attempt_id": f"attempt-{engine}-{outcome}", + "parent_attempt_id": None, + "depth": 0, + "route": Route.MESSAGES.value, + "context": { + "route": Route.MESSAGES.value, + "provider": "anthropic", + "model": "synthetic", + "target": Target.MESSAGES.value, + "features": [], + }, + "decision": decision.value, + "outcome": outcome, + "reason": None if outcome == "success" else "execution_error", + "engine": engine, + } + + +def _event_trace( + directory: Path, + *, + trace_id: str, + rollout: Rollout, + events: tuple[dict[str, object], ...], +) -> Path: + destination: Final = directory / f"trace-{trace_id}.json" + destination.write_text( + json.dumps( + { + "trace_id": trace_id, + "policy": [ + { + "route": Route.MESSAGES.value, + "rollout": rollout.value, + "providers": None, + "models": None, + "targets": None, + "features": None, + } + ], + "events": events, + } + ) + ) + return destination + + +@pytest.mark.parametrize( + ("trace_id", "rollout", "events", "headers"), + ( + ("missing-dispatch", Rollout.PYTHON_ONLY, (_request_event("missing-dispatch", "no_dispatch"),), {}), + ( + "detector-failure", + Rollout.PYTHON_ONLY, + (_early_event("detector-failure", "detector_failure"), _request_event("detector-failure", "no_dispatch")), + {}, + ), + ( + "projection-failure", + Rollout.PYTHON_ONLY, + ( + _early_event("projection-failure", "projection_failure"), + _request_event("projection-failure", "no_dispatch"), + ), + {}, + ), + ( + "failed-only", + Rollout.RUST_REQUIRED, + ( + _terminal_event("failed-only", "failure", Decision.RUST_REQUIRED, "rust"), + _request_event("failed-only", "completed"), + ), + {"x-litellm-rust": "true"}, + ), + ( + "false-early-marker", + Rollout.PYTHON_ONLY, + (_early_event("false-early-marker", "no_candidate"), _request_event("false-early-marker", "no_dispatch")), + {"x-litellm-rust": "true"}, + ), + ( + "optional", + Rollout.RUST_OPT_IN, + (_terminal_event("optional", "success", Decision.RUST_WITH_FALLBACK, "rust"),), + {"x-litellm-rust": "true"}, + ), + ( + "ambiguous", + Rollout.RUST_REQUIRED, + ( + _terminal_event("ambiguous", "success", Decision.RUST_REQUIRED, "rust"), + _terminal_event("ambiguous", "success", Decision.RUST_REQUIRED, "rust-2"), + _request_event("ambiguous", "completed"), + ), + {"x-litellm-rust": "true"}, + ), + ( + "trace-mismatch", + Rollout.PYTHON_ONLY, + ( + _terminal_event("different-trace", "success", Decision.PYTHON, "python"), + _request_event("trace-mismatch", "completed"), + ), + {}, + ), + ), +) +def test_policy_assertion_rejects_unproven_dispatch_outcomes( + tmp_path: Path, + trace_id: str, + rollout: Rollout, + events: tuple[dict[str, object], ...], + headers: dict[str, str], +) -> None: + trace: Final = _event_trace(tmp_path, trace_id=trace_id, rollout=rollout, events=events) + + with pytest.raises(AssertionError): + assert_routing_matches_policy({"x-integration-routing-id": trace_id, **headers}, trace) + + +def test_failed_root_attempt_can_precede_one_successful_root_attempt(tmp_path: Path) -> None: + trace_id: Final = "retry" + trace: Final = _event_trace( + tmp_path, + trace_id=trace_id, + rollout=Rollout.RUST_REQUIRED, + events=( + _terminal_event(trace_id, "failure", Decision.RUST_REQUIRED, "rust"), + _terminal_event(trace_id, "success", Decision.RUST_REQUIRED, "rust"), + _request_event(trace_id, "completed"), + ), + ) + + assert_routing_matches_policy( + {"x-integration-routing-id": trace_id, "x-litellm-rust": "true"}, + trace, + )