diff --git a/litellm/proxy/middleware/security_headers_middleware.py b/litellm/proxy/middleware/security_headers_middleware.py index 01e8ade8241..0e551eaba8a 100644 --- a/litellm/proxy/middleware/security_headers_middleware.py +++ b/litellm/proxy/middleware/security_headers_middleware.py @@ -51,11 +51,11 @@ class SecurityHeadersMiddleware: async def send_with_security_headers(message: Message) -> None: if message["type"] == "http.response.start": headers: Final = MutableHeaders(scope=message) - applied: list[tuple[str, str]] = list(STATIC_SECURITY_HEADERS) - if _hsts_enabled(): - applied.append(HSTS_HEADER) - if _no_store_enabled(): - applied.append(NO_STORE_HEADER) + applied: Final = ( + *STATIC_SECURITY_HEADERS, + *((HSTS_HEADER,) if _hsts_enabled() else ()), + *((NO_STORE_HEADER,) if _no_store_enabled() else ()), + ) for name, value in applied: headers.setdefault(name, value) await send(message) diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 2a7da58f39c..5d5b62c8f80 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -32,6 +32,7 @@ EXCLUDED_ROLLOUT_FLAGS = { "LITELLM_USE_RUST_OCR", "LITELLM_RUST", "LITELLM_BUILTIN_PASS_THROUGH_ROUTES_FIRST", + "LITELLM_ENABLE_NO_STORE", } # Internal infrastructure tuning parameters for streaming/queue management diff --git a/tests/unit/proxy/middleware/test_security_headers_middleware.py b/tests/unit/proxy/middleware/test_security_headers_middleware.py index cebe3645879..c85768b7a58 100644 --- a/tests/unit/proxy/middleware/test_security_headers_middleware.py +++ b/tests/unit/proxy/middleware/test_security_headers_middleware.py @@ -6,6 +6,8 @@ that HSTS is opt-in via LITELLM_ENABLE_HSTS and Cache-Control: no-store is opt-in via LITELLM_ENABLE_NO_STORE. """ +from typing import Final + from starlette.applications import Starlette from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import JSONResponse, Response @@ -74,19 +76,19 @@ def test_does_not_override_existing_header(monkeypatch): def test_no_store_absent_by_default(monkeypatch): monkeypatch.delenv("LITELLM_ENABLE_NO_STORE", raising=False) - resp = _make_client(_ok).get("/") + resp: Final = _make_client(_ok).get("/") assert "cache-control" not in resp.headers def test_no_store_present_when_enabled(monkeypatch): monkeypatch.setenv("LITELLM_ENABLE_NO_STORE", "true") - resp = _make_client(_ok).get("/") + resp: Final = _make_client(_ok).get("/") assert resp.headers["cache-control"] == "no-store" def test_no_store_not_enabled_by_arbitrary_value(monkeypatch): monkeypatch.setenv("LITELLM_ENABLE_NO_STORE", "1") - resp = _make_client(_ok).get("/") + resp: Final = _make_client(_ok).get("/") assert "cache-control" not in resp.headers @@ -97,7 +99,7 @@ def test_no_store_does_not_override_existing_header(monkeypatch): return Response("hi", headers={"Cache-Control": "private, max-age=60"}) monkeypatch.setenv("LITELLM_ENABLE_NO_STORE", "true") - resp = _make_client(custom).get("/") + resp: Final = _make_client(custom).get("/") assert resp.headers["cache-control"] == "private, max-age=60" assert resp.headers["x-content-type-options"] == "nosniff" @@ -105,7 +107,7 @@ def test_no_store_does_not_override_existing_header(monkeypatch): def test_both_hsts_and_no_store_enabled(monkeypatch): monkeypatch.setenv("LITELLM_ENABLE_HSTS", "true") monkeypatch.setenv("LITELLM_ENABLE_NO_STORE", "true") - resp = _make_client(_ok).get("/") + resp: Final = _make_client(_ok).get("/") assert resp.headers["strict-transport-security"] == ( "max-age=31536000; includeSubDomains" )