fix(inference-gateway): preserve backend authentication errors#895
Open
anastasia-nesterenko wants to merge 1 commit into
Open
Conversation
Signed-off-by: anastasia-nesterenko <anesterenko@nvidia.com>
anastasia-nesterenko
requested review from
benmccown,
ironcommit and
mckornfield
July 24, 2026 20:42
Contributor
📝 WalkthroughWalkthroughChangesThe inference gateway now preserves backend 401/403 responses and bodies, wraps backend 404 responses as 502, filters rewritten response headers, treats typed model results as non-streaming, and bypasses VirtualModel rewriting and middleware for authentication failures. Tests cover these behaviors across proxy and provider routes. Inference gateway error propagation
Sequence Diagram(s)sequenceDiagram
participant Client
participant virtual_model_proxy
participant fetch_proxy_response
participant Backend
participant ResponseMiddleware
Client->>virtual_model_proxy: submit inference request
virtual_model_proxy->>fetch_proxy_response: fetch provider response
fetch_proxy_response->>Backend: forward request
Backend-->>fetch_proxy_response: status, headers, body
fetch_proxy_response-->>virtual_model_proxy: auth bytes or response result
alt 401 or 403
virtual_model_proxy-->>Client: buffered upstream error response
else successful response
virtual_model_proxy->>ResponseMiddleware: parse and process response
ResponseMiddleware-->>Client: rewritten model response
end
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
services/core/inference-gateway/src/nmp/core/inference_gateway/api/proxy.py (1)
369-394: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared backend-error policy.
Lines 369-394 and 516-540 are identical apart from the auth return shape, and the mock branch at 898-911 restates the same 401/403/404 policy a third time. A single helper keeps them from drifting:
async def _handle_backend_error( response: aiohttp.ClientResponse, url: str ) -> tuple[bytes, CIMultiDict[str], int]: """Return (raw_body, headers, status) for auth failures; raise otherwise.""" status = response.status headers = _filter_response_headers(response.headers) if status in _BACKEND_AUTH_ERROR_STATUSES else CIMultiDict() raw_body, error_body = await _read_error_body(response) _close_response(response) logger.warning("Backend error %d from %s: %s", status, url, error_body) if status in _BACKEND_AUTH_ERROR_STATUSES: return raw_body, headers, status if status == 404: detail = f"Backend returned 404: {error_body}" if error_body else "Backend returned 404" raise HTTPException(status_code=502, detail=detail) raise HTTPException(status_code=status, detail=error_body or str(status))Callers then become
raw, hdrs, st = await _handle_backend_error(response, next_request_info.url)plus the local return shape.🤖 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 `@services/core/inference-gateway/src/nmp/core/inference_gateway/api/proxy.py` around lines 369 - 394, Extract the duplicated backend-error handling from the proxy request paths into a shared async helper named _handle_backend_error. Have it read and close the response, log the backend error, preserve filtered headers and raw-body returns for auth statuses, and raise the existing 404/other-status HTTPException policies otherwise. Replace the handling in the affected callers, including the mock branch, while preserving each caller’s local response shape.
🤖 Prompt for all review comments with 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.
Inline comments:
In `@services/core/inference-gateway/src/nmp/core/inference_gateway/api/proxy.py`:
- Around line 324-331: Update _raw_response to preserve duplicate upstream
response headers by filtering _BODY_FRAMING_HEADERS_TO_STRIP and appending the
remaining header pairs to the Response raw_headers representation, rather than
converting headers to a dict. Keep the existing body, status code, and
framing-header filtering behavior unchanged.
---
Nitpick comments:
In `@services/core/inference-gateway/src/nmp/core/inference_gateway/api/proxy.py`:
- Around line 369-394: Extract the duplicated backend-error handling from the
proxy request paths into a shared async helper named _handle_backend_error. Have
it read and close the response, log the backend error, preserve filtered headers
and raw-body returns for auth statuses, and raise the existing 404/other-status
HTTPException policies otherwise. Replace the handling in the affected callers,
including the mock branch, while preserving each caller’s local response shape.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 329076ab-09fe-4255-afe9-7cf4c7cf2fb9
📒 Files selected for processing (3)
services/core/inference-gateway/src/nmp/core/inference_gateway/api/proxy.pyservices/core/inference-gateway/tests/unit/test_provider_router.pyservices/core/inference-gateway/tests/unit/test_proxy.py
Comment on lines
+324
to
+331
| def _raw_response( | ||
| body: bytes | bytearray | memoryview | str, | ||
| status_code: int, | ||
| headers: CIMultiDict[str], | ||
| ) -> Response: | ||
| """Return a buffered body with upstream metadata and corrected framing headers.""" | ||
| safe_headers = {k: v for k, v in headers.items() if k.lower() not in _BODY_FRAMING_HEADERS_TO_STRIP} | ||
| return Response(content=body, status_code=status_code, headers=safe_headers) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -nP -A25 'def _filter_response_headers' services/core/inference-gateway/srcRepository: NVIDIA-NeMo/nemo-platform
Length of output: 3313
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE=services/core/inference-gateway/src/nmp/core/inference_gateway/api/proxy.py
sed -n '1,340p' "$FILE" | nl -ba | sed -n '1,180p'
printf '\n--- RESPONSE_HEADERS_TO_DROP occurrences ---\n'
rg -n 'RESPONSE_HEADERS_TO_DROP|BODY_FRAMING_HEADERS_TO_STRIP|www-authenticate|set-cookie' "$FILE" services/core/inference-gateway/src -nRepository: NVIDIA-NeMo/nemo-platform
Length of output: 203
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
try:
from starlette.responses import Response
from multidict import CIMultiDict
except Exception as e:
print(f"IMPORT_ERROR: {e}")
raise SystemExit(0)
headers = CIMultiDict([
("WWW-Authenticate", "Bearer realm=one"),
("WWW-Authenticate", "Basic realm=two"),
("X-Test", "a"),
("X-Test", "b"),
])
safe_headers = {k: v for k, v in headers.items() if k.lower() not in {"content-length", "content-encoding", "transfer-encoding"}}
resp = Response(content=b"x", headers=safe_headers)
print("safe_headers:", safe_headers)
print("dict items from CIMultiDict:", list(headers.items()))
print("response.raw_headers:", resp.raw_headers)
PYRepository: NVIDIA-NeMo/nemo-platform
Length of output: 491
Preserve repeated response headers. dict collapses duplicate names, so backend WWW-Authenticate challenges and Set-Cookie entries are reduced to the last value. Append the filtered headers to raw_headers instead.
🤖 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 `@services/core/inference-gateway/src/nmp/core/inference_gateway/api/proxy.py`
around lines 324 - 331, Update _raw_response to preserve duplicate upstream
response headers by filtering _BODY_FRAMING_HEADERS_TO_STRIP and appending the
remaining header pairs to the Response raw_headers representation, rather than
converting headers to a dict. Keep the existing body, status code, and
framing-header filtering behavior unchanged.
Contributor
|
mckornfield
reviewed
Jul 24, 2026
| response_headers = ( | ||
| _filter_response_headers(response.headers) | ||
| if response_status in _BACKEND_AUTH_ERROR_STATUSES | ||
| else CIMultiDict() |
Contributor
There was a problem hiding this comment.
confused, if we're NOT in error states, don't we still want the headers?
| response_headers = ( | ||
| _filter_response_headers(response.headers) | ||
| if response_status in _BACKEND_AUTH_ERROR_STATUSES | ||
| else CIMultiDict() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Preserve backend
401 Unauthorizedand403 Forbiddenresponses in the inference gateway instead of wrapping them in502 Bad Gateway.This allows downstream consumers to identify authentication failures correctly, avoid retrying them as transient errors, and surface the original backend error to users.
Changes
401and403status codes.Content-TypeandWWW-Authenticate.404responses mapped to502.Test coverage
Added coverage for:
401and403responsesContent-TypeandWWW-AuthenticatepropagationVerification
557 passed, 5 skipped117 passedgit diff --cached --check: passedLinear
AIRCORE-345 — Inference gateway wraps backend 401 Unauthorized in 502
Summary by CodeRabbit
Bug Fixes
Tests