From 7586a4f8e1b788282d9c2fa650aa56b980d5ebec Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 02:05:11 +0000 Subject: [PATCH] refactor(auth): move Claude normalization into native read stages --- litellm-rust/crates/python-bridge/src/lib.rs | 14 +- .../crates/python-bridge/src/request_auth.rs | 24 ++ .../src/claude_membership.rs | 100 +++++++ .../src/claude_normalization.rs | 176 ++++++++++++ .../request-auth-core/src/codex_models.rs | 79 ++++++ .../crates/request-auth-core/src/lib.rs | 2 + litellm/_auth/claude_normalization.py | 74 +++++ .../auth/claude_normalization_adapter.py | 174 ++++++++++++ litellm/proxy/auth/codex_model_adapter.py | 46 +-- .../proxy/auth/codex_normalization_adapter.py | 73 ++--- litellm/proxy/auth/codex_turn_adapter.py | 16 +- .../proxy/auth/key_model_access_adapter.py | 6 +- litellm/proxy/auth/model_policy_read.py | 11 + litellm/rust_bridge/_native.pyi | 2 + tests/unit/_auth/support/claude_reference.py | 103 +++++++ tests/unit/_auth/test_claude_normalization.py | 264 ++++++++++++++++++ tests/unit/_auth/test_codex_models.py | 8 +- tests/unit/_auth/test_codex_normalization.py | 4 +- tests/unit/_auth/test_codex_turn.py | 6 +- .../unit/proxy/auth/test_model_policy_read.py | 18 ++ 20 files changed, 1118 insertions(+), 82 deletions(-) create mode 100644 litellm-rust/crates/request-auth-core/src/claude_membership.rs create mode 100644 litellm-rust/crates/request-auth-core/src/claude_normalization.rs create mode 100644 litellm/_auth/claude_normalization.py create mode 100644 litellm/proxy/auth/claude_normalization_adapter.py create mode 100644 litellm/proxy/auth/model_policy_read.py create mode 100644 tests/unit/_auth/support/claude_reference.py create mode 100644 tests/unit/_auth/test_claude_normalization.py create mode 100644 tests/unit/proxy/auth/test_model_policy_read.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 81d4454b58a..4a7345f2d84 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -68,13 +68,13 @@ mod _native { request_auth_agent_grants, request_auth_agent_group_load, request_auth_agent_invocation, request_auth_budget_alert, request_auth_budget_alert_plan, request_auth_budget_exemption, request_auth_budget_price, request_auth_budget_route, request_auth_budget_window_start, - request_auth_budgets, request_auth_codex_model, request_auth_codex_normalization, - request_auth_codex_turn, request_auth_common_budget, request_auth_common_budget_scopes, - request_auth_common_flow, request_auth_common_read, request_auth_context, - request_auth_context_end_user, request_auth_context_member, request_auth_context_models, - request_auth_context_percentage, request_auth_context_temporary, - request_auth_context_throttle, request_auth_context_token, request_auth_custom_flow, - request_auth_end_user_apply, request_auth_end_user_assignments, + request_auth_budgets, request_auth_claude_membership, request_auth_claude_normalization, + request_auth_codex_model, request_auth_codex_normalization, request_auth_codex_turn, + request_auth_common_budget, request_auth_common_budget_scopes, request_auth_common_flow, + request_auth_common_read, request_auth_context, request_auth_context_end_user, + request_auth_context_member, request_auth_context_models, request_auth_context_percentage, + request_auth_context_temporary, request_auth_context_throttle, request_auth_context_token, + request_auth_custom_flow, request_auth_end_user_apply, request_auth_end_user_assignments, request_auth_end_user_default_id, request_auth_end_user_lookup, request_auth_end_user_normalize_id, request_auth_end_user_registry, request_auth_end_user_validation, request_auth_expired, request_auth_fallback_access, diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index 3f8378fc66a..1fca317400d 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -291,6 +291,30 @@ pub(crate) fn request_auth_codex_normalization( to_py(py, &decision) } +#[pyfunction] +pub(crate) fn request_auth_claude_membership( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::claude_membership::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::claude_membership::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_claude_normalization( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::claude_normalization::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::claude_normalization::decide(&facts)), + ) +} + #[pyfunction] pub(crate) fn request_auth_codex_turn( py: Python<'_>, diff --git a/litellm-rust/crates/request-auth-core/src/claude_membership.rs b/litellm-rust/crates/request-auth-core/src/claude_membership.rs new file mode 100644 index 00000000000..2840e4aaa9b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/claude_membership.rs @@ -0,0 +1,100 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] +pub struct Facts { + pub name: String, + pub router_present: bool, + pub alias_count: usize, + pub global_alias_exists: Option, + pub alias_entries: Vec>, + pub registry_loaded: bool, + pub registry_alias_exists: Option, + pub deployment: Option, + pub names: Option>, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { present: bool }, + ReadGlobal, + ReadAlias { index: usize }, + ReadRegistry, + ReadAliases, + ReadDeployment, + ReadNames, +} + +pub fn decide(facts: &Facts) -> Decision { + let Some(global) = &facts.global_alias_exists else { + return Decision::ReadGlobal; + }; + if *global { + return Decision::Complete { present: true }; + } + for index in 0..facts.alias_count { + let Some(aliases) = facts.alias_entries.get(index) else { + return Decision::ReadAlias { index }; + }; + if *aliases == Some(true) { + return Decision::Complete { present: true }; + } + } + if !facts.router_present { + return Decision::Complete { present: false }; + } + if !facts.registry_loaded { + return Decision::ReadRegistry; + } + let Some(aliases) = &facts.registry_alias_exists else { + return Decision::ReadAliases; + }; + if *aliases { + return Decision::Complete { present: true }; + } + match facts.deployment { + None => return Decision::ReadDeployment, + Some(true) => return Decision::Complete { present: true }, + Some(false) => (), + } + let Some(names) = &facts.names else { + return Decision::ReadNames; + }; + Decision::Complete { + present: !names.is_empty(), + } +} + +#[cfg(test)] +mod tests { + use super::{Decision, Facts, decide}; + use rstest::rstest; + + #[rstest] + #[case::global(Some(true), Decision::Complete { present: true })] + #[case::missing(None, Decision::ReadGlobal)] + #[case::next_alias(Some(false), Decision::ReadAlias { index: 0 })] + fn alias_short_circuit(#[case] global_alias_exists: Option, #[case] expected: Decision) { + let facts = Facts { + name: "model".into(), + router_present: true, + alias_count: 3, + global_alias_exists, + ..Facts::default() + }; + assert_eq!(decide(&facts), expected); + } + + #[rstest] + #[case::no_router(false, Decision::Complete { present: false })] + #[case::registry(true, Decision::ReadRegistry)] + fn registry_after_all_aliases(#[case] router_present: bool, #[case] expected: Decision) { + let facts = Facts { + name: "model".into(), + router_present, + global_alias_exists: Some(false), + ..Facts::default() + }; + assert_eq!(decide(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/claude_normalization.rs b/litellm-rust/crates/request-auth-core/src/claude_normalization.rs new file mode 100644 index 00000000000..7a853e5ae2d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/claude_normalization.rs @@ -0,0 +1,176 @@ +use serde::{Deserialize, Serialize}; + +use crate::claude_membership; + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct Facts { + pub route: String, + pub request_present: bool, + pub router_present: Option, + pub normalized: Option, + pub model_loaded: bool, + pub requested: Option, + pub body_loaded: bool, + pub body_model: Option, + pub settings_loaded: bool, + pub membership: Vec, + pub marked: bool, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete, + ReadRouter, + ReadNormalized, + ReadModel, + ReadBodyModel, + ReadSettings, + ReadMembership { name: String, index: usize }, + MarkNormalized, + Rewrite { target: String }, +} + +fn unmarked(name: &str) -> &str { + if name.to_ascii_lowercase().ends_with("[1m]") { + &name[..name.len() - 4] + } else { + name + } +} + +fn decoded(name: &str) -> Option { + let encoded = unmarked(name).strip_prefix("claude-router-")?; + if !encoded.len().is_multiple_of(2) + || !encoded + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + { + return None; + } + let bytes: Option> = encoded + .as_bytes() + .as_chunks::<2>() + .0 + .iter() + .map(|pair| u8::from_str_radix(std::str::from_utf8(pair).ok()?, 16).ok()) + .collect(); + String::from_utf8(bytes?) + .ok() + .filter(|value| !value.is_empty()) +} + +fn membership(facts: &Facts, index: usize, name: &str) -> Result { + if let Some(read) = facts.membership.get(index).filter(|read| read.name == name) + && let claude_membership::Decision::Complete { present } = claude_membership::decide(read) + { + return Ok(present); + } + Err(Decision::ReadMembership { + name: name.into(), + index, + }) +} + +fn target(facts: &Facts, name: &str) -> Result, Decision> { + if membership(facts, 0, name)? { + return Ok(None); + } + let plain = unmarked(name); + let marked = plain != name; + if marked && membership(facts, 1, plain)? { + return Ok(Some(plain.into())); + } + let Some(model) = decoded(name) else { + return Ok(None); + }; + Ok(membership(facts, if marked { 2 } else { 1 }, &model)?.then_some(model)) +} + +pub fn decide(facts: &Facts) -> Decision { + if !matches!( + facts.route.as_str(), + "/messages" + | "/v1/messages" + | "/chat/completions" + | "/v1/chat/completions" + | "/responses" + | "/v1/responses" + ) { + return Decision::Complete; + } + match facts.router_present { + None => return Decision::ReadRouter, + Some(false) => return Decision::Complete, + Some(true) => (), + } + if facts.request_present { + match facts.normalized { + None => return Decision::ReadNormalized, + Some(true) => return Decision::Complete, + Some(false) => (), + } + } + if !facts.model_loaded { + return Decision::ReadModel; + } + let Some(requested) = facts.requested.as_deref() else { + return Decision::Complete; + }; + if !facts.body_loaded { + return Decision::ReadBodyModel; + } + if Some(requested) != facts.body_model.as_deref() { + return Decision::Complete; + } + if !requested.starts_with("claude-router-") && !requested.to_ascii_lowercase().ends_with("[1m]") + { + return Decision::Complete; + } + if !facts.settings_loaded { + return Decision::ReadSettings; + } + let result = match target(facts, requested) { + Err(read) => return read, + Ok(value) => value, + }; + if facts.request_present && !facts.marked { + return Decision::MarkNormalized; + } + match result { + Some(target) => Decision::Rewrite { target }, + None => Decision::Complete, + } +} + +#[cfg(test)] +mod tests { + use super::{Decision, Facts, decide, decoded}; + use rstest::rstest; + + #[rstest] + #[case::canonical("claude-router-6d6f64656c", Some("model"))] + #[case::marked("claude-router-6d6f64656c[1M]", Some("model"))] + #[case::uppercase("claude-router-6D6F64656C", None)] + #[case::whitespace("claude-router-6d 6f64656c", None)] + #[case::invalid_utf8("claude-router-ff", None)] + #[case::empty("claude-router-", None)] + #[case::odd("claude-router-1", None)] + #[case::unicode("claude-router-e6a8a1e59e8b", Some("模型"))] + fn canonical_picker_decoding(#[case] name: &str, #[case] expected: Option<&str>) { + assert_eq!(decoded(name).as_deref(), expected); + } + + #[rstest] + #[case::irrelevant("/models", Decision::Complete)] + #[case::eligible("/v1/messages", Decision::ReadRouter)] + fn route_before_reads(#[case] route: &str, #[case] expected: Decision) { + assert_eq!( + decide(&Facts { + route: route.into(), + ..Facts::default() + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/codex_models.rs b/litellm-rust/crates/request-auth-core/src/codex_models.rs index a1e8daed5dd..1dbcbbc902c 100644 --- a/litellm-rust/crates/request-auth-core/src/codex_models.rs +++ b/litellm-rust/crates/request-auth-core/src/codex_models.rs @@ -251,6 +251,85 @@ mod tests { use super::*; use rstest::{fixture, rstest}; + fn helper(model: &str, grants: &[String]) -> model_helpers::Facts { + model_helpers::Facts { + kind: model_helpers::Kind::Helper, + model: model.into(), + grants: Some(grants.to_owned()), + team_id: None, + team_aliases: vec![], + router_present: false, + groups: None, + names: None, + provider_loaded: false, + provider: None, + provider_reads: None, + } + } + + #[rstest] + #[case::all_candidates_both_grants(true, true, true)] + #[case::mixed_candidates(false, true, false)] + #[case::team_denies(true, false, false)] + #[case::mixed_and_team_denies(false, false, false)] + fn prefix_requires_all_chatgpt_deployments_and_both_grants( + base: Facts, + #[case] all_chatgpt: bool, + #[case] team_allows: bool, + #[case] allowed: bool, + ) { + let grants = Grants { + key: vec!["chatgpt/model".into()], + team: vec![ + if team_allows { + "chatgpt/model" + } else { + "other" + } + .into(), + ], + }; + let first = Candidate { + provider_loaded: true, + provider: Some("chatgpt".into()), + model: Some("chatgpt/model".into()), + name: Some("chatgpt/model".into()), + public_loaded: false, + public_name: None, + }; + let second = Candidate { + provider: Some(if all_chatgpt { "chatgpt" } else { "openai" }.into()), + model: Some( + if all_chatgpt { + "chatgpt/model" + } else { + "openai/model" + } + .into(), + ), + ..first.clone() + }; + let facts = Facts { + kind: Kind::Chatgpt, + candidates: Some(vec![first, second]), + bare: vec![helper("model", &grants.key)], + prefixed: vec![ + helper("chatgpt/model", &grants.key), + helper("chatgpt/model", &grants.team), + ], + grants: Some(grants.clone()), + prefixed_grants: Some(grants), + ..base + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::Complete { + allowed, + target: allowed.then(|| "chatgpt/model".into()) + } + ); + } + #[fixture] fn base() -> Facts { Facts { diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index 6363f426296..47659f3ca91 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -10,6 +10,8 @@ pub mod budget_alerts; pub mod budget_context; pub mod budget_lookup; pub mod budget_pricing; +pub mod claude_membership; +pub mod claude_normalization; pub mod codex_models; pub mod codex_normalization; pub mod codex_turn; diff --git a/litellm/_auth/claude_normalization.py b/litellm/_auth/claude_normalization.py new file mode 100644 index 00000000000..f3fd934df04 --- /dev/null +++ b/litellm/_auth/claude_normalization.py @@ -0,0 +1,74 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ClaudeMembershipFacts: + name: str + router_present: bool + alias_count: int + global_alias_exists: bool | None = None + alias_entries: tuple[bool | None, ...] = () + registry_loaded: bool = False + registry_alias_exists: bool | None = None + deployment: bool | None = None + names: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ClaudeMembershipDecision: + kind: Literal["Complete", "ReadGlobal", "ReadAlias", "ReadRegistry", "ReadAliases", "ReadDeployment", "ReadNames"] + present: bool = False + index: int = 0 + + +@dataclass(frozen=True, slots=True) +class ClaudeNormalizationFacts: + route: str + request_present: bool + router_present: bool | None = None + normalized: bool | None = None + model_loaded: bool = False + requested: str | None = None + body_loaded: bool = False + body_model: str | None = None + settings_loaded: bool = False + membership: tuple[ClaudeMembershipFacts, ...] = () + marked: bool = False + + +@dataclass(frozen=True, slots=True) +class ClaudeNormalizationDecision: + kind: Literal[ + "Complete", + "ReadRouter", + "ReadNormalized", + "ReadModel", + "ReadBodyModel", + "ReadSettings", + "ReadMembership", + "MarkNormalized", + "Rewrite", + ] + name: str | None = None + index: int = 0 + target: str | None = None + + +_MEMBERSHIP: Final = TypeAdapter(ClaudeMembershipDecision) +_NORMALIZATION: Final = TypeAdapter(ClaudeNormalizationDecision) + + +@native_boundary +def claude_membership(facts: ClaudeMembershipFacts) -> ClaudeMembershipDecision: + return _MEMBERSHIP.validate_python(_native.request_auth_claude_membership(asdict(facts))) + + +@native_boundary +def claude_normalization(facts: ClaudeNormalizationFacts) -> ClaudeNormalizationDecision: + return _NORMALIZATION.validate_python(_native.request_auth_claude_normalization(asdict(facts))) diff --git a/litellm/proxy/auth/claude_normalization_adapter.py b/litellm/proxy/auth/claude_normalization_adapter.py new file mode 100644 index 00000000000..25275edfa22 --- /dev/null +++ b/litellm/proxy/auth/claude_normalization_adapter.py @@ -0,0 +1,174 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from itertools import accumulate, repeat +from typing import ( + Final, + assert_never, + cast, # noqa: TID251 # checked Mapping view avoids eager reads of unselected values +) + +from starlette.requests import Request + +import litellm +from litellm._auth.claude_normalization import ( + ClaudeMembershipFacts, + ClaudeNormalizationDecision, + ClaudeNormalizationFacts, + claude_membership, + claude_normalization, +) +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.model_policy_read import require_native_fact +from litellm.proxy.common_utils.http_parsing_utils import rewrite_request_model +from litellm.router import Router +from litellm.router_api import RegistryReader, RouterRegistry + +_NORMALIZED: Final = "litellm.claude_model_normalized" + + +def _alias_entry_exists(value: object, name: str) -> bool | None: + if not isinstance(value, Mapping): + return None + return name in value + + +def _settings_aliases(value: object) -> object: + if not isinstance(value, Mapping): + return None + entries: Final = cast(Mapping[object, object], value) # cast-ok: Mapping checked without coercing unrelated values + return entries.get("model_group_alias") + + +@dataclass(frozen=True, slots=True) +class _MembershipRead: + facts: ClaudeMembershipFacts + registry: RouterRegistry | None = None + + +def read_claude_membership_facts( + name: str, router: RegistryReader | None, team_id: str | None, alias_maps: tuple[object, ...] +) -> ClaudeMembershipFacts: + initial: Final = _MembershipRead(ClaudeMembershipFacts(name, router is not None, len(alias_maps))) + + def advance(state: _MembershipRead, _unused: None) -> _MembershipRead: + facts: Final = state.facts + decision: Final = claude_membership(facts) + match decision.kind: + case "ReadGlobal": + return replace(state, facts=replace(facts, global_alias_exists=name in litellm.model_alias_map)) + case "ReadAlias": + alias_exists: Final = _alias_entry_exists(alias_maps[decision.index], name) + return replace(state, facts=replace(facts, alias_entries=(*facts.alias_entries, alias_exists))) + case "ReadRegistry": + return _MembershipRead(replace(facts, registry_loaded=True), require_native_fact(router).registry()) + case "ReadAliases": + alias_in_registry: Final = name in require_native_fact(state.registry).aliases() + return replace(state, facts=replace(facts, registry_alias_exists=alias_in_registry)) + case "ReadDeployment": + return replace( + state, facts=replace(facts, deployment=require_native_fact(state.registry).has_deployment(name)) + ) + case "ReadNames": + names: Final = tuple( + require_native_fact(state.registry).names(name, kind="route_candidates", team_id=team_id) + ) + return replace(state, facts=replace(facts, names=names)) + case "Complete": + raise NativeAuthRuntimeError("completed Claude membership requested IO") + case _: + assert_never(decision.kind) + + return next( + state.facts + for state in accumulate(repeat(None), advance, initial=initial) + if claude_membership(state.facts).kind == "Complete" + ) + + +@dataclass(frozen=True, slots=True) +class _Source: + data: dict[str, object] # mutable-ok: request-local rewrite effect + token: UserAPIKeyAuth + request: Request | None + router: Router | None = None + alias_maps: tuple[object, ...] = () + + +def _membership( + facts: ClaudeNormalizationFacts, source: _Source, decision: ClaudeNormalizationDecision +) -> ClaudeMembershipFacts: + if decision.index != len(facts.membership): + raise NativeAuthRuntimeError("Claude membership read out of order") + return read_claude_membership_facts( + require_native_fact(decision.name), source.router, source.token.team_id, source.alias_maps + ) + + +async def _normalize(facts: ClaudeNormalizationFacts, source: _Source) -> None: + decision: Final = claude_normalization(facts) + match decision.kind: + case "Complete": + return + case "ReadRouter": + from litellm.proxy.proxy_server import llm_router + + return await _normalize( + replace(facts, router_present=llm_router is not None), replace(source, router=llm_router) + ) + case "ReadNormalized": + return await _normalize( + replace(facts, normalized=require_native_fact(source.request).scope.get(_NORMALIZED) is True), source + ) + case "ReadModel": + from litellm.proxy.auth.user_api_key_auth import ( + _get_model_from_request_context, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared raw request extractor has an unparameterized legacy dict argument + ) + + requested: Final = _get_model_from_request_context( + source.data, facts.route, source.request, source.router, source.token.team_id + ) + return await _normalize( + replace( + facts, + model_loaded=True, + requested=requested if isinstance(requested, str) else None, + ), + source, + ) + case "ReadBodyModel": + body: Final = source.data.get("model") + return await _normalize( + replace(facts, body_loaded=True, body_model=body if isinstance(body, str) else None), source + ) + case "ReadSettings": + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=source.token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + aliases: Final = _settings_aliases(settings) + return await _normalize( + replace(facts, settings_loaded=True), + replace(source, alias_maps=(source.token.aliases, source.token.team_model_aliases, aliases)), # pyright: ignore[reportUnknownMemberType] # aliases is an unparameterized legacy mapping, queried only at native-selected names + ) + case "ReadMembership": + membership: Final = _membership(facts, source, decision) + return await _normalize(replace(facts, membership=(*facts.membership, membership)), source) + case "MarkNormalized": + request: Final = require_native_fact(source.request) + request.scope[_NORMALIZED] = True + return await _normalize(replace(facts, marked=True), source) + case "Rewrite": + rewrite_request_model(source.data, source.request, require_native_fact(decision.target)) + case _: + assert_never(decision.kind) + + +async def normalize_claude_model( + request_data: dict[str, object], # mutable-ok: native-selected request-local rewrite + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + await _normalize(ClaudeNormalizationFacts(route, request is not None), _Source(request_data, valid_token, request)) diff --git a/litellm/proxy/auth/codex_model_adapter.py b/litellm/proxy/auth/codex_model_adapter.py index 8f3ea7f5207..25aa3f5e8ea 100644 --- a/litellm/proxy/auth/codex_model_adapter.py +++ b/litellm/proxy/auth/codex_model_adapter.py @@ -17,6 +17,7 @@ from litellm._auth.failure import NativeAuthRuntimeError, native_boundary from litellm._auth.model_helpers import ModelHelperFacts from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.model_helper_adapter import read_model_helper_facts +from litellm.proxy.auth.model_policy_read import require_native_fact from litellm.proxy.common_utils.codex_model_catalog import bundled_codex_models from litellm.router_api import RegistryReader from litellm.types.router import Deployment @@ -73,48 +74,45 @@ def read_codex_model_facts( case "ReadBundled": return replace(state, facts=replace(facts, bundled=tuple(bundled_codex_models()))) case "ReadCost": - if decision.model is None: - raise NativeAuthRuntimeError("Codex price lookup without model") - return replace(state, facts=replace(facts, priced=decision.model in litellm.model_cost)) # pyright: ignore[reportUnknownMemberType] # dynamically loaded registry, only raw key presence is read + priced: Final = require_native_fact(decision.model) in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # raw key presence in the dynamically loaded cost registry + return replace(state, facts=replace(facts, priced=priced)) case "ReadCandidates": - if router is None or decision.model is None: - raise NativeAuthRuntimeError("Codex deployment lookup without router or model") - registry: Final = router.registry() + registry: Final = require_native_fact(router).registry() loaded: Final = tuple( registry.deployment(model_id) - for model_id in registry.names(decision.model, kind="route_candidates", team_id=team_id) + for model_id in registry.names( + require_native_fact(decision.model), kind="route_candidates", team_id=team_id + ) ) deployments: Final = tuple(deployment for deployment in loaded if deployment is not None) return _State(replace(facts, candidates=tuple(CodexCandidate() for _ in deployments)), deployments) case "ReadAlias": - if router is None or decision.model is None: - raise NativeAuthRuntimeError("Codex alias lookup without router or model") return replace( state, facts=replace( - facts, alias_loaded=True, alias=router.registry().resolve(decision.model, kind="alias") + facts, + alias_loaded=True, + alias=require_native_fact(router) + .registry() + .resolve(require_native_fact(decision.model), kind="alias"), ), ) case "ReadCandidate": - if facts.candidates is None or decision.field is None: - raise NativeAuthRuntimeError("Codex candidate projection without deployment") + candidates: Final = require_native_fact(facts.candidates) updated: Final = _candidate( - decision.field, facts.candidates[decision.index], state.deployments[decision.index] + require_native_fact(decision.field), candidates[decision.index], state.deployments[decision.index] ) return replace( state, facts=replace( facts, candidates=tuple( - updated if index == decision.index else value - for index, value in enumerate(facts.candidates) + updated if index == decision.index else value for index, value in enumerate(candidates) ), ), ) case "ReadGrants": - if token is None: - raise NativeAuthRuntimeError("Codex model grants without token") - grant_read: Final = _grants(token) + grant_read: Final = _grants(require_native_fact(token)) return replace( state, facts=( @@ -124,12 +122,16 @@ def read_codex_model_facts( ), ) case "ReadHelper": - selected: Final = facts.prefixed_grants if decision.prefixed else facts.grants - if selected is None or decision.model is None: - raise NativeAuthRuntimeError("Codex helper lookup without model grants") + selected: Final = require_native_fact(facts.prefixed_grants if decision.prefixed else facts.grants) grants: Final = (selected.key, selected.team)[decision.index] scope: Final = read_model_helper_facts( - ModelHelperFacts("Helper", decision.model, grants, team_id=team_id, router_present=bool(router)), + ModelHelperFacts( + "Helper", + require_native_fact(decision.model), + grants, + team_id=team_id, + router_present=bool(router), + ), router, ) updated_facts: Final = ( diff --git a/litellm/proxy/auth/codex_normalization_adapter.py b/litellm/proxy/auth/codex_normalization_adapter.py index 12fe38af6d5..e207b512c54 100644 --- a/litellm/proxy/auth/codex_normalization_adapter.py +++ b/litellm/proxy/auth/codex_normalization_adapter.py @@ -1,16 +1,20 @@ from collections.abc import Mapping from dataclasses import dataclass, replace -from typing import Final, cast +from typing import ( + Final, + cast, # noqa: TID251 # checked Mapping view preserves raw lookup timing without eager validation +) from fastapi import Request -from litellm._auth.codex_normalization import CodexNormalizationFacts, codex_normalization +from litellm._auth.codex_normalization import CodexNormalizationDecision, CodexNormalizationFacts, codex_normalization from litellm._auth.failure import NativeAuthRuntimeError from litellm._logging import verbose_proxy_logger from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY, MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.codex_model_adapter import read_codex_model_facts from litellm.proxy.auth.codex_turn_adapter import read_native_turn_facts +from litellm.proxy.auth.model_policy_read import require_native_fact from litellm.proxy.auth.route_policy_adapter import read_eu_jev_configuration, route_context from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body, rewrite_request_model from litellm.router_api import RegistryReader @@ -27,9 +31,8 @@ class _Source: def _value(value: object, key: str) -> object: if not isinstance(value, Mapping): return None - return cast(Mapping[object, object], value).get( - key - ) # cast-ok: Mapping shape was checked without coercing unrelated keys + entries: Final = cast(Mapping[object, object], value) # cast-ok: runtime Mapping shape was checked + return entries.get(key) def _rewrite(facts: CodexNormalizationFacts, source: _Source, target: str | None) -> None: @@ -45,18 +48,10 @@ def _rewrite(facts: CodexNormalizationFacts, source: _Source, target: str | None rewrite_request_model(source.data, source.request, target) -async def _alias(facts: CodexNormalizationFacts, source: _Source) -> None: - from litellm.proxy.auth.auth_utils import ( - request_dispatched_to_pass_through_endpoint, - request_dispatched_to_provider_pass_through, - ) - from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj +async def _alias_scope(facts: CodexNormalizationFacts, source: _Source, decision: CodexNormalizationDecision) -> None: + from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_endpoint - decision: Final = codex_normalization(facts) - if decision.kind == "Complete": - return - if source.request is None: - raise NativeAuthRuntimeError("alias normalization read without request") + request: Final = require_native_fact(source.request) match decision.kind: case "ReadContext": return await _alias(replace(facts, context=route_context(facts.route)), source) @@ -64,10 +59,10 @@ async def _alias(facts: CodexNormalizationFacts, source: _Source) -> None: return await _alias(replace(facts, eu_jev=read_eu_jev_configuration()), source) case "ReadResolved": return await _alias( - replace(facts, resolved=source.request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is True), source + replace(facts, resolved=request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is True), source ) case "MarkResolved": - source.request.scope[MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY] = True + request.scope[MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY] = True return await _alias(replace(facts, marked=True), source) case "ReadModel": model: Final = source.data.get("model") @@ -75,41 +70,49 @@ async def _alias(facts: CodexNormalizationFacts, source: _Source) -> None: replace(facts, model_loaded=True, model=model if isinstance(model, str) else None), source ) case "ReadExplicit": - return await _alias( - replace(facts, explicit=request_dispatched_to_pass_through_endpoint(source.request)), source - ) + return await _alias(replace(facts, explicit=request_dispatched_to_pass_through_endpoint(request)), source) + case _: + raise NativeAuthRuntimeError("invalid alias scope read stage") + + +async def _alias(facts: CodexNormalizationFacts, source: _Source) -> None: + from litellm.proxy.auth.auth_utils import request_dispatched_to_provider_pass_through + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + decision: Final = codex_normalization(facts) + if decision.kind == "Complete": + return + request: Final = require_native_fact(source.request) + if decision.kind in ("ReadContext", "ReadEuJev", "ReadResolved", "MarkResolved", "ReadModel", "ReadExplicit"): + return await _alias_scope(facts, source, decision) + match decision.kind: case "ReadProvider": - return await _alias( - replace(facts, provider=request_dispatched_to_provider_pass_through(source.request)), source - ) + return await _alias(replace(facts, provider=request_dispatched_to_provider_pass_through(request)), source) case "ReadRouting": routing: Final = read_codex_model_facts(facts.model, source.router, source.token.team_id) return await _alias(replace(facts, routing=routing), source) case "ReadBody": return await _alias( - replace(facts, body_present=await read_raw_json_body(request=source.request) is not None), source + replace(facts, body_present=await read_raw_json_body(request=request) is not None), source ) case "ReadSettings": settings: Final = await proxy_config.get_hierarchical_router_settings( user_api_key_dict=source.token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) - if facts.model is None: - raise NativeAuthRuntimeError("alias settings read without requested model") - entry: Final = _value(_value(settings, "model_group_alias"), facts.model) - target: Final = ( - _value(cast(object, entry), "model") if isinstance(entry, Mapping) else entry - ) # cast-ok: retain the object boundary after Mapping narrowing + entry: Final = _value(_value(settings, "model_group_alias"), require_native_fact(facts.model)) + mapped_entry: Final[object] = entry + target: Final = _value(mapped_entry, "model") if isinstance(entry, Mapping) else entry return await _alias( replace(facts, settings_loaded=True, target=target if isinstance(target, str) else None), source ) case "ReadProviderAfter": return await _alias( - replace(facts, provider_after=request_dispatched_to_provider_pass_through(source.request)), source + replace(facts, provider_after=request_dispatched_to_provider_pass_through(request)), source ) case "ReadTurn": - if facts.model is None or source.router is None: - raise NativeAuthRuntimeError("alias turn read without router model") - turn: Final = await read_native_turn_facts(facts.model, source.router, source.request, source.token) + turn: Final = await read_native_turn_facts( + require_native_fact(facts.model), require_native_fact(source.router), request, source.token + ) return await _alias(replace(facts, turn=turn), source) case "Rewrite": _rewrite(facts, source, decision.target) diff --git a/litellm/proxy/auth/codex_turn_adapter.py b/litellm/proxy/auth/codex_turn_adapter.py index c73d1d7c694..b6121ccddc3 100644 --- a/litellm/proxy/auth/codex_turn_adapter.py +++ b/litellm/proxy/auth/codex_turn_adapter.py @@ -21,6 +21,7 @@ from litellm._auth.codex_turn import ( from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.codex_model_adapter import read_codex_model_facts +from litellm.proxy.auth.model_policy_read import require_native_fact from litellm.proxy.auth.route_enable_adapter import read_route_enable_facts from litellm.proxy.auth.route_policy_adapter import read_virtual_route_facts from litellm.proxy.auth.route_scope_adapter import read_route_scope_facts @@ -32,9 +33,10 @@ if TYPE_CHECKING: ) RESPONSES_PATH: Final = "/v1/responses" +_KEY_FACTS: Final = CodexKeyFacts() -def read_codex_key_facts(token: UserAPIKeyAuth, facts: CodexKeyFacts = CodexKeyFacts()) -> CodexKeyFacts: +def read_codex_key_facts(token: UserAPIKeyAuth, facts: CodexKeyFacts = _KEY_FACTS) -> CodexKeyFacts: decision: Final = codex_key(facts) match decision.kind: case "Complete": @@ -176,9 +178,15 @@ async def _read(facts: CodexTurnFacts, source: _Source) -> CodexTurnFacts: routing: Final = read_codex_model_facts(facts.model, source.router, source.token.team_id) return await _read(replace(facts, routing=routing), source) case "ReadJwt": - if source.access is None or facts.model is None: - raise NativeAuthRuntimeError("Codex JWT read without access context") - return await _read(replace(facts, access=await read_codex_access_facts(source.access, facts.model)), source) + return await _read( + replace( + facts, + access=await read_codex_access_facts( + require_native_fact(source.access), require_native_fact(facts.model) + ), + ), + source, + ) case "ReadCatalog": if source.headers is None or source.caller is None or source.cache is None: raise NativeAuthRuntimeError("Codex catalog read without IO context") diff --git a/litellm/proxy/auth/key_model_access_adapter.py b/litellm/proxy/auth/key_model_access_adapter.py index fdb32b58a0a..be311767faf 100644 --- a/litellm/proxy/auth/key_model_access_adapter.py +++ b/litellm/proxy/auth/key_model_access_adapter.py @@ -41,11 +41,9 @@ async def read_key_and_fallback_model_access_facts( llm_model_list: Sequence[object] | None, llm_router: Router | None, ) -> KeyAndFallbackModelAccessFacts: - from litellm.proxy.auth.user_api_key_auth import ( - _normalize_claude_model, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # legacy helper takes the mutable request body - ) + from litellm.proxy.auth.claude_normalization_adapter import normalize_claude_model - await _normalize_claude_model(request_data, valid_token, request, route) + await normalize_claude_model(request_data, valid_token, request, route) await resolve_router_settings_model_group_alias(request_data, valid_token, request, route) resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) facts: Final = KeyAndFallbackModelAccessFacts( diff --git a/litellm/proxy/auth/model_policy_read.py b/litellm/proxy/auth/model_policy_read.py new file mode 100644 index 00000000000..77e0a6c13db --- /dev/null +++ b/litellm/proxy/auth/model_policy_read.py @@ -0,0 +1,11 @@ +from typing import TypeVar + +from litellm._auth.failure import NativeAuthRuntimeError + +_T = TypeVar("_T") # rebind-ok: TypeVar must remain a bare static type-variable binding + + +def require_native_fact(value: _T | None) -> _T: + if value is None: + raise NativeAuthRuntimeError("missing fact for native model authorization read") + return value diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8f5c3fcb809..659278024af 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -72,6 +72,8 @@ def request_auth_team_override(facts: object) -> object: ... def request_auth_model_helpers(facts: object) -> object: ... def request_auth_codex_model(facts: object) -> object: ... def request_auth_codex_normalization(facts: object) -> object: ... +def request_auth_claude_membership(facts: object) -> object: ... +def request_auth_claude_normalization(facts: object) -> object: ... def request_auth_codex_turn(facts: object) -> object: ... def request_auth_route_scope(facts: object) -> object: ... def request_auth_personal_owner(facts: object) -> object: ... diff --git a/tests/unit/_auth/support/claude_reference.py b/tests/unit/_auth/support/claude_reference.py new file mode 100644 index 00000000000..a72792a4abb --- /dev/null +++ b/tests/unit/_auth/support/claude_reference.py @@ -0,0 +1,103 @@ +"""Request-auth Claude policy frozen from main 2f430873aead937b9112e683703956eb1ab179af.""" + +from collections.abc import Container, Mapping +from dataclasses import dataclass +from typing import Final + +from starlette.requests import Request + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.http_parsing_utils import rewrite_request_model +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class RoutingNames: + router: RegistryReader | None + team_id: str | None + alias_maps: tuple[object, ...] + + def __contains__(self, name: object) -> bool: + if not isinstance(name, str): + return False + if name in litellm.model_alias_map or any( + isinstance(aliases, Mapping) and name in aliases for aliases in self.alias_maps + ): + return True + if self.router is None: + return False + registry: Final = self.router.registry() + return ( + name in registry.aliases() + or registry.has_deployment(name) + or bool(registry.names(name, kind="route_candidates", team_id=self.team_id)) + ) + + +def unmarked(name: str) -> str: + return name[:-4] if name.lower().endswith("[1m]") else name + + +def decoded(view_id: str) -> str | None: + encoded: Final = unmarked(view_id).removeprefix("claude-router-") + if encoded == unmarked(view_id): + return None + try: + model_id: Final = bytes.fromhex(encoded).decode() + except (ValueError, UnicodeDecodeError): + return None + return model_id if f"claude-router-{model_id.encode().hex()}" == unmarked(view_id) else None + + +def group_name(view_id: str, routing_names: Container[str]) -> str | None: + if view_id in routing_names: + return None + plain: Final = unmarked(view_id) + if plain != view_id and plain in routing_names: + return plain + model_id: Final = decoded(view_id) + return model_id if model_id and model_id in routing_names else None + + +async def normalize( + data: dict[str, object], # mutable-ok: native and reference both apply the request-local rewrite + token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + from litellm.proxy.auth.user_api_key_auth import ( + _get_model_from_request_context, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # raw request extraction is not the frozen policy + ) + from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj + + if ( + route + not in frozenset( # comprehension-ok: frozen main route set construction + f"/{prefix}{endpoint}" + for prefix in ("", "v1/") + for endpoint in ("messages", "chat/completions", "responses") + ) + or llm_router is None + ): + return + if request is not None and request.scope.get("litellm.claude_model_normalized") is True: + return + requested: Final = _get_model_from_request_context(data, route, request, llm_router, token.team_id) + if not isinstance(requested, str) or requested != data.get("model"): + return + if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"): + return + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + aliases: Final[object] = settings.get("model_group_alias") if isinstance(settings, Mapping) else None + source: Final = group_name( + requested, + RoutingNames(llm_router, token.team_id, (token.aliases, token.team_model_aliases, aliases)), # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # frozen raw alias maps retain legacy unparameterized mappings + ) + if request is not None: + request.scope["litellm.claude_model_normalized"] = True # rebind-ok: frozen request-local effect + if source is None: + return + rewrite_request_model(data, request, source) diff --git a/tests/unit/_auth/test_claude_normalization.py b/tests/unit/_auth/test_claude_normalization.py new file mode 100644 index 00000000000..57ca8e71ce3 --- /dev/null +++ b/tests/unit/_auth/test_claude_normalization.py @@ -0,0 +1,264 @@ +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, field +from typing import Final, cast + +import pytest +from starlette.requests import Request + +import litellm +from litellm._auth.claude_normalization import ClaudeNormalizationFacts, claude_membership, claude_normalization +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.claude_normalization_adapter import normalize_claude_model, read_claude_membership_facts +from litellm.router_api import RouterRegistry +from litellm.rust_bridge import _native +from tests.unit._auth.support import claude_reference as reference + + +@dataclass(frozen=True, slots=True) +class Trace: + fail_at: str | None = None + calls: list[tuple[object, ...]] = field(default_factory=list) # mutable-ok: ordered raw IO observations + + def read(self, operation: str, *values: object) -> None: + self.calls.append((operation, *values)) + if operation == self.fail_at: + raise ValueError(operation) + + +@dataclass(frozen=True, slots=True) +class Registry: + trace: Trace + alias_names: tuple[str, ...] = () + deployments: tuple[str, ...] = () + candidates: tuple[str, ...] = () + + def aliases(self) -> Mapping[str, object]: + self.trace.read("aliases") + return dict.fromkeys(self.alias_names) + + def has_deployment(self, name: str) -> bool: + self.trace.read("deployment", name) + return name in self.deployments + + def names(self, name: str, *, kind: str, team_id: str | None) -> tuple[str, ...]: + self.trace.read("names", name, kind, team_id) + return ("deployment-id",) if name in self.candidates else () + + +@dataclass(frozen=True, slots=True) +class Router: + reader: Registry + + def registry(self) -> RouterRegistry: + self.reader.trace.read("registry") + return cast(RouterRegistry, self.reader) # cast-ok: raw registry fixture supplies the selected IO + + +@dataclass(frozen=True, slots=True) +class Settings: + trace: Trace + aliases: object = None + + async def get_hierarchical_router_settings( + self, *, user_api_key_dict: UserAPIKeyAuth, prisma_client: object, proxy_logging_obj: object + ) -> Mapping[str, object]: + self.trace.read("settings") + return {"model_group_alias": self.aliases} + + +@dataclass(frozen=True, slots=True) +class LookupOnly(Mapping[str, object]): + entries: Mapping[str, object] + trace: Trace + + def __getitem__(self, key: str) -> object: + self.trace.read("lookup", key) + return self.entries[key] + + def __len__(self) -> int: + return len(self.entries) + + def __iter__(self) -> Iterator[str]: + raise RuntimeError("unrequested whole-map read") + + +@pytest.mark.parametrize("present", (False, True)) +def test_membership_only_reads_the_native_selected_alias_name(monkeypatch: pytest.MonkeyPatch, present: bool) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {}) + trace: Final = Trace() + aliases: Final = LookupOnly({"picker": None} if present else {}, trace) + facts: Final = read_claude_membership_facts("picker", None, None, (aliases,)) + assert claude_membership(facts).present is present + assert trace.calls == [("lookup", "picker")] + + +_PICKER: Final = "claude-router-" + "model".encode().hex() +_MODEL_NAMES: Final = ( + "model", + "model[1m]", + "model[1M]", + _PICKER, + _PICKER + "[1m]", + "claude-router-6D6F64656C", + "claude-router-6d 6f64656c", + "claude-router-ff", + "claude-router-1", + "claude-router-", + "claude-router-zz[1m]", + "claude-router-" + "模型".encode().hex(), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", _MODEL_NAMES) +@pytest.mark.parametrize("source", ("global", "key", "team", "settings", "registry", "deployment", "candidates")) +@pytest.mark.parametrize("claim", ("requested", "target", "missing")) +async def test_normalization_matches_frozen_main_rewrite_scope_and_read_order( + monkeypatch: pytest.MonkeyPatch, model: str, source: str, claim: str +) -> None: + name: Final = model if claim == "requested" else (reference.decoded(model) or reference.unmarked(model)) + names: Final = () if claim == "missing" else (name,) + monkeypatch.setattr(litellm, "model_alias_map", dict.fromkeys(names, "target") if source == "global" else {}) + token: Final = UserAPIKeyAuth( + team_id="team", + aliases=dict.fromkeys(names, "target") if source == "key" else {}, + team_model_aliases=dict.fromkeys(names, "target") if source == "team" else {}, + ) + + async def run(native: bool) -> tuple[Mapping[str, object], object, tuple[tuple[object, ...], ...]]: + trace: Final = Trace() + router: Final = Router( + Registry( + trace, + names if source == "registry" else (), + names if source == "deployment" else (), + names if source == "candidates" else (), + ) + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr( + proxy_server, "proxy_config", Settings(trace, dict.fromkeys(names) if source == "settings" else None) + ) + request: Final = Request({"type": "http", "path": "/v1/messages", "method": "POST", "headers": []}) + data: Final[dict[str, object]] = {"model": model, "messages": ()} # mutable-ok: real request rewrite target + await (normalize_claude_model if native else reference.normalize)(data, token, request, "/v1/messages") + return data, request.scope.get("litellm.claude_model_normalized"), tuple(trace.calls) + + assert await run(True) == await run(False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + ( + "/messages", + "/v1/messages", + "/chat/completions", + "/v1/chat/completions", + "/responses", + "/v1/responses", + "/models", + "/backend-api/codex/responses", + ), +) +@pytest.mark.parametrize("router_present", (False, True)) +@pytest.mark.parametrize("normalized", (False, True, 1)) +@pytest.mark.parametrize("request_present", (False, True)) +async def test_normalization_early_exits_and_idempotency_match_main( + monkeypatch: pytest.MonkeyPatch, route: str, router_present: bool, normalized: object, request_present: bool +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {}) + + async def run(native: bool) -> tuple[Mapping[str, object], object, tuple[tuple[object, ...], ...]]: + trace: Final = Trace() + monkeypatch.setattr( + proxy_server, "llm_router", Router(Registry(trace, deployments=("model",))) if router_present else None + ) + monkeypatch.setattr(proxy_server, "proxy_config", Settings(trace)) + request: Final = ( + Request( + { + "type": "http", + "path": route, + "method": "POST", + "headers": [], + "litellm.claude_model_normalized": normalized, + } + ) + if request_present + else None + ) + data: Final[dict[str, object]] = {"model": _PICKER} # mutable-ok: real request rewrite target + await (normalize_claude_model if native else reference.normalize)(data, UserAPIKeyAuth(), request, route) + return ( + data, + request.scope.get("litellm.claude_model_normalized") if request is not None else None, + tuple(trace.calls), + ) + + assert await run(True) == await run(False) + + +@pytest.mark.parametrize("aliases", (None, [], {"picker": None}, {"other": "picker"}, {"picker": 42})) +@pytest.mark.parametrize("position", (0, 1, 2)) +def test_raw_alias_membership_preserves_keys_not_targets( + monkeypatch: pytest.MonkeyPatch, aliases: object, position: int +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {}) + maps: Final = tuple(aliases if index == position else None for index in range(3)) + first: Final = Trace() + second: Final = Trace() + facts: Final = read_claude_membership_facts("picker", Router(Registry(first)), "team", maps) + expected: Final = "picker" in reference.RoutingNames(Router(Registry(second)), "team", maps) + assert claude_membership(facts).present == expected + assert first.calls == second.calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_at", ("settings", "registry", "aliases", "deployment", "names")) +async def test_raw_read_errors_preserve_order_and_do_not_mark_request( + monkeypatch: pytest.MonkeyPatch, fail_at: str +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {}) + + async def run(native: bool) -> tuple[tuple[object, ...], ...]: + trace: Final = Trace(fail_at) + monkeypatch.setattr(proxy_server, "llm_router", Router(Registry(trace))) + monkeypatch.setattr(proxy_server, "proxy_config", Settings(trace)) + request: Final = Request({"type": "http", "path": "/messages", "method": "POST", "headers": []}) + data: Final[dict[str, object]] = {"model": _PICKER} # mutable-ok: real request rewrite target + with pytest.raises(ValueError, match=fail_at): + await (normalize_claude_model if native else reference.normalize)( + data, UserAPIKeyAuth(), request, "/messages" + ) + assert data == {"model": _PICKER} + assert "litellm.claude_model_normalized" not in request.scope + return tuple(trace.calls) + + assert await run(True) == await run(False) + + +def test_non_string_request_never_reads_body_model_or_settings() -> None: + facts: Final = ClaudeNormalizationFacts("/messages", False, router_present=True, model_loaded=True) + assert claude_normalization(facts).kind == "Complete" + + +@pytest.mark.asyncio +async def test_native_membership_failure_is_terminal_before_registry_io(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = NativeAuthRuntimeError("native membership failed") + trace: Final = Trace() + monkeypatch.setattr(proxy_server, "llm_router", Router(Registry(trace))) + monkeypatch.setattr(proxy_server, "proxy_config", Settings(trace)) + + def fail(_facts: object) -> object: + raise failure + + monkeypatch.setattr(_native, "request_auth_claude_membership", fail) + data: Final[dict[str, object]] = {"model": _PICKER} # mutable-ok: real request rewrite target + with pytest.raises(NativeAuthRuntimeError) as raised: + await normalize_claude_model(data, UserAPIKeyAuth(), None, "/messages") + assert raised.value is failure + assert data == {"model": _PICKER} + assert trace.calls == [("settings",)] diff --git a/tests/unit/_auth/test_codex_models.py b/tests/unit/_auth/test_codex_models.py index 2bec858e576..3dba18e0894 100644 --- a/tests/unit/_auth/test_codex_models.py +++ b/tests/unit/_auth/test_codex_models.py @@ -19,9 +19,7 @@ class Registry: deployments: tuple[Deployment | None, ...] alias: str | None = None groups: Mapping[str, tuple[str, ...]] = field(default_factory=dict[str, tuple[str, ...]]) - calls: list[tuple[object, ...]] = field( - default_factory=list[tuple[object, ...]] - ) # mutable-ok: ordered raw registry IO trace + calls: list[tuple[object, ...]] = field(default_factory=list) # mutable-ok: ordered raw registry IO trace def names(self, model: str | None = None, *, kind: str = "all", team_id: str | None = None) -> tuple[str, ...]: self.calls.append(("names", model, kind, team_id)) @@ -50,9 +48,7 @@ class Router: def registry(self) -> RouterRegistry: self.reader.calls.append(("registry",)) - return cast( - RouterRegistry, self.reader - ) # cast-ok: narrow raw registry double implements only the exercised IO methods + return cast(RouterRegistry, self.reader) # cast-ok: raw registry double supplies the selected IO methods def deployment( diff --git a/tests/unit/_auth/test_codex_normalization.py b/tests/unit/_auth/test_codex_normalization.py index 47ef7b1a07d..3108271f83b 100644 --- a/tests/unit/_auth/test_codex_normalization.py +++ b/tests/unit/_auth/test_codex_normalization.py @@ -156,10 +156,10 @@ def test_prefix_adapter_preserves_client_model_and_rewrites_request( request: Final = Request({"type": "http", "path": _NATIVE_ROUTE, "headers": [], "method": "POST"}) if requested_scope is not None: request.scope[CLIENT_REQUESTED_MODEL_SCOPE_KEY] = requested_scope - data: Final[dict[str, object]] = { + data: Final[dict[str, object]] = { # mutable-ok: production mutates request-local body "model": "model", "input": "hello", - } # mutable-ok: production mutates request-local body + } resolve_codex_native_chatgpt_model(data, token, request, _NATIVE_ROUTE) assert data == {"model": "chatgpt/model", "input": "hello"} assert request.scope[CLIENT_REQUESTED_MODEL_SCOPE_KEY] == (requested_scope or "model") diff --git a/tests/unit/_auth/test_codex_turn.py b/tests/unit/_auth/test_codex_turn.py index 6253f2185fc..132f4377edb 100644 --- a/tests/unit/_auth/test_codex_turn.py +++ b/tests/unit/_auth/test_codex_turn.py @@ -1,4 +1,4 @@ -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import replace from typing import Final @@ -210,7 +210,9 @@ async def test_jwt_wrapper_never_turns_a_native_fault_into_subscription_fallback ) -> None: class Validator: @staticmethod - async def auth_builder(**_kwargs: object) -> bool: + async def auth_builder( + *, api_key: str, request_data: Mapping[str, object], general_settings: Mapping[str, object], route: str + ) -> bool: if failure is not None: raise failure return False diff --git a/tests/unit/proxy/auth/test_model_policy_read.py b/tests/unit/proxy/auth/test_model_policy_read.py new file mode 100644 index 00000000000..0bafa72e79b --- /dev/null +++ b/tests/unit/proxy/auth/test_model_policy_read.py @@ -0,0 +1,18 @@ +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.proxy.auth.model_policy_read import require_native_fact + + +@pytest.mark.parametrize("value", (False, 0, "", (), object())) +def test_read_precondition_preserves_falsy_and_opaque_facts(value: object) -> None: + assert require_native_fact(value) is value + + +def test_missing_read_fact_is_terminal() -> None: + with pytest.raises(NativeAuthRuntimeError) as raised: + require_native_fact(None) + error: Final = raised.value + assert isinstance(error, RuntimeError)