mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
c80f1c2ec8
* fix(responses): write one spend-log row per native WebSocket turn Each response.create turn on a native Responses WebSocket connection now logs its own call as it finishes, with its own usage and the resp_ id the client saw as request_id. A failed turn is booked alone and no longer flips the other turns. * fix(responses): bill the usage a failed native WebSocket turn reports * fix(proxy): pin the limiter slot owner to the proxy call id and refund unsettled reservations on disconnect The v3 limiter stashed whichever litellm_call_id the request body carried as the slot owner, so a client-supplied id on a pass-through body left the proxy's own success and failure events unable to release the slot. The owner is now the proxy's logging call id. A WebSocket connection whose turns log under their own ids never settles the TPM, ITPM, or OTPM estimate the connection reserved, so the disconnect release now refunds any reservation no callback settled, once. * fix(proxy): bill failed WebSocket turns and refund only the reserved window on disconnect A failed turn's reported usage now counts toward the key's TPM and the project ITPM/OTPM scopes the same way a success does, and closing the socket hands back only the estimate still unused in the window that holds it, so a window that already rolled is left alone. An id-less error event is booked against the newest pending turn, the one the backend rejected at parse time, and a turn that nests its input under response logs its user message. A connection's one budget reservation goes to the first turn that settles it, every later turn charging its cost directly. * test(proxy): expect the settling call id on a pass-through budget reservation * fix(proxy): charge successful Responses WebSocket turns to TPM and ITPM/OTPM A native Responses WebSocket turn reaches the rate limiter's success callback with its raw event list as the response and no combined_usage_object, so the turn's tokens never reached the key's TPM window or the project's input and output token limits. Read the usage from the turn's events when the call type is a Responses WebSocket turn. * fix(proxy): release a rejected pass-through's slot and count each websocket turn once toward project quotas The limiter's failure hook looked up the parallel slot by the request body's litellm_call_id, while the slot belongs to the proxy's logging call id, so a pass-through request a pre-call hook rejected held its slot for an hour when its body carried an id. Both now resolve the id the same way. Each response.create frame already charges its estimate to project ITPM/OTPM, so the turn's own usage no longer adds a second count. The cost callback's failure hook read the body's litellm_call_id through a strict model and raised on a non-string value, which dropped the failure row. It now reads a string id or none. * fix(proxy): keep per-turn WebSocket limits, request bodies and collector reservations right Free a rejected request's parallel slot through the request it runs in, and skip native Responses WebSocket connections, whose disconnect frees the slot. Count WebSocket turns toward the connection's tag TPM. In collector mode, claim the shared budget reservation before the spend event is built, so only the first turn settles it. Store each turn's own frame as its request body, and send an id-less error to the oldest waiting turn unless an earlier turn is streaming. * fix(proxy): store a failed WebSocket turn's own request and count tag TPM only for calls that carry the tag * fix(responses): log the client's turn frame, keep the payload id and reservation shape for non-WebSocket calls * fix(proxy): keep the legacy limiter's WebSocket slot until the socket closes * fix(router): count native Responses WebSocket turns toward deployment limits and budgets * fix(router): re-check deployment budget and RPM on each Responses WebSocket turn * test(router): list the WebSocket turn admission method as covered through its partial * fix(responses): give a refused WebSocket turn's error frame a 429 status so Codex ends the turn * fix(router): leave ITPM/OTPM to socket open and record a refused WebSocket turn as the socket's failure * fix(proxy): log the routed model group on a WebSocket turn's failure row * fix(router): count WebSocket turn tokens toward TPM and re-run routing strategy checks on each turn * test(integration): cover per-turn admission on Responses WebSockets Adds integration cells for budget, RPM, TPM, a mid-socket TPM edit, usage-based-routing-v2, the project quota frame, and the open-time failure frame on the native and managed relays, plus HTTP controls for TPM counting and the failure row model_group * fix(responses): count WebSocket turn tokens once and end a socket whose turn check fails A native Responses WebSocket turn is now logged as an async call, so a proxy with a sync success callback such as the response cache no longer runs each CustomLogger's sync success hook on top of the async one. That second run counted every turn's tokens twice toward deployment TPM. A turn whose admission check raises an unexpected error now gets a 500 error frame, the socket closes with 1011, and the error becomes the socket's failure, on both the native and the managed relay, instead of the socket dropping with no frame and no failure row. A refused turn's failure row now carries the request's tags from the router's metadata, the same way it already carries the model group. * test(integration): cover a crashed WebSocket turn check and failure-row tags A turn whose admission check raises gets a 500 server_error frame, the socket closes with 1011 and its one failure row carries the error, on both relays. The refused turn's failure row keeps the caller's tags, and HTTP failure rows on chat, responses and messages keep them too. * fix(responses): keep a WebSocket open after a crashed turn check and keep only string tags on failure rows A turn whose admission check raises an unexpected error still gets a 500 error frame, but the socket now stays open, like a refused turn. Closing it cut any turn still streaming on the same socket, which a client that sends its next response.create early would lose. A refused turn's failure row now takes the router's tags only when every tag is a string, so the fallback never copies a caller's unvalidated tags onto the row. * test(integration): cover a turn still streaming when the next WebSocket turn's check crashes A native turn still streaming when the next turn's admission check raises now has its own cell. It must complete and be billed at the turn's price. The crashed-check cell also keeps the fault set for one more turn, so a relay that stopped checking turns after a crash goes red. * fix(responses): run per-turn guardrails before the deployment's turn check on WebSockets A turn after the first now runs the project quota check, then the proxy's pre_call guardrails, then the deployment's turn admission, on both the native and the managed relay, the order HTTP calls use. A turn a guardrail blocks gets the guardrail's error even past the budget and no longer counts toward the deployment's rpm or runs its pre-call checks. * fix(responses): bill a WebSocket turn's guardrail cost when the deployment refuses it A turn the per-turn guardrail passes and the deployment then refuses now writes its own failed spend row with that turn's guardrail information, on both relays, the way a guardrail-blocked frame and an HTTP refusal already do * fix(responses): bill a refused WebSocket turn on one row only A refused turn whose guardrails recorded a cost is already billed on its own row, so it no longer also becomes the connection's close-time failure row * fix(responses): bill a refused WebSocket turn's guardrail cost on the socket's failure row The turns a deployment refuses after their guardrails passed now add their guardrail records to the socket's close-time failure row instead of a row of their own, so that row's failure hooks, model group and call type cover them * fix(responses): bill a WebSocket's refused turns when the session limit closes it * fix(responses): bill refused WebSocket turns when the session limit cancels the relay * fix(responses): free the WebSocket slot at the session limit without waiting on the failure row * fix(responses): keep the session limit's second cancel so a relay that swallows the first frees the slot * fix(responses): free the WebSocket slot before the close-time failure row and drain pending rows at shutdown * fix(proxy): commit spend before waiting on pending logging tasks at shutdown and again after * fix(responses): start the close-time failure row before the slot release so a late cancel cannot drop it * fix(responses): build the close-time failure row inside its task so a bad guardrail record cannot skip the crash close * fix(proxy): always commit spend again after the shutdown wait for pending logging tasks --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
1382 lines
53 KiB
Python
1382 lines
53 KiB
Python
"""Behavior pins for proxy_server lifecycle, helpers, and small utilities.
|
|
|
|
Pins covered:
|
|
- ``proxy_startup_event``
|
|
- ``proxy_shutdown_event``
|
|
- ``_initialize_shared_aiohttp_session``
|
|
- ``cleanup_router_config_variables``
|
|
- ``save_worker_config``
|
|
- ``initialize``
|
|
- ``load_from_azure_key_vault``
|
|
- ``cost_tracking``
|
|
- ``_resolve_typed_dict_type``
|
|
- ``_resolve_pydantic_type``
|
|
- ``get_litellm_model_info``
|
|
- ``run_ollama_serve``
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from pydantic import BaseModel
|
|
from typing_extensions import TypedDict
|
|
|
|
import litellm.proxy.proxy_server as ps
|
|
from litellm.proxy.proxy_server import (
|
|
ProxyStartupEvent,
|
|
initialize_shared_aiohttp_session,
|
|
_resolve_pydantic_type,
|
|
_resolve_typed_dict_type,
|
|
cleanup_router_config_variables,
|
|
cost_tracking,
|
|
get_litellm_model_info,
|
|
initialize,
|
|
initialize_from_worker_config,
|
|
load_from_azure_key_vault,
|
|
proxy_shutdown_event,
|
|
proxy_startup_event,
|
|
run_ollama_serve,
|
|
save_worker_config,
|
|
)
|
|
|
|
from litellm.proxy.management_helpers.audit_logs import track_audit_task
|
|
|
|
from .conftest import normalize
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# cleanup_router_config_variables
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_cleanup_router_config_variables_resets_globals(monkeypatch):
|
|
monkeypatch.setattr(ps, "master_key", "sk-sentinel", raising=False)
|
|
monkeypatch.setattr(ps, "user_config_file_path", "/tmp/config.yaml", raising=False)
|
|
monkeypatch.setattr(ps, "user_custom_auth", lambda x: x, raising=False)
|
|
monkeypatch.setattr(ps, "health_check_interval", 42, raising=False)
|
|
monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
|
|
monkeypatch.setattr(ps, "heuristic_v1_tuning_baselines", {"router": "baseline"}, raising=False)
|
|
|
|
cleanup_router_config_variables()
|
|
|
|
observed = {
|
|
"master_key": ps.master_key,
|
|
"user_config_file_path": ps.user_config_file_path,
|
|
"user_custom_auth": ps.user_custom_auth,
|
|
"health_check_interval": ps.health_check_interval,
|
|
"prisma_client": ps.prisma_client,
|
|
"heuristic_v1_tuning_baselines": ps.heuristic_v1_tuning_baselines,
|
|
}
|
|
assert normalize(observed) == {
|
|
"master_key": None,
|
|
"user_config_file_path": None,
|
|
"user_custom_auth": None,
|
|
"health_check_interval": None,
|
|
"prisma_client": None,
|
|
"heuristic_v1_tuning_baselines": None,
|
|
}
|
|
|
|
|
|
def test_cleanup_router_config_variables_fails_on_unknown_attr_raises():
|
|
"""The function only writes documented globals — accessing a non-existent
|
|
one after cleanup should still raise AttributeError."""
|
|
cleanup_router_config_variables()
|
|
with pytest.raises(AttributeError):
|
|
_ = ps.this_attribute_should_not_exist_xyz
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# proxy_shutdown_event
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
|
|
fake_prisma = MagicMock()
|
|
fake_prisma.disconnect = AsyncMock()
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
monkeypatch.setattr(ps, "master_key", "sk-x", raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
await proxy_shutdown_event()
|
|
|
|
observed = {
|
|
"disconnect_called": fake_prisma.disconnect.await_count == 1,
|
|
"jwt_closed": fake_jwt.close.await_count == 1,
|
|
"master_key_reset": ps.master_key,
|
|
"prisma_reset": ps.prisma_client,
|
|
}
|
|
assert normalize(observed) == {
|
|
"disconnect_called": True,
|
|
"jwt_closed": True,
|
|
"master_key_reset": None,
|
|
"prisma_reset": None,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_flushes_every_langfuse_export_channel(monkeypatch):
|
|
"""A generation finished just before a graceful restart is still queued in its batch
|
|
processor, so shutdown must flush every acquired export channel."""
|
|
from litellm.integrations.langfuse import langfuse_sdk
|
|
|
|
flushed = MagicMock(return_value=True)
|
|
monkeypatch.setattr(langfuse_sdk, "flush_langfuse_tracing", flushed)
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
monkeypatch.setattr(ps, "jwt_handler", MagicMock(close=AsyncMock()), raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
await proxy_shutdown_event()
|
|
|
|
assert flushed.call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_flushes_langfuse_off_the_event_loop_and_logs_a_timeout(monkeypatch, caplog):
|
|
"""The flush blocks on OTLP exports for up to its deadline, so it must run on a worker thread
|
|
with the shutdown deadline, and a channel that misses it is reported instead of ignored."""
|
|
import threading
|
|
|
|
from litellm.constants import LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS
|
|
from litellm.integrations.langfuse import langfuse_sdk
|
|
|
|
ran_on = MagicMock()
|
|
|
|
def flushed(timeout_millis: int) -> bool:
|
|
ran_on(threading.current_thread(), timeout_millis)
|
|
return False
|
|
|
|
monkeypatch.setattr(langfuse_sdk, "flush_langfuse_tracing", flushed)
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
monkeypatch.setattr(ps, "jwt_handler", MagicMock(close=AsyncMock()), raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
with caplog.at_level("WARNING", logger="LiteLLM Proxy"):
|
|
await proxy_shutdown_event()
|
|
|
|
(flush_thread, timeout_millis), _ = ran_on.call_args
|
|
assert flush_thread is not threading.main_thread()
|
|
assert timeout_millis == LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS
|
|
assert any("Langfuse shutdown flush incomplete" in record.getMessage() for record in caplog.records)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monkeypatch):
|
|
"""
|
|
The gateway request fold lives in memory, so shutdown drains it to the database.
|
|
|
|
That drain has to happen while prisma is still connected: a write attempted
|
|
after ``disconnect()`` raises ClientNotConnectedError, the flush swallows it
|
|
and merges the counts back onto an accumulator the process is about to
|
|
discard, and the final interval is lost silently on every restart. Ordering is
|
|
the whole behavior here, so assert the order rather than that both ran.
|
|
"""
|
|
calls: list = [] # mutable-ok: records call order, which is the assertion
|
|
|
|
fake_prisma = MagicMock()
|
|
fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
|
|
async def _record_flush(client, accumulator):
|
|
calls.append("flush")
|
|
assert client is fake_prisma
|
|
|
|
monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
await proxy_shutdown_event()
|
|
|
|
assert calls == ["flush", "disconnect"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_before_disconnecting(monkeypatch):
|
|
"""
|
|
A management call returns before its audit row is written, so shutdown waits
|
|
for that write while prisma is still connected. It waits beside the gateway
|
|
flush rather than before it, so the two waits overlap inside the teardown budget.
|
|
"""
|
|
calls: list[str] = [] # mutable-ok: records call order, which is the assertion
|
|
gateway_flush_started: Final = asyncio.Event()
|
|
|
|
fake_prisma = MagicMock()
|
|
fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
|
|
async def _record_flush(client, accumulator):
|
|
gateway_flush_started.set()
|
|
calls.append("flush")
|
|
|
|
async def _audit_write_that_finishes_after_the_flush_starts() -> None:
|
|
await gateway_flush_started.wait()
|
|
calls.append("audit")
|
|
|
|
monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts()))
|
|
await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
|
|
|
|
assert calls == ["flush", "audit", "disconnect"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch):
|
|
"""No prisma client means nothing to drain to, and no attempt is made."""
|
|
flush = AsyncMock()
|
|
monkeypatch.setattr(ps, "flush_gateway_requests", flush, raising=False)
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
await proxy_shutdown_event()
|
|
|
|
assert flush.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_the_rest_of_teardown(monkeypatch):
|
|
calls: list[str] = [] # mutable-ok: records call order, which is the assertion
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
monkeypatch.setattr(
|
|
ps, "drain_post_call_redis_batches", AsyncMock(side_effect=lambda: calls.append("redis")), raising=False
|
|
)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
async def _audit_callback_still_sending() -> None:
|
|
await asyncio.sleep(0.05)
|
|
calls.append("audit")
|
|
|
|
track_audit_task(asyncio.create_task(_audit_callback_still_sending()))
|
|
await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
|
|
|
|
assert calls == ["audit", "redis"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
|
|
fake_prisma = MagicMock()
|
|
fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone"))
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
with pytest.raises(RuntimeError, match="db gone"):
|
|
await proxy_shutdown_event()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _flush_spend_logs_queue_on_shutdown
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_spend_logs_queue_on_shutdown_drains_before_disconnect(monkeypatch):
|
|
fake_prisma = MagicMock()
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
drain = AsyncMock()
|
|
import litellm.proxy.utils as utils_mod
|
|
|
|
monkeypatch.setattr(utils_mod, "drain_spend_logs_queue", drain)
|
|
|
|
await ps._flush_spend_logs_queue_on_shutdown()
|
|
|
|
observed = {
|
|
"drain_calls": drain.await_count,
|
|
"drain_prisma": drain.await_args.kwargs["prisma_client"] is fake_prisma,
|
|
}
|
|
assert observed == {
|
|
"drain_calls": 1,
|
|
"drain_prisma": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_spend_logs_queue_on_shutdown_swallows_drain_errors(monkeypatch):
|
|
monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm.proxy.utils as utils_mod
|
|
|
|
monkeypatch.setattr(
|
|
utils_mod,
|
|
"drain_spend_logs_queue",
|
|
AsyncMock(side_effect=RuntimeError("db gone")),
|
|
)
|
|
|
|
await ps._flush_spend_logs_queue_on_shutdown()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_spend_counters_on_shutdown_commits_buffered_spend(monkeypatch):
|
|
fake_prisma = MagicMock()
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
commit = AsyncMock()
|
|
monkeypatch.setattr(ps.proxy_logging_obj.db_spend_update_writer, "db_update_spend_transaction_handler", commit)
|
|
|
|
await ps.flush_spend_counters_on_shutdown()
|
|
|
|
observed = {
|
|
"commit_calls": commit.await_count,
|
|
"commit_prisma": commit.await_args.kwargs["prisma_client"] is fake_prisma,
|
|
"commit_proxy_logging": commit.await_args.kwargs["proxy_logging_obj"] is ps.proxy_logging_obj,
|
|
}
|
|
assert observed == {"commit_calls": 1, "commit_prisma": True, "commit_proxy_logging": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_spend_counters_on_shutdown_logs_and_swallows_commit_errors(monkeypatch, caplog):
|
|
monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
|
|
monkeypatch.setattr(
|
|
ps.proxy_logging_obj.db_spend_update_writer,
|
|
"db_update_spend_transaction_handler",
|
|
AsyncMock(side_effect=RuntimeError("db gone")),
|
|
)
|
|
|
|
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
|
|
await ps.flush_spend_counters_on_shutdown()
|
|
|
|
assert "Error flushing spend counters on shutdown: db gone" in caplog.text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _initialize_shared_aiohttp_session
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_shared_aiohttp_session_returns_client_session():
|
|
from aiohttp import ClientSession
|
|
|
|
session = await initialize_shared_aiohttp_session()
|
|
try:
|
|
observed = {
|
|
"is_client_session": isinstance(session, ClientSession),
|
|
"is_closed": session.closed,
|
|
"has_connector": session.connector is not None,
|
|
}
|
|
assert normalize(observed) == {
|
|
"is_client_session": True,
|
|
"is_closed": False,
|
|
"has_connector": True,
|
|
}
|
|
finally:
|
|
if session is not None:
|
|
await session.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_shared_aiohttp_session_aiohttp_missing_returns_none_on_failure(
|
|
monkeypatch,
|
|
):
|
|
"""If aiohttp import fails, the function catches and returns None — no raise."""
|
|
import builtins
|
|
|
|
real_import = builtins.__import__
|
|
|
|
def _raise_for_aiohttp(name, *args, **kwargs):
|
|
if name == "aiohttp":
|
|
raise ImportError("simulated missing aiohttp")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", _raise_for_aiohttp)
|
|
result = await initialize_shared_aiohttp_session()
|
|
assert result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# save_worker_config
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_save_worker_config_writes_json_to_environ(monkeypatch):
|
|
monkeypatch.delenv("WORKER_CONFIG", raising=False)
|
|
|
|
save_worker_config(model="gpt-4", config="/tmp/c.yaml", debug=True)
|
|
|
|
payload = json.loads(os.environ["WORKER_CONFIG"])
|
|
assert normalize(payload) == {
|
|
"model": "gpt-4",
|
|
"config": "/tmp/c.yaml",
|
|
"debug": True,
|
|
}
|
|
|
|
|
|
def test_save_worker_config_invalid_no_kwargs_yields_empty(monkeypatch):
|
|
monkeypatch.delenv("WORKER_CONFIG", raising=False)
|
|
|
|
save_worker_config()
|
|
assert os.environ["WORKER_CONFIG"] == "{}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _redact_worker_config_for_logging (LIT-4152)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_LIT4152_SECRETS = (
|
|
"sk-lit4152-regression-master-key-abcdef1234567890",
|
|
"leak_password_9090",
|
|
"sk-lit4152-provider-api-key-abcdef",
|
|
"postgresql://leak_user:leak_password_9090@leak-host.internal:5432/leak_db",
|
|
)
|
|
|
|
|
|
def _lit4152_worker_config_dict():
|
|
return {
|
|
"model": "openai/gpt-4o-mini",
|
|
"config": "/tmp/c.yaml",
|
|
"master_key": _LIT4152_SECRETS[0],
|
|
"database_url": _LIT4152_SECRETS[3],
|
|
"api_key": _LIT4152_SECRETS[2],
|
|
"drop_params": True,
|
|
}
|
|
|
|
|
|
def test__redact_worker_config_for_logging_dict_masks_all_secret_shapes():
|
|
"""LIT-4152 regression: dict-form worker_config must not embed any raw
|
|
secret. Covers the segment-matched fields (`master_key`, `api_key`) and the
|
|
URL-with-credentials field (`database_url`), which the segment masker
|
|
misses because neither segment matches its sensitive-pattern set.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
redacted = _redact_worker_config_for_logging(_lit4152_worker_config_dict())
|
|
rendered = repr(redacted)
|
|
for secret in _LIT4152_SECRETS:
|
|
assert secret not in rendered, f"leak: {secret} in {rendered!r}"
|
|
assert isinstance(redacted, dict)
|
|
assert redacted["model"] == "openai/gpt-4o-mini"
|
|
assert redacted["drop_params"] is True
|
|
|
|
|
|
def test__redact_worker_config_for_logging_json_string_round_trips_masked():
|
|
"""Docker/K8s deployments hand the proxy a JSON string via ``WORKER_CONFIG``.
|
|
Confirm the string path also masks and that the returned value re-parses
|
|
into a dict with the sensitive fields masked.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
payload = json.dumps(_lit4152_worker_config_dict())
|
|
redacted = _redact_worker_config_for_logging(payload)
|
|
assert isinstance(redacted, str)
|
|
for secret in _LIT4152_SECRETS:
|
|
assert secret not in redacted, f"leak: {secret} in {redacted!r}"
|
|
parsed = json.loads(redacted)
|
|
assert parsed["model"] == "openai/gpt-4o-mini"
|
|
|
|
|
|
def test__redact_worker_config_for_logging_passthrough_for_none_and_non_json_string():
|
|
"""Non-dict, non-JSON-parseable string is passed through verbatim (nothing
|
|
to mask) and ``None`` returns ``None``.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
assert _redact_worker_config_for_logging(None) is None
|
|
assert _redact_worker_config_for_logging("/tmp/some_config.yaml") == "/tmp/some_config.yaml"
|
|
|
|
|
|
def test__redact_worker_config_for_logging_masks_non_string_url_webhook_values():
|
|
"""The URL/webhook fields the segment masker cannot catch by key name
|
|
(``alert_to_webhook_url``, ``pass_through_endpoints``,
|
|
``database_extra_connection_params``) can hold non-string shapes:
|
|
``alert_to_webhook_url`` is typed as ``Optional[Dict]`` and can nest
|
|
secret query params under keys the segment masker also misses. Confirm
|
|
the whole value is replaced regardless of shape so a nested webhook or
|
|
Bearer token under a non-segment-matched key does not slip through.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
nested_webhook_secret = "https://hooks.slack.com/services/T0/B0/nested-webhook-secret-xyz"
|
|
data = {
|
|
"master_key": "sk-should-be-masked",
|
|
"alert_to_webhook_url": {"budget_alerts": nested_webhook_secret},
|
|
"pass_through_endpoints": [
|
|
{
|
|
"path": "/upstream",
|
|
"target": "https://api.provider.com",
|
|
"headers": {"Authorization": "Bearer nested-token-should-be-gone"},
|
|
}
|
|
],
|
|
"database_extra_connection_params": {"password": "extra-db-password-abc"},
|
|
}
|
|
redacted = _redact_worker_config_for_logging(data)
|
|
rendered = repr(redacted)
|
|
for secret in (
|
|
"sk-should-be-masked",
|
|
nested_webhook_secret,
|
|
"nested-token-should-be-gone",
|
|
"extra-db-password-abc",
|
|
):
|
|
assert secret not in rendered, f"leak: {secret} in {rendered!r}"
|
|
|
|
|
|
def test__redact_worker_config_for_logging_masks_nested_secret_fields():
|
|
"""LIT-4152 nested regression: the URL/webhook credential fields the segment
|
|
masker cannot catch by name (``database_url``,
|
|
``database_extra_connection_params``, ``pass_through_endpoints``,
|
|
``alert_to_webhook_url``) must be redacted at any depth, not just the top
|
|
level. A worker_config that nests ``general_settings`` under a parent key
|
|
must not leak a nested ``database_url`` or webhook secret; the earlier
|
|
top-level-only redaction would have passed these through raw.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
nested_db_url = "postgresql://nested_user:nested_pw_4152@nested-host:5432/db"
|
|
nested_webhook = "https://hooks.slack.com/services/T0/B0/nested-4152-webhook"
|
|
nested_extra_pw = "nested-extra-conn-pw-4152"
|
|
nested_bearer = "Bearer nested-passthrough-token-4152"
|
|
data = {
|
|
"config": {
|
|
"general_settings": {
|
|
"database_url": nested_db_url,
|
|
"database_extra_connection_params": {"password": nested_extra_pw},
|
|
"alert_to_webhook_url": {"budget_alerts": nested_webhook},
|
|
"pass_through_endpoints": [{"path": "/up", "headers": {"Authorization": nested_bearer}}],
|
|
}
|
|
}
|
|
}
|
|
redacted = _redact_worker_config_for_logging(data)
|
|
rendered = repr(redacted)
|
|
for secret in (nested_db_url, nested_webhook, nested_extra_pw, nested_bearer):
|
|
assert secret not in rendered, f"nested leak: {secret} in {rendered!r}"
|
|
|
|
inner = redacted["config"]["general_settings"]
|
|
assert inner["database_url"] == "REDACTED"
|
|
assert inner["pass_through_endpoints"] == "REDACTED"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# initialize
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_initialize_signature_is_async_with_expected_params():
|
|
sig = inspect.signature(initialize)
|
|
# Hard-coded so a signature change (param added/removed) trips the gate.
|
|
expected_param_count = 16
|
|
observed = {
|
|
"is_async": inspect.iscoroutinefunction(initialize),
|
|
"param_count": len(sig.parameters),
|
|
"has_model": "model" in sig.parameters,
|
|
"has_config": "config" in sig.parameters,
|
|
}
|
|
assert normalize(observed) == {
|
|
"is_async": True,
|
|
"param_count": expected_param_count,
|
|
"has_model": True,
|
|
"has_config": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_invalid_unexpected_kwarg_raises_type_error():
|
|
with pytest.raises(TypeError):
|
|
await initialize(this_is_not_a_real_kwarg=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_from_worker_config_drops_legacy_telemetry_key():
|
|
with pytest.raises(TypeError):
|
|
await initialize(telemetry=True)
|
|
await initialize_from_worker_config({"telemetry": True, "request_timeout": 77})
|
|
assert ps.user_request_timeout == 77
|
|
with pytest.raises(TypeError):
|
|
await initialize_from_worker_config({"this_is_not_a_real_kwarg": True})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# load_from_azure_key_vault
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_load_from_azure_key_vault_disabled_no_side_effect(monkeypatch):
|
|
import litellm
|
|
|
|
sentinel_secret_mgr = object()
|
|
monkeypatch.setattr(litellm, "secret_manager_client", sentinel_secret_mgr, raising=False)
|
|
|
|
result = load_from_azure_key_vault(use_azure_key_vault=False)
|
|
|
|
observed = {
|
|
"return_value": result,
|
|
"secret_manager_unchanged": litellm.secret_manager_client is sentinel_secret_mgr,
|
|
"called_with": False,
|
|
}
|
|
assert normalize(observed) == {
|
|
"return_value": None,
|
|
"secret_manager_unchanged": True,
|
|
"called_with": False,
|
|
}
|
|
|
|
|
|
def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch):
|
|
"""Enabled but AZURE_KEY_VAULT_URI unset / azure libs likely unavailable —
|
|
function catches Exception and does not raise."""
|
|
monkeypatch.delenv("AZURE_KEY_VAULT_URI", raising=False)
|
|
|
|
result = load_from_azure_key_vault(use_azure_key_vault=True)
|
|
assert result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# cost_tracking
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(monkeypatch):
|
|
import litellm
|
|
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
|
|
|
|
fake_prisma = MagicMock()
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
|
|
monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
|
|
|
|
before_callbacks = len(litellm.callbacks)
|
|
before_async = len(litellm._async_success_callback)
|
|
|
|
cost_tracking()
|
|
cost_tracking()
|
|
|
|
observed = {
|
|
"added_to_callbacks": len(litellm.callbacks) - before_callbacks,
|
|
"added_to_async_success": len(litellm._async_success_callback) - before_async,
|
|
"shadow_eval_loggers": sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks),
|
|
"prisma_was_set": True,
|
|
}
|
|
assert normalize(observed) == {
|
|
"added_to_callbacks": 2,
|
|
"added_to_async_success": 1,
|
|
"shadow_eval_loggers": 1,
|
|
"prisma_was_set": True,
|
|
}
|
|
|
|
|
|
def test_cost_tracking_no_op_when_prisma_missing(monkeypatch):
|
|
"""Without a prisma_client cost_tracking is a no-op — not an error."""
|
|
import litellm
|
|
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
|
|
monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
|
|
|
|
cost_tracking()
|
|
|
|
assert litellm.callbacks == []
|
|
assert litellm._async_success_callback == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _resolve_typed_dict_type
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _SampleTD(TypedDict):
|
|
a: int
|
|
b: str
|
|
|
|
|
|
def test_resolve_typed_dict_type_finds_class_in_optional():
|
|
typ = _SampleTD | None
|
|
result = _resolve_typed_dict_type(typ)
|
|
|
|
observed = {
|
|
"input_repr": "_SampleTD | None",
|
|
"result_is_sample_td": result is _SampleTD,
|
|
"result_is_class": isinstance(result, type),
|
|
}
|
|
assert normalize(observed) == {
|
|
"input_repr": "_SampleTD | None",
|
|
"result_is_sample_td": True,
|
|
"result_is_class": True,
|
|
}
|
|
|
|
|
|
def test_resolve_typed_dict_type_invalid_plain_type_returns_none():
|
|
"""A non-TypedDict, non-Union input returns None — not an error."""
|
|
assert _resolve_typed_dict_type(int) is None
|
|
assert _resolve_typed_dict_type(str) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _resolve_pydantic_type
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _SampleModelA(BaseModel):
|
|
x: int
|
|
|
|
|
|
class _SampleModelB(BaseModel):
|
|
y: str
|
|
|
|
|
|
def test_resolve_pydantic_type_extracts_non_none_args_from_union():
|
|
typ = _SampleModelA | _SampleModelB | None
|
|
result = _resolve_pydantic_type(typ)
|
|
|
|
observed = {
|
|
"result_type": type(result).__name__,
|
|
"result_len": len(result),
|
|
"contains_a": _SampleModelA in result,
|
|
"contains_b": _SampleModelB in result,
|
|
}
|
|
assert normalize(observed) == {
|
|
"result_type": "list",
|
|
"result_len": 2,
|
|
"contains_a": True,
|
|
"contains_b": True,
|
|
}
|
|
|
|
|
|
def test_resolve_pydantic_type_invalid_non_union_non_model_returns_empty():
|
|
"""When given a non-Union and non-BaseModel input the function returns [].
|
|
|
|
This is the silent-empty fallback path — error-ish by behavior."""
|
|
result = _resolve_pydantic_type(int)
|
|
assert result == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# get_litellm_model_info
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_get_litellm_model_info_uses_base_model_for_lookup(monkeypatch):
|
|
import litellm
|
|
|
|
expected_info = {"max_tokens": 8192, "input_cost_per_token": 0.00003}
|
|
fake_get = MagicMock(return_value=expected_info)
|
|
monkeypatch.setattr(litellm, "get_model_info", fake_get, raising=False)
|
|
|
|
model = {
|
|
"model_info": {"base_model": "gpt-4"},
|
|
"litellm_params": {"model": "azure/my-deployment"},
|
|
}
|
|
result = get_litellm_model_info(model=model)
|
|
|
|
observed = {
|
|
"called_arg": (
|
|
fake_get.call_args.args[0] if fake_get.call_args.args else fake_get.call_args.kwargs.get("model")
|
|
),
|
|
"returned_max_tokens": result.get("max_tokens"),
|
|
"returned_cost": result.get("input_cost_per_token"),
|
|
}
|
|
assert normalize(observed) == {
|
|
"called_arg": "gpt-4",
|
|
"returned_max_tokens": 8192,
|
|
"returned_cost": 0.00003,
|
|
}
|
|
|
|
|
|
def test_get_litellm_model_info_invalid_empty_dict_returns_empty():
|
|
"""Empty input means model_to_lookup is None — internal exception is caught
|
|
and the function returns {}."""
|
|
result = get_litellm_model_info(model={})
|
|
assert result == {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# run_ollama_serve
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_run_ollama_serve_invokes_subprocess_popen(monkeypatch):
|
|
fake_popen = MagicMock()
|
|
monkeypatch.setattr(ps.subprocess, "Popen", fake_popen)
|
|
|
|
run_ollama_serve()
|
|
|
|
args, kwargs = fake_popen.call_args
|
|
observed = {
|
|
"popen_called": fake_popen.call_count == 1,
|
|
"command": args[0] if args else kwargs.get("args"),
|
|
"has_stdout_kw": "stdout" in kwargs,
|
|
"has_stderr_kw": "stderr" in kwargs,
|
|
}
|
|
assert normalize(observed) == {
|
|
"popen_called": True,
|
|
"command": ["ollama", "serve"],
|
|
"has_stdout_kw": True,
|
|
"has_stderr_kw": True,
|
|
}
|
|
|
|
|
|
def test_run_ollama_serve_popen_failure_is_swallowed(monkeypatch):
|
|
"""Popen raising OSError must NOT propagate — function logs and returns."""
|
|
monkeypatch.setattr(ps.subprocess, "Popen", MagicMock(side_effect=OSError("no ollama binary")))
|
|
|
|
result = run_ollama_serve()
|
|
assert result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# proxy_startup_event
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_startup_event_is_async_context_manager_with_expected_signature():
|
|
"""proxy_startup_event is the FastAPI lifespan. Verify its surface without
|
|
actually running the heavy init path (DB, Router, OTEL, etc.)."""
|
|
sig = inspect.signature(proxy_startup_event)
|
|
wrapped = getattr(proxy_startup_event, "__wrapped__", None)
|
|
observed = {
|
|
"param_count": len(sig.parameters),
|
|
"has_app_param": "app" in sig.parameters,
|
|
"wrapped_is_async": inspect.iscoroutinefunction(wrapped) or inspect.isasyncgenfunction(wrapped),
|
|
"has_asynccontextmanager_wrapper": wrapped is not None,
|
|
}
|
|
assert normalize(observed) == {
|
|
"param_count": 1,
|
|
"has_app_param": True,
|
|
"wrapped_is_async": True,
|
|
"has_asynccontextmanager_wrapper": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_startup_event_invalid_missing_app_arg_raises():
|
|
"""Calling the lifespan with no FastAPI app argument must fail."""
|
|
with pytest.raises(TypeError):
|
|
# Intentionally invoke the underlying async generator function with
|
|
# no arguments — the decorator preserves the missing-arg TypeError.
|
|
async with proxy_startup_event(): # type: ignore[call-arg]
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_startup_event_prunes_dead_workers_live_gauges(tmp_path):
|
|
"""With PROMETHEUS_MULTIPROC_DIR set, a booting worker drops the live-gauge files of pids that no longer
|
|
exist, so a crashed worker's in-flight samples leave the aggregate as soon as its replacement starts."""
|
|
exited = subprocess.Popen(["true"])
|
|
assert exited.wait(timeout=30) == 0
|
|
stale = tmp_path / f"gauge_livesum_{exited.pid}.db"
|
|
stale.touch()
|
|
counter = tmp_path / f"counter_{exited.pid}.db"
|
|
counter.touch()
|
|
|
|
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
|
|
clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path)
|
|
with patch.dict(os.environ, clean_env, clear=True):
|
|
try:
|
|
async with proxy_startup_event(app=None):
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
assert not stale.exists()
|
|
assert counter.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(("disable_model_info_refresh", "job_scheduled"), [(True, False), (False, True)])
|
|
async def test_proxy_startup_event_honors_disable_model_info_refresh(
|
|
disable_model_info_refresh: bool, job_scheduled: bool
|
|
) -> None:
|
|
"""``general_settings.disable_model_info_refresh: true`` keeps the proxy from polling every
|
|
OpenAI-compatible deployment's ``/v1/models`` in the background, so a proxy fronting a replay
|
|
fixture (or a metered upstream) makes only the calls its clients asked for."""
|
|
scheduler = AsyncIOScheduler()
|
|
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} | {
|
|
"LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true"
|
|
}
|
|
with (
|
|
patch.dict(os.environ, clean_env, clear=True),
|
|
patch.object(ps, "scheduler", scheduler),
|
|
patch.dict(ps.general_settings, {"disable_model_info_refresh": disable_model_info_refresh}),
|
|
):
|
|
try:
|
|
async with proxy_startup_event(app=None):
|
|
job = scheduler.get_job("refresh_model_info")
|
|
finally:
|
|
if scheduler.running:
|
|
scheduler.shutdown(wait=False)
|
|
|
|
assert (job is not None) is job_scheduled, (
|
|
f"disable_model_info_refresh={disable_model_info_refresh} but refresh_model_info job is {job}"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_commits_spend_before_waiting_on_a_slow_logging_task_and_again_after_it() -> None:
|
|
"""A worker killed while shutdown waits on a slow logging task keeps the spend it had already counted,
|
|
and a worker that finishes the wait also commits the spend that task adds."""
|
|
from litellm.responses import streaming_iterator as streaming_module
|
|
|
|
task_may_finish: Final = asyncio.Event()
|
|
task_finished: Final = asyncio.Event()
|
|
flushes: Final[list[tuple[str, bool]]] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion
|
|
|
|
async def slow_logging_task() -> None:
|
|
await task_may_finish.wait()
|
|
task_finished.set()
|
|
|
|
def recording_flush(name: str) -> Callable[[], Awaitable[None]]:
|
|
async def flush() -> None:
|
|
flushes.append((name, task_finished.is_set()))
|
|
task_may_finish.set()
|
|
|
|
return flush
|
|
|
|
scheduler: Final = AsyncIOScheduler()
|
|
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} | {
|
|
"LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true"
|
|
}
|
|
with (
|
|
patch.dict(os.environ, clean_env, clear=True),
|
|
patch.object(ps, "scheduler", scheduler),
|
|
patch.object(ps, "flush_spend_counters_on_shutdown", recording_flush("spend counters")),
|
|
patch.object(ps, "_flush_spend_logs_queue_on_shutdown", recording_flush("spend logs")),
|
|
):
|
|
try:
|
|
async with proxy_startup_event(app=None):
|
|
streaming_module.spawn_logging_task(
|
|
asyncio.get_running_loop(), slow_logging_task(), task_name="slow spend row"
|
|
)
|
|
finally:
|
|
if scheduler.running:
|
|
scheduler.shutdown(wait=False)
|
|
|
|
assert flushes == [
|
|
("spend counters", False),
|
|
("spend logs", False),
|
|
("spend counters", True),
|
|
("spend logs", True),
|
|
]
|
|
|
|
|
|
def test_otel_global_provider_published_after_callback_init():
|
|
"""The OTel V2 global-provider publish must run after callback
|
|
initialization in ``proxy_startup_event``.
|
|
|
|
Regression for the orphan span: a preset (arize, langfuse, …) builds its
|
|
single folded logger during ``_initialize_startup_logging``. Publishing the
|
|
global ``TracerProvider`` before that ran found no logger and built a second
|
|
generic one whose provider became the global, so the FastAPI server span and
|
|
the preset's gen-ai spans exported through different providers and the LLM
|
|
span was orphaned. The publish (``publish_global_otel_v2_provider``) must
|
|
therefore appear after ``_initialize_startup_logging`` in the lifespan source.
|
|
"""
|
|
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
|
|
source = inspect.getsource(wrapped)
|
|
init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(")
|
|
publish_pos = source.find("publish_global_otel_v2_provider(")
|
|
assert init_pos != -1, "callback init call not found in proxy_startup_event"
|
|
assert publish_pos != -1, "OTEL global publish not found in proxy_startup_event"
|
|
assert init_pos < publish_pos, (
|
|
"OTEL global provider is published before callbacks are initialized; a "
|
|
"preset logger will not exist yet and a second generic logger will own "
|
|
"the global provider, orphaning gen-ai spans"
|
|
)
|
|
|
|
|
|
def test_startup_warns_for_global_budget_without_database(caplog):
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
ProxyStartupEvent.warn_budget_without_db(max_budget=100.0, prisma_client=None)
|
|
|
|
assert "litellm.max_budget=100.0" in caplog.text
|
|
assert "will NOT be enforced" in caplog.text
|
|
assert "requests will never be blocked" in caplog.text
|
|
|
|
|
|
def test_startup_does_not_warn_for_global_budget_with_database(caplog):
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
ProxyStartupEvent.warn_budget_without_db(max_budget=100.0, prisma_client=MagicMock())
|
|
|
|
assert "litellm.max_budget" not in caplog.text
|
|
|
|
|
|
@pytest.mark.parametrize("max_budget", [0, None])
|
|
def test_startup_does_not_warn_without_global_budget(caplog, max_budget):
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
ProxyStartupEvent.warn_budget_without_db(max_budget=max_budget, prisma_client=None)
|
|
|
|
assert "litellm.max_budget" not in caplog.text
|
|
|
|
|
|
def test_startup_warns_for_fail_closed_rate_limits_without_redis(caplog):
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
ProxyStartupEvent._warn_fail_closed_rate_limits_without_redis(
|
|
fail_closed_rate_limit_enforcement=True, redis_usage_cache=None
|
|
)
|
|
|
|
assert "fail_closed_rate_limit_enforcement" in caplog.text
|
|
assert "rejects nothing" in caplog.text
|
|
|
|
|
|
@pytest.mark.parametrize("fail_closed, redis_usage_cache", [(True, MagicMock()), (False, None)])
|
|
def test_startup_does_not_warn_for_fail_closed_rate_limits_when_nothing_is_lost(caplog, fail_closed, redis_usage_cache):
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
ProxyStartupEvent._warn_fail_closed_rate_limits_without_redis(
|
|
fail_closed_rate_limit_enforcement=fail_closed, redis_usage_cache=redis_usage_cache
|
|
)
|
|
|
|
assert "fail_closed_rate_limit_enforcement" not in caplog.text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_startup_event_warns_for_fail_closed_rate_limits_without_redis(caplog):
|
|
scheduler = AsyncIOScheduler()
|
|
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} | {
|
|
"LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true"
|
|
}
|
|
with (
|
|
patch.dict(os.environ, clean_env, clear=True),
|
|
patch.object(ps, "scheduler", scheduler),
|
|
patch.dict(ps.general_settings, {"fail_closed_rate_limit_enforcement": True}),
|
|
caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"),
|
|
):
|
|
try:
|
|
async with proxy_startup_event(app=None):
|
|
pass
|
|
finally:
|
|
if scheduler.running:
|
|
scheduler.shutdown(wait=False)
|
|
|
|
assert "fail_closed_rate_limit_enforcement" in caplog.text
|
|
assert "rejects nothing" in caplog.text
|
|
|
|
|
|
def test_proxy_startup_event_warns_for_global_budget_without_database():
|
|
"""Pin the lifespan call that prevents silent DB-less budgets.
|
|
|
|
The call must follow Prisma setup so DB-backed deployments do not false-positive.
|
|
Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
|
|
"""
|
|
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
|
|
source = inspect.getsource(wrapped)
|
|
budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
|
|
warn_pos = source.find("warn_budget_without_db(")
|
|
next_startup_section_pos = source.find(
|
|
"await ProxyStartupEvent.initialize_scheduled_background_jobs(",
|
|
budget_check_pos,
|
|
)
|
|
|
|
assert budget_check_pos != -1, "global budget startup block not found"
|
|
assert warn_pos != -1, "DB-less budget warning call not found"
|
|
assert next_startup_section_pos != -1, "startup section after budget block not found"
|
|
assert budget_check_pos < warn_pos < next_startup_section_pos, (
|
|
"DB-less budget warning must run after Prisma setup and the DB-backed budget block"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row():
|
|
from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT
|
|
|
|
prisma_client = MagicMock()
|
|
prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
|
prisma_client.db.litellm_config.create = AsyncMock()
|
|
deployment = {
|
|
"model_name": "a",
|
|
"litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": {}},
|
|
}
|
|
|
|
result = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, [deployment])
|
|
|
|
assert result == {'yaml:["a",[]]': DEFAULT_TUNING_FINGERPRINT}
|
|
assert prisma_client.db.litellm_config.create.await_args.kwargs["data"] == {
|
|
"param_name": "auto_router_tuning_baseline_v3",
|
|
"param_value": json.dumps(dict(result)),
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scorer_baseline_upgrade_preserves_existing_routers_and_is_not_refreshed_on_restart():
|
|
from litellm.router_utils.auto_router_tuning_baseline import mutable_tuned_identities, snapshot_tuning_baselines
|
|
|
|
deployments = [
|
|
{
|
|
"model_name": name,
|
|
"litellm_params": {
|
|
"model": "auto_router/complexity_router",
|
|
"complexity_router_config": {"tiers": {"SIMPLE": name}, "code_keywords": [name]},
|
|
},
|
|
}
|
|
for name in ("a", "b")
|
|
]
|
|
prisma_client = MagicMock()
|
|
prisma_client.db.litellm_config.find_unique = AsyncMock(
|
|
side_effect=lambda where: (
|
|
MagicMock(param_value='{"legacy-router":"old-combined-hash"}')
|
|
if where["param_name"] == "auto_router_tuning_baseline_v2"
|
|
else None
|
|
)
|
|
)
|
|
prisma_client.db.litellm_config.create = AsyncMock()
|
|
|
|
baseline = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, deployments)
|
|
|
|
assert baseline == snapshot_tuning_baselines(deployments)
|
|
assert mutable_tuned_identities(deployments, baseline) == frozenset()
|
|
prisma_client.db.litellm_config.create.assert_awaited_once_with(
|
|
data={"param_name": "auto_router_tuning_baseline_v3", "param_value": json.dumps(dict(baseline))}
|
|
)
|
|
prisma_client.db.litellm_config.find_unique.side_effect = None
|
|
prisma_client.db.litellm_config.find_unique.return_value = MagicMock(param_value=json.dumps(dict(baseline)))
|
|
changed = [
|
|
{
|
|
"model_name": "a",
|
|
"litellm_params": {
|
|
"model": "auto_router/complexity_router",
|
|
"complexity_router_config": {"tiers": {"SIMPLE": "different-model"}, "code_keywords": ["new-rule"]},
|
|
},
|
|
}
|
|
]
|
|
|
|
reloaded = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, changed)
|
|
|
|
assert reloaded == baseline
|
|
assert mutable_tuned_identities(changed, reloaded) == frozenset({'yaml:["a",[]]'})
|
|
prisma_client.db.litellm_config.create.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tuning_baseline_waits_for_a_complete_db_model_census(monkeypatch):
|
|
prisma_client = MagicMock()
|
|
monkeypatch.setattr(ps.proxy_config, "get_models_from_db", AsyncMock(return_value=None))
|
|
|
|
result = await ProxyStartupEvent.enforce_heuristic_v1_tuning_baseline(
|
|
prisma_client=prisma_client,
|
|
llm_router=None,
|
|
limit=1,
|
|
)
|
|
|
|
assert result is None
|
|
prisma_client.db.litellm_config.find_unique.assert_not_called()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _initialize_slack_alerting_jobs — spend-report pod locking (issue #14809)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
SlackAlertingJobs = dict[str, Callable[[], Awaitable[None]]]
|
|
|
|
|
|
def _make_slack_alerting_proxy_logging(acquire_lock_result: bool | None) -> MagicMock:
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.slack_alerting_instance.alerting = ["slack"]
|
|
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report = AsyncMock()
|
|
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report = AsyncMock()
|
|
proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus = AsyncMock()
|
|
pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
|
|
pod_lock_manager.acquire_lock = AsyncMock(return_value=acquire_lock_result)
|
|
pod_lock_manager.release_lock = AsyncMock()
|
|
return proxy_logging_obj
|
|
|
|
|
|
async def _init_slack_alerting_jobs(
|
|
acquire_lock_result: bool | None,
|
|
spend_report_frequency: str = "7d",
|
|
) -> tuple[SlackAlertingJobs, MagicMock]:
|
|
scheduler = MagicMock()
|
|
proxy_logging_obj = _make_slack_alerting_proxy_logging(acquire_lock_result)
|
|
|
|
await ProxyStartupEvent._initialize_slack_alerting_jobs(
|
|
scheduler=scheduler,
|
|
general_settings={"spend_report_frequency": spend_report_frequency},
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
prisma_client=MagicMock(),
|
|
)
|
|
|
|
jobs = {call.kwargs["id"]: call.args[0] for call in scheduler.add_job.call_args_list}
|
|
return jobs, proxy_logging_obj
|
|
|
|
|
|
@pytest.mark.parametrize("spend_report_frequency", ["0d", "-1d", "7h"])
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_slack_alerting_jobs_invalid_frequency_raises(spend_report_frequency: str):
|
|
"""A non-positive window used to become an every-second APScheduler interval, and now also
|
|
computes a negative lock TTL that expires instantly and suppresses the report for good.
|
|
match= is load-bearing: drop the guard and "-1d" still raises, but from duration_in_seconds."""
|
|
with pytest.raises(ValueError, match="positive number of days"):
|
|
await _init_slack_alerting_jobs(
|
|
acquire_lock_result=True,
|
|
spend_report_frequency=spend_report_frequency,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_weekly_spend_report_skipped_when_another_pod_holds_the_lock():
|
|
"""regression: issue #14809 - every pod ran its own weekly spend report job."""
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
|
|
|
|
await jobs["weekly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_not_awaited()
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
|
|
cronjob_id="weekly_spend_report_job",
|
|
ttl=7 * 86400 - 3600,
|
|
allow_reentrant=False,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("acquire_lock_result", [True, None])
|
|
@pytest.mark.asyncio
|
|
async def test_weekly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result):
|
|
"""None means redis isn't configured; a single-pod deploy must still report."""
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
|
|
|
|
await jobs["weekly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("7d")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_weekly_spend_report_lock_ttl_tracks_the_configured_window():
|
|
"""TTL is the window less an hour: long enough that no second pod re-sends inside the
|
|
window, short enough that the lock is gone before the next one opens. A fixed TTL would
|
|
break one end or the other as soon as spend_report_frequency changes."""
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True, spend_report_frequency="1d")
|
|
|
|
await jobs["weekly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
|
|
cronjob_id="weekly_spend_report_job",
|
|
ttl=86400 - 3600,
|
|
allow_reentrant=False,
|
|
)
|
|
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("1d")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_monthly_spend_report_skipped_when_another_pod_holds_the_lock():
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
|
|
|
|
await jobs["monthly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_not_awaited()
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
|
|
cronjob_id="monthly_spend_report_job",
|
|
ttl=3600,
|
|
allow_reentrant=False,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("acquire_lock_result", [True, None])
|
|
@pytest.mark.asyncio
|
|
async def test_monthly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result):
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
|
|
|
|
await jobs["monthly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_awaited_once_with()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_spend_report_locks_are_never_released():
|
|
"""The lock is a per-window marker, not a mutex: releasing it lets the next pod re-send."""
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True)
|
|
|
|
await jobs["weekly_spend_report_job"]()
|
|
await jobs["monthly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch):
|
|
"""The boot-time send goes through the same gate, so a losing pod sends nothing at all:
|
|
startup and the scheduled job both stay at zero."""
|
|
monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid")
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
|
|
send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus
|
|
assert send_fallback_stats.await_count == 0
|
|
|
|
await jobs["prometheus_fallback_stats_job"]()
|
|
|
|
assert send_fallback_stats.await_count == 0
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_with(
|
|
cronjob_id="prometheus_fallback_stats_job",
|
|
ttl=3600,
|
|
allow_reentrant=False,
|
|
)
|
|
assert proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.await_count == 2
|
|
|
|
|
|
@pytest.mark.parametrize("acquire_lock_result", [True, None])
|
|
@pytest.mark.asyncio
|
|
async def test_prometheus_fallback_stats_job_runs_when_the_lock_is_free_or_absent(monkeypatch, acquire_lock_result):
|
|
monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid")
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
|
|
send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus
|
|
assert send_fallback_stats.await_count == 1
|
|
|
|
await jobs["prometheus_fallback_stats_job"]()
|
|
|
|
assert send_fallback_stats.await_count == 2
|