fix(proxy): flush virtual-key model_max budget spend to Redis after success logging

[milan-berri]

Linear ticket

Resolves LIT-2893

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added tests in tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py (same location as the existing model max budget limiter suite). Note: the repo template points at tests/test_litellm/; these tests follow the existing pattern for this hook in proxy_unit_tests. If maintainers want a mirror under tests/test_litellm/, I can add one.
• My PR passes all unit tests on make test-unit (not run in this environment where make/uv is unavailable; locally: pytest tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py — 13 passed.)
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Screenshots / Proof of Fix

Manual repro (two proxy workers, one Redis):

• Config used the usual pattern: litellm_settings.cache + Redis cache_params, and enable_redis_auth_cache: true so user_api_key_cache shares the same Redis-backed DualCache as _PROXY_VirtualKeyModelMaxBudgetLimiter.
• A virtual key with model_max_budget set to a very small cap for a model alias (e.g. model group name matching enforcement).
• Proxy A (e.g. port 4058): first chat completion → HTTP 200 (spend recorded).
• Proxy B (e.g. port 4059), same Redis and same key: second completion → HTTP 400 with budget_exceeded, proving spend visible across workers after the fix.

Before the fix: Redis could show budget window metadata (e.g. virtual_key_budget_start_time:*) while virtual_key_spend:* was missing or stale across workers, because increments stayed in the in-memory queue and were never flushed (subclass never runs RouterBudgetLimiting.__init__, so the periodic Redis sync task never starts).

After the fix: async_log_success_event calls _push_in_memory_increments_to_redis() when dual_cache.redis_cache is set, so the increment pipeline runs and multi-worker enforcement matches intent.

Type

🐛 Bug Fix

Changes

• litellm/proxy/hooks/model_max_budget_limiter.py: After virtual-key and end-user model budget increments in async_log_success_event, await self._push_in_memory_increments_to_redis() when Redis is attached to dual_cache, so queued increments are written with the same mechanism as RouterBudgetLimiting.
• tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py:
  • test_async_log_success_event_pushes_redis_increments_when_redis_configured — asserts _push_in_memory_increments_to_redis is awaited once when redis_cache is set.
  • test_async_log_success_event_skips_redis_push_without_redis — asserts it is not awaited when Redis is not configured.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/proxy/hooks/model_max_budget_limiter.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR fixes a bug where _PROXY_VirtualKeyModelMaxBudgetLimiter never propagated per-request spend increments to Redis because its __init__ skips super().__init__(), meaning the periodic Redis-sync task from RouterBudgetLimiting was never scheduled. The fix eagerly flushes the in-memory increment queue to Redis at the end of each async_log_success_event call when a Redis backend is present.

• model_max_budget_limiter.py: Adds a 3-line guard after both increment blocks — if dual_cache.redis_cache is set, _push_in_memory_increments_to_redis() is awaited immediately. The underlying method already guards against an empty queue and uses asyncio.create_task for the actual pipeline write, so the extra call is non-blocking and safe.
• test_unit_test_max_model_budget_limiter.py: Two new async tests verify the Redis-present and Redis-absent branches using AsyncMock patches, confirming the method is or is not awaited as expected.

Confidence Score: 5/5

Safe to merge; the change is a targeted 3-line addition that restores expected multi-worker budget enforcement without altering any existing code paths.

The fix directly addresses the documented root cause, is well-tested with two new mock-based unit tests covering both branches, and the underlying flush method already guards against empty queues and uses non-blocking task creation. No existing behavior is altered for deployments without Redis.

No files require special attention; both changed files are straightforward.

Important Files Changed

Filename	Overview
litellm/proxy/hooks/model_max_budget_limiter.py	Adds a per-request Redis flush after spend increments in async_log_success_event; correctly targets the root cause (parent init never runs so the periodic task is never scheduled).
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py	Two new unit tests cover the Redis-configured and no-Redis branches of the fix; both use proper AsyncMock patching and assert the correct invocation behavior.

_{Reviews (1): Last reviewed commit: "Add unit tests for virtual-key model max..." | Re-trigger Greptile}