From 8576cf74c2ef4eaa15f48eaaf99447f9d6fb394d Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Mon, 7 Sep 2026 14:09:37 +0200 Subject: [PATCH 01/23] feat(auth): add disable_env_credential_login setting with admin ui warning Env-credential login (UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset) is always live today. This adds a general_settings flag to turn that login path off once real admin accounts exist, and a warning banner shown to any admin while it remains enabled. The banner flag is served through /health/readiness/details and stays quiet when disable_password_login_when_sso_enabled already makes the env path unreachable. --- litellm/proxy/_types.py | 12 ++ litellm/proxy/auth/login_utils.py | 32 +++- .../health_endpoints/_health_endpoints.py | 20 +++ .../proxy/auth/test_login_utils.py | 154 ++++++++++++++++++ .../health_endpoints/test_health_endpoints.py | 27 +++ .../useHealthReadinessDetails.ts | 1 + .../src/app/(dashboard)/layout.test.tsx | 4 + .../src/app/(dashboard)/layout.tsx | 3 + .../EnvCredentialLoginWarningBanner.test.tsx | 76 +++++++++ .../EnvCredentialLoginWarningBanner.tsx | 35 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 11 files changed, 364 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx create mode 100644 ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c28ac8848ba..c30a3c5bea8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2817,6 +2817,18 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "UI username/password login. Default is False." ), ) + disable_env_credential_login: bool | None = Field( + None, + description=( + "If True, disables signing in to the Admin UI with the environment credentials: " + "UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback " + "means env-credential login is always live by default). Database users with passwords " + "are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password " + "before enabling, or nobody can sign in to the UI. A locked-out admin can still " + "administer the proxy over the API with the master key, and can unset this setting " + "and restart the proxy to restore env-credential login. Default is False." + ), + ) disable_budget_reservation: bool | None = Field( None, description=( diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 8d4f6f81363..6ca7ea5074d 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -85,6 +85,29 @@ def get_ui_credentials(master_key: str | None) -> tuple[str, str]: return ui_username, ui_password +def _matches_env_credentials(username: str, password: str, master_key: str | None) -> bool: + ui_username, ui_password = get_ui_credentials(master_key) + return secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( + password.encode("utf-8"), ui_password.encode("utf-8") + ) + + +def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool: + """Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed. + + Two settings can turn it off: `disable_env_credential_login` unconditionally, and + `disable_password_login_when_sso_enabled` as a side effect, since its gate rejects + every username/password login before the env comparison runs. Feeds both the + `authenticate_user` gate and the Admin UI warning banner, so the banner never nags + about a login path that is already unreachable. + """ + if general_settings.get("disable_env_credential_login") is True: + return False + if general_settings.get("disable_password_login_when_sso_enabled") is True and is_sso_provider_fully_configured(): + return False + return True + + class LoginResult: """Result object containing authentication data from login.""" @@ -129,7 +152,8 @@ async def authenticate_user( master_key: Master key for the proxy (required) prisma_client: Prisma database client (optional) general_settings: Proxy general_settings, checked for - `disable_password_login_when_sso_enabled` + `disable_password_login_when_sso_enabled` and + `disable_env_credential_login` Returns: LoginResult: Object containing authentication data @@ -170,8 +194,6 @@ async def authenticate_user( code=500, ) - ui_username, ui_password = get_ui_credentials(master_key) - # Check if we can find the `username` in the db. On the UI, users can enter username=their email _user_row: LiteLLM_UserTable | None = None user_role: ( @@ -197,8 +219,8 @@ async def authenticate_user( - Login with UI_USERNAME and UI_PASSWORD - Login with Invite Link `user_email` and `password` combination """ - if secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( - password.encode("utf-8"), ui_password.encode("utf-8") + if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials( + username, password, master_key ): # Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin user_role = LitellmUserRoles.PROXY_ADMIN diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1785a2f0992..e5e4f89234b 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1582,6 +1582,23 @@ async def _show_no_redis_warning() -> bool: return await count_live_proxy_workers(prisma_client) != 1 +def _show_env_credential_login_warning() -> bool: + """ + Whether the UI should warn admins that env-credential login is still enabled. + + UI_USERNAME/UI_PASSWORD (or the master key, when UI_PASSWORD is unset) grant + proxy-admin access with a shared static secret: no per-person identity, no + audit trail, no password policy, and it stays valid until the env var or + master key rotates. That is fine for first-time setup, so it is on by + default, but once real admin accounts exist it should be turned off with + `general_settings.disable_env_credential_login`. + """ + from litellm.proxy.auth.login_utils import is_env_credential_login_enabled + from litellm.proxy.proxy_server import general_settings + + return is_env_credential_login_enabled(general_settings) + + async def _get_health_readiness_details( response: Response | None = None, ) -> dict[str, Any]: @@ -1623,6 +1640,7 @@ async def _get_health_readiness_details( log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel()) is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG) show_no_redis_warning: Final = await _show_no_redis_warning() + show_env_credential_login_warning: Final = _show_env_credential_login_warning() # check DB if prisma_client is not None: # if db passed in, check if it's connected @@ -1650,6 +1668,7 @@ async def _get_health_readiness_details( "log_level": log_level_name, "is_detailed_debug": is_detailed_debug, "show_no_redis_warning": show_no_redis_warning, + "show_env_credential_login_warning": show_env_credential_login_warning, } else: return { @@ -1662,6 +1681,7 @@ async def _get_health_readiness_details( "log_level": log_level_name, "is_detailed_debug": is_detailed_debug, "show_no_redis_warning": show_no_redis_warning, + "show_env_credential_login_warning": show_env_credential_login_warning, } except Exception as e: raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e})") diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 8d93d801bfd..72cc6e04a12 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -23,6 +23,7 @@ from litellm.proxy.auth.login_utils import ( LoginResult, authenticate_user, get_ui_credentials, + is_env_credential_login_enabled, ) @@ -799,3 +800,156 @@ class TestDisablePasswordLoginWhenSSOEnabled: assert isinstance(result, LoginResult) assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + +class TestDisableEnvCredentialLogin: + """`disable_env_credential_login` must reject a login with the env + credentials (UI_USERNAME/UI_PASSWORD, or the master-key fallback when + UI_PASSWORD is unset) while leaving database-user password logins + untouched, so admins with real accounts keep a way in.""" + + @pytest.mark.asyncio + async def test_rejects_correct_env_credentials_when_disabled(self): + master_key = "sk-1234" + ui_username = "admin" + ui_password = "env-only-password" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=ui_username, + password=ui_password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "401" + + @pytest.mark.asyncio + async def test_rejects_master_key_fallback_when_disabled(self): + """With UI_PASSWORD unset, the master key IS the env password, so the + setting must reject it too or it protects nothing by default.""" + master_key = "sk-1234" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": "admin"}, clear=True): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username="admin", + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert exc_info.value.code == "401" + + @pytest.mark.asyncio + async def test_db_user_login_still_works_when_disabled(self): + master_key = "sk-1234" + user_email = "admin@example.com" + password = "Str0ng!Passw0rd" + + mock_user = LiteLLM_UserTable( + user_id="db-admin-1", + user_email=user_email, + password=hash_token(token=password), + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user) + + with patch.dict( + os.environ, + { + "UI_USERNAME": "admin", + "UI_PASSWORD": "env-password", + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + stack.enter_context( + patch( # test-quality-ok: internal orchestration, no HTTP boundary; matches pre-existing tests + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "db-user-token"}, + ) + ) + result = await authenticate_user( + username=user_email, + password=password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == "db-admin-1" + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + @pytest.mark.asyncio + async def test_env_login_still_works_when_setting_absent(self): + """Env-credential login is the bootstrap path on a fresh install and + must stay on by default.""" + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": master_key, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + _patch_successful_admin_login_deps(stack) + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + +class TestIsEnvCredentialLoginEnabled: + """Drives the Admin UI warning banner: it must be True exactly when a + login with the env credentials could actually succeed.""" + + def test_enabled_by_default(self): + assert is_env_credential_login_enabled({}) is True + + def test_disabled_by_dedicated_setting(self): + assert is_env_credential_login_enabled({"disable_env_credential_login": True}) is False + + def test_explicit_false_keeps_it_enabled(self): + assert is_env_credential_login_enabled({"disable_env_credential_login": False}) is True + + def test_disabled_when_sso_gate_blocks_all_password_logins(self): + """`disable_password_login_when_sso_enabled` with SSO configured + rejects every username/password login before the env comparison runs, + so the banner must not nag about an already-unreachable path.""" + with ExitStack() as stack: + _patch_sso_configured(stack, configured=True) + assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is False + + def test_enabled_when_sso_gate_is_set_but_sso_not_configured(self): + with ExitStack() as stack: + _patch_sso_configured(stack, configured=False) + assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is True diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e9e58347337..b06d87ac67e 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1301,6 +1301,33 @@ def test_health_readiness_details_returns_diagnostic_fields(monkeypatch): assert "cache" in response_data +@pytest.mark.parametrize( + "general_settings, expected_warning", + [ + ({}, True), + ({"disable_env_credential_login": True}, False), + ], +) +def test_health_readiness_details_reports_env_credential_login_warning(monkeypatch, general_settings, expected_warning): + """ + The Admin UI banner is driven by this flag: it must be True while + env-credential login is possible and False once + `disable_env_credential_login` turns that login path off. + """ + app = FastAPI() + app.include_router(_health_endpoints_module.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + client = TestClient(app) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + response = client.get("/health/readiness/details") + + assert response.status_code == 200, response.text + assert response.json()["show_env_credential_login_warning"] is expected_warning + + def test_health_readiness_allows_explicit_legacy_public_details(monkeypatch): """ Operators can explicitly preserve the legacy public readiness payload. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts index 307fa9e1691..44d9092df34 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts @@ -14,6 +14,7 @@ export interface HealthReadinessDetailsResponse { log_level?: string; is_detailed_debug?: boolean; show_no_redis_warning?: boolean; + show_env_credential_login_warning?: boolean; } const fetchHealthReadinessDetails = async (accessToken: string): Promise => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 7f1f4cc4bd5..3fe34610260 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -29,6 +29,10 @@ vi.mock("@/components/NoRedisWarningBanner", () => ({ NoRedisWarningBanner: () => null, })); +vi.mock("@/components/EnvCredentialLoginWarningBanner", () => ({ + EnvCredentialLoginWarningBanner: () => null, +})); + vi.mock("@/components/LicenseExpiryBanner", () => ({ LicenseExpiryBanner: () => null, })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 98f2a36d6f3..fa6df7f176a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -10,6 +10,7 @@ import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import { useRouter, useSearchParams } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner"; +import { EnvCredentialLoginWarningBanner } from "@/components/EnvCredentialLoginWarningBanner"; import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner"; import { UserBanner } from "@/components/UserBanner"; import { uiHref } from "@/utils/uiHref"; @@ -113,6 +114,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
@@ -132,6 +134,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
{children}
diff --git a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx new file mode 100644 index 00000000000..535694f14da --- /dev/null +++ b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx @@ -0,0 +1,76 @@ +import { renderWithProviders, screen } from "../../tests/test-utils"; +import { vi } from "vitest"; +import { EnvCredentialLoginWarningBanner } from "./EnvCredentialLoginWarningBanner"; +import type { HealthReadinessDetailsResponse } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import type { UseQueryResult } from "@tanstack/react-query"; + +vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails", () => ({ + useHealthReadinessDetails: vi.fn(), +})); +vi.mock("@/contexts/AuthContext", () => ({ + useAuth: vi.fn(), +})); + +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { useAuth } from "@/contexts/AuthContext"; + +const mockDetails = (data: Partial | undefined) => { + vi.mocked(useHealthReadinessDetails).mockReturnValue({ data } as UseQueryResult); +}; + +const mockRole = (userRole: string) => { + vi.mocked(useAuth).mockReturnValue({ userRole } as ReturnType); +}; + +describe("EnvCredentialLoginWarningBanner", () => { + it("should warn an admin when env-credential login is enabled", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(screen.getByText("Environment-credential login is enabled")).toBeInTheDocument(); + }); + + it("should tell the admin to create a regular admin account before disabling", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByText(/First create a regular admin account/i)).toBeInTheDocument(); + expect(screen.getByText("general_settings.disable_env_credential_login: true")).toBeInTheDocument(); + }); + + it("should warn an admin viewer too", () => { + mockRole("Admin Viewer"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + }); + + it("should render nothing for a non-admin even when the proxy reports the warning", () => { + mockRole("Internal User"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should render nothing when env-credential login is disabled", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: false }); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should render nothing when readiness details are unavailable", () => { + mockRole("Admin"); + mockDetails(undefined); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should pass the access token to the readiness hook", () => { + mockRole("Admin"); + mockDetails(undefined); + renderWithProviders(); + expect(useHealthReadinessDetails).toHaveBeenCalledWith("my-token"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx new file mode 100644 index 00000000000..3a9d50011f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx @@ -0,0 +1,35 @@ +"use client"; + +import React from "react"; +import { TriangleAlert } from "lucide-react"; +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { useAuth } from "@/contexts/AuthContext"; +import { isAdminRole } from "@/utils/roles"; + +export const EnvCredentialLoginWarningBanner: React.FC<{ accessToken: string | null }> = ({ accessToken }) => { + const { userRole } = useAuth(); + const { data: healthData } = useHealthReadinessDetails(accessToken); + + if (!isAdminRole(userRole) || !healthData?.show_env_credential_login_warning) { + return null; + } + + return ( +
+
+ ); +}; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1534c19670..dae0af3d3b1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25778,6 +25778,11 @@ export interface components { * @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). A proxy-level WARNING is logged on every request while this flag is active as a reminder that hard enforcement is relaxed. */ disable_budget_reservation?: boolean | null; + /** + * Disable Env Credential Login + * @description If True, disables signing in to the Admin UI with the environment credentials: UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback means env-credential login is always live by default). Database users with passwords are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password before enabling, or nobody can sign in to the UI. A locked-out admin can still administer the proxy over the API with the master key, and can unset this setting and restart the proxy to restore env-credential login. Default is False. + */ + disable_env_credential_login?: boolean | null; /** * Disable Password Login When Sso Enabled * @description If True and SSO is configured (MICROSOFT_CLIENT_ID, GOOGLE_CLIENT_ID, GENERIC_CLIENT_ID, or SAML_IDP_METADATA_URL/XML), disables username/password login on /login, /v2/login, and /v3/login so SSO is the only way to reach the Admin UI. An admin locked out of the UI can still administer the proxy over the API with the master key; unset this setting and restart the proxy to restore UI username/password login. Default is False. From 78893723f7b65cb73d2c5285cb412b25764a3857 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Tue, 8 Sep 2026 13:05:38 +0200 Subject: [PATCH 02/23] fix(auth): drop env-credential hint from 401 when env login is disabled --- litellm/proxy/auth/login_utils.py | 7 ++++++- tests/test_litellm/proxy/auth/test_login_utils.py | 3 +++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 6ca7ea5074d..c0a76a4fc20 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -362,8 +362,13 @@ async def authenticate_user( code=401, ) else: + env_credentials_hint: Final = ( + "\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file" + if is_env_credential_login_enabled(general_settings) + else "" + ) raise ProxyException( - message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file", + message=f"Invalid credentials used to access UI.{env_credentials_hint}", type=ProxyErrorTypes.auth_error, param="invalid_credentials", code=401, diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 72cc6e04a12..e209a491b0a 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -186,6 +186,7 @@ async def test_authenticate_user_invalid_credentials(): assert exc_info.value.type == ProxyErrorTypes.auth_error assert exc_info.value.code == "401" assert "Invalid credentials" in exc_info.value.message + assert "UI_USERNAME" in exc_info.value.message @pytest.mark.asyncio @@ -829,6 +830,8 @@ class TestDisableEnvCredentialLogin: assert exc_info.value.type == ProxyErrorTypes.auth_error assert exc_info.value.code == "401" + assert "UI_USERNAME" not in exc_info.value.message + assert "UI_PASSWORD" not in exc_info.value.message @pytest.mark.asyncio async def test_rejects_master_key_fallback_when_disabled(self): From 992ee37258e7f3c9ed13c374610899388bc8d446 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Tue, 8 Sep 2026 13:21:45 +0200 Subject: [PATCH 03/23] Update litellm/proxy/health_endpoints/_health_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/health_endpoints/_health_endpoints.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index e5e4f89234b..747cfa526d9 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1583,16 +1583,6 @@ async def _show_no_redis_warning() -> bool: def _show_env_credential_login_warning() -> bool: - """ - Whether the UI should warn admins that env-credential login is still enabled. - - UI_USERNAME/UI_PASSWORD (or the master key, when UI_PASSWORD is unset) grant - proxy-admin access with a shared static secret: no per-person identity, no - audit trail, no password policy, and it stays valid until the env var or - master key rotates. That is fine for first-time setup, so it is on by - default, but once real admin accounts exist it should be turned off with - `general_settings.disable_env_credential_login`. - """ from litellm.proxy.auth.login_utils import is_env_credential_login_enabled from litellm.proxy.proxy_server import general_settings From 3caa3b60d5a319efa78871198081783dd4c70ad1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:58:18 -0700 Subject: [PATCH 04/23] feat(guardrails): run post_call policy pipelines on background Responses retrieval A POST /v1/responses with background: true returns a queued response, so the post_call pipelines attached at submit time had nothing to inspect. They now defer on queued and in_progress responses and run on GET /v1/responses/{id} instead: the retrieval resolves the response id back to its deployment, re-attaches the policies that governed the original model, and reports them in the x-litellm-applied-* headers of the retrieval response. --- litellm/proxy/common_request_processing.py | 8 +- .../proxy/policy_engine/response_retrieval.py | 116 +++++++++++++ litellm/proxy/utils.py | 43 +++-- .../policy_engine/test_response_retrieval.py | 157 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 109 ++++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 95 +++++++---- 6 files changed, 485 insertions(+), 43 deletions(-) create mode 100644 litellm/proxy/policy_engine/response_retrieval.py create mode 100644 tests/test_litellm/proxy/policy_engine/test_response_retrieval.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9720e4b1cf8..1c1eda73c0c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -184,6 +184,7 @@ from litellm.proxy.litellm_pre_call_utils import ( refresh_proxy_server_request_body_snapshot, reject_url_valued_destination, ) +from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -1883,7 +1884,6 @@ class ProxyBaseLLMRequestProcessing: data=self.data, user_api_key_dict=user_api_key_dict, ) - # Calculate request queue time after add_litellm_data_to_request # which sets arrival_time in proxy_server_request. Ends at start_time # (not a freshly captured time.time() here) so this window is exactly @@ -2031,6 +2031,12 @@ class ProxyBaseLLMRequestProcessing: data=self.data, call_type=route_type, ) + if route_type == "aget_responses": + attach_post_call_pipelines_to_retrieval( + data=self.data, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py new file mode 100644 index 00000000000..54b9acf6150 --- /dev/null +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -0,0 +1,116 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + add_policy_sources_to_metadata, + add_policy_to_applied_policies_header, +) +from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.proxy.policy_engine import PolicyMatchContext +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router + +PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...] + +_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines) + + +def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | None: + if llm_router is None or not isinstance(response_id, str): + return None + model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) + if model_id is None: + return None + deployment: Final = llm_router.get_deployment(model_id) + return deployment.model_name if deployment is not None else None + + +def _retrieval_context( + data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str +) -> PolicyMatchContext: + team_alias: Final = user_api_key_dict.team_alias + key_alias: Final = user_api_key_dict.key_alias + return PolicyMatchContext( + team_alias=team_alias if isinstance(team_alias, str) else None, + key_alias=key_alias if isinstance(key_alias, str) else None, + model=model_group, + tags=get_tags_from_request_body(data) or None, + ) + + +def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]: + matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context) + if not matches: + return (), MappingProxyType({}) + applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( + policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list + context=context, + ) + post_call_pipelines: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context( + context=context, policy_names=applied_policy_names + ) + if pipeline.mode == "post_call" + ) + return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches}) + + +def attach_post_call_pipelines_to_retrieval( + data: dict, # mutable-ok: the proxy's request-state dict, written in place like every other policy engine hook + user_api_key_dict: "UserAPIKeyAuth", + llm_router: "Router | None", +) -> None: + if not get_policy_registry().is_initialized(): + return + model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router) + if model_group is None: + return + context: Final = _retrieval_context(data, user_api_key_dict, model_group) + post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context) + _, bucket = get_or_create_metadata_bucket(data) + already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ()) + attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached) + added: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in post_call_pipelines + if policy_name not in attached_policy_names + ) + if not added: + return + pipelines: Final = (*already_attached, *added) + bucket["_guardrail_pipelines"] = pipelines + bucket["_pipeline_managed_guardrails"] = frozenset( + step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps + ) + for policy_name, _pipeline in added: + add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name) + for _policy_name, pipeline in added: + for step in pipeline.steps: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail) + add_policy_sources_to_metadata( + request_data=data, + policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict + policy_name: policy_sources[policy_name] for policy_name, _pipeline in added + }, + ) + verbose_proxy_logger.debug( + "Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s", + data.get("response_id"), + model_group, + ", ".join(policy_name for policy_name, _pipeline in added), + ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b36061c6be..3d03d3858ee 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -171,6 +171,7 @@ from litellm.repositories.verification_token_repository import ( ) from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import ( MCPDuringCallResponseObject, MCPPreCallRequestObject, @@ -525,15 +526,21 @@ def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardr ) -def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None: - if data.get("background") is not True: - return +_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress")) + + +def _is_pending_background_response(response: LLMResponseTypes) -> bool: + return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES + + +def _log_deferred_post_call_pipelines(data: Mapping[str, object], response: ResponsesAPIResponse) -> None: policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data)) if not policy_names: return - verbose_proxy_logger.warning( - "Policies with post_call guardrail pipelines do not run on background responses yet; " - "the response is released ungoverned by them: %s", + verbose_proxy_logger.debug( + "Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s", + response.id, + response.status, ", ".join(policy_names), ) @@ -1956,8 +1963,6 @@ class ProxyLogging: ) try: - _warn_background_skips_post_call_pipelines(data) - # Execute guardrail pipelines before the normal callback loop data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below data=data, @@ -2927,6 +2932,24 @@ class ProxyLogging: daemon=True, ).start() + async def _run_post_call_pipelines( + self, + data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + user_api_key_dict: UserAPIKeyAuth, + response: LLMResponseTypes, + ) -> LLMResponseTypes | None: + if _is_pending_background_response(response): + _log_deferred_post_call_pipelines(data, response) + return None + _, pipeline_response = await self._maybe_execute_pipelines( + data=data, + user_api_key_dict=user_api_key_dict, + call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", + event_hook="post_call", + response=response, + ) + return pipeline_response + async def post_call_success_hook( self, data: dict, @@ -2946,11 +2969,9 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks - _, pipeline_response = await self._maybe_execute_pipelines( + pipeline_response: Final = await self._run_post_call_pipelines( data=data, user_api_key_dict=user_api_key_dict, - call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", - event_hook="post_call", response=response, ) if pipeline_response is not None: diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py new file mode 100644 index 00000000000..f7f29d6b46e --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -0,0 +1,157 @@ +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.router import Deployment, LiteLLM_Params + +GOVERNED_MODEL_GROUP = "gpt-5.4-mini" +GOVERNED_MODEL_ID = "deployment-governed" +UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini" +UNGOVERNED_MODEL_ID = "deployment-ungoverned" + + +class FakeRouter: + def __init__(self, deployments: dict[str, Deployment]): + self._deployments = deployments + + def get_deployment(self, model_id: str) -> Deployment | None: + return self._deployments.get(model_id) + + +def _deployment(model_group: str, model_id: str) -> Deployment: + return Deployment( + model_name=model_group, + litellm_params=LiteLLM_Params(model=f"openai/{model_group}"), + model_info={"id": model_id}, + ) + + +def _router() -> FakeRouter: + return FakeRouter( + { + GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID), + UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID), + } + ) + + +def _encoded_response_id(model_id: str) -> str: + return ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=model_id, response_id="resp_upstream" + ) + + +def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, object]: + return { + "guardrails": {"add": [guardrail]}, + "pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": "block"}]}, + } + + +@pytest.fixture +def policy_engine(): + policy_registry = get_policy_registry() + attachment_registry = get_attachment_registry() + policy_registry.load_policies( + { + "response-governance": _pipeline_policy("output-word-filter"), + "input-governance": _pipeline_policy("input-word-filter", mode="pre_call"), + "team-governance": _pipeline_policy("team-word-filter"), + } + ) + attachment_registry.load_attachments( + [ + {"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]}, + {"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]}, + {"policy": "team-governance", "teams": ["governed-team"]}, + ] + ) + yield + policy_registry.clear() + attachment_registry.clear() + + +def _retrieval_data(model_id: str) -> dict: + return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}} + + +def _attached_pipelines(data: dict) -> tuple[tuple[str, str], ...]: + return tuple( + (policy_name, ",".join(step.guardrail for step in pipeline.steps)) + for policy_name, pipeline in data["litellm_metadata"]["_guardrail_pipelines"] + ) + + +def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine): + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert data["litellm_metadata"]["_pipeline_managed_guardrails"] == frozenset({"output-word-filter"}) + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + assert data["litellm_metadata"]["applied_guardrails"] == ["output-word-filter"] + assert data["litellm_metadata"]["policy_sources"] == {"response-governance": "model:gpt-5.4-mini"} + assert "model" not in data + assert "guardrails" not in data["litellm_metadata"] + + +def test_key_and_team_context_also_governs_retrieval(policy_engine): + data = _retrieval_data(UNGOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval( + data=data, user_api_key_dict=UserAPIKeyAuth(team_alias="governed-team"), llm_router=_router() + ) + + assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),) + + +def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine): + data = _retrieval_data(UNGOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(UNGOVERNED_MODEL_ID) + + +def test_already_attached_policy_is_not_attached_twice(policy_engine): + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router() + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + + +@pytest.mark.parametrize( + "response_id", + ["resp_plain_upstream_id", _encoded_response_id("deployment-missing-from-router"), None], +) +def test_unresolvable_response_id_attaches_nothing(policy_engine, response_id): + data = {"response_id": response_id, "litellm_metadata": {}} + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == {"response_id": response_id, "litellm_metadata": {}} + + +def test_without_a_router_attaches_nothing(policy_engine): + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) + + assert data == _retrieval_data(GOVERNED_MODEL_ID) + + +def test_without_policy_engine_attaches_nothing(): + get_policy_registry().clear() + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(GOVERNED_MODEL_ID) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 6acd9d7258e..9283c1ed865 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8259,3 +8259,112 @@ class TestPassthroughHeadersAcceptImmutableMappings: assert merged["content-type"] == "text/event-stream" # the excluded hop-by-hop header is still dropped assert "transfer-encoding" not in merged + + +class TestBackgroundResponseRetrievalGovernance: + """LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines.""" + + GOVERNED_MODEL_GROUP = "gpt-5.4-mini" + GOVERNED_MODEL_ID = "deployment-governed" + + @pytest.fixture + def policy_engine(self): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + get_policy_registry().load_policies( + { + "response-governance": { + "guardrails": {"add": ["output-word-filter"]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": "output-word-filter", "on_pass": "allow", "on_fail": "block"}], + }, + } + } + ) + get_attachment_registry().load_attachments( + [{"policy": "response-governance", "models": [self.GOVERNED_MODEL_GROUP]}] + ) + yield + get_policy_registry().clear() + get_attachment_registry().clear() + + def _router(self) -> MagicMock: + from litellm.types.router import Deployment, LiteLLM_Params + + router = MagicMock() + router.get_deployment.side_effect = lambda model_id: ( + Deployment( + model_name=self.GOVERNED_MODEL_GROUP, + litellm_params=LiteLLM_Params(model=f"openai/{self.GOVERNED_MODEL_GROUP}"), + model_info={"id": model_id}, + ) + if model_id == self.GOVERNED_MODEL_ID + else None + ) + return router + + async def _pre_call(self, route_type: str, monkeypatch) -> dict: + from litellm.responses.utils import ResponsesAPIRequestUtils + + client_facing_response_id = "resp_opaque-client-facing-id" + encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=self.GOVERNED_MODEL_ID, response_id="resp_upstream" + ) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"response_id": client_facing_response_id, "litellm_metadata": {}} + ) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def passthrough_add_litellm_data_to_request(data, **kwargs): + return data + + async def decrypting_pre_call_hook(user_api_key_dict, data, call_type): + if data.get("response_id") == client_facing_response_id: + data["response_id"] = encoded_response_id + return data + + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + passthrough_add_litellm_data_to_request, + ) + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=decrypting_pre_call_hook) + proxy_config = MagicMock(spec=ProxyConfig) + proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(), + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type=route_type, + llm_router=self._router(), + ) + return returned_data + + @pytest.mark.asyncio + async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline( + self, policy_engine, monkeypatch + ): + data = await self._pre_call("aget_responses", monkeypatch) + + assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f") + pipelines = data["litellm_metadata"]["_guardrail_pipelines"] + assert [(policy_name, [step.guardrail for step in pipeline.steps]) for policy_name, pipeline in pipelines] == [ + ("response-governance", ["output-word-filter"]) + ] + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + assert data["model"] is None + + @pytest.mark.asyncio + async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id( + self, policy_engine, monkeypatch + ): + data = await self._pre_call("aresponses", monkeypatch) + + assert "_guardrail_pipelines" not in data["litellm_metadata"] + assert "applied_policies" not in data["litellm_metadata"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 73dd6746b29..7d2c232ad97 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -30,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_g from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, @@ -1419,8 +1420,69 @@ async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_ver assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog)) +def _background_response(status: str, text: str = "") -> ResponsesAPIResponse: + output = ( + [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}] if text else [] + ) + return ResponsesAPIResponse(id="resp_bg", created_at=0, output=output, status=status) + + +def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail]: + class OutputBlockingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + return [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] + + @pytest.mark.asyncio -async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline( +@pytest.mark.parametrize("pending_status", ["queued", "in_progress"]) +async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog, pending_status +): + seen: dict[str, object] = {} + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(background=True) + response = _background_response(pending_status) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert out is response + assert "response" not in seen + assert not _warnings(caplog) + assert any( + "response-governance" in record.getMessage() and pending_status in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final_status", ["completed", "incomplete"]) +async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response( + proxy_logging, make_user_api_key_auth, monkeypatch, final_status +): + seen: dict[str, object] = {} + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = _background_response(final_status, text="kumquat") + + with pytest.raises(HTTPException) as info: + await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert info.value.detail["error"] == "output blocked" + assert seen["response"] is response + + +@pytest.mark.asyncio +async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline( proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): monkeypatch.setattr(litellm, "callbacks", []) @@ -1436,36 +1498,7 @@ async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline( assert out is not None assert out.get("background") is True - assert any("response-governance" in message and "background" in message for message in _warnings(caplog)) - - -@pytest.mark.asyncio -async def test_pre_call_hook_stays_quiet_on_background_request_without_post_call_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, caplog -): - seen: Dict[str, Any] = {} - monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) - pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) - data = { - "model": "m", - "messages": [{"role": "user", "content": "hi"}], - "background": True, - "metadata": { - "_guardrail_pipelines": [("request-governance", pre_call)], - "_pipeline_managed_guardrails": {"gr-post"}, - }, - } - - with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - out = await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), - data=data, - call_type="aresponses", - guardrails_only=True, - ) - - assert out is not None - assert not any("background" in message for message in _warnings(caplog)) + assert not _warnings(caplog) # --------------------------------------------------------------------------- From 2823d09e1bcfc99cee126583f3552e5d5b32dac5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:30:05 -0700 Subject: [PATCH 05/23] refactor(policy_engine): type the request-state parameter of the retrieval hook --- litellm/proxy/policy_engine/response_retrieval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 54b9acf6150..0aa9579ca71 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -71,7 +71,7 @@ def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[Polic def attach_post_call_pipelines_to_retrieval( - data: dict, # mutable-ok: the proxy's request-state dict, written in place like every other policy engine hook + data: dict[str, object], # mutable-ok: request-state dict the policy engine hooks all write in place user_api_key_dict: "UserAPIKeyAuth", llm_router: "Router | None", ) -> None: From 91091fd93e7b7be6576be26e1ef1216391af667e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:59:58 -0700 Subject: [PATCH 06/23] fix: withdraw policy header claims while a background response is pending and log ungoverned retrievals --- .../proxy/policy_engine/response_retrieval.py | 31 ++++-- litellm/proxy/utils.py | 62 +++++++++++- .../policy_engine/test_response_retrieval.py | 36 +++++-- .../proxy_logging/test_guardrail_pipeline.py | 97 +++++++++++++++++++ 4 files changed, 205 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 0aa9579ca71..8fc75e40070 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -1,6 +1,7 @@ from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Final, TypeAlias +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import TypeAdapter @@ -29,14 +30,23 @@ PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...] _POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines) -def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | None: - if llm_router is None or not isinstance(response_id, str): - return None - model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) +@dataclass(frozen=True, slots=True) +class UngovernedRetrieval: + reason: Literal["no router", "response id names no deployment", "deployment no longer in the router"] + + +def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | UngovernedRetrieval: + if llm_router is None: + return UngovernedRetrieval("no router") + model_id: Final = ( + ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) if isinstance(response_id, str) else None + ) if model_id is None: - return None + return UngovernedRetrieval("response id names no deployment") deployment: Final = llm_router.get_deployment(model_id) - return deployment.model_name if deployment is not None else None + if deployment is None: + return UngovernedRetrieval("deployment no longer in the router") + return deployment.model_name def _retrieval_context( @@ -78,7 +88,12 @@ def attach_post_call_pipelines_to_retrieval( if not get_policy_registry().is_initialized(): return model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router) - if model_group is None: + if isinstance(model_group, UngovernedRetrieval): + verbose_proxy_logger.warning( + "Policy engine: background response %s is retrieved without its post_call policy pipelines (%s)", + data.get("response_id"), + model_group.reason, + ) return context: Final = _retrieval_context(data, user_api_key_dict, model_group) post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3d03d3858ee..1972d92b610 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -93,6 +93,7 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, + get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, ) @@ -156,6 +157,8 @@ from litellm.proxy.hooks.sensitive_data_routing import ( from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( @@ -533,16 +536,65 @@ def _is_pending_background_response(response: LLMResponseTypes) -> bool: return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES -def _log_deferred_post_call_pipelines(data: Mapping[str, object], response: ResponsesAPIResponse) -> None: - policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data)) - if not policy_names: +def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]: + resolved: Final = PolicyResolver.resolve_policy_guardrails( + policy_name=policy_name, policies=get_policy_registry().get_all_policies() + ) + return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps) + + +def _without_names( + bucket: dict, # mutable-ok: the applied_* header slots live in the request-state dict every hook writes in place + slot: str, + names: frozenset[str], +) -> None: + claimed: Final = bucket.get(slot) + if not isinstance(claimed, list): + return + remaining: Final = [name for name in claimed if name not in names] + if remaining: + bucket[slot] = remaining + else: + bucket.pop(slot) + + +def _withdraw_deferred_claims( + data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + deferred: Sequence[tuple[str, "GuardrailPipeline"]], +) -> None: + outside_by_policy: Final = MappingProxyType( + {policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred} + ) + running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(*outside_by_policy.values()) + withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside) + withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere + _, bucket = get_or_create_metadata_bucket(data) + _without_names(bucket, "applied_policies", withdrawn_policies) + _without_names(bucket, "applied_guardrails", withdrawn_guardrails) + sources: Final = bucket.get("policy_sources") + if not isinstance(sources, dict): + return + remaining_sources: Final = {name: reason for name, reason in sources.items() if name not in withdrawn_policies} + if remaining_sources: + bucket["policy_sources"] = remaining_sources + else: + bucket.pop("policy_sources") + + +def _defer_post_call_pipelines( + data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + response: ResponsesAPIResponse, +) -> None: + deferred: Final = _post_call_pipelines(data) + if not deferred: return verbose_proxy_logger.debug( "Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s", response.id, response.status, - ", ".join(policy_names), + ", ".join(policy_name for policy_name, _pipeline in deferred), ) + _withdraw_deferred_claims(data, deferred) def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: @@ -2939,7 +2991,7 @@ class ProxyLogging: response: LLMResponseTypes, ) -> LLMResponseTypes | None: if _is_pending_background_response(response): - _log_deferred_post_call_pipelines(data, response) + _defer_post_call_pipelines(data, response) return None _, pipeline_response = await self._maybe_execute_pipelines( data=data, diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index f7f29d6b46e..8c00b18eb63 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -1,3 +1,5 @@ +import logging + import pytest from litellm.proxy._types import UserAPIKeyAuth @@ -128,30 +130,48 @@ def test_already_attached_policy_is_not_attached_twice(policy_engine): assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] +def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "retrieved without its post_call policy pipelines" in record.getMessage() + ] + + @pytest.mark.parametrize( - "response_id", - ["resp_plain_upstream_id", _encoded_response_id("deployment-missing-from-router"), None], + ("response_id", "reason"), + [ + ("resp_plain_upstream_id", "response id names no deployment"), + (_encoded_response_id("deployment-missing-from-router"), "deployment no longer in the router"), + (None, "response id names no deployment"), + ], ) -def test_unresolvable_response_id_attaches_nothing(policy_engine, response_id): +def test_unresolvable_response_id_attaches_nothing_and_warns(policy_engine, caplog, response_id, reason): data = {"response_id": response_id, "litellm_metadata": {}} - attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) assert data == {"response_id": response_id, "litellm_metadata": {}} + assert [message.endswith(f"({reason})") for message in _ungoverned_retrieval_warnings(caplog)] == [True] -def test_without_a_router_attaches_nothing(policy_engine): +def test_without_a_router_attaches_nothing_and_warns(policy_engine, caplog): data = _retrieval_data(GOVERNED_MODEL_ID) - attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) assert data == _retrieval_data(GOVERNED_MODEL_ID) + assert [message.endswith("(no router)") for message in _ungoverned_retrieval_warnings(caplog)] == [True] -def test_without_policy_engine_attaches_nothing(): +def test_without_policy_engine_attaches_nothing_quietly(caplog): get_policy_registry().clear() data = _retrieval_data(GOVERNED_MODEL_ID) - attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) assert data == _retrieval_data(GOVERNED_MODEL_ID) + assert _ungoverned_retrieval_warnings(caplog) == [] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 7d2c232ad97..8da07f39fd8 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1481,6 +1481,103 @@ async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_resp assert seen["response"] is response +def _output_passing_callbacks() -> list[CustomGuardrail]: + class OutputPassingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + return [OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] + + +def _claimed_post_call_pipeline_data(*policy_names: str, extra_guardrails: dict[str, list[str]] | None = None): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + step = {"guardrail": "gr-post", "on_pass": "allow", "on_fail": "block"} + get_policy_registry().load_policies( + { + policy_name: { + "guardrails": {"add": ["gr-post", *(extra_guardrails or {}).get(policy_name, [])]}, + "pipeline": {"mode": "post_call", "steps": [step]}, + } + for policy_name in policy_names + } + ) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(**step)]) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [(policy_name, pipeline) for policy_name in policy_names], + "_pipeline_managed_guardrails": {"gr-post"}, + "applied_policies": list(policy_names), + "applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)], + "policy_sources": {policy_name: "model:m" for policy_name in policy_names}, + }, + } + + +@pytest.fixture +def clear_policy_registry(): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + yield + get_policy_registry().clear() + + +@pytest.mark.asyncio +async def test_pending_background_response_withdraws_the_deferred_policy_claims( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "applied_policies" not in data["metadata"] + assert "policy_sources" not in data["metadata"] + assert "applied_guardrails" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data( + "input-and-output-governance", "response-governance", extra_guardrails={"input-and-output-governance": ["gr-pre"]} + ) + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("in_progress"), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_policies"] == ["input-and-output-governance"] + assert data["metadata"]["applied_guardrails"] == ["gr-pre"] + assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"} + + +@pytest.mark.asyncio +async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + monkeypatch.setattr(litellm, "callbacks", _output_passing_callbacks()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("completed", text="fine"), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_policies"] == ["response-governance"] + assert data["metadata"]["policy_sources"] == {"response-governance": "model:m"} + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + @pytest.mark.asyncio async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline( proxy_logging, make_user_api_key_auth, monkeypatch, caplog From 8ce3fe4fcf5ba43278c4eb7ec09c58df513d36c5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:13:36 -0700 Subject: [PATCH 07/23] style: give the header slot rebuilds and the in-place slot write their lint reasons --- litellm/proxy/utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1972d92b610..e48bd3851fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -551,9 +551,11 @@ def _without_names( claimed: Final = bucket.get(slot) if not isinstance(claimed, list): return - remaining: Final = [name for name in claimed if name not in names] + remaining: Final = [ # mutable-ok: the slot stays a list, the shape every applied_* header writer appends to + name for name in claimed if name not in names + ] if remaining: - bucket[slot] = remaining + bucket[slot] = remaining # rebind-ok: the slot lives in the shared request-state dict, rewritten in place else: bucket.pop(slot) @@ -574,7 +576,9 @@ def _withdraw_deferred_claims( sources: Final = bucket.get("policy_sources") if not isinstance(sources, dict): return - remaining_sources: Final = {name: reason for name, reason in sources.items() if name not in withdrawn_policies} + remaining_sources: Final = { # mutable-ok: policy_sources stays a dict, the shape its writer updates in place + name: reason for name, reason in sources.items() if name not in withdrawn_policies + } if remaining_sources: bucket["policy_sources"] = remaining_sources else: From 071f83980c319b9d9a21f77179781052c050b517 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:26:17 -0700 Subject: [PATCH 08/23] refactor(proxy): type the post_call deferral helpers' request-state parameters --- litellm/proxy/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e48bd3851fb..e653ca6b62c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -544,7 +544,7 @@ def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline" def _without_names( - bucket: dict, # mutable-ok: the applied_* header slots live in the request-state dict every hook writes in place + bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write slot: str, names: frozenset[str], ) -> None: @@ -561,7 +561,7 @@ def _without_names( def _withdraw_deferred_claims( - data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data deferred: Sequence[tuple[str, "GuardrailPipeline"]], ) -> None: outside_by_policy: Final = MappingProxyType( @@ -586,7 +586,7 @@ def _withdraw_deferred_claims( def _defer_post_call_pipelines( - data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data response: ResponsesAPIResponse, ) -> None: deferred: Final = _post_call_pipelines(data) @@ -2990,7 +2990,7 @@ class ProxyLogging: async def _run_post_call_pipelines( self, - data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes, ) -> LLMResponseTypes | None: From 547f81c1a52722f487597d93975df1fa79e9464e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:46:23 -0700 Subject: [PATCH 09/23] test: type the background response retrieval test helpers --- .../proxy/policy_engine/test_response_retrieval.py | 9 ++++++--- .../test_litellm/proxy/test_common_request_processing.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index 8c00b18eb63..b704f666646 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -1,4 +1,5 @@ import logging +from collections.abc import Mapping import pytest @@ -76,14 +77,16 @@ def policy_engine(): attachment_registry.clear() -def _retrieval_data(model_id: str) -> dict: +def _retrieval_data(model_id: str) -> dict[str, object]: return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}} -def _attached_pipelines(data: dict) -> tuple[tuple[str, str], ...]: +def _attached_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + bucket = data["litellm_metadata"] + assert isinstance(bucket, dict) return tuple( (policy_name, ",".join(step.guardrail for step in pipeline.steps)) - for policy_name, pipeline in data["litellm_metadata"]["_guardrail_pipelines"] + for policy_name, pipeline in bucket["_guardrail_pipelines"] ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 08007e22cfb..c451a3b4cb0 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8340,7 +8340,7 @@ class TestBackgroundResponseRetrievalGovernance: ) return router - async def _pre_call(self, route_type: str, monkeypatch) -> dict: + async def _pre_call(self, route_type: str, monkeypatch: pytest.MonkeyPatch) -> dict[str, object]: from litellm.responses.utils import ResponsesAPIRequestUtils client_facing_response_id = "resp_opaque-client-facing-id" From 345298f3c90871a7099952cb08e1e19d396883ee Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:04:19 -0700 Subject: [PATCH 10/23] fix(policy_engine): warn when a poll cannot re-match the submitted model name and keep default_on pre_call claims --- .../proxy/policy_engine/response_retrieval.py | 21 +++++++ litellm/proxy/utils.py | 14 ++++- .../policy_engine/test_response_retrieval.py | 56 ++++++++++++++++++- .../proxy_logging/test_guardrail_pipeline.py | 24 ++++++++ 4 files changed, 111 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 8fc75e40070..d284c44397e 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -18,6 +18,7 @@ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.types.proxy.policy_engine import PolicyMatchContext from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline @@ -46,9 +47,29 @@ def _model_group_for_response_id(response_id: object, llm_router: "Router | None deployment: Final = llm_router.get_deployment(model_id) if deployment is None: return UngovernedRetrieval("deployment no longer in the router") + hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias) + if hidden_by is not None: + verbose_proxy_logger.warning( + "Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), " + "so a policy attached to the model name it was submitted as does not run on it", + response_id, + deployment.model_name, + hidden_by, + ) return deployment.model_name +def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None: + if "*" in model_group: + return "a wildcard deployment" + aliases: Final = tuple( + alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group + ) + if not aliases: + return None + return f"the target of model_group_alias {', '.join(aliases)}" + + def _retrieval_context( data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str ) -> PolicyMatchContext: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e653ca6b62c..37ca52d6fc9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -543,6 +543,16 @@ def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline" return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps) +def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]: + return frozenset( + callback.guardrail_name + for callback in litellm.callbacks + if isinstance(callback, CustomGuardrail) + and callback.guardrail_name is not None + and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + ) + + def _without_names( bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write slot: str, @@ -567,7 +577,9 @@ def _withdraw_deferred_claims( outside_by_policy: Final = MappingProxyType( {policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred} ) - running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(*outside_by_policy.values()) + running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union( + _guardrails_run_standalone_pre_call(data), *outside_by_policy.values() + ) withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside) withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere _, bucket = get_or_create_metadata_bucket(data) diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index b704f666646..9dcbef888ce 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -14,11 +14,14 @@ GOVERNED_MODEL_GROUP = "gpt-5.4-mini" GOVERNED_MODEL_ID = "deployment-governed" UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini" UNGOVERNED_MODEL_ID = "deployment-ungoverned" +WILDCARD_MODEL_GROUP = "openai/*" +WILDCARD_MODEL_ID = "deployment-wildcard" class FakeRouter: - def __init__(self, deployments: dict[str, Deployment]): + def __init__(self, deployments: dict[str, Deployment], model_group_alias: dict[str, object] | None = None): self._deployments = deployments + self.model_group_alias = model_group_alias or {} def get_deployment(self, model_id: str) -> Deployment | None: return self._deployments.get(model_id) @@ -32,12 +35,14 @@ def _deployment(model_group: str, model_id: str) -> Deployment: ) -def _router() -> FakeRouter: +def _router(model_group_alias: dict[str, object] | None = None) -> FakeRouter: return FakeRouter( { GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID), UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID), - } + WILDCARD_MODEL_ID: _deployment(WILDCARD_MODEL_GROUP, WILDCARD_MODEL_ID), + }, + model_group_alias, ) @@ -133,6 +138,51 @@ def test_already_attached_policy_is_not_attached_twice(policy_engine): assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] +def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "the model name it was submitted as" in record.getMessage() + ] + + +def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns(policy_engine, caplog): + data = _retrieval_data(WILDCARD_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(WILDCARD_MODEL_ID) + assert [ + "as model group openai/* (a wildcard deployment)" in message for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_aliased_model_group_still_attaches_its_own_policies_and_warns(policy_engine, caplog): + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}}) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert [ + "(the target of model_group_alias gpt-mini, gpt-hidden)" in message + for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model(policy_engine, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval( + data=_retrieval_data(GOVERNED_MODEL_ID), + user_api_key_dict=UserAPIKeyAuth(), + llm_router=_router({"other-alias": UNGOVERNED_MODEL_GROUP}), + ) + + assert _hidden_submit_model_warnings(caplog) == [] + + def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: return [ record.getMessage() diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 8da07f39fd8..7e0736d4e67 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1561,6 +1561,30 @@ async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"} +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + class DualStageGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-post", event_hook=["pre_call", "post_call"], default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert "applied_policies" not in data["metadata"] + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + @pytest.mark.asyncio async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry From 0c58346ba925d0a15ebcb38b162b9224446a50d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:07:37 -0700 Subject: [PATCH 11/23] fix(proxy): warn when a deferred background policy was matched through a request tag Retrieval re-matches only the key, team, and model scopes, so a post_call policy that reached a pending background response through a request-body tag does not govern the completed response. Log that at submit, next to the deferral, and cover the retrieval re-match with tag-scoped tests. --- litellm/proxy/utils.py | 22 +++++++++ .../policy_engine/test_response_retrieval.py | 22 +++++++++ .../proxy_logging/test_guardrail_pipeline.py | 47 ++++++++++++++++++- 3 files changed, 89 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4b62fc4cdae..d144aa113d6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -613,9 +613,31 @@ def _defer_post_call_pipelines( response.status, ", ".join(policy_name for policy_name, _pipeline in deferred), ) + tag_matched: Final = _tag_matched_deferrals(data, deferred) + if tag_matched: + verbose_proxy_logger.warning( + "Policy engine: background response %s matched post_call policies through a request tag at submit; " + "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " + "does not govern the completed response: %s", + response.id, + ", ".join(f"{policy_name} ({source})" for policy_name, source in tag_matched), + ) _withdraw_deferred_claims(data, deferred) +def _tag_matched_deferrals( + data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] +) -> tuple[tuple[str, str], ...]: + sources: Final = _policy_state_metadata(data).get("policy_sources") + if not isinstance(sources, dict): + return () + return tuple( + (policy_name, str(sources[policy_name])) + for policy_name, _pipeline in deferred + if policy_name in sources and "tag:" in str(sources[policy_name]) + ) + + def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: unsupported: Final = tuple( dict.fromkeys( diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index 9dcbef888ce..1fa1161191d 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -68,6 +68,7 @@ def policy_engine(): "response-governance": _pipeline_policy("output-word-filter"), "input-governance": _pipeline_policy("input-word-filter", mode="pre_call"), "team-governance": _pipeline_policy("team-word-filter"), + "tag-governance": _pipeline_policy("tag-word-filter"), } ) attachment_registry.load_attachments( @@ -75,6 +76,7 @@ def policy_engine(): {"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]}, {"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]}, {"policy": "team-governance", "teams": ["governed-team"]}, + {"policy": "tag-governance", "tags": ["governed"]}, ] ) yield @@ -119,6 +121,26 @@ def test_key_and_team_context_also_governs_retrieval(policy_engine): assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),) +def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine): + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + + +def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine): + data: dict[str, object] = { + "response_id": _encoded_response_id(UNGOVERNED_MODEL_ID), + "litellm_metadata": {"tags": ["governed"]}, + } + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("tag-governance", "tag-word-filter"),) + assert data["litellm_metadata"]["policy_sources"] == {"tag-governance": "tag:governed"} + + def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine): data = _retrieval_data(UNGOVERNED_MODEL_ID) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index d6a74d3bbd9..63afce5d968 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1489,7 +1489,9 @@ def _output_passing_callbacks() -> list[CustomGuardrail]: return [OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] -def _claimed_post_call_pipeline_data(*policy_names: str, extra_guardrails: dict[str, list[str]] | None = None): +def _claimed_post_call_pipeline_data( + *policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str = "model:m" +): from litellm.proxy.policy_engine.policy_registry import get_policy_registry step = {"guardrail": "gr-post", "on_pass": "allow", "on_fail": "block"} @@ -1511,7 +1513,7 @@ def _claimed_post_call_pipeline_data(*policy_names: str, extra_guardrails: dict[ "_pipeline_managed_guardrails": {"gr-post"}, "applied_policies": list(policy_names), "applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)], - "policy_sources": {policy_name: "model:m" for policy_name in policy_names}, + "policy_sources": {policy_name: policy_source for policy_name in policy_names}, }, } @@ -1542,6 +1544,47 @@ async def test_pending_background_response_withdraws_the_deferred_policy_claims( assert "applied_guardrails" not in data["metadata"] +@pytest.mark.asyncio +async def test_pending_background_response_warns_when_the_deferred_policy_was_matched_through_a_tag( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry, caplog +): + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance", policy_source="tag:governed+model:m") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "policy_sources" not in data["metadata"] + assert [ + message for message in _warnings(caplog) if "response-governance (tag:governed+model:m)" in message + ] == [ + "Policy engine: background response resp_bg matched post_call policies through a request tag at submit; " + "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " + "does not govern the completed response: response-governance (tag:governed+model:m)" + ] + + +@pytest.mark.asyncio +async def test_pending_background_response_matched_through_its_model_does_not_warn( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry, caplog +): + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await proxy_logging.post_call_success_hook( + data=_claimed_post_call_pipeline_data("response-governance"), + response=_background_response("queued"), + user_api_key_dict=make_user_api_key_auth(), + ) + + assert _warnings(caplog) == [] + + @pytest.mark.asyncio async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline( proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry From 94f9230d13900897f764498169bdfa3c5fa2a4ea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:59:09 -0700 Subject: [PATCH 12/23] fix(proxy): type the new pipeline tests and keep tag values out of the deferral warning Every test this PR adds now annotates its fixture and parametrize parameters. The submit-time warning for a tag-matched deferred policy names only the policies, since a wildcard attachment pattern would let caller-provided tag text reach the log. --- litellm/proxy/utils.py | 6 +- .../policy_engine/test_response_retrieval.py | 42 ++-- .../proxy_logging/test_guardrail_pipeline.py | 216 ++++++++++-------- 3 files changed, 155 insertions(+), 109 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d144aa113d6..1604dea1003 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -620,19 +620,19 @@ def _defer_post_call_pipelines( "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " "does not govern the completed response: %s", response.id, - ", ".join(f"{policy_name} ({source})" for policy_name, source in tag_matched), + ", ".join(tag_matched), ) _withdraw_deferred_claims(data, deferred) def _tag_matched_deferrals( data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] -) -> tuple[tuple[str, str], ...]: +) -> tuple[str, ...]: sources: Final = _policy_state_metadata(data).get("policy_sources") if not isinstance(sources, dict): return () return tuple( - (policy_name, str(sources[policy_name])) + policy_name for policy_name, _pipeline in deferred if policy_name in sources and "tag:" in str(sources[policy_name]) ) diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index 1fa1161191d..37849101b3a 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -1,5 +1,5 @@ import logging -from collections.abc import Mapping +from collections.abc import Iterator, Mapping import pytest @@ -60,7 +60,7 @@ def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, objec @pytest.fixture -def policy_engine(): +def policy_engine() -> Iterator[None]: policy_registry = get_policy_registry() attachment_registry = get_attachment_registry() policy_registry.load_policies( @@ -97,7 +97,7 @@ def _attached_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, str], .. ) -def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine): +def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine: None) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) @@ -111,7 +111,7 @@ def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine): assert "guardrails" not in data["litellm_metadata"] -def test_key_and_team_context_also_governs_retrieval(policy_engine): +def test_key_and_team_context_also_governs_retrieval(policy_engine: None) -> None: data = _retrieval_data(UNGOVERNED_MODEL_ID) attach_post_call_pipelines_to_retrieval( @@ -121,7 +121,7 @@ def test_key_and_team_context_also_governs_retrieval(policy_engine): assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),) -def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine): +def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine: None) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) @@ -129,7 +129,7 @@ def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) -def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine): +def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine: None) -> None: data: dict[str, object] = { "response_id": _encoded_response_id(UNGOVERNED_MODEL_ID), "litellm_metadata": {"tags": ["governed"]}, @@ -141,7 +141,7 @@ def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag( assert data["litellm_metadata"]["policy_sources"] == {"tag-governance": "tag:governed"} -def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine): +def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine: None) -> None: data = _retrieval_data(UNGOVERNED_MODEL_ID) attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) @@ -149,7 +149,7 @@ def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine): assert data == _retrieval_data(UNGOVERNED_MODEL_ID) -def test_already_attached_policy_is_not_attached_twice(policy_engine): +def test_already_attached_policy_is_not_attached_twice(policy_engine: None) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) router = _router() attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) @@ -168,7 +168,9 @@ def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str] ] -def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns(policy_engine, caplog): +def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: data = _retrieval_data(WILDCARD_MODEL_ID) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -176,11 +178,14 @@ def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns( assert data == _retrieval_data(WILDCARD_MODEL_ID) assert [ - "as model group openai/* (a wildcard deployment)" in message for message in _hidden_submit_model_warnings(caplog) + "as model group openai/* (a wildcard deployment)" in message + for message in _hidden_submit_model_warnings(caplog) ] == [True] -def test_aliased_model_group_still_attaches_its_own_policies_and_warns(policy_engine, caplog): +def test_aliased_model_group_still_attaches_its_own_policies_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}}) @@ -194,7 +199,9 @@ def test_aliased_model_group_still_attaches_its_own_policies_and_warns(policy_en ] == [True] -def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model(policy_engine, caplog): +def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): attach_post_call_pipelines_to_retrieval( data=_retrieval_data(GOVERNED_MODEL_ID), @@ -209,7 +216,8 @@ def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str return [ record.getMessage() for record in caplog.records - if record.levelno == logging.WARNING and "retrieved without its post_call policy pipelines" in record.getMessage() + if record.levelno == logging.WARNING + and "retrieved without its post_call policy pipelines" in record.getMessage() ] @@ -221,7 +229,9 @@ def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str (None, "response id names no deployment"), ], ) -def test_unresolvable_response_id_attaches_nothing_and_warns(policy_engine, caplog, response_id, reason): +def test_unresolvable_response_id_attaches_nothing_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture, response_id: str, reason: str +) -> None: data = {"response_id": response_id, "litellm_metadata": {}} with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -231,7 +241,7 @@ def test_unresolvable_response_id_attaches_nothing_and_warns(policy_engine, capl assert [message.endswith(f"({reason})") for message in _ungoverned_retrieval_warnings(caplog)] == [True] -def test_without_a_router_attaches_nothing_and_warns(policy_engine, caplog): +def test_without_a_router_attaches_nothing_and_warns(policy_engine: None, caplog: pytest.LogCaptureFixture) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -241,7 +251,7 @@ def test_without_a_router_attaches_nothing_and_warns(policy_engine, caplog): assert [message.endswith("(no router)") for message in _ungoverned_retrieval_warnings(caplog)] == [True] -def test_without_policy_engine_attaches_nothing_quietly(caplog): +def test_without_policy_engine_attaches_nothing_quietly(caplog: pytest.LogCaptureFixture) -> None: get_policy_registry().clear() data = _retrieval_data(GOVERNED_MODEL_ID) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 63afce5d968..3bbdbbe2ccd 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -12,6 +12,7 @@ from __future__ import annotations import asyncio import json import logging +from collections.abc import Iterator from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -155,9 +156,7 @@ async def test_execute_guardrail_hook_unknown_hook_type_raises(proxy_logging, ma @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_routes_through_router( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_routes_through_router(proxy_logging, make_user_api_key_auth): cb = _make_guardrail() router = MagicMock() router.get_available_guardrail = MagicMock(return_value={"callback": cb}) @@ -173,9 +172,7 @@ async def test_execute_guardrail_with_load_balancing_routes_through_router( @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_router_none_raises( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_router_none_raises(proxy_logging, make_user_api_key_auth): with patch("litellm.proxy.proxy_server.llm_router", None): with pytest.raises(ValueError, match="Router not initialized"): await proxy_logging._execute_guardrail_with_load_balancing( @@ -188,9 +185,7 @@ async def test_execute_guardrail_with_load_balancing_router_none_raises( @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_no_callback_raises( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_no_callback_raises(proxy_logging, make_user_api_key_auth): router = MagicMock() router.get_available_guardrail = MagicMock(return_value={"callback": None}) with patch("litellm.proxy.proxy_server.llm_router", router): @@ -210,9 +205,7 @@ async def test_execute_guardrail_with_load_balancing_no_callback_raises( @pytest.mark.asyncio -async def test_process_guardrail_callback_skipped_when_should_run_false( - proxy_logging, make_user_api_key_auth -): +async def test_process_guardrail_callback_skipped_when_should_run_false(proxy_logging, make_user_api_key_auth): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=False) out = await proxy_logging._process_guardrail_callback( @@ -226,9 +219,7 @@ async def test_process_guardrail_callback_skipped_when_should_run_false( @pytest.mark.asyncio -async def test_process_guardrail_callback_returns_data_on_success( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_process_guardrail_callback_returns_data_on_success(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=True) proxy_logging._should_use_guardrail_load_balancing = MagicMock(return_value=False) @@ -343,14 +334,14 @@ async def test_maybe_execute_pipelines_no_pipelines_returns_data(proxy_logging, @pytest.mark.asyncio -async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode( + proxy_logging, make_user_api_key_auth, monkeypatch +): pipeline = MagicMock() pipeline.mode = "post_call" # not pre_call data = {"metadata": {"_guardrail_pipelines": [("p1", pipeline)]}, "model": "m", "messages": []} executed = MagicMock() - monkeypatch.setattr( - "litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed - ) + monkeypatch.setattr("litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed) out, replacement = await proxy_logging._maybe_execute_pipelines( data=data, user_api_key_dict=make_user_api_key_auth(), @@ -538,9 +529,7 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode(): litellm.callbacks = [cb] try: with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result( - result=result, data={"model": "m"}, policy_name="p" - ) + ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") finally: litellm.callbacks = saved @@ -652,9 +641,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch monkeypatch.setattr(litellm, "callbacks", [prom]) with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_with_metrics( - callback=cb, coro=task(), hook_type="post_call" - ) + await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call") assert detail["guardrail_name"] == "presidio" recorded = prom._record_guardrail_metrics.call_args.kwargs @@ -682,9 +669,7 @@ def _moderation_guardrail() -> MagicMock: @pytest.mark.asyncio -async def test_during_call_hook_records_latency_metric( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail() prom = _prometheus_callback() monkeypatch.setattr(litellm, "callbacks", [prom, cb]) @@ -703,9 +688,7 @@ async def test_during_call_hook_records_latency_metric( @pytest.mark.asyncio -async def test_post_call_success_hook_records_latency_metric( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail() prom = _prometheus_callback() monkeypatch.setattr(litellm, "callbacks", [prom, cb]) @@ -733,9 +716,7 @@ async def test_post_call_success_hook_records_latency_metric( async def test_process_prompt_template_no_op_when_no_prompt_spec(proxy_logging, monkeypatch): from litellm.proxy.prompts import prompt_registry - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None) data: Dict[str, Any] = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} await proxy_logging._process_prompt_template( data=data, @@ -760,9 +741,7 @@ async def test_process_prompt_template_applies_when_spec_resolves(proxy_logging, "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock( @@ -810,9 +789,7 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=RuntimeError("bad prompt")) with pytest.raises(RuntimeError): @@ -913,9 +890,7 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock( @@ -1124,9 +1099,7 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe }, } - await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" - ) + await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion") assert seen["count"] == 1 @@ -1289,11 +1262,7 @@ async def test_post_call_pipeline_block_keeps_guardrail_metadata_writes( monkeypatch.setattr( litellm, "callbacks", - [ - BlockingWriterGuardrail( - guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False - ) - ], + [BlockingWriterGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data() @@ -1379,9 +1348,7 @@ async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once( }, } - await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" - ) + await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion") assert seen["count"] == 1 @@ -1433,14 +1400,20 @@ def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail] seen["response"] = response raise HTTPException(status_code=400, detail={"error": "output blocked"}) - return [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] + return [ + OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + ] @pytest.mark.asyncio @pytest.mark.parametrize("pending_status", ["queued", "in_progress"]) async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, caplog, pending_status -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + pending_status: str, +) -> None: seen: dict[str, object] = {} monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) @@ -1464,8 +1437,11 @@ async def test_post_call_success_hook_waits_for_pending_background_response_befo @pytest.mark.asyncio @pytest.mark.parametrize("final_status", ["completed", "incomplete"]) async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response( - proxy_logging, make_user_api_key_auth, monkeypatch, final_status -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + final_status: str, +) -> None: seen: dict[str, object] = {} monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) @@ -1486,7 +1462,9 @@ def _output_passing_callbacks() -> list[CustomGuardrail]: async def async_post_call_success_hook(self, data, user_api_key_dict, response): return response - return [OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] + return [ + OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + ] def _claimed_post_call_pipeline_data( @@ -1519,7 +1497,7 @@ def _claimed_post_call_pipeline_data( @pytest.fixture -def clear_policy_registry(): +def clear_policy_registry() -> Iterator[None]: from litellm.proxy.policy_engine.policy_registry import get_policy_registry yield @@ -1528,8 +1506,11 @@ def clear_policy_registry(): @pytest.mark.asyncio async def test_pending_background_response_withdraws_the_deferred_policy_claims( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _claimed_post_call_pipeline_data("response-governance") @@ -1546,8 +1527,12 @@ async def test_pending_background_response_withdraws_the_deferred_policy_claims( @pytest.mark.asyncio async def test_pending_background_response_warns_when_the_deferred_policy_was_matched_through_a_tag( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry, caplog -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _claimed_post_call_pipeline_data("response-governance", policy_source="tag:governed+model:m") @@ -1559,19 +1544,21 @@ async def test_pending_background_response_warns_when_the_deferred_policy_was_ma assert out.status == "queued" assert "policy_sources" not in data["metadata"] - assert [ - message for message in _warnings(caplog) if "response-governance (tag:governed+model:m)" in message - ] == [ + assert [message for message in _warnings(caplog) if "through a request tag" in message] == [ "Policy engine: background response resp_bg matched post_call policies through a request tag at submit; " "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " - "does not govern the completed response: response-governance (tag:governed+model:m)" + "does not govern the completed response: response-governance" ] @pytest.mark.asyncio async def test_pending_background_response_matched_through_its_model_does_not_warn( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry, caplog -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) @@ -1587,12 +1574,17 @@ async def test_pending_background_response_matched_through_its_model_does_not_wa @pytest.mark.asyncio async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _claimed_post_call_pipeline_data( - "input-and-output-governance", "response-governance", extra_guardrails={"input-and-output-governance": ["gr-pre"]} + "input-and-output-governance", + "response-governance", + extra_guardrails={"input-and-output-governance": ["gr-pre"]}, ) await proxy_logging.post_call_success_hook( @@ -1606,8 +1598,11 @@ async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs @pytest.mark.asyncio async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: class DualStageGuardrail(CustomGuardrail): async def async_post_call_success_hook(self, data, user_api_key_dict, response): return response @@ -1630,8 +1625,11 @@ async def test_pending_background_response_keeps_the_claim_of_a_default_on_guard @pytest.mark.asyncio async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_passing_callbacks()) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _claimed_post_call_pipeline_data("response-governance") @@ -1647,8 +1645,11 @@ async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pip @pytest.mark.asyncio async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, caplog -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: monkeypatch.setattr(litellm, "callbacks", []) data = _post_call_pipeline_data(background=True) @@ -1709,7 +1710,9 @@ def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported( steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")], ) pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")]) - data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}} + data = { + "metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]} + } with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions")) @@ -2009,7 +2012,9 @@ def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str, async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): return {**inputs, **transform(inputs)} - return RewritingStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + return RewritingStreamGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) def _tool_call_stream_chunks() -> List[Any]: @@ -2188,11 +2193,38 @@ async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvab def _anthropic_sse_chunks() -> List[bytes]: events = [ - ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), - ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}), + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "m", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}, + ), ("content_block_stop", {"type": "content_block_stop", "index": 0}), - ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + ), ("message_stop", {"type": "message_stop"}), ] return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] @@ -2259,7 +2291,13 @@ async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthrop assert "hello [MASKED]" in raw assert "hello world" not in raw assert raw.count("event: content_block_delta") == 1 - for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + for expected_event in ( + "message_start", + "content_block_start", + "content_block_stop", + "message_delta", + "message_stop", + ): assert f"event: {expected_event}" in raw @@ -2332,9 +2370,7 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( managed = UnifiedRecordingGuardrail( guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True ) - free = RecordingGuardrail( - guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True - ) + free = RecordingGuardrail(guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True) monkeypatch.setattr(litellm, "callbacks", [managed, free]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) From 25a6c2a2749ba802514fb3a5736b540c21a1e630 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 06:17:59 +0000 Subject: [PATCH 13/23] fix(cli): skip remote model cost map fetch in lite CLI processes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/get_model_cost_map.py | 14 ++++- .../test_get_model_cost_map.py | 25 +++++++++ .../proxy/client/cli/test_global_options.py | 56 ++++++++++++++++++- 3 files changed, 91 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index cdc4810ff04..7d3b67e58c8 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -2,6 +2,7 @@ Pulls the cost + context window + provider route for known models from https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json This can be disabled by setting the LITELLM_LOCAL_MODEL_COST_MAP environment variable to True. +The ``lite`` and ``litellm-proxy`` CLI entry points also use the bundled map without fetching. ``` export LITELLM_LOCAL_MODEL_COST_MAP=True @@ -13,11 +14,13 @@ import hashlib import json import os import random +import sys import time from collections.abc import Awaitable, Callable from dataclasses import dataclass, replace from datetime import datetime, timezone from importlib.resources import files +from pathlib import Path from typing import Final, Protocol import httpx @@ -33,6 +36,12 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" +_CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) + + +def _is_cli_process() -> bool: + return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES + # Reserved top-level keys that are not model entries. They must be excluded # from the model-count integrity check so a real upstream shrink can't be masked. @@ -531,7 +540,8 @@ def get_model_cost_map( """ Public entry point — returns the model cost map dict. - 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. + 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set or this is a ``lite`` / + ``litellm-proxy`` CLI process, uses the local backup only. 2. Otherwise fetches from ``url``, retrying transient HTTP errors (429/5xx/transport) with Retry-After-aware backoff, validates integrity, and falls back to the local backup on any failure. @@ -543,7 +553,7 @@ def get_model_cost_map( _cost_map_source_info.loaded_at = datetime.now(timezone.utc) # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. - if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": + if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 0495440c51c..d092b387259 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -6,6 +6,7 @@ count actual model entries, not reserved meta keys) and the extraction of the import json import os +import sys import pytest @@ -711,3 +712,27 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej assert source["etag"] is None assert source["source_revision"] == _bundled_blob_id() assert source["source_revision"] != git_blob_id(shrunk_body) + + +@pytest.mark.parametrize( + ("argv0", "request_count"), + [ + ("/some/venv/bin/lite", 0), + ("/some/venv/bin/lite.exe", 0), + ("/some/venv/bin/python", 1), + ], +) +def test_boot_load_skips_remote_fetch_for_cli_processes(monkeypatch, argv0, request_count): + monkeypatch.setattr(sys, "argv", [argv0, "--version"]) + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + + cost_map = get_model_cost_map(url=_URL, client=client) + + assert calls["count"] == request_count + assert cost_map + source = get_model_cost_map_source_info() + if request_count == 0: + assert source["source"] == "local" + else: + assert source["source"] == "remote" diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 0dd388919a5..02b1cdf4296 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -1,14 +1,16 @@ # stdlib imports import json import os +import shutil +import subprocess +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from unittest.mock import Mock, patch import pytest from click.testing import CliRunner - - import litellm.proxy.client.cli from litellm._version import version as litellm_version from litellm.proxy.client.cli import cli @@ -35,6 +37,56 @@ def test_cli_version_flag(cli_runner): assert "LiteLLM Proxy Server Version: 1.2.3" in result.output +def test_lite_version_does_not_fetch_model_cost_map(): + lite_path = shutil.which("lite") + if lite_path is None: + pytest.skip("lite executable is unavailable") + + requests = [] + + class _CostMapHandler(BaseHTTPRequestHandler): + def do_GET(self): + requests.append(self.path) + body = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format, *args): + return + + server = ThreadingHTTPServer(("127.0.0.1", 0), _CostMapHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + env = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} + source_root = str(Path(__file__).resolve().parents[5]) + env["PYTHONPATH"] = os.pathsep.join(filter(None, (source_root, env.get("PYTHONPATH")))) + env.update( + { + "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", + "LITELLM_PROXY_URL": "http://127.0.0.1:9", + } + ) + result = subprocess.run( + [lite_path, "--version"], + capture_output=True, + text=True, + timeout=120, + env=env, + ) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=10) + + assert result.returncode == 0 + assert "LiteLLM Proxy CLI Version" in result.stdout + assert requests == [] + + def test_cli_source_is_ascii_only(): """Non-ASCII output (emoji, box-drawing chars) raises UnicodeEncodeError on legacy Windows consoles (cp1252), so the whole CLI package must stay ASCII-only.""" From ffae447649e173214a2081b4890410a11dfdf085 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 06:30:50 +0000 Subject: [PATCH 14/23] test(cli): type cost map bypass regression handlers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/test_get_model_cost_map.py | 4 +++- tests/test_litellm/proxy/client/cli/test_global_options.py | 6 +++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index d092b387259..6ab3f9a21af 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -722,7 +722,9 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej ("/some/venv/bin/python", 1), ], ) -def test_boot_load_skips_remote_fetch_for_cli_processes(monkeypatch, argv0, request_count): +def test_boot_load_skips_remote_fetch_for_cli_processes( + monkeypatch: pytest.MonkeyPatch, argv0: str, request_count: int +) -> None: monkeypatch.setattr(sys, "argv", [argv0, "--version"]) monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 02b1cdf4296..6de56573ffb 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -42,10 +42,10 @@ def test_lite_version_does_not_fetch_model_cost_map(): if lite_path is None: pytest.skip("lite executable is unavailable") - requests = [] + requests: list[str] = [] class _CostMapHandler(BaseHTTPRequestHandler): - def do_GET(self): + def do_GET(self) -> None: requests.append(self.path) body = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' self.send_response(200) @@ -54,7 +54,7 @@ def test_lite_version_does_not_fetch_model_cost_map(): self.end_headers() self.wfile.write(body) - def log_message(self, format, *args): + def log_message(self, format: str, *args: object) -> None: return server = ThreadingHTTPServer(("127.0.0.1", 0), _CostMapHandler) From 0c0a1dd76f31b56059508dd674fee1764c83828f Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 06:34:36 +0000 Subject: [PATCH 15/23] test(cli): avoid mutable request tracking Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/client/cli/test_global_options.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 6de56573ffb..9ed9f91ab9f 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -37,16 +37,17 @@ def test_cli_version_flag(cli_runner): assert "LiteLLM Proxy Server Version: 1.2.3" in result.output -def test_lite_version_does_not_fetch_model_cost_map(): +def test_lite_version_does_not_fetch_model_cost_map(tmp_path: Path) -> None: lite_path = shutil.which("lite") if lite_path is None: pytest.skip("lite executable is unavailable") - requests: list[str] = [] + request_log: Path = tmp_path / "requests.log" class _CostMapHandler(BaseHTTPRequestHandler): def do_GET(self) -> None: - requests.append(self.path) + with request_log.open("a", encoding="utf-8") as log_file: + log_file.write(f"{self.path}\n") body = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' self.send_response(200) self.send_header("Content-Type", "application/json") @@ -84,7 +85,8 @@ def test_lite_version_does_not_fetch_model_cost_map(): assert result.returncode == 0 assert "LiteLLM Proxy CLI Version" in result.stdout - assert requests == [] + request_count: int = request_log.read_text(encoding="utf-8").count("\n") if request_log.exists() else 0 + assert request_count == 0 def test_cli_source_is_ascii_only(): From 1fc1aaabdaa3542cd2d410103ee01ebc94025961 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 16:00:35 +0000 Subject: [PATCH 16/23] test(cli): drop lite --version subprocess regression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/client/cli/test_global_options.py | 56 ------------------- 1 file changed, 56 deletions(-) diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 9ed9f91ab9f..b73d1acc6e3 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -1,10 +1,6 @@ # stdlib imports import json import os -import shutil -import subprocess -import threading -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from unittest.mock import Mock, patch @@ -37,58 +33,6 @@ def test_cli_version_flag(cli_runner): assert "LiteLLM Proxy Server Version: 1.2.3" in result.output -def test_lite_version_does_not_fetch_model_cost_map(tmp_path: Path) -> None: - lite_path = shutil.which("lite") - if lite_path is None: - pytest.skip("lite executable is unavailable") - - request_log: Path = tmp_path / "requests.log" - - class _CostMapHandler(BaseHTTPRequestHandler): - def do_GET(self) -> None: - with request_log.open("a", encoding="utf-8") as log_file: - log_file.write(f"{self.path}\n") - body = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def log_message(self, format: str, *args: object) -> None: - return - - server = ThreadingHTTPServer(("127.0.0.1", 0), _CostMapHandler) - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - env = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} - source_root = str(Path(__file__).resolve().parents[5]) - env["PYTHONPATH"] = os.pathsep.join(filter(None, (source_root, env.get("PYTHONPATH")))) - env.update( - { - "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", - "LITELLM_PROXY_URL": "http://127.0.0.1:9", - } - ) - result = subprocess.run( - [lite_path, "--version"], - capture_output=True, - text=True, - timeout=120, - env=env, - ) - finally: - server.shutdown() - server.server_close() - thread.join(timeout=10) - - assert result.returncode == 0 - assert "LiteLLM Proxy CLI Version" in result.stdout - request_count: int = request_log.read_text(encoding="utf-8").count("\n") if request_log.exists() else 0 - assert request_count == 0 - - def test_cli_source_is_ascii_only(): """Non-ASCII output (emoji, box-drawing chars) raises UnicodeEncodeError on legacy Windows consoles (cp1252), so the whole CLI package must stay ASCII-only.""" From 2ece8735384e063b88ca91cf9d2f8f6d2e94b30a Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 17:05:38 +0000 Subject: [PATCH 17/23] test(e2e): lite CLI never fetches the model cost map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/coverage_registry/other.yaml | 2 + tests/e2e/other/test_cli_cost_map_e2e.py | 109 +++++++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 tests/e2e/other/test_cli_cost_map_e2e.py diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 814ebae2e0b..8ac4fc0762c 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -2,6 +2,8 @@ # PROMOTION NOTE: the auth cluster (~14 cells) is a candidate to promote to its own module once stable. - {id: other.auth.master_key.valid_allows, module: other, tier: P0, area: auth, assertions: [valid_allows], source: "user_api_key_auth.py:1569-1588", rationale: "Master key authenticates; timing-safe compare"} - {id: other.auth.master_key.invalid_denied, module: other, tier: P0, area: auth, assertions: [invalid_denied], source: "user_api_key_auth.py:1580", rationale: "Invalid master key rejected"} +- {id: other.cli.model_cost_map.version_skips_fetch, module: other, tier: P1, area: cli, assertions: [version_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite version command succeeds without requesting the remote model-cost map"} +- {id: other.cli.model_cost_map.models_list_skips_fetch, module: other, tier: P1, area: cli, assertions: [models_list_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite models list command uses the proxy without requesting the remote model-cost map"} - {id: other.auth.llm_chat.missing_header_denied, module: other, tier: P0, area: auth, assertions: [missing_header_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Chat with no Authorization header is 401/403"} - {id: other.auth.llm_chat.invalid_bearer_denied, module: other, tier: P0, area: auth, assertions: [invalid_bearer_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Bearer invalid_token on chat is 401/403"} - {id: other.auth.llm_chat.no_bearer_prefix_denied, module: other, tier: P0, area: auth, assertions: [no_bearer_prefix_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Token without Bearer scheme on chat is 401/403"} diff --git a/tests/e2e/other/test_cli_cost_map_e2e.py b/tests/e2e/other/test_cli_cost_map_e2e.py new file mode 100644 index 00000000000..6022d5a0a4c --- /dev/null +++ b/tests/e2e/other/test_cli_cost_map_e2e.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import os +import shutil +import subprocess +import threading +from collections.abc import Mapping +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final + +import pytest +from e2e_config import MASTER_KEY, PROXY_BASE_URL +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + + +def _start_cost_map_server(request_log: Path) -> tuple[ThreadingHTTPServer, threading.Thread]: + class CostMapHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + with request_log.open("a", encoding="utf-8") as log_file: + log_file.write(f"{self.path}\n") + body: Final = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + return + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), CostMapHandler) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return server, thread + + +def _run_lite( + args: tuple[str, ...], + server: ThreadingHTTPServer, + env: Mapping[str, str], +) -> subprocess.CompletedProcess[str]: + lite_path: Final = shutil.which("lite") + assert lite_path is not None, "the installed lite executable is required for e2e coverage" + try: + return subprocess.run( + [lite_path, *args], + capture_output=True, + text=True, + timeout=60, + env=env, + ) + finally: + server.shutdown() + server.server_close() + + +def _request_count(request_log: Path) -> int: + return request_log.read_text(encoding="utf-8").count("\n") if request_log.exists() else 0 + + +class TestLiteCliCostMapFetch: + @pytest.mark.covers("other.cli.model_cost_map.version_skips_fetch") + def test_lite_version_makes_no_cost_map_request(self, tmp_path: Path) -> None: + request_log: Final = tmp_path / "requests.log" + server, thread = _start_cost_map_server(request_log) + source_root: Final = str(Path(__file__).resolve().parents[3]) + pythonpath: Final = os.pathsep.join(filter(None, (source_root, os.environ.get("PYTHONPATH")))) + env: Final = { + **{key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}, + "PYTHONPATH": pythonpath, + "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", + "LITELLM_PROXY_URL": PROXY_BASE_URL, + } + try: + result: Final = _run_lite(("--version",), server, env) + finally: + thread.join(timeout=10) + + assert result.returncode == 0 + assert "LiteLLM Proxy CLI Version" in result.stdout + assert _request_count(request_log) == 0 + + @pytest.mark.covers("other.cli.model_cost_map.models_list_skips_fetch") + def test_lite_models_list_uses_proxy_not_cost_map(self, tmp_path: Path, proxy: ProxyClient) -> None: + model_names: Final = tuple(entry.model_name for entry in proxy.model_info()) + assert model_names + request_log: Final = tmp_path / "requests.log" + server, thread = _start_cost_map_server(request_log) + source_root: Final = str(Path(__file__).resolve().parents[3]) + pythonpath: Final = os.pathsep.join(filter(None, (source_root, os.environ.get("PYTHONPATH")))) + env: Final = { + **{key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}, + "PYTHONPATH": pythonpath, + "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", + "LITELLM_PROXY_URL": PROXY_BASE_URL, + "LITELLM_PROXY_API_KEY": MASTER_KEY, + } + try: + result: Final = _run_lite(("models", "list"), server, env) + finally: + thread.join(timeout=10) + + assert result.returncode == 0 + assert result.stdout.strip() + assert any(model_name in result.stdout for model_name in model_names) + assert _request_count(request_log) == 0 From 6e71b90a888b91fe3aa281fa14da58f258b9f056 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 17:07:34 +0000 Subject: [PATCH 18/23] test(e2e): dedupe lite env setup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/coverage_registry/other.yaml | 4 ++-- tests/e2e/other/test_cli_cost_map_e2e.py | 29 ++++++++++-------------- 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 8ac4fc0762c..a7ba0dbb8cb 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -2,8 +2,6 @@ # PROMOTION NOTE: the auth cluster (~14 cells) is a candidate to promote to its own module once stable. - {id: other.auth.master_key.valid_allows, module: other, tier: P0, area: auth, assertions: [valid_allows], source: "user_api_key_auth.py:1569-1588", rationale: "Master key authenticates; timing-safe compare"} - {id: other.auth.master_key.invalid_denied, module: other, tier: P0, area: auth, assertions: [invalid_denied], source: "user_api_key_auth.py:1580", rationale: "Invalid master key rejected"} -- {id: other.cli.model_cost_map.version_skips_fetch, module: other, tier: P1, area: cli, assertions: [version_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite version command succeeds without requesting the remote model-cost map"} -- {id: other.cli.model_cost_map.models_list_skips_fetch, module: other, tier: P1, area: cli, assertions: [models_list_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite models list command uses the proxy without requesting the remote model-cost map"} - {id: other.auth.llm_chat.missing_header_denied, module: other, tier: P0, area: auth, assertions: [missing_header_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Chat with no Authorization header is 401/403"} - {id: other.auth.llm_chat.invalid_bearer_denied, module: other, tier: P0, area: auth, assertions: [invalid_bearer_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Bearer invalid_token on chat is 401/403"} - {id: other.auth.llm_chat.no_bearer_prefix_denied, module: other, tier: P0, area: auth, assertions: [no_bearer_prefix_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Token without Bearer scheme on chat is 401/403"} @@ -50,3 +48,5 @@ - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} - {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} +- {id: other.cli.model_cost_map.version_skips_fetch, module: other, tier: P1, area: cli, assertions: [version_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite version command succeeds without requesting the remote model-cost map"} +- {id: other.cli.model_cost_map.models_list_skips_fetch, module: other, tier: P1, area: cli, assertions: [models_list_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite models list command uses the proxy without requesting the remote model-cost map"} diff --git a/tests/e2e/other/test_cli_cost_map_e2e.py b/tests/e2e/other/test_cli_cost_map_e2e.py index 6022d5a0a4c..62cd3d6c3f5 100644 --- a/tests/e2e/other/test_cli_cost_map_e2e.py +++ b/tests/e2e/other/test_cli_cost_map_e2e.py @@ -37,6 +37,16 @@ def _start_cost_map_server(request_log: Path) -> tuple[ThreadingHTTPServer, thre return server, thread +def _lite_env(server: ThreadingHTTPServer, api_key: str | None) -> dict[str, str]: + base_env: Final = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} + return { + **base_env, + "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", + "LITELLM_PROXY_URL": PROXY_BASE_URL, + **({"LITELLM_PROXY_API_KEY": api_key} if api_key is not None else {}), + } + + def _run_lite( args: tuple[str, ...], server: ThreadingHTTPServer, @@ -66,14 +76,7 @@ class TestLiteCliCostMapFetch: def test_lite_version_makes_no_cost_map_request(self, tmp_path: Path) -> None: request_log: Final = tmp_path / "requests.log" server, thread = _start_cost_map_server(request_log) - source_root: Final = str(Path(__file__).resolve().parents[3]) - pythonpath: Final = os.pathsep.join(filter(None, (source_root, os.environ.get("PYTHONPATH")))) - env: Final = { - **{key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}, - "PYTHONPATH": pythonpath, - "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", - "LITELLM_PROXY_URL": PROXY_BASE_URL, - } + env: Final = _lite_env(server, None) try: result: Final = _run_lite(("--version",), server, env) finally: @@ -89,15 +92,7 @@ class TestLiteCliCostMapFetch: assert model_names request_log: Final = tmp_path / "requests.log" server, thread = _start_cost_map_server(request_log) - source_root: Final = str(Path(__file__).resolve().parents[3]) - pythonpath: Final = os.pathsep.join(filter(None, (source_root, os.environ.get("PYTHONPATH")))) - env: Final = { - **{key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}, - "PYTHONPATH": pythonpath, - "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", - "LITELLM_PROXY_URL": PROXY_BASE_URL, - "LITELLM_PROXY_API_KEY": MASTER_KEY, - } + env: Final = _lite_env(server, MASTER_KEY) try: result: Final = _run_lite(("models", "list"), server, env) finally: From 34d2d010c3f0f08ab9c13700c7004d98377309fa Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 19:33:59 +0000 Subject: [PATCH 19/23] test(cli): drop lite e2e tests, the e2e runner does not install the package Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/coverage_registry/other.yaml | 2 - tests/e2e/other/test_cli_cost_map_e2e.py | 104 ----------------------- 2 files changed, 106 deletions(-) delete mode 100644 tests/e2e/other/test_cli_cost_map_e2e.py diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index a7ba0dbb8cb..814ebae2e0b 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -48,5 +48,3 @@ - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} - {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} -- {id: other.cli.model_cost_map.version_skips_fetch, module: other, tier: P1, area: cli, assertions: [version_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite version command succeeds without requesting the remote model-cost map"} -- {id: other.cli.model_cost_map.models_list_skips_fetch, module: other, tier: P1, area: cli, assertions: [models_list_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite models list command uses the proxy without requesting the remote model-cost map"} diff --git a/tests/e2e/other/test_cli_cost_map_e2e.py b/tests/e2e/other/test_cli_cost_map_e2e.py deleted file mode 100644 index 62cd3d6c3f5..00000000000 --- a/tests/e2e/other/test_cli_cost_map_e2e.py +++ /dev/null @@ -1,104 +0,0 @@ -from __future__ import annotations - -import os -import shutil -import subprocess -import threading -from collections.abc import Mapping -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path -from typing import Final - -import pytest -from e2e_config import MASTER_KEY, PROXY_BASE_URL -from proxy_client import ProxyClient - -pytestmark = pytest.mark.e2e - - -def _start_cost_map_server(request_log: Path) -> tuple[ThreadingHTTPServer, threading.Thread]: - class CostMapHandler(BaseHTTPRequestHandler): - def do_GET(self) -> None: - with request_log.open("a", encoding="utf-8") as log_file: - log_file.write(f"{self.path}\n") - body: Final = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def log_message(self, format: str, *args: object) -> None: - return - - server: Final = ThreadingHTTPServer(("127.0.0.1", 0), CostMapHandler) - thread: Final = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - return server, thread - - -def _lite_env(server: ThreadingHTTPServer, api_key: str | None) -> dict[str, str]: - base_env: Final = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} - return { - **base_env, - "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", - "LITELLM_PROXY_URL": PROXY_BASE_URL, - **({"LITELLM_PROXY_API_KEY": api_key} if api_key is not None else {}), - } - - -def _run_lite( - args: tuple[str, ...], - server: ThreadingHTTPServer, - env: Mapping[str, str], -) -> subprocess.CompletedProcess[str]: - lite_path: Final = shutil.which("lite") - assert lite_path is not None, "the installed lite executable is required for e2e coverage" - try: - return subprocess.run( - [lite_path, *args], - capture_output=True, - text=True, - timeout=60, - env=env, - ) - finally: - server.shutdown() - server.server_close() - - -def _request_count(request_log: Path) -> int: - return request_log.read_text(encoding="utf-8").count("\n") if request_log.exists() else 0 - - -class TestLiteCliCostMapFetch: - @pytest.mark.covers("other.cli.model_cost_map.version_skips_fetch") - def test_lite_version_makes_no_cost_map_request(self, tmp_path: Path) -> None: - request_log: Final = tmp_path / "requests.log" - server, thread = _start_cost_map_server(request_log) - env: Final = _lite_env(server, None) - try: - result: Final = _run_lite(("--version",), server, env) - finally: - thread.join(timeout=10) - - assert result.returncode == 0 - assert "LiteLLM Proxy CLI Version" in result.stdout - assert _request_count(request_log) == 0 - - @pytest.mark.covers("other.cli.model_cost_map.models_list_skips_fetch") - def test_lite_models_list_uses_proxy_not_cost_map(self, tmp_path: Path, proxy: ProxyClient) -> None: - model_names: Final = tuple(entry.model_name for entry in proxy.model_info()) - assert model_names - request_log: Final = tmp_path / "requests.log" - server, thread = _start_cost_map_server(request_log) - env: Final = _lite_env(server, MASTER_KEY) - try: - result: Final = _run_lite(("models", "list"), server, env) - finally: - thread.join(timeout=10) - - assert result.returncode == 0 - assert result.stdout.strip() - assert any(model_name in result.stdout for model_name in model_names) - assert _request_count(request_log) == 0 From 8bb6d8c120e2c116a719495ecf2401a2bab9dd5a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:34:08 -0700 Subject: [PATCH 20/23] chore(lint): bring ANN202 and BLE001 back under the strict-rule budget Annotate the return types of dispatch_async and transform_then_dispatch in llm_http_handler and _send_batch in azure_sentinel, and mark four legitimate broad catches with the repo's noqa convention, so the promote PR's lint job passes the strict gate again. Supersedes #40328. --- litellm/integrations/azure_sentinel/azure_sentinel.py | 4 +++- litellm/litellm_core_utils/get_model_cost_map.py | 2 +- litellm/litellm_core_utils/llm_cost_calc/utils.py | 2 +- litellm/llms/custom_httpx/llm_http_handler.py | 4 ++-- litellm/proxy/_experimental/mcp_server/server.py | 2 +- litellm/proxy/management_endpoints/cost_tracking_settings.py | 2 +- 6 files changed, 9 insertions(+), 7 deletions(-) diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index eba6c862f7a..db5f790615f 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -21,6 +21,8 @@ from types import MappingProxyType from typing import Final, TypeVar from urllib.parse import urlparse +import httpx + from litellm._logging import verbose_logger from litellm.integrations.batch_utils import ( BatchSendCancelled, @@ -418,7 +420,7 @@ class AzureSentinelLogger(CustomBatchLogger): "Content-Type": "application/json", } - async def _send_batch(batch: Sequence[_QueuedPayload]): + async def _send_batch(batch: Sequence[_QueuedPayload]) -> httpx.Response: body: Final = safe_dumps(batch) return await self.async_httpx_client.post( url=api_endpoint, diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index f81ddbfee2e..37375f25e46 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -585,7 +585,7 @@ def _retry_remote_fetch_in_background( _cost_map_source_info.fallback_reason = None _cost_map_source_info.loaded_at = datetime.now(timezone.utc) adopt_model_cost_map(finalized) - except Exception as e: + except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 5675c59733d..e5977ca4156 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1444,7 +1444,7 @@ def get_billed_token_rates( return _custom_pricing_rates(custom_cost_per_token) try: model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: + except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates return None return _cost_map_billed_rates( model_info=model_info, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7587b963a38..84b03e1955a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -579,7 +579,7 @@ class BaseLLMHTTPHandler: data: dict[str, object], # mutable-ok: async_completion takes dict signed_headers: dict[str, object], # mutable-ok: async_completion takes dict signed_json_body: bytes | None, - ): + ) -> Coroutine[object, object, ModelResponse | CustomStreamWrapper]: async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None if stream is True: return self.acompletion_stream_function( @@ -626,7 +626,7 @@ class BaseLLMHTTPHandler: if acompletion is True and provider_config.uses_async_transform_request: - async def transform_then_dispatch(): + async def transform_then_dispatch() -> ModelResponse | CustomStreamWrapper: transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict "dict[str, object]", await provider_config.async_transform_request( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 9a52da0cab1..f8f1c563811 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1053,7 +1053,7 @@ if MCP_AVAILABLE: route="/mcp/call_tool", traceback_str=failure_traceback, ) - except Exception: + except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error verbose_logger.exception("Error logging failed MCP proxy tool call") raise if proxy_logging_obj is not None: diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 493f75008c3..dc0da63555f 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -643,7 +643,7 @@ async def estimate_cost( custom_cost_per_token=resolved.custom_cost_per_token, litellm_logging_obj=litellm_logging_obj, ) - except Exception as e: + except Exception as e: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriceable model raise HTTPException( status_code=404, detail={ From c4bd3763e87ededa07ac552bb5fab67691bd3d49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:17:19 -0700 Subject: [PATCH 21/23] test(proxy): type the background retrieval governance tests The two test methods, the policy_engine fixture, and the two inner stubs in TestBackgroundResponseRetrievalGovernance now carry full parameter and return annotations, closing the Greptile thread that 94f9230d13 left open. --- .../proxy/test_common_request_processing.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index c451a3b4cb0..f119926e4f4 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3,7 +3,7 @@ import copy import datetime import json from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Optional +from typing import AsyncGenerator, Callable, Final, Iterator, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -8303,7 +8303,7 @@ class TestBackgroundResponseRetrievalGovernance: GOVERNED_MODEL_ID = "deployment-governed" @pytest.fixture - def policy_engine(self): + def policy_engine(self) -> Iterator[None]: from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry @@ -8353,10 +8353,14 @@ class TestBackgroundResponseRetrievalGovernance: mock_request = MagicMock(spec=Request) mock_request.headers = {} - async def passthrough_add_litellm_data_to_request(data, **kwargs): + async def passthrough_add_litellm_data_to_request( + data: dict[str, object], **kwargs: object + ) -> dict[str, object]: return data - async def decrypting_pre_call_hook(user_api_key_dict, data, call_type): + async def decrypting_pre_call_hook( + user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str + ) -> dict[str, object]: if data.get("response_id") == client_facing_response_id: data["response_id"] = encoded_response_id return data @@ -8383,8 +8387,8 @@ class TestBackgroundResponseRetrievalGovernance: @pytest.mark.asyncio async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline( - self, policy_engine, monkeypatch - ): + self, policy_engine: None, monkeypatch: pytest.MonkeyPatch + ) -> None: data = await self._pre_call("aget_responses", monkeypatch) assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f") @@ -8397,8 +8401,8 @@ class TestBackgroundResponseRetrievalGovernance: @pytest.mark.asyncio async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id( - self, policy_engine, monkeypatch - ): + self, policy_engine: None, monkeypatch: pytest.MonkeyPatch + ) -> None: data = await self._pre_call("aresponses", monkeypatch) assert "_guardrail_pipelines" not in data["litellm_metadata"] From c5e93aff132685582679ed1ce7097e7e855630a7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:45:26 -0700 Subject: [PATCH 22/23] fix(proxy): warn at submit when a body-selected post_call policy is deferred --- litellm/proxy/utils.py | 17 ++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 31 +++++++++++++++++-- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1604dea1003..1f8a4d5ba33 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -622,6 +622,15 @@ def _defer_post_call_pipelines( response.id, ", ".join(tag_matched), ) + body_selected: Final = _body_selected_deferrals(data, deferred) + if body_selected: + verbose_proxy_logger.warning( + "Policy engine: background response %s matched post_call policies through the request body's policies " + "list at submit; retrieval carries no request body, so those policies do not govern the completed " + "response: %s", + response.id, + ", ".join(body_selected), + ) _withdraw_deferred_claims(data, deferred) @@ -638,6 +647,14 @@ def _tag_matched_deferrals( ) +def _body_selected_deferrals( + data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] +) -> tuple[str, ...]: + sources: Final = _policy_state_metadata(data).get("policy_sources") + attributed: Final = frozenset(sources) if isinstance(sources, dict) else frozenset() + return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed) + + def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: unsupported: Final = tuple( dict.fromkeys( diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 3bbdbbe2ccd..cc40706c67c 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -941,6 +941,7 @@ def _post_call_pipeline_data( "metadata": { "_guardrail_pipelines": [("response-governance", pipeline)], "_pipeline_managed_guardrails": {guardrail}, + "policy_sources": {"response-governance": "model:m"}, }, **extra, } @@ -1468,7 +1469,7 @@ def _output_passing_callbacks() -> list[CustomGuardrail]: def _claimed_post_call_pipeline_data( - *policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str = "model:m" + *policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str | None = "model:m" ): from litellm.proxy.policy_engine.policy_registry import get_policy_registry @@ -1491,7 +1492,7 @@ def _claimed_post_call_pipeline_data( "_pipeline_managed_guardrails": {"gr-post"}, "applied_policies": list(policy_names), "applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)], - "policy_sources": {policy_name: policy_source for policy_name in policy_names}, + "policy_sources": {policy_name: policy_source for policy_name in policy_names if policy_source is not None}, }, } @@ -1551,6 +1552,32 @@ async def test_pending_background_response_warns_when_the_deferred_policy_was_ma ] +@pytest.mark.asyncio +async def test_pending_background_response_warns_when_the_deferred_policy_came_from_the_request_body( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("body-governance", policy_source=None) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "policy_sources" not in data["metadata"] + assert _warnings(caplog) == [ + "Policy engine: background response resp_bg matched post_call policies through the request body's policies " + "list at submit; retrieval carries no request body, so those policies do not govern the completed " + "response: body-governance" + ] + + @pytest.mark.asyncio async def test_pending_background_response_matched_through_its_model_does_not_warn( proxy_logging: ProxyLogging, From 2c836f473c8ef6574ae7a04ffa8434f7715e7e2f Mon Sep 17 00:00:00 2001 From: "cursor[bot]" <206951365+cursor[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:14:24 -0700 Subject: [PATCH 23/23] test(ui): derive reasoning-effort assertion from the anthropic preset (#40456) Preset #40341 pointed the Anthropic family REASONING tier at claude-fable-5-1, but this test still hardcoded claude-opus-5, so the payload it saw no longer matched. Rebase the assertion on ANTHROPIC_PRESET.complexity_router_config.tier_model_configs so a preset refresh flows through instead of redding the suite on staging. Co-authored-by: Cursor Agent Co-authored-by: Krrish Dholakia --- .../src/components/add_model/add_auto_router_tab.test.tsx | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 014854ac712..f6e619c5e84 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -1046,9 +1046,7 @@ describe("AddAutoRouterTab", () => { await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]).toMatchObject({ complexity_router_config: { - tier_model_configs: { - REASONING: [{ model_name: "claude-opus-5", litellm_params: { reasoning_effort: "high" } }], - }, + tier_model_configs: ANTHROPIC_PRESET.complexity_router_config.tier_model_configs, }, }); });