🚧 initial implementation of forwarding x headers

Author: m-misiura

nemoguardrails/actions/llm/utils.py

@@ -25,6 +25,8 @@
 from nemoguardrails.colang.v2_x.lang.colang_ast import Flow
 from nemoguardrails.colang.v2_x.runtime.flows import InternalEvent, InternalEvents
 from nemoguardrails.context import (
+    api_request_headers_var,
+    get_llm_needs_runtime_auth,
     llm_call_info_var,
     llm_response_metadata_var,
     reasoning_trace_var,
@@ -54,6 +56,45 @@
 ]
 
 
+_INFRA_PREFIXES = ("x-forwarded", "x-real-", "x-request-id", "x-remote-")
+
+
+def get_extra_headers_from_request(forward_auth: bool = True) -> Optional[Dict[str, str]]:
+    """Forward X-* headers from the incoming request to the LLM call.
+
+    Excludes proxy/infra headers and the inbound Authorization header (which
+    typically carries K8s/proxy auth and must never be forwarded to the LLM).
+    When forward_auth is True, forwards X-Authorization as Authorization to
+    the LLM provider.
+    """
+    request_headers = api_request_headers_var.get()
+    if not request_headers:
+        return None
+
+    extra_headers = {}
+
+    for k, v in request_headers.items():
+        lower = k.lower()
+        if lower in ("authorization", "x-authorization"):
+            continue
+        if lower.startswith("x-") and not lower.startswith(_INFRA_PREFIXES):
+            extra_headers[k] = v
+
+    if forward_auth:
+        auth = request_headers.get("x-authorization")
+        if auth:
+            extra_headers["Authorization"] = auth
+
+    return extra_headers or None
+
+
+def get_extra_headers_for_llm(llm: "BaseLanguageModel") -> Dict[str, str]:
+    """Return extra_headers dict for an LLM based on its runtime auth needs."""
+    needs_runtime_auth = get_llm_needs_runtime_auth(llm)
+    extra_headers = get_extra_headers_from_request(forward_auth=needs_runtime_auth)
+    return extra_headers or {}
+
+
 def _infer_provider_from_module(llm: BaseLanguageModel) -> Optional[str]:
     """Infer provider name from the LLM's module path.
 
@@ -224,6 +265,15 @@ async def llm_call(
         llm_params_with_stop = llm_params
 
     filtered_params = _filter_params_for_openai_reasoning_models(llm, llm_params_with_stop)
+
+    # Forward X-* headers from the incoming request. Only forward Authorization
+    # to models that need runtime auth (i.e. had no valid static API key at init).
+    needs_runtime_auth = get_llm_needs_runtime_auth(llm)
+    extra_headers = get_extra_headers_from_request(forward_auth=needs_runtime_auth)
+    if extra_headers:
+        filtered_params = filtered_params.copy() if filtered_params else {}
+        filtered_params["extra_headers"] = extra_headers
+
     generation_llm: Union[BaseLanguageModel, Runnable] = llm.bind(**filtered_params) if filtered_params else llm
 
     if streaming_handler:

nemoguardrails/context.py

@@ -62,3 +62,24 @@
 llm_response_metadata_var: contextvars.ContextVar[Optional[dict]] = contextvars.ContextVar(
     "llm_response_metadata", default=None
 )
+
+# The request headers from the incoming API request (set by the server layer).
+# Used to forward Authorization and X-* headers to outgoing LLM calls.
+api_request_headers_var: contextvars.ContextVar[Optional[Dict[str, str]]] = contextvars.ContextVar(
+    "api_request_headers", default=None
+)
+
+# Registry tracking which LLM instances need runtime auth (i.e. had no valid
+# static API key at init time). Keyed by id(llm) — safe because LLM objects
+# are server-lifetime singletons created once in _init_llms().
+_llm_needs_runtime_auth: Dict[int, bool] = {}
+
+
+def set_llm_needs_runtime_auth(llm: Any, needs_auth: bool) -> None:
+    """Register whether an LLM instance needs runtime auth from request headers."""
+    _llm_needs_runtime_auth[id(llm)] = needs_auth
+
+
+def get_llm_needs_runtime_auth(llm: Any) -> bool:
+    """Check whether an LLM instance needs runtime auth from request headers."""
+    return _llm_needs_runtime_auth.get(id(llm), False)

nemoguardrails/library/hallucination/actions.py

@@ -22,6 +22,7 @@
 from nemoguardrails import RailsConfig
 from nemoguardrails.actions import action
 from nemoguardrails.actions.llm.utils import (
+    get_extra_headers_for_llm,
     get_multiline_response,
     llm_call,
     strip_quotes,
@@ -79,7 +80,11 @@ async def self_check_hallucination(
 
         # Generate multiple responses with temperature 1.
         # Bind the config parameters to the LLM for this call
-        llm_with_config = llm.bind(temperature=1.0, n=num_responses)
+        bind_kwargs = {"temperature": 1.0, "n": num_responses}
+        extra_headers = get_extra_headers_for_llm(llm)
+        if extra_headers:
+            bind_kwargs["extra_headers"] = extra_headers
+        llm_with_config = llm.bind(**bind_kwargs)
         extra_llm_response = await llm_with_config.agenerate(
             [formatted_prompt],
             callbacks=logging_callback_manager_for_chain.handlers,

nemoguardrails/logging/callbacks.py

@@ -12,7 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import copy
 import logging
+import re
 import uuid
 from time import time
 from typing import Any, Dict, List, Optional, cast
@@ -36,6 +38,20 @@
 
 log = logging.getLogger(__name__)
 
+_BEARER_RE = re.compile(r"(Bearer\s+)\S+", re.IGNORECASE)
+
+
+def _redact_invocation_params(params: Dict[str, Any]) -> Dict[str, Any]:
+    """Return a copy of invocation params with auth tokens redacted."""
+    if "extra_headers" not in params:
+        return params
+    params = copy.copy(params)
+    headers = params["extra_headers"]
+    params["extra_headers"] = {
+        k: _BEARER_RE.sub(r"\1[REDACTED]", v) if "authorization" in k.lower() else v for k, v in headers.items()
+    }
+    return params
+
 
 class LoggingCallbackHandler(AsyncCallbackHandler):
     """Async callback handler that can be used to handle callbacks from langchain."""
@@ -64,7 +80,7 @@ async def on_llm_start(
         if explain_info:
             explain_info.llm_calls.append(llm_call_info)
 
-        log.info("Invocation Params :: %s", kwargs.get("invocation_params", {}))
+        log.info("Invocation Params :: %s", _redact_invocation_params(kwargs.get("invocation_params", {})))
         log.info(
             "Prompt :: %s",
             prompts[0],
@@ -123,7 +139,7 @@ async def on_chat_model_start(
             ]
         )
 
-        log.info("Invocation Params :: %s", kwargs.get("invocation_params", {}))
+        log.info("Invocation Params :: %s", _redact_invocation_params(kwargs.get("invocation_params", {})))
         log.info(
             "Prompt Messages :: %s",
             prompt,

nemoguardrails/rails/llm/llmrails.py

@@ -66,6 +66,7 @@
     generation_options_var,
     llm_stats_var,
     raw_llm_request,
+    set_llm_needs_runtime_auth,
     streaming_handler_var,
 )
 from nemoguardrails.embeddings.index import EmbeddingsIndex
@@ -369,19 +370,31 @@ def _prepare_model_kwargs(self, model_config):
             model_config: The model configuration object
 
         Returns:
-            dict: The prepared kwargs for model initialization
+            dict: kwargs dict for model initialization
         """
         # Make a copy to avoid modifying the original model config
         kwargs = dict(model_config.parameters) if model_config.parameters else {}
 
-        # If the optional API Key Environment Variable is set, add it to kwargs
         if model_config.api_key_env_var:
             api_key = os.environ.get(model_config.api_key_env_var)
             if api_key:
                 kwargs["api_key"] = api_key
+                kwargs["openai_api_key"] = api_key
+        elif "api_key" not in kwargs:
+            # Placeholder to satisfy LangChain constructors. Real auth arrives
+            # via forwarded request headers at call time. If no auth header is
+            # provided, the LLM call will fail with a 401 — this is the intended
+            # fail-safe: requests without valid auth are rejected by the provider.
+            kwargs["api_key"] = "runtime-provided"
+            kwargs["openai_api_key"] = "runtime-provided"
 
         return kwargs
 
+    @staticmethod
+    def _needs_runtime_auth(kwargs):
+        """Check if the model needs auth provided at request time."""
+        return kwargs.get("api_key") == "runtime-provided"
+
     def _init_llms(self):
         """
         Initializes the right LLM engines based on the configuration.
@@ -421,6 +434,7 @@ def _init_llms(self):
                     mode="chat",
                     kwargs=kwargs,
                 )
+                set_llm_needs_runtime_auth(self.llm, self._needs_runtime_auth(kwargs))
                 self.runtime.register_action_param("llm", self.llm)
 
             else:
@@ -453,6 +467,7 @@ def _init_llms(self):
                     mode=mode,
                     kwargs=kwargs,
                 )
+                set_llm_needs_runtime_auth(llm_model, self._needs_runtime_auth(kwargs))
 
                 # Configure the model based on its type
                 if llm_config.type == "main":

nemoguardrails/server/api.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import asyncio
-import contextvars
 import copy
 import importlib.util
 import json
@@ -37,6 +36,7 @@
 from starlette.staticfiles import StaticFiles
 
 from nemoguardrails import LLMRails, RailsConfig, utils
+from nemoguardrails.context import api_request_headers_var
 from nemoguardrails.rails.llm.config import Model
 from nemoguardrails.rails.llm.options import (
     ActivatedRail,
@@ -91,9 +91,6 @@ def __init__(self, *args, **kwargs):
 
 api_description = """Guardrails Sever API."""
 
-# The headers for each request
-api_request_headers: contextvars.ContextVar = contextvars.ContextVar("headers")
-
 # The datastore that the Server should use.
 # This is currently used only for storing threads.
 # TODO: refactor to wrap the FastAPI instance inside a RailsServer class
@@ -306,9 +303,13 @@ def _update_models_in_config(config: RailsConfig, main_model: Model) -> RailsCon
             break
 
     if main_model_index is not None:
-        parameters = {**models[main_model_index].parameters, **main_model.parameters}
+        existing = models[main_model_index]
+        parameters = {**existing.parameters, **main_model.parameters}
+        # Preserve api_key_env_var from the original config if the override doesn't set one
+        api_key_env_var = main_model.api_key_env_var or existing.api_key_env_var
         models[main_model_index] = main_model
         models[main_model_index].parameters = parameters
+        models[main_model_index].api_key_env_var = api_key_env_var
     else:
         models.append(main_model)
 
@@ -479,7 +480,7 @@ async def chat_completion(body: GuardrailsChatCompletionRequest, request: Reques
         asyncio.get_event_loop().create_task(logger({"endpoint": "/v1/chat/completions", "body": body.json()}))
 
     # Save the request headers in a context variable.
-    api_request_headers.set(request.headers)
+    api_request_headers_var.set(dict(request.headers))
 
     # Use Request config_ids if set, otherwise use the FastAPI default config.
     # If neither is available we can't generate any completions as we have no config_id
@@ -1103,7 +1104,7 @@ async def guardrail_checks(body: GuardrailsChatCompletionRequest, request: Reque
             )
         )
 
-    api_request_headers.set(request.headers)
+    api_request_headers_var.set(dict(request.headers))
 
     async def process_checks():
         """Process guardrail checks and yield results.

tests/test_header_forwarding.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for header forwarding and log redaction."""
+
+from nemoguardrails.actions.llm.utils import get_extra_headers_from_request
+from nemoguardrails.context import api_request_headers_var
+from nemoguardrails.logging.callbacks import _redact_invocation_params
+
+
+def _set_headers(headers):
+    api_request_headers_var.set(headers)
+
+
+def test_no_headers_returns_none():
+    api_request_headers_var.set(None)
+    assert get_extra_headers_from_request() is None
+
+
+def test_mixed_headers_full_scenario():
+    """Core test: infra filtered, x-auth wins, non-x ignored, custom forwarded."""
+    _set_headers(
+        {
+            "authorization": "Bearer oauth-token",
+            "x-authorization": "Bearer llm-key",
+            "x-forwarded-for": "1.2.3.4",
+            "x-remote-user": "admin",
+            "x-real-ip": "5.6.7.8",
+            "x-request-id": "abc",
+            "x-maas-subscription": "sub-key",
+            "content-type": "application/json",
+        }
+    )
+    result = get_extra_headers_from_request(forward_auth=True)
+    assert result == {
+        "Authorization": "Bearer llm-key",
+        "x-maas-subscription": "sub-key",
+    }
+
+
+def test_forward_auth_false_skips_auth():
+    _set_headers({"authorization": "Bearer key", "x-authorization": "Bearer key2"})
+    assert get_extra_headers_from_request(forward_auth=False) is None
+
+
+def test_authorization_never_forwarded_to_llm():
+    """Authorization header (K8s/proxy auth) must never reach the LLM."""
+    _set_headers({"authorization": "Bearer k8s-token"})
+    result = get_extra_headers_from_request(forward_auth=True)
+    assert result is None
+
+
+def test_x_authorization_forwarded_without_authorization():
+    _set_headers({"x-authorization": "Bearer llm-key"})
+    result = get_extra_headers_from_request(forward_auth=True)
+    assert result == {"Authorization": "Bearer llm-key"}
+
+
+def test_redact_hides_bearer_tokens():
+    params = {
+        "model": "gpt-4",
+        "extra_headers": {
+            "Authorization": "Bearer sk-secret-123",
+            "x-custom": "visible",
+        },
+    }
+    result = _redact_invocation_params(params)
+    assert result["extra_headers"]["Authorization"] == "Bearer [REDACTED]"
+    assert result["extra_headers"]["x-custom"] == "visible"
+    # original unchanged
+    assert params["extra_headers"]["Authorization"] == "Bearer sk-secret-123"
+
+
+def test_redact_no_extra_headers_passthrough():
+    params = {"model": "gpt-4"}
+    assert _redact_invocation_params(params) is params