From c468fcc58d9a4ae9276ad64acabe2a4acce55445 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 5 Oct 2026 16:17:17 +0200 Subject: [PATCH 01/11] feat(server): bind guarded streams to HTTP Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_http_kernel.py | 18 +- .../server/experimental/_streaming_http.py | 210 +++++++++++ .../experimental/test_streaming_http.py | 352 ++++++++++++++++++ 3 files changed, 575 insertions(+), 5 deletions(-) create mode 100644 nemoguardrails/server/experimental/_streaming_http.py create mode 100644 tests/server/experimental/test_streaming_http.py diff --git a/nemoguardrails/server/experimental/_http_kernel.py b/nemoguardrails/server/experimental/_http_kernel.py index 508038e262..82a7c51146 100644 --- a/nemoguardrails/server/experimental/_http_kernel.py +++ b/nemoguardrails/server/experimental/_http_kernel.py @@ -374,17 +374,25 @@ async def _buffer_request(request: Request, max_body_bytes: int) -> BufferedHttp ) -def _render_response(value: BufferedHttpResponse, request_method: str | None = None) -> Response: - """Preserve end-to-end values while generating safe downstream framing.""" +def _end_to_end_headers(headers: HttpHeaders) -> list[tuple[bytes, bytes]]: + """Remove connection-specific metadata from buffered and streamed responses.""" connection_tokens = { token.strip().lower() - for name, content in value.headers + for name, content in headers if name.lower() == b"connection" for token in content.split(b",") } - excluded = _HOP_BY_HOP_HEADERS | connection_tokens | {b"content-length"} - headers = [(name, content) for name, content in value.headers if name.lower() not in excluded] + excluded = _HOP_BY_HOP_HEADERS | connection_tokens + return [(name, content) for name, content in headers if name.lower() not in excluded] + + +def _render_response(value: BufferedHttpResponse, request_method: str | None = None) -> Response: + """Preserve end-to-end values while generating safe downstream framing.""" + + headers = [ + (name, content) for name, content in _end_to_end_headers(value.headers) if name.lower() != b"content-length" + ] bodyless = request_method == "HEAD" or value.status_code < 200 or value.status_code in {204, 205, 304} response = Response(content=b"" if bodyless else value.body, status_code=value.status_code) if value.status_code == 205: diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py new file mode 100644 index 0000000000..4f97eb33cb --- /dev/null +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -0,0 +1,210 @@ +# 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. + +"""Bind guarded provider streams to an injected HTTP exchange.""" + +import inspect +import re +from collections.abc import AsyncIterator, Awaitable, Callable +from dataclasses import dataclass + +from starlette.responses import Response, StreamingResponse + +from nemoguardrails.server.experimental._buffered_kernel import InspectionStage, OperationProjectionFailed +from nemoguardrails.server.experimental._content_checker import ContentChecker, StreamBufferingPolicy +from nemoguardrails.server.experimental._guarded_operation import UnsupportedGuardedPayload +from nemoguardrails.server.experimental._guarded_stream import ( + StreamOutcomeRenderer, + StreamProcessingFailed, + StreamUpstreamFailed, + guard_provider_stream, + validate_streaming_policy, +) +from nemoguardrails.server.experimental._http_kernel import ( + BufferedHttpRequest, + BufferedHttpResponse, + HttpHeaders, + ResponseBodyTooLarge, + _end_to_end_headers, + _render_response, +) +from nemoguardrails.server.experimental.provider.stream import ProviderStreamAdapter +from nemoguardrails.server.experimental.provider.types import GuardedMessage + + +@dataclass(frozen=True, slots=True) +class StreamingHttpResponse: + """Carry one successful response without choosing an HTTP client.""" + + status_code: int + headers: HttpHeaders + body: AsyncIterator[bytes] + + def __post_init__(self) -> None: + if type(self.status_code) is not int or not 200 <= self.status_code <= 299: + raise ValueError("A streaming HTTP response must have a successful status code.") + if not hasattr(self.body, "__aiter__"): + raise TypeError("A streaming HTTP response body must be an async iterator.") + if any(not isinstance(name, bytes) or not isinstance(value, bytes) for name, value in self.headers): + raise TypeError("Streaming HTTP response headers must contain byte pairs.") + + +StreamingHttpDispatch = Callable[[BufferedHttpRequest], Awaitable[StreamingHttpResponse | BufferedHttpResponse]] + + +def _streaming_response( + value: StreamingHttpResponse, + body: AsyncIterator[bytes], + *, + may_modify: bool, +) -> StreamingResponse: + response = StreamingResponse(body, status_code=value.status_code) + response.raw_headers = [ + (name, header_value) + for name, header_value in _end_to_end_headers(value.headers) + if not may_modify or name.lower() != b"content-length" + ] + return response + + +def _header_values(headers: HttpHeaders, name: bytes) -> tuple[bytes, ...]: + return tuple(value for header_name, value in headers if header_name.lower() == name) + + +def _is_event_stream(headers: HttpHeaders) -> bool: + values = _header_values(headers, b"content-type") + return len(values) == 1 and values[0].partition(b";")[0].strip().lower() == b"text/event-stream" + + +def _has_identity_encoding(headers: HttpHeaders) -> bool: + values = _header_values(headers, b"content-encoding") + return not values or all(token.strip().lower() == b"identity" for value in values for token in value.split(b",")) + + +async def _close_source(source: AsyncIterator[bytes]) -> None: + close = getattr(source, "aclose", None) + if callable(close): + result = close() + if inspect.isawaitable(result): + await result + + +async def _relay(source: AsyncIterator[bytes], expected_length: bytes | None) -> AsyncIterator[bytes]: + observed = 0 + try: + async for chunk in source: + if not isinstance(chunk, bytes): + raise TypeError("A provider stream must yield bytes.") + observed += len(chunk) + if expected_length is not None: + digits = str(observed).encode() + if len(digits) > len(expected_length) or ( + len(digits) == len(expected_length) and digits > expected_length + ): + raise ValueError("The provider stream exceeds its declared Content-Length.") + yield chunk + if expected_length is not None and str(observed).encode() != expected_length: + raise ValueError("The provider stream does not match its declared Content-Length.") + finally: + await _close_source(source) + + +async def execute_streaming_http( + request: BufferedHttpRequest, + *, + dispatch: StreamingHttpDispatch, + checker: ContentChecker, + streaming_policy: StreamBufferingPolicy | None, + input_message: GuardedMessage, + adapter: ProviderStreamAdapter, + render_outcome: StreamOutcomeRenderer, + max_event_bytes: int, + max_pending_bytes: int, +) -> Response: + """Dispatch one prepared stream and preserve its HTTP lifecycle.""" + + if any(type(limit) is not int or limit <= 0 for limit in (max_event_bytes, max_pending_bytes)): + raise ValueError("Streaming HTTP byte limits must be positive.") + if streaming_policy is not None: + validate_streaming_policy(streaming_policy) + + try: + upstream = await dispatch(request) + except Exception as failure: + return _render_response(render_outcome(StreamUpstreamFailed(failure)), request.method) + + if isinstance(upstream, BufferedHttpResponse): + if 200 <= upstream.status_code <= 299: + return _render_response( + render_outcome( + OperationProjectionFailed( + InspectionStage.OUTPUT, + UnsupportedGuardedPayload("A successful streaming response must have an open body."), + ) + ), + request.method, + ) + if len(upstream.body) > max_pending_bytes: + return _render_response(render_outcome(StreamUpstreamFailed(ResponseBodyTooLarge())), request.method) + return _render_response(upstream, request.method) + if not isinstance(upstream, StreamingHttpResponse): + return _render_response( + render_outcome(StreamProcessingFailed(TypeError("Unsupported streaming dispatch result."))), request.method + ) + headers = tuple(_end_to_end_headers(upstream.headers)) + content_lengths = _header_values(headers, b"content-length") + if ( + request.method == "HEAD" + or upstream.status_code in {204, 205} + or not _is_event_stream(headers) + or len(content_lengths) > 1 + or any(re.fullmatch(rb"[0-9]+", length) is None for length in content_lengths) + ): + await _close_source(upstream.body) + return _render_response( + render_outcome( + OperationProjectionFailed( + InspectionStage.OUTPUT, + UnsupportedGuardedPayload("A successful provider stream requires valid SSE body framing."), + ) + ), + request.method, + ) + if streaming_policy is not None and not _has_identity_encoding(upstream.headers): + await _close_source(upstream.body) + return _render_response( + render_outcome( + OperationProjectionFailed( + InspectionStage.OUTPUT, + UnsupportedGuardedPayload("An inspected provider stream must use identity content encoding."), + ) + ), + request.method, + ) + if streaming_policy is None: + expected_length = (content_lengths[0].lstrip(b"0") or b"0") if content_lengths else None + return _streaming_response(upstream, _relay(upstream.body, expected_length), may_modify=False) + + body = guard_provider_stream( + upstream.body, + checker=checker, + streaming_policy=streaming_policy, + input_message=input_message, + adapter=adapter, + render_outcome=render_outcome, + max_event_bytes=max_event_bytes, + max_pending_bytes=max_pending_bytes, + ) + return _streaming_response(upstream, body, may_modify=True) diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py new file mode 100644 index 0000000000..cab9dfd930 --- /dev/null +++ b/tests/server/experimental/test_streaming_http.py @@ -0,0 +1,352 @@ +# 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. + +import json + +import pytest + +from nemoguardrails.server.experimental._buffered_kernel import OperationProjectionFailed +from nemoguardrails.server.experimental._content_checker import ContentAllowed, StreamBufferingPolicy +from nemoguardrails.server.experimental._guarded_stream import StreamProcessingFailed, StreamUpstreamFailed +from nemoguardrails.server.experimental._http_kernel import BufferedHttpRequest, BufferedHttpResponse +from nemoguardrails.server.experimental._streaming_http import StreamingHttpResponse, execute_streaming_http +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent +from nemoguardrails.server.experimental.provider.stream import ( + ClassifiedStreamAdapter, + GuardedStreamEvent, + StreamCapabilityProfile, + StreamEventRole, + StreamProjectionContract, + StreamShapeCoverage, +) +from nemoguardrails.server.experimental.provider.types import GuardedMessage + +CONTRACT = StreamProjectionContract( + projection_id="test.stream", + profile=StreamCapabilityProfile.SINGLE_TEXT_DELTA_V1, + shapes=StreamShapeCoverage( + guarded_shapes=frozenset({"text"}), + opaque_shapes=frozenset({"end"}), + ), +) + + +class Checker: + def __init__(self): + self.calls = [] + + async def check_output(self, check): + self.calls.append(check) + return ContentAllowed() + + +class Classifier: + contract = CONTRACT + + def classify_event(self, event: ServerSentEvent): + if event.data == b"[END]": + return GuardedStreamEvent("end", StreamEventRole.OPAQUE_METADATA) + payload = json.loads(event.data or b"") + return GuardedStreamEvent("text", StreamEventRole.GUARDED_TEXT, payload["text"]) + + +class Hooks: + def observe_event(self, event, classified): + pass + + def validate_end_of_stream(self): + pass + + def encode_error(self, body): + return (b"data: " + body + b"\n\n", b"data: [END]\n\n") + + +class Source: + def __init__(self, chunks): + self.chunks = iter(chunks) + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self.chunks) + except StopIteration: + raise StopAsyncIteration + + async def aclose(self): + self.closed = True + + +REQUEST = BufferedHttpRequest("POST", "/v1/generate", b"/v1/generate", b"", (), b"request") + + +def render_outcome(outcome): + if isinstance(outcome, OperationProjectionFailed): + code = "projection_failed" + elif isinstance(outcome, StreamUpstreamFailed): + code = "upstream_failed" + elif isinstance(outcome, StreamProcessingFailed): + code = "processing_failed" + else: + code = "stream_failure" + return BufferedHttpResponse(502, ((b"content-type", b"text/plain"),), code.encode()) + + +async def response_body(response): + body = getattr(response, "body", None) + if isinstance(body, bytes): + return body + return b"".join([chunk async for chunk in response.body_iterator]) + + +async def execute(dispatch, *, policy=StreamBufferingPolicy(1, 0), checker=None, **limits): + return await execute_streaming_http( + REQUEST, + dispatch=dispatch, + checker=checker or Checker(), + streaming_policy=policy, + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(Classifier(), Hooks()), + render_outcome=render_outcome, + max_event_bytes=limits.get("max_event_bytes", 1024), + max_pending_bytes=limits.get("max_pending_bytes", 1024), + ) + + +@pytest.mark.asyncio +async def test_inspected_stream_preserves_events_and_removes_stale_content_length(): + stream = b'data: {"text":"answer"}\n\ndata: [END]\n\n' + source = Source([stream[:7], stream[7:]]) + checker = Checker() + + async def dispatch(_request): + return StreamingHttpResponse( + 200, + ( + (b"content-type", b"text/event-stream"), + (b"content-length", str(len(stream)).encode()), + (b"x-provider-id", b"stream-id"), + ), + source, + ) + + response = await execute(dispatch, checker=checker) + + assert await response_body(response) == stream + assert (b"x-provider-id", b"stream-id") in response.raw_headers + assert all(name.lower() != b"content-length" for name, _ in response.raw_headers) + assert checker.calls[0].output_content == "answer" + assert source.closed is True + + +@pytest.mark.asyncio +async def test_uninspected_stream_preserves_opaque_bytes_and_original_length(): + stream = b"opaque provider bytes" + source = Source([stream]) + + async def dispatch(_request): + return StreamingHttpResponse( + 200, + ((b"content-type", b"text/event-stream"), (b"content-length", str(len(stream)).encode())), + source, + ) + + response = await execute(dispatch, policy=None) + + assert await response_body(response) == stream + assert (b"content-length", str(len(stream)).encode()) in response.raw_headers + assert source.closed is True + + +@pytest.mark.asyncio +async def test_provider_http_error_is_bounded_and_preserved_without_event_inspection(): + async def dispatch(_request): + return BufferedHttpResponse(429, ((b"content-type", b"application/json"),), b'{"error":"busy"}') + + response = await execute(dispatch) + + assert response.status_code == 429 + assert await response_body(response) == b'{"error":"busy"}' + + oversized = await execute(dispatch, max_pending_bytes=4) + assert oversized.status_code == 502 + assert await response_body(oversized) == b"upstream_failed" + + +@pytest.mark.asyncio +async def test_buffered_success_cannot_bypass_stream_inspection(): + async def dispatch(_request): + return BufferedHttpResponse(200, ((b"content-type", b"text/plain"),), b"unchecked success") + + response = await execute(dispatch) + + assert response.status_code == 502 + assert await response_body(response) == b"projection_failed" + + +@pytest.mark.asyncio +async def test_dispatch_failure_is_rendered_before_streaming_starts(): + async def dispatch(_request): + raise RuntimeError("private upstream detail") + + response = await execute(dispatch) + + assert response.status_code == 502 + assert await response_body(response) == b"upstream_failed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "headers", + [ + ((b"content-type", b"application/json"),), + ((b"content-type", b"text/event-stream"), (b"content-encoding", b"gzip")), + ((b"content-type", b"text/event-stream"), (b"content-type", b"text/event-stream")), + ], +) +async def test_unsupported_inspected_success_closes_source_before_headers(headers): + source = Source([b"hidden provider content"]) + + async def dispatch(_request): + return StreamingHttpResponse(200, headers, source) + + response = await execute(dispatch) + + assert response.status_code == 502 + assert await response_body(response) == b"projection_failed" + assert source.closed is True + + +@pytest.mark.asyncio +async def test_unsafe_stream_policy_is_rejected_before_dispatch(): + calls = [] + + async def dispatch(_request): + calls.append(True) + raise AssertionError("unsafe policy must stop before dispatch") + + with pytest.raises(ValueError, match="cannot release"): + await execute(dispatch, policy=StreamBufferingPolicy(1, 0, True)) + + assert calls == [] + + +@pytest.mark.asyncio +async def test_buffered_provider_error_uses_shared_response_framing(): + body = b'{"error":"busy"}' + + async def dispatch(_request): + return BufferedHttpResponse( + 429, + ( + (b"content-type", b"application/json"), + (b"content-length", b"999"), + (b"connection", b"x-private"), + (b"x-private", b"hidden"), + (b"transfer-encoding", b"chunked"), + (b"set-cookie", b"a=1"), + (b"set-cookie", b"b=2"), + ), + body, + ) + + response = await execute(dispatch) + + assert await response_body(response) == body + assert response.headers["content-length"] == str(len(body)) + assert response.headers.getlist("set-cookie") == ["a=1", "b=2"] + assert all(name not in response.headers for name in ("connection", "transfer-encoding", "x-private")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", [None, StreamBufferingPolicy(1, 0)]) +async def test_stream_filters_connection_headers_and_preserves_provider_headers(policy): + stream = b'data: {"text":"answer"}\n\ndata: [END]\n\n' + source = Source([stream]) + + async def dispatch(_request): + return StreamingHttpResponse( + 200, + ( + (b"content-type", b"text/event-stream"), + (b"connection", b"X-Private"), + (b"x-private", b"hidden"), + (b"transfer-encoding", b"chunked"), + (b"x-provider-id", b"stream-id"), + (b"set-cookie", b"a=1"), + (b"set-cookie", b"b=2"), + ), + source, + ) + + response = await execute(dispatch, policy=policy) + + assert await response_body(response) == stream + assert response.headers["x-provider-id"] == "stream-id" + assert response.headers.getlist("set-cookie") == ["a=1", "b=2"] + assert all(name not in response.headers for name in ("connection", "transfer-encoding", "x-private")) + assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [204, 205]) +async def test_bodyless_success_cannot_be_used_as_a_provider_stream(status): + source = Source([b"hidden"]) + + async def dispatch(_request): + return StreamingHttpResponse(status, ((b"content-type", b"text/event-stream"),), source) + + response = await execute(dispatch) + + assert response.status_code == 502 + assert await response_body(response) == b"projection_failed" + assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lengths", [(b"-1",), (b"3", b"3"), (b"3,3",)]) +async def test_invalid_stream_framing_closes_source_before_headers(lengths): + source = Source([b"hidden"]) + + async def dispatch(_request): + return StreamingHttpResponse( + 200, + ((b"content-type", b"text/event-stream"),) + tuple((b"content-length", value) for value in lengths), + source, + ) + + response = await execute(dispatch, policy=None) + + assert response.status_code == 502 + assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("length", [b"2", b"4"]) +async def test_uninspected_stream_length_mismatch_aborts_and_closes_source(length): + source = Source([b"abc"]) + + async def dispatch(_request): + return StreamingHttpResponse( + 200, ((b"content-type", b"text/event-stream"), (b"content-length", length)), source + ) + + response = await execute(dispatch, policy=None) + + with pytest.raises(ValueError, match="declared Content-Length"): + await response_body(response) + assert source.closed is True From 080857ae293e6c0a2cda5a28935142e9ff48eb13 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:26:47 +0200 Subject: [PATCH 02/11] fix(server): close provider streams however responses end Make the streaming response the owner of the upstream body. It closes the body iterator and the source in a shielded cleanup, so a client disconnect or send failure before the first chunk no longer leaks the provider connection. Reject a non-user input message before dispatch instead of failing after the response starts. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_streaming_http.py | 31 ++- .../experimental/test_streaming_http.py | 227 +++++++++++++++++- 2 files changed, 254 insertions(+), 4 deletions(-) diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index 4f97eb33cb..708f775479 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -20,7 +20,9 @@ from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass +import anyio from starlette.responses import Response, StreamingResponse +from starlette.types import Receive, Scope, Send from nemoguardrails.server.experimental._buffered_kernel import InspectionStage, OperationProjectionFailed from nemoguardrails.server.experimental._content_checker import ContentChecker, StreamBufferingPolicy @@ -46,7 +48,11 @@ @dataclass(frozen=True, slots=True) class StreamingHttpResponse: - """Carry one successful response without choosing an HTTP client.""" + """Carry one successful response without choosing an HTTP client. + + The returned HTTP response closes ``body`` when it ends, possibly more than + once, so closing ``body`` must be idempotent. + """ status_code: int headers: HttpHeaders @@ -64,13 +70,32 @@ def __post_init__(self) -> None: StreamingHttpDispatch = Callable[[BufferedHttpRequest], Awaitable[StreamingHttpResponse | BufferedHttpResponse]] +class _ClosingStreamingResponse(StreamingResponse): + """Close the provider body however the response ends, even before iteration.""" + + def __init__(self, body: AsyncIterator[bytes], source: AsyncIterator[bytes], *, status_code: int): + super().__init__(body, status_code=status_code) + self.owned_body = body + self.source = source + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + try: + await super().__call__(scope, receive, send) + finally: + with anyio.CancelScope(shield=True): + try: + await _close_source(self.owned_body) + finally: + await _close_source(self.source) + + def _streaming_response( value: StreamingHttpResponse, body: AsyncIterator[bytes], *, may_modify: bool, ) -> StreamingResponse: - response = StreamingResponse(body, status_code=value.status_code) + response = _ClosingStreamingResponse(body, value.body, status_code=value.status_code) response.raw_headers = [ (name, header_value) for name, header_value in _end_to_end_headers(value.headers) @@ -137,6 +162,8 @@ async def execute_streaming_http( if any(type(limit) is not int or limit <= 0 for limit in (max_event_bytes, max_pending_bytes)): raise ValueError("Streaming HTTP byte limits must be positive.") + if input_message.role != "user": + raise ValueError("A guarded stream requires a user input message.") if streaming_policy is not None: validate_streaming_policy(streaming_policy) diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index cab9dfd930..9b1e96328d 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -13,16 +13,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import json +import anyio import pytest from nemoguardrails.server.experimental._buffered_kernel import OperationProjectionFailed from nemoguardrails.server.experimental._content_checker import ContentAllowed, StreamBufferingPolicy from nemoguardrails.server.experimental._guarded_stream import StreamProcessingFailed, StreamUpstreamFailed from nemoguardrails.server.experimental._http_kernel import BufferedHttpRequest, BufferedHttpResponse -from nemoguardrails.server.experimental._streaming_http import StreamingHttpResponse, execute_streaming_http -from nemoguardrails.server.experimental.provider.sse import ServerSentEvent +from nemoguardrails.server.experimental._streaming_http import ( + StreamingHttpResponse, + _ClosingStreamingResponse, + execute_streaming_http, +) +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent, iter_sse_events from nemoguardrails.server.experimental.provider.stream import ( ClassifiedStreamAdapter, GuardedStreamEvent, @@ -32,6 +38,9 @@ StreamShapeCoverage, ) from nemoguardrails.server.experimental.provider.types import GuardedMessage +from nemoguardrails.server.experimental.providers.openai.chat_completions.stream_classifier import STREAM_CLASSIFIER +from nemoguardrails.server.experimental.providers.openai.chat_completions.stream_hooks import ChatCompletionsStreamHooks +from nemoguardrails.server.experimental.providers.openai.errors import render_openai_error CONTRACT = StreamProjectionContract( projection_id="test.stream", @@ -77,17 +86,20 @@ class Source: def __init__(self, chunks): self.chunks = iter(chunks) self.closed = False + self.read_calls = 0 def __aiter__(self): return self async def __anext__(self): + self.read_calls += 1 try: return next(self.chunks) except StopIteration: raise StopAsyncIteration async def aclose(self): + await anyio.sleep(0) self.closed = True @@ -350,3 +362,214 @@ async def dispatch(_request): with pytest.raises(ValueError, match="declared Content-Length"): await response_body(response) assert source.closed is True + + +@pytest.mark.asyncio +async def test_assistant_input_message_is_rejected_before_dispatch(): + async def dispatch(_request): + raise AssertionError("dispatch must not run") + + with pytest.raises(ValueError, match="user input message"): + await execute_streaming_http( + REQUEST, + dispatch=dispatch, + checker=Checker(), + streaming_policy=StreamBufferingPolicy(1, 0), + input_message=GuardedMessage("assistant", "answer"), + adapter=ClassifiedStreamAdapter(Classifier(), Hooks()), + render_outcome=render_outcome, + max_event_bytes=1024, + max_pending_bytes=1024, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", [None, StreamBufferingPolicy(1, 0)]) +@pytest.mark.parametrize("spec_version", ["2.0", "2.4"]) +@pytest.mark.parametrize("failure_type", ["http.response.start", "http.response.body"]) +async def test_http_send_failure_closes_source(policy, spec_version, failure_type): + source = Source([b'data: {"text":"answer"}\n\n']) + + async def dispatch(_request): + return StreamingHttpResponse(200, ((b"content-type", b"text/event-stream"),), source) + + response = await execute(dispatch, policy=policy) + + async def receive(): + await asyncio.Event().wait() + + async def send(message): + if message["type"] == failure_type: + raise OSError("downstream disconnected") + + with pytest.raises(Exception): + await response({"type": "http", "asgi": {"spec_version": spec_version}}, receive, send) + + assert source.closed is True + if failure_type == "http.response.start": + assert source.read_calls == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", [None, StreamBufferingPolicy(1, 0)]) +async def test_http_disconnect_before_first_iteration_closes_source(policy): + source = Source([b"unused"]) + started = asyncio.Event() + + async def dispatch(_request): + return StreamingHttpResponse(200, ((b"content-type", b"text/event-stream"),), source) + + response = await execute(dispatch, policy=policy) + + async def receive(): + await started.wait() + return {"type": "http.disconnect"} + + async def send(_message): + started.set() + await asyncio.Event().wait() + + await response({"type": "http", "asgi": {"spec_version": "2.0"}}, receive, send) + + assert source.closed is True + assert source.read_calls == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", [None, StreamBufferingPolicy(1, 0)]) +async def test_http_task_cancellation_before_first_iteration_closes_source(policy): + source = Source([b"unused"]) + started = asyncio.Event() + + async def dispatch(_request): + return StreamingHttpResponse(200, ((b"content-type", b"text/event-stream"),), source) + + response = await execute(dispatch, policy=policy) + + async def receive(): + await asyncio.Event().wait() + + async def send(_message): + started.set() + await asyncio.Event().wait() + + task = asyncio.create_task(response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send)) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert source.closed is True + assert source.read_calls == 0 + + +def openai_block(text, ending=b"\n\n"): + payload = {"object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {"content": text}}]} + return b"data: " + json.dumps(payload).encode() + ending + + +async def openai_response(source, checker, policy=StreamBufferingPolicy(2, 0)): + async def dispatch(_request): + return StreamingHttpResponse(200, ((b"content-type", b"text/event-stream"),), source) + + return await execute_streaming_http( + REQUEST, + dispatch=dispatch, + checker=checker, + streaming_policy=policy, + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(STREAM_CLASSIFIER, ChatCompletionsStreamHooks()), + render_outcome=render_openai_error, + max_event_bytes=1024, + max_pending_bytes=4096, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ending", [b"\n\r\n", b"\r\n\n", b"\r\r"]) +@pytest.mark.parametrize("fragmented", [False, True]) +async def test_openai_classifier_receives_correct_blocks_and_preserves_bom_and_metadata(ending, fragmented): + raw = ( + b"\xef\xbb\xbf" + + openai_block("one ", ending) + + b": comment\n\n\n" + + openai_block("two", ending) + + b"data: [DONE]\n\n" + ) + chunks = [bytes([byte]) for byte in raw] if fragmented else [raw] + source = Source(chunks) + checker = Checker() + + response = await openai_response(source, checker) + + assert await response_body(response) == raw + assert [call.output_content for call in checker.calls] == ["one two"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tail", [openai_block("hidden")[:-2], b"event: unknown", b" "]) +async def test_openai_truncation_hides_pending_text_and_emits_complete_native_error(tail): + source = Source([openai_block("pending"), tail]) + checker = Checker() + + response = await openai_response(source, checker) + result = await response_body(response) + + assert b"pending" not in result + assert b"hidden" not in result + assert result.startswith(b'data: {"error":') + assert result.endswith(b"\n\ndata: [DONE]\n\n") + assert checker.calls == [] + + +@pytest.mark.asyncio +async def test_openai_comment_tail_is_dropped_after_approved_complete_block(): + block = openai_block("answer", b"\n\r\n") + b"data: [DONE]\n\n" + source = Source([block, b": unfinished comment"]) + checker = Checker() + + response = await openai_response(source, checker) + + assert await response_body(response) == block + assert [call.output_content for call in checker.calls] == ["answer"] + + +@pytest.mark.asyncio +async def test_openai_error_encoding_is_complete_sse_blocks(): + chunks = ChatCompletionsStreamHooks().encode_error(b'{"error":{"message":"blocked"}}') + + async def source(): + for chunk in chunks: + yield chunk + + events = [event async for event in iter_sse_events(source())] + + assert b"".join(event.raw for event in events) == b"".join(chunks) + assert [event.data for event in events] == [b'{"error":{"message":"blocked"}}', b"[DONE]"] + + +@pytest.mark.asyncio +async def test_source_is_closed_when_closing_the_body_fails(): + class FailingBody: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + async def aclose(self): + raise RuntimeError("body cleanup failed") + + source = Source([]) + response = _ClosingStreamingResponse(FailingBody(), source, status_code=200) + + async def receive(): + await asyncio.Event().wait() + + async def send(_message): + pass + + with pytest.raises(RuntimeError, match="body cleanup failed"): + await response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send) + + assert source.closed is True From 68f43fad4001c8d7b08d142af9e5c3683f7fda4b Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 6 Oct 2026 19:17:44 +0200 Subject: [PATCH 03/11] fix(server): drop body-dependent headers from inspected streams An inspected stream may withhold events, emit a provider-native error, or drop an unfinished tail, so validators computed over the provider body no longer describe the bytes sent downstream. Remove ETag, Content-MD5, Digest, Content-Digest, and Repr-Digest alongside Content-Length whenever the stream may change. Uninspected relays keep every end-to-end header. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_streaming_http.py | 7 ++++- .../experimental/test_streaming_http.py | 27 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index 708f775479..dedf1b9fd0 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -89,6 +89,11 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: await _close_source(self.source) +_BODY_DEPENDENT_HEADERS = frozenset( + {b"content-length", b"content-md5", b"content-digest", b"digest", b"etag", b"repr-digest"} +) + + def _streaming_response( value: StreamingHttpResponse, body: AsyncIterator[bytes], @@ -99,7 +104,7 @@ def _streaming_response( response.raw_headers = [ (name, header_value) for name, header_value in _end_to_end_headers(value.headers) - if not may_modify or name.lower() != b"content-length" + if not may_modify or name.lower() not in _BODY_DEPENDENT_HEADERS ] return response diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index 9b1e96328d..8bc38635f6 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -573,3 +573,30 @@ async def send(_message): await response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send) assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", [None, StreamBufferingPolicy(1, 0)]) +async def test_body_dependent_headers_are_removed_only_when_the_stream_may_change(policy): + stream = b'data: {"text":"answer"}\n\ndata: [END]\n\n' + body_dependent = ( + (b"ETag", b'"provider"'), + (b"content-md5", b"md5"), + (b"digest", b"sha-256=digest"), + (b"content-digest", b"sha-256=:digest:"), + (b"repr-digest", b"sha-256=:digest:"), + ) + + async def dispatch(_request): + return StreamingHttpResponse( + 200, + ((b"content-type", b"text/event-stream"), (b"x-provider-id", b"stream-id"), *body_dependent), + Source([stream]), + ) + + response = await execute(dispatch, policy=policy) + + assert await response_body(response) == stream + assert (b"x-provider-id", b"stream-id") in response.raw_headers + forwarded = [header for header in body_dependent if header in response.raw_headers] + assert forwarded == ([] if policy is not None else list(body_dependent)) From 07e54786b4e1c8e87e104f459ccadb41b34e0e5f Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:04:30 +0200 Subject: [PATCH 04/11] fix(server): request identity encoding for inspected streams Clients such as the OpenAI SDK send Accept-Encoding: gzip, and inspected streams reject a provider stream that is not identity-encoded. A compressing provider therefore turned every inspected streaming request into a proxy error. When the stream is inspected, replace Accept-Encoding with identity before dispatch, matching buffered output inspection. Uninspected relays forward the client's preference unchanged, and a still-encoded inspected stream remains rejected. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_streaming_http.py | 5 +++- .../experimental/test_streaming_http.py | 28 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index dedf1b9fd0..fb6b4cc490 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -18,7 +18,7 @@ import inspect import re from collections.abc import AsyncIterator, Awaitable, Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace import anyio from starlette.responses import Response, StreamingResponse @@ -41,6 +41,7 @@ ResponseBodyTooLarge, _end_to_end_headers, _render_response, + _request_identity_encoding, ) from nemoguardrails.server.experimental.provider.stream import ProviderStreamAdapter from nemoguardrails.server.experimental.provider.types import GuardedMessage @@ -171,6 +172,8 @@ async def execute_streaming_http( raise ValueError("A guarded stream requires a user input message.") if streaming_policy is not None: validate_streaming_policy(streaming_policy) + # An encoded stream cannot be inspected; uninspected relays keep the client's preference. + request = replace(request, headers=_request_identity_encoding(request.headers)) try: upstream = await dispatch(request) diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index 8bc38635f6..6e02d98546 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -15,6 +15,7 @@ import asyncio import json +from dataclasses import replace import anyio import pytest @@ -600,3 +601,30 @@ async def dispatch(_request): assert (b"x-provider-id", b"stream-id") in response.raw_headers forwarded = [header for header in body_dependent if header in response.raw_headers] assert forwarded == ([] if policy is not None else list(body_dependent)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("policy", "expected"), [(StreamBufferingPolicy(1, 0), [b"identity"]), (None, [b"gzip, br"])]) +async def test_identity_encoding_is_requested_only_for_inspected_streams(policy, expected): + """Inspected streams need readable bytes; uninspected relays forward the client's preference.""" + dispatched = [] + + async def dispatch(request): + dispatched.append(request) + return StreamingHttpResponse(200, ((b"content-type", b"text/event-stream"),), Source([b"data: [END]\n\n"])) + + request = replace(REQUEST, headers=((b"accept-encoding", b"gzip, br"),)) + response = await execute_streaming_http( + request, + dispatch=dispatch, + checker=Checker(), + streaming_policy=policy, + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(Classifier(), Hooks()), + render_outcome=render_outcome, + max_event_bytes=1024, + max_pending_bytes=1024, + ) + await response_body(response) + + assert [value for name, value in dispatched[0].headers if name.lower() == b"accept-encoding"] == expected From a6cfdb9c046b424b573a7367b76e809d0d8fb8b7 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:32:17 +0200 Subject: [PATCH 05/11] fix(server): shield provider stream rejection cleanup Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_streaming_http.py | 12 +++--- .../experimental/test_streaming_http.py | 41 +++++++++++++++++++ 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index fb6b4cc490..2dc904d983 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -125,11 +125,13 @@ def _has_identity_encoding(headers: HttpHeaders) -> bool: async def _close_source(source: AsyncIterator[bytes]) -> None: - close = getattr(source, "aclose", None) - if callable(close): - result = close() - if inspect.isawaitable(result): - await result + """Finish owned-body cleanup even when a request cancel scope has fired.""" + with anyio.CancelScope(shield=True): + close = getattr(source, "aclose", None) + if callable(close): + result = close() + if inspect.isawaitable(result): + await result async def _relay(source: AsyncIterator[bytes], expected_length: bytes | None) -> AsyncIterator[bytes]: diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index 6e02d98546..f28ce80da3 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -244,6 +244,47 @@ async def dispatch(_request): assert source.closed is True +@pytest.mark.asyncio +@pytest.mark.parametrize("encoded", [False, True]) +async def test_rejection_cleanup_finishes_under_cancellation(encoded): + """Cancellation during pre-response rejection cannot abandon an open body.""" + started = anyio.Event() + release = anyio.Event() + + class DelayedClose(Source): + async def aclose(self): + started.set() + await release.wait() + self.closed = True + + source = DelayedClose([b"unread provider bytes"]) + headers = ( + ((b"content-type", b"text/event-stream"), (b"content-encoding", b"gzip")) + if encoded + else ((b"content-type", b"text/plain"),) + ) + + async def dispatch(_request): + return StreamingHttpResponse(200, headers, source) + + scope = anyio.CancelScope() + + async def request(): + with scope: + await execute(dispatch) + + with anyio.fail_after(2): + async with anyio.create_task_group() as group: + group.start_soon(request) + await started.wait() + scope.cancel() + await anyio.sleep(0) + release.set() + + assert source.closed is True + assert source.read_calls == 0 + + @pytest.mark.asyncio async def test_unsafe_stream_policy_is_rejected_before_dispatch(): calls = [] From 2fe2f6692210806fb5ed48f8e39dbd6b88117ac4 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:10:22 +0200 Subject: [PATCH 06/11] fix(server): bound shielded stream cleanup waits Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_streaming_http.py | 31 +++-- .../experimental/test_streaming_http.py | 111 +++++++++++++++++- 2 files changed, 131 insertions(+), 11 deletions(-) diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index 2dc904d983..4e641cf209 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -16,6 +16,7 @@ """Bind guarded provider streams to an injected HTTP exchange.""" import inspect +import logging import re from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass, replace @@ -46,13 +47,21 @@ from nemoguardrails.server.experimental.provider.stream import ProviderStreamAdapter from nemoguardrails.server.experimental.provider.types import GuardedMessage +log = logging.getLogger(__name__) + +_STREAM_CLOSE_TIMEOUT_SECONDS = 5.0 + @dataclass(frozen=True, slots=True) class StreamingHttpResponse: """Carry one successful response without choosing an HTTP client. The returned HTTP response closes ``body`` when it ends, possibly more than - once, so closing ``body`` must be idempotent. + once, so closing ``body`` must be idempotent. HTTP cleanup attempts are + shielded from enclosing cancellation scopes for up to five seconds each. + A timeout is logged without replacing the response or its original failure. + Bodies must cooperate with cancellation; cleanup cannot forcibly stop + synchronous blocking code or a closer that suppresses cancellation. """ status_code: int @@ -83,11 +92,10 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: try: await super().__call__(scope, receive, send) finally: - with anyio.CancelScope(shield=True): - try: - await _close_source(self.owned_body) - finally: - await _close_source(self.source) + try: + await _close_source(self.owned_body) + finally: + await _close_source(self.source) _BODY_DEPENDENT_HEADERS = frozenset( @@ -125,13 +133,20 @@ def _has_identity_encoding(headers: HttpHeaders) -> bool: async def _close_source(source: AsyncIterator[bytes]) -> None: - """Finish owned-body cleanup even when a request cancel scope has fired.""" - with anyio.CancelScope(shield=True): + """Bound shielded cleanup, logging its timeout without masking the outcome. + + Only this scope's timeout is suppressed. Other close failures propagate as + before. Each call gets its own budget, including a final provider-body close + after the response iterator's cleanup, so the limit is not per request. + """ + with anyio.move_on_after(_STREAM_CLOSE_TIMEOUT_SECONDS, shield=True) as scope: close = getattr(source, "aclose", None) if callable(close): result = close() if inspect.isawaitable(result): await result + if scope.cancelled_caught: + log.warning("Stream cleanup timed out; resources may remain open.") async def _relay(source: AsyncIterator[bytes], expected_length: bytes | None) -> AsyncIterator[bytes]: diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index f28ce80da3..9695b151fe 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -19,7 +19,9 @@ import anyio import pytest +from starlette.requests import ClientDisconnect +from nemoguardrails.server.experimental import _streaming_http from nemoguardrails.server.experimental._buffered_kernel import OperationProjectionFailed from nemoguardrails.server.experimental._content_checker import ContentAllowed, StreamBufferingPolicy from nemoguardrails.server.experimental._guarded_stream import StreamProcessingFailed, StreamUpstreamFailed @@ -104,6 +106,20 @@ async def aclose(self): self.closed = True +class StalledClose(Source): + def __init__(self): + super().__init__([]) + self.close_started = anyio.Event() + self.close_cancelled = False + + async def aclose(self): + self.close_started.set() + try: + await anyio.sleep_forever() + finally: + self.close_cancelled = True + + REQUEST = BufferedHttpRequest("POST", "/v1/generate", b"/v1/generate", b"", (), b"request") @@ -285,6 +301,94 @@ async def request(): assert source.read_calls == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("encoded", [False, True]) +@pytest.mark.parametrize("cancel_request", [False, True]) +async def test_stalled_rejection_cleanup_preserves_rejection(monkeypatch, caplog, encoded, cancel_request): + monkeypatch.setattr(_streaming_http, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01, raising=False) + source = StalledClose() + headers = ( + ((b"content-type", b"text/event-stream"), (b"content-encoding", b"gzip")) + if encoded + else ((b"content-type", b"text/plain"),) + ) + + async def dispatch(_request): + return StreamingHttpResponse(200, headers, source) + + scope = anyio.CancelScope() + + async def request(): + with scope: + return await execute(dispatch) + + task = asyncio.create_task(request()) + await source.close_started.wait() + if cancel_request: + scope.cancel() + response = await asyncio.wait_for(task, timeout=1) + + assert response.status_code == 502 + assert await response_body(response) == b"projection_failed" + assert source.close_cancelled is True + assert source.closed is False + assert source.read_calls == 0 + assert "Stream cleanup timed out; resources may remain open." in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stalled_body", [False, True]) +async def test_stalled_response_cleanup_preserves_send_failure(monkeypatch, caplog, stalled_body): + monkeypatch.setattr(_streaming_http, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01, raising=False) + body = StalledClose() if stalled_body else Source([]) + source = Source([]) if stalled_body else StalledClose() + response = _ClosingStreamingResponse(body, source, status_code=200) + + async def receive(): + await anyio.sleep_forever() + + async def send(_message): + raise OSError("downstream disconnected") + + with pytest.raises(ClientDisconnect) as failure: + await asyncio.wait_for(response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send), timeout=1) + + assert isinstance(failure.value.__context__, OSError) + assert str(failure.value.__context__) == "downstream disconnected" + stalled, healthy = (body, source) if stalled_body else (source, body) + assert stalled.close_cancelled is True + assert healthy.closed is True + assert "Stream cleanup timed out; resources may remain open." in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stalled_body", [False, True]) +async def test_stalled_response_cleanup_preserves_task_cancellation(monkeypatch, caplog, stalled_body): + monkeypatch.setattr(_streaming_http, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01) + body = StalledClose() if stalled_body else Source([]) + source = Source([]) if stalled_body else StalledClose() + response = _ClosingStreamingResponse(body, source, status_code=200) + started = anyio.Event() + + async def receive(): + await anyio.sleep_forever() + + async def send(_message): + started.set() + await anyio.sleep_forever() + + task = asyncio.create_task(response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send)) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + + stalled, healthy = (body, source) if stalled_body else (source, body) + assert stalled.close_cancelled is True + assert healthy.closed is True + assert "Stream cleanup timed out; resources may remain open." in caplog.text + + @pytest.mark.asyncio async def test_unsafe_stream_policy_is_rejected_before_dispatch(): calls = [] @@ -591,7 +695,8 @@ async def source(): @pytest.mark.asyncio -async def test_source_is_closed_when_closing_the_body_fails(): +@pytest.mark.parametrize("failure_type", [RuntimeError, TimeoutError]) +async def test_source_is_closed_when_closing_the_body_fails(failure_type): class FailingBody: def __aiter__(self): return self @@ -600,7 +705,7 @@ async def __anext__(self): raise StopAsyncIteration async def aclose(self): - raise RuntimeError("body cleanup failed") + raise failure_type("body cleanup failed") source = Source([]) response = _ClosingStreamingResponse(FailingBody(), source, status_code=200) @@ -611,7 +716,7 @@ async def receive(): async def send(_message): pass - with pytest.raises(RuntimeError, match="body cleanup failed"): + with pytest.raises(failure_type, match="body cleanup failed"): await response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send) assert source.closed is True From f255129a4007b0c9d5013ce0c9fdb1be97c46f8f Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 18:22:15 +0200 Subject: [PATCH 07/11] fix(server): apply one Content-Encoding rule to guarded bodies Buffered requests and responses read only the last Content-Encoding header, so "gzip" followed by "identity" passed as unencoded, while streams combined every header but rejected an empty value. Both paths now share one rule: duplicate headers and comma lists form one coding list, empty elements are ignored, and the body is unencoded only when every coding is identity. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_proxy.py | 7 +++--- .../server/experimental/_http_kernel.py | 16 +++++++++++++ .../server/experimental/_streaming_http.py | 6 +---- tests/server/experimental/test_http_kernel.py | 24 +++++++++++++++++++ .../test_openai_chat_integration.py | 13 ++++++++++ .../experimental/test_streaming_http.py | 22 +++++++++++++++++ 6 files changed, 79 insertions(+), 9 deletions(-) diff --git a/nemoguardrails/server/experimental/_guarded_proxy.py b/nemoguardrails/server/experimental/_guarded_proxy.py index c5a774d52b..865490a250 100644 --- a/nemoguardrails/server/experimental/_guarded_proxy.py +++ b/nemoguardrails/server/experimental/_guarded_proxy.py @@ -31,6 +31,7 @@ BufferedHttpResponse, GuardedHttpOperation, GuardedOperationPath, + _has_identity_encoding, ) from nemoguardrails.server.experimental._json_payload import InvalidJson, UnsupportedJsonShape, parse_json_object from nemoguardrails.server.experimental.provider.endpoint import GuardedJsonEndpoint @@ -67,8 +68,7 @@ def _require_json_representation(request: BufferedHttpRequest) -> None: "Content-Type must be application/json.", "unsupported_media_type", ) - content_encoding = _header(request.headers, b"content-encoding") - if content_encoding is not None and content_encoding.strip().lower() not in {"", "identity"}: + if not _has_identity_encoding(request.headers): raise UnsupportedGuardedRepresentation( "Encoded request bodies are not supported by the guarded path.", "unsupported_content_encoding", @@ -137,8 +137,7 @@ def project_response( raise UnsupportedGuardedPayload( "The successful provider response must use application/json.", endpoint.unsupported_response_code ) - content_encoding = _header(payload.headers, b"content-encoding") - if content_encoding is not None and content_encoding.strip().lower() not in {"", "identity"}: + if not _has_identity_encoding(payload.headers): raise UnsupportedGuardedPayload( "Encoded successful provider responses are not supported.", endpoint.unsupported_response_code ) diff --git a/nemoguardrails/server/experimental/_http_kernel.py b/nemoguardrails/server/experimental/_http_kernel.py index 82a7c51146..05aa9061f4 100644 --- a/nemoguardrails/server/experimental/_http_kernel.py +++ b/nemoguardrails/server/experimental/_http_kernel.py @@ -409,6 +409,22 @@ def _render_response(value: BufferedHttpResponse, request_method: str | None = N return response +def _has_identity_encoding(headers: HttpHeaders) -> bool: + """Return whether every Content-Encoding coding is identity. + + Duplicate headers and comma lists combine into one coding list, so a body + is unencoded only when no coding other than identity appears anywhere. + Empty list elements carry no coding and are ignored. + """ + return all( + token.strip().lower() == b"identity" + for name, value in headers + if name.lower() == b"content-encoding" + for token in value.split(b",") + if token.strip() + ) + + def _request_identity_encoding(headers: HttpHeaders) -> HttpHeaders: """Ask the provider for an unencoded response that output checks can read. diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index 4e641cf209..c68761e677 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -41,6 +41,7 @@ HttpHeaders, ResponseBodyTooLarge, _end_to_end_headers, + _has_identity_encoding, _render_response, _request_identity_encoding, ) @@ -127,11 +128,6 @@ def _is_event_stream(headers: HttpHeaders) -> bool: return len(values) == 1 and values[0].partition(b";")[0].strip().lower() == b"text/event-stream" -def _has_identity_encoding(headers: HttpHeaders) -> bool: - values = _header_values(headers, b"content-encoding") - return not values or all(token.strip().lower() == b"identity" for value in values for token in value.split(b",")) - - async def _close_source(source: AsyncIterator[bytes]) -> None: """Bound shielded cleanup, logging its timeout without masking the outcome. diff --git a/tests/server/experimental/test_http_kernel.py b/tests/server/experimental/test_http_kernel.py index 5b1e5fc420..ad4be6821f 100644 --- a/tests/server/experimental/test_http_kernel.py +++ b/tests/server/experimental/test_http_kernel.py @@ -49,6 +49,7 @@ HttpOperationFailed, HttpRouteRejected, HttpRouteRejectionKind, + _has_identity_encoding, _render_response, _request_path, create_http_proxy_router, @@ -1473,3 +1474,26 @@ def test_guarded_operation_requires_both_request_adapters(guarded_operation, ada operation=guarded_operation.operation, **{adapter: lambda request: request}, ) + + +@pytest.mark.parametrize( + ("values", "identity"), + [ + ((), True), + ((b"",), True), + ((b"identity",), True), + ((b" IDENTITY ",), True), + ((b"identity, ,identity",), True), + ((b"identity", b""), True), + ((b"gzip",), False), + ((b"gzip, identity",), False), + ((b"identity, gzip",), False), + ((b"identity", b"gzip"), False), + ((b"gzip", b"identity"), False), + ((b"br", b""), False), + ], +) +def test_identity_encoding_combines_duplicate_headers_and_comma_lists(values, identity): + """Any non-identity coding in any Content-Encoding header marks the body encoded.""" + headers = ((b"content-type", b"application/json"), *((b"Content-Encoding", value) for value in values)) + assert _has_identity_encoding(headers) is identity diff --git a/tests/server/experimental/test_openai_chat_integration.py b/tests/server/experimental/test_openai_chat_integration.py index 62ff16be72..144413eb0d 100644 --- a/tests/server/experimental/test_openai_chat_integration.py +++ b/tests/server/experimental/test_openai_chat_integration.py @@ -498,6 +498,12 @@ async def dispatch(request): 415, "unsupported_content_encoding", ), + ( + _request_body(), + [("content-type", "application/json"), ("content-encoding", "gzip"), ("content-encoding", "identity")], + 415, + "unsupported_content_encoding", + ), ], ) async def test_request_representation_failures_are_native_and_stop_before_dispatch( @@ -557,6 +563,13 @@ def test_openai_error_preserves_provider_binding_failure_code(): 502, "example_response_shape", ), + ( + "response", + _response_body(), + _json_headers((b"content-encoding", b"gzip"), (b"content-encoding", b"identity")), + 502, + "example_response_shape", + ), ], ) async def test_endpoint_error_codes_control_guarded_http_failures(stage, body, headers, status, code): diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index 9695b151fe..5d961d0630 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -182,6 +182,27 @@ async def dispatch(_request): assert source.closed is True +@pytest.mark.asyncio +@pytest.mark.parametrize("encoding", [b"", b"identity", b" Identity , "]) +async def test_inspected_stream_accepts_identity_or_empty_content_encoding(encoding): + """Empty Content-Encoding elements carry no coding, as in buffered bodies.""" + stream = b'data: {"text":"answer"}\n\ndata: [END]\n\n' + checker = Checker() + + async def dispatch(_request): + return StreamingHttpResponse( + 200, + ((b"content-type", b"text/event-stream"), (b"content-encoding", encoding)), + Source([stream]), + ) + + response = await execute(dispatch, checker=checker) + + assert response.status_code == 200 + assert await response_body(response) == stream + assert checker.calls[0].output_content == "answer" + + @pytest.mark.asyncio async def test_uninspected_stream_preserves_opaque_bytes_and_original_length(): stream = b"opaque provider bytes" @@ -244,6 +265,7 @@ async def dispatch(_request): [ ((b"content-type", b"application/json"),), ((b"content-type", b"text/event-stream"), (b"content-encoding", b"gzip")), + ((b"content-type", b"text/event-stream"), (b"content-encoding", b"gzip"), (b"content-encoding", b"identity")), ((b"content-type", b"text/event-stream"), (b"content-type", b"text/event-stream")), ], ) From b4f7bb8c55f374ce8e588412751c9d60984eb242 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:12:38 +0200 Subject: [PATCH 08/11] fix(server): reject ambiguous guarded Content-Type headers Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_proxy.py | 16 ++-------- .../server/experimental/_http_kernel.py | 10 +++++++ .../server/experimental/_streaming_http.py | 8 ++--- .../openai/chat_completions/README.md | 6 ++++ tests/server/experimental/test_http_kernel.py | 20 +++++++++++++ .../test_openai_chat_integration.py | 30 +++++++++++++++++++ 6 files changed, 71 insertions(+), 19 deletions(-) diff --git a/nemoguardrails/server/experimental/_guarded_proxy.py b/nemoguardrails/server/experimental/_guarded_proxy.py index 865490a250..a59325f907 100644 --- a/nemoguardrails/server/experimental/_guarded_proxy.py +++ b/nemoguardrails/server/experimental/_guarded_proxy.py @@ -32,6 +32,7 @@ GuardedHttpOperation, GuardedOperationPath, _has_identity_encoding, + _has_media_type, ) from nemoguardrails.server.experimental._json_payload import InvalidJson, UnsupportedJsonShape, parse_json_object from nemoguardrails.server.experimental.provider.endpoint import GuardedJsonEndpoint @@ -50,20 +51,9 @@ class GuardableProviderRequest: target: GuardedMessageTarget -def _header(headers: tuple[tuple[bytes, bytes], ...], name: bytes) -> str | None: - """Return the last value of a case-insensitive HTTP header.""" - values = [value.decode("latin-1") for key, value in headers if key.lower() == name] - return values[-1] if values else None - - -def _media_type(value: str | None) -> str | None: - """Return a normalized media type without parameters.""" - return value.split(";", 1)[0].strip().lower() if value is not None else None - - def _require_json_representation(request: BufferedHttpRequest) -> None: """Require an unencoded JSON representation for guarded request parsing.""" - if _media_type(_header(request.headers, b"content-type")) != "application/json": + if not _has_media_type(request.headers, b"application/json"): raise UnsupportedGuardedRepresentation( "Content-Type must be application/json.", "unsupported_media_type", @@ -133,7 +123,7 @@ def project_response( """Select guarded output text from a successful JSON response.""" if not 200 <= payload.status_code < 300: return ContentInspectionNotApplicable() - if _media_type(_header(payload.headers, b"content-type")) != "application/json": + if not _has_media_type(payload.headers, b"application/json"): raise UnsupportedGuardedPayload( "The successful provider response must use application/json.", endpoint.unsupported_response_code ) diff --git a/nemoguardrails/server/experimental/_http_kernel.py b/nemoguardrails/server/experimental/_http_kernel.py index 05aa9061f4..a6c1e36cca 100644 --- a/nemoguardrails/server/experimental/_http_kernel.py +++ b/nemoguardrails/server/experimental/_http_kernel.py @@ -409,6 +409,16 @@ def _render_response(value: BufferedHttpResponse, request_method: str | None = N return response +def _has_media_type(headers: HttpHeaders, expected: bytes) -> bool: + """Require one Content-Type with the expected lowercase media type. + + Header names and media types are case-insensitive; parameters are ignored. + Reject even identical duplicates so readers cannot select different values. + """ + values = [value for name, value in headers if name.lower() == b"content-type"] + return len(values) == 1 and values[0].partition(b";")[0].strip().lower() == expected + + def _has_identity_encoding(headers: HttpHeaders) -> bool: """Return whether every Content-Encoding coding is identity. diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index c68761e677..3a463b30ef 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -42,6 +42,7 @@ ResponseBodyTooLarge, _end_to_end_headers, _has_identity_encoding, + _has_media_type, _render_response, _request_identity_encoding, ) @@ -123,11 +124,6 @@ def _header_values(headers: HttpHeaders, name: bytes) -> tuple[bytes, ...]: return tuple(value for header_name, value in headers if header_name.lower() == name) -def _is_event_stream(headers: HttpHeaders) -> bool: - values = _header_values(headers, b"content-type") - return len(values) == 1 and values[0].partition(b";")[0].strip().lower() == b"text/event-stream" - - async def _close_source(source: AsyncIterator[bytes]) -> None: """Bound shielded cleanup, logging its timeout without masking the outcome. @@ -216,7 +212,7 @@ async def execute_streaming_http( if ( request.method == "HEAD" or upstream.status_code in {204, 205} - or not _is_event_stream(headers) + or not _has_media_type(headers, b"text/event-stream") or len(content_lengths) > 1 or any(re.fullmatch(rb"[0-9]+", length) is None for length in content_lengths) ): diff --git a/nemoguardrails/server/experimental/providers/openai/chat_completions/README.md b/nemoguardrails/server/experimental/providers/openai/chat_completions/README.md index ba1eae510e..cebb42ba79 100644 --- a/nemoguardrails/server/experimental/providers/openai/chat_completions/README.md +++ b/nemoguardrails/server/experimental/providers/openai/chat_completions/README.md @@ -151,6 +151,12 @@ unchanged without output checks, but only after the full response passes the closed projection. Null content alongside refusal, reasoning, tool, or citation content is still rejected. +Guarded requests require exactly one `Content-Type: application/json` header. +Inspected successful buffered responses require the same; inspected streams +require exactly one `Content-Type: text/event-stream` header. Media types are +case-insensitive and may have parameters. Duplicate headers are rejected even +when their values agree, rather than selecting a first or last value. + When output checks run, guarded requests ask the provider for `accept-encoding: identity`, replacing any client value, because an encoded response cannot be inspected. A successful response that is still encoded is diff --git a/tests/server/experimental/test_http_kernel.py b/tests/server/experimental/test_http_kernel.py index ad4be6821f..f544c0c0c5 100644 --- a/tests/server/experimental/test_http_kernel.py +++ b/tests/server/experimental/test_http_kernel.py @@ -50,6 +50,7 @@ HttpRouteRejected, HttpRouteRejectionKind, _has_identity_encoding, + _has_media_type, _render_response, _request_path, create_http_proxy_router, @@ -57,6 +58,25 @@ from nemoguardrails.server.experimental.provider.types import GuardedMessage +@pytest.mark.parametrize( + ("values", "accepted"), + [ + ((), False), + ((b"",), False), + ((b"text/plain",), False), + ((b"application/json",), True), + ((b"Application/JSON; charset=utf-8",), True), + ((b"application/json, application/json",), False), + ((b"application/json", b"application/json"), False), + ((b"application/json", b"text/plain"), False), + ((b"text/plain", b"application/json"), False), + ], +) +def test_media_type_requires_one_unambiguous_header(values, accepted): + headers = tuple((b"Content-Type", value) for value in values) + assert _has_media_type(headers, b"application/json") is accepted + + class StaticChecker: """Provide configurable content-check results and observable calls.""" diff --git a/tests/server/experimental/test_openai_chat_integration.py b/tests/server/experimental/test_openai_chat_integration.py index 144413eb0d..c3c02e3191 100644 --- a/tests/server/experimental/test_openai_chat_integration.py +++ b/tests/server/experimental/test_openai_chat_integration.py @@ -552,10 +552,40 @@ def test_openai_error_preserves_provider_binding_failure_code(): ("request", b'{"messages":[],"stream":true}', _json_headers(), 422, "example_request_shape"), ("request", b'{"private-payload":', _json_headers(), 400, "invalid_json"), ("request", _request_body(), ((b"content-type", b"text/plain"),), 415, "unsupported_media_type"), + ( + "request", + _request_body(), + _json_headers((b"Content-Type", b"application/json")), + 415, + "unsupported_media_type", + ), + ("request", _request_body(), _json_headers((b"Content-Type", b"text/plain")), 415, "unsupported_media_type"), + ( + "request", + _request_body(), + ((b"Content-Type", b"text/plain"), *_json_headers()), + 415, + "unsupported_media_type", + ), ("response", b'{"choices":[]}', _json_headers(), 502, "example_response_shape"), ("response", b'["private-payload"]', _json_headers(), 502, "example_response_shape"), ("response", b'{"private-payload":', _json_headers(), 502, "example_response_shape"), ("response", b"private-payload", ((b"content-type", b"text/plain"),), 502, "example_response_shape"), + ( + "response", + _response_body(), + _json_headers((b"Content-Type", b"application/json")), + 502, + "example_response_shape", + ), + ("response", _response_body(), _json_headers((b"Content-Type", b"text/plain")), 502, "example_response_shape"), + ( + "response", + _response_body(), + ((b"Content-Type", b"text/plain"), *_json_headers()), + 502, + "example_response_shape", + ), ( "response", b"private-payload", From b944c65a68e3a6e3981626dbdc33b3b5eedf4428 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 12 Oct 2026 11:48:42 +0200 Subject: [PATCH 09/11] fix(server): close upstream resources independently of stream iterators Closing the body iterator does not prove the upstream connection is closed. A client disconnect can cancel an async generator's own cleanup part-way, and later aclose() calls cannot resume it. StreamingHttpResponse now carries an optional close callback that releases the upstream resource directly. Iterators that own their resource may keep exposing aclose(); a bare async generator must supply the callback and is otherwise rejected. Dispatch owns resources until it returns the response and must clean up if construction fails. From then on the executor owns cleanup on every path, including rejection and disconnect before iteration. Closing must be idempotent; each attempt is bounded and failures are logged. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_streaming_http.py | 78 +++++++------ .../experimental/test_streaming_http.py | 110 ++++++++++++++++-- 2 files changed, 143 insertions(+), 45 deletions(-) diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index 3a463b30ef..7a1f8c5b0b 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -16,12 +16,10 @@ """Bind guarded provider streams to an injected HTTP exchange.""" import inspect -import logging import re from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass, replace -import anyio from starlette.responses import Response, StreamingResponse from starlette.types import Receive, Scope, Send @@ -46,29 +44,32 @@ _render_response, _request_identity_encoding, ) +from nemoguardrails.server.experimental._stream_cleanup import StreamClose, close_stream from nemoguardrails.server.experimental.provider.stream import ProviderStreamAdapter from nemoguardrails.server.experimental.provider.types import GuardedMessage -log = logging.getLogger(__name__) - -_STREAM_CLOSE_TIMEOUT_SECONDS = 5.0 - @dataclass(frozen=True, slots=True) class StreamingHttpResponse: """Carry one successful response without choosing an HTTP client. - The returned HTTP response closes ``body`` when it ends, possibly more than - once, so closing ``body`` must be idempotent. HTTP cleanup attempts are - shielded from enclosing cancellation scopes for up to five seconds each. - A timeout is logged without replacing the response or its original failure. - Bodies must cooperate with cancellation; cleanup cannot forcibly stop - synchronous blocking code or a closer that suppresses cancellation. + ``close`` releases the upstream response/connection independently of body + iteration. An ordinary self-owning iterator may instead expose ``aclose``; + a bare async generator must supply an independent resource closer. Never + rely solely on its finally block: cancellation can interrupt it permanently. + + Dispatch owns resources until it returns this value and must clean up if + acquisition or construction fails. The executor then owns cleanup through + rejection or response completion, including disconnect before iteration. + Closing must be idempotent: iterator and resource cleanup may be retried. + Each attempt is shielded and bounded; failures are logged, not substituted + for the primary outcome. Cancellation-suppressing closers cannot be forced. """ status_code: int headers: HttpHeaders body: AsyncIterator[bytes] + close: StreamClose | None = None def __post_init__(self) -> None: if type(self.status_code) is not int or not 200 <= self.status_code <= 299: @@ -77,6 +78,11 @@ def __post_init__(self) -> None: raise TypeError("A streaming HTTP response body must be an async iterator.") if any(not isinstance(name, bytes) or not isinstance(value, bytes) for name, value in self.headers): raise TypeError("Streaming HTTP response headers must contain byte pairs.") + if self.close is None: + if inspect.isasyncgen(self.body) or not callable(getattr(self.body, "aclose", None)): + raise TypeError("A streaming body needs an independent resource close callback") + elif not callable(self.close): + raise TypeError("The streaming resource close callback must be callable") StreamingHttpDispatch = Callable[[BufferedHttpRequest], Awaitable[StreamingHttpResponse | BufferedHttpResponse]] @@ -85,19 +91,27 @@ def __post_init__(self) -> None: class _ClosingStreamingResponse(StreamingResponse): """Close the provider body however the response ends, even before iteration.""" - def __init__(self, body: AsyncIterator[bytes], source: AsyncIterator[bytes], *, status_code: int): + def __init__( + self, + body: AsyncIterator[bytes], + source: AsyncIterator[bytes], + *, + status_code: int, + close: StreamClose | None = None, + ): super().__init__(body, status_code=status_code) self.owned_body = body self.source = source + self.close = close async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: try: await super().__call__(scope, receive, send) finally: try: - await _close_source(self.owned_body) + await close_stream(self.owned_body) finally: - await _close_source(self.source) + await close_stream(self.source, close=self.close) _BODY_DEPENDENT_HEADERS = frozenset( @@ -111,7 +125,7 @@ def _streaming_response( *, may_modify: bool, ) -> StreamingResponse: - response = _ClosingStreamingResponse(body, value.body, status_code=value.status_code) + response = _ClosingStreamingResponse(body, value.body, status_code=value.status_code, close=value.close) response.raw_headers = [ (name, header_value) for name, header_value in _end_to_end_headers(value.headers) @@ -124,24 +138,9 @@ def _header_values(headers: HttpHeaders, name: bytes) -> tuple[bytes, ...]: return tuple(value for header_name, value in headers if header_name.lower() == name) -async def _close_source(source: AsyncIterator[bytes]) -> None: - """Bound shielded cleanup, logging its timeout without masking the outcome. - - Only this scope's timeout is suppressed. Other close failures propagate as - before. Each call gets its own budget, including a final provider-body close - after the response iterator's cleanup, so the limit is not per request. - """ - with anyio.move_on_after(_STREAM_CLOSE_TIMEOUT_SECONDS, shield=True) as scope: - close = getattr(source, "aclose", None) - if callable(close): - result = close() - if inspect.isawaitable(result): - await result - if scope.cancelled_caught: - log.warning("Stream cleanup timed out; resources may remain open.") - - -async def _relay(source: AsyncIterator[bytes], expected_length: bytes | None) -> AsyncIterator[bytes]: +async def _relay( + source: AsyncIterator[bytes], expected_length: bytes | None, *, close: StreamClose | None +) -> AsyncIterator[bytes]: observed = 0 try: async for chunk in source: @@ -158,7 +157,7 @@ async def _relay(source: AsyncIterator[bytes], expected_length: bytes | None) -> if expected_length is not None and str(observed).encode() != expected_length: raise ValueError("The provider stream does not match its declared Content-Length.") finally: - await _close_source(source) + await close_stream(source, close=close) async def execute_streaming_http( @@ -216,7 +215,7 @@ async def execute_streaming_http( or len(content_lengths) > 1 or any(re.fullmatch(rb"[0-9]+", length) is None for length in content_lengths) ): - await _close_source(upstream.body) + await close_stream(upstream.body, close=upstream.close) return _render_response( render_outcome( OperationProjectionFailed( @@ -227,7 +226,7 @@ async def execute_streaming_http( request.method, ) if streaming_policy is not None and not _has_identity_encoding(upstream.headers): - await _close_source(upstream.body) + await close_stream(upstream.body, close=upstream.close) return _render_response( render_outcome( OperationProjectionFailed( @@ -239,7 +238,9 @@ async def execute_streaming_http( ) if streaming_policy is None: expected_length = (content_lengths[0].lstrip(b"0") or b"0") if content_lengths else None - return _streaming_response(upstream, _relay(upstream.body, expected_length), may_modify=False) + return _streaming_response( + upstream, _relay(upstream.body, expected_length, close=upstream.close), may_modify=False + ) body = guard_provider_stream( upstream.body, @@ -250,5 +251,6 @@ async def execute_streaming_http( render_outcome=render_outcome, max_event_bytes=max_event_bytes, max_pending_bytes=max_pending_bytes, + close=upstream.close, ) return _streaming_response(upstream, body, may_modify=True) diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index 5d961d0630..53fba456f7 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -21,7 +21,7 @@ import pytest from starlette.requests import ClientDisconnect -from nemoguardrails.server.experimental import _streaming_http +from nemoguardrails.server.experimental import _stream_cleanup from nemoguardrails.server.experimental._buffered_kernel import OperationProjectionFailed from nemoguardrails.server.experimental._content_checker import ContentAllowed, StreamBufferingPolicy from nemoguardrails.server.experimental._guarded_stream import StreamProcessingFailed, StreamUpstreamFailed @@ -327,7 +327,7 @@ async def request(): @pytest.mark.parametrize("encoded", [False, True]) @pytest.mark.parametrize("cancel_request", [False, True]) async def test_stalled_rejection_cleanup_preserves_rejection(monkeypatch, caplog, encoded, cancel_request): - monkeypatch.setattr(_streaming_http, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01, raising=False) + monkeypatch.setattr(_stream_cleanup, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01) source = StalledClose() headers = ( ((b"content-type", b"text/event-stream"), (b"content-encoding", b"gzip")) @@ -361,7 +361,7 @@ async def request(): @pytest.mark.asyncio @pytest.mark.parametrize("stalled_body", [False, True]) async def test_stalled_response_cleanup_preserves_send_failure(monkeypatch, caplog, stalled_body): - monkeypatch.setattr(_streaming_http, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01, raising=False) + monkeypatch.setattr(_stream_cleanup, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01) body = StalledClose() if stalled_body else Source([]) source = Source([]) if stalled_body else StalledClose() response = _ClosingStreamingResponse(body, source, status_code=200) @@ -386,7 +386,7 @@ async def send(_message): @pytest.mark.asyncio @pytest.mark.parametrize("stalled_body", [False, True]) async def test_stalled_response_cleanup_preserves_task_cancellation(monkeypatch, caplog, stalled_body): - monkeypatch.setattr(_streaming_http, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01) + monkeypatch.setattr(_stream_cleanup, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01) body = StalledClose() if stalled_body else Source([]) source = Source([]) if stalled_body else StalledClose() response = _ClosingStreamingResponse(body, source, status_code=200) @@ -718,7 +718,7 @@ async def source(): @pytest.mark.asyncio @pytest.mark.parametrize("failure_type", [RuntimeError, TimeoutError]) -async def test_source_is_closed_when_closing_the_body_fails(failure_type): +async def test_source_is_closed_when_closing_the_body_fails(failure_type, caplog): class FailingBody: def __aiter__(self): return self @@ -738,10 +738,10 @@ async def receive(): async def send(_message): pass - with pytest.raises(failure_type, match="body cleanup failed"): - await response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send) + await response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send) assert source.closed is True + assert "Stream cleanup failed" in caplog.text @pytest.mark.asyncio @@ -796,3 +796,99 @@ async def dispatch(request): await response_body(response) assert [value for name, value in dispatched[0].headers if name.lower() == b"accept-encoding"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", [None, StreamBufferingPolicy(1, 0)]) +async def test_normal_eof_with_stalled_cleanup_sends_final_asgi_body(monkeypatch, policy): + monkeypatch.setattr(_stream_cleanup, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01) + source = StalledClose() + source.chunks = iter([b'data: {"text":"answer"}\n\n']) + messages = [] + + async def dispatch(_request): + return StreamingHttpResponse(200, ((b"content-type", b"text/event-stream"),), source) + + async def receive(): + await anyio.sleep_forever() + + async def send(message): + messages.append(message) + + response = await execute(dispatch, policy=policy) + with anyio.fail_after(1): + await response({"type": "http", "asgi": {"spec_version": "2.4"}}, receive, send) + assert messages[-1] == {"type": "http.response.body", "body": b"", "more_body": False} + assert source.close_cancelled + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", [None, StreamBufferingPolicy(1, 0)]) +@pytest.mark.parametrize("cancel_at", ["source_read", "checker", "startup"]) +async def test_disconnect_closes_resource_independently_of_generator_finally(policy, cancel_at): + if policy is None and cancel_at == "checker": + return + started = anyio.Event() + finalizer_started = False + finalizer_finished = False + resource_closed = False + + async def source(): + nonlocal finalizer_started, finalizer_finished + try: + yield b'data: {"text":"answer"}\n\n' + started.set() + await anyio.sleep_forever() + finally: + finalizer_started = True + await anyio.sleep(0) + finalizer_finished = True + + async def close_resource(): + nonlocal resource_closed + await anyio.sleep(0) + resource_closed = True + + class WaitingChecker(Checker): + async def check_output(self, check): + if cancel_at == "checker": + started.set() + await anyio.sleep_forever() + return await super().check_output(check) + + body = source() + + async def dispatch(_request): + return StreamingHttpResponse(200, ((b"content-type", b"text/event-stream"),), body, close=close_resource) + + async def receive(): + await started.wait() + return {"type": "http.disconnect"} + + async def send(message): + if cancel_at == "startup" and message["type"] == "http.response.start": + started.set() + await anyio.sleep_forever() + + response = await execute(dispatch, policy=policy, checker=WaitingChecker()) + with anyio.fail_after(1): + await response({"type": "http", "asgi": {"spec_version": "2.3"}}, receive, send) + + assert resource_closed + if cancel_at == "source_read": + assert finalizer_started and not finalizer_finished + await body.aclose() + assert resource_closed + + +@pytest.mark.asyncio +async def test_bare_generator_requires_independent_resource_ownership(): + async def source(): + yield b"data: [END]\n\n" + + body = source() + try: + with pytest.raises(TypeError, match="independent resource close"): + StreamingHttpResponse(200, (), body) + finally: + await body.aclose() From 0082da888de76f7b99bd9bd39047c98b0d1f2a04 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:44:42 +0200 Subject: [PATCH 10/11] fix(server): report only typed dispatch failures as upstream outages The streaming executor treated any exception from dispatch as an upstream failure, so a bug in a dispatcher or in our own code was reported to the client and operators as a provider outage. It now catches only HttpDispatchFailed, like the buffered path; other exceptions propagate. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_streaming_http.py | 5 ++++- .../experimental/test_streaming_http.py | 19 +++++++++++++++++-- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/nemoguardrails/server/experimental/_streaming_http.py b/nemoguardrails/server/experimental/_streaming_http.py index 7a1f8c5b0b..8ce765299f 100644 --- a/nemoguardrails/server/experimental/_streaming_http.py +++ b/nemoguardrails/server/experimental/_streaming_http.py @@ -36,6 +36,7 @@ from nemoguardrails.server.experimental._http_kernel import ( BufferedHttpRequest, BufferedHttpResponse, + HttpDispatchFailed, HttpHeaders, ResponseBodyTooLarge, _end_to_end_headers, @@ -185,7 +186,9 @@ async def execute_streaming_http( try: upstream = await dispatch(request) - except Exception as failure: + except HttpDispatchFailed as failure: + # Match the buffered path: only a typed dispatch failure is an upstream + # outage. Other exceptions are bugs and propagate. return _render_response(render_outcome(StreamUpstreamFailed(failure)), request.method) if isinstance(upstream, BufferedHttpResponse): diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index 53fba456f7..044bd000a1 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -25,7 +25,11 @@ from nemoguardrails.server.experimental._buffered_kernel import OperationProjectionFailed from nemoguardrails.server.experimental._content_checker import ContentAllowed, StreamBufferingPolicy from nemoguardrails.server.experimental._guarded_stream import StreamProcessingFailed, StreamUpstreamFailed -from nemoguardrails.server.experimental._http_kernel import BufferedHttpRequest, BufferedHttpResponse +from nemoguardrails.server.experimental._http_kernel import ( + BufferedHttpRequest, + BufferedHttpResponse, + HttpDispatchFailed, +) from nemoguardrails.server.experimental._streaming_http import ( StreamingHttpResponse, _ClosingStreamingResponse, @@ -251,7 +255,7 @@ async def dispatch(_request): @pytest.mark.asyncio async def test_dispatch_failure_is_rendered_before_streaming_starts(): async def dispatch(_request): - raise RuntimeError("private upstream detail") + raise HttpDispatchFailed("private upstream detail") response = await execute(dispatch) @@ -259,6 +263,17 @@ async def dispatch(_request): assert await response_body(response) == b"upstream_failed" +@pytest.mark.asyncio +async def test_unexpected_dispatch_errors_are_not_reported_as_upstream_failures(): + """As in the buffered path, a dispatcher bug propagates instead of looking like an outage.""" + + async def dispatch(_request): + raise TypeError("dispatcher bug") + + with pytest.raises(TypeError, match="dispatcher bug"): + await execute(dispatch) + + @pytest.mark.asyncio @pytest.mark.parametrize( "headers", From ef094d6ba2a1e2fab2a9857c33ea7b30f98d8f03 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:45:51 +0200 Subject: [PATCH 11/11] test(server): cover streaming HTTP validation paths Adds cases for invalid StreamingHttpResponse values, non-byte chunks in an uninspected relay, non-positive byte limits, and an unsupported dispatch result. No runtime behavior changes. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../experimental/test_streaming_http.py | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/tests/server/experimental/test_streaming_http.py b/tests/server/experimental/test_streaming_http.py index 044bd000a1..4b1303e671 100644 --- a/tests/server/experimental/test_streaming_http.py +++ b/tests/server/experimental/test_streaming_http.py @@ -907,3 +907,53 @@ async def source(): StreamingHttpResponse(200, (), body) finally: await body.aclose() + + +@pytest.mark.parametrize( + ("kwargs", "error", "message"), + [ + ({"status_code": 404}, ValueError, "successful status code"), + ({"body": [b"data: x\n\n"]}, TypeError, "must be an async iterator"), + ({"headers": (("content-type", "text/event-stream"),)}, TypeError, "byte pairs"), + ({"close": "not callable"}, TypeError, "close callback must be callable"), + ], +) +def test_streaming_response_rejects_invalid_values(kwargs, error, message): + values = {"status_code": 200, "headers": (), "body": Source([]), **kwargs} + with pytest.raises(error, match=message): + StreamingHttpResponse(**values) + + +@pytest.mark.asyncio +async def test_uninspected_relay_rejects_non_byte_chunks_and_closes_source(): + source = Source(["data: text\n\n"]) + + async def dispatch(_request): + return StreamingHttpResponse(200, ((b"content-type", b"text/event-stream"),), source) + + response = await execute(dispatch, policy=None) + + with pytest.raises(TypeError, match="must yield bytes"): + await response_body(response) + assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limits", [{"max_event_bytes": 0}, {"max_pending_bytes": -1}]) +async def test_byte_limits_must_be_positive(limits): + async def dispatch(_request): + raise AssertionError("dispatch must not run") + + with pytest.raises(ValueError, match="byte limits must be positive"): + await execute(dispatch, **limits) + + +@pytest.mark.asyncio +async def test_unsupported_dispatch_result_is_a_processing_failure(): + async def dispatch(_request): + return "not a response" + + response = await execute(dispatch) + + assert response.status_code == 502 + assert await response_body(response) == b"processing_failed"