Merge eb77fc1651 into 72228db5c6

litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs

@@ -206,16 +206,27 @@ pub(crate) struct PyAtomicOverLimit {
     index: usize,
     current: i64,
     limit: i64,
+    window_start: Option<i64>,
+    window_observed_at: Option<i64>,
 }
 
 #[pymethods]
 impl PyAtomicOverLimit {
     #[new]
-    fn new(index: usize, current: i64, limit: i64) -> Self {
+    #[pyo3(signature = (index, current, limit, window_start = None, window_observed_at = None))]
+    fn new(
+        index: usize,
+        current: i64,
+        limit: i64,
+        window_start: Option<i64>,
+        window_observed_at: Option<i64>,
+    ) -> Self {
         Self {
             index,
             current,
             limit,
+            window_start,
+            window_observed_at,
         }
     }
 }
@@ -313,10 +324,14 @@ fn atomic(reply: AtomicReply) -> PyAtomicReply {
             index,
             current,
             limit,
+            window_start,
+            window_observed_at,
         } => PyAtomicReply::OverLimit(PyAtomicOverLimit {
             index,
             current,
             limit,
+            window_start: Some(window_start),
+            window_observed_at: Some(window_observed_at),
         }),
     }
 }

litellm-rust/crates/redis-store/src/rate_limits/lua.rs

@@ -60,7 +60,7 @@ pub(super) const CHECK_AND_INCREMENT_BY_N: &str = r#"
 --
 -- Return on success:
 --     { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... }
--- Return on over-limit: { 1, descriptor_index, current_counter, limit }
+-- Return on over-limit: { 1, descriptor_index, current_counter, limit, window_start, now }
 local time_reply = redis.call('TIME')
 local now = tonumber(time_reply[1])
 local descriptor_count = #KEYS / 2
@@ -94,7 +94,8 @@ for i = 1, descriptor_count do
         blocked = current_counter >= limit
     end
     if blocked then
-        return { 1, i, current_counter, limit }
+        local active_window_start = window_expired and now or tonumber(window_start)
+        return { 1, i, current_counter, limit, active_window_start, now }
     end
 
     descriptor_state[i] = { window_expired, current_counter, window_start }

litellm-rust/crates/redis-store/src/rate_limits/scripts.rs

@@ -264,10 +264,21 @@ fn atomic_reply(reply: Value) -> Result<AtomicReply, Error> {
             })
             .collect::<Result<_, _>>()
             .map(AtomicReply::Applied),
-        Some((Value::Int(1), [index_value, current, limit])) => Ok(AtomicReply::OverLimit {
+        Some((
+            Value::Int(1),
+            [
+                index_value,
+                current,
+                limit,
+                window_start,
+                window_observed_at,
+            ],
+        )) => Ok(AtomicReply::OverLimit {
             index: index(index_value)?,
             current: integer(current)?,
             limit: integer(limit)?,
+            window_start: integer(window_start)?,
+            window_observed_at: integer(window_observed_at)?,
         }),
         _ => Err(unexpected_reply("rate limiter script reply has no status")),
     }
@@ -301,3 +312,31 @@ fn acquire_reply(reply: Value) -> Result<AcquireReply, Error> {
         _ => Err(unexpected_reply("rate limiter script reply has no status")),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::two_seconds_remaining(9_998)]
+    #[case::one_second_remaining(9_999)]
+    fn atomic_rejection_keeps_the_server_window_clock(#[case] observed_at: i64) {
+        let reply = Value::Array(
+            [1, 1, 10, 10, 9_910, observed_at]
+                .into_iter()
+                .map(Value::Int)
+                .collect(),
+        );
+        assert_eq!(
+            atomic_reply(reply).unwrap(),
+            AtomicReply::OverLimit {
+                index: 0,
+                current: 10,
+                limit: 10,
+                window_start: 9_910,
+                window_observed_at: observed_at,
+            }
+        );
+    }
+}

litellm-rust/crates/redis-store/src/rate_limits/types.rs

@@ -30,6 +30,8 @@ pub enum AtomicReply {
         index: usize,
         current: i64,
         limit: i64,
+        window_start: i64,
+        window_observed_at: i64,
     },
 }
 

litellm-rust/crates/redis-store/tests/rate_limits.rs

@@ -397,18 +397,25 @@ async fn check_and_increment_is_all_or_nothing(
     assert_eq!(applied.iter().map(|a| a.count).collect::<Vec<_>>(), [1, 60]);
     assert!(applied.iter().all(|a| (a.window_start - now).abs() <= 2));
 
-    assert_eq!(
-        limits
-            .check_and_increment(RedisSend::Now, &both(1, 41))
-            .unwrap()
-            .await
-            .unwrap(),
-        AtomicReply::OverLimit {
-            index: 1,
-            current: 60,
-            limit: 100
-        }
-    );
+    let denied = limits
+        .check_and_increment(RedisSend::Now, &both(1, 41))
+        .unwrap()
+        .await
+        .unwrap();
+    let AtomicReply::OverLimit {
+        index,
+        current,
+        limit,
+        window_start,
+        window_observed_at,
+    } = denied
+    else {
+        panic!("{denied:?}");
+    };
+    assert_eq!((index, current, limit), (1, 60, 100));
+    assert_eq!(window_start, applied[1].window_start);
+    assert!(window_observed_at >= now);
+    assert!(window_observed_at <= redis_now(&raw).await);
     assert_eq!(
         raw_text(&raw, "{model_per_key:k1:gpt}:requests").await,
         Some("1".to_owned())
@@ -426,18 +433,21 @@ async fn check_and_increment_is_all_or_nothing(
     assert!(
         matches!(at_limit, AtomicReply::Applied(ref applied) if applied[0].count == 2 && applied[1].count == 100)
     );
-    assert_eq!(
-        limits
-            .check_and_increment(RedisSend::Now, &both(0, 0))
-            .unwrap()
-            .await
-            .unwrap(),
+    let read_only_denied = limits
+        .check_and_increment(RedisSend::Now, &both(0, 0))
+        .unwrap()
+        .await
+        .unwrap();
+    assert!(matches!(
+        read_only_denied,
         AtomicReply::OverLimit {
             index: 0,
             current: 2,
-            limit: 2
-        }
-    );
+            limit: 2,
+            window_start: start,
+            window_observed_at: observed,
+        } if start == applied[0].window_start && observed >= window_observed_at && observed <= redis_now(&raw).await
+    ));
 }
 
 #[rstest]

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -15,6 +15,7 @@ from contextlib import asynccontextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
+from math import ceil
 from types import MappingProxyType
 from typing import (
     TYPE_CHECKING,
@@ -412,6 +413,9 @@ class RateLimitStatus(TypedDict):
     # models' project ITPM/OTPM in one call, or a request carrying multiple
     # rate-limited tags, produces statuses sharing the same descriptor_key.
     descriptor_value: NotRequired[ReadOnly[str]]
+    window_start: NotRequired[ReadOnly[float]]
+    window_size: NotRequired[ReadOnly[int]]
+    window_observed_at: NotRequired[ReadOnly[float]]
 
 
 class RateLimitResponse(TypedDict):
@@ -1200,6 +1204,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         cache_values: Sequence[CounterState],
         key_metadata: Mapping[RateLimitScope, WindowKeyMetadata],
         read_only: bool = False,
+        observed_at: float | None = None,
     ) -> RateLimitResponse:
         """
         Check if the cache values are over the limit.
@@ -1208,7 +1213,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         statuses: Final[list[RateLimitStatus]] = []
         overall_code = "OK"
 
-        for counter, (_window_start, counter_value) in zip(counters, cache_values):
+        for counter, (window_start, counter_value) in zip(counters, cache_values):
             item_code = "OK"
             metadata = key_metadata[counter.scope]
             current_limit = metadata["requests_limit"] if counter.kind == "requests" else metadata["tokens_limit"]
@@ -1221,6 +1226,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
 
             limit_remaining = current_limit - int(counter_value) if counter_value is not None else current_limit
 
+            if item_code == "OVER_LIMIT" and window_start is not None:
+                statuses.append(
+                    RateLimitStatus(
+                        code=item_code,
+                        current_limit=current_limit,
+                        limit_remaining=limit_remaining,
+                        rate_limit_type=counter.kind,
+                        descriptor_key=metadata["descriptor_key"],
+                        descriptor_value=metadata["descriptor_value"],
+                        window_start=float(window_start),
+                        window_size=self.window_size,
+                        window_observed_at=(
+                            observed_at if observed_at is not None else self._get_current_time().timestamp()
+                        ),
+                    )
+                )
+                continue
             statuses.append(
                 {
                     "code": item_code,
@@ -1412,7 +1434,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         windowed_response = RateLimitResponse(overall_code="OK", statuses=[])
         if counters:
             local_response: Final = self.is_cache_list_over_limit(
-                counters, self._local_counter_states(counters), key_metadata, read_only=read_only
+                counters, self._local_counter_states(counters), key_metadata, read_only=read_only, observed_at=now
             )
             if local_response["overall_code"] == "OVER_LIMIT":
                 return local_response
@@ -1433,7 +1455,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                     window_size=self.window_size,
                 )
 
-            windowed_response = self.is_cache_list_over_limit(counters, cache_values, key_metadata, read_only=read_only)
+            windowed_response = self.is_cache_list_over_limit(
+                counters, cache_values, key_metadata, read_only=read_only, observed_at=now
+            )
             if windowed_response["overall_code"] == "OVER_LIMIT":
                 return windowed_response
 
@@ -2074,20 +2098,32 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         level when modifying this code.
         """
         if isinstance(reply, AtomicOverLimit):
-            meta = per_counter_meta[reply.index]
-            return RateLimitResponse(
-                overall_code="OVER_LIMIT",
-                statuses=[
-                    RateLimitStatus(
-                        code="OVER_LIMIT",
-                        current_limit=reply.limit,
-                        limit_remaining=max(0, reply.limit - reply.current),
-                        rate_limit_type=meta["rate_limit_type"],
-                        descriptor_key=meta["descriptor_key"],
-                        descriptor_value=meta["descriptor_value"],
-                    )
-                ],
+            over_limit_meta: Final = per_counter_meta[reply.index]
+            window_start: Final = reply.window_start
+            window_observed_at: Final = reply.window_observed_at
+            status: Final = (
+                RateLimitStatus(
+                    code="OVER_LIMIT",
+                    current_limit=reply.limit,
+                    limit_remaining=max(0, reply.limit - reply.current),
+                    rate_limit_type=over_limit_meta["rate_limit_type"],
+                    descriptor_key=over_limit_meta["descriptor_key"],
+                    descriptor_value=over_limit_meta["descriptor_value"],
+                    window_start=float(window_start),
+                    window_size=over_limit_meta["window_size"],
+                    window_observed_at=float(window_observed_at),
+                )
+                if window_start is not None and window_observed_at is not None
+                else RateLimitStatus(
+                    code="OVER_LIMIT",
+                    current_limit=reply.limit,
+                    limit_remaining=max(0, reply.limit - reply.current),
+                    rate_limit_type=over_limit_meta["rate_limit_type"],
+                    descriptor_key=over_limit_meta["descriptor_key"],
+                    descriptor_value=over_limit_meta["descriptor_value"],
+                )
             )
+            return RateLimitResponse(overall_code="OVER_LIMIT", statuses=[status])
         statuses: Final[list[RateLimitStatus]] = [
             RateLimitStatus(
                 code="OK",
@@ -2121,7 +2157,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         """
         # Use a single 'now' for the duration of this critical section so all
         # descriptors evaluate window expiry consistently.
-        now_int: Final = int(self._get_current_time().timestamp())
+        now: Final = self._get_current_time().timestamp()
+        now_int: Final = int(now)
 
         # Pass 1: read state, validate.
         descriptor_state: Final[list[AtomicCounterState]] = []
@@ -2147,6 +2184,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                             rate_limit_type=meta["rate_limit_type"],
                             descriptor_key=meta["descriptor_key"],
                             descriptor_value=meta["descriptor_value"],
+                            window_start=float(now_int if window_expired else int(window_start)),
+                            window_size=meta["window_size"],
+                            window_observed_at=now,
                         )
                     ],
                 )
@@ -3299,9 +3339,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                 )
                 descriptor_value = matching_descriptor["value"] if matching_descriptor is not None else "unknown"
 
-                now = self._get_current_time().timestamp()
-                reset_time = now + self.window_size
-                reset_time_formatted = datetime.fromtimestamp(reset_time, tz=timezone.utc).strftime(
+                window_start: Final = status.get("window_start")
+                window_size: Final = status.get("window_size")
+                window_observed_at: Final = status.get("window_observed_at")
+                reset_time: Final = (
+                    window_start + window_size
+                    if window_start is not None and window_size is not None and window_observed_at is not None
+                    else self._get_current_time().timestamp() + self.window_size
+                )
+                retry_after: Final = (
+                    max(1, ceil(reset_time - window_observed_at))
+                    if window_observed_at is not None
+                    else self.window_size
+                )
+                reset_time_formatted: Final = datetime.fromtimestamp(reset_time, tz=timezone.utc).strftime(
                     "%Y-%m-%d %H:%M:%S UTC"
                 )
 
@@ -3316,14 +3367,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                     f"Limit resets at: {reset_time_formatted}"
                 )
 
+                headers: Final = {
+                    "retry-after": str(retry_after),
+                    "rate_limit_type": str(status["rate_limit_type"]),
+                    "reset_at": reset_time_formatted,
+                }
                 resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(requested_model)
                 raise ProxyRateLimitError(
                     detail=detail,
-                    headers={
-                        "retry-after": str(self.window_size),
-                        "rate_limit_type": str(status["rate_limit_type"]),
-                        "reset_at": reset_time_formatted,
-                    },
+                    headers=headers,
                     rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]),
                     model=resolved_model,
                     llm_provider=llm_provider,

litellm/rust_bridge/_native.pyi

@@ -204,13 +204,24 @@ class AtomicApplied:
 
 @final
 class AtomicOverLimit:
-    def __new__(cls, index: int, current: int, limit: int) -> AtomicOverLimit: ...
+    def __new__(
+        cls,
+        index: int,
+        current: int,
+        limit: int,
+        window_start: int | None = None,
+        window_observed_at: int | None = None,
+    ) -> AtomicOverLimit: ...
     @property
     def index(self) -> int: ...
     @property
     def current(self) -> int: ...
     @property
     def limit(self) -> int: ...
+    @property
+    def window_start(self) -> int | None: ...
+    @property
+    def window_observed_at(self) -> int | None: ...
 
 @final
 class GuardedReply:

tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py

@@ -56,6 +56,7 @@ from litellm.proxy.spend_tracking.input_tokens import models_counted_for
 from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
 from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
 from litellm.redis.stores.rate_limits import (
+    AtomicOverLimit,
     CounterIncrement,
     CounterKind,
     LocalValue,
@@ -6854,6 +6855,123 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension():
     assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"]
 
 
+@pytest.mark.asyncio
+async def test_rpm_rate_limit_headers_follow_fixed_window_reset(time_controller):
+    time_controller._current = datetime.fromtimestamp(1_000, tz=timezone.utc)
+    cache: Final = DualCache()
+    handler: Final = _PROXY_MaxParallelRequestsHandler(
+        internal_usage_cache=InternalUsageCache(cache),
+        time_provider=time_controller.now,
+    )
+    auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-reset-rpm"), rpm_limit=1)
+
+    await handler.async_pre_call_hook(auth, cache, {}, "")
+    time_controller.advance(1)
+    with pytest.raises(HTTPException) as first_denial:
+        await handler.async_pre_call_hook(auth, cache, {}, "")
+
+    time_controller.advance(58)
+    with pytest.raises(HTTPException) as second_denial:
+        await handler.async_pre_call_hook(auth, cache, {}, "")
+
+    reset_at: Final = datetime.fromtimestamp(1_060, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
+    assert first_denial.value.headers["reset_at"] == reset_at
+    assert second_denial.value.headers["reset_at"] == reset_at
+    assert first_denial.value.headers["retry-after"] == "59"
+    assert second_denial.value.headers["retry-after"] == "1"
+
+
+@pytest.mark.asyncio
+async def test_rpm_headers_use_the_enforced_handler_window(time_controller):
+    time_controller._current = datetime.fromtimestamp(3_000, tz=timezone.utc)
+    handler: Final = _PROXY_MaxParallelRequestsHandler(
+        internal_usage_cache=InternalUsageCache(DualCache()),
+        time_provider=time_controller.now,
+    )
+    descriptor: Final = RateLimitDescriptor(
+        key="api_key",
+        value="rpm-handler-window",
+        rate_limit={"requests_per_unit": 1, "window_size": 90},
+    )
+
+    await handler.should_rate_limit(descriptors=[descriptor])
+    time_controller.advance(59)
+    denied: Final = await handler.should_rate_limit(descriptors=[descriptor])
+    with pytest.raises(HTTPException) as exc_info:
+        handler._handle_rate_limit_error(denied, [descriptor])
+
+    assert exc_info.value.headers["retry-after"] == "1"
+    assert exc_info.value.headers["reset_at"] == datetime.fromtimestamp(3_060, tz=timezone.utc).strftime(
+        "%Y-%m-%d %H:%M:%S UTC"
+    )
+
+    time_controller.advance(1.01)
+    assert (await handler.should_rate_limit(descriptors=[descriptor]))["overall_code"] == "OK"
+
+
+@pytest.mark.asyncio
+async def test_atomic_tpm_in_memory_rejection_has_window_timing(time_controller):
+    time_controller._current = datetime.fromtimestamp(2_059, tz=timezone.utc)
+    cache: Final = DualCache()
+    handler: Final = _PROXY_MaxParallelRequestsHandler(
+        internal_usage_cache=InternalUsageCache(cache),
+        time_provider=time_controller.now,
+    )
+    descriptor: Final = RateLimitDescriptor(
+        key="api_key",
+        value="atomic-tpm-reset",
+        rate_limit={"tokens_per_unit": 10, "window_size": 60},
+    )
+    counter: Final = handler.rate_limit_counter("api_key", "atomic-tpm-reset", "tokens")
+    handler.rate_limits.set_local(counter.window, "2000", 60)
+    handler.rate_limits.set_local(counter, 10, 60)
+
+    response: Final = await handler.atomic_check_and_increment_by_n(
+        descriptors=[descriptor],
+        increments=[{"tokens": 1}],
+    )
+    with pytest.raises(HTTPException) as exc_info:
+        handler._handle_rate_limit_error(response, [descriptor])
+
+    assert exc_info.value.headers["retry-after"] == "1"
+    assert exc_info.value.headers["reset_at"] == datetime.fromtimestamp(2_060, tz=timezone.utc).strftime(
+        "%Y-%m-%d %H:%M:%S UTC"
+    )
+
+    time_controller.advance(1)
+    assert (await handler.atomic_check_and_increment_by_n(descriptors=[descriptor], increments=[{"tokens": 1}]))[
+        "overall_code"
+    ] == "OK"
+
+
+@pytest.mark.parametrize("redis_now, expected_retry_after", [(9_998, "2"), (9_999, "1")])
+def test_atomic_redis_rejection_uses_server_window_clock(redis_now: int, expected_retry_after: str) -> None:
+    handler: Final = _PROXY_MaxParallelRequestsHandler(
+        internal_usage_cache=InternalUsageCache(DualCache()),
+        time_provider=lambda: datetime.fromtimestamp(100_000, tz=timezone.utc),
+    )
+    descriptor: Final = RateLimitDescriptor(
+        key="api_key",
+        value="redis-tpm-reset",
+        rate_limit={"tokens_per_unit": 10, "window_size": 90},
+    )
+    _counters, meta = handler._build_descriptor_atomic_payload(
+        descriptor=descriptor, increment_amounts={"tokens": 1}
+    )
+    response: Final = handler._build_atomic_response(
+        reply=AtomicOverLimit(0, 10, 10, 9_910, redis_now),
+        per_counter_meta=meta,
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        handler._handle_rate_limit_error(response, [descriptor])
+
+    assert exc_info.value.headers["retry-after"] == expected_retry_after
+    assert exc_info.value.headers["reset_at"] == datetime.fromtimestamp(10_000, tz=timezone.utc).strftime(
+        "%Y-%m-%d %H:%M:%S UTC"
+    )
+
+
 STATIC_OUTPUT_FLOOR = 1024
 ONE_TOKEN_PROMPT = [{"role": "user", "content": "hello"}]
 ONE_TOKEN_PROMPT_INPUT_ESTIMATE = 1