From 8576cf74c2ef4eaa15f48eaaf99447f9d6fb394d Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Mon, 7 Sep 2026 14:09:37 +0200 Subject: [PATCH 01/29] 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/29] 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/29] 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 cbe340a31ca81c40be088520ebbaaaca644977ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:38:18 -0700 Subject: [PATCH 04/29] feat(guardrails): deliver tool-call rewrites into buffered streams A post_call pipeline guardrail that rewrites a streamed tool call (its arguments or its name) now has that rewrite written back across the buffered chunks on chat, Responses, and Messages streams, so the client receives the rewritten tool call instead of the original. The chat handler rewrites the first fragment of each tool-call index and blanks the rest, the Responses handler syncs the function_call output items and their argument events, and the Messages handler rewrites the tool_use content_block_start and input_json_delta events in both dict and SSE-bytes chunks. The delivers_ended_stream_text_rewrites flag becomes delivers_ended_stream_rewrites, since the write-back now covers both text and tool calls, and the executor only discards a tool-call rewrite on translations without write-back or on a shape the translation refuses. --- .../chat/guardrail_translation/handler.py | 176 +++++++++++++++--- .../guardrail_translation/base_translation.py | 17 +- .../chat/guardrail_translation/handler.py | 96 +++++++++- .../guardrail_translation/handler.py | 111 ++++++++++- .../proxy/policy_engine/pipeline_executor.py | 22 +-- litellm/proxy/utils.py | 10 +- .../test_anthropic_guardrail_handler.py | 66 +++++++ .../test_openai_guardrail_handler.py | 73 ++++++++ ...test_openai_responses_guardrail_handler.py | 88 +++++++++ .../policy_engine/test_pipeline_executor.py | 22 ++- .../proxy_logging/test_guardrail_pipeline.py | 9 +- 11 files changed, 618 insertions(+), 72 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e486be12fe2..2049866b444 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,10 +13,11 @@ Pattern Overview: """ import json -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass from itertools import chain, repeat +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable from typing_extensions import ReadOnly, TypedDict, assert_never @@ -41,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_field, stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( @@ -153,6 +155,28 @@ class ExtractedInput: EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) +@dataclass(frozen=True, slots=True) +class _ToolCallShape: + name: str | None + arguments: str + + +_SSEEventRewriter = Callable[[Mapping[str, object]], Mapping[str, object] | None] + + +def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: + """The guardrail-visible shape of each tool call, whether the guardrail handed + back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts.""" + functions: Final = tuple(stream_item_field(tool_call, "function") for tool_call in tool_calls) + return tuple( + _ToolCallShape( + name=name if isinstance(name := stream_item_field(function, "name"), str) else None, + arguments=arguments if isinstance(arguments := stream_item_field(function, "arguments"), str) else "", + ) + for function in functions + ) + + class _AnthropicSSEDelta(TypedDict, total=False): type: ReadOnly[str] text: ReadOnly[str] @@ -170,7 +194,7 @@ class AnthropicMessagesHandler(BaseTranslation): them through guardrail rewrites; downstream provider handling is out of scope. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True def __init__(self): super().__init__() @@ -1050,6 +1074,7 @@ class AnthropicMessagesHandler(BaseTranslation): first_choice.message.tool_calls, ) string_so_far = first_choice.message.content + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ()) guardrail_inputs: Final = GenericGuardrailAPIInputs() if string_so_far: guardrail_inputs["texts"] = [string_so_far] @@ -1084,6 +1109,19 @@ class AnthropicMessagesHandler(BaseTranslation): and guardrailed_texts[0] != string_so_far ): self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0]) + if deliver_ended_stream_rewrites: + returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls") + self._write_ended_stream_tool_call_rewrites( + responses_so_far, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=_tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) + and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tool_calls_list or () + ), + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) else: verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far @@ -1212,38 +1250,120 @@ class AnthropicMessagesHandler(BaseTranslation): """Deliver an ended-stream guardrail text rewrite by rewriting the buffered chunks in place: the first ``text_delta`` carries the full rewritten text and every later one is blanked, leaving the surrounding - message and content-block framing untouched. Handles both chunk formats - this stream carries (parsed event dicts and raw SSE bytes).""" + message and content-block framing untouched.""" replacements: Final = chain((rewritten_text,), repeat("")) - for idx, item in enumerate(responses_so_far): - if isinstance(item, dict): - delta = item.get("delta") - if item.get("type") == "content_block_delta" and isinstance(delta, dict): - if delta.get("type") == "text_delta": - delta["text"] = next(replacements) - elif isinstance(item, (bytes, bytearray)): - responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer - AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements) - ) + + def rewrite_text_delta(event: Mapping[str, object]) -> Mapping[str, object] | None: + delta: Final = event.get("delta") + if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping): + return None + if delta.get("type") != "text_delta": + return None + return {**event, "delta": {**delta, "text": next(replacements)}} + + AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) + + @classmethod + def _write_ended_stream_tool_call_rewrites( + cls, + responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + *, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str, + ) -> None: + """Deliver ended-stream guardrail tool-call rewrites by rewriting the + buffered chunks in place: the rebuilt response lists tool calls in the + order of the stream's ``tool_use`` blocks, so the nth rewritten call lands + on the nth block, its first ``input_json_delta`` carrying the full rewritten + arguments, every later one blanked, and ``content_block_start`` carrying the + rewritten name. Blocks that do not line up with the rebuilt tool calls make + the rewrite undeliverable, so the pipeline executor discards it and releases + the original chunks.""" + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + block_indices: Final = tuple( + index + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + and isinstance(index := event.get("index"), int) + ) + if len(block_indices) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + rewrites_by_block: Final = MappingProxyType( + { + index: after + for index, before, after in zip(block_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + } + ) + argument_replacements: Final = MappingProxyType( + {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()} + ) + + def rewrite_tool_use(event: Mapping[str, object]) -> Mapping[str, object] | None: + index: Final = event.get("index") + if not isinstance(index, int) or index not in rewrites_by_block: + return None + match event.get("type"): + case "content_block_start": + block: Final = event.get("content_block") + name: Final = rewrites_by_block[index].name + if not isinstance(block, Mapping) or name is None: + return None + return {**event, "content_block": {**block, "name": name}} + case "content_block_delta": + delta: Final = event.get("delta") + if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta": + return None + return {**event, "delta": {**delta, "partial_json": next(argument_replacements[index])}} + case _: + return None + + cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use) @staticmethod - def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes: - """Rewrite every ``text_delta`` data line in one SSE chunk with the next - replacement text, leaving all other events and framing byte-identical.""" + def _rewrite_ended_stream_events( + responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + rewrite_event: _SSEEventRewriter, + ) -> None: + """Replace every buffered event ``rewrite_event`` returns a rewrite for, in + both chunk formats this stream carries (parsed event dicts and raw SSE + bytes), leaving every other event and the framing untouched.""" + rewritten_items: Final = tuple( + AnthropicMessagesHandler._rewrite_buffered_item(item, rewrite_event) for item in responses_so_far + ) + responses_so_far[:] = rewritten_items # rebind-ok: delivers the rewrites into the caller's buffer + + @staticmethod + def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object: + if isinstance(item, dict): + rewritten: Final = rewrite_event(_as_str_mapping(item)) + return item if rewritten is None else dict(rewritten) + if isinstance(item, (bytes, bytearray)): + return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event) + return item + + @staticmethod + def _rewrite_sse_events(sse_bytes: bytes, rewrite_event: _SSEEventRewriter) -> bytes: + """Rewrite the data lines of one SSE chunk that ``rewrite_event`` rewrites, + leaving all other events and framing byte-identical.""" try: decoded: Final = sse_bytes.decode("utf-8") except UnicodeDecodeError: return sse_bytes return "\n\n".join( - AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n") + "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, rewrite_event) for line in block.split("\n")) + for block in decoded.split("\n\n") ).encode("utf-8") @staticmethod - def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str: - return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n")) - - @staticmethod - def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str: + def _rewrite_sse_line(line: str, rewrite_event: _SSEEventRewriter) -> str: if not line.startswith("data:"): return line try: @@ -1252,14 +1372,10 @@ class AnthropicMessagesHandler(BaseTranslation): ) except json.JSONDecodeError: return line - if not isinstance(data, dict) or data.get("type") != "content_block_delta": + if not isinstance(data, dict): return line - delta: Final = data.get("delta") - if not isinstance(delta, dict) or delta.get("type") != "text_delta": - return line - return "data: " + json.dumps( - {**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts - ) + rewritten: Final = rewrite_event(_as_str_mapping(data)) + return line if rewritten is None else "data: " + json.dumps(rewritten) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index afd8e0f67f7..6d1a9ab1c3e 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -52,13 +52,14 @@ class StreamingScanKey: class BaseTranslation(ABC): - delivers_ended_stream_text_rewrites: ClassVar[bool] = False + delivers_ended_stream_rewrites: ClassVar[bool] = False """Whether ``process_output_streaming_response`` accepts ``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered) - stream, writes guardrail text rewrites back across ``responses_so_far`` so - a buffered pipeline can release rewritten chunks. Tool-call rewrites, and - text rewrites on every other translation, are undeliverable: the pipeline - executor discards them and releases the original chunks.""" + stream, writes guardrail text and tool-call rewrites back across + ``responses_so_far`` so a buffered pipeline can release rewritten chunks, + raising ``UndeliverableStreamRewrite`` for a shape it cannot place. Rewrites + on every other translation are undeliverable: the pipeline executor + discards them and releases the original chunks.""" @staticmethod def transform_user_api_key_dict_to_metadata( @@ -175,9 +176,9 @@ class BaseTranslation(ABC): transformations (see ``StreamTransformSink``); base handlers ignore it. ``deliver_ended_stream_rewrites`` is passed True only when the caller holds the whole buffered stream and the subclass declares - ``delivers_ended_stream_text_rewrites``: the handler then writes - guardrail text rewrites back across ``responses_so_far`` instead of - discarding them. + ``delivers_ended_stream_rewrites``: the handler then writes + guardrail text and tool-call rewrites back across ``responses_so_far`` + instead of discarding them. """ return responses_so_far diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 80292aef2cf..42c95ac3316 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -49,6 +49,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import coerce_stream_holdback_value, ) from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + ChatCompletionMessageToolCall, Choices, GenericGuardrailAPIInputs, ModelResponse, @@ -78,7 +80,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ @@ -610,13 +612,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): deliver_ended_stream_rewrites: bool, ) -> None: """Ended-stream path: rebuild the full response, run the non-streaming - output guardrail against it, and (when opted in) write any text rewrite - back across the buffered chunks.""" + output guardrail against it, and (when opted in) write any text or + tool-call rewrite back across the buffered chunks.""" model_response: Final = cast( ModelResponse, stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), ) pre_guardrail_texts: Final = self._string_choice_contents(model_response) + pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response) await self.process_output_response( response=model_response, guardrail_to_apply=guardrail_to_apply, @@ -624,13 +627,21 @@ class OpenAIChatCompletionsHandler(BaseTranslation): user_api_key_dict=user_api_key_dict, request_data=request_data, ) - if deliver_ended_stream_rewrites: - await self._write_ended_stream_text_rewrites( - responses_so_far=responses_so_far, - guardrailed_response=model_response, - pre_guardrail_texts=pre_guardrail_texts, - guardrail_name=guardrail_to_apply.guardrail_name or "unknown", - ) + if not deliver_ended_stream_rewrites: + return + guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown" + await self._write_ended_stream_text_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_texts=pre_guardrail_texts, + guardrail_name=guardrail_name, + ) + self._write_ended_stream_tool_call_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_name, + ) def build_stream_error_items( self, @@ -1043,6 +1054,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists ) + @staticmethod + def _function_tool_call_shapes(response: "ModelResponse") -> tuple[tuple[str | None, str], ...]: + return tuple( + (tool_call.function.name, tool_call.function.arguments) + for choice in response.choices + for tool_call in choice.message.tool_calls or () + if isinstance(tool_call, ChatCompletionMessageToolCall) + ) + + @staticmethod + def _function_tool_call_fragments( + responses_so_far: Sequence["ModelResponseStream"], + ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]: + """Group the stream's function tool-call fragments by their tool-call index, in + the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping + only the indices the builder keeps (an id and a name somewhere in the stream).""" + fragments: Final = tuple( + tool_call + for response in responses_so_far + for choice in response.choices + for tool_call in choice.delta.tool_calls or () + if isinstance(tool_call, ChatCompletionDeltaToolCall) + ) + identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id) + named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name) + return tuple( + tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named) + ) + + def _write_ended_stream_tool_call_rewrites( + self, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrailed_response: "ModelResponse", + pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail tool-call rewrites back across the buffered + chunks: the rewritten name and full arguments land in the tool call's first + fragment and the arguments of its later fragments are blanked, mirroring the + text write-back. A rewrite on a stream carrying more than one distinct choice + index, or whose fragments do not line up with the rebuilt tool calls, is + reported as undeliverable, so the pipeline executor discards it and releases + the original chunks.""" + post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response) + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) + fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) + if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + for before, (name, arguments), fragments in zip( + pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call + ): + if (name, arguments) == before: + continue + head, *tail = fragments + head.function.name = name + head.function.arguments = arguments + for fragment in tail: + fragment.function.arguments = "" + async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b0f79552bc5..447870fc0be 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -101,6 +101,18 @@ if TYPE_CHECKING: from litellm.types.llms.openai import ResponseInputParam +class _ToolCallShape(NamedTuple): + name: str | None + arguments: str + + +def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]: + return tuple( + _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", "")) + for tool_call in tool_calls + ) + + class ResponseOutputEnvelope(TypedDict, total=False): """Dict form of a Responses API response, as far as guardrail write-back reads it.""" @@ -340,7 +352,7 @@ class OpenAIResponsesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ @@ -754,6 +766,7 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -762,6 +775,12 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") + post_guardrail_tool_calls: Final = _tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) + else tool_calls_to_check + ) # Write guardrailed texts back into the output items in-place. # final_chunk is a reference into responses_so_far so this @@ -784,6 +803,13 @@ class OpenAIResponsesHandler(BaseTranslation): stream_events=responses_so_far[:-1], rewrites_by_position=rewrites_by_position, ) + self._deliver_ended_stream_tool_call_rewrites( + responses_so_far=responses_so_far, + outputs=outputs, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=post_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) return responses_so_far # ------------------------------------------------------------------ # @@ -894,6 +920,89 @@ class OpenAIResponsesHandler(BaseTranslation): continue OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) + def _deliver_ended_stream_tool_call_rewrites( + self, + responses_so_far: Sequence[object], + outputs: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail tool-call rewrites into the completed + envelope's ``function_call`` items and sync the earlier stream events, + keyed by ``output_index``. The guardrail sees the envelope's function + calls in output order, which is how a rewritten call finds its item; a + rewrite whose calls do not line up with the envelope is reported as + undeliverable, so the pipeline executor discards it and releases the + original events.""" + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + function_call_indices: Final = tuple( + output_idx + for output_idx, output_item in enumerate(outputs) + if stream_item_field(output_item, "type") == "function_call" + ) + if len(function_call_indices) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + rewrites_by_output_index: Final = MappingProxyType( + { + output_idx: after + for output_idx, before, after in zip( + function_call_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls + ) + if after != before + } + ) + for output_idx, rewrite in rewrites_by_output_index.items(): + self._write_function_call_item(outputs[output_idx], rewrite.name, rewrite.arguments) + self._sync_stream_events_with_tool_call_rewrites( + stream_events=responses_so_far[:-1], + rewrites_by_output_index=rewrites_by_output_index, + ) + + def _sync_stream_events_with_tool_call_rewrites( + self, + stream_events: Sequence[object], + rewrites_by_output_index: Mapping[int, _ToolCallShape], + ) -> None: + """Sync pre-completion function-call events with the rewritten completed + response: the first ``function_call_arguments.delta`` for a rewritten call + carries the full rewritten arguments and the rest are blanked, while + ``function_call_arguments.done`` and ``output_item.done`` carry the full + rewritten arguments and ``output_item.added`` / ``output_item.done`` the + rewritten name, so every event a client may read agrees with the + rewritten ``response.completed`` payload.""" + delta_replacements: Final = MappingProxyType( + {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_output_index.items()} + ) + for event in stream_events: + output_index = stream_item_field(event, "output_index") + if not isinstance(output_index, int) or output_index not in rewrites_by_output_index: + continue + rewrite = rewrites_by_output_index[output_index] + match stream_item_field(event, "type"): + case "response.function_call_arguments.delta": + self._write_event_field(event, "delta", next(delta_replacements[output_index])) + case "response.function_call_arguments.done": + self._write_event_field(event, "arguments", rewrite.arguments) + case "response.output_item.added": + self._write_function_call_item(stream_item_field(event, "item"), rewrite.name, None) + case "response.output_item.done": + self._write_function_call_item(stream_item_field(event, "item"), rewrite.name, rewrite.arguments) + case _: + pass + + @staticmethod + def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None: + if not (isinstance(item, dict) or hasattr(item, "get")): + return + if name is not None: + OpenAIResponsesHandler._write_event_field(item, "name", name) + if arguments is not None: + OpenAIResponsesHandler._write_event_field(item, "arguments", arguments) + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 9bc10949e9f..3e112bf8e67 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -85,10 +85,10 @@ def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT: class _StreamRewriteObserver(CustomGuardrail): """Stand-in handed to the endpoint translation in place of a streaming pipeline step's guardrail. It records whether the guardrail returned different output than it was given, - which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text - rewrites are deliverable on translations that write them back across the buffered chunks - (``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any - other translation are discarded by the executor, which releases the original chunks. + which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and + tool-call rewrites are deliverable on translations that write them back across the + buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation + are discarded by the executor, which releases the original chunks. The inner guardrail's ``apply_guardrail`` already records the guardrail information and span, so the observer's stays out of ``log_guardrail_information``.""" @@ -290,13 +290,13 @@ class PipelineExecutor: litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> None: """Run one streaming post_call step through the endpoint translation, delivering - text rewrites on translations that support ended-stream write-back. A rewrite that - cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation - without write-back, or one the translation refused with - ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the - originals and the step passes, so the client gets the stream the merge base sent.""" + text and tool-call rewrites on translations that support ended-stream write-back. A + rewrite that cannot reach the client yet (one on a translation without write-back, or + one the translation refused with ``UndeliverableStreamRewrite``) is discarded: the + buffered chunks go back to the originals and the step passes, so the client gets the + stream the merge base sent.""" observer: Final = _StreamRewriteObserver(callback) - deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites + deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites originals: Final = copy.deepcopy(streaming_chunks) try: if deliver_rewrites: @@ -319,7 +319,7 @@ class PipelineExecutor: except UndeliverableStreamRewrite: _release_original_chunks(step.guardrail, streaming_chunks, originals) else: - if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): + if not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls): _release_original_chunks(step.guardrail, streaming_chunks, originals) if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b36061c6be..76b9b4b44ed 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3445,11 +3445,11 @@ class ProxyLogging: assembled output through the endpoint guardrail translation, the same machinery flat post_call guardrails use at end of stream. An allow releases the buffered chunks: verbatim when no guardrail rewrote the - output, rewritten in place when one rewrote text and the translation - delivers ended-stream rewrites (later steps then re-scan the rewritten - chunks, so rewrites chain). A rewrite the translation cannot deliver - yet (a tool-call rewrite, or a text rewrite on a route without - write-back) is discarded by the executor and the original chunks are + output, rewritten in place when one rewrote text or a tool call and the + translation delivers ended-stream rewrites (later steps then re-scan the + rewritten chunks, so rewrites chain). A rewrite the translation cannot + deliver yet (one on a route without write-back, or a shape the route + refuses) is discarded by the executor and the original chunks are released, as is a buffered shape no translation resolves; a block or modify_response terminates with the translation's block chunks or the raised error. diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index bf40f781fa3..bd2dff9b33c 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -315,6 +315,72 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert "event: message_start" in raw and "event: message_stop" in raw assert '"stop_reason": "end_turn"' in raw + @staticmethod + def _ended_tool_use_sse_chunks() -> list: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "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] + + @staticmethod + def _argument_masking_guardrail() -> CustomGuardrail: + class MaskArguments(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call.function.arguments = '{"fruit": "[MASKED]"}' + return inputs + + return MaskArguments(guardrail_name="test") + + @staticmethod + def _partial_jsons(chunks: list) -> list: + return [ + json.loads(line[len("data:") :].strip())["delta"]["partial_json"] + for chunk in chunks + for line in chunk.decode().split("\n") + if line.startswith("data:") and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta" + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_use_input_back_into_sse_chunks(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', "", ""] + raw = b"".join(chunks).decode() + assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw + assert '"stop_reason": "tool_use"' in raw + assert "persim" not in raw + + @pytest.mark.asyncio + async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + original = [bytes(chunk) for chunk in chunks] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert chunks == original + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index aff0530ee8b..be168ec83e6 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1113,6 +1113,79 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert chunks[1].choices[0].delta.content in (None, "") assert chunks[1].choices[0].finish_reason == "stop" + @staticmethod + def _ended_tool_call_stream_chunks() -> list: + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None): + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(tool_calls=[tool_call] if tool_call else None), + finish_reason=finish_reason, + ) + ], + ) + + def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None): + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + return [ + chunk(fragment("", name="lookup_fruit", call_id="call_1")), + chunk(fragment('{"fruit":')), + chunk(fragment(' "persimmon"}')), + chunk(None, finish_reason="tool_calls"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_call_arguments_back_into_chunks(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_tool_call_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] + assert [fragment[0].function.arguments for fragment in fragments] == ['{"fruit": "PERSIMMON"}', "", ""] + assert fragments[0][0].function.name == "lookup_fruit" + assert fragments[0][0].id == "call_1" + assert chunks[3].choices[0].delta.tool_calls is None + assert chunks[3].choices[0].finish_reason == "tool_calls" + + @pytest.mark.asyncio + async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] + assert [fragment[0].function.arguments for fragment in fragments] == ["", '{"fruit":', ' "persimmon"}'] + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 33c8c97fea7..a024fc7ff81 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1195,6 +1195,94 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + @staticmethod + def _ended_function_call_stream_events() -> List[dict]: + def item(arguments: str, status: str) -> dict: + return { + "type": "function_call", + "id": "fc_123", + "call_id": "call_123", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.done", + "item_id": "fc_123", + "output_index": 0, + "arguments": '{"fruit": "persimmon"}', + }, + {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "model": "gpt-4o", + "output": [item('{"fruit": "persimmon"}', "completed")], + "status": "completed", + }, + }, + ] + + @staticmethod + def _argument_masking_guardrail() -> CustomGuardrail: + class MaskArguments(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + tool_calls = [ + {**tool_call, "function": {**tool_call["function"], "arguments": '{"fruit": "[MASKED]"}'}} + for tool_call in inputs.get("tool_calls", []) + ] + return {**inputs, "tool_calls": tool_calls} + + return MaskArguments(guardrail_name="test-mask-arguments") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_function_call_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["item"]["arguments"] == "" + assert events[1]["delta"] == '{"fruit": "[MASKED]"}' + assert events[2]["delta"] == "" + assert events[3]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["response"]["output"][0]["name"] == "lookup_fruit" + + @pytest.mark.asyncio + async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + ) + + assert events[1]["delta"] == '{"fruit":' + assert events[3]["arguments"] == '{"fruit": "persimmon"}' + assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "persimmon"}' + @pytest.mark.asyncio @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"]) async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type): diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 54ef1f79f4d..ee75003db9f 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -924,7 +924,7 @@ class _TextReturningGuardrail(CustomGuardrail): class _TextTranslation: - delivers_ended_stream_text_rewrites = False + delivers_ended_stream_rewrites = False def __init__(self): self.seen_guardrail_names = [] @@ -946,7 +946,7 @@ class _WritingTranslation: """Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the chat/Responses/Messages handlers do on an ended stream.""" - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True async def process_output_streaming_response( self, @@ -970,7 +970,7 @@ class _WritingTranslation: class _RefusingTranslation: - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True async def process_output_streaming_response( self, @@ -1088,13 +1088,27 @@ async def test_streaming_step_delivers_text_rewrite_through_writing_translation( @pytest.mark.asyncio -async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_text(monkeypatch, caplog): +async def test_streaming_step_delivers_tool_call_rewrite_through_writing_translation(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) chunks = [_chunk()] with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_WritingTranslation(), chunks) + assert result.terminal_action == "allow" + assert chunks[0]["text"] == "hello [MASKED]" + assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "[MASKED]"}' + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_TextTranslation(), chunks) + _assert_passed_with_discard_warning(result, caplog) assert chunks == [_chunk()] 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..d27f504ba6d 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 @@ -1836,7 +1836,7 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: @pytest.mark.asyncio @pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) -async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_tool_call_rewrite( +async def test_streaming_iterator_hook_pipeline_delivers_runtime_tool_call_rewrite( proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error, caplog ): transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731 @@ -1855,9 +1855,12 @@ async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_to delivered.append(item) assert len(delivered) == 2 - assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}' + delivered_tool_call = delivered[0].choices[0].delta.tool_calls[0] + assert delivered_tool_call.function.arguments == '{"ssn": "[MASKED]"}' + assert delivered_tool_call.function.name == "lookup" + assert delivered_tool_call.id == "call_1" assert delivered[1].choices[0].finish_reason == "tool_calls" - assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + assert not any("discarded" in message for message in _warnings(caplog)) @pytest.mark.asyncio From e65c56876f9aae1f94120212a2092592792f6538 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:42:51 -0700 Subject: [PATCH 05/29] fix(guardrails): write tool-call rewrites into typed Responses envelopes The response.completed envelope carries its function_call output items as SDK objects without a get shim, so the write-back skipped them and the envelope still showed the original arguments after every stream event had been rewritten. Write the item whenever one is present, and cover the typed event shape the live proxy carries in the handler test. --- .../guardrail_translation/handler.py | 2 +- ...test_openai_responses_guardrail_handler.py | 46 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 447870fc0be..208ceb959e6 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -996,7 +996,7 @@ class OpenAIResponsesHandler(BaseTranslation): @staticmethod def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None: - if not (isinstance(item, dict) or hasattr(item, "get")): + if item is None: return if name is not None: OpenAIResponsesHandler._write_event_field(item, "name", name) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index a024fc7ff81..bdf7f83757c 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1222,6 +1222,7 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: "type": "response.completed", "response": { "id": "resp_123", + "created_at": 1, "model": "gpt-4o", "output": [item('{"fruit": "persimmon"}', "completed")], "status": "completed", @@ -1268,6 +1269,51 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' assert events[5]["response"]["output"][0]["name"] == "lookup_fruit" + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_typed_function_call_events(self): + from litellm.types.llms.openai import ( + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, + OutputItemAddedEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ) + + handler = OpenAIResponsesHandler() + typed_events: List[Any] = [ + model.model_validate(event) + for model, event in zip( + ( + OutputItemAddedEvent, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ), + self._ended_function_call_stream_events(), + ) + ] + completed_event = typed_events[5] + assert isinstance(completed_event, ResponseCompletedEvent) + assert isinstance(completed_event.response, ResponsesAPIResponse) + assert isinstance(completed_event.response.output[0], ResponseFunctionToolCall) + + await handler.process_output_streaming_response( + responses_so_far=typed_events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert typed_events[1].delta == '{"fruit": "[MASKED]"}' + assert typed_events[2].delta == "" + assert typed_events[3].arguments == '{"fruit": "[MASKED]"}' + assert typed_events[4].item.arguments == '{"fruit": "[MASKED]"}' + assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}' + assert completed_event.response.output[0].name == "lookup_fruit" + @pytest.mark.asyncio async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): handler = OpenAIResponsesHandler() 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 06/29] 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 07/29] 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 08/29] 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 09/29] 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 f59354d09b642bf50d888a82b5fac5b96d34169b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:22:29 -0700 Subject: [PATCH 10/29] refactor(guardrails): patch Anthropic SSE rewrites by field The ended-stream rewriters now describe the one field they change as a frozen _SSEFieldRewrite and a single applier builds the patched event, so the handler adds no mutable-collection constructions over staging's total --- .../chat/guardrail_translation/handler.py | 42 +++++++++++++------ 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 2049866b444..cbbccac17f3 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,7 +13,7 @@ Pattern Overview: """ import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass from itertools import chain, repeat @@ -161,7 +161,25 @@ class _ToolCallShape: arguments: str -_SSEEventRewriter = Callable[[Mapping[str, object]], Mapping[str, object] | None] +@dataclass(frozen=True, slots=True) +class _SSEFieldRewrite: + """One field of one nested section of a buffered SSE event, rewritten.""" + + section: str + field: str + value: object + + +class _SSEEventRewriter(Protocol): + def __call__(self, event: Mapping[str, object]) -> _SSEFieldRewrite | None: ... + + +def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewriter) -> Mapping[str, object]: + rewrite: Final = rewrite_event(event) + section: Final = None if rewrite is None else event.get(rewrite.section) + if rewrite is None or not isinstance(section, Mapping): + return event + return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} # mutable-ok: json.dumps needs a dict def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: @@ -1253,13 +1271,13 @@ class AnthropicMessagesHandler(BaseTranslation): message and content-block framing untouched.""" replacements: Final = chain((rewritten_text,), repeat("")) - def rewrite_text_delta(event: Mapping[str, object]) -> Mapping[str, object] | None: + def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None: delta: Final = event.get("delta") if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping): return None if delta.get("type") != "text_delta": return None - return {**event, "delta": {**delta, "text": next(replacements)}} + return _SSEFieldRewrite("delta", "text", next(replacements)) AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) @@ -1306,22 +1324,21 @@ class AnthropicMessagesHandler(BaseTranslation): {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()} ) - def rewrite_tool_use(event: Mapping[str, object]) -> Mapping[str, object] | None: + def rewrite_tool_use(event: Mapping[str, object]) -> _SSEFieldRewrite | None: index: Final = event.get("index") if not isinstance(index, int) or index not in rewrites_by_block: return None match event.get("type"): case "content_block_start": - block: Final = event.get("content_block") name: Final = rewrites_by_block[index].name - if not isinstance(block, Mapping) or name is None: + if name is None: return None - return {**event, "content_block": {**block, "name": name}} + return _SSEFieldRewrite("content_block", "name", name) case "content_block_delta": delta: Final = event.get("delta") if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta": return None - return {**event, "delta": {**delta, "partial_json": next(argument_replacements[index])}} + return _SSEFieldRewrite("delta", "partial_json", next(argument_replacements[index])) case _: return None @@ -1343,8 +1360,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object: if isinstance(item, dict): - rewritten: Final = rewrite_event(_as_str_mapping(item)) - return item if rewritten is None else dict(rewritten) + return _rewritten_event(_as_str_mapping(item), rewrite_event) if isinstance(item, (bytes, bytearray)): return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event) return item @@ -1374,8 +1390,8 @@ class AnthropicMessagesHandler(BaseTranslation): return line if not isinstance(data, dict): return line - rewritten: Final = rewrite_event(_as_str_mapping(data)) - return line if rewritten is None else "data: " + json.dumps(rewritten) + rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event) + return line if rewritten is data else "data: " + json.dumps(rewritten) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) From 89e11949c834d0b466375b0442b27304946568a9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:22:30 -0700 Subject: [PATCH 11/29] fix(guardrails): key Responses stream tool-call rewrites by call_id Bridged Responses streams give reasoning and message items output_index 0 and start function calls at 1, so keying rewrites by output_index rewrote the wrong items. Rewrites now follow each function call's call_id through the buffered item and argument events, refuse when an event cannot be resolved to a rewritten call, and the refusal branches on all three handlers get regression tests --- .../guardrail_translation/handler.py | 132 ++++++++++++------ .../test_anthropic_guardrail_handler.py | 25 ++++ .../test_openai_guardrail_handler.py | 56 ++++++++ ...test_openai_responses_guardrail_handler.py | 90 ++++++++++++ 4 files changed, 259 insertions(+), 44 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 208ceb959e6..280a8670c36 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -140,6 +140,10 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( ) +_FUNCTION_CALL_ARGUMENT_EVENT_TYPES: Final = frozenset( + {"response.function_call_arguments.delta", "response.function_call_arguments.done"} +) +_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( {"function_call_output": "output", "message": "content"} ) @@ -930,70 +934,110 @@ class OpenAIResponsesHandler(BaseTranslation): ) -> None: """Write ended-stream guardrail tool-call rewrites into the completed envelope's ``function_call`` items and sync the earlier stream events, - keyed by ``output_index``. The guardrail sees the envelope's function - calls in output order, which is how a rewritten call finds its item; a - rewrite whose calls do not line up with the envelope is reported as - undeliverable, so the pipeline executor discards it and releases the - original events.""" + keyed by ``call_id``. The guardrail sees the envelope's function calls + in output order, which is how a rewritten call finds its ``call_id``; + the stream events find their call through the ``call_id`` on + ``output_item`` events and the ``item_id`` on argument events, since an + event's ``output_index`` need not match the envelope's (the chat bridge + numbers tool calls from 1 while the envelope lists them after the + message). A rewrite whose calls do not line up with the envelope, or + whose events cannot be found, is reported as undeliverable, so the + pipeline executor discards it and releases the original events.""" if post_guardrail_tool_calls == pre_guardrail_tool_calls: return - function_call_indices: Final = tuple( - output_idx - for output_idx, output_item in enumerate(outputs) - if stream_item_field(output_item, "type") == "function_call" + function_call_items: Final = tuple( + output_item for output_item in outputs if stream_item_field(output_item, "type") == "function_call" ) - if len(function_call_indices) != len(post_guardrail_tool_calls): - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - - raise UndeliverableStreamRewrite(guardrail_name) - rewrites_by_output_index: Final = MappingProxyType( + call_ids: Final = tuple( + call_id + for output_item in function_call_items + if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id + ) + stream_events: Final = responses_so_far[:-1] + call_id_by_item_id: Final = self._function_call_ids_by_item_id(stream_events) + event_call_ids: Final = tuple( + self._function_call_event_call_id(event, call_id_by_item_id) for event in stream_events + ) + rewrites_by_call_id: Final = MappingProxyType( { - output_idx: after - for output_idx, before, after in zip( - function_call_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls - ) + call_id: after + for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls) if after != before } ) - for output_idx, rewrite in rewrites_by_output_index.items(): - self._write_function_call_item(outputs[output_idx], rewrite.name, rewrite.arguments) - self._sync_stream_events_with_tool_call_rewrites( - stream_events=responses_so_far[:-1], - rewrites_by_output_index=rewrites_by_output_index, + unresolved_argument_event: Final = any( + call_id is None and stream_item_field(event, "type") in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES + for event, call_id in zip(stream_events, event_call_ids) ) + if ( + len(call_ids) != len(function_call_items) + or len(frozenset(call_ids)) != len(call_ids) + or len(call_ids) != len(post_guardrail_tool_calls) + or unresolved_argument_event + or not rewrites_by_call_id.keys() <= frozenset(event_call_ids) + ): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - def _sync_stream_events_with_tool_call_rewrites( - self, - stream_events: Sequence[object], - rewrites_by_output_index: Mapping[int, _ToolCallShape], - ) -> None: - """Sync pre-completion function-call events with the rewritten completed - response: the first ``function_call_arguments.delta`` for a rewritten call - carries the full rewritten arguments and the rest are blanked, while - ``function_call_arguments.done`` and ``output_item.done`` carry the full - rewritten arguments and ``output_item.added`` / ``output_item.done`` the - rewritten name, so every event a client may read agrees with the - rewritten ``response.completed`` payload.""" + raise UndeliverableStreamRewrite(guardrail_name) + for output_item, rewrite in ( + (output_item, rewrites_by_call_id[call_id]) + for output_item, call_id in zip(function_call_items, call_ids) + if call_id in rewrites_by_call_id + ): + self._write_function_call_item(output_item, rewrite.name, rewrite.arguments) delta_replacements: Final = MappingProxyType( - {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_output_index.items()} + {call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()} ) - for event in stream_events: - output_index = stream_item_field(event, "output_index") - if not isinstance(output_index, int) or output_index not in rewrites_by_output_index: + for event, call_id in zip(stream_events, event_call_ids): + if call_id not in rewrites_by_call_id: continue - rewrite = rewrites_by_output_index[output_index] match stream_item_field(event, "type"): case "response.function_call_arguments.delta": - self._write_event_field(event, "delta", next(delta_replacements[output_index])) + self._write_event_field(event, "delta", next(delta_replacements[call_id])) case "response.function_call_arguments.done": - self._write_event_field(event, "arguments", rewrite.arguments) + self._write_event_field(event, "arguments", rewrites_by_call_id[call_id].arguments) case "response.output_item.added": - self._write_function_call_item(stream_item_field(event, "item"), rewrite.name, None) + self._write_function_call_item( + stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None + ) case "response.output_item.done": - self._write_function_call_item(stream_item_field(event, "item"), rewrite.name, rewrite.arguments) + self._write_function_call_item( + stream_item_field(event, "item"), + rewrites_by_call_id[call_id].name, + rewrites_by_call_id[call_id].arguments, + ) case _: pass + @staticmethod + def _function_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]: + items: Final = tuple( + stream_item_field(event, "item") + for event in stream_events + if stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES + ) + return MappingProxyType( + { + item_id: call_id + for item in items + if stream_item_field(item, "type") == "function_call" + and isinstance(item_id := stream_item_field(item, "id"), str) + and isinstance(call_id := stream_item_field(item, "call_id"), str) + } + ) + + @staticmethod + def _function_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None: + event_type: Final = stream_item_field(event, "type") + if event_type in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES: + item_id: Final = stream_item_field(event, "item_id") + return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None + if event_type not in _OUTPUT_ITEM_EVENT_TYPES: + return None + item: Final = stream_item_field(event, "item") + call_id: Final = stream_item_field(item, "call_id") + return call_id if stream_item_field(item, "type") == "function_call" and isinstance(call_id, str) else None + @staticmethod def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None: if item is None: diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index bd2dff9b33c..64c243343d5 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -381,6 +381,31 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert chunks == original + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = AnthropicMessagesHandler() + server_tool_use = [ + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ] + tool_use = self._ended_tool_use_sse_chunks() + chunks = ( + tool_use[:1] + + [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in server_tool_use] + + [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]] + ) + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index be168ec83e6..c7011a6f00e 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1252,6 +1252,62 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + @staticmethod + def _two_choice_tool_call_stream_chunks() -> list: + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk( + choice_index: int, tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None + ) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=choice_index, + delta=Delta(tool_calls=[tool_call] if tool_call else None), + finish_reason=finish_reason, + ) + ], + ) + + def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None): + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + return [ + chunk(0, fragment("", name="lookup_fruit", call_id="call_1")), + chunk(1, fragment("", name="lookup_fruit", call_id="call_2")), + chunk(0, fragment('{"fruit": "persimmon"}')), + chunk(1, fragment('{"fruit": "durian"}')), + chunk(0, None, finish_reason="tool_calls"), + chunk(1, None, finish_reason="tool_calls"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index bdf7f83757c..fa1e969b277 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1314,6 +1314,96 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}' assert completed_event.response.output[0].name == "lookup_fruit" + @staticmethod + def _bridged_function_call_stream_events() -> List[dict]: + reasoning = {"type": "reasoning", "id": "rs_1", "summary": []} + text = {"type": "output_text", "text": "Looking that up", "annotations": []} + message = {"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", "content": [text]} + + def function_call(arguments: str, status: str) -> dict: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": dict(reasoning)}, + {"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning)}, + {"type": "response.output_item.added", "output_index": 0, "item": {**message, "status": "in_progress", "content": []}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": "Looking that up"}, + {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": [dict(text)]}}, + {"type": "response.output_item.added", "output_index": 1, "item": function_call("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.done", + "item_id": "fc_1", + "output_index": 1, + "arguments": '{"fruit": "persimmon"}', + }, + {"type": "response.output_item.done", "output_index": 1, "item": function_call('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_1", + "model": "claude-haiku-4-5", + "output": [ + dict(reasoning), + {**message, "content": [dict(text)]}, + function_call('{"fruit": "persimmon"}', "completed"), + ], + }, + }, + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_keys_bridged_function_call_events_by_call_id(self): + handler = OpenAIResponsesHandler() + events = self._bridged_function_call_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert events[6]["delta"] == '{"fruit": "[MASKED]"}' + assert events[7]["delta"] == "" + assert events[8]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["item"]["name"] == "lookup_fruit" + assert events[9]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[10]["response"]["output"][2]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[3]["delta"] == "Looking that up" + assert events[4]["item"]["content"][0]["text"] == "Looking that up" + assert events[10]["response"]["output"][1]["content"][0]["text"] == "Looking that up" + assert events[1]["item"] == {"type": "reasoning", "id": "rs_1", "summary": []} + + @pytest.mark.asyncio + @pytest.mark.parametrize("mismatch", ["orphan_call_id", "duplicate_call_id"]) + async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed(self, mismatch): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + envelope_item = events[5]["response"]["output"][0] + if mismatch == "orphan_call_id": + events[5]["response"]["output"] = [{**envelope_item, "call_id": "call_999"}] + else: + events[5]["response"]["output"] = [dict(envelope_item), dict(envelope_item)] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + @pytest.mark.asyncio async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): handler = OpenAIResponsesHandler() 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 12/29] 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 13/29] 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 14/29] 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 3ea65e761ee2f021e967c87202267cd4d70ff3f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:30:06 -0700 Subject: [PATCH 15/29] test(guardrails): cover Responses and Messages pipeline tool-call delivery --- .../proxy_logging/test_guardrail_pipeline.py | 104 ++++++++++++++++++ 1 file changed, 104 insertions(+) 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 f753d2ab690..bf71780ad7b 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 @@ -25,6 +25,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines @@ -2185,3 +2186,106 @@ async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_str assert result is not None assert seen["count"] == 1 assert seen["response"] == "hello " + + +def _mask_tool_call_arguments(inputs: Dict[str, Any]) -> Dict[str, Any]: + return { + "tool_calls": [ + { + "id": stream_item_field(tool_call, "id"), + "type": "function", + "function": { + "name": stream_item_field(stream_item_field(tool_call, "function"), "name"), + "arguments": '{"fruit": "[MASKED]"}', + }, + } + for tool_call in inputs.get("tool_calls", []) + ] + } + + +def _anthropic_tool_use_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": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "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] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_tool_use_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert '{\\"fruit\\": \\"[MASKED]\\"}' in raw + assert "persim" not in raw + assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw + assert '"stop_reason": "tool_use"' in raw + assert raw.count("event: content_block_delta") == 2 + + +def _responses_function_call_events() -> List[Dict[str, Any]]: + def item(arguments: str, status: str) -> Dict[str, Any]: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": ' "persimmon"}'}, + {"type": "response.function_call_arguments.done", "item_id": "fc_1", "output_index": 0, "arguments": '{"fruit": "persimmon"}'}, + {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": {"id": "resp_1", "created_at": 1, "model": "m", "output": [item('{"fruit": "persimmon"}', "completed")], "status": "completed"}, + }, + ] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_function_call_rewrite_on_responses_events( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"), + response=_async_chunk_iter(_responses_function_call_events()), + request_data=data, + ) + ] + + assert [event["type"] for event in delivered] == [event["type"] for event in _responses_function_call_events()] + assert [event["delta"] for event in delivered if event["type"] == "response.function_call_arguments.delta"] == ['{"fruit": "[MASKED]"}', ""] + assert delivered[3]["arguments"] == '{"fruit": "[MASKED]"}' + assert delivered[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert delivered[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' + assert "persimmon" not in json.dumps(delivered) From 9cb5d9b76c0c6b6e7a20e852e9673c9ea410c3f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:55:19 -0700 Subject: [PATCH 16/29] fix(proxy): gate streaming pipelines on a route-resolved guardrail translation The per-chunk hook skipped pipeline-managed guardrails whenever the request route was empty, while the gated stream could still fail to resolve a translation and release the buffered stream ungoverned. The gate now needs a translation resolved from the route, the iterator hook resolves it once and hands it to the gated stream, and the ungoverned release branch is gone. --- litellm/proxy/utils.py | 27 ++++------ .../proxy_logging/test_guardrail_pipeline.py | 53 +++++++++++++++---- 2 files changed, 54 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c38eb900a05..37d077d607c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -192,6 +192,7 @@ if TYPE_CHECKING: from prisma.types import HttpConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -559,7 +560,7 @@ def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool: - return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None + return resolve_endpoint_translation(user_api_key_dict, None) is not None def _stream_gated_guardrail_names( @@ -3435,12 +3436,16 @@ class ProxyLogging: ), ) - if post_call_pipelines: + pipeline_translation: Final = ( + resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None + ) + if pipeline_translation is not None: current_response = self._pipeline_gated_stream( response=current_response, user_api_key_dict=user_api_key_dict, request_data=request_data, pipelines=post_call_pipelines, + translation=pipeline_translation, ) try: @@ -3464,6 +3469,7 @@ class ProxyLogging: user_api_key_dict: UserAPIKeyAuth, request_data: dict, # mutable-ok: same request-payload shape the hooks mutate pipelines: "tuple[tuple[str, GuardrailPipeline], ...]", + translation: "tuple[str, BaseTranslation]", ) -> "AsyncGenerator[Any, None]": """ Execute post_call policy pipelines against a streamed response. @@ -3478,9 +3484,8 @@ class ProxyLogging: rewritten chunks, so rewrites chain). A rewrite the translation cannot deliver yet (one on a route without write-back, or a shape the route refuses) is discarded by the executor and the original chunks are - released, as is a buffered shape no translation resolves; a block or - modify_response terminates with the translation's block chunks or the - raised error. + released; a block or modify_response terminates with the translation's + block chunks or the raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: @@ -3488,17 +3493,7 @@ class ProxyLogging: if not buffered: return - resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0]) - if resolved is None: - verbose_proxy_logger.warning( - "Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; " - "the stream is released ungoverned by them: %s", - ", ".join(policy_name for policy_name, _pipeline in pipelines), - ) - for buffered_item in buffered: - yield buffered_item - return - call_type, endpoint_translation = resolved + call_type, endpoint_translation = translation for policy_name, pipeline in pipelines: result: PipelineExecutionResult = await PipelineExecutor.execute_steps( 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 bf71780ad7b..3b59bbac3fd 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 @@ -1969,28 +1969,61 @@ async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_anothe @pytest.mark.asyncio -async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvable_response_shape( +async def test_streaming_iterator_hook_skips_pipeline_and_warns_without_request_route( proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): seen: Dict[str, Any] = {} monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) - chunks = [object(), object()] - delivered: List[Any] = [] + chunks = _stream_chunks() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - async for item in proxy_logging.async_post_call_streaming_iterator_hook( - user_api_key_dict=make_user_api_key_auth(), - response=_async_chunk_iter(chunks), - request_data=data, - ): - delivered.append(item) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] assert len(delivered) == 2 assert seen.get("count") is None - assert any("response-governance" in message and "shape" in message for message in _warnings(caplog)) + assert any("response-governance" in message and "route None" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_runs_pipeline_managed_guardrail_without_request_route( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class UnifiedRecordingGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 + return None + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + monkeypatch.setattr( + litellm, + "callbacks", + [UnifiedRecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(), + ) + + assert result is not None + assert seen["gr-post"] == 1 def _anthropic_sse_chunks() -> List[bytes]: From 88d2d775534457b3397c9332ca968b7a8707a161 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:59:55 -0700 Subject: [PATCH 17/29] feat(hosted_vllm): add image edit support Register a HostedVLLMImageEditConfig so hosted_vllm/ deployments route POST /v1/images/edits to the vLLM-Omni OpenAI-compatible endpoint instead of failing with 'image edit is not supported for hosted_vllm' before any request is sent --- .../llms/hosted_vllm/image_edit/__init__.py | 9 ++ .../hosted_vllm/image_edit/transformation.py | 42 +++++++ litellm/utils.py | 4 + ...t_hosted_vllm_image_edit_transformation.py | 114 ++++++++++++++++++ 4 files changed, 169 insertions(+) create mode 100644 litellm/llms/hosted_vllm/image_edit/__init__.py create mode 100644 litellm/llms/hosted_vllm/image_edit/transformation.py create mode 100644 tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py diff --git a/litellm/llms/hosted_vllm/image_edit/__init__.py b/litellm/llms/hosted_vllm/image_edit/__init__.py new file mode 100644 index 00000000000..27e005e8a0d --- /dev/null +++ b/litellm/llms/hosted_vllm/image_edit/__init__.py @@ -0,0 +1,9 @@ +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + +from .transformation import HostedVLLMImageEditConfig + +__all__ = ("HostedVLLMImageEditConfig",) + + +def get_hosted_vllm_image_edit_config(model: str) -> BaseImageEditConfig: + return HostedVLLMImageEditConfig() diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py new file mode 100644 index 00000000000..bd27c962da8 --- /dev/null +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -0,0 +1,42 @@ +"""Image edits for Hosted VLLM (vLLM-Omni OpenAI-compatible /v1/images/edits).""" + +from typing import Final + +from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig +from litellm.secret_managers.main import get_secret_str + + +class HostedVLLMImageEditConfig(OpenAIImageEditConfig): + """ + vLLM-Omni images edits API follows the OpenAI multipart contract. + + https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/images_api/ + """ + + def validate_environment( + self, + headers: dict, # mutable-ok: BaseImageEditConfig contract + model: str, + api_key: str | None = None, + litellm_params: dict | None = None, # mutable-ok: BaseImageEditConfig contract + api_base: str | None = None, + ) -> dict: # mutable-ok: BaseImageEditConfig contract + resolved_key: Final = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" + return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: BaseImageEditConfig contract + ) -> str: + resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE") + if resolved_api_base is None: + raise ValueError( + "api_base not set for Hosted VLLM images edits API. " + "Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable" + ) + trimmed: Final = resolved_api_base.rstrip("/") + if trimmed.endswith("/v1"): + return f"{trimmed}/images/edits" + return f"{trimmed}/v1/images/edits" diff --git a/litellm/utils.py b/litellm/utils.py index 141ec323776..0f2b2f9042d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9239,6 +9239,10 @@ class ProviderConfigManager: from litellm.llms.openai.image_edit import get_openai_image_edit_config return get_openai_image_edit_config(model=model) + if LlmProviders.HOSTED_VLLM == provider: + from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config + + return get_hosted_vllm_image_edit_config(model=model) elif LlmProviders.AZURE == provider: from litellm.llms.azure.image_edit.transformation import ( AzureImageEditConfig, diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py new file mode 100644 index 00000000000..e1167b92553 --- /dev/null +++ b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py @@ -0,0 +1,114 @@ +"""Tests for hosted_vllm image edits (vLLM-Omni /v1/images/edits).""" + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config +from litellm.llms.hosted_vllm.image_edit.transformation import HostedVLLMImageEditConfig +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +PNG_BYTES = b"\x89PNG\r\n\x1a\nfakepng" +MODEL = "Qwen/Qwen-Image-Edit-2511" + + +@pytest.fixture(autouse=True) +def _clear_hosted_vllm_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + + +def test_provider_config_registration(): + config = ProviderConfigManager.get_provider_image_edit_config( + model=f"hosted_vllm/{MODEL}", + provider=LlmProviders.HOSTED_VLLM, + ) + + assert isinstance(config, HostedVLLMImageEditConfig) + assert isinstance(get_hosted_vllm_image_edit_config(MODEL), HostedVLLMImageEditConfig) + + +@pytest.mark.parametrize( + "api_base", + ["http://localhost:8091", "http://localhost:8091/", "http://localhost:8091/v1", "http://localhost:8091/v1/"], +) +def test_get_complete_url_appends_images_edits(api_base: str): + config = HostedVLLMImageEditConfig() + + assert ( + config.get_complete_url(model=MODEL, api_base=api_base, litellm_params={}) + == "http://localhost:8091/v1/images/edits" + ) + + +def test_get_complete_url_falls_back_to_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://vllm-omni:8000/v1") + config = HostedVLLMImageEditConfig() + + assert ( + config.get_complete_url(model=MODEL, api_base=None, litellm_params={}) + == "http://vllm-omni:8000/v1/images/edits" + ) + + +def test_get_complete_url_requires_api_base(): + config = HostedVLLMImageEditConfig() + + with pytest.raises(ValueError, match="api_base not set"): + config.get_complete_url(model=MODEL, api_base=None, litellm_params={}) + + +def test_validate_environment_defaults_to_fake_api_key(): + headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL) + + assert headers == {"Authorization": "Bearer fake-api-key"} + + +def test_validate_environment_uses_provided_api_key_and_keeps_headers(): + headers = HostedVLLMImageEditConfig().validate_environment( + headers={"X-Test": "1"}, + model=MODEL, + api_key="my-custom-key", + ) + + assert headers == {"X-Test": "1", "Authorization": "Bearer my-custom-key"} + + +def test_validate_environment_falls_back_to_env_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOSTED_VLLM_API_KEY", "env-key") + + headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL) + + assert headers["Authorization"] == "Bearer env-key" + + +def test_image_edit_posts_multipart_to_vllm_omni(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + response = litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + api_key="test-key", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + seed=42, + ) + + assert response.data + assert len(captured) == 1 + request = captured[0] + assert str(request.url) == "http://localhost:8091/v1/images/edits" + assert request.headers["authorization"] == "Bearer test-key" + assert request.headers["content-type"].startswith("multipart/form-data") + assert b'name="image[]"' in request.content + assert PNG_BYTES in request.content + assert f'name="model"\r\n\r\n{MODEL}'.encode() in request.content + assert b'name="prompt"\r\n\r\nadd a hat' in request.content + assert b'name="seed"\r\n\r\n42' in request.content 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 18/29] 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 260097b2dde7f12f77784892e96e9956b0b19cc1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:16:34 -0700 Subject: [PATCH 19/29] refactor(hosted_vllm): drop redundant image edit docstrings --- litellm/llms/hosted_vllm/image_edit/transformation.py | 8 -------- .../test_hosted_vllm_image_edit_transformation.py | 2 -- 2 files changed, 10 deletions(-) diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py index bd27c962da8..fa8c1dbfcfd 100644 --- a/litellm/llms/hosted_vllm/image_edit/transformation.py +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -1,5 +1,3 @@ -"""Image edits for Hosted VLLM (vLLM-Omni OpenAI-compatible /v1/images/edits).""" - from typing import Final from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig @@ -7,12 +5,6 @@ from litellm.secret_managers.main import get_secret_str class HostedVLLMImageEditConfig(OpenAIImageEditConfig): - """ - vLLM-Omni images edits API follows the OpenAI multipart contract. - - https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/images_api/ - """ - def validate_environment( self, headers: dict, # mutable-ok: BaseImageEditConfig contract diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py index e1167b92553..5646f3d9d13 100644 --- a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py @@ -1,5 +1,3 @@ -"""Tests for hosted_vllm image edits (vLLM-Omni /v1/images/edits).""" - import httpx import pytest From 113350756595b7954d5f830f8a5454895de12526 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:46:55 -0700 Subject: [PATCH 20/29] refactor: type the buffered stream rewrite helpers without Any --- .../llms/anthropic/chat/guardrail_translation/handler.py | 8 ++++---- .../responses/test_openai_responses_guardrail_handler.py | 3 ++- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index cbbccac17f3..6f966ae1a02 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,7 +13,7 @@ Pattern Overview: """ import json -from collections.abc import Mapping, Sequence +from collections.abc import Mapping, MutableSequence, Sequence from copy import deepcopy from dataclasses import dataclass from itertools import chain, repeat @@ -1262,7 +1262,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _write_ended_stream_text_rewrite( - responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place rewritten_text: str, ) -> None: """Deliver an ended-stream guardrail text rewrite by rewriting the @@ -1284,7 +1284,7 @@ class AnthropicMessagesHandler(BaseTranslation): @classmethod def _write_ended_stream_tool_call_rewrites( cls, - responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place *, pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], post_guardrail_tool_calls: tuple[_ToolCallShape, ...], @@ -1346,7 +1346,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _rewrite_ended_stream_events( - responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place rewrite_event: _SSEEventRewriter, ) -> None: """Replace every buffered event ``rewrite_event`` returns a rewrite for, in diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index fa1e969b277..6ed4ec6618f 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -17,6 +17,7 @@ from fastapi import HTTPException from openai.types.responses import ResponseFunctionToolCall from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, @@ -1238,7 +1239,7 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, + logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: tool_calls = [ {**tool_call, "function": {**tool_call["function"], "arguments": '{"fruit": "[MASKED]"}'}} From ddedb4867b8cedea4ac2223dbc636f813e8bf997 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:09:44 -0700 Subject: [PATCH 21/29] fix: discard a streamed rewrite that drops or adds a tool call A guardrail that removes or adds a tool call on an ended stream used to be silently ignored: every handler substitutes the original list on a count mismatch and the executor skipped its observer once the translation could deliver rewrites. The executor now tracks the count change on the observer and releases the original chunks with the discard warning on every translation, matching what the merge base did for any tool call rewrite --- .../proxy/policy_engine/pipeline_executor.py | 20 ++++-- .../policy_engine/test_pipeline_executor.py | 23 +++++- .../proxy_logging/test_guardrail_pipeline.py | 71 +++++++++++++++++++ 3 files changed, 108 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index c870bd7d2ec..a51468cc0fb 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -78,6 +78,10 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non return sent is not None and returned is not None and returned != sent +def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: + return sent is not None and returned is not None and len(returned) != len(sent) + + _GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object]) @@ -91,8 +95,9 @@ class _StreamRewriteObserver(CustomGuardrail): guardrail. It records whether the guardrail returned different output than it was given, which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and tool-call rewrites are deliverable on translations that write them back across the - buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation - are discarded by the executor, which releases the original chunks. + buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation, + and a rewrite that drops or adds a tool call on any translation, are discarded by the + executor, which releases the original chunks. The inner guardrail's ``apply_guardrail`` already records the guardrail information and span, so the observer's stays out of ``log_guardrail_information``.""" @@ -101,6 +106,7 @@ class _StreamRewriteObserver(CustomGuardrail): self.inner: Final = inner self.rewrote_texts = False self.rewrote_tool_calls = False + self.changed_tool_call_count = False def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() @@ -118,9 +124,11 @@ class _StreamRewriteObserver(CustomGuardrail): outputs: Final = await self.inner.apply_guardrail( inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj ) + returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls")) self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) - self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote( - sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls")) + self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes) + self.changed_tool_call_count = self.changed_tool_call_count or _changed_count( + sent_tool_shapes, returned_tool_shapes ) return outputs @@ -325,7 +333,9 @@ class PipelineExecutor: except UndeliverableStreamRewrite: _release_original_chunks(step.guardrail, streaming_chunks, originals) else: - if not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls): + if observer.changed_tool_call_count or ( + not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls) + ): _release_original_chunks(step.guardrail, streaming_chunks, originals) if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 5d042d3c1ad..16401ccdaad 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1106,7 +1106,8 @@ class _WritingTranslation: logging_obj=litellm_logging_obj, ) responses_so_far[0]["text"] = outputs["texts"][0] - responses_so_far[0]["tool_call"] = outputs["tool_calls"][0] + if len(outputs["tool_calls"]) == 1: + responses_so_far[0]["tool_call"] = outputs["tool_calls"][0] return responses_so_far @@ -1242,6 +1243,26 @@ async def test_streaming_step_delivers_tool_call_rewrite_through_writing_transla assert not any("discarded" in record.getMessage() for record in caplog.records) +class _ToolCallDroppingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": []} + + +@pytest.mark.asyncio +async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool_call(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_ToolCallDroppingGuardrail()]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_WritingTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + @pytest.mark.asyncio async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) 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 3b59bbac3fd..d46881d045b 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 @@ -2322,3 +2322,74 @@ async def test_streaming_iterator_hook_pipeline_delivers_function_call_rewrite_o assert delivered[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' assert delivered[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' assert "persimmon" not in json.dumps(delivered) + + +def _drop_tool_calls(inputs: Dict[str, Any]) -> Dict[str, Any]: + return {"tool_calls": []} + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_chat_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_tool_call_stream_chunks()), + request_data=data, + ) + ] + + assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}' + assert delivered[1].choices[0].finish_reason == "tool_calls" + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()), + request_data=data, + ) + ] + + assert delivered == _anthropic_tool_use_sse_chunks() + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_responses_events( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"), + response=_async_chunk_iter(_responses_function_call_events()), + request_data=data, + ) + ] + + assert delivered == _responses_function_call_events() + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) 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 22/29] 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 c66ae07e3bb97e875d88096458e5de80b351828c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:11:01 -0700 Subject: [PATCH 23/29] fix(hosted_vllm): reject image edit params vLLM-Omni ignores --- .../hosted_vllm/image_edit/transformation.py | 9 ++++ litellm/utils.py | 2 +- ...t_hosted_vllm_image_edit_transformation.py | 43 +++++++++++++++++++ 3 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py index fa8c1dbfcfd..3b8cc437168 100644 --- a/litellm/llms/hosted_vllm/image_edit/transformation.py +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -3,8 +3,17 @@ from typing import Final from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str +PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT: Final = frozenset({"mask", "quality", "input_fidelity"}) + class HostedVLLMImageEditConfig(OpenAIImageEditConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseImageEditConfig contract + return [ # mutable-ok: BaseImageEditConfig returns list + param + for param in super().get_supported_openai_params(model) + if param not in PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT + ] + def validate_environment( self, headers: dict, # mutable-ok: BaseImageEditConfig contract diff --git a/litellm/utils.py b/litellm/utils.py index 0f2b2f9042d..fc530f14e28 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9239,7 +9239,7 @@ class ProviderConfigManager: from litellm.llms.openai.image_edit import get_openai_image_edit_config return get_openai_image_edit_config(model=model) - if LlmProviders.HOSTED_VLLM == provider: + elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config return get_hosted_vllm_image_edit_config(model=model) diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py index 5646f3d9d13..bc0ea23e249 100644 --- a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py @@ -110,3 +110,46 @@ def test_image_edit_posts_multipart_to_vllm_omni(): assert f'name="model"\r\n\r\n{MODEL}'.encode() in request.content assert b'name="prompt"\r\n\r\nadd a hat' in request.content assert b'name="seed"\r\n\r\n42' in request.content + + +@pytest.mark.parametrize("param", ["mask", "quality", "input_fidelity"]) +def test_params_vllm_omni_ignores_are_not_advertised(param: str): + supported = HostedVLLMImageEditConfig().get_supported_openai_params(MODEL) + + assert param not in supported + assert {"image", "prompt", "n", "size", "response_format", "background", "user"} <= set(supported) + + +def test_image_edit_rejects_quality_unless_dropped(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + + with pytest.raises(litellm.UnsupportedParamsError, match="quality"): + litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + client=client, + quality="low", + ) + assert captured == [] + + litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + client=client, + quality="low", + drop_params=True, + ) + + assert len(captured) == 1 + assert b'name="quality"' not in captured[0].content + assert b'name="prompt"\r\n\r\nadd a hat' in captured[0].content 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 24/29] 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 25/29] 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 26/29] 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, }, }); }); From 0a053d2c8146e4a4739f92d48f1e2bdec75202e5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:21:40 -0700 Subject: [PATCH 27/29] test(guardrails): cover streamed tool-call name rewrites on chat and Messages Skipping the name write-back in either handler left every test green; a guardrail that renames a tool call now has a regression test on both the chat chunk path and the Anthropic SSE path --- .../test_anthropic_guardrail_handler.py | 23 +++++++++++++++++++ .../test_openai_guardrail_handler.py | 23 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 64c243343d5..e091355b69d 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -367,6 +367,29 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert '"stop_reason": "tool_use"' in raw assert "persim" not in raw + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_use_name_back_into_sse_chunks(self): + class RenameTool(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call.function.name = "lookup_fruit_reviewed" + return inputs + + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=RenameTool(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + raw = b"".join(chunks).decode() + assert '"name": "lookup_fruit_reviewed"' in raw and '"id": "toolu_1"' in raw + assert '"name": "lookup_fruit"' not in raw + assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"} + @pytest.mark.asyncio async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index c7011a6f00e..5a29a96829f 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1171,6 +1171,29 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert chunks[3].choices[0].delta.tool_calls is None assert chunks[3].choices[0].finish_reason == "tool_calls" + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_call_name_back_into_chunks(self): + class RenameTool(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call["function"]["name"] = "lookup_fruit_reviewed" + return inputs + + handler = OpenAIChatCompletionsHandler() + chunks = self._ended_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=RenameTool(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]] + assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None] + assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"} + assert fragments[0].id == "call_1" + @pytest.mark.asyncio async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self): handler = OpenAIChatCompletionsHandler() From 2488f84b0293956b3736fcd9cf639491ab68e65e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:53:35 -0700 Subject: [PATCH 28/29] fix(proxy): keep a body litellm_session_id in SpendLogs under missing_session_id omit (#40379) * fix(proxy): keep a body litellm_session_id in SpendLogs under missing_session_id omit Under general_settings.missing_session_id: omit, apply_missing_session_id_policy now mirrors a client-supplied top-level litellm_session_id into metadata.session_id when the client did not set one there, so SpendLogs.session_id and Langfuse agree with the session callbacks already report through StandardLoggingPayload.session_id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep client metadata.session_id ahead of body litellm_session_id on litellm_metadata routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop docstrings from the missing_session_id omit regression tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 10 ++ .../proxy/test_litellm_pre_call_utils.py | 101 ++++++++++++++++++ 2 files changed, 111 insertions(+) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3250ae5cca9..2925226f235 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -773,6 +773,16 @@ def apply_missing_session_id_policy( return if policy == "omit": metadata[SESSION_ID_OMITTED_METADATA_KEY] = True + requester_metadata: Final = data.get("metadata") + requester_session_id: Final = ( + requester_metadata.get("session_id") if isinstance(requester_metadata, dict) else None + ) + if ( + (body_session_id := data.get("litellm_session_id")) + and not metadata.get("session_id") + and not requester_session_id + ): + metadata["session_id"] = body_session_id return if data.get("litellm_session_id") or metadata.get("session_id"): return diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 892fa484ab4..7564fa91c43 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7678,6 +7678,107 @@ async def test_missing_session_id_omit_keeps_client_supplied_session_id(): assert _spend_log_session_id(updated) == "client-session-1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_body", + [ + {"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1"}, + {"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1", "metadata": {"trace_id": "trace-1"}}, + ], +) +async def test_missing_session_id_omit_keeps_body_litellm_session_id( + monkeypatch: pytest.MonkeyPatch, client_body: dict[str, object] +): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + updated = await add_litellm_data_to_request( + data=client_body, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + callback_session_id = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=SimpleNamespace(litellm_session_id=""), + litellm_params=get_litellm_params(litellm_session_id="cust-sess-1", metadata=updated["metadata"]), + ) + assert callback_session_id == "cust-sess-1" + assert updated["metadata"]["session_id"] == "cust-sess-1" + assert _spend_log_session_id(updated) == "cust-sess-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_body_litellm_session_id_does_not_override_metadata_session_id(): + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "messages": [], + "litellm_session_id": "cust-sess-1", + "metadata": {"session_id": "meta-sess-1"}, + }, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["metadata"]["session_id"] == "meta-sess-1" + assert _spend_log_session_id(updated) == "meta-sess-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +async def test_missing_session_id_omit_keeps_metadata_session_id_on_litellm_metadata_routes(path: str): + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "input": "hi", + "litellm_session_id": "cust-sess-1", + "metadata": {"session_id": "meta-sess-1"}, + }, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "meta-sess-1" + assert _spend_log_session_id(updated, "litellm_metadata") == "meta-sess-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +async def test_missing_session_id_omit_keeps_body_litellm_session_id_on_litellm_metadata_routes(path: str): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "input": "hi", "litellm_session_id": "cust-sess-1"}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "cust-sess-1" + assert _spend_log_session_id(updated, "litellm_metadata") == "cust-sess-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_ignores_empty_body_litellm_session_id(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], "litellm_session_id": ""}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert "session_id" not in updated["metadata"] + assert _spend_log_session_id(updated) is None + + @pytest.mark.asyncio async def test_missing_session_id_generate_reuses_traceparent_trace_id(): """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" From 264fc82dc1bb66e840b625830c9d8be03d655a0b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:14:10 -0700 Subject: [PATCH 29/29] fix(rate_limiter): attach v3 priority rate limit headers on /v1/messages (#37228) * fix(rate_limiter): attach v3 priority rate limit headers on /v1/messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(playground): honor the Stream responses toggle for /v1/messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rate_limiter): drop explanatory docstrings from v3 dict response tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 18 ++--- .../hooks/parallel_request_limiter_v3.py | 33 +++------- .../add_retry_fallback_headers.py | 6 ++ .../hooks/test_dynamic_rate_limiter_v3.py | 57 ++++++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 65 ++++++++++++++++++ .../components/chat_ui/ChatUI.test.tsx | 53 +++++++++++++++ .../playground/components/chat_ui/ChatUI.tsx | 6 +- .../llm_calls/anthropic_messages.test.tsx | 66 ++++++++++++++++++- .../llm_calls/anthropic_messages.tsx | 33 +++++++--- 9 files changed, 291 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index de8834449de..a074f02f4e8 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -32,6 +32,10 @@ from litellm.proxy.hooks.rate_limiter_utils import ( resolve_llm_provider_for_rate_limit, ) from litellm.proxy.utils import InternalUsageCache +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral @@ -659,22 +663,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): data=data, user_api_key_dict=user_api_key_dict, response=response ) - # Add additional priority-specific headers - if isinstance(response, ModelResponse): + if response_has_hidden_params(response): priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) - - # Get existing additional headers - additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {} - - # Add priority information + additional_headers: Final = ensure_response_additional_headers(response) additional_headers["x-litellm-priority"] = priority or "default" additional_headers["x-litellm-rate-limiter-version"] = "v3" - # Update response - if not hasattr(response, "_hidden_params"): - response._hidden_params = {} - response._hidden_params["additional_headers"] = additional_headers - return response except Exception as e: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 31437af7770..a72ae3bb1ea 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -52,6 +52,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.utils import ( @@ -4677,34 +4681,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Post-call hook to update rate limit headers in the response. """ try: - from pydantic import BaseModel - stash: Final = get_request_stash() litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None - if litellm_proxy_rate_limit_response is not None: - # Update response headers - if hasattr(response, "_hidden_params"): - _hidden_params = getattr(response, "_hidden_params") - else: - _hidden_params = None - - if _hidden_params is not None and ( - isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict) - ): - if isinstance(_hidden_params, BaseModel): - _hidden_params = _hidden_params.model_dump() - - _additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers( - additional_headers=_hidden_params.get("additional_headers", {}) or {}, + if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response): + additional_headers: Final = ensure_response_additional_headers(response) + additional_headers.update( + self._merge_ratelimit_statuses_into_additional_headers( + additional_headers={}, statuses=litellm_proxy_rate_limit_response["statuses"], ) - - setattr( - response, - "_hidden_params", - {**_hidden_params, "additional_headers": _additional_headers}, - ) + ) except Exception as e: verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e) diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 3ec92ad226a..3251ea457cf 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -50,6 +50,12 @@ def prepare_response_for_header_attachment(response: object) -> object | None: return response +def response_has_hidden_params(response: object) -> bool: + if isinstance(response, dict): + return "_hidden_params" in response + return hasattr(response, "_hidden_params") + + def ensure_response_additional_headers(response: object) -> dict[str, object]: hidden_params: Final = get_hidden_params_dict(response, create=isinstance(response, dict)) _write_hidden_params(response, hidden_params) diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 0ff8b67b1a7..0cd6b4ede9c 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1861,3 +1861,60 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch): ) assert capacity_blocked.value.status_code == 429 assert "Model capacity reached" in capacity_blocked.value.detail["error"] + + +@pytest.mark.asyncio +async def test_post_call_success_hook_attaches_priority_headers_to_dict_response(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + RateLimitResponse, + RateLimitStatus, + get_or_create_request_stash, + ) + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=75, + limit_remaining=74, + rate_limit_type="requests", + descriptor_key="priority_model", + ) + ], + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}}, + } + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-litellm-attempted-retries"] == 0 + assert additional_headers["x-ratelimit-priority_model-limit-requests"] == 75 + assert additional_headers["x-ratelimit-priority_model-remaining-requests"] == 74 + assert additional_headers["x-litellm-priority"] == "premium" + assert additional_headers["x-litellm-rate-limiter-version"] == "v3" + + +@pytest.mark.asyncio +async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}), + response=response, + ) + + assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 4003286d887..6d382370f5f 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6171,3 +6171,68 @@ async def test_success_hook_leaves_stash_untouched_for_non_batch_responses(): data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5)) ) assert get_request_stash().batch_enqueued_reservation == reservation + + +@pytest.mark.asyncio +async def test_post_call_success_hook_attaches_ratelimit_headers_to_dict_response(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=100, + limit_remaining=99, + rate_limit_type="requests", + descriptor_key="model_saturation_check", + ) + ], + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}}, + } + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-dict-response")), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-litellm-attempted-retries"] == 0 + assert additional_headers["x-ratelimit-model_saturation_check-limit-requests"] == 100 + assert additional_headers["x-ratelimit-model_saturation_check-remaining-requests"] == 99 + + +@pytest.mark.asyncio +async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=100, + limit_remaining=99, + rate_limit_type="requests", + descriptor_key="model_saturation_check", + ) + ], + ) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-raw-dict")), + response=response, + ) + + assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index f07b66efdf8..bf6b092a2b0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import ChatUI from "./ChatUI"; import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages"; vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn(), @@ -14,6 +15,10 @@ vi.mock("@/components/llm_calls/chat_completion", () => ({ makeOpenAIChatCompletionRequest: vi.fn().mockResolvedValue(undefined), })); +vi.mock("../../llm_calls/anthropic_messages", () => ({ + makeAnthropicMessagesRequest: vi.fn().mockResolvedValue(undefined), +})); + vi.mock("@/components/networking", () => ({ tagListCall: vi.fn().mockResolvedValue({}), vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), @@ -32,6 +37,8 @@ beforeEach(() => { const CHAT_REQUEST_ARG_COUNT = 26; const STREAMING_ENABLED_ARG_INDEX = 25; +const MESSAGES_REQUEST_ARG_COUNT = 19; +const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18; async function openComboboxByPlaceholder(placeholder: string) { const user = userEvent.setup(); @@ -378,6 +385,52 @@ describe("ChatUI", () => { expect(requestArgs[STREAMING_ENABLED_ARG_INDEX]).toBe(false); }); + it("should send the /v1/messages request non-streaming after Stream responses is unchecked", async () => { + const user = userEvent.setup(); + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/messages"); + await selectComboboxOption("Select a Model", "Model 1"); + + await user.click(await screen.findByTestId("model-settings-button")); + + const streamingCheckbox = await screen.findByRole("checkbox", { name: /Stream responses/i }); + expect(streamingCheckbox).toBeChecked(); + await user.click(streamingCheckbox); + + await waitFor(() => { + expect(screen.getByRole("checkbox", { name: /Stream responses/i })).not.toBeChecked(); + }); + + const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)"); + await act(async () => { + fireEvent.change(messageInput, { target: { value: "hello" } }); + }); + await act(async () => { + fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" }); + }); + + await waitFor(() => { + expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1); + }); + + const requestArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0]; + expect(requestArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT); + expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false); + }); + it("should force streaming in simplified mode even when the playground setting is off", async () => { sessionStorage.setItem("streamingEnabled", "false"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 35378d3d4e7..eca9e2323ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -1025,6 +1025,7 @@ const ChatUI: React.FC = ({ mcpServers, mcpServerToolRestrictions, mcpToolsets, + streamingEnabled, ); } else if (endpointType === EndpointType.EMBEDDINGS) { await makeOpenAIEmbeddingsRequest( @@ -1174,7 +1175,10 @@ const ChatUI: React.FC = ({ return !model.mode || model.mode === "chat"; }; - const supportsStreamingToggle = endpointType === EndpointType.CHAT || endpointType === EndpointType.RESPONSES; + const supportsStreamingToggle = + endpointType === EndpointType.CHAT || + endpointType === EndpointType.RESPONSES || + endpointType === EndpointType.ANTHROPIC_MESSAGES; const modelsForEndpoint = useMemo( () => filterModelsForEndpoint(modelInfo, endpointType as EndpointType), [modelInfo, endpointType], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx index 96ace129f87..9f030d8b2d4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx @@ -7,13 +7,27 @@ vi.mock("@/components/networking", () => ({ })); const mockMessagesStream = vi.fn(); +const mockMessagesCreate = vi.fn(); vi.mock("@anthropic-ai/sdk", () => ({ default: vi.fn(function () { - return { messages: { stream: mockMessagesStream } }; + return { messages: { stream: mockMessagesStream, create: mockMessagesCreate } }; }), })); +const NON_STREAMING_ARGS = [ + undefined, // traceId + undefined, // vector_store_ids + undefined, // guardrails + undefined, // policies + undefined, // selectedMCPServers + undefined, // customBaseUrl + undefined, // mcpServers + undefined, // mcpServerToolRestrictions + undefined, // mcpToolsets + false, // streamingEnabled +] as const; + describe("anthropic_messages prompt cache usage", () => { const captureUsage = async (usage: Record): Promise => { async function* mockStream() { @@ -59,3 +73,53 @@ describe("anthropic_messages prompt cache usage", () => { expect(usageData.promptTokens).toBe(5000); }); }); + +describe("anthropic_messages non-streaming", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it("sends stream:false through messages.create and renders the full reply at once", async () => { + mockMessagesCreate.mockResolvedValue({ + content: [ + { type: "thinking", thinking: "considering" }, + { type: "text", text: "OK" }, + ], + usage: { input_tokens: 12, output_tokens: 3, cache_read_input_tokens: 7 }, + }); + const updateTextUI = vi.fn(); + const onReasoningContent = vi.fn(); + const onUsageData = vi.fn(); + + await makeAnthropicMessagesRequest( + [{ role: "user", content: "Hello" }], + updateTextUI, + "claude-haiku-4-5", + "test-token", + undefined, + undefined, + onReasoningContent, + undefined, + onUsageData, + ...NON_STREAMING_ARGS, + ); + + expect(mockMessagesStream).not.toHaveBeenCalled(); + expect(mockMessagesCreate).toHaveBeenCalledTimes(1); + expect(mockMessagesCreate.mock.calls[0][0]).toMatchObject({ model: "claude-haiku-4-5", stream: false }); + expect(updateTextUI).toHaveBeenCalledWith("assistant", "OK", "claude-haiku-4-5"); + expect(onReasoningContent).toHaveBeenCalledWith("considering"); + const expectedUsage: TokenUsage = { completionTokens: 3, promptTokens: 12, totalTokens: 15, cacheReadTokens: 7 }; + expect(onUsageData).toHaveBeenCalledWith(expectedUsage); + }); + + it("keeps streaming as the default when the flag is omitted", async () => { + async function* emptyStream() {} + mockMessagesStream.mockReturnValue(emptyStream()); + + await makeAnthropicMessagesRequest([{ role: "user", content: "Hello" }], vi.fn(), "claude-haiku-4-5", "test-token"); + + expect(mockMessagesCreate).not.toHaveBeenCalled(); + expect(mockMessagesStream.mock.calls[0][0]).toMatchObject({ stream: true }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx index 14afa013768..9dd2f675c44 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx @@ -7,6 +7,13 @@ import { getProxyBaseUrl } from "@/components/networking"; import { toast } from "@/lib/toast"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +const toTokenUsage = (usage: Anthropic.Usage): TokenUsage => ({ + completionTokens: usage.output_tokens, + promptTokens: usage.input_tokens, + totalTokens: usage.input_tokens + usage.output_tokens, + ...extractPromptCacheTokens(usage), +}); + export async function makeAnthropicMessagesRequest( messages: MessageType[], updateTextUI: (role: string, delta: string, model?: string) => void, @@ -26,6 +33,7 @@ export async function makeAnthropicMessagesRequest( mcpServers?: MCPServer[], mcpServerToolRestrictions?: Record, mcpToolsets?: MCPToolset[], + streamingEnabled: boolean = true, ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -58,7 +66,7 @@ export async function makeAnthropicMessagesRequest( const requestBody: any = { model: selectedModel, messages: messages.map((m) => ({ role: m.role, content: m.content })), - stream: true, + stream: streamingEnabled, max_tokens: 1024, // @ts-ignore - litellm specific parameter litellm_trace_id: traceId, @@ -74,6 +82,20 @@ export async function makeAnthropicMessagesRequest( if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids; if (guardrails) requestBody.guardrails = guardrails; if (policies) requestBody.policies = policies; + + if (!streamingEnabled) { + const message: Anthropic.Message = await client.messages.create({ ...requestBody, stream: false }, { signal }); + for (const block of message.content) { + if (block.type === "text") { + updateTextUI("assistant", block.text, selectedModel); + } else if (block.type === "thinking" && onReasoningContent) { + onReasoningContent(block.thinking); + } + } + onUsageData?.(toTokenUsage(message.usage)); + return; + } + // Use the streaming helper method for cleaner async iteration // @ts-ignore - The SDK types might not include all litellm-specific parameters const stream = client.messages.stream(requestBody, { signal }); @@ -105,14 +127,7 @@ export async function makeAnthropicMessagesRequest( // Process usage data from message_delta events if (messageStreamEvent.type === "message_delta" && (messageStreamEvent as any).usage && onUsageData) { - const usage = (messageStreamEvent as any).usage; - const usageData: TokenUsage = { - completionTokens: usage.output_tokens, - promptTokens: usage.input_tokens, - totalTokens: usage.input_tokens + usage.output_tokens, - ...extractPromptCacheTokens(usage), - }; - onUsageData(usageData); + onUsageData(toTokenUsage((messageStreamEvent as any).usage)); } } } catch (error) {