litellm/tests/test_litellm/proxy/hooks
devin-ai-integration[bot] ae01882535
feat(proxy): offload spend tracking to a pod-local collector sidecar (#40545)
* feat(proxy): offload spend tracking to a pod-local spend worker sidecar

py-spy on the gateway showed the post-response _PROXY_track_cost_callback,
spend-log and DBSpendUpdateWriter work running on the inference workers'
event loop, so a DB or Redis stall backed up the request path.

When LITELLM_SPEND_WORKER_ENABLED=true, _ProxyDBLogger serializes one compact
typed SpendEvent per success and hands it to a SpendEventProducer that ships
it over a unix socket (default) or loopback-only TCP to a sidecar started as
`python -m gateway.spend_worker`. The sidecar runs the unchanged
_ProxyDBLogger pipeline against the pod's PgBouncer (pooled_database_url).
When the sidecar is unreachable, the buffer is full, or the gateway shuts
down with events still queued or in flight, the producer applies
LITELLM_SPEND_WORKER_ON_UNAVAILABLE (fallback in-process, or drop). The
sidecar half-closes producers on SIGTERM and drains, the producer treats
EOF as unavailable, and the gateway flushes buffered spend counters on
shutdown. The sidecar honors LITELLM_LOG so its writes are visible in its
own process log.

Helm: both charts gain an opt-in spend-worker sidecar container sharing an
emptyDir socket dir, and the componentized chart's HPA uses a
ContainerResource CPU metric scoped to the gateway container so sidecar
CPU does not drive inference scaling.

* feat(terraform): opt-in spend-worker sidecar for the AWS and GCP gateway stacks

Adds spend_worker_* inputs to both modules. On ECS Fargate the sidecar is a second, non-essential container in the gateway task; on Cloud Run it is a second container in the gateway service. Both listen on loopback TCP, share the gateway's DB/Redis/secret env, and set LITELLM_JOB_ROLE=spend_worker. Disabled by default. Plan-only tests cover both, and the terraform CI workflow now runs the gcp module too

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

* test(proxy): retrieve a completed batch in the in-process spend path test

The base now defers cost tracking for batches that are still in flight, so an in_progress batch never reaches update_database

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

* refactor(proxy): rename the spend worker sidecar to collector

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

* fix(proxy): run the collector from the installed litellm package and finish in-flight fallbacks on shutdown

The sidecar command becomes python -m litellm.proxy.collector so the classic image, whose runtime
stage copies only the installed package, can run it. The module now assembles DATABASE_URL and the
pod-local pgbouncer URL itself, replacing gateway/collector.py

The componentized collector sidecar inherits gateway.volumeMounts so custom CA mounts reach it.
SpendEventProducer shields an in-progress fallback from the writer task cancellation so close()
no longer loses an event already handed to the in-process pipeline

Helpers used across modules (address_argument, should_store_prompts_and_responses_in_spend_logs,
flush_spend_counters_on_shutdown) become public so the change adds no reportPrivateUsage errors

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

* ci(terraform): drop the gcp job duplicated by the aws/gcp matrix

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

* fix(collector): keep metrics env off the classic sidecar and reject shared loopback ports

The classic chart no longer hands PROMETHEUS_METRICS_PORT and the billing metrics env to the collector container, and gives it the same /.npm scratch mount as the proxy on a read-only root. AWS and GCP now refuse a plan where the spend collector and the metrics sidecar bind the same loopback port. A regression test drives a sidecar crash mid-stream on asyncio and uvloop and checks no event is billed by both the sidecar and the in-process fallback; the producer docstring spells out why a failed drain() cannot double count

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

* style(proxy): format pooled_database_url after the pgbouncer rebase

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

* fix(proxy): keep the cache-hit preset key and survive dead producers on collector drain

Cache hits updated the logging object after the early return, so the offloaded spend event carried
preset_cache_key=None and the collector re-hashed reconstructed kwargs. Also guard write_eof() against
producer transports uvloop already closed so one dead connection cannot abort the drain

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

* fix(terraform): keep the gcp collector port off the metrics sidecar health port

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

* fix(proxy): collector connects to Postgres directly under IAM or Entra token auth

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

* fix(proxy): mark the collector's DATABASE_URL as pooled when it uses the pod's pgbouncer

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

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-10 17:14:13 -07:00
..
litellm_skills test: use local model cost map in import-isolation subprocess 2026-08-30 12:47:15 -07:00
test_async_post_call_streaming_iterator_hook.py test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 2026-08-22 09:25:58 -07:00
test_batch_enqueued_tokens.py test(batch): prove an open Redis breaker keeps enqueued-token reservations quiet 2026-09-10 14:26:11 -07:00
test_batch_file_validation.py Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_batch_enqueued_token_limit 2026-08-19 19:31:46 -07:00
test_dynamic_rate_limiter.py fix(dynamic_rate_limiter): inject clock so active-project window is stable within a request (#32299) 2026-07-06 18:12:47 -07:00
test_dynamic_rate_limiter_v3.py fix(rate_limiter): attach v3 priority rate limit headers on /v1/messages (#37228) 2026-09-09 15:14:10 -07:00
test_image_generation_guardrails.py test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 2026-08-22 09:25:58 -07:00
test_key_management_event_hooks.py test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 2026-08-22 09:25:58 -07:00
test_max_budget_limiter.py feat(proxy): add apply_user_budget_to_team_keys opt-in (#36102) 2026-08-07 15:40:13 +00:00
test_max_budget_per_session_limiter.py fix(redis): quiet every per-request Redis fallback while the breaker is open 2026-09-10 14:18:53 -07:00
test_max_iterations_limiter.py Agents - add max budget + tpm/rpm limiting per agent AND per agent session (#22849) 2026-03-07 19:12:42 -08:00
test_parallel_request_limiter.py Litellm oss 090626 (#30021) 2026-06-10 10:34:07 -07:00
test_parallel_request_limiter_v3.py test(limiter): inject the open-breaker Redis double instead of replacing script attributes 2026-09-10 14:37:39 -07:00
test_post_call_failure_hook_integration.py test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 2026-08-22 09:25:58 -07:00
test_post_call_response_headers_hook.py test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 2026-08-22 09:25:58 -07:00
test_post_call_streaming_hook_integration.py test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 2026-08-22 09:25:58 -07:00
test_post_call_success_hook_integration.py test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 2026-08-22 09:25:58 -07:00
test_prompt_injection_detection.py fix: prompt injection not working (#16701) 2025-11-17 20:04:57 -08:00
test_proxy_hooks_init.py [Fix] Proxy: Address Greptile feedback on hook-cycle PR 2026-05-04 20:13:31 -07:00
test_proxy_rate_limit_provider_field.py refactor(rate-limits): move the v3 limiter per-request stash off request metadata onto a ContextVar 2026-07-30 14:00:20 -07:00
test_proxy_track_cost_callback.py feat(proxy): offload spend tracking to a pod-local collector sidecar (#40545) 2026-09-10 17:14:13 -07:00
test_rate_limiter_toctou.py test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 2026-08-22 09:25:58 -07:00
test_send_invite_email.py feat(proxy)!: default audit logs on for enterprise licenses (#37518) 2026-08-19 18:49:21 -07:00
test_sensitive_data_routing.py fix(redis): quiet every per-request Redis fallback while the breaker is open 2026-09-10 14:18:53 -07:00
test_tpm_concurrent.py test(token_counter): type the project ITPM lag test's request body 2026-09-09 14:20:56 -07:00
test_user_management_event_hooks.py test(proxy): cover the user-created audit hook's database read-back 2026-07-31 19:48:30 +00:00