chore: merge latest main into usage cost branch

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm-rust/Cargo.lock

@@ -894,6 +894,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
 dependencies = [
  "fastrand",
+ "tokio",
 ]
 
 [[package]]
@@ -3657,19 +3658,15 @@ dependencies = [
 name = "litellm-cache-s3"
 version = "0.1.0"
 dependencies = [
- "aws-credential-types",
  "aws-sdk-s3",
- "aws-smithy-runtime-api",
  "aws-smithy-types",
- "aws-types",
  "futures-util",
- "http 1.4.2",
  "litellm-auth-aws",
  "litellm-cache",
  "litellm-cache-testing",
  "litellm-error",
  "litellm-http",
- "reqwest",
+ "litellm-s3",
  "rstest",
  "serde_json",
  "tokio",
@@ -4346,12 +4343,14 @@ dependencies = [
  "jsonschema",
  "litellm-cost",
  "litellm-error",
+ "litellm-tracing",
  "macro_rules_attribute",
  "rstest",
  "schemars 1.2.2",
  "serde",
  "serde_json",
  "serde_with",
+ "sha1 0.10.7",
  "strum",
  "thiserror 2.0.19",
  "time",
@@ -4361,9 +4360,13 @@ dependencies = [
 name = "litellm-model-catalog-refresh"
 version = "0.1.0"
 dependencies = [
+ "aws-credential-types",
+ "backon",
+ "bytes",
  "litellm-error",
  "litellm-http",
- "litellm-model-catalog",
+ "litellm-s3",
+ "litellm-tracing",
  "reqwest",
  "rstest",
  "thiserror 2.0.19",
@@ -4402,6 +4405,7 @@ dependencies = [
  "litellm-llms",
  "litellm-llms-types",
  "litellm-model-catalog",
+ "litellm-model-catalog-refresh",
  "litellm-python-compat",
  "litellm-redis-client",
  "litellm-redis-connect",
@@ -4609,6 +4613,23 @@ dependencies = [
  "serde_with",
 ]
 
+[[package]]
+name = "litellm-s3"
+version = "0.1.0"
+dependencies = [
+ "aws-config",
+ "aws-credential-types",
+ "aws-sdk-s3",
+ "aws-smithy-runtime-api",
+ "aws-smithy-types",
+ "aws-types",
+ "http 1.4.2",
+ "litellm-http",
+ "reqwest",
+ "rstest",
+ "thiserror 2.0.19",
+]
+
 [[package]]
 name = "litellm-secret-value"
 version = "0.1.0"

litellm-rust/Cargo.toml

@@ -60,6 +60,7 @@ litellm-cache = { path = "crates/cache" }
 litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
 litellm-cache-memory = { path = "crates/cache-memory" }
 litellm-cache-s3 = { path = "crates/cache-s3" }
+litellm-s3 = { path = "crates/s3" }
 litellm-cache-gcs = { path = "crates/cache-gcs" }
 litellm-cache-disk = { path = "crates/cache-disk" }
 litellm-cache-response = { path = "crates/cache-response" }
@@ -90,6 +91,8 @@ axum = { version = "0.8.9", default-features = false, features = ["http1", "toki
 axum-login = "0.18.0"
 tower-sessions = { version = "0.14.0", features = ["memory-store"] }
 bytes = "1"
+backon = { version = "1.6.0", default-features = false, features = ["std", "tokio-sleep"] }
+aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rt-tokio"] }
 http = "1"
 google-cloud-auth = { version = "1.16.0", default-features = false }
 jsonwebtoken = { version = "11.1.0", default-features = false }

litellm-rust/crates/auth-aws/src/lib.rs

@@ -1,9 +1,11 @@
 mod aws;
 pub mod constants;
 mod error;
+mod provider;
 mod signer;
 
 pub use aws::*;
 pub use aws_credential_types::Credentials;
 pub use error::{CredentialFlow, Error, SigningStage};
+pub use provider::AwsCredentialsProvider;
 pub use signer::SigV4Signer;

litellm-rust/crates/auth-aws/src/provider.rs

@@ -0,0 +1,135 @@
+use std::sync::Arc;
+
+use aws_credential_types::Credentials;
+use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
+
+use crate::{AwsAuthConfig, AwsAuthService};
+
+type Environment = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
+
+#[derive(Clone)]
+pub struct AwsCredentialsProvider {
+    auth: AwsAuthService,
+    config: Box<AwsAuthConfig>,
+    environment: Environment,
+}
+
+impl AwsCredentialsProvider {
+    /// Explicit keys in `config` as given, otherwise LiteLLM's `AWS_*` settings read through `environment`
+    pub fn configured(
+        auth: AwsAuthService,
+        config: AwsAuthConfig,
+        environment: impl Fn(&str) -> Option<String> + Send + Sync + 'static,
+    ) -> Self {
+        Self {
+            auth,
+            config: Box::new(config),
+            environment: Arc::new(environment),
+        }
+    }
+
+    async fn credentials(&self) -> Result<Credentials, crate::Error> {
+        match (&self.config.access_key_id, &self.config.secret_access_key) {
+            (Some(access_key_id), Some(secret_access_key)) => Ok(Credentials::new(
+                access_key_id.expose(),
+                secret_access_key.expose(),
+                self.config
+                    .session_token
+                    .as_ref()
+                    .map(|token| token.expose().to_owned()),
+                None,
+                "litellm-configured",
+            )),
+            _ => {
+                self.auth
+                    .resolve_credentials(
+                        AwsAuthConfig::clone(&self.config),
+                        self.environment.as_ref(),
+                    )
+                    .await
+            }
+        }
+    }
+}
+
+impl ProvideCredentials for AwsCredentialsProvider {
+    fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
+    where
+        Self: 'a,
+    {
+        future::ProvideCredentials::new(async {
+            self.credentials()
+                .await
+                .map_err(CredentialsError::provider_error)
+        })
+    }
+}
+
+impl std::fmt::Debug for AwsCredentialsProvider {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("AwsCredentialsProvider")
+            .finish_non_exhaustive()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
+        AwsAuthConfig {
+            access_key_id: Some("key".into()),
+            secret_access_key: Some("secret".into()),
+            session_token: session_token.map(Into::into),
+            region_name: Some("us-east-1".to_string()),
+            ..Default::default()
+        }
+    }
+
+    fn ambient_token(name: &str) -> Option<String> {
+        (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
+    }
+
+    fn environment_keys(name: &str) -> Option<String> {
+        match name {
+            "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
+            "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
+            "AWS_SESSION_TOKEN" => Some("env-token".to_string()),
+            _ => None,
+        }
+    }
+
+    #[rstest]
+    #[case::explicit_keys_ignore_an_ambient_session_token(
+        explicit(None), ambient_token, ("key", "secret", None)
+    )]
+    #[case::explicit_keys_keep_their_session_token(
+        explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
+    )]
+    #[case::environment_keys_resolve_with_their_session_token(
+        AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
+    )]
+    #[tokio::test]
+    async fn configured_credentials_read_the_environment_they_are_given(
+        #[case] config: AwsAuthConfig,
+        #[case] environment: fn(&str) -> Option<String>,
+        #[case] expected: (&str, &str, Option<&str>),
+    ) {
+        let credentials =
+            AwsCredentialsProvider::configured(AwsAuthService::default(), config, environment)
+                .provide_credentials()
+                .await
+                .unwrap();
+        assert_eq!(
+            (
+                credentials.access_key_id(),
+                credentials.secret_access_key(),
+                credentials.session_token(),
+            ),
+            expected
+        );
+    }
+}

litellm-rust/crates/cache-s3/Cargo.toml

@@ -9,14 +9,10 @@ repository.workspace = true
 litellm-http.workspace = true
 litellm-cache.workspace = true
 litellm-auth-aws.workspace = true
-aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
-aws-credential-types = "1.3.0"
-aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] }
-aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] }
-aws-types = "1.6.0"
+litellm-s3.workspace = true
+aws-sdk-s3.workspace = true
+aws-smithy-types = "1.6.0"
 futures-util.workspace = true
-http.workspace = true
-reqwest.workspace = true
 tokio.workspace = true
 
 [dev-dependencies]

litellm-rust/crates/cache-s3/src/auth.rs

@@ -1,61 +0,0 @@
-use aws_credential_types::{
-    Credentials as AwsCredentials,
-    provider::{ProvideCredentials, error::CredentialsError, future},
-};
-use litellm_auth_aws::{AwsAuthConfig, AwsAuthService};
-
-#[derive(Clone)]
-pub struct S3Credentials {
-    auth: AwsAuthService,
-    config: AwsAuthConfig,
-    env: fn(&str) -> Option<String>,
-}
-
-impl S3Credentials {
-    pub fn new(config: AwsAuthConfig) -> Self {
-        Self::with_env(config, |name| std::env::var(name).ok())
-    }
-
-    pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
-        Self {
-            auth: AwsAuthService::default(),
-            config,
-            env,
-        }
-    }
-}
-
-impl ProvideCredentials for S3Credentials {
-    fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
-    where
-        Self: 'a,
-    {
-        future::ProvideCredentials::new(async {
-            if let (Some(access_key_id), Some(secret_access_key)) = (
-                self.config.access_key_id.clone(),
-                self.config.secret_access_key.clone(),
-            ) {
-                return Ok(AwsCredentials::new(
-                    access_key_id.expose(),
-                    secret_access_key.expose(),
-                    self.config
-                        .session_token
-                        .as_ref()
-                        .map(|value| value.expose().to_owned()),
-                    None,
-                    "litellm-s3-cache",
-                ));
-            }
-            self.auth
-                .resolve_credentials(self.config.clone(), &self.env)
-                .await
-                .map_err(CredentialsError::provider_error)
-        })
-    }
-}
-
-impl std::fmt::Debug for S3Credentials {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("S3Credentials").finish_non_exhaustive()
-    }
-}

litellm-rust/crates/cache-s3/src/cache.rs

@@ -5,20 +5,18 @@ use std::{
 };
 
 use aws_sdk_s3::{
-    config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
+    config::{RequestChecksumCalculation, ResponseChecksumValidation},
     error::SdkError,
     primitives::ByteStream,
 };
 use aws_smithy_types::{DateTime, date_time::Format};
 use futures_util::future::try_join_all;
-use litellm_auth_aws::AwsAuthConfig;
+use litellm_auth_aws::AwsCredentialsProvider;
 use litellm_cache::{
     BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
 };
 use tokio::runtime::Handle;
 
-use crate::{auth::S3Credentials, transport::ReqwestHttpClient};
-
 pub struct S3Endpoint {
     pub url: String,
 }
@@ -28,7 +26,7 @@ pub struct S3CacheConfig {
     pub key_prefix: String,
     pub region: String,
     pub endpoint: Option<S3Endpoint>,
-    pub auth: AwsAuthConfig,
+    pub credentials: AwsCredentialsProvider,
 }
 
 pub struct S3Cache<C: CacheCodec> {
@@ -49,17 +47,14 @@ impl<C: CacheCodec> S3Cache<C> {
         runtime: Handle,
     ) -> Self {
         let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
-        let base = aws_sdk_s3::Config::builder()
-            .behavior_version(BehaviorVersion::latest())
-            .region(Region::new(config.region.clone()))
-            .http_client(ReqwestHttpClient(http))
-            .credentials_provider(S3Credentials::new(config.auth))
-            .request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
-            .response_checksum_validation(ResponseChecksumValidation::WhenRequired);
-        let builder = match &endpoint_url {
-            Some(url) => base.endpoint_url(url).force_path_style(true),
-            None => base,
-        };
+        let builder = litellm_s3::config(
+            config.region.clone(),
+            endpoint_url.as_deref(),
+            http,
+            config.credentials,
+        )
+        .request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
+        .response_checksum_validation(ResponseChecksumValidation::WhenRequired);
         Self {
             client: aws_sdk_s3::Client::from_conf(builder.build()),
             codec,

litellm-rust/crates/cache-s3/src/lib.rs

@@ -1,6 +1,3 @@
-mod auth;
 mod cache;
-mod transport;
 
-pub use auth::S3Credentials;
 pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

litellm-rust/crates/cache-s3/tests/auth.rs

@@ -1,57 +0,0 @@
-use aws_credential_types::provider::ProvideCredentials;
-use litellm_auth_aws::AwsAuthConfig;
-use litellm_cache_s3::S3Credentials;
-use rstest::rstest;
-
-fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
-    AwsAuthConfig {
-        access_key_id: Some("key".into()),
-        secret_access_key: Some("secret".into()),
-        session_token: session_token.map(Into::into),
-        region_name: Some("us-east-1".to_string()),
-        ..Default::default()
-    }
-}
-
-fn ambient_token(name: &str) -> Option<String> {
-    (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
-}
-
-fn environment_keys(name: &str) -> Option<String> {
-    match name {
-        "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
-        "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
-        "AWS_SESSION_TOKEN" => Some("env-token".to_string()),
-        _ => None,
-    }
-}
-
-#[rstest]
-#[case::explicit_keys_ignore_an_ambient_session_token(
-    explicit(None), ambient_token, ("key", "secret", None)
-)]
-#[case::explicit_keys_keep_their_session_token(
-    explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
-)]
-#[case::environment_keys_resolve_with_their_session_token(
-    AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
-)]
-#[tokio::test]
-async fn credentials_resolve(
-    #[case] config: AwsAuthConfig,
-    #[case] env: fn(&str) -> Option<String>,
-    #[case] expected: (&str, &str, Option<&str>),
-) {
-    let credentials = S3Credentials::with_env(config, env)
-        .provide_credentials()
-        .await
-        .unwrap();
-    assert_eq!(
-        (
-            credentials.access_key_id(),
-            credentials.secret_access_key(),
-            credentials.session_token(),
-        ),
-        expected
-    );
-}

litellm-rust/crates/cache-s3/tests/support/mod.rs

@@ -5,7 +5,7 @@ use std::{
     sync::{Arc, Mutex},
 };
 
-use litellm_auth_aws::AwsAuthConfig;
+use litellm_auth_aws::{AwsAuthConfig, AwsAuthService, AwsCredentialsProvider};
 use litellm_cache::JsonCodec;
 use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
 use serde_json::Value;
@@ -22,12 +22,16 @@ pub fn config(endpoint: &str) -> S3CacheConfig {
         endpoint: Some(S3Endpoint {
             url: endpoint.to_string(),
         }),
-        auth: AwsAuthConfig {
-            access_key_id: Some("key".into()),
-            secret_access_key: Some("secret".into()),
-            region_name: Some("us-east-1".to_string()),
-            ..Default::default()
-        },
+        credentials: AwsCredentialsProvider::configured(
+            AwsAuthService::default(),
+            AwsAuthConfig {
+                access_key_id: Some("key".into()),
+                secret_access_key: Some("secret".into()),
+                region_name: Some("us-east-1".to_string()),
+                ..Default::default()
+            },
+            |_| None,
+        ),
     }
 }
 

litellm-rust/crates/host-python/src/lib.rs

@@ -36,7 +36,7 @@ pub use marshal::{
 pub use owned::PythonOwned;
 pub use runtime::{
     ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value,
-    ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync,
+    ready_future, reserve_process_for_forking, run_async, run_async_value, run_isolated, run_sync,
     run_sync_value, runtime_handle, runtime_started,
 };
 pub use services::PythonHostCalls;

litellm-rust/crates/host-python/src/runtime.rs

@@ -68,6 +68,18 @@ fn runtime() -> PyResult<&'static Runtime> {
     Ok(pyo3_async_runtimes::tokio::get_runtime())
 }
 
+/// Runs `future` on a current-thread runtime built and shut down inside the call, its blocking pool joined before
+/// it returns. No thread outlives the call, so unlike every other door this one never claims the process in the fork
+/// gate: a process that only used it can still fork. The future must not hold resources tied to another runtime
+pub fn run_isolated<F: Future>(future: F) -> std::io::Result<F::Output> {
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()?;
+    let output = runtime.block_on(future);
+    runtime.shutdown_timeout(Duration::from_secs(5));
+    Ok(output)
+}
+
 /// The native runtime, for work that must outlive the call that schedules it.
 pub fn runtime_handle() -> PyResult<Handle> {
     Ok(runtime()?.handle().clone())
@@ -818,4 +830,21 @@ asyncio.run(exercise())
                 .expect("async route contract should hold");
         });
     }
+
+    #[rstest]
+    fn an_isolated_run_joins_blocking_work_it_left_running_before_returning() {
+        let finished = Arc::new(AtomicBool::new(false));
+        let flag = Arc::clone(&finished);
+        let output = run_isolated(async move {
+            drop(tokio::task::spawn_blocking(move || {
+                thread::sleep(Duration::from_millis(50));
+                flag.store(true, Ordering::SeqCst);
+            }));
+            7
+        })
+        .unwrap();
+
+        assert_eq!(output, 7);
+        assert!(finished.load(Ordering::SeqCst));
+    }
 }

litellm-rust/crates/host-python/tests/isolated_run.rs

@@ -0,0 +1,9 @@
+use litellm_host_python::{reserve_process_for_forking, run_isolated, runtime_started};
+
+#[test]
+fn an_isolated_run_leaves_the_process_free_to_reserve_for_forking() {
+    run_isolated(async { tokio::task::spawn_blocking(|| ()).await.unwrap() }).unwrap();
+
+    assert!(!runtime_started());
+    assert!(reserve_process_for_forking().is_ok());
+}

litellm-rust/crates/http/src/pool.rs

@@ -55,36 +55,13 @@ impl HttpClientPool {
         config: &HttpClientConfig,
         variant: ClientVariant,
     ) -> Result<Client, Error> {
-        let effective = match variant {
-            ClientVariant::Media => HttpClientConfig {
-                client_certificate: None,
-                proxies: EnvironmentProxies::default(),
-                ..config.clone()
-            },
-            ClientVariant::UnpinnedMedia => HttpClientConfig {
-                client_certificate: None,
-                ..config.clone()
-            },
-            ClientVariant::WebSocket => HttpClientConfig {
-                http2: false,
-                verify: match &config.verify {
-                    Verify::BuiltInRoots => Verify::NativeRoots,
-                    verify => verify.clone(),
-                },
-                ..config.clone()
-            },
-            ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
-        };
-        let key = (effective, variant);
+        let key = (effective(config, variant), variant);
         if let Some(pooled) = self.lock().get(&key)
             && pooled.built_at.elapsed() < self.ttl
         {
             return Ok(Client::new(pooled.client.clone()));
         }
-        let client = self
-            .apply(variant, reqwest::ClientBuilder::try_from(&key.0)?)
-            .build()
-            .map_err(Repr::Client)?;
+        let client = self.build(&key.0, variant)?;
         self.lock().insert(
             key,
             PooledClient {
@@ -95,6 +72,26 @@ impl HttpClientPool {
         Ok(Client::new(client))
     }
 
+    /// A client nobody else holds, for a caller that drives it on a runtime of its own and drops both together
+    pub fn unpooled(
+        &self,
+        config: &HttpClientConfig,
+        variant: ClientVariant,
+    ) -> Result<Client, Error> {
+        self.build(&effective(config, variant), variant)
+            .map(Client::new)
+    }
+
+    fn build(
+        &self,
+        config: &HttpClientConfig,
+        variant: ClientVariant,
+    ) -> Result<reqwest::Client, Error> {
+        self.apply(variant, reqwest::ClientBuilder::try_from(config)?)
+            .build()
+            .map_err(|error| Error::from(Repr::Client(error)))
+    }
+
     fn lock(&self) -> MutexGuard<'_, Clients> {
         self.clients.lock().unwrap_or_else(PoisonError::into_inner)
     }
@@ -116,6 +113,29 @@ impl HttpClientPool {
     }
 }
 
+fn effective(config: &HttpClientConfig, variant: ClientVariant) -> HttpClientConfig {
+    match variant {
+        ClientVariant::Media => HttpClientConfig {
+            client_certificate: None,
+            proxies: EnvironmentProxies::default(),
+            ..config.clone()
+        },
+        ClientVariant::UnpinnedMedia => HttpClientConfig {
+            client_certificate: None,
+            ..config.clone()
+        },
+        ClientVariant::WebSocket => HttpClientConfig {
+            http2: false,
+            verify: match &config.verify {
+                Verify::BuiltInRoots => Verify::NativeRoots,
+                verify => verify.clone(),
+            },
+            ..config.clone()
+        },
+        ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::{

litellm-rust/crates/model-catalog-refresh/AGENTS.md

@@ -0,0 +1,9 @@
+# Model catalog refresh
+
+- `fetch` fetches one document over http(s) or `s3://` and never parses it
+- Only 429, 500, 502, 503, 504 and transport errors are retried. S3 errors are final
+- `Failed.retryable` means only the attempt limit ended the fetch
+- A server's Retry-After overrides the backoff, capped at `max_delay`
+- `Start::After(failed)` resumes the backoff schedule where another caller's first attempt left off
+- Error reasons use the URL with userinfo, query and fragment stripped, since presigned URLs carry credentials
+- Bodies over `MAX_BODY_BYTES` are rejected

litellm-rust/crates/model-catalog-refresh/Cargo.toml

@@ -6,14 +6,19 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
+backon.workspace = true
+bytes.workspace = true
+litellm-error.workspace = true
 litellm-http.workspace = true
-litellm-model-catalog.workspace = true
+litellm-s3.workspace = true
+litellm-tracing.workspace = true
+reqwest.workspace = true
 thiserror.workspace = true
+tokio = { workspace = true, features = ["time"] }
 url.workspace = true
 
 [dev-dependencies]
-litellm-error.workspace = true
+aws-credential-types = "1.3.0"
 litellm-http = { workspace = true, features = ["test-support"] }
-reqwest.workspace = true
 rstest.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] }

litellm-rust/crates/model-catalog-refresh/src/attempt.rs

@@ -0,0 +1,153 @@
+use std::time::Duration;
+
+use bytes::Bytes;
+
+use litellm_http::Client;
+use litellm_s3::S3Uri;
+use litellm_s3::sdk::config::retry::RetryConfig;
+use litellm_s3::sdk::config::timeout::TimeoutConfig;
+use reqwest::header::{ETAG, HeaderMap, RETRY_AFTER};
+use url::Url;
+
+use crate::{Error, RETRYABLE_STATUSES, S3Access};
+
+pub(crate) const MAX_BODY_BYTES: u64 = 100 * 1024 * 1024;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Fetched {
+    pub body: Vec<u8>,
+    pub etag: Option<String>,
+}
+
+#[derive(Debug)]
+pub(crate) enum Cause {
+    Status(u16),
+    Transport(Error),
+}
+
+#[derive(Debug)]
+pub(crate) struct Rejected {
+    pub(crate) cause: Cause,
+    pub(crate) retryable: bool,
+    pub(crate) retry_after: Option<Duration>,
+}
+
+impl Rejected {
+    pub(crate) fn fatal(error: Error) -> Self {
+        Self {
+            cause: Cause::Transport(error),
+            retryable: false,
+            retry_after: None,
+        }
+    }
+}
+
+/// Any status below 400 is the document
+pub(crate) async fn http(
+    client: &Client,
+    url: Url,
+    timeout: Duration,
+) -> Result<Fetched, Rejected> {
+    let transport = |error: litellm_http::Error| Rejected {
+        cause: Cause::Transport(Error::from(error)),
+        retryable: true,
+        retry_after: None,
+    };
+    let mut response = tokio::time::timeout(timeout, client.get(url).send())
+        .await
+        .map_err(litellm_http::Error::read_timeout)
+        .and_then(|sent| sent.map_err(litellm_http::Error::from))
+        .map_err(transport)?;
+    let status = response.status().as_u16();
+    let etag = etag(response.headers());
+    let retry_after = response
+        .headers()
+        .get(RETRY_AFTER)
+        .and_then(|value| value.to_str().ok())
+        .and_then(|value| value.trim().parse().ok())
+        .map(Duration::from_secs);
+    if status >= 400 {
+        return Err(Rejected {
+            cause: Cause::Status(status),
+            retryable: RETRYABLE_STATUSES.contains(&status),
+            retry_after,
+        });
+    }
+    if response
+        .content_length()
+        .is_some_and(|length| length > MAX_BODY_BYTES)
+    {
+        return Err(Rejected::fatal(Error::body_too_large(MAX_BODY_BYTES)));
+    }
+    let mut body = Vec::new();
+    while let Some(chunk) = tokio::time::timeout(timeout, response.chunk())
+        .await
+        .map_err(litellm_http::Error::body_read_timeout)
+        .and_then(|read| read.map_err(litellm_http::Error::from))
+        .map_err(transport)?
+    {
+        append(&mut body, &chunk).map_err(Rejected::fatal)?;
+    }
+    Ok(Fetched { body, etag })
+}
+
+/// S3 failures are final: the SDK already classifies its own errors and a missing or denied object stays so
+pub(crate) async fn s3(
+    client: &Client,
+    s3: &S3Access,
+    object: &S3Uri,
+    timeout: Duration,
+) -> Result<Fetched, Rejected> {
+    let fetched = async {
+        let config = match s3 {
+            S3Access::Environment { region } => {
+                litellm_s3::environment_config(client.clone(), region.clone()).await
+            }
+            S3Access::Config(config) => config.to_builder(),
+        }
+        .retry_config(RetryConfig::disabled())
+        .timeout_config(
+            TimeoutConfig::builder()
+                .operation_attempt_timeout(timeout)
+                .build(),
+        )
+        .build();
+        let output = litellm_s3::sdk::Client::from_conf(config)
+            .get_object()
+            .bucket(&object.bucket)
+            .key(&object.key)
+            .send()
+            .await?;
+        if output
+            .content_length()
+            .is_some_and(|length| length.unsigned_abs() > MAX_BODY_BYTES)
+        {
+            return Err(Error::body_too_large(MAX_BODY_BYTES));
+        }
+        let etag = output.e_tag().map(str::to_owned);
+        let mut stream = output.body;
+        let mut body = Vec::new();
+        while let Some(chunk) = tokio::time::timeout(timeout, stream.try_next()).await?? {
+            append(&mut body, &chunk)?;
+        }
+        Ok::<_, Error>(Fetched { body, etag })
+    };
+    fetched.await.map_err(Rejected::fatal)
+}
+
+fn append(body: &mut Vec<u8>, chunk: &Bytes) -> Result<(), Error> {
+    if (body.len() + chunk.len()) as u64 > MAX_BODY_BYTES {
+        return Err(Error::body_too_large(MAX_BODY_BYTES));
+    }
+    body.extend_from_slice(chunk);
+    Ok(())
+}
+
+fn etag(headers: &HeaderMap) -> Option<String> {
+    let values: Vec<&str> = headers
+        .get_all(ETAG)
+        .iter()
+        .filter_map(|value| value.to_str().ok())
+        .collect();
+    (!values.is_empty()).then(|| values.join(", "))
+}

litellm-rust/crates/model-catalog-refresh/src/error.rs

@@ -1,6 +1,9 @@
-use std::io;
-
+use litellm_s3::sdk::config::http::HttpResponse;
+use litellm_s3::sdk::error::SdkError;
+use litellm_s3::sdk::operation::get_object::GetObjectError;
+use litellm_s3::sdk::primitives::ByteStreamError;
 use thiserror::Error;
+use tokio::time::error::Elapsed;
 
 #[derive(Debug, Error)]
 #[error(transparent)]
@@ -10,18 +13,48 @@ pub struct Error(Box<Repr>);
 enum Repr {
     #[error(transparent)]
     Http(#[from] litellm_http::Error),
-    #[error("could not read the catalog")]
-    Read(#[source] io::Error),
     #[error("invalid catalog URL")]
     InvalidUrl(#[source] url::ParseError),
     #[error("unsupported catalog URL scheme {scheme}")]
     UnsupportedScheme { scheme: String },
+    #[error("invalid S3 URL")]
+    S3Uri(#[source] litellm_s3::Error),
+    #[error("S3 GetObject failed")]
+    S3GetObject(#[source] Box<SdkError<GetObjectError, HttpResponse>>),
+    #[error("S3 GetObject body read failed")]
+    S3Body(#[source] ByteStreamError),
+    #[error("S3 GetObject timed out")]
+    S3Timeout(#[source] Elapsed),
+    #[error("response body is over the {limit} byte limit")]
+    BodyTooLarge { limit: u64 },
 }
 
 impl Error {
     pub fn unsupported_scheme(scheme: String) -> Self {
         Self(Box::new(Repr::UnsupportedScheme { scheme }))
     }
+
+    pub(crate) fn body_too_large(limit: u64) -> Self {
+        Self(Box::new(Repr::BodyTooLarge { limit }))
+    }
+}
+
+impl From<SdkError<GetObjectError, HttpResponse>> for Error {
+    fn from(cause: SdkError<GetObjectError, HttpResponse>) -> Self {
+        Self(Box::new(Repr::S3GetObject(Box::new(cause))))
+    }
+}
+
+impl From<ByteStreamError> for Error {
+    fn from(cause: ByteStreamError) -> Self {
+        Self(Box::new(Repr::S3Body(cause)))
+    }
+}
+
+impl From<Elapsed> for Error {
+    fn from(cause: Elapsed) -> Self {
+        Self(Box::new(Repr::S3Timeout(cause)))
+    }
 }
 
 impl From<litellm_http::Error> for Error {
@@ -30,9 +63,9 @@ impl From<litellm_http::Error> for Error {
     }
 }
 
-impl From<io::Error> for Error {
-    fn from(cause: io::Error) -> Self {
-        Self(Box::new(Repr::Read(cause)))
+impl From<litellm_s3::Error> for Error {
+    fn from(cause: litellm_s3::Error) -> Self {
+        Self(Box::new(Repr::S3Uri(cause)))
     }
 }
 
@@ -44,29 +77,18 @@ impl From<url::ParseError> for Error {
 
 #[cfg(test)]
 mod tests {
+    use std::io;
+
     use litellm_error::{Report, find};
     use rstest::rstest;
 
     use super::*;
-    use crate::{DoneReason, Reply, Step};
 
     #[rstest]
-    fn exhausted_transport_keeps_its_typed_cause() {
-        let reply = Reply::Transport {
-            error: io::Error::from(io::ErrorKind::ConnectionReset).into(),
-        };
-        let Reply::Transport { error } = reply else {
-            panic!("transport failure")
-        };
-        let step = Step::Done {
-            reason: DoneReason::Exhausted(error),
-        };
-        let Step::Done {
-            reason: DoneReason::Exhausted(error),
-        } = step
-        else {
-            panic!("exhausted attempts")
-        };
+    fn an_s3_failure_keeps_its_typed_cause_once() {
+        let error = Error::from(SdkError::<GetObjectError, HttpResponse>::timeout_error(
+            io::Error::from(io::ErrorKind::ConnectionReset),
+        ));
         let cause = find::<io::Error>(&error).expect("typed I/O cause");
         assert_eq!(cause.kind(), io::ErrorKind::ConnectionReset);
         assert_eq!(

litellm-rust/crates/model-catalog-refresh/src/fetch.rs

@@ -0,0 +1,122 @@
+use std::borrow::Cow;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::time::Duration;
+
+use backon::Retryable;
+use litellm_error::Report;
+use litellm_http::Client;
+use url::Url;
+
+use crate::attempt::{self, Cause, Fetched, Rejected};
+use crate::{RetryPolicy, Source};
+
+/// How an `s3://` fetch reaches S3
+#[derive(Clone, Debug)]
+pub enum S3Access {
+    /// What an AWS SDK finds in the environment and shared config files, loaded on each `s3://` attempt, with
+    /// `region` winning over the region it finds
+    Environment { region: Option<String> },
+    /// A config the caller built with its region, endpoint and credentials
+    Config(Box<litellm_s3::sdk::Config>),
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Failed {
+    pub reason: String,
+    /// Another attempt could succeed. Only the attempt limit ended the fetch
+    pub retryable: bool,
+    pub retry_after: Option<Duration>,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum Start {
+    Fresh,
+    /// Continues a fetch whose first attempt failed retryably elsewhere, waiting before its next attempt
+    After(Failed),
+}
+
+pub async fn fetch(
+    client: &Client,
+    s3: &S3Access,
+    url: &str,
+    policy: RetryPolicy,
+    start: Start,
+) -> Result<Fetched, Failed> {
+    let shown = redacted(url);
+    let source =
+        Source::parse(url).map_err(|error| Failed::of(&shown, &Rejected::fatal(error), 1))?;
+    let mut backoff = policy.backoff();
+    let made = AtomicU32::new(0);
+    if let Start::After(failed) = start {
+        let Some(wait) = backoff.next() else {
+            return Err(failed);
+        };
+        tokio::time::sleep(policy.delay(wait, failed.retry_after)).await;
+        made.store(1, Ordering::Relaxed);
+    }
+    let attempt = || {
+        made.fetch_add(1, Ordering::Relaxed);
+        async {
+            match &source {
+                Source::Http(http) => attempt::http(client, http.clone(), policy.timeout).await,
+                Source::S3(object) => attempt::s3(client, s3, object, policy.timeout).await,
+            }
+        }
+    };
+    attempt
+        .retry(backoff)
+        .sleep(tokio::time::sleep)
+        .when(|rejected| rejected.retryable)
+        .adjust(|rejected, wait| wait.map(|wait| policy.delay(wait, rejected.retry_after)))
+        .notify(|rejected, wait| {
+            litellm_tracing::warn!(
+                "LiteLLM: model cost map fetch attempt {}/{} failed ({}); retrying in {:.1}s",
+                made.load(Ordering::Relaxed),
+                policy.max_attempts,
+                reason(&shown, &rejected.cause),
+                wait.as_secs_f64(),
+            );
+        })
+        .await
+        .map_err(|rejected| Failed::of(&shown, &rejected, made.load(Ordering::Relaxed)))
+}
+
+impl Failed {
+    fn of(url: &str, rejected: &Rejected, attempts: u32) -> Self {
+        let reason = reason(url, &rejected.cause);
+        Self {
+            reason: match rejected.retryable && attempts > 1 {
+                true => format!("{reason} (after {attempts} attempts)"),
+                false => reason,
+            },
+            retryable: rejected.retryable,
+            retry_after: rejected.retry_after,
+        }
+    }
+}
+
+fn reason(url: &str, cause: &Cause) -> String {
+    match cause {
+        Cause::Status(status) => format!("HTTP {status} from {url}"),
+        Cause::Transport(error) => format!("error fetching {url}: {}", Report::new(error)),
+    }
+}
+
+/// Drops userinfo, query and fragment, where presigned URLs and tokens carry credentials
+fn redacted(url: &str) -> Cow<'_, str> {
+    let Ok(mut parsed) = Url::parse(url) else {
+        return Cow::Borrowed(url);
+    };
+    if parsed.username().is_empty()
+        && parsed.password().is_none()
+        && parsed.query().is_none()
+        && parsed.fragment().is_none()
+    {
+        return Cow::Borrowed(url);
+    }
+    let _ = parsed.set_username("");
+    let _ = parsed.set_password(None);
+    parsed.set_query(None);
+    parsed.set_fragment(None);
+    Cow::Owned(parsed.into())
+}

litellm-rust/crates/model-catalog-refresh/src/lib.rs

@@ -1,9 +1,11 @@
+mod attempt;
 mod error;
+mod fetch;
 mod policy;
 mod source;
-mod step;
 
+pub use attempt::Fetched;
 pub use error::Error;
+pub use fetch::*;
 pub use policy::*;
 pub use source::*;
-pub use step::*;

litellm-rust/crates/model-catalog-refresh/src/policy.rs

@@ -1,20 +1,76 @@
 use std::time::Duration;
 
-use litellm_model_catalog::IntegrityLimits;
+use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder};
 
-/// Main's defaults, 3 attempts, 30 s max wait, 5 s timeout, backoff base 2.0, jitter 0.0 to 1.0 s
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct RetryPolicy {
     pub max_attempts: u32,
-    pub max_wait: Duration,
     pub timeout: Duration,
-    pub backoff_base_seconds: f64,
-    pub jitter_min_seconds: f64,
-    pub jitter_max_seconds: f64,
-    /// Main's 50 model floor and 0.5 shrink ratio
-    pub limits: IntegrityLimits,
+    pub min_delay: Duration,
+    pub max_delay: Duration,
 }
 
+pub const RETRYABLE_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
+
 /// Main's default `LITELLM_MODEL_COST_MAP_URL`
 pub const DEFAULT_REGISTRY_URL: &str =
     "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json";
+
+impl RetryPolicy {
+    pub(crate) fn backoff(&self) -> ExponentialBackoff {
+        ExponentialBuilder::new()
+            .with_min_delay(self.min_delay)
+            .with_max_delay(self.max_delay)
+            .with_max_times(self.max_attempts.saturating_sub(1) as usize)
+            .with_jitter()
+            .build()
+    }
+
+    pub(crate) fn delay(&self, backoff: Duration, retry_after: Option<Duration>) -> Duration {
+        retry_after.map_or(backoff, |server| server.min(self.max_delay))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    const POLICY: RetryPolicy = RetryPolicy {
+        max_attempts: 3,
+        timeout: Duration::from_secs(5),
+        min_delay: Duration::from_secs(1),
+        max_delay: Duration::from_secs(30),
+    };
+
+    #[rstest]
+    #[case::backoff_without_retry_after(None, 4)]
+    #[case::retry_after_wins(Some(7), 7)]
+    #[case::retry_after_zero(Some(0), 0)]
+    #[case::retry_after_capped(Some(120), 30)]
+    fn the_server_retry_after_overrides_the_backoff_up_to_max_delay(
+        #[case] retry_after: Option<u64>,
+        #[case] seconds: u64,
+    ) {
+        assert_eq!(
+            POLICY.delay(Duration::from_secs(4), retry_after.map(Duration::from_secs)),
+            Duration::from_secs(seconds)
+        );
+    }
+
+    #[rstest]
+    #[case::one_attempt(1, 0)]
+    #[case::three_attempts(3, 2)]
+    #[case::zero_attempts(0, 0)]
+    fn the_backoff_allows_one_retry_fewer_than_the_attempts(
+        #[case] max_attempts: u32,
+        #[case] retries: usize,
+    ) {
+        let policy = RetryPolicy {
+            max_attempts,
+            ..POLICY
+        };
+        assert_eq!(policy.backoff().count(), retries);
+    }
+}

litellm-rust/crates/model-catalog-refresh/src/source.rs

@@ -1,17 +1,59 @@
-use std::path::PathBuf;
-
+use litellm_s3::S3Uri;
 use url::Url;
 
+use crate::Error;
+
 /// Where a fetched catalog comes from
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum Source {
     Http(Url),
-    File(PathBuf),
-    S3(S3Object),
+    S3(S3Uri),
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct S3Object {
-    pub bucket: String,
-    pub key: String,
+impl Source {
+    pub fn parse(url: &str) -> Result<Self, Error> {
+        let parsed = Url::parse(url)?;
+        match parsed.scheme() {
+            "http" | "https" => Ok(Self::Http(parsed)),
+            "s3" => Ok(Self::S3(S3Uri::try_from(url)?)),
+            scheme => Err(Error::unsupported_scheme(scheme.to_owned())),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    #[rstest]
+    #[case::https("https://example.com/cost.json")]
+    #[case::http("http://127.0.0.1:8000/cost.json")]
+    fn http_urls_are_fetched_over_http(#[case] url: &str) {
+        assert_eq!(
+            Source::parse(url).unwrap(),
+            Source::Http(Url::parse(url).unwrap())
+        );
+    }
+
+    #[rstest]
+    fn s3_urls_name_a_bucket_and_key() {
+        assert_eq!(
+            Source::parse("s3://bucket/prices/cost.json").unwrap(),
+            Source::S3(S3Uri {
+                bucket: "bucket".into(),
+                key: "prices/cost.json".into(),
+            })
+        );
+    }
+
+    #[rstest]
+    #[case::s3_without_a_key("s3://bucket/")]
+    #[case::unsupported("ftp://example.com/cost.json")]
+    #[case::file("file:///tmp/cost.json")]
+    #[case::not_a_url("not a url")]
+    fn other_urls_are_refused(#[case] url: &str) {
+        assert!(Source::parse(url).is_err());
+    }
 }

litellm-rust/crates/model-catalog-refresh/src/step.rs

@@ -1,66 +0,0 @@
-use std::time::Duration;
-
-use litellm_model_catalog::Catalog;
-
-use crate::Error;
-use crate::source::Source;
-
-/// One attempt. Where to fetch from is the `Source`, so the request carries no URL
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct FetchRequest {
-    pub if_none_match: Option<String>,
-    pub timeout: Duration,
-}
-
-#[derive(Debug)]
-pub enum Reply {
-    /// An HTTP, S3 or file answer. The body is read only on 200
-    Response {
-        status: u16,
-        etag: Option<String>,
-        retry_after: Option<Duration>,
-        body: Option<Vec<u8>>,
-    },
-    /// No answer reached us
-    Transport { error: Error },
-}
-
-/// What the fetch state machine asks its driver to do next
-#[derive(Debug)]
-pub enum Step {
-    Fetch(FetchRequest),
-    Sleep(Duration),
-    /// A parsed document to install, with its rejected row errors
-    Adopt {
-        catalog: Box<Catalog>,
-        installed: Installed,
-        errors: Vec<litellm_model_catalog::Error>,
-    },
-    /// 304, a final status, or retries exhausted
-    Done {
-        reason: DoneReason,
-    },
-}
-
-#[derive(Debug)]
-pub enum DoneReason {
-    NotModified,
-    Status(u16),
-    Exhausted(Error),
-}
-
-/// What is installed now, the input to the next revalidation
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct Installed {
-    pub origin: Origin,
-    pub etag: Option<String>,
-    /// `git_blob_id` of the installed bytes
-    pub revision: String,
-    pub model_count: usize,
-}
-
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub enum Origin {
-    Bundled,
-    Remote { source: Source },
-}

litellm-rust/crates/model-catalog-refresh/tests/fetch.rs

@@ -0,0 +1,541 @@
+use std::io::{BufRead, BufReader, Write};
+use std::net::TcpListener;
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+
+use aws_credential_types::Credentials;
+use aws_credential_types::provider::error::CredentialsError;
+use aws_credential_types::provider::{ProvideCredentials, future};
+use litellm_http::Client;
+use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start, fetch};
+use rstest::rstest;
+
+const POLICY: RetryPolicy = RetryPolicy {
+    max_attempts: 3,
+    timeout: Duration::from_secs(5),
+    min_delay: Duration::from_millis(10),
+    max_delay: Duration::from_millis(50),
+};
+
+const OK: &str =
+    "HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}";
+const UNAVAILABLE: &str = "HTTP/1.1 503 Service Unavailable\r\nRetry-After: 0\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
+const NOT_FOUND: &str = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
+
+const SLOW_BODY: &[u8] = b"{      }";
+
+#[derive(Clone, Copy)]
+enum Answer {
+    Raw(&'static str),
+    Stall,
+    StallBody,
+    SlowBody,
+    Oversized,
+}
+
+/// A scripted upstream: the n-th request gets the n-th answer, repeating the last, and every request head is kept
+struct Upstream {
+    base: String,
+    heads: Arc<Mutex<Vec<String>>>,
+}
+
+impl Upstream {
+    fn serve(answers: Vec<Answer>) -> Self {
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let base = format!("http://{}", listener.local_addr().unwrap());
+        let heads = Arc::new(Mutex::new(Vec::new()));
+        let recorded = Arc::clone(&heads);
+        std::thread::spawn(move || {
+            for (index, stream) in listener.incoming().enumerate() {
+                let answer = answers[index.min(answers.len() - 1)];
+                let recorded = Arc::clone(&recorded);
+                std::thread::spawn(move || {
+                    let mut stream = stream.unwrap();
+                    let mut reader = BufReader::new(stream.try_clone().unwrap());
+                    let mut head = String::new();
+                    loop {
+                        let mut line = String::new();
+                        if reader.read_line(&mut line).unwrap() <= 2 {
+                            break;
+                        }
+                        head.push_str(&line);
+                    }
+                    recorded.lock().unwrap().push(head);
+                    match answer {
+                        Answer::Raw(raw) => stream.write_all(raw.as_bytes()).unwrap(),
+                        Answer::Stall => std::thread::sleep(Duration::from_secs(5)),
+                        Answer::Oversized => {
+                            let chunk = vec![b' '; 1024 * 1024];
+                            let written = stream
+                                .write_all(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n")
+                                .and_then(|()| {
+                                    (0..=100).try_for_each(|_| stream.write_all(&chunk))
+                                });
+                            drop(written);
+                        }
+                        Answer::StallBody => {
+                            stream
+                                .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{")
+                                .unwrap();
+                            std::thread::sleep(Duration::from_secs(5));
+                        }
+                        Answer::SlowBody => {
+                            stream.set_nodelay(true).unwrap();
+                            stream
+                                .write_all(
+                                    b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 8\r\nConnection: close\r\n\r\n",
+                                )
+                                .unwrap();
+                            for byte in SLOW_BODY {
+                                std::thread::sleep(Duration::from_millis(60));
+                                stream.write_all(&[*byte]).unwrap();
+                            }
+                        }
+                    }
+                });
+            }
+        });
+        Self { base, heads }
+    }
+
+    fn url(&self) -> String {
+        format!("{}/cost.json", self.base)
+    }
+
+    fn heads(&self) -> Vec<String> {
+        self.heads.lock().unwrap().clone()
+    }
+}
+
+fn header<'a>(head: &'a str, name: &str) -> Option<&'a str> {
+    head.lines().find_map(|line| {
+        let (key, value) = line.split_once(':')?;
+        key.eq_ignore_ascii_case(name).then(|| value.trim())
+    })
+}
+
+#[derive(Clone, Debug)]
+struct NoCredentials;
+
+impl ProvideCredentials for NoCredentials {
+    fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
+    where
+        Self: 'a,
+    {
+        future::ProvideCredentials::ready(Err(CredentialsError::not_loaded("no credentials")))
+    }
+}
+
+fn keys() -> Credentials {
+    Credentials::new(
+        "AKIDMODELCOST",
+        "model-cost-secret",
+        None,
+        None,
+        "model-cost-test",
+    )
+}
+
+async fn fetched(url: &str, policy: RetryPolicy, start: Start) -> Result<Fetched, Failed> {
+    fetch(
+        &Client::for_test(),
+        &S3Access::Environment { region: None },
+        url,
+        policy,
+        start,
+    )
+    .await
+}
+
+#[rstest]
+#[case::every_etag_joined(
+    "HTTP/1.1 200 OK\r\nETag: \"a\"\r\nETag: \"b\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
+    Some("\"a\", \"b\"")
+)]
+#[case::no_etag(
+    "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
+    None
+)]
+#[tokio::test]
+async fn a_200_returns_the_body_and_its_etags(
+    #[case] response: &'static str,
+    #[case] etag: Option<&str>,
+) {
+    let upstream = Upstream::serve(vec![Answer::Raw(response)]);
+    assert_eq!(
+        fetched(&upstream.url(), POLICY, Start::Fresh).await,
+        Ok(Fetched {
+            body: b"{}".to_vec(),
+            etag: etag.map(str::to_owned),
+        })
+    );
+    assert_eq!(upstream.heads().len(), 1);
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn a_fetch_runs_as_a_task_on_a_multi_thread_runtime() {
+    let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
+    let url = upstream.url();
+    let fetched = tokio::spawn(async move {
+        fetch(
+            &Client::for_test(),
+            &S3Access::Environment { region: None },
+            &url,
+            POLICY,
+            Start::Fresh,
+        )
+        .await
+    })
+    .await
+    .unwrap();
+    assert!(fetched.is_ok());
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_redirect_is_not_followed_and_its_body_is_the_document() {
+    let upstream = Upstream::serve(vec![
+        Answer::Raw(
+            "HTTP/1.1 301 Moved Permanently\r\nLocation: /elsewhere.json\r\nContent-Length: 7\r\nConnection: close\r\n\r\n{\"a\":1}",
+        ),
+        Answer::Raw(OK),
+    ]);
+    assert_eq!(
+        fetched(&upstream.url(), POLICY, Start::Fresh).await,
+        Ok(Fetched {
+            body: b"{\"a\":1}".to_vec(),
+            etag: None,
+        })
+    );
+    assert_eq!(upstream.heads().len(), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_retry_after_shorter_than_the_backoff_is_honored() {
+    let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE), Answer::Raw(OK)]);
+    let policy = RetryPolicy {
+        min_delay: Duration::from_secs(10),
+        max_delay: Duration::from_secs(10),
+        ..POLICY
+    };
+    let fetched = tokio::time::timeout(
+        Duration::from_secs(2),
+        fetched(&upstream.url(), policy, Start::Fresh),
+    )
+    .await
+    .expect("Retry-After: 0 replaces the 10s backoff");
+    assert!(fetched.is_ok(), "{fetched:?}");
+    assert_eq!(upstream.heads().len(), 2);
+}
+
+#[rstest]
+#[tokio::test]
+async fn retryable_failures_stop_at_the_attempt_limit() {
+    let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]);
+    let url = upstream.url();
+    assert_eq!(
+        fetched(&url, POLICY, Start::Fresh).await,
+        Err(Failed {
+            reason: format!("HTTP 503 from {url} (after 3 attempts)"),
+            retryable: true,
+            retry_after: Some(Duration::ZERO),
+        })
+    );
+    assert_eq!(upstream.heads().len(), 3);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_404_fails_on_the_first_attempt() {
+    let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]);
+    let url = upstream.url();
+    assert_eq!(
+        fetched(&url, POLICY, Start::Fresh).await,
+        Err(Failed {
+            reason: format!("HTTP 404 from {url}"),
+            retryable: false,
+            retry_after: None,
+        })
+    );
+    assert_eq!(upstream.heads().len(), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_stalled_response_times_out_and_is_retried() {
+    let upstream = Upstream::serve(vec![Answer::Stall]);
+    let url = upstream.url();
+    let policy = RetryPolicy {
+        max_attempts: 2,
+        timeout: Duration::from_millis(100),
+        ..POLICY
+    };
+    let failed = fetched(&url, policy, Start::Fresh).await.unwrap_err();
+    assert!(
+        failed
+            .reason
+            .starts_with(&format!("error fetching {url}: "))
+            && failed.reason.ends_with(" (after 2 attempts)"),
+        "{}",
+        failed.reason
+    );
+    assert!(failed.retryable);
+    assert_eq!(upstream.heads().len(), 2);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_resumed_fetch_makes_only_the_attempts_left() {
+    let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]);
+    let url = upstream.url();
+    let first = Failed {
+        reason: "first".into(),
+        retryable: true,
+        retry_after: Some(Duration::ZERO),
+    };
+    let failed = fetched(&url, POLICY, Start::After(first))
+        .await
+        .unwrap_err();
+    assert_eq!(
+        failed.reason,
+        format!("HTTP 503 from {url} (after 3 attempts)")
+    );
+    assert_eq!(upstream.heads().len(), 2);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_resumed_fetch_with_no_attempts_left_returns_the_failure_it_resumed() {
+    let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
+    let first = Failed {
+        reason: "first".into(),
+        retryable: true,
+        retry_after: None,
+    };
+    let policy = RetryPolicy {
+        max_attempts: 1,
+        ..POLICY
+    };
+    assert_eq!(
+        fetched(&upstream.url(), policy, Start::After(first.clone())).await,
+        Err(first)
+    );
+    assert!(upstream.heads().is_empty());
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_failure_reason_drops_url_credentials() {
+    let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]);
+    let address = upstream.base.trim_start_matches("http://");
+    let url = format!("http://user:hunter2@{address}/cost.json?X-Amz-Signature=sig-secret#frag");
+    let failed = fetched(&url, POLICY, Start::Fresh).await.unwrap_err();
+    assert_eq!(failed.reason, format!("HTTP 404 from {}", upstream.url()));
+}
+
+#[rstest]
+#[case::unsupported_scheme("ftp://example.com/cost.json")]
+#[case::not_a_url("not a url")]
+#[case::s3_without_a_key("s3://prices")]
+#[tokio::test]
+async fn an_unfetchable_url_fails_without_a_request(#[case] url: &str) {
+    let failed = fetched(url, POLICY, Start::Fresh).await.unwrap_err();
+    assert!(
+        failed
+            .reason
+            .starts_with(&format!("error fetching {url}: ")),
+        "{}",
+        failed.reason
+    );
+    assert!(!failed.retryable);
+}
+
+fn s3(upstream: &Upstream, credentials: impl ProvideCredentials + 'static) -> S3Access {
+    S3Access::Config(Box::new(
+        litellm_s3::config(
+            "us-west-2".into(),
+            Some(&upstream.base),
+            Client::for_test(),
+            credentials,
+        )
+        .build(),
+    ))
+}
+
+async fn fetched_from_s3(
+    url: &str,
+    upstream: &Upstream,
+    credentials: impl ProvideCredentials + 'static,
+) -> Result<Fetched, Failed> {
+    fetch(
+        &Client::for_test(),
+        &s3(upstream, credentials),
+        url,
+        POLICY,
+        Start::Fresh,
+    )
+    .await
+}
+
+#[rstest]
+#[tokio::test]
+async fn an_s3_fetch_sends_a_signed_get_object_to_the_endpoint() {
+    let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
+    assert_eq!(
+        fetched_from_s3("s3://prices/maps/cost.json", &upstream, keys()).await,
+        Ok(Fetched {
+            body: b"{}".to_vec(),
+            etag: Some("\"v1\"".into()),
+        })
+    );
+    let heads = upstream.heads();
+    let [head] = heads.as_slice() else {
+        panic!("expected one request, got {heads:?}");
+    };
+    assert!(head.starts_with("GET /prices/maps/cost.json"), "{head}");
+    let authorization = header(head, "authorization").unwrap_or_default();
+    assert!(
+        authorization.starts_with("AWS4-HMAC-SHA256 Credential=AKIDMODELCOST/")
+            && authorization.contains("/us-west-2/s3/aws4_request"),
+        "{authorization}"
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn an_s3_body_that_stalls_after_the_headers_times_out() {
+    let upstream = Upstream::serve(vec![Answer::StallBody]);
+    let url = "s3://prices/cost.json";
+    let policy = RetryPolicy {
+        timeout: Duration::from_millis(200),
+        ..POLICY
+    };
+    let failed = tokio::time::timeout(
+        Duration::from_secs(2),
+        fetch(
+            &Client::for_test(),
+            &s3(&upstream, keys()),
+            url,
+            policy,
+            Start::Fresh,
+        ),
+    )
+    .await
+    .expect("the attempt timeout bounds the body read")
+    .unwrap_err();
+    assert!(
+        failed
+            .reason
+            .starts_with(&format!("error fetching {url}: ")),
+        "{}",
+        failed.reason
+    );
+    assert!(!failed.retryable);
+    assert_eq!(upstream.heads().len(), 1);
+}
+
+#[rstest]
+#[case::http(|upstream: &Upstream| (S3Access::Environment { region: None }, upstream.url()))]
+#[case::s3(|upstream: &Upstream| (s3(upstream, keys()), "s3://prices/cost.json".to_owned()))]
+#[tokio::test]
+async fn a_body_that_keeps_arriving_finishes_after_the_timeout(
+    #[case] source: fn(&Upstream) -> (S3Access, String),
+) {
+    let upstream = Upstream::serve(vec![Answer::SlowBody]);
+    let (access, url) = source(&upstream);
+    let policy = RetryPolicy {
+        timeout: Duration::from_millis(300),
+        ..POLICY
+    };
+    assert_eq!(
+        fetch(&Client::for_test(), &access, &url, policy, Start::Fresh).await,
+        Ok(Fetched {
+            body: SLOW_BODY.to_vec(),
+            etag: Some("\"v1\"".into()),
+        })
+    );
+    assert_eq!(upstream.heads().len(), 1);
+}
+
+#[rstest]
+#[case::missing_key(NOT_FOUND)]
+#[case::denied("HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")]
+#[tokio::test]
+async fn an_s3_error_fails_without_a_retry(#[case] response: &'static str) {
+    let upstream = Upstream::serve(vec![Answer::Raw(response), Answer::Raw(OK)]);
+    let url = "s3://prices/cost.json";
+    let failed = fetched_from_s3(url, &upstream, keys()).await.unwrap_err();
+    assert!(
+        failed
+            .reason
+            .starts_with(&format!("error fetching {url}: ")),
+        "{}",
+        failed.reason
+    );
+    assert!(!failed.retryable);
+    assert_eq!(upstream.heads().len(), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn an_s3_fetch_without_credentials_fails_without_a_request() {
+    let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
+    let url = "s3://prices/cost.json";
+    let failed = fetched_from_s3(url, &upstream, NoCredentials)
+        .await
+        .unwrap_err();
+    assert!(
+        failed
+            .reason
+            .starts_with(&format!("error fetching {url}: ")),
+        "{}",
+        failed.reason
+    );
+    assert!(!failed.retryable);
+    assert!(upstream.heads().is_empty());
+}
+
+const DECLARED_OVERSIZED: &str =
+    "HTTP/1.1 200 OK\r\nContent-Length: 104857601\r\nConnection: close\r\n\r\n";
+
+#[rstest]
+#[case::http_declared(Answer::Raw(DECLARED_OVERSIZED), false)]
+#[case::http_streamed(Answer::Oversized, false)]
+#[case::s3_declared(Answer::Raw(DECLARED_OVERSIZED), true)]
+#[case::s3_streamed(Answer::Oversized, true)]
+#[tokio::test]
+async fn a_body_over_100_mib_fails_without_a_retry(#[case] answer: Answer, #[case] s3: bool) {
+    let upstream = Upstream::serve(vec![answer, Answer::Raw(OK)]);
+    let failed = match s3 {
+        true => fetched_from_s3("s3://prices/cost.json", &upstream, keys()).await,
+        false => fetched(&upstream.url(), POLICY, Start::Fresh).await,
+    }
+    .unwrap_err();
+    assert!(
+        failed
+            .reason
+            .contains("response body is over the 104857600 byte limit"),
+        "{}",
+        failed.reason
+    );
+    assert!(!failed.retryable);
+    assert_eq!(upstream.heads().len(), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_body_of_exactly_100_mib_is_fetched() {
+    let upstream = Upstream::serve(vec![Answer::Raw(Box::leak(
+        format!(
+            "HTTP/1.1 200 OK\r\nContent-Length: 104857600\r\nConnection: close\r\n\r\n{}",
+            " ".repeat(100 * 1024 * 1024)
+        )
+        .into_boxed_str(),
+    ))]);
+    let fetched = fetched(&upstream.url(), POLICY, Start::Fresh)
+        .await
+        .unwrap();
+    assert_eq!(fetched.body.len(), 100 * 1024 * 1024);
+}

litellm-rust/crates/model-catalog/AGENTS.md

@@ -46,7 +46,14 @@ Every price in a row lowers to one `Price { meter, channel, when }` and its `Uni
 
 `store/` is the catalog's in-memory store and the only place that indexes model names. Its rules are in [src/store/AGENTS.md](src/store/AGENTS.md)
 
-`layers.rs`, `resolve/`, `model_info.rs` and `register.rs` hold types only. The functions over them land with their first caller
+`cost_map.rs` holds a host's inputs, the installed document and the `register_model` overlay in `register.rs`. It serves the document in the key order and alias sharing Python's `litellm.model_cost` has, without the overlay, since where a registration lands is a lookup in the published table. It never holds the live table, which belongs to the host
+
+- `CostMap` is immutable behind an `Arc`. A new document is a new `CostMap`
+- `Registrations<V>` is generic over the host's value type, because Python registrations can hold `datetime` and other values JSON cannot. The host keeps its own objects (`Py<PyDict>` in python-bridge) and does the merging
+  - Entry order is latest-registration order. `landed_on` records which registered keys landed on each row key
+- `write_pass` stops at the first `place` that returns `None` and reports `Pass::Swapped`. Retrying is the host's job
+
+`resolve/` and `model_info.rs` hold types only. The functions over them land with their first caller
 
 Catalog errors use one opaque `Error` with typed causes and JSON Pointer accessors (`path`, `key`, `field`). `Error::rejects` alone decides whether an error fails the document, rejects its row or neither, and `ErrorKind` is the stable name hosts report it by
 

litellm-rust/crates/model-catalog/Cargo.toml

@@ -11,6 +11,7 @@ schema = ["dep:schemars", "litellm-cost/schema"]
 [dependencies]
 litellm-cost = { workspace = true, features = ["serde"] }
 litellm-error.workspace = true
+litellm-tracing.workspace = true
 
 indexmap = { version = "2.14.0", features = ["serde"] }
 jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] }
@@ -19,6 +20,7 @@ schemars = { workspace = true, optional = true }
 serde.workspace = true
 serde_json = { workspace = true, features = ["preserve_order"] }
 serde_with.workspace = true
+sha1.workspace = true
 strum.workspace = true
 thiserror.workspace = true
 time.workspace = true

litellm-rust/crates/model-catalog/src/catalog.rs

@@ -12,7 +12,8 @@ use crate::store::{Rows, Store};
 use crate::wire::{EntryWire, Fields};
 
 const SAMPLE_SPEC: &str = "sample_spec";
-const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations";
+/// The top-level key of the rules block, which is never a model row
+pub const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations";
 
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub struct IntegrityLimits {
@@ -111,6 +112,11 @@ impl Rows for CatalogRows {
 }
 
 impl IntegrityLimits {
+    pub const UNCHECKED: Self = Self {
+        min_model_count: 0,
+        min_reference_ratio: 0.0,
+    };
+
     pub fn new(min_model_count: usize, min_reference_ratio: f64) -> Option<Self> {
         (0.0..=1.0).contains(&min_reference_ratio).then_some(Self {
             min_model_count,
@@ -418,6 +424,20 @@ impl Catalog {
         })
     }
 
+    /// The row under an exact model or alias key, an alias reading its owner's row
+    pub fn get(&self, key: &str) -> Option<&Fields> {
+        self.rows().fields(key)
+    }
+
+    /// A top-level value that is not a model row, such as `sample_spec`
+    pub fn section(&self, key: &str) -> Option<&Value> {
+        self.sections
+            .iter()
+            .filter(|section| is_section(&section.key))
+            .find(|section| section.key == key)
+            .map(|section| &section.value)
+    }
+
     pub fn model_count(&self) -> usize {
         self.rows().entries.len()
     }

litellm-rust/crates/model-catalog/src/cost_map.rs

@@ -0,0 +1,130 @@
+use std::sync::{Arc, LazyLock};
+
+use indexmap::IndexMap;
+use serde_json::Value;
+use sha1::{Digest, Sha1};
+
+use crate::bundled::BUNDLED_REGISTRY;
+use crate::catalog::{Catalog, FALLBACK_GENERALIZATIONS, IntegrityLimits};
+use crate::error::Error;
+
+/// Which document is installed: the `git_blob_id` of its bytes and the ETag it was served with
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Provenance {
+    pub revision: String,
+    pub etag: Option<String>,
+}
+
+/// An installed cost map document. A host builds its live table from it
+#[derive(Clone, Debug)]
+pub struct CostMap {
+    base: Arc<Base>,
+    provenance: Provenance,
+}
+
+#[derive(Debug)]
+struct Base {
+    catalog: Catalog,
+    /// Each owner's names, the owner first, for owners with at least one alias
+    groups: IndexMap<String, Vec<String>>,
+}
+
+impl Base {
+    fn new(catalog: Catalog) -> Self {
+        let groups = catalog.aliases().iter().fold(
+            IndexMap::<String, Vec<String>>::new(),
+            |mut groups, (alias, owner)| {
+                groups
+                    .entry(owner.clone())
+                    .or_insert_with(|| vec![owner.clone()])
+                    .push(alias.clone());
+                groups
+            },
+        );
+        Self { catalog, groups }
+    }
+}
+
+static BUNDLED: LazyLock<CostMap> = LazyLock::new(|| CostMap {
+    base: Arc::new(Base::new(
+        Catalog::parse(BUNDLED_REGISTRY)
+            .expect("the bundled registry parses")
+            .value,
+    )),
+    provenance: Provenance {
+        revision: git_blob_id(BUNDLED_REGISTRY),
+        etag: None,
+    },
+});
+
+impl CostMap {
+    pub fn bundled() -> Self {
+        BUNDLED.clone()
+    }
+
+    /// Reads a fetched document. One that does not parse, or that has fewer models than `limits` allow against the
+    /// bundled one, is refused. Every alias problem it has is logged
+    pub fn read(body: &[u8], etag: Option<String>, limits: IntegrityLimits) -> Result<Self, Error> {
+        let loaded = Catalog::parse(body)?;
+        loaded
+            .value
+            .check_size(BUNDLED.base.catalog.model_count(), limits)?;
+        for problem in loaded.errors.iter().filter(|error| error.is_alias()) {
+            litellm_tracing::warn!(
+                "LiteLLM: model cost map alias problem: {}",
+                litellm_error::Report::new(problem)
+            );
+        }
+        Ok(Self {
+            base: Arc::new(Base::new(loaded.value)),
+            provenance: Provenance {
+                revision: git_blob_id(body),
+                etag,
+            },
+        })
+    }
+
+    pub fn catalog(&self) -> &Catalog {
+        &self.base.catalog
+    }
+
+    pub fn provenance(&self) -> &Provenance {
+        &self.provenance
+    }
+
+    /// Every name the row under `key` is served under, the owner first. Empty for a row without aliases
+    pub fn names(&self, key: &str) -> &[String] {
+        let owner = self
+            .catalog()
+            .aliases()
+            .get(key)
+            .map_or(key, String::as_str);
+        self.base
+            .groups
+            .get(owner)
+            .map(Vec::as_slice)
+            .unwrap_or_default()
+    }
+
+    /// Each owner with its names, the owner first, in the order the document declared the aliases
+    pub fn alias_groups(&self) -> impl Iterator<Item = (&str, &[String])> {
+        self.base
+            .groups
+            .iter()
+            .map(|(owner, names)| (owner.as_str(), names.as_slice()))
+    }
+
+    /// The `fallback_generalizations` block as the document wrote it
+    pub fn fallback_generalizations(&self) -> Option<&Value> {
+        self.catalog().section(FALLBACK_GENERALIZATIONS)
+    }
+}
+
+/// `git hash-object` of the bytes, which main reports as `source_revision`
+pub fn git_blob_id(body: &[u8]) -> String {
+    let digest = Sha1::new()
+        .chain_update(format!("blob {}\0", body.len()))
+        .chain_update(body)
+        .finalize();
+    format!("{digest:x}")
+}

litellm-rust/crates/model-catalog/src/error.rs

@@ -167,6 +167,74 @@ impl Error {
     }
 }
 
+impl Error {
+    /// The document parsed but holds fewer models than the integrity limits trust
+    pub fn is_integrity(&self) -> bool {
+        match self.repr() {
+            Repr::Empty | Repr::BelowMinimum { .. } | Repr::Shrunk { .. } => true,
+            Repr::Json(_)
+            | Repr::NotAnObject
+            | Repr::Invalid(_)
+            | Repr::UnknownBlockKey
+            | Repr::UnknownTimezone { .. }
+            | Repr::AliasShadowsModel { .. }
+            | Repr::AliasTaken { .. }
+            | Repr::AliasTargetMissing { .. }
+            | Repr::AliasRowField
+            | Repr::DuplicatePrice { .. }
+            | Repr::SuffixInsideBlock { .. }
+            | Repr::TierRange
+            | Repr::TierReversed
+            | Repr::TierGap { .. }
+            | Repr::MixedTierBasis
+            | Repr::NotPricedByResults
+            | Repr::BandsFromTwoSources
+            | Repr::NotAnOffPeakRate
+            | Repr::NoOffPeakPrices
+            | Repr::NoOffPeakHours
+            | Repr::EmptyList
+            | Repr::InvalidDate { .. }
+            | Repr::MissingProvider
+            | Repr::UnreadField
+            | Repr::NoFallbackRules => false,
+        }
+    }
+
+    /// A problem with a row's aliases, which the served map skips
+    pub fn is_alias(&self) -> bool {
+        match self.repr() {
+            Repr::AliasShadowsModel { .. }
+            | Repr::AliasTaken { .. }
+            | Repr::AliasTargetMissing { .. }
+            | Repr::AliasRowField => true,
+            Repr::Invalid(_) => matches!(self.field().as_deref(), Some("aliases" | "alias_of")),
+            Repr::Json(_)
+            | Repr::Empty
+            | Repr::BelowMinimum { .. }
+            | Repr::Shrunk { .. }
+            | Repr::NotAnObject
+            | Repr::UnknownBlockKey
+            | Repr::UnknownTimezone { .. }
+            | Repr::DuplicatePrice { .. }
+            | Repr::SuffixInsideBlock { .. }
+            | Repr::TierRange
+            | Repr::TierReversed
+            | Repr::TierGap { .. }
+            | Repr::MixedTierBasis
+            | Repr::NotPricedByResults
+            | Repr::BandsFromTwoSources
+            | Repr::NotAnOffPeakRate
+            | Repr::NoOffPeakPrices
+            | Repr::NoOffPeakHours
+            | Repr::EmptyList
+            | Repr::InvalidDate { .. }
+            | Repr::MissingProvider
+            | Repr::UnreadField
+            | Repr::NoFallbackRules => false,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct Loaded<T> {
     pub value: T,

litellm-rust/crates/model-catalog/src/layers.rs

@@ -1,32 +0,0 @@
-use std::sync::Arc;
-
-use indexmap::IndexMap;
-use serde_json::Value;
-
-use crate::catalog::Catalog;
-use crate::register::Registrations;
-use crate::wire::Fields;
-
-/// One field of a fetched row over the layer below, RFC 7386. An absent field inherits, JSON `null` is `Remove`
-#[derive(Clone, Debug, PartialEq)]
-pub enum Patch {
-    Set(Value),
-    Remove,
-}
-
-/// The four layers served rows are composed from, lowest first
-#[derive(Clone, Debug)]
-#[expect(
-    dead_code,
-    reason = "02 and 07 add the functions that read these layers"
-)]
-pub struct Layers {
-    /// 0, the wheel's embedded map
-    bundled: Arc<Catalog>,
-    /// 1, the last installed fetch, rejected rows left out
-    fetched: Option<Arc<Catalog>>,
-    /// 2, `register(persist=False)`, cleared by every install and refilled by the router replay
-    transient: IndexMap<String, Fields>,
-    /// 3, `register(persist=True)`, applied over every install
-    persisted: Registrations,
-}

litellm-rust/crates/model-catalog/src/lib.rs

@@ -17,6 +17,9 @@
 //!   failing the document.
 //! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known
 //!   families. They are data in the file, not entries.
+//! - `CostMap`: an installed document with its provenance, which a host builds its live table from.
+//!   `Registrations` is the ordered `register_model` overlay over the host's own values, and
+//!   `write_pass` writes registrations into a host's table through `RegistrationTable`.
 //! - Validation: `Catalog::check_size` checks a catalog against `IntegrityLimits` (a remote copy
 //!   that lost most of its models is refused), and `Catalog::row_errors` validates one row.
 //!
@@ -37,9 +40,9 @@ macro_rules_attribute::attribute_alias! {
 
 mod bundled;
 mod catalog;
+mod cost_map;
 mod error;
 mod fallback;
-mod layers;
 mod lower;
 mod model;
 mod model_info;
@@ -51,9 +54,9 @@ pub mod wire;
 
 pub use bundled::*;
 pub use catalog::*;
+pub use cost_map::*;
 pub use error::*;
 pub use fallback::*;
-pub use layers::*;
 pub use model::*;
 pub use model_info::*;
 pub use register::*;

litellm-rust/crates/model-catalog/src/register.rs

@@ -1,26 +1,252 @@
-use std::collections::{BTreeSet, HashMap};
+use std::collections::BTreeMap;
 
 use indexmap::IndexMap;
 
-use crate::wire::Fields;
-
-/// Where one `register` call writes and which row it inherits from
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct RegistrationTarget {
-    /// Always the caller's key, never a case variant of it
-    pub write: String,
-    /// The first candidate that is an exact key whose row passes `provider_matches`
-    pub inherit_from: Option<String>,
+/// `register_model(persist_across_reloads=True)` calls, replayed above every installed document
+#[derive(Clone, Debug, PartialEq)]
+pub struct Registrations<V> {
+    /// Each registered key's merged value, ordered by its latest registration
+    entries: IndexMap<String, V>,
+    /// The registered keys whose write landed on each row key, in landing order
+    landed_on: IndexMap<String, Vec<String>>,
 }
 
-/// Python's inheritance candidates in main's order, `get_model_info`'s then `_resolve_builtin_model_cost_entry`'s
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct RegistrationCandidates(pub Vec<String>);
-
-/// Layer 3, and the fields `register` last wrote, which litellm owns again
-#[derive(Clone, Debug, Default)]
-#[expect(dead_code, reason = "07 adds the functions that read registrations")]
-pub struct Registrations {
-    persisted: IndexMap<String, Fields>,
-    registered_fields: HashMap<String, BTreeSet<String>>,
+impl<V> Default for Registrations<V> {
+    fn default() -> Self {
+        Self {
+            entries: IndexMap::new(),
+            landed_on: IndexMap::new(),
+        }
+    }
+}
+
+/// One registration written to a host's table: the key it was registered under and the row key it landed on
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Landed {
+    pub registered: String,
+    pub written: String,
+}
+
+impl<V> Registrations<V> {
+    /// Keeps `value` as the key's registration, already merged with `get(key)` by the host, and makes it the latest
+    pub fn register(&mut self, key: &str, value: V) {
+        self.entries.shift_remove(key);
+        self.entries.insert(key.to_owned(), value);
+    }
+
+    pub fn get(&self, key: &str) -> Option<&V> {
+        self.entries.get(key)
+    }
+
+    pub fn record(&mut self, landed: &[Landed]) {
+        for Landed {
+            registered,
+            written,
+        } in landed
+        {
+            let names = self.landed_on.entry(written.clone()).or_default();
+            if !names.contains(registered) {
+                names.push(registered.clone());
+            }
+        }
+    }
+
+    /// Every registration, oldest first
+    pub fn entries(&self) -> impl Iterator<Item = (&str, &V)> {
+        self.entries
+            .iter()
+            .map(|(key, value)| (key.as_str(), value))
+    }
+
+    /// The registrations that apply to the row under `key`, oldest first: `key`'s own and every one that landed on it
+    pub fn over<'a>(&'a self, key: &str) -> Vec<&'a V> {
+        let landed = self.landed_on.get(key).into_iter().flatten();
+        let by_order: BTreeMap<usize, &V> = landed
+            .map(String::as_str)
+            .chain([key])
+            .filter_map(|name| self.entries.get_full(name))
+            .map(|(order, _, value)| (order, value))
+            .collect();
+        by_order.into_values().collect()
+    }
+
+    pub fn map<W>(&self, mut convert: impl FnMut(&V) -> W) -> Registrations<W> {
+        Registrations {
+            entries: self
+                .entries
+                .iter()
+                .map(|(key, value)| (key.clone(), convert(value)))
+                .collect(),
+            landed_on: self.landed_on.clone(),
+        }
+    }
+}
+
+/// A host's live table that registrations are written into
+pub trait RegistrationTable {
+    type Registration;
+    type Error;
+
+    /// Writes `registration` where it lands, `None` when this table is no longer the one the host publishes
+    fn place(
+        &mut self,
+        key: &str,
+        registration: &Self::Registration,
+    ) -> Result<Option<Placed>, Self::Error>;
+}
+
+/// Where a registration was written and whether a built-in row matched it
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Placed {
+    pub key: String,
+    pub matched_builtin: bool,
+}
+
+#[derive(Debug, PartialEq)]
+pub enum Pass<'a, R> {
+    /// Every entry was written. `unmatched` are those no built-in row matched
+    Finished {
+        landed: Vec<Landed>,
+        unmatched: Vec<(&'a str, &'a R)>,
+    },
+    /// The table stopped being the published one before every entry was written
+    Swapped { landed: Vec<Landed> },
+}
+
+impl<R> Pass<'_, R> {
+    pub fn landed(&self) -> &[Landed] {
+        match self {
+            Pass::Finished { landed, .. } | Pass::Swapped { landed } => landed,
+        }
+    }
+}
+
+/// Writes `entries` into `table` in order
+pub fn write_pass<'a, T: RegistrationTable>(
+    table: &mut T,
+    entries: &'a [(String, T::Registration)],
+) -> Result<Pass<'a, T::Registration>, T::Error> {
+    let mut landed = Vec::with_capacity(entries.len());
+    let mut unmatched = Vec::new();
+    for (key, registration) in entries {
+        let Some(written) = table.place(key, registration)? else {
+            return Ok(Pass::Swapped { landed });
+        };
+        if !written.matched_builtin {
+            unmatched.push((key.as_str(), registration));
+        }
+        landed.push(Landed {
+            registered: key.clone(),
+            written: written.key,
+        });
+    }
+    Ok(Pass::Finished { landed, unmatched })
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    fn landed(registered: &str, written: &str) -> Landed {
+        Landed {
+            registered: registered.to_owned(),
+            written: written.to_owned(),
+        }
+    }
+
+    #[rstest]
+    fn reregistering_a_key_makes_it_the_latest() {
+        let mut registrations = Registrations::default();
+        registrations.register("a", 1);
+        registrations.register("b", 2);
+        registrations.register("a", 3);
+        assert_eq!(
+            registrations.entries().collect::<Vec<_>>(),
+            [("b", &2), ("a", &3)]
+        );
+    }
+
+    #[rstest]
+    fn over_applies_names_landed_on_the_row_in_registration_order() {
+        let mut registrations = Registrations::default();
+        registrations.register("openai/o3-pro", 1);
+        registrations.register("o3-pro", 2);
+        registrations.register("other", 3);
+        registrations.record(&[
+            landed("openai/o3-pro", "o3-pro"),
+            landed("o3-pro", "o3-pro"),
+        ]);
+        registrations.register("openai/o3-pro", 4);
+        assert_eq!(registrations.over("o3-pro"), [&2, &4]);
+    }
+
+    #[rstest]
+    fn over_applies_a_keys_own_registration_that_landed_on_another_row() {
+        let mut registrations = Registrations::default();
+        registrations.register("openai/o3-pro", 1);
+        registrations.record(&[landed("openai/o3-pro", "o3-pro")]);
+        assert_eq!(
+            (
+                registrations.over("openai/o3-pro"),
+                registrations.over("o3-pro")
+            ),
+            (vec![&1], vec![&1])
+        );
+    }
+
+    /// A table published until `swaps_after` writes, whose rows match a built-in only when the key starts with `b`
+    struct Table {
+        writes: Vec<(String, i32)>,
+        swaps_after: usize,
+    }
+
+    impl RegistrationTable for Table {
+        type Registration = i32;
+        type Error = ();
+
+        fn place(&mut self, key: &str, registration: &i32) -> Result<Option<Placed>, ()> {
+            if self.writes.len() >= self.swaps_after {
+                return Ok(None);
+            }
+            let written = key.to_uppercase();
+            self.writes.push((written.clone(), *registration));
+            Ok(Some(Placed {
+                matched_builtin: written.starts_with('B'),
+                key: written,
+            }))
+        }
+    }
+
+    #[rstest]
+    #[case::published_throughout(
+        usize::MAX,
+        Pass::Finished { landed: vec![landed("a", "A"), landed("b", "B")], unmatched: vec![("a", &1)] },
+        &[("A", 1), ("B", 2)],
+    )]
+    #[case::swapped_after_the_first_write(
+        1,
+        Pass::Swapped { landed: vec![landed("a", "A")] },
+        &[("A", 1)],
+    )]
+    fn a_pass_writes_entries_in_order_until_the_table_is_swapped(
+        #[case] swaps_after: usize,
+        #[case] expected: Pass<'static, i32>,
+        #[case] writes: &[(&str, i32)],
+    ) {
+        let entries = [("a".to_owned(), 1), ("b".to_owned(), 2)];
+        let mut table = Table {
+            writes: Vec::new(),
+            swaps_after,
+        };
+        assert_eq!(write_pass(&mut table, &entries), Ok(expected));
+        assert_eq!(
+            table.writes,
+            writes
+                .iter()
+                .map(|(key, value)| ((*key).to_owned(), *value))
+                .collect::<Vec<_>>()
+        );
+    }
 }

litellm-rust/crates/model-catalog/src/store/AGENTS.md

@@ -24,7 +24,7 @@
 | Source | Stamp | Blind spot |
 |---|---|---|
 | `Catalog`'s `CatalogRows` | `()`, every change goes through `Store::edit`, which drops the index | none |
-| python-bridge's bound `litellm.model_cost` | the dict object and its `len` | a delete plus an add between two reads, or a `litellm_provider` rewrite on an existing key |
+| python-bridge's bound `litellm.model_cost` | the dict object, its `len` and its last key | deletes plus adds between two reads that put the old last key back last, or a `litellm_provider` rewrite on an existing key |
 
 A source with a blind spot calls `Store::invalidate` when it knows keys changed. Until then a key added in the blind spot is not found, and a key moved to another provider stays listed under the old one. Removed keys are still never returned
 

litellm-rust/crates/model-catalog/tests/cost_map.rs

@@ -0,0 +1,53 @@
+use litellm_model_catalog::{
+    BUNDLED_REGISTRY, CostMap, ErrorKind, IntegrityLimits, Provenance, git_blob_id,
+};
+use rstest::rstest;
+use serde_json::json;
+
+fn any_size() -> IntegrityLimits {
+    IntegrityLimits::new(0, 0.0).unwrap()
+}
+
+#[rstest]
+#[case::not_json(b"{".as_slice(), false, ErrorKind::Json)]
+#[case::empty(b"{}".as_slice(), true, ErrorKind::Empty)]
+fn a_refused_document_says_whether_it_failed_the_integrity_limits(
+    #[case] body: &[u8],
+    #[case] integrity: bool,
+    #[case] kind: ErrorKind,
+) {
+    let error = CostMap::read(body, None, any_size()).unwrap_err();
+    assert_eq!((error.is_integrity(), error.kind()), (integrity, kind));
+}
+
+#[rstest]
+fn provenance_is_the_read_bytes_and_etag() {
+    let body = br#"{"m":{"litellm_provider":"p"}}"#;
+    let cost_map = CostMap::read(body, Some("\"v1\"".into()), any_size()).unwrap();
+    assert_eq!(
+        cost_map.provenance(),
+        &Provenance {
+            revision: "eaa747f210a1527d5ea7138be0f768c47578ca68".to_owned(),
+            etag: Some("\"v1\"".into()),
+        }
+    );
+    assert_eq!(
+        CostMap::bundled().provenance().revision,
+        git_blob_id(BUNDLED_REGISTRY)
+    );
+}
+
+#[rstest]
+fn every_name_of_an_aliased_row_lists_its_owner_first(
+    #[values("owner", "short", "owner-latest")] key: &str,
+) {
+    let body = json!({
+        "owner": {"litellm_provider": "p", "aliases": ["short"]},
+        "owner-latest": {"alias_of": "owner"},
+        "plain": {"litellm_provider": "p"},
+    })
+    .to_string();
+    let cost_map = CostMap::read(body.as_bytes(), None, any_size()).unwrap();
+    assert_eq!(cost_map.names(key), ["owner", "short", "owner-latest"]);
+    assert!(cost_map.names("plain").is_empty());
+}

litellm-rust/crates/python-bridge/Cargo.toml

@@ -60,6 +60,7 @@ litellm-router-types.workspace = true
 litellm-router-core.workspace = true
 litellm-host-python.workspace = true
 litellm-model-catalog.workspace = true
+litellm-model-catalog-refresh.workspace = true
 litellm-token-counter = { path = "../token-counter", default-features = false }
 indexmap = "2.14.0"
 pyo3 = { workspace = true, features = ["indexmap"] }

litellm-rust/crates/python-bridge/src/http.rs

@@ -108,13 +108,35 @@ pub(crate) fn call_config(
             Some(snapshot) => configured(&snapshot)?,
             None => HttpSettingsLayer::default(),
         },
-    ])
-    .without_missing_files(&|path: &Path| path.exists());
-    let resolution = Resolution::from(&settings);
+    ]);
+    Ok(resolved(py, settings))
+}
+
+/// The client config for a value of `litellm.rust_bridge.host.settings.http_settings()` the caller handed in, under
+/// the process environment. `None` reads the environment alone
+pub(crate) fn settings_config(
+    py: Python<'_>,
+    http_settings: Option<&Bound<'_, PyAny>>,
+) -> PyResult<HttpClientConfig> {
+    let configured = match http_settings {
+        Some(value) => configured(&PythonSettings::Http.snapshot(value.clone()))?,
+        None => HttpSettingsLayer::default(),
+    };
+    let settings = HttpSettings::from_layers([
+        HttpSettingsLayer::from_environment(&ProcessEnvironment),
+        configured,
+    ]);
+    Ok(resolved(py, settings))
+}
+
+/// The client config `settings` resolve to, warning once per process about each setting the client cannot apply
+fn resolved(py: Python<'_>, settings: HttpSettings) -> HttpClientConfig {
+    let resolution =
+        Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists()));
     for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) {
         crate::logger::capture(py).scope(|| litellm_tracing::warn!("{unsupported}"));
     }
-    Ok(resolution.config)
+    resolution.config
 }
 
 pub(crate) fn provider_client(

litellm-rust/crates/python-bridge/src/lib.rs

@@ -9,6 +9,7 @@ mod execution;
 mod http;
 mod logger;
 mod marshal;
+mod model_cost;
 mod model_registry;
 mod preflight;
 mod python_settings;
@@ -36,6 +37,12 @@ mod _native {
     #[pymodule_export]
     use crate::logger::NativeDiagnosticProcessor;
     #[pymodule_export]
+    use crate::model_cost::{
+        NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched, NativeInstalled,
+        NativeModelCostMap, NativeModelCostMapState, NativeProvenance, NativeRegistered,
+        NativeRejected, NativeTableSwapped,
+    };
+    #[pymodule_export]
     use crate::model_registry::bundled_model_cost_map;
     #[pymodule_export]
     use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states};
@@ -166,6 +173,17 @@ mod tests {
                 "NativeDatabase",
                 "trace_encode_error",
                 "bundled_model_cost_map",
+                "NativeModelCostMap",
+                "NativeModelCostMapState",
+                "NativeInstalled",
+                "NativeRejected",
+                "NativeRegistered",
+                "NativeTableSwapped",
+                "NativeProvenance",
+                "NativeCostMapSource",
+                "NativeFetchPolicy",
+                "NativeFetched",
+                "NativeFetchFailed",
                 "TokenCounter",
                 "Tokenizer",
                 "RouterCoreHandle",

litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md

@@ -0,0 +1,16 @@
+# Model cost bridge
+
+- `NativeModelCostMap` holds inputs only, the installed `CostMap`, kept `Registrations` and the name index. Python's `litellm.model_cost` is the only copy of the live rows
+- Every call that reads or writes rows takes the dict as an argument. Never store rows or build a second table here
+- `LiveDict` holds a reference to that same dict and reads it under the GIL on every query, so a Python write is seen on the next read
+  - It needs the GIL, so Rust code running without the GIL cannot use it
+  - Its stamp is the dict's address, length and last key. [model-catalog/src/store/AGENTS.md](../../../model-catalog/src/store/AGENTS.md) says what that misses
+- `host.rs` is the only code that calls Python, by name through `MODULE`. Calls happen with the GIL held and no `NativeModelCostMap` lock held
+  - `resolve` reaches `find_key` on this same object, so holding a lock across a host call deadlocks
+  - Test `a_host_lookup_that_reaches_find_key_on_the_same_map_finishes` guards this
+- Rows and registrations cross as Python objects, never JSON, because router `model_info` holds `datetime`
+- `table::build` serves the document as main's `_expand_model_aliases` did, so all names of a model share one row object made by `host.entry`
+- `install` touches no state until `CostMap::read` succeeds. `NativeRejected` means nothing changed
+- `register` keeps a registration only after its pass finished, so an exception in a host call keeps nothing
+- `rows=None` writes each entry into whichever table is published at that moment and never reports a swap
+- `fetch_blocking` runs on a throwaway runtime with an unpooled client so an importing process can still fork. `fetch` is for workers after the fork

litellm-rust/crates/python-bridge/src/model_cost/fetch.rs

@@ -0,0 +1,239 @@
+use std::time::Duration;
+
+use litellm_http::{ClientVariant, HttpClientConfig};
+use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start};
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::PyBytes;
+
+/// A cost map URL and how to connect to it. An `s3://` URL reaches S3 as an AWS SDK finds it in the environment,
+/// with LiteLLM's `AWS_REGION_NAME` winning over the region it finds
+#[pyclass(
+    frozen,
+    name = "NativeCostMapSource",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeCostMapSource {
+    url: String,
+    http: HttpClientConfig,
+}
+
+#[pymethods]
+impl NativeCostMapSource {
+    /// `http_settings` is `litellm.rust_bridge.host.settings.http_settings()`, or `None` for the environment alone
+    #[new]
+    #[pyo3(signature = (url, *, http_settings))]
+    fn new(py: Python<'_>, url: String, http_settings: Option<Bound<'_, PyAny>>) -> PyResult<Self> {
+        Ok(Self {
+            url,
+            http: crate::http::settings_config(py, http_settings.as_ref())?,
+        })
+    }
+
+    /// Fetches on the shared runtime with a pooled client
+    fn fetch<'py>(
+        &self,
+        py: Python<'py>,
+        policy: PyRef<'_, NativeFetchPolicy>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let client = crate::http::pool()
+            .client(&self.http, ClientVariant::Provider)
+            .map_err(crate::http::client_error)?;
+        let url = self.url.clone();
+        let policy = policy.0;
+        crate::execution::run_async_value(py, async move {
+            Ok(Outcome(
+                litellm_model_catalog_refresh::fetch(&client, &s3(), &url, policy, Start::Fresh)
+                    .await,
+            ))
+        })
+    }
+
+    /// Fetches on a throwaway runtime with a client of its own. `resume_from` continues a fetch whose first attempt
+    /// failed retryably
+    #[pyo3(signature = (policy, resume_from = None))]
+    fn fetch_blocking(
+        &self,
+        py: Python<'_>,
+        policy: PyRef<'_, NativeFetchPolicy>,
+        resume_from: Option<PyRef<'_, NativeFetchFailed>>,
+    ) -> PyResult<Py<PyAny>> {
+        let start = resume_from.map_or(Start::Fresh, |failed| Start::After(failed.0.clone()));
+        let (config, url, policy) = (self.http.clone(), self.url.clone(), policy.0);
+        let logger = crate::logger::capture(py);
+        let fetched = litellm_host_python::release_gil(py, move || {
+            logger.scope(|| {
+                let client = crate::http::pool()
+                    .unpooled(&config, ClientVariant::Provider)
+                    .map_err(crate::http::client_error)?;
+                litellm_host_python::run_isolated(litellm_model_catalog_refresh::fetch(
+                    &client,
+                    &s3(),
+                    &url,
+                    policy,
+                    start,
+                ))
+                .map_err(PyErr::from)
+            })
+        })?;
+        Outcome(fetched).into_pyobject(py).map(Bound::unbind)
+    }
+}
+
+fn s3() -> S3Access {
+    S3Access::Environment {
+        region: std::env::var("AWS_REGION_NAME").ok(),
+    }
+}
+
+#[pyclass(
+    frozen,
+    name = "NativeFetchPolicy",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeFetchPolicy(RetryPolicy);
+
+#[pymethods]
+impl NativeFetchPolicy {
+    #[new]
+    /// Retries wait an exponential backoff with jitter from `min_delay_seconds` up to `max_wait_seconds`, or the
+    /// server's `Retry-After` up to that cap
+    #[pyo3(signature = (*, max_attempts, max_wait_seconds, timeout_seconds, min_delay_seconds))]
+    fn new(
+        max_attempts: u32,
+        max_wait_seconds: f64,
+        timeout_seconds: f64,
+        min_delay_seconds: f64,
+    ) -> PyResult<Self> {
+        if max_attempts == 0 {
+            return Err(PyValueError::new_err("max_attempts must be at least 1"));
+        }
+        let seconds = |name: &str, value: f64| {
+            Duration::try_from_secs_f64(value).map_err(|_| {
+                PyValueError::new_err(format!("{name} must be a non-negative number of seconds"))
+            })
+        };
+        let min_delay = seconds("min_delay_seconds", min_delay_seconds)?;
+        let max_delay = seconds("max_wait_seconds", max_wait_seconds)?;
+        if min_delay > max_delay {
+            return Err(PyValueError::new_err(
+                "min_delay_seconds must not exceed max_wait_seconds",
+            ));
+        }
+        Ok(Self(RetryPolicy {
+            max_attempts,
+            timeout: seconds("timeout_seconds", timeout_seconds)?,
+            min_delay,
+            max_delay,
+        }))
+    }
+}
+
+#[pyclass(frozen, name = "NativeFetched", module = "litellm.rust_bridge._native")]
+pub(crate) struct NativeFetched {
+    #[pyo3(get)]
+    body: Py<PyBytes>,
+    #[pyo3(get)]
+    etag: Option<String>,
+}
+
+#[pyclass(
+    frozen,
+    skip_from_py_object,
+    name = "NativeFetchFailed",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeFetchFailed(Failed);
+
+#[pymethods]
+impl NativeFetchFailed {
+    #[getter]
+    fn reason(&self) -> &str {
+        &self.0.reason
+    }
+
+    /// Another attempt could succeed. Only the attempt limit ended the fetch
+    #[getter]
+    fn retryable(&self) -> bool {
+        self.0.retryable
+    }
+
+    #[getter]
+    fn retry_after(&self) -> Option<f64> {
+        self.0.retry_after.map(|wait| wait.as_secs_f64())
+    }
+}
+
+struct Outcome(Result<Fetched, Failed>);
+
+impl<'py> IntoPyObject<'py> for Outcome {
+    type Target = PyAny;
+    type Output = Bound<'py, PyAny>;
+    type Error = PyErr;
+
+    fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
+        match self.0 {
+            Ok(fetched) => Bound::new(
+                py,
+                NativeFetched {
+                    body: PyBytes::new(py, &fetched.body).unbind(),
+                    etag: fetched.etag,
+                },
+            )
+            .map(Bound::into_any),
+            Err(failed) => Bound::new(py, NativeFetchFailed(failed)).map(Bound::into_any),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::{BufRead, BufReader, Write};
+    use std::net::TcpListener;
+
+    use rstest::rstest;
+
+    use super::*;
+
+    #[rstest]
+    #[case::equal_bounds(2.0, 2.0, true)]
+    #[case::min_below_max(2.0, 30.0, true)]
+    #[case::min_above_max(31.0, 30.0, false)]
+    fn a_policy_accepts_a_min_delay_only_up_to_the_max_wait(
+        #[case] min_delay_seconds: f64,
+        #[case] max_wait_seconds: f64,
+        #[case] accepted: bool,
+    ) {
+        Python::initialize();
+        let policy = NativeFetchPolicy::new(3, max_wait_seconds, 5.0, min_delay_seconds);
+        assert_eq!(policy.is_ok(), accepted);
+    }
+
+    #[rstest]
+    fn a_blocking_fetch_hands_python_the_document_without_starting_the_shared_runtime() {
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let url = format!("http://{}/cost.json", listener.local_addr().unwrap());
+        std::thread::spawn(move || {
+            let mut stream = listener.incoming().next().unwrap().unwrap();
+            let mut reader = BufReader::new(stream.try_clone().unwrap());
+            let mut line = String::new();
+            while reader.read_line(&mut line).unwrap() > 2 {
+                line.clear();
+            }
+            stream
+                .write_all(b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}")
+                .unwrap();
+        });
+        Python::initialize();
+        Python::attach(|py| {
+            let source = NativeCostMapSource::new(py, url, None).unwrap();
+            let policy =
+                Bound::new(py, NativeFetchPolicy::new(3, 30.0, 5.0, 2.0).unwrap()).unwrap();
+            let outcome = source.fetch_blocking(py, policy.borrow(), None).unwrap();
+            let fetched = outcome.bind(py).cast::<NativeFetched>().unwrap().get();
+            assert_eq!(fetched.body.bind(py).as_bytes(), b"{}");
+            assert_eq!(fetched.etag.as_deref(), Some("\"v1\""));
+        });
+        assert!(!litellm_host_python::runtime_started());
+    }
+}

litellm-rust/crates/python-bridge/src/model_cost/host.rs

@@ -0,0 +1,81 @@
+//! Calls into `litellm.rust_bridge.host.model_cost`
+
+use pyo3::call::PyCallArgs;
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+
+pub(crate) const MODULE: &str = "litellm.rust_bridge.host.model_cost";
+
+/// Where a registration lands in the published table and the row it starts from
+#[derive(FromPyObject)]
+pub(crate) struct Landing<'py> {
+    #[pyo3(attribute)]
+    pub(crate) key: String,
+    #[pyo3(attribute)]
+    pub(crate) base: Bound<'py, PyAny>,
+    #[pyo3(attribute)]
+    pub(crate) matched_builtin: bool,
+}
+
+fn call<'py>(
+    py: Python<'py>,
+    name: &str,
+    args: impl PyCallArgs<'py>,
+) -> PyResult<Bound<'py, PyAny>> {
+    py.import(MODULE)?.getattr(name)?.call1(args)
+}
+
+/// The row object every name in `names` shares, the owner first
+pub(crate) fn entry<'py>(
+    fields: Bound<'py, PyAny>,
+    names: &[String],
+) -> PyResult<Bound<'py, PyAny>> {
+    let py = fields.py();
+    call(py, "entry", (fields, PyTuple::new(py, names)?))
+}
+
+/// Where `registration` lands in `rows`, `None` when `rows` is no longer the published table. `rows` of `None` is the
+/// published table
+pub(crate) fn resolve<'py>(
+    rows: Option<&Bound<'py, PyDict>>,
+    key: &str,
+    registration: &Bound<'py, PyAny>,
+    custom_llm_provider: Option<&str>,
+) -> PyResult<Option<Landing<'py>>> {
+    call(
+        registration.py(),
+        "resolve",
+        (rows, key, registration, custom_llm_provider),
+    )?
+    .extract()
+}
+
+/// The row `registration` writes over `current`
+pub(crate) fn merged<'py>(
+    current: Option<&Bound<'py, PyAny>>,
+    landing: &Landing<'py>,
+    registration: &Bound<'py, PyAny>,
+) -> PyResult<Bound<'py, PyAny>> {
+    call(
+        registration.py(),
+        "merged",
+        (current, &landing.base, registration),
+    )
+}
+
+/// The registration kept for a key: `registration` merged into the one kept before it
+pub(crate) fn kept<'py>(
+    previous: Option<&Bound<'py, PyDict>>,
+    registration: &Bound<'py, PyAny>,
+) -> PyResult<Bound<'py, PyDict>> {
+    Ok(call(registration.py(), "kept", (previous, registration))?.cast_into::<PyDict>()?)
+}
+
+pub(crate) fn published(py: Python<'_>) -> PyResult<Bound<'_, PyDict>> {
+    Ok(call(py, "published", ())?.cast_into::<PyDict>()?)
+}
+
+/// Writes `written` into `rows`, the published table, through the host, which tracks the table's version
+pub(crate) fn write(rows: &Bound<'_, PyDict>, written: &Bound<'_, PyDict>) -> PyResult<()> {
+    call(rows.py(), "write", (rows, written)).map(drop)
+}

litellm-rust/crates/python-bridge/src/model_cost/live.rs

@@ -0,0 +1,176 @@
+use std::borrow::Cow;
+
+use litellm_model_catalog::Rows;
+use pyo3::intern;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+
+/// The dict Python last handed in as `litellm.model_cost`, read under the GIL on every query
+pub(crate) struct LiveDict(Py<PyDict>);
+
+/// One row exactly as the dict holds it
+pub(crate) struct LiveRow(Py<PyAny>);
+
+impl Clone for LiveRow {
+    fn clone(&self) -> Self {
+        Python::attach(|py| Self(self.0.clone_ref(py)))
+    }
+}
+
+impl LiveDict {
+    pub(crate) fn new(dict: &Bound<'_, PyDict>) -> Self {
+        Self(dict.clone().unbind())
+    }
+
+    pub(crate) fn is(&self, dict: &Bound<'_, PyDict>) -> bool {
+        self.0.bind(dict.py()).is(dict)
+    }
+}
+
+impl Rows for LiveDict {
+    type Row = LiveRow;
+    /// The dict's address, length and last key
+    type Stamp = (usize, usize, Option<String>);
+
+    fn fold(name: &str) -> String {
+        name.to_lowercase()
+    }
+
+    fn stamp(&self) -> Self::Stamp {
+        Python::attach(|py| {
+            let rows = self.0.bind(py);
+            let last = rows
+                .call_method0(intern!(py, "__reversed__"))
+                .and_then(|keys| keys.try_iter())
+                .ok()
+                .and_then(|mut keys| keys.next())
+                .and_then(Result::ok)
+                .and_then(|key| key.extract::<String>().ok());
+            (self.0.as_ptr() as usize, rows.len(), last)
+        })
+    }
+
+    fn contains(&self, key: &str) -> bool {
+        Python::attach(|py| self.0.bind(py).contains(key).unwrap_or(false))
+    }
+
+    fn row(&self, key: &str) -> Option<Cow<'_, LiveRow>> {
+        Python::attach(|py| {
+            let row = self.0.bind(py).get_item(key).ok().flatten()?;
+            Some(Cow::Owned(LiveRow(row.unbind())))
+        })
+    }
+
+    /// Groups each key by its row's `litellm_provider` when that is a string
+    fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) {
+        Python::attach(|py| {
+            for (key, row) in self.0.bind(py).iter() {
+                let Ok(key) = key.extract::<String>() else {
+                    continue;
+                };
+                let provider = row
+                    .cast::<PyDict>()
+                    .ok()
+                    .and_then(|row| row.get_item("litellm_provider").ok().flatten())
+                    .and_then(|provider| provider.extract::<String>().ok());
+                visit(&key, provider.as_deref());
+            }
+        });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use litellm_model_catalog::Store;
+    use pyo3::ffi::c_str;
+    use rstest::rstest;
+
+    use super::*;
+
+    fn dict<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
+        py.eval(source, None, None)
+            .unwrap()
+            .cast_into::<PyDict>()
+            .unwrap()
+    }
+
+    #[rstest]
+    fn a_key_added_to_the_dict_is_found_on_the_next_query() {
+        Python::initialize();
+        Python::attach(|py| {
+            let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}"));
+            let store = Store::new(LiveDict::new(&rows));
+            assert_eq!(store.find_key("GPT-4O"), Some("gpt-4o".to_owned()));
+            rows.set_item("Claude-X", PyDict::new(py)).unwrap();
+            assert_eq!(store.find_key("claude-x"), Some("Claude-X".to_owned()));
+        });
+    }
+
+    #[rstest]
+    fn a_removed_key_is_never_returned() {
+        Python::initialize();
+        Python::attach(|py| {
+            let rows = dict(
+                py,
+                c_str!("{'gpt-4o': {'litellm_provider': 'openai'}, 'o3': {}}"),
+            );
+            let store = Store::new(LiveDict::new(&rows));
+            assert_eq!(store.keys_in("openai"), ["gpt-4o"]);
+            rows.del_item("gpt-4o").unwrap();
+            rows.set_item("o4", PyDict::new(py)).unwrap();
+            assert_eq!(store.find_key("GPT-4O"), None);
+            assert!(store.keys_in("openai").is_empty());
+        });
+    }
+
+    #[rstest]
+    fn a_key_swapped_in_for_a_removed_one_is_found() {
+        Python::initialize();
+        Python::attach(|py| {
+            let rows = dict(py, c_str!("{'gpt-4o': {}, 'Lit-Gone': {}}"));
+            let store = Store::new(LiveDict::new(&rows));
+            assert_eq!(store.find_key("lit-gone"), Some("Lit-Gone".to_owned()));
+            rows.del_item("Lit-Gone").unwrap();
+            rows.set_item("Lit-New", PyDict::new(py)).unwrap();
+            assert_eq!(store.find_key("lit-new"), Some("Lit-New".to_owned()));
+            assert_eq!(store.find_key("lit-gone"), None);
+        });
+    }
+
+    #[rstest]
+    fn two_keys_folding_alike_resolve_to_the_later_one() {
+        Python::initialize();
+        Python::attach(|py| {
+            let rows = dict(py, c_str!("{'Model': {}, 'MODEL': {}}"));
+            assert_eq!(
+                Store::new(LiveDict::new(&rows)).find_key("model"),
+                Some("MODEL".to_owned())
+            );
+        });
+    }
+
+    #[rstest]
+    fn rows_are_the_objects_the_dict_holds() {
+        Python::initialize();
+        Python::attach(|py| {
+            let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}"));
+            let store = Store::new(LiveDict::new(&rows));
+            let row = store.get("gpt-4o").unwrap().into_owned();
+            assert!(row.0.bind(py).is(rows.get_item("gpt-4o").unwrap().unwrap()));
+        });
+    }
+
+    #[rstest]
+    fn providers_group_only_string_providers() {
+        Python::initialize();
+        Python::attach(|py| {
+            let rows = dict(
+                py,
+                c_str!(
+                    "{'a': {'litellm_provider': 'p'}, 'b': {'litellm_provider': None}, 'c': 1, 1: {'litellm_provider': 'p'}}"
+                ),
+            );
+            assert_eq!(Store::new(LiveDict::new(&rows)).keys_in("p"), ["a"]);
+        });
+    }
+}

litellm-rust/crates/python-bridge/src/model_cost/mod.rs

@@ -0,0 +1,497 @@
+mod fetch;
+mod host;
+mod live;
+mod table;
+
+use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+use litellm_host_python::to_py;
+use litellm_model_catalog::{
+    CostMap, IntegrityLimits, Landed, Pass, Placed, Provenance, RegistrationTable, Registrations,
+    Store, write_pass,
+};
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+
+pub(crate) use fetch::{NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched};
+use live::LiveDict;
+
+/// The installed document and the `register_model` calls kept above it, each as the dict Python merged
+struct Inputs {
+    cost_map: CostMap,
+    registrations: Registrations<Py<PyDict>>,
+}
+
+impl Inputs {
+    fn clone_ref(&self, py: Python<'_>) -> Self {
+        Self {
+            cost_map: self.cost_map.clone(),
+            registrations: self.registrations.map(|kept| kept.clone_ref(py)),
+        }
+    }
+}
+
+/// The cost map's inputs and the name index over `litellm.model_cost`
+#[pyclass(
+    frozen,
+    name = "NativeModelCostMap",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeModelCostMap {
+    inputs: RwLock<Inputs>,
+    limits: IntegrityLimits,
+    index: RwLock<Option<Arc<Store<LiveDict>>>>,
+}
+
+#[pyclass(
+    frozen,
+    name = "NativeInstalled",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeInstalled {
+    /// The new live table, the document without registrations
+    #[pyo3(get)]
+    rows: Py<PyDict>,
+    /// How many keys the document itself serves, aliases included and registrations left out
+    #[pyo3(get)]
+    model_count: usize,
+    /// The document's `fallback_generalizations` block as written, `None` when it has none
+    #[pyo3(get)]
+    fallback_generalizations: Py<PyAny>,
+}
+
+#[pyclass(
+    frozen,
+    name = "NativeRejected",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeRejected {
+    #[pyo3(get)]
+    reason: String,
+    /// The document parsed but has fewer models than the integrity limits allow
+    #[pyo3(get)]
+    integrity: bool,
+}
+
+/// A finished registration pass, with the registrations no built-in row matched, by registered key
+#[pyclass(
+    frozen,
+    name = "NativeRegistered",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeRegistered {
+    #[pyo3(get)]
+    unmatched: Py<PyDict>,
+}
+
+/// A registration pass that stopped because the table it was handed is no longer the published one
+#[pyclass(
+    frozen,
+    name = "NativeTableSwapped",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeTableSwapped;
+
+/// The inputs at one moment, for a caller that has to put them back, such as test isolation
+#[pyclass(
+    frozen,
+    name = "NativeModelCostMapState",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeModelCostMapState(Inputs);
+
+#[pyclass(
+    frozen,
+    name = "NativeProvenance",
+    module = "litellm.rust_bridge._native"
+)]
+pub(crate) struct NativeProvenance {
+    #[pyo3(get)]
+    revision: String,
+    #[pyo3(get)]
+    etag: Option<String>,
+}
+
+#[pymethods]
+impl NativeModelCostMap {
+    /// Starts from the bundled document. A fetched document with fewer than `min_model_count` models, or fewer than
+    /// `max_shrink_ratio` times the bundled count, is rejected
+    #[staticmethod]
+    fn bundled(min_model_count: i64, max_shrink_ratio: f64) -> PyResult<Self> {
+        let min_model_count = usize::try_from(min_model_count).map_err(|_| {
+            PyValueError::new_err(format!(
+                "min_model_count must be 0 or more, got {min_model_count}"
+            ))
+        })?;
+        let limits = IntegrityLimits::new(min_model_count, max_shrink_ratio).ok_or_else(|| {
+            PyValueError::new_err(format!(
+                "max_shrink_ratio must be between 0 and 1, got {max_shrink_ratio}"
+            ))
+        })?;
+        Ok(Self {
+            inputs: RwLock::new(Inputs {
+                cost_map: CostMap::bundled(),
+                registrations: Registrations::default(),
+            }),
+            limits,
+            index: RwLock::new(None),
+        })
+    }
+
+    /// The live table built from a fetched document, or why it was rejected, in which case nothing changed
+    #[pyo3(signature = (body, etag = None))]
+    fn install(&self, py: Python<'_>, body: &[u8], etag: Option<String>) -> PyResult<Py<PyAny>> {
+        match crate::logger::capture(py).scope(|| CostMap::read(body, etag, self.limits)) {
+            Ok(cost_map) => self.install_document(py, cost_map),
+            Err(error) => rejected(py, &error),
+        }
+    }
+
+    /// A caller's document as a table, read without size limits and installed nowhere
+    #[staticmethod]
+    fn finalize(py: Python<'_>, body: &[u8]) -> PyResult<Py<PyAny>> {
+        match crate::logger::capture(py)
+            .scope(|| CostMap::read(body, None, IntegrityLimits::UNCHECKED))
+        {
+            Ok(cost_map) => Ok(Py::new(py, built(py, &cost_map)?)?.into_any()),
+            Err(error) => rejected(py, &error),
+        }
+    }
+
+    /// The live table built from the bundled document, which is never size checked
+    fn install_bundled(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
+        self.install_document(py, CostMap::bundled())
+    }
+
+    /// Writes each entry into `rows`, or into the published table when `rows` is `None`
+    #[pyo3(signature = (rows, entries, persist, custom_llm_provider = None))]
+    fn register(
+        &self,
+        py: Python<'_>,
+        rows: Option<&Bound<'_, PyDict>>,
+        entries: &Bound<'_, PyDict>,
+        persist: bool,
+        custom_llm_provider: Option<&str>,
+    ) -> PyResult<Py<PyAny>> {
+        let entries: Vec<(String, Bound<'_, PyAny>)> = entries
+            .iter()
+            .map(|(key, registration)| Ok((key.str()?.to_string(), registration)))
+            .collect::<PyResult<_>>()?;
+        let kept = match persist {
+            true => self.merged_with_kept(&entries)?,
+            false => Vec::new(),
+        };
+        let pass = self.write_pass(rows, &entries, custom_llm_provider)?;
+        if persist {
+            self.keep(kept, pass.landed());
+        }
+        into_py(py, pass)
+    }
+
+    /// Writes every kept registration into `rows` again, oldest first
+    fn reapply(&self, rows: &Bound<'_, PyDict>) -> PyResult<Py<PyAny>> {
+        let py = rows.py();
+        let entries: Vec<(String, Bound<'_, PyAny>)> = self
+            .read_inputs()
+            .registrations
+            .entries()
+            .map(|(key, kept)| (key.to_owned(), kept.bind(py).clone().into_any()))
+            .collect();
+        let pass = self.write_pass(Some(rows), &entries, None)?;
+        self.write_inputs().registrations.record(pass.landed());
+        into_py(py, pass)
+    }
+
+    /// The installed row under `key` with the kept registrations that apply to it, without anything the router wrote
+    fn registered_row<'py>(
+        &self,
+        py: Python<'py>,
+        key: &str,
+    ) -> PyResult<Option<Bound<'py, PyDict>>> {
+        let (cost_map, applied) = {
+            let inputs = self.read_inputs();
+            let applied: Vec<Bound<'py, PyDict>> = inputs
+                .registrations
+                .over(key)
+                .into_iter()
+                .map(|kept| kept.bind(py).clone())
+                .collect();
+            (inputs.cost_map.clone(), applied)
+        };
+        table::registered_row(py, &cost_map, key, &applied)
+    }
+
+    fn registered_keys(&self) -> Vec<String> {
+        let inputs = self.read_inputs();
+        table::registered_keys(
+            &inputs.cost_map,
+            inputs.registrations.entries().map(|(key, _)| key),
+        )
+    }
+
+    /// The installed document as a table, without any registration
+    fn loaded<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
+        let cost_map = self.read_inputs().cost_map.clone();
+        table::build(py, &cost_map)
+    }
+
+    /// The key itself when `rows` has it, otherwise the key that differs from `name` only in case
+    fn find_key(&self, rows: &Bound<'_, PyDict>, name: &str) -> Option<String> {
+        self.index_over(rows).find_key(name)
+    }
+
+    /// Keys whose row's `litellm_provider` is exactly `provider`, in row order
+    fn keys_in(&self, rows: &Bound<'_, PyDict>, provider: &str) -> Vec<String> {
+        self.index_over(rows).keys_in(provider)
+    }
+
+    /// Drops the name index, for a caller that changed keys in a way the dict's length does not show
+    fn invalidate(&self) {
+        if let Some(index) = self
+            .index
+            .read()
+            .unwrap_or_else(PoisonError::into_inner)
+            .as_ref()
+        {
+            index.invalidate();
+        }
+    }
+
+    fn snapshot(&self, py: Python<'_>) -> NativeModelCostMapState {
+        NativeModelCostMapState(self.read_inputs().clone_ref(py))
+    }
+
+    /// Puts back the document, registrations and provenance a `snapshot` saw
+    fn restore(&self, py: Python<'_>, state: PyRef<'_, NativeModelCostMapState>) {
+        *self.write_inputs() = state.0.clone_ref(py);
+        self.invalidate();
+    }
+
+    #[getter]
+    fn provenance(&self) -> NativeProvenance {
+        let Provenance { revision, etag } = self.read_inputs().cost_map.provenance().clone();
+        NativeProvenance { revision, etag }
+    }
+}
+
+impl NativeModelCostMap {
+    fn read_inputs(&self) -> RwLockReadGuard<'_, Inputs> {
+        self.inputs.read().unwrap_or_else(PoisonError::into_inner)
+    }
+
+    fn write_inputs(&self) -> RwLockWriteGuard<'_, Inputs> {
+        self.inputs.write().unwrap_or_else(PoisonError::into_inner)
+    }
+
+    fn install_document(&self, py: Python<'_>, cost_map: CostMap) -> PyResult<Py<PyAny>> {
+        let installed = built(py, &cost_map)?;
+        self.write_inputs().cost_map = cost_map;
+        *self.index.write().unwrap_or_else(PoisonError::into_inner) =
+            Some(Arc::new(Store::new(LiveDict::new(installed.rows.bind(py)))));
+        Ok(Py::new(py, installed)?.into_any())
+    }
+
+    /// Each entry merged into the registration kept for its key
+    fn merged_with_kept<'py>(
+        &self,
+        entries: &[(String, Bound<'py, PyAny>)],
+    ) -> PyResult<Vec<(String, Bound<'py, PyDict>)>> {
+        entries
+            .iter()
+            .map(|(key, registration)| {
+                let previous = self
+                    .read_inputs()
+                    .registrations
+                    .get(key)
+                    .map(|kept| kept.bind(registration.py()).clone());
+                Ok((key.clone(), host::kept(previous.as_ref(), registration)?))
+            })
+            .collect()
+    }
+
+    fn keep(&self, kept: Vec<(String, Bound<'_, PyDict>)>, landed: &[Landed]) {
+        let mut inputs = self.write_inputs();
+        for (key, registration) in kept {
+            inputs.registrations.register(&key, registration.unbind());
+        }
+        inputs.registrations.record(landed);
+    }
+
+    fn write_pass<'a, 'py>(
+        &self,
+        rows: Option<&Bound<'py, PyDict>>,
+        entries: &'a [(String, Bound<'py, PyAny>)],
+        custom_llm_provider: Option<&str>,
+    ) -> PyResult<Pass<'a, Bound<'py, PyAny>>> {
+        let mut table = PublishedTable {
+            map: self,
+            rows,
+            custom_llm_provider,
+        };
+        write_pass(&mut table, entries)
+    }
+
+    /// The index over `rows`, rebuilt when `rows` is not the dict it was built over
+    fn index_over(&self, rows: &Bound<'_, PyDict>) -> Arc<Store<LiveDict>> {
+        let current = self
+            .index
+            .read()
+            .unwrap_or_else(PoisonError::into_inner)
+            .clone();
+        if let Some(index) = current.filter(|index| index.rows().is(rows)) {
+            return index;
+        }
+        let index = Arc::new(Store::new(LiveDict::new(rows)));
+        *self.index.write().unwrap_or_else(PoisonError::into_inner) = Some(Arc::clone(&index));
+        index
+    }
+}
+
+/// The dict Python publishes as `litellm.model_cost`, written through the host
+struct PublishedTable<'a, 'py> {
+    map: &'a NativeModelCostMap,
+    rows: Option<&'a Bound<'py, PyDict>>,
+    custom_llm_provider: Option<&'a str>,
+}
+
+impl<'py> RegistrationTable for PublishedTable<'_, 'py> {
+    type Registration = Bound<'py, PyAny>;
+    type Error = PyErr;
+
+    fn place(&mut self, key: &str, registration: &Self::Registration) -> PyResult<Option<Placed>> {
+        let Some(landing) = host::resolve(self.rows, key, registration, self.custom_llm_provider)?
+        else {
+            return Ok(None);
+        };
+        let rows = match self.rows {
+            Some(rows) => rows.clone(),
+            None => host::published(registration.py())?,
+        };
+        table::write(&rows, &landing, registration)?;
+        self.map.invalidate();
+        Ok(Some(Placed {
+            key: landing.key,
+            matched_builtin: landing.matched_builtin,
+        }))
+    }
+}
+
+fn built(py: Python<'_>, cost_map: &CostMap) -> PyResult<NativeInstalled> {
+    let rows = table::build(py, cost_map)?;
+    let fallback_generalizations = match cost_map.fallback_generalizations() {
+        Some(block) => to_py(py, block)?,
+        None => py.None(),
+    };
+    Ok(NativeInstalled {
+        model_count: rows.len(),
+        rows: rows.unbind(),
+        fallback_generalizations,
+    })
+}
+
+fn rejected(py: Python<'_>, error: &litellm_model_catalog::Error) -> PyResult<Py<PyAny>> {
+    let rejected = NativeRejected {
+        reason: litellm_error::Report::new(error).to_string(),
+        integrity: error.is_integrity(),
+    };
+    Ok(Py::new(py, rejected)?.into_any())
+}
+
+fn into_py<'py>(py: Python<'py>, pass: Pass<'_, Bound<'py, PyAny>>) -> PyResult<Py<PyAny>> {
+    match pass {
+        Pass::Swapped { .. } => Ok(Py::new(py, NativeTableSwapped)?.into_any()),
+        Pass::Finished { unmatched, .. } => {
+            let table = PyDict::new(py);
+            for (key, registration) in unmatched {
+                table.set_item(key, registration)?;
+            }
+            let registered = NativeRegistered {
+                unmatched: table.unbind(),
+            };
+            Ok(Py::new(py, registered)?.into_any())
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::mpsc;
+    use std::time::Duration;
+
+    use pyo3::ffi::c_str;
+    use rstest::rstest;
+
+    use super::*;
+
+    const HOST: &std::ffi::CStr = c_str!(
+        r#"
+from types import SimpleNamespace
+
+class Row(dict):
+    names = ()
+
+found = []
+
+def resolve(rows, key, registration, custom_llm_provider):
+    found.append(cost_map.find_key(rows, key))
+    return SimpleNamespace(key=key, base={}, matched_builtin=True)
+
+def merged(current, base, registration):
+    return Row({**(current or {}), **base, **registration})
+
+def write(rows, written):
+    rows.update(written)
+"#
+    );
+
+    #[rstest]
+    fn a_host_lookup_that_reaches_find_key_on_the_same_map_finishes() {
+        Python::initialize();
+        let (sender, receiver) = mpsc::channel();
+        std::thread::spawn(move || {
+            let outcome = Python::attach(|py| -> PyResult<(Vec<Option<String>>, f64)> {
+                let cost_map = Py::new(py, NativeModelCostMap::bundled(0, 0.0)?)?;
+                let module_name = std::ffi::CString::new(host::MODULE).unwrap();
+                let host =
+                    PyModule::from_code(py, HOST, c_str!("host_model_cost.py"), &module_name)?;
+                host.setattr("cost_map", &cost_map)?;
+                let modules = py.import("sys")?.getattr("modules")?;
+                let packages = ["litellm", "litellm.rust_bridge", "litellm.rust_bridge.host"];
+                for package in packages {
+                    modules.set_item(package, PyModule::new(py, package)?)?;
+                }
+                modules.set_item(host::MODULE, &host)?;
+                let rows = PyDict::new(py);
+                rows.set_item("Lit-Reentry", PyDict::new(py))?;
+                let entry = PyDict::new(py);
+                entry.set_item("input_cost_per_token", 1.5)?;
+                let entries = PyDict::new(py);
+                entries.set_item("lit-reentry", entry)?;
+                let registered = cost_map
+                    .bind(py)
+                    .call_method1("register", (&rows, &entries, false));
+                for name in [host::MODULE].into_iter().chain(packages) {
+                    modules.del_item(name)?;
+                }
+                registered?;
+                let price = rows
+                    .get_item("lit-reentry")?
+                    .expect("the registration landed")
+                    .get_item("input_cost_per_token")?
+                    .extract()?;
+                Ok((host.getattr("found")?.extract()?, price))
+            });
+            sender
+                .send(outcome.map_err(|error| error.to_string()))
+                .unwrap();
+        });
+
+        let (found, price) = receiver
+            .recv_timeout(Duration::from_secs(30))
+            .expect("register finished instead of deadlocking on its own lock")
+            .unwrap();
+        assert_eq!(found, [Some("Lit-Reentry".to_owned())]);
+        assert!((price - 1.5).abs() < f64::EPSILON);
+    }
+}

litellm-rust/crates/python-bridge/src/model_cost/table.rs

@@ -0,0 +1,166 @@
+//! Builds `litellm.model_cost` from an installed document
+
+use std::borrow::Cow;
+use std::collections::HashMap;
+
+use indexmap::IndexSet;
+use litellm_host_python::to_py;
+use litellm_model_catalog::wire::Fields;
+use litellm_model_catalog::{CostMap, FALLBACK_GENERALIZATIONS, Written};
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use serde_json::Value;
+
+use super::host::{self, Landing};
+
+const ALIASES: &str = "aliases";
+
+enum Served<'a> {
+    Model,
+    AliasOf(&'a str),
+    /// A top-level value that is not a model row, such as `sample_spec`
+    Section(&'a Value),
+}
+
+/// Every served key in dict order. An `alias_of` row repeats among the aliases and keeps its document position
+fn served(cost_map: &CostMap) -> impl Iterator<Item = (&str, Served<'_>)> {
+    let catalog = cost_map.catalog();
+    let documented = catalog
+        .written()
+        .filter(|(key, _)| *key != FALLBACK_GENERALIZATIONS)
+        .filter_map(|(key, written)| {
+            let served = match written {
+                Written::Section(value) => Served::Section(value),
+                Written::Entry(_) => match catalog.aliases().get(key) {
+                    Some(owner) => Served::AliasOf(owner),
+                    None => catalog.get(key).map(|_| Served::Model)?,
+                },
+            };
+            Some((key, served))
+        });
+    let listed = cost_map.alias_groups().flat_map(|(owner, names)| {
+        names
+            .iter()
+            .skip(1)
+            .map(move |alias| (alias.as_str(), Served::AliasOf(owner)))
+    });
+    documented.chain(listed)
+}
+
+/// The installed document as a new dict
+pub(crate) fn build<'py>(py: Python<'py>, cost_map: &CostMap) -> PyResult<Bound<'py, PyDict>> {
+    let table = PyDict::new(py);
+    let mut rows = HashMap::new();
+    for (key, served) in served(cost_map) {
+        let value = match served {
+            Served::Model => row(py, cost_map, key, &mut rows)?,
+            Served::AliasOf(owner) => row(py, cost_map, owner, &mut rows)?,
+            Served::Section(value) => to_py(py, value)?.into_bound(py),
+        };
+        table.set_item(key, value)?;
+    }
+    Ok(table)
+}
+
+fn row<'a, 'py>(
+    py: Python<'py>,
+    cost_map: &'a CostMap,
+    owner: &'a str,
+    rows: &mut HashMap<&'a str, Bound<'py, PyAny>>,
+) -> PyResult<Bound<'py, PyAny>> {
+    if let Some(row) = rows.get(owner) {
+        return Ok(row.clone());
+    }
+    let fields = cost_map
+        .catalog()
+        .get(owner)
+        .expect("a served model or alias owner is a catalog row");
+    let fields = to_py(py, without_alias_list(fields).as_ref())?.into_bound(py);
+    let row = match cost_map.names(owner) {
+        [] => fields,
+        names => host::entry(fields, names)?,
+    };
+    rows.insert(owner, row.clone());
+    Ok(row)
+}
+
+fn without_alias_list(fields: &Fields) -> Cow<'_, Fields> {
+    if fields.get(ALIASES).is_none_or(Value::is_null) {
+        return Cow::Borrowed(fields);
+    }
+    Cow::Owned(
+        fields
+            .iter()
+            .filter(|(name, _)| *name != ALIASES)
+            .map(|(name, value)| (name.clone(), value.clone()))
+            .collect(),
+    )
+}
+
+/// The installed row under `key` with `applied` written over it in order, `None` when neither has one
+pub(crate) fn registered_row<'py>(
+    py: Python<'py>,
+    cost_map: &CostMap,
+    key: &str,
+    applied: &[Bound<'py, PyDict>],
+) -> PyResult<Option<Bound<'py, PyDict>>> {
+    let catalog = cost_map.catalog();
+    let installed = catalog
+        .get(key)
+        .or_else(|| match catalog.section(key) {
+            Some(Value::Object(fields)) if key != FALLBACK_GENERALIZATIONS => Some(fields),
+            _ => None,
+        })
+        .map(without_alias_list);
+    let row = match (installed, applied) {
+        (None, []) => return Ok(None),
+        (Some(fields), _) => to_py(py, fields.as_ref())?
+            .into_bound(py)
+            .cast_into::<PyDict>()?,
+        (None, _) => PyDict::new(py),
+    };
+    for registration in applied {
+        row.update(registration.as_mapping())?;
+    }
+    Ok(Some(row))
+}
+
+/// The keys `registered_row` answers for, installed keys first
+pub(crate) fn registered_keys<'a>(
+    cost_map: &'a CostMap,
+    registered: impl Iterator<Item = &'a str>,
+) -> Vec<String> {
+    let served: IndexSet<&str> = served(cost_map).map(|(key, _)| key).collect();
+    let only_registered: Vec<&str> = registered.filter(|key| !served.contains(key)).collect();
+    served
+        .into_iter()
+        .chain(only_registered)
+        .map(str::to_owned)
+        .collect()
+}
+
+/// Writes `registration` over the row under `landing.key`, and under every name that shares that row object
+pub(crate) fn write<'py>(
+    rows: &Bound<'py, PyDict>,
+    landing: &Landing<'py>,
+    registration: &Bound<'py, PyAny>,
+) -> PyResult<()> {
+    let current = rows.get_item(&landing.key)?;
+    let row = host::merged(current.as_ref(), landing, registration)?;
+    let names: Vec<String> = row.getattr("names")?.extract()?;
+    let sharing: Vec<&String> = names
+        .iter()
+        .filter(|name| {
+            let served = rows.get_item(name.as_str()).ok().flatten();
+            served
+                .zip(current.as_ref())
+                .is_some_and(|(served, current)| served.is(current))
+        })
+        .collect();
+    let written = PyDict::new(rows.py());
+    written.set_item(&landing.key, &row)?;
+    sharing
+        .into_iter()
+        .try_for_each(|name| written.set_item(name, &row))?;
+    host::write(rows, &written)
+}

litellm-rust/crates/python-bridge/src/python_settings.rs

@@ -56,7 +56,7 @@ impl PythonSettings {
         }
     }
 
-    #[cfg(test)]
+    /// A snapshot of a value Python handed in, read with the same field rules as the accessor's
     pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> {
         Snapshot { group: self, value }
     }

litellm-rust/crates/router-core/src/call.rs

@@ -1037,8 +1037,13 @@ impl Core {
                 self.log_retry(target, at, 0, &failed.error);
             }
             let named = failure.failed_deployment_id.as_ref();
-            if let (Some(status), Some(id)) =
-                (failure.status, named.or(failed.deployment_id.as_ref()))
+            let failed_on = named.or(failed.deployment_id.as_ref());
+            if let Some(id) = failed_on
+                && !inputs.failed_ids.contains(id)
+            {
+                inputs.failed_ids.push(id.clone());
+            }
+            if let (Some(status), Some(id)) = (failure.status, failed_on)
                 && !should_retry_status(status)
                 && !inputs.skipped_ids.contains(id)
             {

litellm-rust/crates/router-core/src/call/stream.rs

@@ -236,10 +236,11 @@ impl Router {
             .unwrap_or((0, 0));
         let (budget, policy_applies) =
             stream_retry_budget(&snapshot, &request, &origin, failure.kind);
-        let inputs = SelectInputs {
+        let mut inputs = SelectInputs {
             target_order: target.order,
             requested: Some(target.group.clone()),
             fallback_hop: at.level > 0,
+            failed_ids: origin.deployment_id.iter().cloned().collect(),
             ..SelectInputs::default()
         };
         let mut last = failure;
@@ -290,6 +291,14 @@ impl Router {
                     // `_stream_fallback_error_for_raised`: a status the pre-stream retry rule
                     // would not retry ends the retries with that failure.
                     Error::Attempt(failure) => {
+                        if let Some(id) = failure
+                            .failed_deployment_id
+                            .clone()
+                            .or(failed.deployment_id)
+                            && !inputs.failed_ids.contains(&id)
+                        {
+                            inputs.failed_ids.push(id);
+                        }
                         let stop = failure.status.is_some_and(|s| !should_retry_status(s));
                         last = failure;
                         if stop {

litellm-rust/crates/router-core/src/select.rs

@@ -108,6 +108,9 @@ impl Candidates {
 pub(crate) struct SelectInputs {
     pub target_order: Option<i64>,
     pub skipped_ids: Vec<String>,
+    /// Every deployment an attempt of this request failed on, whatever the status: the callback
+    /// filters read it as `_retry_skipped_deployment_ids`, so an affinity pin can leave one.
+    pub failed_ids: Vec<String>,
     /// A weighted failover hop's `_excluded_deployment_ids`: never picked.
     pub excluded_ids: Arc<[String]>,
     /// Stop before the plugins, order and skip filters: the group less cooldowns and blocked
@@ -470,7 +473,7 @@ impl Core {
                 StrategyScope::Default,
             );
             ctx.excluded_ids = &inputs.excluded_ids;
-            ctx.skipped_ids = &inputs.skipped_ids;
+            ctx.skipped_ids = &inputs.failed_ids;
             ctx.target_order = inputs.target_order.or(request_order);
             let filtered = self.host.callback_filters(&ctx, healthy).await?;
             healthy = filtered.kept;

litellm-rust/crates/router-core/tests/call_loop.rs

@@ -5,7 +5,7 @@ use std::{collections::VecDeque, sync::Arc, time::Duration};
 
 use async_trait::async_trait;
 use litellm_router_core::{
-    Attempt, AttemptError, Deployment, Error, FailureKind, FailureReport, Host,
+    Attempt, AttemptError, CallbackFiltered, Deployment, Error, FailureKind, FailureReport, Host,
     InMemoryRouterCache, Op, Output, Payload, Router, RouterCache, RouterConfig,
     RouterDynamicConfig, RouterEvent, RouterRequest, RoutingContext, SeededRandomness, SelectStage,
     SuccessReport,
@@ -37,6 +37,9 @@ struct ScriptedHost {
     pick_failures: Mutex<VecDeque<AttemptError>>,
     /// The next attempt's failure, whatever group it serves.
     first_failure: Mutex<Option<AttemptError>>,
+    /// The callback filters pin to this deployment until an attempt of the request failed on it,
+    /// as `encrypted_content_affinity` does.
+    pinned: Option<&'static str>,
 }
 
 impl ScriptedHost {
@@ -112,6 +115,24 @@ impl Host for ScriptedHost {
         }
     }
 
+    async fn callback_filters(
+        &self,
+        ctx: &RoutingContext<'_>,
+        candidates: Vec<Arc<Deployment>>,
+    ) -> Result<CallbackFiltered, Error> {
+        let kept = match self.pinned {
+            Some(pin) if !ctx.skipped_ids.iter().any(|id| id == pin) => {
+                candidates.into_iter().filter(|d| d.id == pin).collect()
+            }
+            Some(pin) => candidates.into_iter().filter(|d| d.id != pin).collect(),
+            None => candidates,
+        };
+        Ok(CallbackFiltered {
+            kept,
+            preferred: None,
+        })
+    }
+
     async fn allow_fallback(
         &self,
         _request: &RouterRequest,
@@ -223,6 +244,32 @@ async fn a_rate_limit_retries_in_the_same_group(config: RouterConfig) {
     );
 }
 
+#[rstest]
+#[case::server_error(server_error())]
+#[case::rate_limit(failure(FailureKind::RateLimit, 429))]
+#[case::connection_error(AttemptError::new(FailureKind::ApiConnection, None))]
+#[tokio::test]
+async fn a_retry_leaves_a_pinned_deployment_that_failed_the_request(
+    mut config: RouterConfig,
+    #[case] origin_failure: AttemptError,
+) {
+    config.callback_filters = true;
+    let host = Arc::new(ScriptedHost {
+        first_failure: Mutex::new(Some(origin_failure)),
+        pinned: Some("a1"),
+        ..ScriptedHost::default()
+    });
+    let router = router(config, json!({"num_retries": 3}), host.clone());
+    let response = router
+        .call(Op::from("acompletion"), request())
+        .await
+        .unwrap();
+    assert_eq!(
+        (host.selected(), response.attempted_retries),
+        (vec!["a1".to_owned(), "a2".to_owned()], 1)
+    );
+}
+
 #[rstest]
 #[tokio::test]
 async fn a_bad_request_falls_back_without_retrying(config: RouterConfig) {
@@ -945,6 +992,56 @@ async fn a_stream_retry_is_reported_before_its_attempt(config: RouterConfig) {
     assert_eq!((response.attempted_retries, host.retries()), (1, 1));
 }
 
+fn pinned_to_a1() -> Arc<ScriptedHost> {
+    Arc::new(ScriptedHost {
+        pinned: Some("a1"),
+        ..ScriptedHost::default()
+    })
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_stream_retry_leaves_the_pinned_deployment_the_stream_failed_on(
+    mut config: RouterConfig,
+) {
+    config.callback_filters = true;
+    let host = pinned_to_a1();
+    let router = router(config, json!({}), host.clone());
+    let origin = StreamOrigin {
+        deployment_id: Some("a1".to_owned()),
+        ..stream_origin()
+    };
+    router
+        .retry_stream(
+            Op::from("acompletion"),
+            stream_request(),
+            origin,
+            server_error(),
+        )
+        .await
+        .unwrap();
+    assert_eq!(host.selected(), vec!["a2".to_owned()]);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_stream_retry_leaves_a_pinned_deployment_a_retry_failed_on(mut config: RouterConfig) {
+    config.callback_filters = true;
+    let host = pinned_to_a1();
+    *host.first_failure.lock() = Some(server_error());
+    let router = router(config, json!({}), host.clone());
+    router
+        .retry_stream(
+            Op::from("acompletion"),
+            stream_request(),
+            stream_origin(),
+            server_error(),
+        )
+        .await
+        .unwrap();
+    assert_eq!(host.selected(), vec!["a1".to_owned(), "a2".to_owned()]);
+}
+
 fn fallback_list(value: serde_json::Value) -> Fallbacks {
     serde_json::from_value(value).unwrap()
 }

litellm-rust/crates/s3/AGENTS.md

@@ -0,0 +1,3 @@
+`litellm-s3` builds `aws-sdk-s3` clients that send through `litellm_http::Client`, and owns the S3 concepts the SDK lacks, such as the `s3://bucket/key` URI. It reads no environment and picks no credentials: the host passes the region, the endpoint and a credentials provider, usually `litellm_auth_aws::AwsCredentialsProvider`
+
+Operation policy such as retries, timeouts and checksum settings stays with the caller, which adjusts the returned config builder

litellm-rust/crates/s3/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "litellm-s3"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+litellm-http.workspace = true
+aws-sdk-s3.workspace = true
+aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] }
+aws-types = "1.4.0"
+aws-credential-types = "1.3.0"
+aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] }
+aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] }
+http.workspace = true
+reqwest.workspace = true
+thiserror.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true

litellm-rust/crates/s3/src/error.rs

@@ -0,0 +1,7 @@
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("expected s3://bucket/key")]
+    InvalidUri,
+}

litellm-rust/crates/s3/src/lib.rs

@@ -0,0 +1,60 @@
+mod error;
+mod transport;
+mod uri;
+
+pub use aws_credential_types::provider::ProvideCredentials;
+pub use aws_sdk_s3 as sdk;
+pub use error::Error;
+pub use uri::S3Uri;
+
+use aws_config::meta::region::RegionProviderChain;
+use aws_sdk_s3::config::{BehaviorVersion, Builder, Region};
+use aws_types::SdkConfig;
+use aws_types::service_config::ServiceConfigKey;
+
+use crate::transport::ReqwestHttpClient;
+
+/// A custom endpoint is addressed path style, as boto3 addresses one
+pub fn config(
+    region: String,
+    endpoint_url: Option<&str>,
+    http: litellm_http::Client,
+    credentials: impl ProvideCredentials + 'static,
+) -> Builder {
+    let builder = aws_sdk_s3::Config::builder()
+        .behavior_version(BehaviorVersion::latest())
+        .region(Region::new(region))
+        .http_client(ReqwestHttpClient(http))
+        .credentials_provider(credentials);
+    match endpoint_url {
+        Some(url) => builder.endpoint_url(url).force_path_style(true),
+        None => builder,
+    }
+}
+
+/// The region, endpoint and credentials an AWS SDK finds in the environment and shared config files, as boto3's
+/// default `s3` client found them. `region` wins over that chain, and `us-east-1` is used when nothing names one
+pub async fn environment_config(http: litellm_http::Client, region: Option<String>) -> Builder {
+    let region = RegionProviderChain::first_try(region.map(Region::new))
+        .or_default_provider()
+        .or_else(Region::from_static("us-east-1"));
+    let shared = aws_config::defaults(BehaviorVersion::latest())
+        .region(region)
+        .http_client(ReqwestHttpClient(http))
+        .load()
+        .await;
+    Builder::from(&shared).force_path_style(has_custom_endpoint(&shared))
+}
+
+fn has_custom_endpoint(shared: &SdkConfig) -> bool {
+    let s3_endpoint = ServiceConfigKey::builder()
+        .service_id("S3")
+        .env("AWS_ENDPOINT_URL")
+        .profile("endpoint_url")
+        .build()
+        .ok();
+    shared.endpoint_url().is_some()
+        || s3_endpoint
+            .and_then(|key| shared.service_config()?.load_config(key))
+            .is_some()
+}

litellm-rust/crates/s3/src/uri.rs

@@ -0,0 +1,76 @@
+use crate::Error;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct S3Uri {
+    pub bucket: String,
+    pub key: String,
+}
+
+/// Splits the way Python's `urllib.parse.urlsplit` does, so the key keeps dot segments and characters a URL parser
+/// would normalize or percent-encode. S3 keys are opaque strings
+impl TryFrom<&str> for S3Uri {
+    type Error = Error;
+
+    fn try_from(uri: &str) -> Result<Self, Error> {
+        let rest = match uri.split_once("://") {
+            Some((scheme, rest)) if scheme.eq_ignore_ascii_case("s3") => rest,
+            _ => return Err(Error::InvalidUri),
+        };
+        let (bucket, path) = rest.split_at(rest.find(['/', '?', '#']).unwrap_or(rest.len()));
+        let key = path
+            .split(['?', '#'])
+            .next()
+            .unwrap_or_default()
+            .trim_start_matches('/');
+        if bucket.is_empty() || key.is_empty() {
+            return Err(Error::InvalidUri);
+        }
+        Ok(Self {
+            bucket: bucket.to_owned(),
+            key: key.to_owned(),
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    #[rstest]
+    #[case::key("s3://bucket/cost.json", "bucket", "cost.json")]
+    #[case::nested_key("s3://bucket/prices/cost.json", "bucket", "prices/cost.json")]
+    #[case::dotted_bucket("s3://my.bucket/cost.json", "my.bucket", "cost.json")]
+    #[case::dot_segments("s3://bucket/archives/../cost.json", "bucket", "archives/../cost.json")]
+    #[case::current_dir_segment("s3://bucket/./cost.json", "bucket", "./cost.json")]
+    #[case::space_kept_raw("s3://bucket/my prices.json", "bucket", "my prices.json")]
+    #[case::query_dropped("s3://bucket/cost.json?versionId=1", "bucket", "cost.json")]
+    #[case::fragment_dropped("s3://bucket/cost.json#top", "bucket", "cost.json")]
+    #[case::leading_slashes_trimmed("s3://bucket//cost.json", "bucket", "cost.json")]
+    #[case::uppercase_scheme("S3://bucket/cost.json", "bucket", "cost.json")]
+    fn an_s3_url_names_its_bucket_and_key(
+        #[case] url: &str,
+        #[case] bucket: &str,
+        #[case] key: &str,
+    ) {
+        assert_eq!(
+            S3Uri::try_from(url).unwrap(),
+            S3Uri {
+                bucket: bucket.into(),
+                key: key.into(),
+            }
+        );
+    }
+
+    #[rstest]
+    #[case::no_key("s3://bucket/")]
+    #[case::no_path("s3://bucket")]
+    #[case::query_only("s3://bucket?cost.json")]
+    #[case::no_bucket("s3:///cost.json")]
+    #[case::other_scheme("https://bucket/cost.json")]
+    #[case::no_scheme("bucket/cost.json")]
+    fn a_url_without_a_bucket_and_key_is_refused(#[case] url: &str) {
+        assert!(matches!(S3Uri::try_from(url), Err(Error::InvalidUri)));
+    }
+}

litellm/litellm_core_utils/AGENTS.md

@@ -1,2 +1,18 @@
 `usage_cost.py` is the one place `usage.cost` is stamped. Reuse its dispatcher and adapters.
 Never inline usage-cost stamping elsewhere. Never delete `usage_cost.py`.
+
+## get_model_cost_map.py
+
+- This module decides when to fetch and install. Fetching, parsing, integrity checks and alias expansion run in Rust
+- The install order is publish, `add_known_models`, router replay, then `model_cost_inputs.reapply`
+  - Registrations land by lookup in the published table, so they are written after publishing
+- Startup and `register_model(url)` fetch with `http_settings=None`, which reads env vars only like main's bare httpx. Reload passes `http_settings()`
+- Startup falls back to the bundled map on any failure. Reload never falls back and keeps the current map
+- `register_model(url)` makes one attempt and registers nothing on failure
+- `utils._register_on_current_map` reruns the whole pass on the new table after `NativeTableSwapped`, and its last attempt passes `rows=None`, so no entry is lost
+- Compat names exist only so code written against main keeps working. New code never calls them, and they never gain behavior of their own
+  - `_finalize_model_cost_map`, use `model_cost_inputs.install`
+  - `GetModelCostMap.validate_model_cost_map` and `_check_*`, since Rust checks fetched documents
+  - `ModelCostMapReloaded.model_cost_map`, use `GetModelCostMap.loaded_model_cost_map()`
+  - the `sleep`, `rng` and `client` parameters, which are ignored with a `DeprecationWarning`
+  - `AliasedCostEntry` re-exported here, which now lives in `rust_bridge/host/model_cost.py`

litellm/llms/fireworks_ai/chat/transformation.py

@@ -19,6 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
 )
 from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning
 from litellm.litellm_core_utils.url_utils import join_api_base
+from litellm.rust_bridge.host.model_cost import ModelCostVersion, model_cost_version
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import (
     AllMessageValues,
@@ -36,7 +37,6 @@ from litellm.types.utils import (
     ProviderSpecificModelInfo,
 )
 from litellm.utils import (
-    get_model_cost_mutation_generation,
     supports_function_calling,
     supports_reasoning,
     supports_tool_choice,
@@ -542,21 +542,15 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
             for message in messages
         ]
 
-    # Cached index of fireworks_ai/* entries from litellm.model_cost. Building
-    # this index requires a full scan of model_cost (tens of thousands of
-    # entries), so we memoize it. The cache key is (id(model_cost),
-    # mutation_generation): the generation counter is bumped on every
-    # register_model / reload path, so add+remove or in-place value
-    # replacement (which can leave id and len unchanged) still invalidates.
-    _fireworks_index_cache: tuple[int, int, list[tuple[str, dict]]] | None = None
+    _fireworks_index_cache: tuple[ModelCostVersion, list[tuple[str, dict]]] | None = None
 
     @classmethod
     def _get_fireworks_index(cls) -> list[tuple[str, dict]]:
         model_cost: Final = litellm.model_cost
-        signature: Final = (id(model_cost), get_model_cost_mutation_generation())
+        version: Final = model_cost_version()
         cached: Final = cls._fireworks_index_cache
-        if cached is not None and cached[0] == signature[0] and cached[1] == signature[1]:
-            return cached[2]
+        if cached is not None and cached[0] == version:
+            return cached[1]
 
         index: Final[list[tuple[str, dict]]] = []
         for key, model_info in (*model_cost.items(),):
@@ -570,7 +564,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
                 continue
             index.append((key_short, model_info))
 
-        cls._fireworks_index_cache = (signature[0], signature[1], index)
+        cls._fireworks_index_cache = (version, index)
         return index
 
     @staticmethod

litellm/proxy/common_utils/model_deprecation.py

@@ -10,7 +10,8 @@ from typing import TYPE_CHECKING, Final, TypeAlias
 
 import litellm
 from litellm._logging import verbose_logger
-from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetModelCostMap
+from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
+from litellm.rust_bridge.host.model_cost import AliasedCostEntry
 from litellm.types.proxy.model_deprecation import (
     DEFAULT_DEPRECATION_WARN_DAYS,
     DeprecationStatus,

litellm/proxy/proxy_server.py

@@ -5682,12 +5682,6 @@ def resolve_classifier_plugin(
     return resolved
 
 
-def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
-    from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map
-
-    return adopt_model_cost_map(new_model_cost_map)
-
-
 def _websearch_handler_params(stored: Mapping[str, object]) -> dict[str, object]:
     """
     Translate stored web search interception settings into handler kwargs.
@@ -9515,7 +9509,7 @@ class ProxyConfig:
                 )
                 return
 
-            models_count = _swap_in_model_cost_map(reload_result.model_cost_map)
+            models_count = reload_result.model_count
             self.model_cost_map_loaded_at = current_time
             await record_reload_run(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time)
             # Adopted last, so neither a failed fetch nor a failed status write is recorded
@@ -21979,7 +21973,7 @@ async def reload_model_cost_map(
                 detail=f"Failed to reload model cost map: {reload_result.reason}. Current pricing data was kept.",
             )
 
-        models_count = _swap_in_model_cost_map(reload_result.model_cost_map)
+        models_count = reload_result.model_count
         current_time = utc_now()
         proxy_config.model_cost_map_loaded_at = current_time
         provenance: Final = get_model_cost_map_provenance()

litellm/router_rust/invoke/streaming.py

@@ -684,7 +684,10 @@ def _mid_stream_continuation_eligible(
     lookup_groups: Final = fallback_lookup_groups(
         request, requested_group if isinstance(requested_group, str) else None
     )
-    targets: Final = continuation_fallback_targets(core.settings().fallbacks, lookup_groups)  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # fallbacks as Python spells them
+    targets: Final = continuation_fallback_targets(
+        request.get("fallbacks", core.settings().fallbacks),  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # the request's list replaces the router's
+        lookup_groups,
+    )
     if any(fallback_target_overrides_request_params(target) for target in targets):
         return False
     if any(_names_deployment_that_cannot_continue(core, target) for target in targets):

litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py

@@ -681,4 +681,4 @@ class EncryptedContentAffinityCheck(CustomLogger):
         )
         ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input)
         strip_encrypted_reasoning_from_messages(anthropic_messages)
-        return candidates
+        return untried_candidates if excluded_ids or retry_skipped_ids else candidates

litellm/rust_bridge/_native.pyi

@@ -8,6 +8,7 @@ from typing_extensions import disjoint_base
 
 from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings
 from litellm.rust_bridge.execution import Execution
+from litellm.rust_bridge.host.settings import HttpSettings
 from litellm.rust_bridge.public_call import NativeCall
 
 _T_co = TypeVar("_T_co", covariant=True)
@@ -22,6 +23,102 @@ def trace_encode_error(message: str) -> bytes: ...
 def bundled_model_cost_map() -> bytes: ...
 def native_redis_breaker_states() -> list[str]: ...
 
+@final
+class NativeCostMapSource:
+    def __new__(cls, url: str, *, http_settings: HttpSettings | None) -> NativeCostMapSource: ...
+    def fetch(self, policy: NativeFetchPolicy) -> Future[NativeFetched | NativeFetchFailed]: ...
+    def fetch_blocking(
+        self, policy: NativeFetchPolicy, resume_from: NativeFetchFailed | None = None
+    ) -> NativeFetched | NativeFetchFailed: ...
+
+@final
+class NativeFetchPolicy:
+    def __new__(
+        cls,
+        *,
+        max_attempts: int,
+        max_wait_seconds: float,
+        timeout_seconds: float,
+        min_delay_seconds: float,
+    ) -> NativeFetchPolicy: ...
+
+@final
+class NativeFetched:
+    @property
+    def body(self) -> bytes: ...
+    @property
+    def etag(self) -> str | None: ...
+
+@final
+class NativeFetchFailed:
+    @property
+    def reason(self) -> str: ...
+    @property
+    def retryable(self) -> bool: ...
+    @property
+    def retry_after(self) -> float | None: ...
+
+@final
+class NativeInstalled:
+    @property
+    def rows(self) -> dict[str, dict[str, object]]: ...
+    @property
+    def model_count(self) -> int: ...
+    @property
+    def fallback_generalizations(self) -> JsonValue: ...
+
+@final
+class NativeRegistered:
+    @property
+    def unmatched(self) -> dict[str, Mapping[str, object]]: ...
+
+@final
+class NativeTableSwapped: ...
+
+@final
+class NativeRejected:
+    @property
+    def reason(self) -> str: ...
+    @property
+    def integrity(self) -> bool: ...
+
+@final
+class NativeProvenance:
+    @property
+    def revision(self) -> str: ...
+    @property
+    def etag(self) -> str | None: ...
+
+@final
+class NativeModelCostMap:
+    @staticmethod
+    def bundled(min_model_count: int, max_shrink_ratio: float) -> NativeModelCostMap: ...
+    def install(self, body: bytes, etag: str | None = None) -> NativeInstalled | NativeRejected: ...
+    def install_bundled(self) -> NativeInstalled: ...
+    @staticmethod
+    def finalize(body: bytes) -> NativeInstalled | NativeRejected: ...
+    def register(
+        self,
+        rows: dict[str, dict[str, object]] | None,
+        entries: Mapping[str, Mapping[str, object]],
+        persist: bool,
+        custom_llm_provider: str | None = None,
+    ) -> NativeRegistered | NativeTableSwapped: ...
+    def reapply(self, rows: dict[str, dict[str, object]]) -> NativeRegistered | NativeTableSwapped: ...
+    def registered_row(self, key: str) -> dict[str, object] | None: ...
+    def registered_keys(self) -> list[str]: ...
+    def loaded(self) -> dict[str, dict[str, object]]: ...
+    def find_key(self, rows: dict[str, dict[str, object]], name: str) -> str | None: ...
+    def keys_in(self, rows: dict[str, dict[str, object]], provider: str) -> list[str]: ...
+    def invalidate(self) -> None: ...
+    def snapshot(self) -> NativeModelCostMapState: ...
+    def restore(self, state: NativeModelCostMapState) -> None: ...
+    @property
+    def provenance(self) -> NativeProvenance: ...
+
+@final
+class NativeModelCostMapState: ...
+
 @final
 class NativeClickHouseSpendConfig:
     def __new__(cls, database: str, url: str, retention_days: int) -> NativeClickHouseSpendConfig: ...
@@ -589,19 +686,30 @@ __all__ = [
     "NativeAuthCache",
     "NativeClickHouseSpendConfig",
     "NativeClickHouseSpendStorage",
+    "NativeCostMapSource",
     "NativeDatabase",
     "NativeDiagnosticProcessor",
+    "NativeFetchFailed",
+    "NativeFetchPolicy",
+    "NativeFetched",
+    "NativeInstalled",
+    "NativeModelCostMap",
+    "NativeModelCostMapState",
     "NativePending",
+    "NativeProvenance",
     "NativeRateLimits",
     "NativeRedisBreaker",
     "NativeRedisClient",
     "NativeRedisStoreError",
+    "NativeRegistered",
+    "NativeRejected",
     "NativeRequestRedis",
     "NativeRouterCooldowns",
     "NativeRouterUsage",
     "NativeSend",
     "NativeSensitiveRoutePins",
     "NativeSpendCounters",
+    "NativeTableSwapped",
     "PostCallSettlement",
     "ProcessReservedForForking",
     "ResponsesWebSocketConnection",

litellm/rust_bridge/host/AGENTS.md

@@ -2,10 +2,23 @@
 
 Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name
 
-- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`), `redis_credentials` (what a caller's redis-py `CredentialProvider` authenticates with)
+- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`), `redis_credentials` (what a caller's redis-py `CredentialProvider` authenticates with), `model_cost` (the live `litellm.model_cost` table, every write to it and its version)
 - A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here
 - Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers
 - The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green
 - Inputs and outputs use values and frozen snapshots. Retain caller-owned clients and caches as opaque references when identity is required; Rust never discovers them through module globals
 - Moving a question into Rust means replacing the function body with a native call, not adding a second resolver. The Python signature stays
 - `routes/<route>.py` holds the questions one route asks, under the same function names on every route (`response`, `map_failure`, `connection_defaults`), so Rust names only the module. A question two routes ask goes in a top-level module here instead
+
+## model_cost
+
+- `publish`, `write` and `changed_in_place` are the only writers. Each moves `ModelCostVersion`
+- The version is `(id(table), len(table), writes)`, so replacing the dict or changing its size moves it without the host
+  - An in-place value edit does not, so code that edits rows in place calls `changed_in_place()`
+- Anything derived from the table is cached under the version (`@VersionCached`, `model_cost_version()`), never cleared by hand
+- `resolve`, `merged` and `kept` are shims until `get_model_info` and a merge generic over Python values move to Rust. Then their bodies become native calls and the signatures stay
+- `resolve` returns `None` when it was handed a table that is no longer published. That is how a pass learns a reload swapped the table
+- Compat names exist only so code written against main keeps working. New code never calls them, and they never gain behavior of their own
+  - `utils._invalidate_model_cost_lowercase_map` forwards to `changed_in_place`
+  - `utils.get_model_cost_mutation_generation` forwards to `model_cost_version().writes`
+  - `utils._model_cost_lowercase_map` is always `None`

litellm/rust_bridge/host/model_capabilities.py

@@ -6,6 +6,7 @@ from typing import Final, TypeAlias
 
 import litellm
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
+from litellm.rust_bridge.host.model_cost import ModelCostVersion, model_cost_version
 
 
 class _Held:
@@ -21,7 +22,7 @@ class _Held:
         return isinstance(other, _Held) and other.value is self.value
 
 
-_Fingerprint: TypeAlias = tuple[str, str | None, _Held, int, int, tuple[_Held, ...]]
+_Fingerprint: TypeAlias = tuple[str, str | None, _Held, ModelCostVersion, int, tuple[_Held, ...]]
 
 
 def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str, str]:
@@ -33,12 +34,11 @@ def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str
 
 
 def _fingerprint(model: str, custom_llm_provider: str | None, resolved_model: str, provider: str) -> _Fingerprint:
-    """Changes whenever an answer could: ``register_model`` bumps the generation, and replacing
+    """Changes whenever an answer could: every write moves the cost map version, and replacing
     any entry the probes look up (``litellm.model_cost[key] = ...``) changes that entry's identity.
     The key holds the entries it names, so a cached answer's entries are never freed and their
     addresses never reused."""
     from litellm.llms.anthropic.common_utils import AnthropicModelInfo
-    from litellm.utils import get_model_cost_mutation_generation
 
     model_cost: Final = as_str_mapping(litellm.model_cost) or {}  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # the model map is an untyped dict of rows
     keys: Final = dict.fromkeys(
@@ -52,7 +52,7 @@ def _fingerprint(model: str, custom_llm_provider: str | None, resolved_model: st
         model,
         custom_llm_provider,
         _Held(model_cost),
-        get_model_cost_mutation_generation(),
+        model_cost_version(),
         len(model_cost),
         tuple(_Held(model_cost.get(key)) for key in keys),
     )

litellm/rust_bridge/host/model_cost.py

@@ -0,0 +1,148 @@
+"""The live ``litellm.model_cost`` table, every write to it, and its version."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Final, Generic, ParamSpec, TypeVar
+
+from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class AliasedCostEntry(dict[str, object]):
+    """A ``litellm.model_cost`` row, the one object every alias name of a model is served under."""
+
+    __slots__ = ("names", "registered_fields")
+
+    def __init__(
+        self, entry: Mapping[str, object], names: tuple[str, ...], registered_fields: frozenset[str] = frozenset()
+    ) -> None:
+        super().__init__(entry)
+        self.names: tuple[str, ...] = names
+        self.registered_fields: frozenset[str] = registered_fields
+
+
+entry: Final = AliasedCostEntry
+
+
+@dataclass(frozen=True, slots=True)
+class ModelCostVersion:
+    table: int
+    size: int
+    writes: int
+
+
+class _Writes:
+    __slots__ = ("count",)
+
+    def __init__(self) -> None:
+        self.count: int = 0
+
+    def record(self) -> None:
+        self.count += 1
+
+
+_WRITES: Final = _Writes()
+
+
+def model_cost_version() -> ModelCostVersion:
+    import litellm
+
+    table: Final = litellm.model_cost
+    return ModelCostVersion(table=id(table), size=len(table), writes=_WRITES.count)
+
+
+def publish(table: dict[str, object]) -> None:
+    """Makes ``table`` the live ``litellm.model_cost``."""
+    import litellm
+
+    litellm.model_cost = table
+    _WRITES.record()
+
+
+def write(rows: dict[str, object], written: Mapping[str, object]) -> None:
+    """Writes one registration's rows into ``rows``, the published table."""
+    rows.update(written)
+    _WRITES.record()
+
+
+def changed_in_place() -> None:
+    """Moves the version and drops the name index after an in-place edit of ``litellm.model_cost``."""
+    from litellm.litellm_core_utils.get_model_cost_map import model_cost_inputs
+
+    model_cost_inputs.invalidate()
+    _WRITES.record()
+
+
+class VersionCached(Generic[P, R]):
+    """Decorates ``function`` with an ``lru_cache`` keyed on the ``litellm.model_cost`` version each call read, and
+    exposes that cache's ``cache_clear`` and ``cache_info``."""
+
+    __slots__ = ("_at", "cache_clear", "cache_info")
+
+    def __init__(self, function: Callable[P, R]) -> None:
+        def at(
+            _version: ModelCostVersion,
+            *args: P.args,
+            **kwargs: P.kwargs,  # kwargs-ok: forwards a ParamSpec unchanged
+        ) -> R:
+            return function(*args, **kwargs)
+
+        cached: Final = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(at)
+        self._at: Final = cached
+        self.cache_clear: Final = cached.cache_clear
+        self.cache_info: Final = cached.cache_info
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:  # kwargs-ok: forwards a ParamSpec unchanged
+        return self._at(model_cost_version(), *args, **kwargs)
+
+
+@dataclass(frozen=True, slots=True)
+class Landing:
+    key: str
+    base: Mapping[str, object]
+    matched_builtin: bool
+
+
+def published() -> dict[str, object]:
+    """The table published as ``litellm.model_cost`` now."""
+    import litellm
+
+    return litellm.model_cost
+
+
+def resolve(
+    rows: Mapping[str, Mapping[str, object]] | None,
+    key: str,
+    registration: Mapping[str, object],
+    custom_llm_provider: str | None,
+) -> Landing | None:
+    """Where a ``register_model`` entry lands in ``rows`` and the row it starts from, ``None`` when ``rows`` is no
+    longer the published table. ``rows`` of ``None`` is the published table."""
+    import litellm
+    from litellm.utils import registration_base
+
+    if rows is not None and rows is not litellm.model_cost:
+        return None
+    return registration_base(key, registration, custom_llm_provider)
+
+
+def merged(
+    current: Mapping[str, object] | None, base: Mapping[str, object], registration: Mapping[str, object]
+) -> AliasedCostEntry:
+    """The row a registration writes over ``current``, the row now under its landing key."""
+    from litellm.utils import model_cost_registration_write
+
+    return model_cost_registration_write(current, base, registration)
+
+
+def kept(previous: dict[str, object] | None, registration: Mapping[str, object]) -> dict[str, object]:
+    """The registration ``register_model(persist_across_reloads=True)`` keeps for a key: ``registration`` merged
+    into the one kept before it."""
+    from litellm.utils import update_dictionary
+
+    return update_dictionary(dict(previous or {}), registration)

litellm/utils.py

@@ -102,7 +102,11 @@ from litellm.litellm_core_utils.fallback_generalizations import (
     match_capability_generalizations,
     match_fill_missing_generalizations,
 )
-from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetModelCostMap
+from litellm.litellm_core_utils.get_model_cost_map import (
+    GetModelCostMap,
+    fetch_registration_document,
+    model_cost_inputs,
+)
 from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases
 from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired
 from litellm.litellm_core_utils.offline_mode import is_offline_mode
@@ -114,8 +118,16 @@ from litellm.litellm_core_utils.tool_fields import (
     tools_without_hosted_mcp,
 )
 from litellm.rust_bridge import tokenizer as tokenizer_dispatch
+from litellm.rust_bridge._native import NativeModelCostMap, NativeRegistered
 from litellm.rust_bridge.catalog import decision
 from litellm.rust_bridge.configuration import Decision
+from litellm.rust_bridge.host.model_cost import (
+    AliasedCostEntry,
+    Landing,
+    VersionCached,
+    changed_in_place,
+    model_cost_version,
+)
 
 _CachingHandlerResponse = None
 _LLMCachingHandler = None
@@ -3887,22 +3899,6 @@ def _builtin_cost_map() -> dict[str, Mapping[str, object]]:
     return cast(dict[str, Mapping[str, object]], litellm.model_cost)
 
 
-def _replace_cost_map_entry(
-    cost_map: dict[str, Mapping[str, object]], key: str, entry: Mapping[str, object], registration: Mapping[str, object]
-) -> None:
-    current: Final = cost_map.get(key)
-    merged: Final = {**(current or {}), **entry}
-    names: Final = current.names if isinstance(current, AliasedCostEntry) else ()
-    previous_fields: Final[frozenset[str]] = (
-        current.registered_fields if isinstance(current, AliasedCostEntry) else frozenset()
-    )
-    registered_fields: Final = previous_fields | frozenset(
-        name for name, value in registration.items() if value is not None
-    )
-    sharing_names: Final = tuple(name for name in names if cost_map.get(name) is current)
-    cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, names, registered_fields)))
-
-
 def _without_synthesized_costs(
     model_info: Mapping[str, object], key: str, value: Mapping[str, object]
 ) -> Mapping[str, object]:
@@ -4143,48 +4139,47 @@ def _skips_builtin_registration_lookup(key: str, provider: object) -> bool:
     )
 
 
-@dataclass(frozen=True, slots=True)
-class _RegistrationWrite:
-    key: str
-    entry: Mapping[str, object]
-
-
 def model_cost_registration_row(key: str, provider: object) -> Mapping[str, object]:
     """The row a ``register_model`` write of a ``key`` entry with this ``litellm_provider`` merges onto: the
     ``litellm.model_cost`` row it resolves to (``openai/o3-pro`` lands on ``o3-pro``), else the row it seeds ``key``
     with from the built-in entry it falls back to (an OpenAI fine-tune takes its base model's row)."""
-    base: Final = _registration_base(
-        key, {"litellm_provider": provider}, custom_llm_provider=None, warning_display_name=None
-    )
-    return _builtin_cost_map().get(base.key) or base.entry
+    base: Final = registration_base(key, {"litellm_provider": provider}, custom_llm_provider=None)
+    return _builtin_cost_map().get(base.key) or base.base
 
 
-def _registration_base(
-    key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
-) -> _RegistrationWrite:
+def registration_base(key: str, value: Mapping[str, object], custom_llm_provider: str | None) -> Landing:
     provider: Final = str(value.get("litellm_provider") or "")
     if _skips_builtin_registration_lookup(key, provider):
-        return _RegistrationWrite(key=key, entry=dict(_builtin_cost_map().get(key, {})))
+        return Landing(key=key, base=dict(_builtin_cost_map().get(key, {})), matched_builtin=True)
     builtin_model_info: Final = _get_builtin_model_info_for_registration(
         model=key, custom_llm_provider=custom_llm_provider
     )
     if not isinstance(builtin_model_info, ModelInfoLookupMiss):
-        return _RegistrationWrite(
+        return Landing(
             key=builtin_model_info["key"],
-            entry=dict(_without_synthesized_costs(builtin_model_info, key=key, value=value)),
+            base=dict(_without_synthesized_costs(builtin_model_info, key=key, value=value)),
+            matched_builtin=True,
         )
     fallback: Final = _fallback_registration(key=key, provider=provider, value=value)
-    if not fallback.matched_builtin and (
+    return Landing(key=key, base=dict(fallback.entry), matched_builtin=fallback.matched_builtin)
+
+
+def _prices_tokens_without_cache_pricing(value: Mapping[str, object]) -> bool:
+    return (
         value.get("cache_creation_input_token_cost") is None
         and value.get("cache_read_input_token_cost") is None
         and value.get("tiered_pricing") is None
         and (value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None)
-    ):
-        verbose_logger.warning(
-            "register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info",
-            warning_display_name or key,
-        )
-    return _RegistrationWrite(key=key, entry=dict(fallback.entry))
+    )
+
+
+def _warn_unpriced_cache(unmatched: Mapping[str, Mapping[str, object]], warning_display_name: str | None) -> None:
+    for key, registration in unmatched.items():
+        if _prices_tokens_without_cache_pricing(registration):
+            verbose_logger.warning(
+                "register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info",
+                warning_display_name or key,
+            )
 
 
 def native_structured_output_support(model_info: Mapping[str, object]) -> bool | None:
@@ -4203,104 +4198,65 @@ def _base_fields_replaced_by(value: Mapping[str, object]) -> frozenset[str]:
     )
 
 
-def _registration_write(
-    key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
-) -> _RegistrationWrite:
-    base: Final = _registration_base(key, value, custom_llm_provider, warning_display_name)
+def model_cost_registration_write(
+    current: Mapping[str, object] | None, base: Mapping[str, object], registration: Mapping[str, object]
+) -> AliasedCostEntry:
+    """The row a ``register_model`` entry writes over ``current``."""
     # ``get_model_info`` returns ``litellm_provider: None`` when the
     # provider is unknown (e.g. custom deployments registered via
     # ``Router.add_deployment``). Persisting that None into
     # ``litellm.model_cost`` causes ``_check_provider_match`` to drop
     # custom pricing on subsequent cost lookups.
     existing_model: Final = {
-        name: field_value
-        for name, field_value in base.entry.items()
-        if name != "litellm_provider" or field_value is not None
+        name: field_value for name, field_value in base.items() if name != "litellm_provider" or field_value is not None
     }
-    return _RegistrationWrite(
-        key=base.key,
-        entry={**update_dictionary(existing_model, value), **dict.fromkeys(_base_fields_replaced_by(value))},
+    entry: Final = {
+        **update_dictionary(existing_model, registration),
+        **dict.fromkeys(_base_fields_replaced_by(registration)),
+    }
+    names: Final = current.names if isinstance(current, AliasedCostEntry) else ()
+    previous_fields: Final[frozenset[str]] = (
+        current.registered_fields if isinstance(current, AliasedCostEntry) else frozenset()
     )
-
-
-def _registration_attempt(
-    key: str,
-    value: Mapping[str, object],
-    custom_llm_provider: str | None,
-    warning_display_name: str | None,
-    final_attempt: bool,
-) -> _RegistrationWrite | None:
-    source: Final = _builtin_cost_map()
-    write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name)
-    _replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry, value)
-    return write if final_attempt or litellm.model_cost is source else None
+    registered_fields: Final = previous_fields | frozenset(
+        name for name, value in registration.items() if value is not None
+    )
+    return AliasedCostEntry({**(current or {}), **entry}, names, registered_fields)
 
 
 def _register_on_current_map(
-    key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
-) -> _RegistrationWrite:
-    """A price data reload that swaps the map while the entry is being built or written would otherwise get the old
-    map's prices written over its own, or drop the entry with the old map, so each attempt writes into the map it was
-    built from and the entry is rebuilt when a reload replaced that map before the write finished. The last attempt
-    writes into the current map, so the entry is never lost."""
-    attempts: Final = (
-        _registration_attempt(
-            key, value, custom_llm_provider, warning_display_name, final_attempt=attempt == _REGISTRATION_ATTEMPTS - 1
-        )
-        for attempt in range(_REGISTRATION_ATTEMPTS)
-    )
-    return next(write for write in attempts if write is not None)
-
-
-_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {}  # mutable-ok: replayed on reload
-_runtime_registration_order: Final[dict[str, int]] = {}  # mutable-ok: tells which of a row's registered names came last
-_runtime_registered_names: Final[dict[str, tuple[str, ...]]] = {}  # mutable-ok: names registered onto each live row
-_REGISTRATION_TICKS: Final = itertools.count()
+    entries: Mapping[str, Mapping[str, object]], persist: bool, custom_llm_provider: str | None
+) -> Mapping[str, Mapping[str, object]]:
+    """Writes ``entries`` into the live table and returns those no built-in row matched."""
+    for _ in range(_REGISTRATION_ATTEMPTS - 1):
+        live = litellm.model_cost
+        written = model_cost_inputs.register(live, entries, persist, custom_llm_provider)
+        if isinstance(written, NativeRegistered) and litellm.model_cost is live:
+            return written.unmatched
+    last: Final = model_cost_inputs.register(None, entries, persist, custom_llm_provider)
+    return last.unmatched if isinstance(last, NativeRegistered) else {}
 
 
 @dataclass(frozen=True, slots=True)
-class _RegistrationsOverCatalog(Mapping[str, Mapping[str, object]]):
-    catalog: Mapping[str, Mapping[str, object]]
-    registrations: Mapping[str, Mapping[str, object]]
-    registration_order: Mapping[str, int]
-    registered_names: Mapping[str, tuple[str, ...]]
+class _RegisteredCatalog(Mapping[str, Mapping[str, object]]):
+    inputs: NativeModelCostMap
 
     def __getitem__(self, key: str) -> Mapping[str, object]:
-        catalog_entry: Final = self.catalog.get(key)
-        names: Final = dict.fromkeys((*self.registered_names.get(key, ()), key))
-        present: Final = tuple(name for name in names if name in self.registrations)
-        if not present:
-            return self.catalog[key]
-        oldest_first: Final = sorted(present, key=lambda name: self.registration_order.get(name, -1))
-        registered: Final = (self.registrations[name] for name in oldest_first)
-        return MappingProxyType(
-            dict(itertools.chain.from_iterable(entry.items() for entry in (catalog_entry or {}, *registered)))
-        )
+        row: Final = self.inputs.registered_row(key)
+        if row is None:
+            raise KeyError(key)
+        return MappingProxyType(row)
 
     def __iter__(self) -> Iterator[str]:
-        return iter({**dict.fromkeys(self.catalog), **dict.fromkeys(self.registrations)})
+        return iter(self.inputs.registered_keys())
 
     def __len__(self) -> int:
-        return len(self.catalog.keys() | self.registrations.keys())
-
-
-def _record_registered_name(written_key: str, registered_key: str) -> None:
-    names: Final = _runtime_registered_names.get(written_key, ())
-    if registered_key not in names:
-        _runtime_registered_names[written_key] = (*names, registered_key)
+        return len(self.inputs.registered_keys())
 
 
 def catalog_with_runtime_registrations() -> Mapping[str, Mapping[str, object]]:
     """The loaded catalog plus ``register_model`` calls, without what routers register for their deployments."""
-    catalog: Final = GetModelCostMap.loaded_model_cost_map()
-    if not _runtime_registered_model_cost:
-        return catalog
-    return _RegistrationsOverCatalog(
-        catalog,
-        MappingProxyType(_runtime_registered_model_cost),
-        MappingProxyType(_runtime_registration_order),
-        MappingProxyType(_runtime_registered_names),
-    )
+    return _RegisteredCatalog(model_cost_inputs)
 
 
 def runtime_model_cost_registrations(
@@ -4374,15 +4330,9 @@ def reapply_runtime_model_cost_registrations() -> None:
         _LiveDeploymentReplay.callback()
     for extra_replay in _LiveDeploymentReplay.extra:
         extra_replay()
-    if _runtime_registered_model_cost:
-        register_model(
-            model_cost=dict(
-                sorted(
-                    _runtime_registered_model_cost.items(),
-                    key=lambda registration: _runtime_registration_order.get(registration[0], -1),
-                )
-            )
-        )
+    reapplied: Final = model_cost_inputs.reapply(litellm.model_cost)
+    if isinstance(reapplied, NativeRegistered):
+        _warn_unpriced_cache(reapplied.unmatched, None)
 
 
 def cost_map_omits_token_price(*keys: object) -> bool:
@@ -4442,34 +4392,19 @@ def register_model(
     keys are opaque ids rather than model names (e.g. router deployment ids).
     """
 
-    loaded_model_cost = {}
-    if isinstance(model_cost, dict):
-        # Convert stringified numbers to appropriate numeric types
-        loaded_model_cost = model_cost
-    elif isinstance(model_cost, str):
-        loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1)
-
-    if persist_across_reloads:
-        _registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost
-        for _registered_key, _registered_value in _registrations.items():
-            _runtime_registered_model_cost[_registered_key] = update_dictionary(
-                dict(_runtime_registered_model_cost.get(_registered_key, {})),
-                _RATE_TABLE.validate_python(_registered_value),
-            )
-            _runtime_registration_order[_registered_key] = next(_REGISTRATION_TICKS)
-
-    for key, value in loaded_model_cost.items():
-        write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name)
-        if persist_across_reloads:
-            _record_registered_name(write.key, str(key))
-
-        # Invalidate case-insensitive lookup map since model_cost was modified
-        _invalidate_model_cost_lowercase_map()
-
-        verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", write.key, write.key)
-        if add_to_provider_model_sets and value.get("litellm_provider") == "openrouter":
-            litellm.openrouter_models.add(str(key).split("/", 1)[-1])
+    loaded_model_cost: Final = (
+        model_cost if isinstance(model_cost, dict) else (fetch_registration_document(model_cost) or {})
+    )
+    entries: Final = {
+        str(key): dict(_RATE_TABLE.validate_python(value) if persist_across_reloads else value)
+        for key, value in loaded_model_cost.items()
+    }
+    unmatched: Final = _register_on_current_map(entries, persist_across_reloads, custom_llm_provider)
+    _warn_unpriced_cache(unmatched, warning_display_name)
     if add_to_provider_model_sets:
+        litellm.openrouter_models.update(
+            key.split("/", 1)[-1] for key, value in entries.items() if value.get("litellm_provider") == "openrouter"
+        )
         litellm.add_known_models(loaded_model_cost)
     return model_cost
 
@@ -6789,123 +6724,24 @@ def _strip_model_name(model: str, custom_llm_provider: str | None) -> str:
         return _strip_dated_snapshot_suffix(model_name=model)
 
 
-# Global case-insensitive lookup map for model_cost (built eagerly at module import)
-_model_cost_lowercase_map: dict[str, str] | None = None
+_model_cost_lowercase_map: Final[None] = None
 
-# Monotonic counter bumped on every model_cost mutation. Consumers that
-# memoize derived state (e.g. provider-specific indices) can include this
-# value in their cache key so they invalidate even when key add+remove or
-# in-place value replacement leaves len/id unchanged.
-_model_cost_mutation_generation: int = 0
+
+def _invalidate_model_cost_lowercase_map() -> None:  # pyright: ignore[reportUnusedFunction]  # downstream code calls main's name
+    """Compat for main's callers. Use ``changed_in_place()``."""
+    changed_in_place()
 
 
 def get_model_cost_mutation_generation() -> int:
-    return _model_cost_mutation_generation
-
-
-def _invalidate_model_cost_lowercase_map() -> None:
-    """Invalidate the case-insensitive lookup map for model_cost.
-
-    Call this whenever litellm.model_cost is modified to ensure the map is rebuilt.
-    Also clears related LRU caches that depend on model_cost data.
-    """
-    global _model_cost_lowercase_map, _model_cost_mutation_generation
-    _model_cost_lowercase_map = None
-    _model_cost_mutation_generation += 1
-
-    # Clear LRU caches that depend on model_cost data
-    _cached_get_model_info.cache_clear()
-    cached_get_model_info_helper.cache_clear()
-    cached_cost_map_key.cache_clear()
-
-
-def _rebuild_model_cost_lowercase_map() -> dict[str, str]:
-    """Rebuild the case-insensitive lookup map from the current model_cost.
-
-    Returns:
-        The rebuilt map (guaranteed to be not None).
-    """
-    global _model_cost_lowercase_map
-    _model_cost_lowercase_map = {k.lower(): k for k in tuple(litellm.model_cost)}
-    return _model_cost_lowercase_map
-
-
-def _handle_stale_map_entry_rebuild(
-    potential_key_lower: str,
-) -> str | None:
-    """
-    Handle stale _model_cost_lowercase_map entry (key was popped).
-
-    Rebuilds the map and retries the lookup.
-
-    Returns:
-        The matched key if found after rebuild, None otherwise.
-    """
-    matched_key: Final = _rebuild_model_cost_lowercase_map().get(potential_key_lower)
-    if matched_key is not None and matched_key in litellm.model_cost:
-        return matched_key
-    return None
-
-
-def _handle_new_key_with_scan(
-    potential_key_lower: str,
-) -> str | None:
-    """
-    Handle new key added to model_cost without invalidating _model_cost_lowercase_map.
-
-    Scans model_cost for case-insensitive match and rebuilds the map if found.
-
-    Returns:
-        The matched key if found, None otherwise.
-    """
-    global _model_cost_lowercase_map
-    for key in tuple(_builtin_cost_map()):
-        if key.lower() == potential_key_lower:
-            _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map()
-            return key
-    return None
+    """Compat for main's callers. Use ``model_cost_version()``."""
+    return model_cost_version().writes
 
 
 def get_model_cost_key(potential_key: str) -> str | None:
-    """
-    Get the actual key from model_cost, with case-insensitive fallback.
-
-    WARNING: Only O(1) lookup operations are acceptable. O(n) lookups will cause severe
-    CPU overhead. This function is called frequently during router operations.
-
-    ALLOWED HELPER FUNCTIONS (conditionally called, O(n) operations are acceptable):
-    - _rebuild_model_cost_lowercase_map: Rebuilds the lookup map (only when map is None)
-    - _handle_stale_map_entry_rebuild: Rebuilds map when stale entry detected (rare case)
-
-    If you need to add a new helper function with O(n) operations that is conditionally
-    called and confirmed not to cause performance issues, add it to the allowed_helpers
-    list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py
-    """
-    # Exact match (O(1))
+    """The key itself when ``litellm.model_cost`` has it, otherwise the key that differs from it only in case."""
     if potential_key in litellm.model_cost:
         return potential_key
-
-    # Case-insensitive lookup via map (O(1)); bound locally because another thread's
-    # registration can invalidate the module-level map at any point
-    cached_lowercase_map: Final = _model_cost_lowercase_map
-    lowercase_map: Final = (
-        cached_lowercase_map if cached_lowercase_map is not None else _rebuild_model_cost_lowercase_map()
-    )
-
-    potential_key_lower: Final = potential_key.lower()
-    matched_key = lowercase_map.get(potential_key_lower)
-
-    # Verify key exists (O(1) - handles model_cost.pop() case)
-    if matched_key is not None and matched_key in litellm.model_cost:
-        return matched_key
-
-    # Rebuild map if stale entry detected (O(n) rebuild, but only when stale entry found)
-    if matched_key is not None:
-        matched_key = _handle_stale_map_entry_rebuild(potential_key_lower)
-        if matched_key is not None:
-            return matched_key
-
-    return None
+    return model_cost_inputs.find_key(litellm.model_cost, potential_key)
 
 
 _get_model_cost_key = get_model_cost_key
@@ -7169,7 +7005,7 @@ def _get_max_position_embeddings(model_name: str) -> int | None:
         return None
 
 
-@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
+@VersionCached
 def cached_get_model_info_helper(
     model: str,
     custom_llm_provider: str | None,
@@ -7177,11 +7013,6 @@ def cached_get_model_info_helper(
     *,
     preserve_unset_costs: bool = False,
 ) -> ModelInfoBase:
-    """
-    _get_model_info_helper wrapped with lru_cache
-
-    Speed Optimization to hit high RPS
-    """
     return get_model_info_helper(
         model=model,
         custom_llm_provider=custom_llm_provider,
@@ -7193,7 +7024,7 @@ def cached_get_model_info_helper(
 _cached_get_model_info_helper = cached_get_model_info_helper
 
 
-@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
+@VersionCached
 def cached_cost_map_key(model: str, custom_llm_provider: str | None) -> str | None:
     try:
         return cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)["key"]
@@ -7753,7 +7584,7 @@ def _build_model_info(
     return ModelInfo(**_model_info, supported_openai_params=supported_openai_params)
 
 
-@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
+@VersionCached
 def _cached_get_model_info(
     model: str,
     custom_llm_provider: str | None = None,

tests/code_coverage_tests/check_get_model_cost_key_performance.py

@@ -151,12 +151,7 @@ def check_get_model_cost_key_performance():
 
         # Check for for loops
         if re.search(r"\bfor\s+\w+\s+in\s+", line):
-            # Allow helper function calls (they're conditional)
-            if not re.search(
-                r"(_rebuild_model_cost_lowercase_map|_handle_stale_map_entry_rebuild|_handle_new_key_with_scan)",
-                line,
-            ):
-                problematic_lines.append((i, "for loop", line_stripped))
+            problematic_lines.append((i, "for loop", line_stripped))
 
         # Check for while loops
         if re.search(r"\bwhile\s+", line):
@@ -187,11 +182,7 @@ def check_get_model_cost_key_performance():
 
         # Check for calls to functions that might have O(n) operations
         # Allow known helper functions that are conditional
-        allowed_helpers = [
-            "_rebuild_model_cost_lowercase_map",
-            "_handle_stale_map_entry_rebuild",
-            "_handle_new_key_with_scan",
-        ]
+        allowed_helpers: list[str] = []
 
         # Check for function calls (pattern: function_name(...), but not function definitions)
         # Skip function definitions (def function_name(...))

tests/integration/pricing/test_price_data_reload_prices.py

@@ -5,16 +5,22 @@ import threading
 import uuid
 from collections.abc import Callable, Mapping
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import AbstractContextManager
 from dataclasses import dataclass
 from pathlib import Path
 from types import MappingProxyType
-from typing import Final
+from typing import Final, TypeAlias
 
 import httpx
 import pytest
-from integration._support.client import Gateway, eventually, object_value, string_value
-from integration._support.database import read_rows
-from integration._support.process import graceful_stop_seconds, owned_proxy
+from integration._support.client import Gateway, eventually, held, object_value, string_value
+from integration._support.database import read_rows, scratch_database
+from integration._support.process import (
+    OwnedProxy,
+    graceful_stop_seconds,
+    owned_proxy,
+    owned_proxy_process,
+)
 from integration._support.wire import Reply, Request, wire_server
 from pydantic import JsonValue
 
@@ -84,6 +90,25 @@ _CATALOG: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType(
 )
 
 
+@dataclass(frozen=True, slots=True)
+class _Published:
+    scale: float
+    added: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class _Failing:
+    status: int
+
+
+_Served: TypeAlias = _Published | _Failing
+
+
+class _CostMapServer:
+    def __init__(self, serving: _Served = _Published(1.0)) -> None:
+        self.serving: _Served = serving
+
+
 @dataclass(frozen=True, slots=True)
 class _Usage:
     prompt: int
@@ -123,8 +148,9 @@ class _Charge:
     expected: float
 
 
-def _catalog(scale: float) -> dict[str, JsonValue]:
+def _catalog(scale: float, added: tuple[str, ...] = ()) -> dict[str, JsonValue]:
     return {
+        **{model: {"litellm_provider": "openai", "mode": "chat"} for model in added},
         **{
             model: {name: value * scale if isinstance(value, float) else value for name, value in row.items()}
             for model, row in _CATALOG.items()
@@ -189,8 +215,13 @@ def _connection(model: str, wire_url: str) -> dict[str, JsonValue]:
     return {"api_base": f"{wire_url}/v1", "api_key": "synthetic-openai-key"}
 
 
-def _config(directory: Path, deployments: tuple[_Deployment, ...], wire_url: str) -> Path:
-    config: Final = directory / "price_data_reload_config.yaml"
+def _config(
+    directory: Path,
+    deployments: tuple[_Deployment, ...],
+    wire_url: str,
+    general: Mapping[str, JsonValue] = MappingProxyType({}),
+) -> Path:
+    config: Final = directory / f"price_data_reload_config_{uuid.uuid4().hex}.yaml"
     config.write_text(
         json.dumps(
             {
@@ -208,6 +239,7 @@ def _config(directory: Path, deployments: tuple[_Deployment, ...], wire_url: str
                 "general_settings": {
                     "master_key": "os.environ/LITELLM_MASTER_KEY",
                     "database_url": "os.environ/DATABASE_URL",
+                    **general,
                 },
                 "router_settings": {"disable_cooldowns": True},
             }
@@ -276,10 +308,18 @@ def _reply(request: Request) -> Reply:
     )
 
 
-def _responder(published: threading.Event) -> Callable[[Request], Reply]:
+def _cost_map_reply(served: _Served) -> Reply:
+    match served:
+        case _Published(scale=scale, added=added):
+            return Reply(body=json.dumps(_catalog(scale, added)).encode())
+        case _Failing(status=status):
+            return Reply(status=status, body=b"{}")
+
+
+def _responder(server: _CostMapServer) -> Callable[[Request], Reply]:
     def respond(request: Request) -> Reply:
         if request.target == "/model_prices.json":
-            return Reply(body=json.dumps(_catalog(_PUBLISHED_SCALE if published.is_set() else 1.0)).encode())
+            return _cost_map_reply(server.serving)
         if request.method == "GET":
             return Reply(body=json.dumps({"object": "list", "data": []}).encode())
         return _reply(request)
@@ -340,9 +380,9 @@ def _proxy_environment(wire_url: str) -> dict[str, str]:
 
 
 def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments: tuple[_Deployment, ...]) -> None:
-    published: Final = threading.Event()
+    server: Final = _CostMapServer()
     with (
-        wire_server(_responder(published)) as wire,
+        wire_server(_responder(server)) as wire,
         owned_proxy(
             gateway,
             tmp_path,
@@ -352,7 +392,7 @@ def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments
         ) as proxy,
     ):
         before: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in deployments)
-        published.set()
+        server.serving = _Published(_PUBLISHED_SCALE)
         reload: Final = proxy.request("POST", "/reload/model_cost_map")
         assert reload.status_code == 200, reload.text
         after: Final = tuple(_charged_as_listed(proxy, deployment, _PUBLISHED_SCALE) for deployment in deployments)
@@ -361,72 +401,87 @@ def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments
 
 
 @pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
-def test_price_data_reload_reprices_catalog_deployments_and_keeps_custom_prices(
+def test_custom_prices_on_one_deployment_never_reach_a_sibling_on_the_same_backend_across_a_reload(
     gateway: Gateway, tmp_path: Path
 ) -> None:
     _assert_priced_across_a_reload(
         gateway,
         tmp_path,
         (
-            _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)),
             _Deployment("mini-custom", f"openai/{_MINI}", _CUSTOM, _Usage(1000, 0, 500), _custom_rates),
+            _Deployment("mini-catalog", f"openai/{_MINI}", {}, _Usage(1000, 0, 500), _catalog_rates(_MINI)),
         ),
     )
 
 
+def _catalog_tier_rates(model: str) -> Callable[[float], _Rates]:
+    return lambda scale: _Rates(
+        input=_rate(model, "input_cost_per_token_above_200k_tokens", scale),
+        output=_rate(model, "output_cost_per_token_above_200k_tokens", scale),
+        cache_read=_rate(model, "cache_read_input_token_cost", scale),
+    )
+
+
 @pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
-@pytest.mark.parametrize(
-    "deployment",
-    (
-        _Deployment(
-            "region-unpriced", f"bedrock/ap-northeast-1.{_ROW}", {}, _Usage(1000, 0, 500), _catalog_rates(_ROW)
-        ),
-        _Deployment(
-            "finetune-unpriced",
-            f"openai/{_FINETUNE}:org:suffix:y",
-            {},
-            _Usage(2000, 1000, 1000),
-            _catalog_rates(_FINETUNE),
-        ),
-        _Deployment(
-            "finetune-custom", f"openai/{_FINETUNE}:org:suffix:z", _CUSTOM, _Usage(2000, 0, 1000), _custom_rates
-        ),
-        _Deployment("cross-provider-unpriced", f"azure/{_CLAUDE}", {}, _Usage(2000, 1000, 1000), _unpriced),
-        _Deployment("catalog-free", f"openai/{_FREE}", {}, _Usage(2000, 1000, 1000), _catalog_rates(_FREE)),
-        _Deployment("rule-free", f"openai/{_RULE_FREE}", {}, _Usage(2000, 1000, 1000), _catalog_rates(_FREE)),
-        _Deployment(
-            "fill-rule-free",
-            f"openai/{_FILL_FREE}",
-            {},
-            _Usage(2000, 1000, 1000),
-            _catalog_rates(_FREE),
-            listed_rates=_unpriced,
-        ),
-        _Deployment("rule-unpriced", f"openai/{_RULE_UNPRICED}", {}, _Usage(2000, 1000, 1000), _unpriced),
-        _Deployment(
-            "provider-reported-free", f"openrouter/{_REPORTED_FREE}", {}, _Usage(2000, 1000, 1000), _unpriced, 0.0
-        ),
-        _Deployment(
-            "cross-provider-free",
-            f"azure/{_CLAUDE}",
-            {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0},
-            _Usage(2000, 1000, 1000),
-            _catalog_rates(_FREE),
-        ),
-        _Deployment(
-            "region-custom-above-threshold",
-            f"bedrock/ap-northeast-1.{_ROW}",
-            _CUSTOM,
-            _Usage(250_000, 0, 1000),
-            _custom_rates,
-        ),
-    ),
-    ids=lambda deployment: deployment.name,
-)
-def test_deployment_bills_the_catalog_row_its_registration_resolves_to_across_a_reload(
-    gateway: Gateway, tmp_path: Path, deployment: _Deployment
+def test_every_deployment_bills_the_catalog_row_its_registration_resolves_to_across_a_reload(
+    gateway: Gateway, tmp_path: Path
 ) -> None:
-    _assert_priced_across_a_reload(gateway, tmp_path, (deployment,))
+    _assert_priced_across_a_reload(
+        gateway,
+        tmp_path,
+        (
+            _Deployment(
+                "region-unpriced", f"bedrock/ap-northeast-1.{_ROW}", {}, _Usage(1000, 0, 500), _catalog_rates(_ROW)
+            ),
+            _Deployment(
+                "finetune-unpriced",
+                f"openai/{_FINETUNE}:org:suffix:y",
+                {},
+                _Usage(2000, 1000, 1000),
+                _catalog_rates(_FINETUNE),
+            ),
+            _Deployment(
+                "finetune-custom", f"openai/{_FINETUNE}:org:suffix:z", _CUSTOM, _Usage(2000, 0, 1000), _custom_rates
+            ),
+            _Deployment("cross-provider-unpriced", f"azure/{_CLAUDE}", {}, _Usage(2000, 1000, 1000), _unpriced),
+            _Deployment("catalog-free", f"openai/{_FREE}", {}, _Usage(2000, 1000, 1000), _catalog_rates(_FREE)),
+            _Deployment("rule-free", f"openai/{_RULE_FREE}", {}, _Usage(2000, 1000, 1000), _catalog_rates(_FREE)),
+            _Deployment(
+                "fill-rule-free",
+                f"openai/{_FILL_FREE}",
+                {},
+                _Usage(2000, 1000, 1000),
+                _catalog_rates(_FREE),
+                listed_rates=_unpriced,
+            ),
+            _Deployment("rule-unpriced", f"openai/{_RULE_UNPRICED}", {}, _Usage(2000, 1000, 1000), _unpriced),
+            _Deployment(
+                "provider-reported-free", f"openrouter/{_REPORTED_FREE}", {}, _Usage(2000, 1000, 1000), _unpriced, 0.0
+            ),
+            _Deployment(
+                "cross-provider-free",
+                f"azure/{_CLAUDE}",
+                {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0},
+                _Usage(2000, 1000, 1000),
+                _catalog_rates(_FREE),
+            ),
+            _Deployment(
+                "region-custom-above-threshold",
+                f"bedrock/ap-northeast-1.{_ROW}",
+                _CUSTOM,
+                _Usage(250_000, 0, 1000),
+                _custom_rates,
+            ),
+            _Deployment(
+                "region-catalog-above-threshold",
+                f"bedrock/ap-northeast-1.{_ROW}",
+                {},
+                _Usage(250_000, 0, 1000),
+                _catalog_tier_rates(_ROW),
+                listed_rates=_catalog_rates(_ROW),
+            ),
+        ),
+    )
 
 
 def _storm_deployment(index: int) -> _Deployment:
@@ -467,7 +522,7 @@ def test_price_data_reloads_under_custom_priced_traffic_keep_every_deployment_pr
     geo: Final = _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO))
     finished: Final = threading.Event()
     with (
-        wire_server(_responder(threading.Event())) as wire,
+        wire_server(_responder(_CostMapServer())) as wire,
         owned_proxy(
             gateway,
             tmp_path,
@@ -492,3 +547,153 @@ def test_price_data_reloads_under_custom_priced_traffic_keep_every_deployment_pr
         settled: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in (geo, *storm))
         for charge in settled:
             assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6)
+
+
+_GEO_CATALOG: Final = _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO))
+_POLL_SECONDS: Final = 1
+_LIFECYCLE_SETTINGS: Final[Mapping[str, JsonValue]] = MappingProxyType(
+    {
+        "proxy_config_reload_interval_seconds": _POLL_SECONDS,
+    }
+)
+
+
+def _lifecycle_proxy(
+    gateway: Gateway,
+    tmp_path: Path,
+    wire_url: str,
+    deployments: tuple[_Deployment, ...],
+    database: Mapping[str, str] = MappingProxyType({}),
+) -> AbstractContextManager[OwnedProxy]:
+    return owned_proxy_process(
+        gateway,
+        tmp_path,
+        {**_proxy_environment(wire_url), **database},
+        config=_config(tmp_path, deployments, wire_url, _LIFECYCLE_SETTINGS),
+        remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",),
+    )
+
+
+def _bill(proxy: Gateway, deployment: _Deployment) -> float | None:
+    return _billed(proxy.client, proxy.key, deployment)[1]
+
+
+def _reload(proxy: Gateway) -> httpx.Response:
+    return proxy.request("POST", "/reload/model_cost_map")
+
+
+def _billed_at(deployment: _Deployment, scale: float) -> Callable[[float | None], bool]:
+    expected: Final = _expected_cost(deployment.rates(scale), deployment.usage)
+    return lambda charged: charged == pytest.approx(expected, rel=1e-6)
+
+
+@pytest.mark.timeout(4 * graceful_stop_seconds() + 180)
+def test_price_data_changes_only_through_a_successful_reload_or_a_restart(gateway: Gateway, tmp_path: Path) -> None:
+    server: Final = _CostMapServer()
+    with wire_server(_responder(server)) as wire, scratch_database("reload_lifecycle") as database_url:
+        database: Final = {"DATABASE_URL": database_url}
+        with _lifecycle_proxy(gateway, tmp_path, wire.url, (_GEO_CATALOG,), database=database) as first:
+            _charged_as_listed(first.gateway, _GEO_CATALOG, 1.0)
+            server.serving = _Published(_PUBLISHED_SCALE)
+            held(
+                lambda: _bill(first.gateway, _GEO_CATALOG),
+                _billed_at(_GEO_CATALOG, 1.0),
+                holding=3 * _POLL_SECONDS,
+            )
+            server.serving = _Failing(404)
+            failed: Final = _reload(first.gateway)
+            assert failed.status_code == 502, failed.text
+            _charged_as_listed(first.gateway, _GEO_CATALOG, 1.0)
+            server.serving = _Published(_PUBLISHED_SCALE)
+        with _lifecycle_proxy(gateway, tmp_path, wire.url, (_GEO_CATALOG,), database=database) as restarted:
+            _charged_as_listed(restarted.gateway, _GEO_CATALOG, _PUBLISHED_SCALE)
+
+
+def _add_model(proxy: Gateway, deployment: _Deployment, wire_url: str) -> None:
+    created: Final = proxy.post(
+        "/model/new",
+        {
+            "model_name": deployment.name,
+            "litellm_params": {
+                "model": deployment.model,
+                **_connection(deployment.model, wire_url),
+                **deployment.prices,
+            },
+            "model_info": {},
+        },
+    )
+    assert object_value(created["model_info"])["id"], created
+
+
+def _custom_with_catalog_cache(model: str) -> Callable[[float], _Rates]:
+    return lambda scale: _Rates(
+        input=_CUSTOM["input_cost_per_token"],
+        output=_CUSTOM["output_cost_per_token"],
+        cache_read=_rate(model, "cache_read_input_token_cost", scale),
+    )
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 180)
+def test_a_model_added_through_the_api_keeps_its_prices_and_inherits_the_reloaded_ones(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    added: Final = _Deployment(
+        "added-finetune-custom",
+        f"openai/{_FINETUNE}",
+        _CUSTOM,
+        _Usage(2000, 1000, 1000),
+        _custom_with_catalog_cache(_FINETUNE),
+    )
+    server: Final = _CostMapServer()
+    with (
+        wire_server(_responder(server)) as wire,
+        _lifecycle_proxy(gateway, tmp_path, wire.url, (_GEO_CATALOG,)) as owned,
+    ):
+        proxy: Final = owned.gateway
+        _add_model(proxy, added, wire.url)
+        before: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in (_GEO_CATALOG, added))
+        server.serving = _Published(_PUBLISHED_SCALE)
+        reload: Final = _reload(proxy)
+        assert reload.status_code == 200, reload.text
+        after: Final = tuple(
+            _charged_as_listed(proxy, deployment, _PUBLISHED_SCALE) for deployment in (_GEO_CATALOG, added)
+        )
+        for charge in (*before, *after):
+            assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6)
+
+
+_WILDCARD: Final = _Deployment("openai/*", "openai/*", {}, _Usage(1000, 0, 500), _unpriced)
+_LISTED_AFTER_RELOAD: Final = "gpt-audit-reload-listed"
+
+
+def _listed_models(proxy: Gateway) -> list[str]:
+    data: Final = proxy.get("/v1/models")["data"]
+    assert isinstance(data, list), data
+    return [string_value(object_value(entry)["id"]) for entry in data]
+
+
+def _deployment_ids(proxy: Gateway) -> list[str]:
+    entries: Final = proxy.get("/model/info")["data"]
+    assert isinstance(entries, list), entries
+    return [string_value(object_value(object_value(entry)["model_info"])["id"]) for entry in entries]
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 180)
+def test_a_wildcard_lists_the_models_a_reload_adds_and_never_a_deployment_id(gateway: Gateway, tmp_path: Path) -> None:
+    custom: Final = _Deployment("listed-custom", f"openai/{_MINI}", _CUSTOM, _Usage(1000, 0, 500), _custom_rates)
+    server: Final = _CostMapServer()
+    with (
+        wire_server(_responder(server)) as wire,
+        _lifecycle_proxy(gateway, tmp_path, wire.url, (_WILDCARD, custom)) as owned,
+    ):
+        proxy: Final = owned.gateway
+        _charged_as_listed(proxy, custom, 1.0)
+        assert f"openai/{_LISTED_AFTER_RELOAD}" not in _listed_models(proxy)
+        server.serving = _Published(1.0, added=(_LISTED_AFTER_RELOAD,))
+        reload: Final = _reload(proxy)
+        assert reload.status_code == 200, reload.text
+        _charged_as_listed(proxy, custom, 1.0)
+        listed: Final = _listed_models(proxy)
+        deployment_ids: Final = _deployment_ids(proxy)
+        assert f"openai/{_LISTED_AFTER_RELOAD}" in listed, listed
+        assert [model for model in listed if any(id_ in model for id_ in deployment_ids)] == [], listed

tests/integration/providers/test_gemini_long_decode_wire.py

@@ -1213,7 +1213,6 @@ def test_request_with_no_output_cap_is_never_continued_and_returns_googles_cut_w
 _NOT_A_NUMBER: Final[tuple[tuple[str, Mapping[str, JsonValue], int | None], ...]] = (
     ("n-string", {"n": "1"}, CAP),
     ("n-list", {"n": [1]}, CAP),
-    ("n-empty", {"n": ""}, CAP),
     ("n-5kb", {"n": "1" * 5120}, CAP),
     ("max_tokens-string", {"max_tokens": "300"}, None),
     ("max_tokens-list", {"max_tokens": [300]}, None),

tests/integration/sdk/_model_cost_sdk_child.py

@@ -0,0 +1,203 @@
+"""One SDK session per case, in a fresh interpreter booted from the cost map at LITELLM_MODEL_COST_MAP_URL.
+
+Each case does what SDK code does to ``litellm.model_cost`` and bills a real ``completion`` after every step, then
+prints the charges as one JSON line. The steps use only names main exposes, so a case runs on main and the branch.
+"""
+
+import asyncio
+import json
+import os
+import sys
+from collections.abc import Callable
+from typing import Final
+
+import litellm
+from litellm.litellm_core_utils.get_model_cost_map import (
+    ModelCostMapReloadUnavailable,
+    _finalize_model_cost_map,  # pyright: ignore[reportPrivateUsage]  # downstream code calls main's finalizer
+    adopt_model_cost_map,
+    refetch_model_cost_map,
+)
+from litellm.utils import (
+    _cached_get_model_info_helper,  # pyright: ignore[reportPrivateUsage]  # downstream code clears it
+    _invalidate_model_cost_lowercase_map,  # pyright: ignore[reportPrivateUsage]  # downstream code calls it
+)
+
+MODEL: Final = "lit-sdk-chat"
+ALIAS: Final = "lit-sdk-chat-alias"
+ADDED: Final = "lit-sdk-added"
+FINETUNE: Final = f"ft:{MODEL}:org:custom:id"
+CUSTOM_INPUT: Final = 9e-06
+CUSTOM_OUTPUT: Final = 7e-06
+BASE: Final = os.environ["MODEL_COST_SDK_BASE_URL"]
+DOCUMENT_P2: Final = os.environ["MODEL_COST_SDK_DOCUMENT"]
+
+
+def bill(model: str) -> float | None:
+    response: Final = litellm.completion(
+        model=f"openai/{model}",
+        api_base=f"{BASE}/v1",
+        api_key="sk-sdk-probe",
+        messages=[{"role": "user", "content": "bill"}],
+    )
+    try:
+        return float(litellm.completion_cost(completion_response=response))
+    except Exception:  # noqa: BLE001  # an unpriced model is an observation here, not a failure
+        return None
+
+
+def reload(document: str) -> bool:
+    result: Final = asyncio.run(refetch_model_cost_map(url=f"{BASE}/{document}"))
+    if isinstance(result, ModelCostMapReloadUnavailable):
+        return False
+    adopt_model_cost_map(result.model_cost_map)
+    return True
+
+
+def row(scale: float) -> dict[str, object]:
+    return {
+        "litellm_provider": "openai",
+        "mode": "chat",
+        "input_cost_per_token": 1e-06 * scale,
+        "output_cost_per_token": 2e-06 * scale,
+        "cache_read_input_token_cost": 5e-07 * scale,
+    }
+
+
+def reload_reprices() -> list[object]:
+    return [bill(MODEL), reload("p2.json"), bill(MODEL)]
+
+
+def failed_reload_keeps_prices() -> list[object]:
+    return [bill(MODEL), reload("missing.json"), bill(MODEL)]
+
+
+def billed_and_listed(model: str) -> list[object]:
+    info: Final = litellm.get_model_info(f"openai/{model}")
+    return [bill(model), info["input_cost_per_token"], info["output_cost_per_token"]]
+
+
+def every_name_bills_and_lists_the_row_across_a_reload() -> list[object]:
+    names: Final = (MODEL, ALIAS, MODEL.upper())
+    before: Final = [value for name in names for value in billed_and_listed(name)]
+    return [*before, reload("p2.json"), *(value for name in names for value in billed_and_listed(name))]
+
+
+def registration_survives_reload() -> list[object]:
+    litellm.register_model({ADDED: {**row(3.0)}})
+    litellm.register_model({"lit-sdk-request-scoped": row(3.0)}, persist_across_reloads=False)
+    before: Final = [bill(ADDED), bill("lit-sdk-request-scoped")]
+    return [*before, reload("p2.json"), bill(ADDED), bill("lit-sdk-request-scoped")]
+
+
+def partial_registration_follows_reload() -> list[object]:
+    litellm.register_model({MODEL: {"input_cost_per_token": CUSTOM_INPUT}})
+    return [bill(MODEL), reload("p2.json"), bill(MODEL)]
+
+
+def partial_registrations_merge_and_follow_reload() -> list[object]:
+    litellm.register_model({MODEL: {"input_cost_per_token": CUSTOM_INPUT}})
+    litellm.register_model({MODEL: {"output_cost_per_token": CUSTOM_OUTPUT}})
+    return [bill(MODEL), reload("p2.json"), bill(MODEL)]
+
+
+def finetune_registration_bills_the_same_after_a_reload() -> list[object]:
+    litellm.register_model({FINETUNE: {"input_cost_per_token": CUSTOM_INPUT, "litellm_provider": "openai"}})
+    before: Final = bill(FINETUNE)
+    return [before is not None, reload("p1.json"), bill(FINETUNE) == before]
+
+
+def finetune_registration_inherits_the_base_prices() -> list[object]:
+    litellm.register_model({FINETUNE: {"input_cost_per_token": CUSTOM_INPUT, "litellm_provider": "openai"}})
+    return [bill(FINETUNE)]
+
+
+def new_row_after_a_miss_without_invalidating() -> list[object]:
+    before: Final = bill(ADDED)
+    litellm.model_cost[ADDED] = row(3.0)
+    return [before, bill(ADDED)]
+
+
+def value_edited_in_place_then_invalidate() -> list[object]:
+    before: Final = bill(MODEL)
+    litellm.model_cost[MODEL]["input_cost_per_token"] = CUSTOM_INPUT
+    _invalidate_model_cost_lowercase_map()
+    return [before, bill(MODEL)]
+
+
+def table_swapped_then_invalidate() -> list[object]:
+    before: Final = bill(MODEL)
+    litellm.model_cost = {**litellm.model_cost, MODEL: row(2.0)}
+    _invalidate_model_cost_lowercase_map()
+    return [before, bill(MODEL)]
+
+
+def table_swapped_then_helper_cache_cleared() -> list[object]:
+    before: Final = bill(MODEL)
+    litellm.model_cost = {**litellm.model_cost, MODEL: row(2.0)}
+    _cached_get_model_info_helper.cache_clear()
+    litellm.get_model_info.cache_clear()
+    return [before, bill(MODEL)]
+
+
+def contents_replaced_then_invalidate() -> list[object]:
+    before: Final = bill(MODEL)
+    saved: Final = {key: dict(value) for key, value in litellm.model_cost.items()}
+    litellm.model_cost.clear()
+    litellm.model_cost.update({**saved, MODEL: row(2.0)})
+    _invalidate_model_cost_lowercase_map()
+    return [before, bill(MODEL)]
+
+
+def finalized_snapshot_replaces_contents() -> list[object]:
+    with open(DOCUMENT_P2) as handle:
+        finalized: Final = _finalize_model_cost_map(json.load(handle))
+    litellm.model_cost.clear()
+    litellm.model_cost.update(finalized)
+    litellm.add_known_models(finalized)
+    return [bill(MODEL)]
+
+
+def alias_registration_after_an_adopted_reload() -> list[object]:
+    reloaded: Final = reload("p2.json")
+    litellm.register_model({ALIAS: {"input_cost_per_token": CUSTOM_INPUT}})
+    return [reloaded, bill(MODEL), bill(ALIAS)]
+
+
+def small_map_finalizes_and_leaves_the_live_map() -> list[object]:
+    finalized: Final = _finalize_model_cost_map({ADDED: row(3.0)})
+    return [list(finalized) == [ADDED], bill(MODEL)]
+
+
+def popped_row_then_registered() -> list[object]:
+    before: Final = bill(MODEL)
+    litellm.model_cost.pop(MODEL)
+    litellm.register_model({MODEL: row(3.0)})
+    return [before, bill(MODEL)]
+
+
+CASES: Final[dict[str, Callable[[], list[object]]]] = {
+    case.__name__: case
+    for case in (
+        reload_reprices,
+        failed_reload_keeps_prices,
+        every_name_bills_and_lists_the_row_across_a_reload,
+        registration_survives_reload,
+        partial_registration_follows_reload,
+        partial_registrations_merge_and_follow_reload,
+        finetune_registration_bills_the_same_after_a_reload,
+        finetune_registration_inherits_the_base_prices,
+        new_row_after_a_miss_without_invalidating,
+        value_edited_in_place_then_invalidate,
+        table_swapped_then_invalidate,
+        table_swapped_then_helper_cache_cleared,
+        contents_replaced_then_invalidate,
+        finalized_snapshot_replaces_contents,
+        alias_registration_after_an_adopted_reload,
+        small_map_finalizes_and_leaves_the_live_map,
+        popped_row_then_registered,
+    )
+}
+
+if __name__ == "__main__":
+    sys.stdout.write(json.dumps(CASES[sys.argv[1]]()) + "\n")

tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py

@@ -10,7 +10,7 @@ from pydantic import JsonValue, TypeAdapter
 
 import litellm
 from litellm import ModelResponse, Router
-from litellm.utils import _invalidate_model_cost_lowercase_map  # pyright: ignore[reportPrivateUsage]  # required cache invalidation
+from litellm.rust_bridge.host.model_cost import publish
 
 _OPUS_4_8: Final = "us.anthropic.claude-opus-4-8"
 _JSON_TOOL: Final = "json_tool_call"
@@ -87,7 +87,7 @@ def opus_4_8_opted_in_for_the_process(monkeypatch: pytest.MonkeyPatch) -> Iterat
     )
     yield
     monkeypatch.undo()
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 def _model(route: str) -> str:

tests/integration/sdk/test_model_cost_fetch_sdk.py

@@ -0,0 +1,594 @@
+import asyncio
+import hashlib
+import json
+import logging
+import os
+import ssl
+import sys
+import threading
+from collections.abc import Generator, Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Final, Literal
+from urllib.parse import urlsplit
+
+import pytest
+from integration._support.tls import server_context, write_self_signed_cert
+from integration._support.wire import Reply, Request, wire_server
+
+import litellm
+from litellm.litellm_core_utils.get_model_cost_map import (
+    GetModelCostMap,
+    ModelCostMapReloaded,
+    ModelCostMapReloadUnavailable,
+    get_model_cost_map_loaded_at,
+    get_model_cost_map_provenance,
+    get_model_cost_map_source_info,
+    model_cost_inputs,
+    refetch_model_cost_map,
+)
+from litellm.rust_bridge.host.model_cost import publish
+from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter
+
+FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations"
+_BUNDLED: Final = GetModelCostMap.read_local_model_cost_map_bytes()
+_IMPORTED_MODEL_COST_INPUTS: Final = model_cost_inputs.snapshot()
+
+
+def _blob_id(body: bytes) -> str:
+    return hashlib.sha1(b"blob %d\0" % len(body) + body, usedforsecurity=False).hexdigest()
+
+
+def _document_with(**rows: Mapping[str, object]) -> bytes:
+    """The bundled document with some rows replaced, so it passes the integrity limits."""
+    return json.dumps({**json.loads(_BUNDLED), **rows}).encode()
+
+
+@dataclass(frozen=True, slots=True)
+class _Reply:
+    status: int = 200
+    body: bytes = _BUNDLED
+    headers: Mapping[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True, slots=True)
+class _Server:
+    url: str
+    received: Sequence[Request]
+
+    @property
+    def targets(self) -> Sequence[str]:
+        return [request.target for request in self.received]
+
+    @property
+    def authorizations(self) -> Sequence[str]:
+        return [request.headers.get("authorization", "") for request in self.received]
+
+    @property
+    def requests(self) -> int:
+        return len(self.received)
+
+
+@contextmanager
+def _serving(*replies: _Reply, tls: ssl.SSLContext | None = None) -> Generator[_Server]:
+    """Answers each GET with the next scripted reply, repeating the last one."""
+    received: Final[list[Request]] = []  # mutable-ok: respond records each request the peer answers
+
+    def respond(request: Request) -> Reply:
+        received.append(request)
+        reply: Final = replies[min(len(received), len(replies)) - 1]
+        return Reply(status=reply.status, body=reply.body, headers=reply.headers)
+
+    with wire_server(respond, tls=tls) as wire:
+        yield _Server(url=f"{wire.url}/cost.json", received=received)
+
+
+@pytest.fixture(autouse=True)
+def isolated_model_cost() -> Iterator[None]:
+    inputs: Final = model_cost_inputs.snapshot()
+    rows: Final = litellm.model_cost
+    entries: Final = dict(rows)
+    model_cost_inputs.restore(_IMPORTED_MODEL_COST_INPUTS)
+    publish(rows)
+    yield
+    model_cost_inputs.restore(inputs)
+    rows.clear()
+    rows.update(entries)
+    publish(rows)
+
+
+@pytest.fixture
+def remote_source(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
+    monkeypatch.setattr(sys, "argv", ["python"])
+
+
+@pytest.fixture
+def own_model_cost(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Restores the binding of ``litellm.model_cost`` after a test that installs a new table."""
+    monkeypatch.setattr(litellm, "model_cost", litellm.model_cost)
+
+
+@pytest.fixture
+def isolated_openai_model_sets(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "open_ai_chat_completion_models", set(litellm.open_ai_chat_completion_models))
+    monkeypatch.setattr(litellm, "open_ai_text_completion_models", set(litellm.open_ai_text_completion_models))
+
+
+@pytest.mark.usefixtures("remote_source")
+def test_boot_load_installs_the_served_document_and_records_its_provenance():
+    body: Final = _document_with(**{"lit-boot-model": {"litellm_provider": "openai", "mode": "chat"}})
+    with _serving(_Reply(body=body, headers={"ETag": '"v1"'})) as server:
+        rows: Final = litellm.get_model_cost_map(url=server.url)
+
+    assert "lit-boot-model" in rows
+    assert FALLBACK_GENERALIZATIONS_KEY not in rows
+    assert get_model_cost_map_provenance() == {"source_revision": _blob_id(body), "etag": '"v1"'}
+    source: Final = get_model_cost_map_source_info()
+    assert (source["source"], source["url"], source["fallback_reason"]) == ("remote", server.url, None)
+
+
+@pytest.mark.usefixtures("remote_source")
+@pytest.mark.parametrize(
+    ("reply", "fallback_reason"),
+    [
+        (_Reply(status=404, body=b"missing"), "Remote fetch failed: HTTP 404 from "),
+        (_Reply(body=b"[1, 2]"), "Remote fetch failed: invalid JSON from "),
+        (
+            _Reply(body=json.dumps({"only": {"litellm_provider": "openai"}}).encode()),
+            "Remote data failed integrity validation",
+        ),
+    ],
+)
+def test_boot_load_falls_back_to_the_bundled_document(reply: _Reply, fallback_reason: str):
+    with _serving(reply) as server:
+        rows: Final = litellm.get_model_cost_map(url=server.url, max_attempts=1)
+
+    assert "gpt-4o" in rows
+    assert get_model_cost_map_provenance() == {"source_revision": _blob_id(_BUNDLED), "etag": None}
+    source: Final = get_model_cost_map_source_info()
+    assert (source["source"], source["url"], source["is_env_forced"]) == ("local", server.url, False)
+    assert source["fallback_reason"] is not None and source["fallback_reason"].startswith(fallback_reason)
+
+
+_MOVED: Final = _Reply(status=302, body=b"", headers={"Location": "/moved.json"})
+_REDIRECTED_DOCUMENT: Final = _document_with(**{"lit-redirected-model": {"litellm_provider": "openai", "mode": "chat"}})
+
+
+def _rows_loaded_through(load: Literal["boot", "reload", "register_model"], url: str) -> Mapping[str, object]:
+    """The rows ``load`` produced from ``url``: the boot load's return value, or the live table after the others."""
+    match load:
+        case "boot":
+            return litellm.get_model_cost_map(url=url, max_attempts=1)
+        case "reload":
+            assert isinstance(asyncio.run(refetch_model_cost_map(url=url, max_attempts=1)), ModelCostMapReloaded)
+            return litellm.model_cost
+        case "register_model":
+            litellm.register_model(url, persist_across_reloads=False)
+            return litellm.model_cost
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+@pytest.mark.parametrize("load", ["boot", "reload", "register_model"])
+def test_every_load_follows_a_redirect_to_the_document(load: Literal["boot", "reload", "register_model"]):
+    with _serving(_MOVED, _Reply(body=_REDIRECTED_DOCUMENT)) as server:
+        rows: Final = _rows_loaded_through(load, server.url)
+
+    assert server.targets == ["/cost.json", "/moved.json"]
+    assert "lit-redirected-model" in rows
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+def test_boot_load_transient_failure_serves_the_bundled_map_then_installs_the_remote_one_in_the_background():
+    body: Final = _document_with(**{"lit-retried-model": {"litellm_provider": "openai", "mode": "chat"}})
+    with _serving(_Reply(status=503, body=b"busy", headers={"Retry-After": "0"}), _Reply(body=body)) as server:
+        rows: Final = litellm.get_model_cost_map(url=server.url)
+        assert "lit-retried-model" not in rows
+        booted_at: Final = get_model_cost_map_loaded_at()
+        retry: Final = next(t for t in threading.enumerate() if t.name == "litellm-model-cost-map-retry")
+        retry.join(timeout=30)
+
+    assert server.requests == 2
+    assert "lit-retried-model" in litellm.model_cost
+    source: Final = get_model_cost_map_source_info()
+    assert (source["source"], source["url"], source["is_env_forced"], source["fallback_reason"]) == (
+        "remote",
+        server.url,
+        False,
+        None,
+    )
+    retried_at: Final = get_model_cost_map_loaded_at()
+    assert booted_at is not None and retried_at is not None and retried_at > booted_at
+
+
+@pytest.mark.usefixtures("remote_source")
+def test_boot_load_does_not_retry_a_final_failure():
+    with _serving(_Reply(status=404, body=b"missing")) as server:
+        litellm.get_model_cost_map(url=server.url)
+
+    assert not any(t.name == "litellm-model-cost-map-retry" for t in threading.enumerate())
+    assert server.requests == 1
+
+
+@pytest.mark.parametrize(
+    ("argv0", "environment"),
+    [
+        ("/usr/bin/litellm-proxy", {}),
+        ("/usr/bin/lite", {}),
+        ("python", {"LITELLM_LOCAL_MODEL_COST_MAP": "True"}),
+        ("python", {"LITELLM_OFFLINE_MODE": "true"}),
+    ],
+    ids=["proxy-cli", "lite-cli", "local-map", "offline"],
+)
+def test_boot_load_serves_the_bundled_map_without_fetching(
+    monkeypatch: pytest.MonkeyPatch, argv0: str, environment: Mapping[str, str]
+):
+    monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
+    monkeypatch.setattr(sys, "argv", [argv0])
+    for name, value in environment.items():
+        monkeypatch.setenv(name, value)
+    with _serving(_Reply()) as server:
+        rows: Final = litellm.get_model_cost_map(url=server.url)
+
+    assert server.requests == 0
+    assert "gpt-4o" in rows
+    assert get_model_cost_map_provenance() == {"source_revision": _blob_id(_BUNDLED), "etag": None}
+    assert get_model_cost_map_source_info()["is_env_forced"] is True
+
+
+def test_importing_litellm_loads_the_map_litellm_model_cost_map_url_names():
+    script: Final = """
+import os
+import litellm
+assert litellm.model_cost_map_url == os.environ["LITELLM_MODEL_COST_MAP_URL"], litellm.model_cost_map_url
+assert "lit-env-url-model" in litellm.model_cost
+"""
+    body: Final = _document_with(**{"lit-env-url-model": {"litellm_provider": "openai", "mode": "chat"}})
+    environment: Final = {
+        key: value
+        for key, value in os.environ.items()
+        if key not in {"LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_OFFLINE_MODE"}
+    }
+    with _serving(_Reply(body=body)) as server:
+        result: Final = run_child_interpreter(
+            script, env={**environment, "LITELLM_MODEL_COST_MAP_URL": server.url}, timeout=120
+        )
+
+    assert result.returncode == 0, result.stderr
+    assert server.targets == ["/cost.json"]
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="fork is the default start method only on Linux")
+def test_importing_litellm_with_a_remote_source_leaves_the_process_free_to_fork():
+    script: Final = """
+import os
+import litellm
+from litellm.rust_bridge._native import process_state_started, reserve_process_for_forking
+assert "lit-fork-model" in litellm.model_cost
+assert not process_state_started()
+reserve_process_for_forking()
+pid = os.fork()
+if pid == 0:
+    os._exit(0 if "lit-fork-model" in litellm.model_cost else 1)
+_, status = os.waitpid(pid, 0)
+assert os.waitstatus_to_exitcode(status) == 0
+"""
+    body: Final = _document_with(**{"lit-fork-model": {"litellm_provider": "openai", "mode": "chat"}})
+    environment: Final = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}
+    with _serving(_Reply(body=body)) as server:
+        result: Final = run_child_interpreter(
+            script, env={**environment, "LITELLM_MODEL_COST_MAP_URL": server.url}, timeout=120
+        )
+
+    assert result.returncode == 0, result.stderr
+
+
+def test_register_model_leaves_the_native_runtime_unstarted() -> None:
+    script: Final = """
+import litellm
+from litellm.rust_bridge._native import process_state_started
+litellm.register_model({"lit-fork-safe-model": {"input_cost_per_token": 1e-06, "litellm_provider": "openai"}})
+assert litellm.model_cost["lit-fork-safe-model"]["input_cost_per_token"] == 1e-06
+assert not process_state_started()
+"""
+    result: Final = run_child_interpreter(
+        script, env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, timeout=120
+    )
+
+    assert result.returncode == 0, result.stderr
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+def test_reload_retries_a_429_after_its_retry_after():
+    with _serving(_Reply(status=429, body=b"slow", headers={"Retry-After": "0"}), _Reply()) as server:
+        result: Final = asyncio.run(refetch_model_cost_map(url=server.url))
+
+    assert isinstance(result, ModelCostMapReloaded)
+    assert server.requests == 2
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+@pytest.mark.parametrize(
+    ("reply", "reason", "requests"),
+    [
+        (_Reply(status=503, body=b"busy", headers={"Retry-After": "0"}), "HTTP 503 from {url} (after 3 attempts)", 3),
+        (_Reply(body=b"not json"), "invalid JSON from {url}", 1),
+        (
+            _Reply(body=json.dumps({"only": {"litellm_provider": "openai"}}).encode()),
+            "model cost map from {url} failed integrity validation",
+            1,
+        ),
+    ],
+    ids=["gives-up-after-its-attempts", "not-json", "below-integrity-limits"],
+)
+def test_a_failed_reload_keeps_the_installed_map(reply: _Reply, reason: str, requests: int):
+    installed: Final = litellm.model_cost
+    provenance: Final = get_model_cost_map_provenance()
+    with _serving(reply) as server:
+        result: Final = asyncio.run(refetch_model_cost_map(url=server.url, max_attempts=3))
+
+    assert isinstance(result, ModelCostMapReloadUnavailable)
+    assert result.reason.startswith(reason.format(url=server.url)), result.reason
+    assert server.requests == requests
+    assert litellm.model_cost is installed
+    assert get_model_cost_map_provenance() == provenance
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+def test_a_reload_rejected_by_the_integrity_limits_logs_the_counts(caplog: pytest.LogCaptureFixture):
+    with (
+        _serving(_Reply(body=json.dumps({"only": {"litellm_provider": "openai"}}).encode())) as server,
+        caplog.at_level(logging.WARNING, logger="LiteLLM"),
+    ):
+        asyncio.run(refetch_model_cost_map(url=server.url))
+
+    assert "has 1 models, below minimum" in caplog.text
+
+
+@pytest.mark.usefixtures("own_model_cost")
+@pytest.mark.parametrize("argv0", ["/usr/bin/litellm-proxy", "/usr/bin/lite"])
+def test_a_cli_process_that_booted_on_the_bundled_map_still_fetches_on_reload(
+    monkeypatch: pytest.MonkeyPatch, argv0: str
+):
+    monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
+    monkeypatch.setattr(sys, "argv", [argv0])
+    with _serving(_Reply()) as server:
+        litellm.get_model_cost_map(url=server.url)
+        booted_requests: Final = server.requests
+        result: Final = asyncio.run(refetch_model_cost_map(url=server.url))
+
+    assert booted_requests == 0
+    assert isinstance(result, ModelCostMapReloaded)
+    assert server.requests == 1
+    assert get_model_cost_map_source_info()["source"] == "remote"
+
+
+@pytest.mark.usefixtures("own_model_cost")
+def test_reload_with_the_local_override_installs_the_bundled_map_without_fetching(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+    with _serving(_Reply()) as server:
+        result: Final = asyncio.run(refetch_model_cost_map(url=server.url))
+
+    assert isinstance(result, ModelCostMapReloaded)
+    assert server.requests == 0
+    assert get_model_cost_map_provenance() == {"source_revision": _blob_id(_BUNDLED), "etag": None}
+    source: Final = get_model_cost_map_source_info()
+    assert (source["source"], source["url"], source["is_env_forced"]) == ("local", None, True)
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+def test_only_a_reload_that_installs_moves_the_source_info_and_load_time():
+    with _serving(_Reply()) as first:
+        asyncio.run(refetch_model_cost_map(url=first.url))
+    installed: Final = get_model_cost_map_source_info()
+    loaded_at: Final = get_model_cost_map_loaded_at()
+
+    with _serving(_Reply(status=404, body=b"missing")) as failing:
+        asyncio.run(refetch_model_cost_map(url=failing.url))
+    assert get_model_cost_map_source_info() == installed
+    assert get_model_cost_map_loaded_at() is loaded_at
+
+    with _serving(_Reply()) as second:
+        asyncio.run(refetch_model_cost_map(url=second.url))
+    source: Final = get_model_cost_map_source_info()
+    assert (source["source"], source["url"], source["is_env_forced"], source["fallback_reason"]) == (
+        "remote",
+        second.url,
+        False,
+        None,
+    )
+    assert get_model_cost_map_loaded_at() is not loaded_at
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+def test_reload_records_the_served_bytes_and_etag():
+    body: Final = _document_with(**{"lit-revision-model": {"litellm_provider": "openai", "mode": "chat"}})
+    with _serving(_Reply(body=body, headers={"ETag": 'W/"etag"'})) as server:
+        asyncio.run(refetch_model_cost_map(url=server.url))
+
+    assert get_model_cost_map_provenance() == {"source_revision": _blob_id(body), "etag": 'W/"etag"'}
+
+
+def _self_signed(directory: Path) -> tuple[ssl.SSLContext, Path]:
+    """A server context for 127.0.0.1 and the path of its certificate, which nothing trusts by default."""
+    certificate, key = write_self_signed_cert(directory)
+    return server_context(certificate, key), certificate
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+def test_a_reload_rejects_an_untrusted_certificate_by_default(tmp_path: Path):
+    context, _ = _self_signed(tmp_path)
+    with _serving(_Reply(), tls=context) as server:
+        result: Final = asyncio.run(refetch_model_cost_map(url=server.url, max_attempts=1))
+
+    assert isinstance(result, ModelCostMapReloadUnavailable)
+    assert result.reason.startswith(f"error fetching {server.url}: ")
+    assert server.requests == 0
+
+
+@pytest.fixture
+def no_tls_environment(monkeypatch: pytest.MonkeyPatch) -> None:
+    for name in ("SSL_VERIFY", "SSL_CERT_FILE", "SSL_CERTIFICATE"):
+        monkeypatch.delenv(name, raising=False)
+
+
+@pytest.mark.usefixtures("remote_source", "no_tls_environment", "own_model_cost")
+@pytest.mark.parametrize(
+    ("load", "setting", "trust"),
+    [
+        ("reload", "litellm.ssl_verify", "off"),
+        ("reload", "litellm.ssl_verify", "certificate"),
+        ("boot", "SSL_VERIFY", "off"),
+        ("boot", "SSL_CERT_FILE", "certificate"),
+    ],
+)
+def test_a_load_trusts_the_certificate_its_settings_allow(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+    load: Literal["boot", "reload"],
+    setting: str,
+    trust: Literal["off", "certificate"],
+):
+    context, certificate = _self_signed(tmp_path)
+    if setting == "litellm.ssl_verify":
+        monkeypatch.setattr(litellm, "ssl_verify", {"off": False, "certificate": str(certificate)}[trust])
+    else:
+        monkeypatch.setenv(setting, {"off": "False", "certificate": str(certificate)}[trust])
+    with _serving(_Reply(), tls=context) as server:
+        _rows_loaded_through(load, server.url)
+
+    assert server.requests == 1
+    assert get_model_cost_map_source_info()["source"] == "remote"
+
+
+@pytest.mark.usefixtures("remote_source", "no_tls_environment")
+def test_a_boot_load_ignores_litellm_ssl_verify_because_it_runs_before_settings_exist(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+):
+    context, _ = _self_signed(tmp_path)
+    monkeypatch.setattr(litellm, "ssl_verify", False)
+    with _serving(_Reply(), tls=context) as server:
+        litellm.get_model_cost_map(url=server.url, max_attempts=1)
+
+    assert server.requests == 0
+    source: Final = get_model_cost_map_source_info()
+    assert source["source"] == "local"
+    assert source["fallback_reason"] is not None
+    assert source["fallback_reason"].startswith(f"Remote fetch failed: error fetching {server.url}: ")
+
+
+_UNREACHABLE_ENDPOINT: Final = "http://127.0.0.1:9"
+
+
+@pytest.mark.usefixtures("remote_source")
+@pytest.mark.parametrize(
+    ("environment", "endpoint_variable", "region"),
+    [
+        (
+            {"AWS_REGION_NAME": "eu-west-3", "AWS_REGION": "ap-south-1", "AWS_DEFAULT_REGION": "sa-east-1"},
+            "AWS_ENDPOINT_URL_S3",
+            "eu-west-3",
+        ),
+        ({"AWS_REGION": "ap-south-1", "AWS_DEFAULT_REGION": "sa-east-1"}, "AWS_ENDPOINT_URL_S3", "ap-south-1"),
+        ({"AWS_DEFAULT_REGION": "sa-east-1"}, "AWS_ENDPOINT_URL_S3", "sa-east-1"),
+        ({"AWS_ENDPOINT_URL": _UNREACHABLE_ENDPOINT}, "AWS_ENDPOINT_URL_S3", "us-east-1"),
+        ({}, "AWS_ENDPOINT_URL", "us-east-1"),
+        ({"AWS_CONFIG_FILE": "profile_region"}, "AWS_ENDPOINT_URL_S3", "eu-north-1"),
+    ],
+)
+def test_an_s3_load_signs_for_the_region_and_reaches_the_endpoint_the_environment_names(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+    environment: Mapping[str, str],
+    endpoint_variable: str,
+    region: str,
+):
+    for name in (
+        "AWS_REGION_NAME",
+        "AWS_REGION",
+        "AWS_DEFAULT_REGION",
+        "AWS_PROFILE",
+        "AWS_SESSION_TOKEN",
+        "AWS_ENDPOINT_URL_S3",
+        "AWS_ENDPOINT_URL",
+    ):
+        monkeypatch.delenv(name, raising=False)
+    profile: Final = tmp_path / "profile_region"
+    profile.write_text("[default]\nregion = eu-north-1\n")
+    monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "no_config"))
+    monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "no_credentials"))
+    monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true")
+    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIDMODELCOST")
+    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "model-cost-secret")
+    for name, value in environment.items():
+        monkeypatch.setenv(name, str(tmp_path / value) if name == "AWS_CONFIG_FILE" else value)
+    with _serving(_Reply()) as server:
+        monkeypatch.setenv(endpoint_variable, server.url.removesuffix("/cost.json"))
+        litellm.get_model_cost_map(url="s3://lit-prices/maps/cost.json", max_attempts=1)
+
+    assert [urlsplit(target).path for target in server.targets] == ["/lit-prices/maps/cost.json"]
+    assert f"/{region}/s3/aws4_request" in server.authorizations[0]
+    assert get_model_cost_map_source_info()["source"] == "remote"
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+@pytest.mark.parametrize("load", ["boot", "reload"])
+def test_both_loads_fetch_through_the_proxy_the_environment_names(
+    monkeypatch: pytest.MonkeyPatch, load: Literal["boot", "reload"]
+):
+    for name in ("http_proxy", "HTTPS_PROXY", "https_proxy", "ALL_PROXY", "all_proxy", "NO_PROXY", "no_proxy"):
+        monkeypatch.delenv(name, raising=False)
+    url: Final = "http://lit-cost-map.invalid/cost.json"
+    with _serving(_Reply()) as proxy:
+        monkeypatch.setenv("HTTP_PROXY", proxy.url.removesuffix("/cost.json"))
+        rows: Final = _rows_loaded_through(load, url)
+
+    assert "gpt-4o" in rows
+    assert get_model_cost_map_source_info()["source"] == "remote"
+    assert proxy.targets == [url]
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+def test_register_model_from_a_url_registers_the_served_rows_without_the_fallback_rules():
+    document: Final = {
+        "lit-url-model": {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 4e-06},
+        FALLBACK_GENERALIZATIONS_KEY: {"rules": []},
+    }
+    with _serving(_Reply(body=json.dumps(document).encode())) as server:
+        litellm.register_model(server.url)
+
+    assert server.requests == 1
+    assert litellm.model_cost["lit-url-model"]["input_cost_per_token"] == 4e-06
+    assert FALLBACK_GENERALIZATIONS_KEY not in litellm.model_cost
+    assert model_cost_inputs.registered_keys() == [*GetModelCostMap.loaded_model_cost_map(), "lit-url-model"]
+
+
+@pytest.mark.usefixtures("remote_source", "own_model_cost")
+@pytest.mark.parametrize(
+    "reply", [_Reply(status=500, body=b"down"), _Reply(body=b"not json")], ids=["server_error", "invalid_json"]
+)
+def test_register_model_from_a_url_that_fails_registers_nothing(reply: _Reply):
+    with _serving(reply) as server:
+        litellm.register_model(server.url)
+
+    assert model_cost_inputs.registered_keys() == list(GetModelCostMap.loaded_model_cost_map())
+
+
+@pytest.mark.usefixtures("isolated_openai_model_sets")
+def test_register_model_url_fetch_uses_single_attempt(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
+    monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
+    before: Final = dict(litellm.model_cost)
+    with wire_server(lambda _: Reply(status=503, body=b"")) as peer:
+        threads_before: Final = {thread.name for thread in threading.enumerate()}
+        litellm.register_model(model_cost=f"{peer.url}/custom_pricing.json")
+
+    threads_after: Final = {thread.name for thread in threading.enumerate()}
+    requests: Final = [request.target for request in peer.drain()]
+    assert requests == ["/custom_pricing.json"]
+    assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"}
+    assert not any(
+        thread.name == "litellm-model-cost-map-retry" and thread.is_alive() for thread in threading.enumerate()
+    )
+    assert litellm.model_cost.keys() >= before.keys()

tests/integration/sdk/test_model_cost_sdk.py

@@ -0,0 +1,188 @@
+import json
+import os
+import subprocess
+import sys
+import uuid
+from collections.abc import Iterator
+from dataclasses import dataclass
+from importlib.resources import files
+from pathlib import Path
+from typing import Final
+
+import pytest
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
+
+_MODEL: Final = "lit-sdk-chat"
+_ALIAS: Final = "lit-sdk-chat-alias"
+_PROMPT: Final = 1000
+_CACHED: Final = 400
+_COMPLETION: Final = 500
+_CUSTOM_INPUT: Final = 9e-06
+_CUSTOM_OUTPUT: Final = 7e-06
+_CHARGES: Final = TypeAdapter(list[float | bool | None])
+_BUNDLED: Final = TypeAdapter(dict[str, JsonValue]).validate_json(
+    files("litellm").joinpath("model_prices_and_context_window_backup.json").read_bytes()
+)
+
+
+class _Requested(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    model: str
+
+
+_REQUEST: Final = TypeAdapter(_Requested)
+
+
+@dataclass(frozen=True, slots=True)
+class _Rates:
+    input: float
+    output: float
+    cache_read: float
+
+    def cost(self) -> float:
+        return (_PROMPT - _CACHED) * self.input + _CACHED * self.cache_read + _COMPLETION * self.output
+
+
+def _catalog(scale: float) -> _Rates:
+    return _Rates(input=1e-06 * scale, output=2e-06 * scale, cache_read=5e-07 * scale)
+
+
+def _document(scale: float) -> bytes:
+    rates: Final = _catalog(scale)
+    return json.dumps(
+        {
+            **_BUNDLED,
+            _MODEL: {
+                "aliases": [_ALIAS],
+                "litellm_provider": "openai",
+                "mode": "chat",
+                "input_cost_per_token": rates.input,
+                "output_cost_per_token": rates.output,
+                "cache_read_input_token_cost": rates.cache_read,
+            },
+        }
+    ).encode()
+
+
+def _respond(request: Request) -> Reply:
+    match request.target:
+        case "/p1.json":
+            return Reply(body=_document(1.0))
+        case "/p2.json":
+            return Reply(body=_document(2.0))
+        case "/missing.json":
+            return Reply(status=404, body=b"{}")
+        case _:
+            return Reply(
+                body=json.dumps(
+                    {
+                        "id": f"chatcmpl-{uuid.uuid4().hex}",
+                        "object": "chat.completion",
+                        "created": 1,
+                        "model": _REQUEST.validate_json(request.body).model,
+                        "choices": [
+                            {"index": 0, "message": {"role": "assistant", "content": "billed"}, "finish_reason": "stop"}
+                        ],
+                        "usage": {
+                            "prompt_tokens": _PROMPT,
+                            "completion_tokens": _COMPLETION,
+                            "total_tokens": _PROMPT + _COMPLETION,
+                            "prompt_tokens_details": {"cached_tokens": _CACHED},
+                        },
+                    }
+                ).encode()
+            )
+
+
+@pytest.fixture(scope="module")
+def cost_map_peer() -> Iterator[Wire]:
+    with wire_server(_respond) as peer:
+        yield peer
+
+
+def _run(case: str, peer: Wire, directory: Path) -> list[float | bool | None]:
+    document: Final = directory / "p2.json"
+    document.write_bytes(_document(2.0))
+    environment: Final = {
+        **{name: value for name, value in os.environ.items() if name != "LITELLM_LOCAL_MODEL_COST_MAP"},
+        "LITELLM_MODEL_COST_MAP_URL": f"{peer.url}/p1.json",
+        "MODEL_COST_SDK_BASE_URL": peer.url,
+        "MODEL_COST_SDK_DOCUMENT": str(document),
+    }
+    result: Final = subprocess.run(  # test-quality-ok: each case is one SDK session in its own interpreter
+        (sys.executable, "-m", "integration.sdk._model_cost_sdk_child", case),
+        env=environment,
+        capture_output=True,
+        text=True,
+        timeout=180,
+        check=False,
+    )
+    assert result.returncode == 0, result.stderr[-4000:]
+    return _CHARGES.validate_json(result.stdout.strip().splitlines()[-1])
+
+
+def _custom_input(rates: _Rates) -> _Rates:
+    return _Rates(input=_CUSTOM_INPUT, output=rates.output, cache_read=rates.cache_read)
+
+
+def _custom_input_and_output(rates: _Rates) -> _Rates:
+    return _Rates(input=_CUSTOM_INPUT, output=_CUSTOM_OUTPUT, cache_read=rates.cache_read)
+
+
+def _charged(rates: _Rates) -> object:
+    return pytest.approx(rates.cost(), rel=1e-9)
+
+
+def _billed_and_listed(rates: _Rates) -> list[object]:
+    return [_charged(rates), pytest.approx(rates.input, rel=1e-9), pytest.approx(rates.output, rel=1e-9)] * 3
+
+
+_P1: Final = _catalog(1.0)
+_P2: Final = _catalog(2.0)
+_P3: Final = _catalog(3.0)
+_FINETUNE_BUG: Final = pytest.mark.xfail(
+    strict=True,
+    reason="BUG: a fine-tune registered with only an input price bills output at 0 and cache reads at input",
+)
+
+
+@pytest.mark.parametrize(
+    ("case", "expected"),
+    [
+        ("reload_reprices", [_charged(_P1), True, _charged(_P2)]),
+        ("failed_reload_keeps_prices", [_charged(_P1), False, _charged(_P1)]),
+        (
+            "every_name_bills_and_lists_the_row_across_a_reload",
+            [*_billed_and_listed(_P1), True, *_billed_and_listed(_P2)],
+        ),
+        ("registration_survives_reload", [_charged(_P3), _charged(_P3), True, _charged(_P3), None]),
+        ("partial_registration_follows_reload", [_charged(_custom_input(_P1)), True, _charged(_custom_input(_P2))]),
+        (
+            "partial_registrations_merge_and_follow_reload",
+            [_charged(_custom_input_and_output(_P1)), True, _charged(_custom_input_and_output(_P2))],
+        ),
+        ("finetune_registration_bills_the_same_after_a_reload", [True, True, True]),
+        pytest.param(
+            "finetune_registration_inherits_the_base_prices", [_charged(_custom_input(_P1))], marks=_FINETUNE_BUG
+        ),
+        ("new_row_after_a_miss_without_invalidating", [None, _charged(_P3)]),
+        ("value_edited_in_place_then_invalidate", [_charged(_P1), _charged(_custom_input(_P1))]),
+        ("table_swapped_then_invalidate", [_charged(_P1), _charged(_P2)]),
+        ("table_swapped_then_helper_cache_cleared", [_charged(_P1), _charged(_P2)]),
+        ("contents_replaced_then_invalidate", [_charged(_P1), _charged(_P2)]),
+        ("finalized_snapshot_replaces_contents", [_charged(_P2)]),
+        (
+            "alias_registration_after_an_adopted_reload",
+            [True, _charged(_custom_input(_P2)), _charged(_custom_input(_P2))],
+        ),
+        ("small_map_finalizes_and_leaves_the_live_map", [True, _charged(_P1)]),
+        ("popped_row_then_registered", [_charged(_P1), _charged(_P3)]),
+    ],
+    ids=lambda value: value if isinstance(value, str) else "",
+)
+def test_sdk_code_that_reads_or_writes_the_cost_map_gets_what_it_set_up(
+    case: str, expected: list[object], cost_map_peer: Wire, tmp_path: Path
+) -> None:
+    assert _run(case, cost_map_peer, tmp_path) == expected

tests/unit/caching/test_llm_caching_handler.py

@@ -28,7 +28,7 @@ from litellm.llms.custom_httpx.http_handler import(
     HTTPHandler,
 )
 from litellm.types.utils import LlmProviders
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
 
@@ -408,7 +408,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 _SCALAR_DEFAULTS = {
     "num_retries": getattr(litellm, "num_retries", None),

tests/unit/conftest.py

@@ -134,11 +134,14 @@ os.environ["PYTHON_DOTENV_DISABLED"] = "1"
 os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
 
 import litellm  # noqa: E402  # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import
-import litellm.utils as litellm_utils_module  # noqa: E402  # same import-time dependency
+from litellm.rust_bridge.host.model_cost import publish  # noqa: E402  # same import-time dependency
 from litellm._logging import ALL_LOGGERS  # noqa: E402  # same import-time dependency
 from litellm.anthropic_beta_headers_manager import (
     reload_beta_headers_config,  # noqa: E402  # same import-time dependency
 )
+from litellm.litellm_core_utils.get_model_cost_map import (  # noqa: E402  # same import-time dependency
+    model_cost_inputs,
+)
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER  # noqa: E402  # same import-time dependency
 from litellm.litellm_core_utils.prompt_templates import (
     factory as prompt_factory_module,  # noqa: E402  # same import-time dependency
@@ -430,22 +433,21 @@ def isolate_litellm_globals() -> Iterator[None]:
     tool_registry_writer_module._tool_policy_registry = original_tool_policy_registry
 
 
+_IMPORTED_MODEL_COST_INPUTS: Final = model_cost_inputs.snapshot()
+
+
 @pytest.fixture(autouse=True)
 def isolate_router_model_cost_state() -> Iterator[None]:
     original_live_registries: Final = frozenset(live_router_registries)
-    original_runtime_registered_model_cost: Final = {
-        model_key: dict(model_value)
-        for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items()
-    }
-    litellm_utils_module._invalidate_model_cost_lowercase_map()
+    original_model_cost_inputs: Final = model_cost_inputs.snapshot()
+    model_cost_inputs.restore(_IMPORTED_MODEL_COST_INPUTS)
+    publish(litellm.model_cost)
     live_router_registries.clear()
-    litellm_utils_module._runtime_registered_model_cost.clear()
     yield
     live_router_registries.clear()
     live_router_registries.update(original_live_registries)
-    litellm_utils_module._runtime_registered_model_cost.clear()
-    litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost)
-    litellm_utils_module._invalidate_model_cost_lowercase_map()
+    model_cost_inputs.restore(original_model_cost_inputs)
+    publish(litellm.model_cost)
     litellm.get_model_info.cache_clear()
 
 
@@ -469,7 +471,7 @@ def isolate_model_catalog() -> Iterator[None]:
     for models, entries in provider_model_sets:
         models.clear()
         models.update(entries)
-    litellm_utils_module._invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     litellm.get_model_info.cache_clear()
 
 

tests/unit/decisions/test_main.py

@@ -26,7 +26,7 @@ import litellm
 from litellm.constants import AWS_SIGNING_MAX_THREADS
 from litellm.cost_calculator import get_response_cost_from_hidden_params
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded, refetch_model_cost_map
+from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded, install_model_cost_map
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers
 from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR
@@ -3963,13 +3963,8 @@ CatalogLoader: TypeAlias = Callable[[Mapping[str, Mapping[str, object]]], Awaita
 def load_catalog_with(monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None) -> Iterator[CatalogLoader]:
     async def load(rows: Mapping[str, Mapping[str, object]]) -> None:
         catalog: Final = {**litellm.get_model_cost_map(url=""), **rows}
-        transport: Final = httpx.MockTransport(lambda _: httpx.Response(200, json=catalog))
-        monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP")
-        async with httpx.AsyncClient(transport=transport) as client:
-            loaded: Final = await refetch_model_cost_map(url="https://catalog.example/prices.json", client=client)
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-        assert isinstance(loaded, ModelCostMapReloaded)
-        monkeypatch.setattr(litellm, "model_cost", loaded.model_cost_map)
+        monkeypatch.setattr(litellm, "model_cost", litellm.model_cost)
+        assert isinstance(install_model_cost_map(json.dumps(catalog).encode()), ModelCostMapReloaded)
         litellm.get_model_info.cache_clear()
 
     yield load

tests/unit/integrations/arize/test_arize.py

@@ -16,7 +16,7 @@ from litellm.integrations.arize.arize import ArizeConfig, ArizeLogger
 from litellm.integrations.opentelemetry import OpenTelemetryConfig
 from litellm._logging import verbose_logger, verbose_proxy_logger
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 
@@ -346,7 +346,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 _SCALAR_DEFAULTS = {
     "num_retries": getattr(litellm, "num_retries", None),

tests/unit/integrations/arize/test_arize_phoenix.py

@@ -17,7 +17,7 @@ from litellm.integrations.arize.arize_phoenix import (
 )
 from litellm._logging import verbose_logger, verbose_proxy_logger
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 
@@ -1044,7 +1044,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 _SCALAR_DEFAULTS = {
     "num_retries": getattr(litellm, "num_retries", None),

tests/unit/integrations/opik/test_opik.py

@@ -12,7 +12,7 @@ import pytest
 import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 
@@ -341,7 +341,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 _SCALAR_DEFAULTS = {

tests/unit/integrations/test_anthropic_cache_control_hook.py

@@ -19,6 +19,7 @@ from openai.types.chat import ChatCompletionMessage
 from pydantic import BaseModel, ConfigDict
 
 import litellm
+from litellm.rust_bridge.host.model_cost import publish
 from litellm import Router
 from litellm._logging import verbose_logger
 from litellm.integrations.anthropic_cache_control_hook import (
@@ -4353,9 +4354,9 @@ class TestPromptCacheBreakpointCapability:
         bundled = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json")
         with open(bundled) as handle:
             monkeypatch.setattr(litellm, "model_cost", json.load(handle))
-        litellm.utils.cached_get_model_info_helper.cache_clear()
+        publish(litellm.model_cost)
         yield
-        litellm.utils.cached_get_model_info_helper.cache_clear()
+        publish(litellm.model_cost)
 
     def test_listed_model_uses_the_model_map_flag(self, monkeypatch):
         flagged = {**litellm.model_cost["gpt-4.1"], "supports_prompt_cache_breakpoint": True}
@@ -4587,9 +4588,9 @@ class TestBedrockMantleGptShipsTheOpenAIDialect:
         bundled = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json")
         with open(bundled) as handle:
             monkeypatch.setattr(litellm, "model_cost", json.load(handle))
-        litellm.utils.cached_get_model_info_helper.cache_clear()
+        publish(litellm.model_cost)
         yield
-        litellm.utils.cached_get_model_info_helper.cache_clear()
+        publish(litellm.model_cost)
 
     def test_shipped_entry_makes_the_deployment_eligible(self):
         assert supports_openai_prompt_cache_breakpoint(self.MANTLE_MODEL) is True

tests/unit/integrations/test_helicone.py

@@ -7,7 +7,7 @@ from collections.abc import Iterator
 from pathlib import Path
 from typing import Final
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 from unittest.mock import MagicMock
 
@@ -157,7 +157,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 _SCALAR_DEFAULTS = {
     "num_retries": getattr(litellm, "num_retries", None),

tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py

@@ -83,10 +83,10 @@ from litellm.types.llms.openai import (
 from litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse, TextCompletionResponse
 from litellm.utils import (
     Rules,
-    _invalidate_model_cost_lowercase_map,
     function_setup,
     validate_and_fix_openai_messages,
 )
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}]
@@ -6190,7 +6190,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 _SCALAR_DEFAULTS = {

tests/unit/litellm_core_utils/test_get_llm_provider_logic.py

@@ -21,7 +21,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases
 from litellm.types.router import LiteLLM_Params
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 from unittest.mock import patch
 
@@ -209,7 +209,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 _SCALAR_DEFAULTS = {

tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py

@@ -34,7 +34,7 @@ from litellm.types.utils import (
     StandardLoggingPayloadErrorInformation,
     StreamingChoices,
 )
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 import unittest.mock
 
 
@@ -1570,7 +1570,7 @@ async def test_messages_honors_a_chat_pin_on_a_prefixed_responses_deployment(
         "supported_endpoints": ["/v1/chat/completions", "/v1/responses"],
     }
     litellm.model_cost["perplexity/pin-test-model"] = {"litellm_provider": "perplexity", "mode": "chat"}
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
     litellm.in_memory_llm_clients_cache.flush_cache()
 
@@ -3017,7 +3017,7 @@ async def test_azure_messages_replay_encrypted_reasoning_only_to_responses(
             },
         },
     )
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     litellm.in_memory_llm_clients_cache.flush_cache()
     chat_reply: Final = {
         "id": "chatcmpl-azure",

tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py

@@ -26,7 +26,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT
 from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config
 from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig
 from litellm.router import Router
-from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params
+from litellm.utils import get_optional_params
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 _MAPPED_PARAMS: Final = TypeAdapter(dict[str, object])
@@ -598,7 +599,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 _SCALAR_DEFAULTS = {
     "num_retries": getattr(litellm, "num_retries", None),

tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py

@@ -15,17 +15,18 @@ from litellm.llms.azure_ai.image_generation.flux_transformation import (
     AzureFoundryFluxImageGenerationConfig,
 )
 from litellm.types.utils import ImageObject, ImageResponse
-from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params_image_gen
+from litellm.utils import get_optional_params_image_gen
+from litellm.rust_bridge.host.model_cost import publish
 
 
 @pytest.fixture(autouse=True)
 def use_local_model_cost_map(monkeypatch):
     monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map())
     litellm.get_model_info.cache_clear()
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     yield
     litellm.get_model_info.cache_clear()
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 @pytest.mark.parametrize(

tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py

@@ -19,7 +19,7 @@ def use_local_model_cost_map():
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
 
     import litellm
-    from litellm.utils import _invalidate_model_cost_lowercase_map
+    from litellm.rust_bridge.host.model_cost import publish
 
     original_model_cost = litellm.model_cost
     litellm.model_cost = json.loads(
@@ -28,13 +28,13 @@ def use_local_model_cost_map():
         .read_text(encoding="utf-8")
     )
     litellm.get_model_info.cache_clear()
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     try:
         yield litellm
     finally:
         litellm.model_cost = original_model_cost
         litellm.get_model_info.cache_clear()
-        _invalidate_model_cost_lowercase_map()
+        publish(litellm.model_cost)
         monkeypatch.undo()
 
 

tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py

@@ -14,7 +14,7 @@ def use_local_model_cost_map():
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
 
     import litellm
-    from litellm.utils import _invalidate_model_cost_lowercase_map
+    from litellm.rust_bridge.host.model_cost import publish
 
     original_model_cost = litellm.model_cost
     litellm.model_cost = json.loads(
@@ -23,13 +23,13 @@ def use_local_model_cost_map():
         .read_text(encoding="utf-8")
     )
     litellm.get_model_info.cache_clear()
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     try:
         yield litellm
     finally:
         litellm.model_cost = original_model_cost
         litellm.get_model_info.cache_clear()
-        _invalidate_model_cost_lowercase_map()
+        publish(litellm.model_cost)
         monkeypatch.undo()
 
 

tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py

@@ -27,7 +27,7 @@ def ensure_nova_canvas_image_edit_model_cost_flags(monkeypatch):
     Full ``model_prices_and_context_window.json`` includes these flags, but CI or
     alternate cost maps may omit them—merge minimal entries so tests match production.
     """
-    from litellm.utils import _invalidate_model_cost_lowercase_map
+    from litellm.rust_bridge.host.model_cost import publish
 
     for key in (
         "amazon.nova-canvas-v1:0",
@@ -46,11 +46,11 @@ def ensure_nova_canvas_image_edit_model_cost_flags(monkeypatch):
                 "supports_nova_canvas_image_edit": True,
             },
         )
-        _invalidate_model_cost_lowercase_map()
+        publish(litellm.model_cost)
 
     yield
 
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 def test_get_config_class_nova_canvas():

tests/unit/llms/custom_httpx/test_http_handler.py

@@ -42,7 +42,7 @@ from litellm.types.llms.custom_http import VerifyTypes
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from litellm.exceptions import Timeout as LitellmTimeout
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 
@@ -2183,7 +2183,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 _SCALAR_DEFAULTS = {

tests/unit/llms/openai/test_openai.py

@@ -42,7 +42,7 @@ from litellm.types.utils import (
     PromptTokensDetails,
 )
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from openai.types.beta.assistant import Assistant
 from openai.types.beta.assistant_deleted import AssistantDeleted
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -1835,7 +1835,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 _SCALAR_DEFAULTS = {

tests/unit/llms/parallel_ai/test_parallel_ai_search.py

@@ -57,13 +57,13 @@ def bundled_cost_map(monkeypatch):
     litellm caches model-info lookups, so swapping ``model_cost`` only takes
     effect once those caches are invalidated -- on the way in and back out.
     """
-    from litellm.utils import _invalidate_model_cost_lowercase_map
+    from litellm.rust_bridge.host.model_cost import publish
 
     monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     yield
     monkeypatch.undo()
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 class TestParallelAISearch:

tests/unit/llms/test_custom_llm.py

@@ -35,7 +35,8 @@ from litellm.types.utils import (
     ModelResponseStream,
     StreamingChoices,
 )
-from litellm.utils import ModelResponseIterator, _invalidate_model_cost_lowercase_map
+from litellm.utils import ModelResponseIterator
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 
@@ -729,7 +730,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 _SCALAR_DEFAULTS = {

tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py

@@ -36,7 +36,8 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
 from litellm.types.llms.base import CachedTokensDetails
 from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata
 from litellm.types.utils import ChoiceLogprobs, ModelInfo, ModelResponseStream, Usage
-from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper
+from litellm.utils import CustomStreamWrapper
+from litellm.rust_bridge.host.model_cost import publish
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.llms.vertex_ai.gemini.transformation import (
@@ -8161,7 +8162,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 
 _SCALAR_DEFAULTS = {
@@ -9676,7 +9677,7 @@ async def test_gemini_pro_json_schema_args_sent_httpx(
     monkeypatch.setenv("GEMINI_API_KEY", "test-gemini-api-key")
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
     monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     monkeypatch.setattr(litellm, "set_verbose", True)
     messages = [{"role": "user", "content": "List 5 cookie recipes"}]
     from litellm.llms.custom_httpx.http_handler import HTTPHandler
@@ -9782,7 +9783,7 @@ async def test_gemini_pro_json_schema_args_sent_httpx_openai_schema(
 
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
     monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
     monkeypatch.setattr(litellm, "set_verbose", True)
     messages = [{"role": "user", "content": "List 5 cookie recipes"}]
     from litellm.llms.custom_httpx.http_handler import HTTPHandler

tests/unit/proxy/common_utils/test_config_sync_pubsub.py

@@ -1086,7 +1086,7 @@ async def test_model_cost_map_reload_does_not_publish_config_change() -> None:
     from litellm.proxy import proxy_server
     from litellm.proxy.proxy_server import ProxyConfig, _set_redis_usage_cache
     from litellm.proxy.utils import litellm_config_cache
-    from litellm.utils import _invalidate_model_cost_lowercase_map
+    from litellm.rust_bridge.host.model_cost import publish
 
     litellm_config_cache.flush_cache()
     prisma_client = _reload_config_prisma_client()
@@ -1100,13 +1100,12 @@ async def test_model_cost_map_reload_does_not_publish_config_change() -> None:
         with patch(
             "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map",
             new=AsyncMock(
-                return_value=ModelCostMapReloaded(model_cost_map={"gpt-5.2": {"input_cost_per_token": 0.001}})
+                return_value=ModelCostMapReloaded(model_count=len({"gpt-5.2": {"input_cost_per_token": 0.001}}))
             ),
         ):
             await ProxyConfig()._check_and_reload_model_cost_map(prisma_client=prisma_client)
     finally:
-        litellm.model_cost = original_model_cost
-        _invalidate_model_cost_lowercase_map()
+        publish(original_model_cost)
         _set_redis_usage_cache(previous_cache)
 
     prisma_client.db.litellm_config.update_many.assert_awaited_once()

tests/unit/proxy/common_utils/test_http_parsing_utils.py

@@ -34,7 +34,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
 )
 from fastapi import Request as Request_http_parsing
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from starlette.types import Message
 from starlette.websockets import WebSocket
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -1561,7 +1561,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 _SCALAR_DEFAULTS = {
     "num_retries": getattr(litellm, "num_retries", None),

tests/unit/proxy/common_utils/test_model_deprecation.py

@@ -15,7 +15,7 @@ import pytest
 import litellm
 from litellm import Router
 from litellm._logging import verbose_logger
-from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
+from litellm.rust_bridge.host.model_cost import AliasedCostEntry
 from litellm.proxy.common_utils.model_deprecation import (
     _classify,
     _parse_deprecation_date,

tests/unit/proxy/hooks/test_dynamic_rate_limiter.py

@@ -14,7 +14,7 @@ from litellm import DualCache as DualCache_dynamic_rate, Router
 from litellm._uuid import uuid
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.proxy._types import UserAPIKeyAuth
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 from typing import Optional
 
@@ -139,7 +139,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 _SCALAR_DEFAULTS = {
     "num_retries": getattr(litellm, "num_retries", None),

tests/unit/proxy/hooks/test_prompt_injection_detection.py

@@ -17,7 +17,7 @@ from litellm.proxy.utils import ProxyLogging
 from litellm.router import Router
 from litellm import Router as Router_prompt_injection
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.utils import _invalidate_model_cost_lowercase_map
+from litellm.rust_bridge.host.model_cost import publish
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 
@@ -279,7 +279,7 @@ def isolate_litellm_state():
     for attr, original_value in original_state.items():
         if hasattr(litellm, attr):
             setattr(litellm, attr, original_value)
-    _invalidate_model_cost_lowercase_map()
+    publish(litellm.model_cost)
 
 _SCALAR_DEFAULTS = {
     "num_retries": getattr(litellm, "num_retries", None),