refactor: centralize response header filtering

Author: ALikesToCode

route_helpers.py

@@ -1,7 +1,7 @@
 import logging
 import os
 from functools import wraps
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, Mapping, Optional
 from urllib.parse import urlparse
 
 from flask import Response, g, jsonify, redirect, request, url_for
@@ -16,6 +16,29 @@
 CORS_ALLOWED_METHODS = "GET, POST, PUT, DELETE, PATCH, OPTIONS"
 CORS_DEFAULT_HEADERS = "Authorization, Content-Type, Accept, Origin, X-Requested-With"
 LOCAL_DEVELOPMENT_HOSTS = {"localhost", "127.0.0.1", "::1"}
+HOP_BY_HOP_RESPONSE_HEADERS = frozenset(
+    {
+        "connection",
+        "content-encoding",
+        "content-length",
+        "keep-alive",
+        "proxy-authenticate",
+        "proxy-authorization",
+        "te",
+        "trailer",
+        "transfer-encoding",
+        "upgrade",
+    }
+)
+
+
+def copy_upstream_response_headers(upstream_headers: Mapping[str, Any]) -> Dict[str, Any]:
+    """Copy only response headers that are safe for Flask to emit downstream."""
+    return {
+        key: value
+        for key, value in upstream_headers.items()
+        if key.lower() not in HOP_BY_HOP_RESPONSE_HEADERS
+    }
 
 
 def mask_secret(value: Optional[str]) -> str:

routes/core.py

@@ -12,27 +12,13 @@
 from config import Config
 from error_handlers import APIError, INTERNAL_ERROR_MESSAGE, get_request_id, internal_error_payload
 from proxy import PROVIDER_DETAILS
-from route_helpers import apply_cors_headers, check_provider, login_required
+from route_helpers import apply_cors_headers, check_provider, copy_upstream_response_headers, login_required
 from services.auth_service import AuthService
 from services.metrics_service import MetricsService
 
 logger = logging.getLogger(__name__)
 
 
-HOP_BY_HOP_RESPONSE_HEADERS = {
-    "connection",
-    "content-encoding",
-    "content-length",
-    "keep-alive",
-    "proxy-authenticate",
-    "proxy-authorization",
-    "te",
-    "trailer",
-    "transfer-encoding",
-    "upgrade",
-}
-
-
 def is_safe_redirect_target(target: str | None) -> bool:
     """Allow only local, absolute-path redirects after login."""
     if not target:
@@ -114,15 +100,6 @@ def build_openrouter_dashboard_headers() -> dict:
     return headers
 
 
-def copy_upstream_response_headers(upstream_headers: requests.structures.CaseInsensitiveDict) -> dict:
-    """Copy only response headers that are safe for Flask to emit downstream."""
-    return {
-        key: value
-        for key, value in upstream_headers.items()
-        if key.lower() not in HOP_BY_HOP_RESPONSE_HEADERS
-    }
-
-
 def register_core_routes(app) -> None:
     @app.errorhandler(CSRFError)
     def handle_csrf_error(error: CSRFError):

routes/proxy.py

@@ -10,32 +10,10 @@
 from error_handlers import APIError
 from providers.registry import get_adapter
 from proxy import PROVIDER_DETAILS
-from route_helpers import api_auth_required, login_required
+from route_helpers import api_auth_required, copy_upstream_response_headers, login_required
 
 logger = logging.getLogger(__name__)
 
-HOP_BY_HOP_RESPONSE_HEADERS = {
-    "connection",
-    "content-encoding",
-    "content-length",
-    "keep-alive",
-    "proxy-authenticate",
-    "proxy-authorization",
-    "te",
-    "trailer",
-    "transfer-encoding",
-    "upgrade",
-}
-
-
-def _safe_response_headers(headers):
-    return {
-        key: value
-        for key, value in headers.items()
-        if key.lower() not in HOP_BY_HOP_RESPONSE_HEADERS
-    }
-
-
 def _dashboard_chat_completions_url(app, provider):
     if provider == "googleai":
         project_id = os.environ.get("PROJECT_ID")
@@ -286,7 +264,7 @@ def proxy_chat_completions():
                 response.content,
                 status=status_code,
                 content_type=response.headers.get("content-type", "application/json"),
-                headers=_safe_response_headers(response.headers),
+                headers=copy_upstream_response_headers(response.headers),
             )
 
         except APIError as error:

routes/unified.py

@@ -7,33 +7,11 @@
 
 from error_handlers import APIError
 from providers.registry import get_adapter
-from route_helpers import api_auth_required, login_required
+from route_helpers import api_auth_required, copy_upstream_response_headers, login_required
 from services.auth_service import AuthService
 from services.model_registry import ModelRegistry
 
 
-HOP_BY_HOP_RESPONSE_HEADERS = {
-    "connection",
-    "content-encoding",
-    "content-length",
-    "keep-alive",
-    "proxy-authenticate",
-    "proxy-authorization",
-    "te",
-    "trailer",
-    "transfer-encoding",
-    "upgrade",
-}
-
-
-def _safe_response_headers(headers):
-    return {
-        key: value
-        for key, value in headers.items()
-        if key.lower() not in HOP_BY_HOP_RESPONSE_HEADERS
-    }
-
-
 def _provider_token(auth_service_cls, provider: str) -> str:
     token = auth_service_cls.get_google_token() if provider == "googleai" else auth_service_cls.get_api_key(provider)
     if not token:
@@ -184,7 +162,7 @@ def unified_chat_completions():
                 response.content,
                 status=response.status_code,
                 content_type=response.headers.get("content-type", "application/json"),
-                headers=_safe_response_headers(response.headers),
+                headers=copy_upstream_response_headers(response.headers),
             )
         except ValueError as error:
             raise APIError(str(error), status_code=400) from error

test_route_helpers.py

@@ -7,6 +7,7 @@
     api_auth_required,
     apply_cors_headers,
     build_cors_preflight_response,
+    copy_upstream_response_headers,
     extract_bearer_token,
     is_cors_origin_allowed,
     mask_authorization_header,
@@ -37,6 +38,24 @@ def test_extract_bearer_token_rejects_missing_or_wrong_scheme(self):
         self.assertIsNone(extract_bearer_token("Basic token-value"))
         self.assertIsNone(extract_bearer_token("Bearer "))
 
+    def test_copy_upstream_response_headers_drops_hop_by_hop_values(self):
+        headers = copy_upstream_response_headers(
+            {
+                "Content-Type": "application/json",
+                "Connection": "keep-alive",
+                "Transfer-Encoding": "chunked",
+                "X-Request-ID": "req_123",
+            }
+        )
+
+        self.assertEqual(
+            headers,
+            {
+                "Content-Type": "application/json",
+                "X-Request-ID": "req_123",
+            },
+        )
+
 
 class RouteHelperCorsTest(unittest.TestCase):
     def setUp(self):