From a278ee0ef11ab62332f09db445bfcba9be5e75b6 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 12:59:51 -0700 Subject: [PATCH] refactor(native): move chat eligibility checks into Rust --- .../crates/core/src/chat_completions/mod.rs | 3 +- .../core/src/chat_completions/prepare.rs | 4 +- .../crates/core/src/chat_completions/tests.rs | 84 ++++++++++++++++++- .../src/chat_completions/transformation.rs | 2 + .../anthropic/chat_completions/tests.rs | 7 +- .../chat_completions/transformation.rs | 10 +++ .../bedrock/chat_completions/tests.rs | 7 +- .../chat_completions/transformation.rs | 6 ++ .../src/routes/chat_completions.rs | 17 +++- litellm/rust_bridge/chat_completions.py | 77 +++++------------ litellm/rust_bridge/protocols.py | 3 + .../chat/test_anthropic_chat_handler.py | 2 +- .../chat/test_bedrock_converse_handler.py | 5 +- .../rust_bridge/test_chat_completions.py | 30 ++++--- 14 files changed, 182 insertions(+), 75 deletions(-) diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index b6139a57b70..7987542eb57 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -45,6 +45,7 @@ pub fn chat_completions_decline_reason( custom_llm_provider: Option<&str>, messages: Value, optional_params: &Map, + context: &LiteLlmRequestContext, ) -> Option<&'static str> { let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else { return Some("provider is not on the rust chat completions path"); @@ -56,7 +57,7 @@ pub fn chat_completions_decline_reason( return Some("empty message list"); } config - .unsupported_reason(&messages, optional_params) + .unsupported_reason(&messages, optional_params, context) .map(|reason| reason.0) } diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 92012e9cf4a..f6ef44c2c0c 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -42,7 +42,7 @@ pub(super) fn parse_messages(messages: Value) -> Result, Error> pub(super) fn resolve_request( request: ChatCompletionsRequest<'_>, options: RequestOptions, - _context: &LiteLlmRequestContext, + context: &LiteLlmRequestContext, ) -> Result { let (model, config) = resolve_provider_config(request.model, options.custom_llm_provider.as_deref()) @@ -52,7 +52,7 @@ pub(super) fn resolve_request( if messages.is_empty() { return Err(Error::Declined("empty message list")); } - if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) { + if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params, context) { return Err(Error::Declined(reason.0)); } Ok(ResolvedChatCompletionsRequest { diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index a02904db716..ec822c1f37b 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -513,7 +513,13 @@ fn decline_reason( Value::Object(map) => map, other => panic!("params must be an object, got {other}"), }; - super::chat_completions_decline_reason(model, provider, messages, ¶ms) + super::chat_completions_decline_reason( + model, + provider, + messages, + ¶ms, + &LiteLlmRequestContext::default(), + ) } #[test] @@ -877,3 +883,79 @@ mod round_trip { )); } } + +#[test] +fn preflight_and_execution_share_provider_metadata_eligibility() { + let messages = json!([{"role": "user", "content": "hi"}]); + let cases = [ + ("anthropic", json!({"user_id": "u-123"}), vec![], true), + ("anthropic", json!({"user_id": null}), vec![], false), + ("anthropic", json!({"trace_id": "t-1"}), vec![], false), + ( + "anthropic", + json!({}), + vec!["user_api_key_team_id".into()], + false, + ), + #[cfg(feature = "bedrock-auth")] + ("bedrock", json!({"user_id": "u-123"}), vec![], false), + #[cfg(feature = "bedrock-auth")] + ( + "bedrock", + json!({}), + vec!["user_api_key_team_id".into()], + true, + ), + ]; + for (provider, metadata, request_metadata_fields, expected_decline) in cases { + let context = LiteLlmRequestContext { + metadata: metadata.as_object().cloned(), + request_metadata_fields, + ..Default::default() + }; + let params = Map::new(); + let preflight = super::chat_completions_decline_reason( + "claude-sonnet-4-5", + Some(provider), + messages.clone(), + ¶ms, + &context, + ); + let execution = resolve_request( + ChatCompletionsRequest { + model: "claude-sonnet-4-5", + messages: messages.clone(), + optional_params: params, + options: RequestOptions { + custom_llm_provider: Some(provider.into()), + ..Default::default() + }, + }, + &context, + ); + assert_eq!( + preflight.is_some(), + expected_decline, + "{provider} preflight" + ); + assert_eq!(execution.is_err(), expected_decline, "{provider} execution"); + } +} + +#[test] +fn anthropic_preflight_preserves_litellm_metadata_scope() { + let context = LiteLlmRequestContext { + litellm_metadata: json!({"user_id": "u-123"}).as_object().cloned(), + ..Default::default() + }; + assert_eq!( + super::chat_completions_decline_reason( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + &Map::new(), + &context, + ), + None, + ); +} diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index d7b9704c46c..b71adb8a8cd 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -1,4 +1,5 @@ use crate::Error; +use crate::request_context::LiteLlmRequestContext; use serde_json::{Map, Value}; use super::types::{ @@ -75,6 +76,7 @@ pub trait ChatCompletionsProviderConfig: Sync { &self, messages: &[ChatMessage], optional_params: &Map, + _context: &LiteLlmRequestContext, ) -> Option { unsupported_param( self.supported_openai_params(), diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs index b22de6c47de..8142374182c 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::Error; +use crate::request_context::LiteLlmRequestContext; use serde_json::json; fn messages(value: Value) -> Vec { @@ -26,7 +27,11 @@ fn transform_response(body: Value) -> Result { } fn reason(msgs: Value, opts: Value) -> Option { - ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) + ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason( + &messages(msgs), + ¶ms(opts), + &LiteLlmRequestContext::default(), + ) } #[test] diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index a7d5a8ad0cf..f4a1d043457 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -1,3 +1,4 @@ +use crate::request_context::LiteLlmRequestContext; use serde_json::{Map, Value, json}; use crate::chat_completions::conversation::{Conversation, build_conversation}; @@ -126,9 +127,18 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { &self, messages: &[ChatMessage], optional_params: &Map, + context: &LiteLlmRequestContext, ) -> Option { unsupported_param(self.supported_openai_params(), &[], optional_params) .or_else(|| messages.iter().find_map(unsupported_message)) + .or_else(|| { + context + .metadata + .as_ref() + .and_then(|metadata| metadata.get("user_id")) + .is_some_and(|value| !value.is_null()) + .then_some(Unsupported("LiteLLM user metadata")) + }) // Anthropic rejects a request whose first turn is not a user turn. // Python only repairs that under `litellm.modify_params`, which the // core cannot observe, so decline instead of guessing. diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs index c86f061b9ca..7746881de49 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::Error; +use crate::request_context::LiteLlmRequestContext; use serde_json::json; fn messages(value: Value) -> Vec { @@ -32,7 +33,11 @@ fn transform_response(body: Value) -> Result { } fn reason(msgs: Value, opts: Value) -> Option { - BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) + BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason( + &messages(msgs), + ¶ms(opts), + &LiteLlmRequestContext::default(), + ) } #[test] diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index 7be3d108d44..39e552f7b0a 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -1,3 +1,4 @@ +use crate::request_context::LiteLlmRequestContext; use serde_json::{Map, Value, json}; use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation}; @@ -176,6 +177,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { &self, messages: &[ChatMessage], optional_params: &Map, + context: &LiteLlmRequestContext, ) -> Option { unsupported_param( self.supported_openai_params(), @@ -183,6 +185,10 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { optional_params, ) .or_else(|| messages.iter().find_map(unsupported_message)) + .or_else(|| { + (!context.request_metadata_fields.is_empty()) + .then_some(Unsupported("LiteLLM request metadata forwarding")) + }) // Python's Converse translation drops blank text blocks instead of // substituting the placeholder the shared conversation builder // applies, so decline blank text rather than diverge. diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index aad6c3c737f..738384033f5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -5,6 +5,7 @@ use litellm_core::chat_completions::chat_completions as run_route; use litellm_core::chat_completions::chat_completions_decline_reason; use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; use litellm_core::request_context::LiteLlmRequestContext; +use litellm_core::request_options::RequestOptions; use pyo3::prelude::*; use serde_json::{Map, Value}; use std::future::Future; @@ -40,13 +41,25 @@ fn prepare_chat_completions( } #[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] +#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None, *, options, context))] fn chat_completions_decline( model: String, #[pyo3(from_py_with = litellm_python_interop::from_py)] messages: Value, #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, custom_llm_provider: Option, + options: NativeRequestOptions, + context: NativeRequestContext, ) -> PyResult> { + if !matches!( + custom_llm_provider.as_deref(), + Some("anthropic" | "bedrock") + ) { + return Ok(Some( + "provider is not on the rust chat completions path".into(), + )); + } + let context: LiteLlmRequestContext = context.into(); + let options: RequestOptions = options.into(); let optional_params = match optional_params { None | Some(Value::Null) => Map::new(), Some(Value::Object(params)) => params, @@ -61,6 +74,8 @@ fn chat_completions_decline( custom_llm_provider.as_deref(), messages, &optional_params, + &options, + &context, ) .map(str::to_string)) } diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index 553b8b60e6a..be5801aa74a 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -10,16 +10,15 @@ from __future__ import annotations import json from collections.abc import Mapping, Sequence +from dataclasses import replace from typing import TYPE_CHECKING, Final, Protocol import httpx -from pydantic import TypeAdapter, ValidationError -from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( convert_to_model_response_object, ) -from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned +from litellm.llms.bedrock.request_metadata import get_bedrock_request_metadata_fields from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.protocols import ( @@ -35,6 +34,8 @@ from litellm.rust_bridge.request import ( NativeRequestContext, NativeRequestOptions, PreparedNativeCall, + anthropic_options, + bedrock_options, call_native, with_capabilities, ) @@ -45,14 +46,6 @@ from litellm.types.utils import ModelResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -# Providers whose `/chat/completions` deployments the Rust core can serve. A -# provider outside this set never reaches the bridge. -RUST_CHAT_COMPLETIONS_PROVIDERS: Final = frozenset({"anthropic", "bedrock"}) - -# `litellm_params` values are `object`, so validate the one this module reads -# rather than narrowing an unparameterized `Mapping` and typing the result Any. -_LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) - RUST_RESPONSE_HEADER: Final = "x-litellm-rust" @@ -128,42 +121,21 @@ def set_rust_chat_completions( _CHAT_PREFLIGHT.override(decline) -def _anthropic_user_id_reaches_the_body(litellm_params: Mapping[str, object] | None) -> bool: - metadata: Final = litellm_params.get("metadata") if litellm_params is not None else None - try: - entries: Final = _LITELLM_METADATA_ADAPTER.validate_python(metadata) - except ValidationError: - return False - return entries.get("user_id") is not None - - -def _litellm_metadata_reaches_the_provider( - custom_llm_provider: str | None, litellm_params: Mapping[str, object] | None -) -> bool: - """Whether the Python transform would promote proxy-owned attribution into the - provider request, below this gate and inside the function the Rust route replaces. - - `AnthropicConfig.transform_request` promotes a valid `metadata["user_id"]` - into the Messages body, so the core never sees the key and would send the - request to Anthropic with the abuse-detection attribution missing. - - `AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the - Converse body whenever the operator armed `bedrock_request_metadata_fields`. - Owning that field also means evicting a caller-supplied one, which the core - cannot do either, so ownership alone is the condition rather than whether - anything resolved. - - Deliberately a superset of Python's condition in both cases: declining a - request Python would not have attributed anyway costs only the Rust path, - while missing one loses the attribution silently. - """ - match custom_llm_provider: - case "anthropic": - return _anthropic_user_id_reaches_the_body(litellm_params) - case "bedrock": - return bedrock_request_metadata_is_owned() - case _: - return False +def _provider_eligibility_options( + provider: str | None, + litellm_params: Mapping[str, object] | None, + optional_params: Mapping[str, object], +) -> NativeRequestOptions: + bedrock: Final = ( + replace( + bedrock_options(optional_params), + request_metadata_fields=get_bedrock_request_metadata_fields(), + ) + if provider == "bedrock" + else None + ) + anthropic: Final = anthropic_options(litellm_params) if provider == "anthropic" else None + return NativeRequestOptions(custom_llm_provider=provider, bedrock=bedrock, anthropic=anthropic) def rust_chat_completions_accepts( @@ -182,13 +154,6 @@ def rust_chat_completions_accepts( capability gate answers the second half; it resolves no credentials and performs no I/O. """ - if custom_llm_provider not in RUST_CHAT_COMPLETIONS_PROVIDERS: - return False - if stream: - return False - if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): - verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") - return False if not rust_enabled(): return False decline: Final = _CHAT_PREFLIGHT.load() @@ -200,6 +165,10 @@ def rust_chat_completions_accepts( messages=messages, optional_params=optional_params, custom_llm_provider=custom_llm_provider, + options=_provider_eligibility_options(custom_llm_provider, litellm_params, optional_params), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities(stream=bool(stream)) + ), ) except Exception as error: # noqa: BLE001 # capability checks perform no provider I/O verbose_logger.debug("Native chat acceptance check failed: %s", error) diff --git a/litellm/rust_bridge/protocols.py b/litellm/rust_bridge/protocols.py index dc779b346ab..0b1055aa01c 100644 --- a/litellm/rust_bridge/protocols.py +++ b/litellm/rust_bridge/protocols.py @@ -31,6 +31,9 @@ class RustChatCompletionsDecline(Protocol): messages: Sequence[object], optional_params: Mapping[str, object] | None, custom_llm_provider: str | None, + *, + options: NativeRequestOptions, + context: NativeRequestContext, ) -> str | None: ... diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index db73a349b79..d448cbafaae 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2427,7 +2427,7 @@ class TestRustChatCompletionsHook: ) except Exception: pass - assert seen["gate"] == [] + assert seen["call"] == [] def test_pre_call_logging_fires_exactly_once_on_the_rust_path(self): from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index bd1d93e7a3a..a2b5714f3f7 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -181,11 +181,14 @@ def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch): def test_streaming_stays_on_the_python_path(): seen = _inject() + bridge.set_rust_chat_completions( + decline=pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline + ) try: _run(optional_params={"maxTokens": 16, "stream": True}) except Exception: pass - assert seen["gate"] == [] + assert seen["call"] == [] def test_a_declined_request_never_reaches_the_native_call(): diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 31163ebaa05..5888b944118 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -176,25 +176,17 @@ class TestGate: def test_declines_streaming_and_providers_off_the_path(self, monkeypatch): monkeypatch.setenv("LITELLM_RUST", "1") - gate = _RecordingDecline() + gate = pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline bridge.set_rust_chat_completions(decline=gate) assert _accepts(stream=True) is False assert _accepts(custom_llm_provider="openai") is False assert _accepts(custom_llm_provider=None) is False - assert gate.calls == [] def test_declines_an_anthropic_request_carrying_a_litellm_metadata_user_id(self, monkeypatch): - """`AnthropicConfig.transform_request` copies a valid `user_id` into the Messages body. - - It does that inside the function the Rust route replaces, and the core is - handed `optional_params` only, so accepting here would send the request - to Anthropic with the abuse-detection attribution silently missing. - """ monkeypatch.setenv("LITELLM_RUST", "1") - gate = _RecordingDecline() + gate = pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline bridge.set_rust_chat_completions(decline=gate) assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False - assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of" # Bedrock's Converse transform reads no `user_id`, and an Anthropic request # whose metadata carries none is one Python would not attribute either. @@ -202,6 +194,7 @@ class TestGate: _accepts( custom_llm_provider="bedrock", model="bedrock/us-east-1/anthropic.claude-v2", + optional_params={"maxTokens": 16}, litellm_params={"metadata": {"user_id": "u-123"}}, ) is True @@ -209,6 +202,19 @@ class TestGate: assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True assert _accepts(litellm_params={"metadata": None}) is True + assert _accepts(litellm_params={"litellm_metadata": {"user_id": "u-123"}}) is True + assert _accepts(litellm_params={"metadata": "invalid"}) is True + assert _accepts(litellm_params={"metadata": {"trace": object()}}) is True + assert _accepts(litellm_params={"metadata": {"user_id": object()}}) is False + assert ( + _accepts( + custom_llm_provider="bedrock", + model="bedrock/us-east-1/anthropic.claude-v2", + optional_params={"maxTokens": 16}, + litellm_params={"metadata": {"user_id": object()}}, + ) + is True + ) def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch): """`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the @@ -217,16 +223,16 @@ class TestGate: who armed `bedrock_request_metadata_fields` keeps the Python path. """ monkeypatch.setenv("LITELLM_RUST", "1") - gate = _RecordingDecline() + gate = pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline bridge.set_rust_chat_completions(decline=gate) bedrock = { "custom_llm_provider": "bedrock", "model": "bedrock/us-east-1/anthropic.claude-v2", + "optional_params": {"maxTokens": 16}, } monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ["user_api_key_team_id"]) assert _accepts(**bedrock) is False - assert gate.calls == [], "the core must not be consulted for a field it cannot write" assert _accepts() is True, "arming Bedrock attribution must not decline Anthropic" monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None)