fix(rust): isolate explicit s3 keys from env tokens and treat 403 misses

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

@@ -1,14 +1,22 @@
-use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
+use aws_credential_types::{
+    Credentials as AwsCredentials,
+    provider::{ProvideCredentials, error::CredentialsError, future},
+};
 use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
 
 #[derive(Clone)]
 pub(crate) struct Credentials {
     config: AwsAuthConfig,
+    env: fn(&str) -> Option<String>,
 }
 
 impl Credentials {
     pub(crate) fn new(config: AwsAuthConfig) -> Self {
-        Self { config }
+        Self::with_env(config, |name| std::env::var(name).ok())
+    }
+
+    pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
+        Self { config, env }
     }
 }
 
@@ -18,7 +26,19 @@ impl ProvideCredentials for Credentials {
         Self: 'a,
     {
         future::ProvideCredentials::new(async {
-            resolve_credentials(self.config.clone(), &|name| std::env::var(name).ok())
+            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,
+                    secret_access_key,
+                    self.config.session_token.clone(),
+                    None,
+                    "litellm-s3-cache",
+                ));
+            }
+            resolve_credentials(self.config.clone(), &self.env)
                 .await
                 .map_err(|_| CredentialsError::provider_error("S3 cache authentication failed"))
         })
@@ -30,3 +50,52 @@ impl std::fmt::Debug for Credentials {
         f.debug_struct("Credentials").finish_non_exhaustive()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn explicit_keys_ignore_an_ambient_session_token() {
+        let provider = Credentials::with_env(
+            AwsAuthConfig {
+                access_key_id: Some("key".to_string()),
+                secret_access_key: Some("secret".to_string()),
+                region_name: Some("us-east-1".to_string()),
+                ..Default::default()
+            },
+            |name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()),
+        );
+        let credentials = provider.provide_credentials().await.unwrap();
+        assert_eq!(credentials.access_key_id(), "key");
+        assert_eq!(credentials.secret_access_key(), "secret");
+        assert_eq!(credentials.session_token(), None);
+    }
+
+    #[tokio::test]
+    async fn explicit_keys_keep_their_session_token() {
+        let provider = Credentials::new(AwsAuthConfig {
+            access_key_id: Some("key".to_string()),
+            secret_access_key: Some("secret".to_string()),
+            session_token: Some("t".to_string()),
+            region_name: Some("us-east-1".to_string()),
+            ..Default::default()
+        });
+        let credentials = provider.provide_credentials().await.unwrap();
+        assert_eq!(credentials.session_token(), Some("t"));
+    }
+
+    #[tokio::test]
+    async fn environment_keys_resolve_with_their_session_token() {
+        let provider = Credentials::with_env(AwsAuthConfig::default(), |name| 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,
+        });
+        let credentials = provider.provide_credentials().await.unwrap();
+        assert_eq!(credentials.access_key_id(), "env-key");
+        assert_eq!(credentials.secret_access_key(), "env-secret");
+        assert_eq!(credentials.session_token(), Some("env-token"));
+    }
+}

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

@@ -134,11 +134,13 @@ impl<C: CacheCodec> S3Cache<C> {
             Ok(output) => output,
             Err(error) => {
                 if let SdkError::ServiceError(service) = &error {
+                    let status = error
+                        .raw_response()
+                        .map(|response| response.status().as_u16());
                     let not_found = service.err().is_no_such_key()
-                        || error
-                            .raw_response()
-                            .map(|response| response.status().as_u16())
-                            == Some(404);
+                        || service.err().meta().code() == Some("AccessDenied")
+                        || status == Some(404)
+                        || status == Some(403);
                     if not_found {
                         return Ok(None);
                     }

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

@@ -125,6 +125,13 @@ async fn get_hit_miss_expired_and_invalid_entries() {
         )
         .mount(&server)
         .await;
+    Mock::given(method("GET"))
+        .and(path("/cache-bucket/team/denied"))
+        .respond_with(
+            ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
+        )
+        .mount(&server)
+        .await;
     Mock::given(method("GET"))
         .and(path("/cache-bucket/team/expired"))
         .respond_with(
@@ -147,6 +154,7 @@ async fn get_hit_miss_expired_and_invalid_entries() {
         Some(json!({"answer": 3}))
     );
     assert_eq!(cache.get_cache("missing", &context).unwrap(), None);
+    assert_eq!(cache.get_cache("denied", &context).unwrap(), None);
     assert_eq!(cache.get_cache("expired", &context).unwrap(), None);
     assert_eq!(
         cache.get_cache("malformed", &context),

litellm/rust_bridge/_native.pyi

@@ -165,6 +165,8 @@ class _CacheTestHandle:
         startup_nodes: Sequence[tuple[str, int]] | None = None,
     ) -> _CacheTestHandle: ...
     @staticmethod
+    def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
+    @staticmethod
     def s3(
         bucket: str,
         *,

tests/test_litellm_rust/support/s3_stub.py

@@ -51,7 +51,7 @@ class S3Stub:
             def do_PUT(self) -> None:
                 body: Final = self._read_body()
                 headers: Final = {name: self.headers[name] for name in _STORED_HEADERS if name in self.headers}
-                stub._objects[self._key()] = S3Object(body=body, headers=headers)
+                stub._objects = {**stub._objects, self._key(): S3Object(body=body, headers=headers)}
                 self.send_response(200)
                 self.send_header("ETag", '"stub"')
                 self.send_header("Content-Length", "0")
@@ -83,7 +83,9 @@ class S3Stub:
                 if send_body:
                     self.wfile.write(entry.body)
 
-            def log_message(self, format: str, *args: object) -> None:  # noqa: A002
+            def log_message(
+                self, format: str, *args: object
+            ) -> None:  # BaseHTTPRequestHandler.log_message names this parameter format
                 pass
 
         self._server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
@@ -100,7 +102,7 @@ class S3Stub:
         return self._objects
 
     def put_object(self, key: str, body: bytes, headers: dict[str, str] | None = None) -> None:
-        self._objects[key] = S3Object(body=body, headers=headers or {})
+        self._objects = {**self._objects, key: S3Object(body=body, headers=headers or {})}
 
     def expires(self, key: str) -> object:
         header: Final = self._objects[key].headers.get("expires")

tests/test_litellm_rust/test_cache.py

@@ -11,6 +11,7 @@ from collections.abc import Generator
 from datetime import datetime
 from types import SimpleNamespace
 from typing import Final, Protocol, cast
+from unittest.mock import Mock
 from urllib.parse import urlparse
 
 import boto3
@@ -24,8 +25,8 @@ import litellm
 from litellm.caching.azure_blob_cache import AzureBlobCache
 from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache
 from litellm.caching.in_memory_cache import InMemoryCache
-from litellm.caching.s3_cache import S3Cache
 from litellm.caching.redis_cluster_cache import RedisClusterCache
+from litellm.caching.s3_cache import S3Cache
 from litellm.rust_bridge import _native
 from litellm.types.caching import LiteLLMCacheType
 from tests.test_litellm_rust.support.isolation import rebound
@@ -416,15 +417,20 @@ def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure
     assert handle.backend == "azure-blob"
     account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}")
     with pytest.raises(TypeError, match="containers must match"):
-        _native._CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade(
-            azure_blob_facade
-        )
+        _native._CacheTestHandle.azure_blob(
+            account_url, f"{backend.container_client.container_name}-other"
+        )._bind_facade(azure_blob_facade)
     handle._bind_facade(azure_blob_facade)
     resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade))
     native: Final = resolver.resolve()
     assert native.kind == "native"
 
-    response: Final = {"choices": [{"text": "caf\u00e9 \u2603"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
+    response: Final = {
+        "choices": [{"text": "caf\u00e9 \u2603"}],
+        "usage": {"total_tokens": 3},
+        "flag": True,
+        "empty": None,
+    }
     native.store({**request("sync"), "ttl_seconds": 0.001}, response)
     native.store(request("sync"), {"choices": [{"text": "second"}]})
     time.sleep(0.01)
@@ -479,7 +485,9 @@ async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_pyt
     await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2})
     time.sleep(0.01)
     assert await binding.async_lookup(request("async")) == {"value": 2}
-    assert await backend.async_get_cache("async") == json.loads(backend.container_client.download_blob("async").readall())
+    assert await backend.async_get_cache("async") == json.loads(
+        backend.container_client.download_blob("async").readall()
+    )
     assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2}
 
     await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}])
@@ -624,11 +632,11 @@ def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_
     binding: Final = resolver.resolve()
     assert binding.kind == "native"
 
-    calls: Final = []
-    facade.cache.s3_client.meta.events.register("before-call.s3.*", lambda **_kwargs: calls.append(1))
+    handler: Final = Mock()
+    facade.cache.s3_client.meta.events.register("before-call.s3.*", handler)
     binding.store(request("native"), {"answer": 1})
     assert binding.lookup(request("native")) == {"answer": 1}
-    assert calls == []
+    assert handler.call_count == 0
     assert "team/native" in s3_stub.objects
 
     with rebound(facade.cache, "bucket_name", "other"):
@@ -756,7 +764,9 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n
 
     await binding.async_flush()
 
-    remaining: Final = tuple(sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)))
+    remaining: Final = tuple(
+        sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node))
+    )
     assert remaining == (), remaining
     assert client.get("unscoped") == b"stays"
     client.delete("unscoped")