Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions nemoguardrails/server/schemas/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,13 +138,20 @@ async def fetch_models(
api_key_env = provider.get("api_key_env", "OPENAI_API_KEY")

headers: Dict[str, str] = {}
forwarded = request_headers.get("Authorization", "")

forwarded = next(
(value for name, value in request_headers.items() if name.lower() == "authorization"),
"",
)
raw_key = os.environ.get(api_key_env, "")
if not raw_key:
raw_key = forwarded.removeprefix("Bearer ").strip() if forwarded else ""
if raw_key:
headers[auth_header_name] = f"Bearer {raw_key}" if use_bearer else raw_key
elif forwarded:
if auth_header_name == "Authorization" and use_bearer:
headers[auth_header_name] = forwarded
else:
raw_key = forwarded.removeprefix("Bearer ").strip()
if raw_key:
headers[auth_header_name] = f"Bearer {raw_key}" if use_bearer else raw_key
Comment on lines +152 to +154

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Handle case-insensitive bearer prefix normalization.

At Line 152, removeprefix("Bearer ") is case-sensitive. A valid header like authorization: bearer user-token won’t be stripped and can become an invalid provider key (e.g., x-api-key: bearer user-token).

Suggested fix
-            raw_key = forwarded.removeprefix("Bearer ").strip()
+            forwarded_value = forwarded.strip()
+            if forwarded_value[:7].lower() == "bearer ":
+                raw_key = forwarded_value[7:].strip()
+            else:
+                raw_key = forwarded_value
             if raw_key:
                 headers[auth_header_name] = f"Bearer {raw_key}" if use_bearer else raw_key
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
raw_key = forwarded.removeprefix("Bearer ").strip()
if raw_key:
headers[auth_header_name] = f"Bearer {raw_key}" if use_bearer else raw_key
forwarded_value = forwarded.strip()
if forwarded_value[:7].lower() == "bearer ":
raw_key = forwarded_value[7:].strip()
else:
raw_key = forwarded_value
if raw_key:
headers[auth_header_name] = f"Bearer {raw_key}" if use_bearer else raw_key
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@nemoguardrails/server/schemas/utils.py` around lines 152 - 154, The code uses
forwarded.removeprefix("Bearer ") which is case-sensitive and will not strip
"bearer " or other case variants; update the logic around raw_key to perform a
case-insensitive prefix check on forwarded (e.g., check
forwarded.lower().startswith("bearer ")) and if true slice off the
original-length prefix from forwarded to produce raw_key (so token casing is
preserved), otherwise just strip forwarded; then continue using raw_key when
setting headers[auth_header_name] with the existing use_bearer handling.

Comment on lines +148 to +154

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 The shortcut branch for the OpenAI-compatible case only triggers when both auth_header_name == "Authorization" and use_bearer are true, so the forwarded value is passed through verbatim. If a client sends a non-Bearer scheme (e.g. Authorization: ApiKey xyz), that value lands in headers["Authorization"] unmodified and the upstream provider will likely reject the request with a 401. Stripping and re-wrapping with Bearer (as the else branch already does) would make the forwarding consistent with the env-key path.

Suggested change
elif forwarded:
if auth_header_name == "Authorization" and use_bearer:
headers[auth_header_name] = forwarded
else:
raw_key = forwarded.removeprefix("Bearer ").strip()
if raw_key:
headers[auth_header_name] = f"Bearer {raw_key}" if use_bearer else raw_key
elif forwarded:
raw_key = forwarded.removeprefix("Bearer ").strip()
if raw_key:
headers[auth_header_name] = f"Bearer {raw_key}" if use_bearer else raw_key
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/server/schemas/utils.py
Line: 148-154

Comment:
The shortcut branch for the OpenAI-compatible case only triggers when both `auth_header_name == "Authorization"` **and** `use_bearer` are true, so the forwarded value is passed through verbatim. If a client sends a non-Bearer scheme (e.g. `Authorization: ApiKey xyz`), that value lands in `headers["Authorization"]` unmodified and the upstream provider will likely reject the request with a 401. Stripping and re-wrapping with `Bearer` (as the `else` branch already does) would make the forwarding consistent with the env-key path.

```suggestion
    elif forwarded:
        raw_key = forwarded.removeprefix("Bearer ").strip()
        if raw_key:
            headers[auth_header_name] = f"Bearer {raw_key}" if use_bearer else raw_key
```

How can I resolve this? If you propose a fix, please make it concise.


headers.update(provider.get("extra_headers", {}))

Expand Down
28 changes: 27 additions & 1 deletion tests/server/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -820,7 +820,7 @@ def test_list_models_forwards_auth_header():
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)

with patch.dict(os.environ, {"MAIN_MODEL_BASE_URL": "http://localhost:8000"}):
with patch.dict(os.environ, {"MAIN_MODEL_BASE_URL": "http://localhost:8000", "OPENAI_API_KEY": ""}):
with patch("httpx.AsyncClient", return_value=mock_client):
response = client.get(
"/v1/models",
Expand All @@ -833,6 +833,32 @@ def test_list_models_forwards_auth_header():
assert call_kwargs.kwargs["headers"]["Authorization"] == "Bearer my-token"


def test_list_models_env_key_precedes_forwarded_auth_header():
"""Test /v1/models uses configured provider key before a request auth placeholder."""
mock_response = _make_httpx_response({"data": []})
mock_client = AsyncMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)

with patch.dict(
os.environ,
{
"MAIN_MODEL_BASE_URL": "http://localhost:8000",
"OPENAI_API_KEY": "sk-test-key",
},
):
with patch("httpx.AsyncClient", return_value=mock_client):
response = client.get(
"/v1/models",
headers={"Authorization": "Bearer not-used"},
)

assert response.status_code == 200
call_kwargs = mock_client.get.call_args
assert call_kwargs.kwargs["headers"]["Authorization"] == "Bearer sk-test-key"


def test_list_models_uses_openai_api_key_fallback():
"""Test /v1/models falls back to OPENAI_API_KEY when no auth header."""
mock_response = _make_httpx_response({"data": []})
Expand Down
30 changes: 29 additions & 1 deletion tests/server/test_schema_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -450,13 +450,41 @@ async def test_fetch_unknown_engine_no_base_url():
async def test_fetch_auth_forwarded():
"""Incoming Authorization header is forwarded for OpenAI-compatible providers."""
mock = _mock_httpx({"data": []})
with patch.dict(os.environ, {"MAIN_MODEL_BASE_URL": "http://localhost:8000"}):
with patch.dict(os.environ, {"MAIN_MODEL_BASE_URL": "http://localhost:8000", "OPENAI_API_KEY": ""}):
with patch("httpx.AsyncClient", return_value=mock):
await fetch_models("openai", {"Authorization": "Bearer user-token"})
call_headers = mock.get.call_args.kwargs["headers"]
assert call_headers["Authorization"] == "Bearer user-token"


@pytest.mark.asyncio
async def test_fetch_env_key_precedes_forwarded_auth():
"""Configured OpenAI key is used before a caller-supplied placeholder token."""
mock = _mock_httpx({"data": []})
with patch.dict(
os.environ,
{
"MAIN_MODEL_BASE_URL": "http://localhost:8000",
"OPENAI_API_KEY": "sk-configured-key",
},
):
with patch("httpx.AsyncClient", return_value=mock):
await fetch_models("openai", {"Authorization": "Bearer not-used"})
call_headers = mock.get.call_args.kwargs["headers"]
assert call_headers["Authorization"] == "Bearer sk-configured-key"


@pytest.mark.asyncio
async def test_fetch_non_openai_provider_uses_forwarded_auth_without_env_key():
"""Forwarded auth falls back to provider-specific auth headers when env key is unset."""
mock = _mock_httpx({"data": []})
with patch.dict(os.environ, {"ANTHROPIC_API_KEY": ""}):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 The Anthropic provider constructs its URL using MAIN_MODEL_BASE_URL when that variable is set. If a test environment (or a previous test) leaves MAIN_MODEL_BASE_URL set, this test would hit a different URL than https://api.anthropic.com/v1/models. Since httpx.AsyncClient is fully mocked, the URL doesn't affect the assertion, but explicitly clearing the variable documents the intent and guards against subtle breakage if the mock scope ever narrows.

Suggested change
with patch.dict(os.environ, {"ANTHROPIC_API_KEY": ""}):
with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "", "MAIN_MODEL_BASE_URL": ""}):
Prompt To Fix With AI
This is a comment left during a code review.
Path: tests/server/test_schema_utils.py
Line: 481

Comment:
The Anthropic provider constructs its URL using `MAIN_MODEL_BASE_URL` when that variable is set. If a test environment (or a previous test) leaves `MAIN_MODEL_BASE_URL` set, this test would hit a different URL than `https://api.anthropic.com/v1/models`. Since `httpx.AsyncClient` is fully mocked, the URL doesn't affect the assertion, but explicitly clearing the variable documents the intent and guards against subtle breakage if the mock scope ever narrows.

```suggestion
    with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "", "MAIN_MODEL_BASE_URL": ""}):
```

How can I resolve this? If you propose a fix, please make it concise.

with patch("httpx.AsyncClient", return_value=mock):
await fetch_models("anthropic", {"Authorization": "Bearer user-token"})
call_headers = mock.get.call_args.kwargs["headers"]
assert call_headers["x-api-key"] == "user-token"


@pytest.mark.asyncio
async def test_fetch_non_dict_items_skipped():
"""Non-dict items in the response data are skipped."""
Expand Down
Loading