From a2ba0e9c6b018eb1a0248795aac9de1116b8d072 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 5 Oct 2026 16:16:48 +0200 Subject: [PATCH 01/16] feat(server): execute guarded provider streams Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_content_checker.py | 22 ++ .../server/experimental/_guarded_stream.py | 306 +++++++++++++++ .../server/experimental/provider/sse.py | 131 +++++++ .../server/experimental/provider/stream.py | 219 +++++++++++ .../provider/stream_classifier.py | 258 +++++++++++++ .../experimental/test_content_checker.py | 25 ++ .../experimental/test_guarded_stream.py | 353 ++++++++++++++++++ tests/server/experimental/test_sse.py | 49 +++ .../experimental/test_stream_classifier.py | 210 +++++++++++ 9 files changed, 1573 insertions(+) create mode 100644 nemoguardrails/server/experimental/_guarded_stream.py create mode 100644 nemoguardrails/server/experimental/provider/sse.py create mode 100644 nemoguardrails/server/experimental/provider/stream.py create mode 100644 nemoguardrails/server/experimental/provider/stream_classifier.py create mode 100644 tests/server/experimental/test_guarded_stream.py create mode 100644 tests/server/experimental/test_sse.py create mode 100644 tests/server/experimental/test_stream_classifier.py diff --git a/nemoguardrails/server/experimental/_content_checker.py b/nemoguardrails/server/experimental/_content_checker.py index 101cc20ac6..e7c535fbad 100644 --- a/nemoguardrails/server/experimental/_content_checker.py +++ b/nemoguardrails/server/experimental/_content_checker.py @@ -33,18 +33,40 @@ class UnsupportedContentModification(RuntimeError): """Report a checker modification that the proxy cannot safely apply.""" +@dataclass(frozen=True, slots=True) +class StreamBufferingPolicy: + """Declare how guarded stream deltas are checked before release.""" + + chunk_size: int + context_size: int + release_before_check: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.chunk_size, int) or self.chunk_size <= 0: + raise ValueError("A stream chunk size must be a positive integer.") + if not isinstance(self.context_size, int) or self.context_size < 0: + raise ValueError("A stream context size must be a non-negative integer.") + if not isinstance(self.release_before_check, bool): + raise TypeError("The stream release policy must be a boolean.") + + @dataclass(frozen=True, slots=True) class ContentInspectionPolicy: """Record whether a checker examines input and output content.""" inspect_input: bool inspect_output: bool + stream_buffering: StreamBufferingPolicy | None = None def __post_init__(self) -> None: """Reject non-boolean inspection flags.""" if not isinstance(self.inspect_input, bool) or not isinstance(self.inspect_output, bool): raise TypeError("Content inspection flags must be booleans.") + if self.stream_buffering is not None and not isinstance(self.stream_buffering, StreamBufferingPolicy): + raise TypeError("Stream buffering must be a StreamBufferingPolicy.") + if not self.inspect_output and self.stream_buffering is not None: + raise ValueError("Stream buffering requires output inspection.") @dataclass(frozen=True, slots=True) diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py new file mode 100644 index 0000000000..c3b9988b0d --- /dev/null +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -0,0 +1,306 @@ +# 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. + +"""Apply buffered output checks to provider-native event streams.""" + +import inspect +from collections.abc import AsyncIterator, Callable +from dataclasses import dataclass + +from nemoguardrails.server.experimental._buffered_kernel import ( + InspectionStage, + OperationBlocked, + OperationCheckFailed, + OperationModificationUnsupported, + OperationProjectionFailed, +) +from nemoguardrails.server.experimental._content_checker import ( + ContentAllowed, + ContentBlocked, + ContentChecker, + ContentCheckFailed, + OutputContentCheck, + StreamBufferingPolicy, + UnsupportedContentModification, + validate_content_check_decision, +) +from nemoguardrails.server.experimental._guarded_operation import UnsupportedGuardedPayload +from nemoguardrails.server.experimental._http_kernel import BufferedHttpResponse +from nemoguardrails.server.experimental.provider.sse import ( + ServerSentEvent, + ServerSentEventTooLarge, + iter_sse_events, +) +from nemoguardrails.server.experimental.provider.stream import ( + ProviderStreamAdapter, + StreamEventRole, + UnsupportedProviderStream, + validate_stream_event, +) +from nemoguardrails.server.experimental.provider.types import GuardedMessage + + +class UnsupportedStreamInspection(ValueError): + """Report an output policy that cannot safely release a guarded stream.""" + + +class PendingStreamBufferTooLarge(ValueError): + """Report unchecked provider events that exceed the buffering limit.""" + + +class StreamSnapshotHistoryTooLarge(ValueError): + """Report snapshot verification state that exceeds the stream limit.""" + + +@dataclass(frozen=True, slots=True) +class StreamUpstreamFailed: + """Stop a stream when its upstream byte iterator fails.""" + + failure: BaseException + + +@dataclass(frozen=True, slots=True) +class StreamProcessingFailed: + """Stop a stream when guarded stream processing fails unexpectedly.""" + + failure: BaseException + + +StreamOutcome = ( + OperationBlocked + | OperationCheckFailed + | OperationModificationUnsupported + | OperationProjectionFailed + | StreamUpstreamFailed + | StreamProcessingFailed +) +StreamOutcomeRenderer = Callable[[StreamOutcome], BufferedHttpResponse] + + +@dataclass(frozen=True, slots=True) +class _PendingStreamEvent: + event: ServerSentEvent + has_guarded_text: bool + + +@dataclass(frozen=True, slots=True) +class _TextWindow: + processing: tuple[str, ...] + release_count: int + + +class _UpstreamStreamFailure(Exception): + def __init__(self, failure: BaseException): + super().__init__(str(failure)) + self.failure = failure + + +def validate_streaming_policy(streaming_policy: StreamBufferingPolicy | None) -> None: + """Require output streams to buffer every guarded delta before release.""" + + if streaming_policy is None: + raise UnsupportedStreamInspection("Guarded output streaming requires a buffering policy.") + if streaming_policy.release_before_check: + raise UnsupportedStreamInspection("Guarded streams cannot release output before checking it.") + + +def _take_checked_events( + pending: list[_PendingStreamEvent], + text_count: int, +) -> tuple[ServerSentEvent, ...]: + checked_count = 0 + end = 0 + for end, event in enumerate(pending, start=1): + if event.has_guarded_text: + checked_count += 1 + if checked_count == text_count: + break + checked = pending[:end] + del pending[:end] + return tuple(event.event for event in checked) + + +async def _text_windows( + source: AsyncIterator[str], + *, + chunk_size: int, + context_size: int, +) -> AsyncIterator[_TextWindow]: + context: list[str] = [] + pending: list[str] = [] + async for text in source: + pending.append(text) + if len(pending) < chunk_size: + continue + processing = context + pending + yield _TextWindow(tuple(processing), len(pending)) + context = processing[-context_size:] if context_size else [] + pending = [] + if pending: + yield _TextWindow(tuple(context + pending), len(pending)) + + +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 _guard_upstream(source: AsyncIterator[bytes]) -> AsyncIterator[bytes]: + try: + async for chunk in source: + if not isinstance(chunk, bytes): + raise TypeError("A provider stream must yield bytes.") + yield chunk + except Exception as failure: + raise _UpstreamStreamFailure(failure) from failure + + +def _encoded_outcome( + outcome: StreamOutcome, + *, + adapter: ProviderStreamAdapter, + render_outcome: StreamOutcomeRenderer, +) -> tuple[bytes, ...]: + response = render_outcome(outcome) + return tuple(adapter.encode_error(response.body)) + + +async def guard_provider_stream( + source: AsyncIterator[bytes], + *, + checker: ContentChecker, + streaming_policy: StreamBufferingPolicy | None, + input_message: GuardedMessage, + adapter: ProviderStreamAdapter, + render_outcome: StreamOutcomeRenderer, + max_event_bytes: int = 1024 * 1024, + max_pending_bytes: int = 10 * 1024 * 1024, +) -> AsyncIterator[bytes]: + """Release provider SSE events only after required output checks pass.""" + + if input_message.role != "user": + raise ValueError("A guarded stream requires a user input message.") + if any(type(limit) is not int or limit <= 0 for limit in (max_event_bytes, max_pending_bytes)): + raise ValueError("Stream byte limits must be positive.") + + try: + if streaming_policy is None: + async for chunk in source: + yield chunk + return + validate_streaming_policy(streaming_policy) + + pending: list[_PendingStreamEvent] = [] + pending_bytes = 0 + observed_text: list[str] = [] + observed_text_bytes = 0 + snapshot_history_complete = True + + async def text_deltas() -> AsyncIterator[str]: + nonlocal observed_text_bytes, pending_bytes, snapshot_history_complete + async for event in iter_sse_events(_guard_upstream(source), max_event_bytes=max_event_bytes): + classified = adapter.classify_event(event) + validate_stream_event(adapter.contract, classified) + if classified.role is StreamEventRole.TEXT_SNAPSHOT: + if not snapshot_history_complete: + raise StreamSnapshotHistoryTooLarge( + "Provider stream text exceeds the configured snapshot verification limit." + ) + if classified.text != "".join(observed_text): + raise UnsupportedProviderStream( + "A provider stream text snapshot does not match the guarded text deltas." + ) + text = classified.text if classified.role is StreamEventRole.GUARDED_TEXT else None + pending.append(_PendingStreamEvent(event, has_guarded_text=bool(text))) + pending_bytes += len(event.raw) + if pending_bytes > max_pending_bytes: + raise PendingStreamBufferTooLarge( + "Unchecked provider stream events exceed the configured buffering limit." + ) + if text: + if snapshot_history_complete: + observed_text_bytes += len(text.encode()) + if observed_text_bytes > max_event_bytes: + observed_text.clear() + snapshot_history_complete = False + else: + observed_text.append(text) + yield text + adapter.validate_end_of_stream() + + async for window in _text_windows( + text_deltas(), + chunk_size=streaming_policy.chunk_size, + context_size=streaming_policy.context_size, + ): + checked_events = _take_checked_events(pending, window.release_count) + pending_bytes -= sum(len(event.raw) for event in checked_events) + try: + decision = validate_content_check_decision( + await checker.check_output(OutputContentCheck(input_message, "".join(window.processing))) + ) + except UnsupportedContentModification as failure: + outcome: StreamOutcome = OperationModificationUnsupported(InspectionStage.OUTPUT, failure) + except Exception as failure: + outcome = OperationCheckFailed( + InspectionStage.OUTPUT, + ContentCheckFailed("The output content check failed.", failure), + ) + else: + if isinstance(decision, ContentAllowed): + for event in checked_events: + yield event.raw + continue + if isinstance(decision, ContentBlocked): + outcome = OperationBlocked(InspectionStage.OUTPUT, decision) + else: + outcome = OperationCheckFailed(InspectionStage.OUTPUT, decision) + for chunk in _encoded_outcome(outcome, adapter=adapter, render_outcome=render_outcome): + yield chunk + return + + for event in pending: + yield event.event.raw + except ( + PendingStreamBufferTooLarge, + ServerSentEventTooLarge, + StreamSnapshotHistoryTooLarge, + UnsupportedProviderStream, + ) as failure: + outcome = OperationProjectionFailed( + InspectionStage.OUTPUT, + UnsupportedGuardedPayload(str(failure) or "The provider stream does not match the guarded content schema."), + ) + for chunk in _encoded_outcome(outcome, adapter=adapter, render_outcome=render_outcome): + yield chunk + except _UpstreamStreamFailure as failure: + for chunk in _encoded_outcome( + StreamUpstreamFailed(failure.failure), + adapter=adapter, + render_outcome=render_outcome, + ): + yield chunk + except Exception as failure: + for chunk in _encoded_outcome( + StreamProcessingFailed(failure), + adapter=adapter, + render_outcome=render_outcome, + ): + yield chunk + finally: + await _close_source(source) diff --git a/nemoguardrails/server/experimental/provider/sse.py b/nemoguardrails/server/experimental/provider/sse.py new file mode 100644 index 0000000000..2bc78268b4 --- /dev/null +++ b/nemoguardrails/server/experimental/provider/sse.py @@ -0,0 +1,131 @@ +# 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. + +"""Parse Server-Sent Events without provider-specific knowledge. + +The parser preserves complete upstream events byte-for-byte while exposing the +event type and joined data fields needed by provider adapters. +""" + +from collections.abc import AsyncIterator +from dataclasses import dataclass + + +class ServerSentEventTooLarge(ValueError): + """Report an SSE event that exceeds the configured buffering limit.""" + + +@dataclass(frozen=True, slots=True) +class ServerSentEvent: + """Represent one complete SSE event. + + Attributes: + raw: Original event bytes, including its terminating blank line when + supplied by the upstream server. + lines: Event field and comment lines without line endings. + """ + + raw: bytes + lines: tuple[bytes, ...] + + @classmethod + def from_bytes(cls, raw: bytes) -> "ServerSentEvent": + """Create an event while retaining its original representation. + + Args: + raw: Complete SSE event bytes. + + Returns: + The parsed event representation. + """ + + body = raw + for separator in (b"\r\n\r\n", b"\n\n", b"\r\r"): + if body.endswith(separator): + body = body[: -len(separator)] + break + return cls(raw=raw, lines=tuple(body.replace(b"\r\n", b"\n").replace(b"\r", b"\n").split(b"\n"))) + + @property + def data(self) -> bytes | None: + """Return joined SSE data fields, or ``None`` when none are present.""" + + values = [] + for line in self.lines: + name, separator, value = line.partition(b":") + if separator and name == b"data": + values.append(value[1:] if value.startswith(b" ") else value) + elif not separator and name == b"data": + values.append(b"") + return b"\n".join(values) if values else None + + @property + def event_type(self) -> bytes | None: + """Return the last SSE event field, or ``None`` when it is absent.""" + + value = None + for line in self.lines: + name, separator, field_value = line.partition(b":") + if name == b"event": + value = field_value[1:] if separator and field_value.startswith(b" ") else field_value + return value + + +def _event_boundary(buffer: bytes) -> tuple[int, int] | None: + boundaries = [] + for separator in (b"\r\n\r\n", b"\n\n", b"\r\r"): + index = buffer.find(separator) + if index >= 0: + boundaries.append((index, len(separator))) + return min(boundaries) if boundaries else None + + +async def iter_sse_events( + source: AsyncIterator[bytes], + *, + max_event_bytes: int = 1024 * 1024, +) -> AsyncIterator[ServerSentEvent]: + """Yield complete SSE events from arbitrary HTTP byte chunks. + + Args: + source: Raw upstream response byte stream. + max_event_bytes: Maximum bytes buffered for one event. + + Yields: + Complete SSE events in upstream order. + + Raises: + ValueError: If ``max_event_bytes`` is not positive. + ServerSentEventTooLarge: If one event exceeds the limit. + """ + + if max_event_bytes <= 0: + raise ValueError("max_event_bytes must be positive.") + + buffer = b"" + async for chunk in source: + buffer += chunk + while boundary := _event_boundary(buffer): + index, separator_length = boundary + event_length = index + separator_length + if event_length > max_event_bytes: + raise ServerSentEventTooLarge + raw, buffer = buffer[:event_length], buffer[event_length:] + yield ServerSentEvent.from_bytes(raw) + if len(buffer) > max_event_bytes: + raise ServerSentEventTooLarge + + if buffer: + yield ServerSentEvent.from_bytes(buffer) diff --git a/nemoguardrails/server/experimental/provider/stream.py b/nemoguardrails/server/experimental/provider/stream.py new file mode 100644 index 0000000000..3bbf071c06 --- /dev/null +++ b/nemoguardrails/server/experimental/provider/stream.py @@ -0,0 +1,219 @@ +# 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. + +"""Define the stream boundary implemented by provider adapters.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Callable, Protocol + +from nemoguardrails.server.experimental.provider.payload import ProjectionFieldCoverage +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent + + +class UnsupportedProviderStream(ValueError): + """Report a provider event that the guarded stream cannot interpret.""" + + +class StreamEventRole(str, Enum): + """Describe how one provider event relates to guarded content.""" + + GUARDED_TEXT = "guarded_text" + TEXT_SNAPSHOT = "text_snapshot" + OPAQUE_METADATA = "opaque_metadata" + PROVIDER_ERROR = "provider_error" + + +class StreamCapabilityProfile(str, Enum): + """Identify a closed streaming capability understood by the framework.""" + + SINGLE_TEXT_DELTA_V1 = "single_text_delta.v1" + + +@dataclass(frozen=True, slots=True) +class StreamShapeCoverage: + """Classify every reviewed native event shape in a provider stream.""" + + guarded_shapes: frozenset[str] = frozenset() + snapshot_shapes: frozenset[str] = frozenset() + opaque_shapes: frozenset[str] = frozenset() + provider_error_shapes: frozenset[str] = frozenset() + + def __post_init__(self) -> None: + classifications = ( + self.guarded_shapes, + self.snapshot_shapes, + self.opaque_shapes, + self.provider_error_shapes, + ) + repeated = set() + for index, shapes in enumerate(classifications): + for other_shapes in classifications[index + 1 :]: + repeated.update(shapes & other_shapes) + if repeated: + raise ValueError(f"Stream shapes have multiple classifications: {sorted(repeated)}") + + @property + def reviewed_shapes(self) -> frozenset[str]: + """Return every shape deliberately covered by the stream contract.""" + + return self.guarded_shapes | self.snapshot_shapes | self.opaque_shapes | self.provider_error_shapes + + +@dataclass(frozen=True, slots=True) +class StreamProjectionContract: + """Describe the runtime semantics of one guarded stream projection.""" + + projection_id: str + profile: StreamCapabilityProfile + shapes: StreamShapeCoverage + fields: ProjectionFieldCoverage | None = None + + def __post_init__(self) -> None: + if not self.projection_id.strip(): + raise ValueError("A stream projection contract must have an identifier.") + if not self.shapes.reviewed_shapes: + raise ValueError("A stream projection contract must classify at least one shape.") + + +@dataclass(frozen=True, slots=True) +class GuardedStreamEvent: + """Describe one provider event at the Guardrails boundary.""" + + shape: str + role: StreamEventRole + text: str | None = None + + def __post_init__(self) -> None: + """Require text exactly for roles that carry generated content.""" + + if not self.shape: + raise ValueError("A guarded stream event must identify its provider shape.") + carries_text = self.role in {StreamEventRole.GUARDED_TEXT, StreamEventRole.TEXT_SNAPSHOT} + if carries_text != (self.text is not None): + raise ValueError("Guarded text and snapshot events must carry text; other event roles must not.") + + +def validate_stream_event(contract: StreamProjectionContract, event: GuardedStreamEvent) -> None: + """Require a classified event to match its declared shape semantics.""" + + roles = { + StreamEventRole.GUARDED_TEXT: contract.shapes.guarded_shapes, + StreamEventRole.TEXT_SNAPSHOT: contract.shapes.snapshot_shapes, + StreamEventRole.OPAQUE_METADATA: contract.shapes.opaque_shapes, + StreamEventRole.PROVIDER_ERROR: contract.shapes.provider_error_shapes, + } + declared_role = next((role for role, shapes in roles.items() if event.shape in shapes), None) + if declared_role is None: + raise UnsupportedProviderStream( + f"Provider stream shape {event.shape!r} is not declared by {contract.projection_id!r}." + ) + if event.role is not declared_role: + raise UnsupportedProviderStream( + f"Provider stream shape {event.shape!r} is classified as {event.role.value!r}; " + f"the contract declares {declared_role.value!r}." + ) + + +class ProviderStreamClassifier(Protocol): + """Classify native events using declarative projection semantics.""" + + contract: StreamProjectionContract + + def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: + """Validate and classify one provider event.""" + + ... + + +class ProviderStreamHooks(Protocol): + """Implement stateful protocol checks and provider-native error framing.""" + + def observe_event(self, event: ServerSentEvent, classified: GuardedStreamEvent) -> None: + """Observe one classified event and update protocol state.""" + + ... + + def validate_end_of_stream(self) -> None: + """Validate provider protocol state when the byte stream ends.""" + + ... + + def encode_error(self, rendered_body: bytes) -> Sequence[bytes]: + """Frame a proxy error and its provider-native body for the stream.""" + + ... + + +class ProviderStreamAdapter(Protocol): + """Expose the complete stream behavior consumed by the shared runtime.""" + + contract: StreamProjectionContract + + def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: + """Validate, classify, and observe one provider event.""" + + ... + + def validate_end_of_stream(self) -> None: + """Validate provider protocol state when the byte stream ends.""" + + ... + + def encode_error(self, rendered_body: bytes) -> Sequence[bytes]: + """Frame a proxy error and its provider-native body for the stream.""" + + ... + + +@dataclass(slots=True) +class ClassifiedStreamAdapter: + """Compose a generated classifier with handwritten protocol hooks.""" + + classifier: ProviderStreamClassifier + hooks: ProviderStreamHooks + + @property + def contract(self) -> StreamProjectionContract: + """Return the classifier's stream projection contract.""" + + return self.classifier.contract + + def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: + """Classify an event and expose it to stateful protocol hooks.""" + + classified = self.classifier.classify_event(event) + self.hooks.observe_event(event, classified) + return classified + + def validate_end_of_stream(self) -> None: + """Delegate end-of-stream validation to provider hooks.""" + + self.hooks.validate_end_of_stream() + + def encode_error(self, rendered_body: bytes) -> Sequence[bytes]: + """Delegate provider-native error framing to provider hooks.""" + + return self.hooks.encode_error(rendered_body) + + +def create_classified_stream_adapter_factory( + classifier: ProviderStreamClassifier, + hooks_factory: Callable[[], ProviderStreamHooks], +) -> Callable[[], ProviderStreamAdapter]: + """Create a zero-argument adapter factory for an endpoint declaration.""" + + return lambda: ClassifiedStreamAdapter(classifier, hooks_factory()) diff --git a/nemoguardrails/server/experimental/provider/stream_classifier.py b/nemoguardrails/server/experimental/provider/stream_classifier.py new file mode 100644 index 0000000000..09939dfc9c --- /dev/null +++ b/nemoguardrails/server/experimental/provider/stream_classifier.py @@ -0,0 +1,258 @@ +# 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. + +"""Classify provider stream events from generated semantic rules.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Literal + +from pydantic import BaseModel, ValidationError + +from nemoguardrails.server.experimental._json_payload import parse_json_object +from nemoguardrails.server.experimental.provider.payload import guarded_schema_error +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent +from nemoguardrails.server.experimental.provider.stream import ( + GuardedStreamEvent, + StreamEventRole, + StreamProjectionContract, + UnsupportedProviderStream, + validate_stream_event, +) + +StreamMatchValue = str | int | float | bool | None + + +@dataclass(frozen=True, slots=True) +class StreamArraySelection: + """Select stream content by provider discriminator fields.""" + + match: tuple[tuple[str, StreamMatchValue], ...] + cardinality: Literal["at_most_one", "exactly_one"] + + def select(self, items: Sequence[object]) -> object | None: + matches = [ + item + for item in items + if all( + type(value := (item.get(name) if isinstance(item, dict) else getattr(item, name, object()))) + is type(expected) + and value == expected + for name, expected in self.match + ) + ] + if len(matches) > 1 or (self.cardinality == "exactly_one" and not matches): + raise ValueError(f"Stream text selection requires {self.cardinality.replace('_', ' ')} matching item.") + return matches[0] if matches else None + + +StreamTextPath = tuple[str | int | StreamArraySelection, ...] + + +@dataclass(frozen=True, slots=True) +class StreamEventRule: + """Declare validation and classification for one provider event shape.""" + + shape: str + model: type[BaseModel] + role: StreamEventRole + match: tuple[tuple[str, StreamMatchValue], ...] = () + required_fields: frozenset[str] = frozenset() + text_path: StreamTextPath | None = None + missing_text_role: StreamEventRole | None = None + missing_text_shape: str | None = None + missing_text_value: str | None = None + + def matches(self, payload: dict[str, Any]) -> bool: + """Return whether a provider payload selects this event rule.""" + + missing = object() + + def resolve(path: str) -> object: + value: object = payload + for segment in path.split("."): + if not isinstance(value, dict): + return missing + value = value.get(segment, missing) + return value + + return self.required_fields <= payload.keys() and all( + type(value := resolve(name)) is type(expected) and value == expected for name, expected in self.match + ) + + +@dataclass(frozen=True, slots=True) +class StreamClassifierDefinition: + """Describe the event-shape semantics of one guarded provider stream.""" + + subject: str + contract: StreamProjectionContract + rules: tuple[StreamEventRule, ...] + sentinels: tuple[tuple[bytes, str], ...] = () + non_data_shape: str | None = None + event_type_field: str | None = None + require_event_type: bool = False + + +@dataclass(frozen=True, slots=True) +class StreamClassifier: + """Classify provider events using validated stream rules.""" + + subject: str + contract: StreamProjectionContract + rules: tuple[StreamEventRule, ...] + sentinels: tuple[tuple[bytes, str], ...] + non_data_shape: str | None + event_type_field: str | None + require_event_type: bool + + def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: + """Validate and classify one native provider event.""" + + data = event.data + if data is None: + if self.non_data_shape is None: + raise UnsupportedProviderStream(f"{self.subject} events must contain data.") + return GuardedStreamEvent(self.non_data_shape, StreamEventRole.OPAQUE_METADATA) + stripped = data.strip() + for sentinel, shape in self.sentinels: + if stripped == sentinel: + return GuardedStreamEvent(shape, StreamEventRole.OPAQUE_METADATA) + try: + payload = parse_json_object(data) + if self.event_type_field is not None: + event_type = payload.get(self.event_type_field) + if not isinstance(event_type, str): + raise UnsupportedProviderStream( + f"{self.subject} events must declare string field {self.event_type_field!r}." + ) + if self.require_event_type and event.event_type is None: + raise UnsupportedProviderStream(f"{self.subject} SSE event field is required.") + if event.event_type is not None and event.event_type != event_type.encode("ascii"): + raise UnsupportedProviderStream( + f"{self.subject} SSE event field must match data field {self.event_type_field!r}." + ) + matching_rules = [candidate for candidate in self.rules if candidate.matches(payload)] + if not matching_rules: + raise UnsupportedProviderStream(f"{self.subject} event is not covered by the guarded profile.") + if len(matching_rules) > 1: + shapes = ", ".join(repr(rule.shape) for rule in matching_rules) + raise UnsupportedProviderStream(f"{self.subject} event ambiguously matches {shapes}.") + rule = matching_rules[0] + projection = rule.model.model_validate(payload) + text = _resolve_text(projection, rule.text_path) + if text is None and rule.missing_text_role is not None: + return GuardedStreamEvent( + rule.missing_text_shape or rule.shape, + rule.missing_text_role, + rule.missing_text_value, + ) + if text is None and rule.role in {StreamEventRole.GUARDED_TEXT, StreamEventRole.TEXT_SNAPSHOT}: + raise UnsupportedProviderStream(f"{self.subject} event {rule.shape!r} has no guarded text value.") + return GuardedStreamEvent(rule.shape, rule.role, text) + except UnsupportedProviderStream: + raise + except ValidationError as error: + raise UnsupportedProviderStream(guarded_schema_error(error, self.subject)) from error + except (IndexError, TypeError, UnicodeError, ValueError) as error: + raise UnsupportedProviderStream(str(error)) from error + + +def _resolve_text(projection: BaseModel, path: StreamTextPath | None) -> str | None: + if path is None: + return None + value: object = projection + for segment in path: + if isinstance(segment, str) and isinstance(value, BaseModel): + value = getattr(value, segment) + elif isinstance(segment, int) and isinstance(value, list): + if segment < 0: + raise ValueError("Stream event text paths cannot use negative indexes.") + if segment >= len(value): + return None + value = value[segment] + elif isinstance(segment, StreamArraySelection) and isinstance(value, list): + value = segment.select(value) + else: + raise TypeError("Stream event text path does not match its projection.") + if value is None: + return None + if not isinstance(value, str): + raise TypeError("Stream event text must resolve to a string.") + return value + + +def build_stream_classifier(definition: StreamClassifierDefinition) -> StreamClassifier: + """Validate generated event rules and return their runtime classifier.""" + + if not definition.subject: + raise ValueError("A stream classifier definition must identify its provider event subject.") + if not definition.rules: + raise ValueError("A stream classifier definition must declare at least one event rule.") + if definition.subject != definition.contract.projection_id: + raise ValueError("A stream classifier subject must match its projection contract identifier.") + if definition.require_event_type and definition.event_type_field is None: + raise ValueError("A required SSE event field needs a data discriminator.") + signatures: set[tuple[tuple[tuple[str, Any], ...], frozenset[str]]] = set() + + def validate_declared_event(shape: str, role: StreamEventRole, text: str | None = None) -> None: + try: + validate_stream_event(definition.contract, GuardedStreamEvent(shape, role, text)) + except UnsupportedProviderStream as error: + raise ValueError(str(error)) from error + + if definition.non_data_shape is not None: + validate_declared_event(definition.non_data_shape, StreamEventRole.OPAQUE_METADATA) + for _, shape in definition.sentinels: + validate_declared_event(shape, StreamEventRole.OPAQUE_METADATA) + for rule in definition.rules: + if not rule.shape: + raise ValueError("A stream event rule must identify its provider shape.") + if not rule.match and not rule.required_fields: + raise ValueError(f"Stream event rule {rule.shape!r} must declare how it is selected.") + signature = (rule.match, rule.required_fields) + if signature in signatures: + raise ValueError(f"Stream event rule {rule.shape!r} duplicates another selection.") + signatures.add(signature) + carries_text = rule.role in {StreamEventRole.GUARDED_TEXT, StreamEventRole.TEXT_SNAPSHOT} + if carries_text != (rule.text_path is not None): + raise ValueError(f"Stream event rule {rule.shape!r} has an inconsistent text role and path.") + if rule.missing_text_role is not None and rule.text_path is None: + raise ValueError(f"Stream event rule {rule.shape!r} cannot classify missing text without a text path.") + if rule.missing_text_shape is not None and rule.missing_text_role is None: + raise ValueError(f"Stream event rule {rule.shape!r} cannot rename missing text without classifying it.") + if rule.missing_text_value is not None and rule.missing_text_role is None: + raise ValueError(f"Stream event rule {rule.shape!r} cannot supply missing text without classifying it.") + if ( + rule.missing_text_role in {StreamEventRole.GUARDED_TEXT, StreamEventRole.TEXT_SNAPSHOT} + and rule.missing_text_value is None + ): + raise ValueError(f"Stream event rule {rule.shape!r} cannot classify missing text as content.") + validate_declared_event(rule.shape, rule.role, "" if carries_text else None) + if rule.missing_text_role is not None: + validate_declared_event( + rule.missing_text_shape or rule.shape, + rule.missing_text_role, + rule.missing_text_value, + ) + return StreamClassifier( + subject=definition.subject, + contract=definition.contract, + rules=definition.rules, + sentinels=definition.sentinels, + non_data_shape=definition.non_data_shape, + event_type_field=definition.event_type_field, + require_event_type=definition.require_event_type, + ) diff --git a/tests/server/experimental/test_content_checker.py b/tests/server/experimental/test_content_checker.py index d4a8a8b67d..339df5ec41 100644 --- a/tests/server/experimental/test_content_checker.py +++ b/tests/server/experimental/test_content_checker.py @@ -25,6 +25,7 @@ InputContentCheck, InvalidContentChecker, OutputContentCheck, + StreamBufferingPolicy, UnsupportedContentCheckerConfiguration, UnsupportedContentModification, validate_content_check_decision, @@ -157,3 +158,27 @@ def test_input_check_carries_the_guarded_message(): message = GuardedMessage("user", "question") assert InputContentCheck(message).message is message + + +def test_stream_buffering_policy_is_part_of_output_inspection(): + buffering = StreamBufferingPolicy(chunk_size=2, context_size=1) + + assert ContentInspectionPolicy(True, True, buffering).stream_buffering is buffering + + +@pytest.mark.parametrize( + "buffering", + [ + pytest.param(lambda: StreamBufferingPolicy(0, 0), id="zero-chunk"), + pytest.param(lambda: StreamBufferingPolicy(1, -1), id="negative-context"), + pytest.param(lambda: StreamBufferingPolicy(1, 0, 1), id="non-boolean-release"), + ], +) +def test_stream_buffering_policy_rejects_invalid_values(buffering): + with pytest.raises((TypeError, ValueError)): + buffering() + + +def test_stream_buffering_requires_output_inspection(): + with pytest.raises(ValueError, match="requires output inspection"): + ContentInspectionPolicy(True, False, StreamBufferingPolicy(1, 0)) diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py new file mode 100644 index 0000000000..c45e5f5a12 --- /dev/null +++ b/tests/server/experimental/test_guarded_stream.py @@ -0,0 +1,353 @@ +# 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 ( + OperationBlocked, + OperationCheckFailed, + OperationModificationUnsupported, + OperationProjectionFailed, +) +from nemoguardrails.server.experimental._content_checker import ( + ContentAllowed, + ContentBlocked, + ContentCheckFailed, + StreamBufferingPolicy, +) +from nemoguardrails.server.experimental._guarded_stream import ( + StreamProcessingFailed, + StreamUpstreamFailed, + UnsupportedStreamInspection, + guard_provider_stream, + validate_streaming_policy, +) +from nemoguardrails.server.experimental._http_kernel import BufferedHttpResponse +from nemoguardrails.server.experimental._json_payload import parse_json_object +from nemoguardrails.server.experimental.provider.stream import ( + ClassifiedStreamAdapter, + GuardedStreamEvent, + StreamCapabilityProfile, + StreamEventRole, + StreamProjectionContract, + StreamShapeCoverage, + UnsupportedProviderStream, +) +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"}), + snapshot_shapes=frozenset({"snapshot"}), + opaque_shapes=frozenset({"metadata", "end"}), + provider_error_shapes=frozenset({"error"}), + ), +) + + +class StaticChecker: + def __init__(self, decisions=()): + self.decisions = list(decisions) + self.calls = [] + + async def check_output(self, check): + self.calls.append(check) + decision = self.decisions.pop(0) if self.decisions else ContentAllowed() + if isinstance(decision, Exception): + raise decision + return decision + + +class Classifier: + contract = _CONTRACT + + def classify_event(self, event): + if event.data == b"[END]": + return GuardedStreamEvent("end", StreamEventRole.OPAQUE_METADATA) + try: + payload = parse_json_object(event.data or b"") + except ValueError as error: + raise UnsupportedProviderStream("unsupported fake stream event") from error + if "error" in payload: + return GuardedStreamEvent("error", StreamEventRole.PROVIDER_ERROR) + if "snapshot" in payload: + return GuardedStreamEvent("snapshot", StreamEventRole.TEXT_SNAPSHOT, payload["snapshot"]) + if "text" in payload: + return GuardedStreamEvent("text", StreamEventRole.GUARDED_TEXT, payload["text"]) + return GuardedStreamEvent("metadata", StreamEventRole.OPAQUE_METADATA) + + +class Hooks: + def __init__(self): + self.observed = [] + self.finished = False + self.encoded = [] + + def observe_event(self, event, classified): + self.observed.append((event, classified)) + + def validate_end_of_stream(self): + self.finished = True + + def encode_error(self, rendered_body): + self.encoded.append(rendered_body) + return (b"data: " + rendered_body + b"\n\n", b"data: [END]\n\n") + + +class Source: + def __init__(self, chunks, failure=None): + self.chunks = iter(chunks) + self.failure = failure + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self.chunks) + except StopIteration: + if self.failure is not None: + failure, self.failure = self.failure, None + raise failure + raise StopAsyncIteration + + async def aclose(self): + self.closed = True + + +def render_outcome(outcome): + if isinstance(outcome, OperationBlocked): + code = "blocked" + elif isinstance(outcome, OperationCheckFailed): + code = "check_failed" + elif isinstance(outcome, OperationModificationUnsupported): + code = "modification_unsupported" + elif isinstance(outcome, OperationProjectionFailed): + code = "unsupported_stream" + elif isinstance(outcome, StreamUpstreamFailed): + code = "upstream_failed" + elif isinstance(outcome, StreamProcessingFailed): + code = "processing_failed" + else: + raise AssertionError("unknown outcome") + return BufferedHttpResponse(200, (), json.dumps({"code": code}).encode()) + + +def stream_event(payload): + return b"data: " + json.dumps(payload, separators=(",", ":")).encode() + b"\n\n" + + +async def collect(source, checker, hooks, policy=StreamBufferingPolicy(1, 0), **limits): + return b"".join( + [ + chunk + async for chunk in guard_provider_stream( + source, + checker=checker, + streaming_policy=policy, + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(Classifier(), hooks), + render_outcome=render_outcome, + **limits, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_guarded_stream_checks_windows_before_releasing_original_events(): + events = [ + stream_event({"meta": 1}), + stream_event({"text": "one "}), + stream_event({"text": "two "}), + stream_event({"text": "three"}), + b"data: [END]\n\n", + ] + source = Source([b"".join(events)[:23], b"".join(events)[23:]]) + checker = StaticChecker() + hooks = Hooks() + + result = await collect(source, checker, hooks, StreamBufferingPolicy(chunk_size=2, context_size=1)) + + assert result == b"".join(events) + assert [call.output_content for call in checker.calls] == ["one two ", "two three"] + assert hooks.finished is True + assert source.closed is True + + +@pytest.mark.asyncio +async def test_zero_context_does_not_retain_the_previous_window(): + events = [stream_event({"text": text}) for text in ("one ", "two ", "three ", "four")] + checker = StaticChecker() + + result = await collect( + Source([b"".join(events)]), + checker, + Hooks(), + StreamBufferingPolicy(chunk_size=2, context_size=0), + ) + + assert result == b"".join(events) + assert [call.output_content for call in checker.calls] == ["one two ", "three four"] + + +@pytest.mark.asyncio +async def test_stream_block_suppresses_the_complete_pending_window(): + first = stream_event({"text": "unsafe "}) + second = stream_event({"text": "content"}) + checker = StaticChecker([ContentBlocked("blocked", "rule")]) + + result = await collect( + Source([first + second + b"data: [END]\n\n"]), + checker, + Hooks(), + StreamBufferingPolicy(2, 0), + ) + + assert first not in result + assert second not in result + assert b'"code": "blocked"' in result + assert result.endswith(b"data: [END]\n\n") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("decision", [ContentCheckFailed("failed"), RuntimeError("checker failed")]) +async def test_stream_check_failure_is_distinct_from_provider_drift(decision): + result = await collect(Source([stream_event({"text": "answer"})]), StaticChecker([decision]), Hooks()) + + assert b'"code": "check_failed"' in result + assert b"unsupported_stream" not in result + + +@pytest.mark.asyncio +async def test_stream_replacement_is_rejected_without_releasing_pending_text(): + event = stream_event({"text": "answer"}) + + result = await collect(Source([event]), StaticChecker([ContentAllowed(replacement="changed")]), Hooks()) + + assert event not in result + assert b'"code": "modification_unsupported"' in result + + +@pytest.mark.asyncio +async def test_provider_error_events_are_preserved_without_being_checked_as_text(): + error = stream_event({"error": {"message": "upstream failed"}}) + terminal = b"data: [END]\n\n" + checker = StaticChecker() + + result = await collect(Source([error + terminal]), checker, Hooks(), StreamBufferingPolicy(2, 0)) + + assert result == error + terminal + assert checker.calls == [] + + +@pytest.mark.asyncio +async def test_snapshot_must_equal_the_guarded_delta_history(): + source = Source([stream_event({"text": "answer"}) + stream_event({"snapshot": "different"})]) + + result = await collect(source, StaticChecker(), Hooks()) + + assert b'"code": "unsupported_stream"' in result + + +@pytest.mark.asyncio +async def test_unsupported_event_and_pending_limit_fail_closed(): + unsupported = await collect(Source([b"data: not-json\n\n"]), StaticChecker(), Hooks()) + oversized = await collect( + Source([stream_event({"meta": "long"})]), + StaticChecker(), + Hooks(), + max_pending_bytes=5, + ) + + assert b'"code": "unsupported_stream"' in unsupported + assert b'"code": "unsupported_stream"' in oversized + + +@pytest.mark.asyncio +async def test_manual_adapter_contract_mismatch_fails_closed(): + class MismatchedClassifier(Classifier): + def classify_event(self, _event): + return GuardedStreamEvent("text", StreamEventRole.OPAQUE_METADATA) + + hooks = Hooks() + result = b"".join( + [ + chunk + async for chunk in guard_provider_stream( + Source([stream_event({"text": "answer"})]), + checker=StaticChecker(), + streaming_policy=StreamBufferingPolicy(1, 0), + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(MismatchedClassifier(), hooks), + render_outcome=render_outcome, + ) + ] + ) + + assert b'"code": "unsupported_stream"' in result + assert hooks.finished is False + + +@pytest.mark.asyncio +async def test_uninspected_stream_is_relayed_by_chunk_without_parsing(): + chunks = [b"not sse", b" and still opaque"] + source = Source(chunks) + + result = await collect(source, StaticChecker(), Hooks(), policy=None) + + assert result == b"".join(chunks) + assert source.closed is True + + +@pytest.mark.asyncio +async def test_cancelling_downstream_iteration_closes_the_upstream_source(): + source = Source([b"first", b"second"]) + stream = guard_provider_stream( + source, + checker=StaticChecker(), + streaming_policy=None, + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(Classifier(), Hooks()), + render_outcome=render_outcome, + ) + + assert await anext(stream) == b"first" + await stream.aclose() + + assert source.closed is True + + +@pytest.mark.asyncio +async def test_upstream_failure_is_distinct_from_processing_failure(): + source = Source([], RuntimeError("upstream disconnected")) + + result = await collect(source, StaticChecker(), Hooks()) + + assert b'"code": "upstream_failed"' in result + assert source.closed is True + + +def test_stream_policy_must_buffer_before_release(): + with pytest.raises(UnsupportedStreamInspection, match="requires a buffering policy"): + validate_streaming_policy(None) + + with pytest.raises(UnsupportedStreamInspection, match="cannot release"): + validate_streaming_policy(StreamBufferingPolicy(1, 0, True)) diff --git a/tests/server/experimental/test_sse.py b/tests/server/experimental/test_sse.py new file mode 100644 index 0000000000..8e871a250d --- /dev/null +++ b/tests/server/experimental/test_sse.py @@ -0,0 +1,49 @@ +# 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 pytest + +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent, ServerSentEventTooLarge, iter_sse_events + + +async def _chunks(*values): + for value in values: + yield value + + +def test_sse_event_exposes_joined_data_and_last_event_type_without_changing_bytes(): + raw = b"event: first\r\nevent: final\r\ndata: one\r\ndata:two\r\n\r\n" + event = ServerSentEvent.from_bytes(raw) + + assert event.raw == raw + assert event.event_type == b"final" + assert event.data == b"one\ntwo" + + +@pytest.mark.asyncio +async def test_sse_parser_preserves_events_split_across_arbitrary_chunks(): + raw = b": metadata\n\ndata: first\n\ndata: second" + + events = [event async for event in iter_sse_events(_chunks(raw[:7], raw[7:19], raw[19:]), max_event_bytes=30)] + + assert [event.raw for event in events] == [b": metadata\n\n", b"data: first\n\n", b"data: second"] + + +@pytest.mark.asyncio +async def test_sse_parser_bounds_complete_and_incomplete_events(): + with pytest.raises(ServerSentEventTooLarge): + _ = [event async for event in iter_sse_events(_chunks(b"data: oversized\n\n"), max_event_bytes=8)] + with pytest.raises(ServerSentEventTooLarge): + _ = [event async for event in iter_sse_events(_chunks(b"data: unfinished"), max_event_bytes=8)] diff --git a/tests/server/experimental/test_stream_classifier.py b/tests/server/experimental/test_stream_classifier.py new file mode 100644 index 0000000000..c91ebb8b27 --- /dev/null +++ b/tests/server/experimental/test_stream_classifier.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. + +from typing import Literal + +import pytest +from pydantic import BaseModel + +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent +from nemoguardrails.server.experimental.provider.stream import ( + GuardedStreamEvent, + StreamCapabilityProfile, + StreamEventRole, + StreamProjectionContract, + StreamShapeCoverage, + UnsupportedProviderStream, + validate_stream_event, +) +from nemoguardrails.server.experimental.provider.stream_classifier import ( + StreamArraySelection, + StreamClassifierDefinition, + StreamEventRule, + build_stream_classifier, +) + + +class Delta(BaseModel): + text: str | None = None + + +class Chunk(BaseModel): + kind: Literal["chunk"] + deltas: list[Delta] + + +class Error(BaseModel): + error: str + + +class SnapshotItem(BaseModel): + kind: str + text: str + + +class Snapshot(BaseModel): + kind: Literal["snapshot"] + items: list[SnapshotItem] + + +CONTRACT = StreamProjectionContract( + projection_id="test.stream", + profile=StreamCapabilityProfile.SINGLE_TEXT_DELTA_V1, + shapes=StreamShapeCoverage( + guarded_shapes=frozenset({"chunk:text"}), + snapshot_shapes=frozenset({"snapshot"}), + opaque_shapes=frozenset({"done", "chunk:metadata"}), + provider_error_shapes=frozenset({"error"}), + ), +) + + +def classifier(): + return build_stream_classifier( + StreamClassifierDefinition( + subject="test.stream", + contract=CONTRACT, + sentinels=((b"[DONE]", "done"),), + rules=( + StreamEventRule( + shape="error", + model=Error, + role=StreamEventRole.PROVIDER_ERROR, + required_fields=frozenset({"error"}), + ), + StreamEventRule( + shape="chunk:text", + model=Chunk, + role=StreamEventRole.GUARDED_TEXT, + match=(("kind", "chunk"),), + text_path=("deltas", 0, "text"), + missing_text_role=StreamEventRole.OPAQUE_METADATA, + missing_text_shape="chunk:metadata", + ), + ), + ) + ) + + +@pytest.mark.parametrize( + ("raw", "shape", "role", "text"), + [ + (b"data: [DONE]\n\n", "done", StreamEventRole.OPAQUE_METADATA, None), + ( + b'data: {"kind":"chunk","deltas":[{"text":"hello"}]}\n\n', + "chunk:text", + StreamEventRole.GUARDED_TEXT, + "hello", + ), + ( + b'data: {"kind":"chunk","deltas":[]}\n\n', + "chunk:metadata", + StreamEventRole.OPAQUE_METADATA, + None, + ), + (b'data: {"error":"failed"}\n\n', "error", StreamEventRole.PROVIDER_ERROR, None), + ], +) +def test_classifier_validates_and_classifies_reviewed_shapes(raw, shape, role, text): + event = classifier().classify_event(ServerSentEvent.from_bytes(raw)) + + assert (event.shape, event.role, event.text) == (shape, role, text) + + +def test_classifier_rejects_unknown_and_ambiguous_events(): + with pytest.raises(UnsupportedProviderStream, match="not covered"): + classifier().classify_event(ServerSentEvent.from_bytes(b'data: {"kind":"future"}\n\n')) + + with pytest.raises(UnsupportedProviderStream, match="ambiguously matches"): + classifier().classify_event( + ServerSentEvent.from_bytes(b'data: {"kind":"chunk","error":"failed","deltas":[]}\n\n') + ) + + +def test_classifier_can_select_one_nested_text_item(): + snapshot_classifier = build_stream_classifier( + StreamClassifierDefinition( + subject="test.stream", + contract=CONTRACT, + rules=( + StreamEventRule( + shape="snapshot", + model=Snapshot, + role=StreamEventRole.TEXT_SNAPSHOT, + match=(("kind", "snapshot"),), + text_path=( + "items", + StreamArraySelection(match=(("kind", "message"),), cardinality="at_most_one"), + "text", + ), + missing_text_role=StreamEventRole.TEXT_SNAPSHOT, + missing_text_value="", + ), + ), + ) + ) + + empty = snapshot_classifier.classify_event(ServerSentEvent.from_bytes(b'data: {"kind":"snapshot","items":[]}\n\n')) + populated = snapshot_classifier.classify_event( + ServerSentEvent.from_bytes(b'data: {"kind":"snapshot","items":[{"kind":"message","text":"hello"}]}\n\n') + ) + + assert empty.text == "" + assert populated.text == "hello" + + +def test_stream_contract_rejects_unknown_or_misclassified_events(): + with pytest.raises(UnsupportedProviderStream, match="contract declares"): + validate_stream_event(CONTRACT, GuardedStreamEvent("chunk:text", StreamEventRole.OPAQUE_METADATA)) + + +def test_classifier_definition_must_agree_with_the_stream_contract(): + definition = StreamClassifierDefinition( + subject="test.stream", + contract=CONTRACT, + rules=( + StreamEventRule( + shape="chunk:text", + model=Chunk, + role=StreamEventRole.OPAQUE_METADATA, + match=(("kind", "chunk"),), + ), + ), + ) + + with pytest.raises(ValueError, match="contract declares 'guarded_text'"): + build_stream_classifier(definition) + + +def test_stream_shape_coverage_rejects_multiple_classifications(): + with pytest.raises(ValueError, match="multiple classifications"): + StreamShapeCoverage( + guarded_shapes=frozenset({"chunk"}), + opaque_shapes=frozenset({"chunk"}), + ) + + +@pytest.mark.parametrize("expected,other", [(False, 0), (True, 1), (0, False), (1, True), (1, 1.0)]) +def test_stream_discriminators_require_matching_types(expected, other): + correct = {"kind": expected} + incorrect = {"kind": other} + selector = StreamArraySelection(match=(("kind", expected),), cardinality="exactly_one") + rule = StreamEventRule( + shape="chunk:text", model=Chunk, role=StreamEventRole.GUARDED_TEXT, match=(("kind", expected),) + ) + + assert selector.select([incorrect, correct]) is correct + assert rule.matches(correct) is True + assert rule.matches(incorrect) is False From 558d6021d4dc57082ad010d7c5117b68b1e9e13c Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:24:58 +0200 Subject: [PATCH 02/16] fix(server): frame guarded SSE streams incrementally Parse inspected streams line by line so mixed CR, LF, and CRLF endings and terminators split across chunks frame correctly in linear time. Strip one initial BOM for field parsing while keeping it in raw bytes. Release complete SSE blocks only, drop unfinished comment tails, and reject unfinished field blocks as truncated provider streams. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_stream.py | 21 ++- .../server/experimental/provider/sse.py | 119 +++++++++++----- .../server/experimental/provider/stream.py | 4 +- .../experimental/test_guarded_stream.py | 56 ++++++++ tests/server/experimental/test_sse.py | 127 +++++++++++++++++- 5 files changed, 286 insertions(+), 41 deletions(-) diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index c3b9988b0d..186062b2fa 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -41,6 +41,7 @@ from nemoguardrails.server.experimental.provider.sse import ( ServerSentEvent, ServerSentEventTooLarge, + TruncatedServerSentEvent, iter_sse_events, ) from nemoguardrails.server.experimental.provider.stream import ( @@ -191,12 +192,21 @@ async def guard_provider_stream( max_event_bytes: int = 1024 * 1024, max_pending_bytes: int = 10 * 1024 * 1024, ) -> AsyncIterator[bytes]: - """Release provider SSE events only after required output checks pass.""" + """Release provider SSE events only after required output checks pass. - if input_message.role != "user": - raise ValueError("A guarded stream requires a user input message.") - if any(type(limit) is not int or limit <= 0 for limit in (max_event_bytes, max_pending_bytes)): - raise ValueError("Stream byte limits must be positive.") + Iteration closes ``source`` when it ends, including when arguments are + rejected. A stream that is never iterated does not run that cleanup, so its + caller must close ``source``. + """ + + try: + if input_message.role != "user": + raise ValueError("A guarded stream requires a user input message.") + if any(type(limit) is not int or limit <= 0 for limit in (max_event_bytes, max_pending_bytes)): + raise ValueError("Stream byte limits must be positive.") + except ValueError: + await _close_source(source) + raise try: if streaming_policy is None: @@ -279,6 +289,7 @@ async def text_deltas() -> AsyncIterator[str]: except ( PendingStreamBufferTooLarge, ServerSentEventTooLarge, + TruncatedServerSentEvent, StreamSnapshotHistoryTooLarge, UnsupportedProviderStream, ) as failure: diff --git a/nemoguardrails/server/experimental/provider/sse.py b/nemoguardrails/server/experimental/provider/sse.py index 2bc78268b4..844b176544 100644 --- a/nemoguardrails/server/experimental/provider/sse.py +++ b/nemoguardrails/server/experimental/provider/sse.py @@ -19,7 +19,7 @@ event type and joined data fields needed by provider adapters. """ -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass @@ -27,6 +27,10 @@ class ServerSentEventTooLarge(ValueError): """Report an SSE event that exceeds the configured buffering limit.""" +class TruncatedServerSentEvent(ValueError): + """Report an unfinished SSE block containing fields at end of stream.""" + + @dataclass(frozen=True, slots=True) class ServerSentEvent: """Represent one complete SSE event. @@ -45,18 +49,17 @@ def from_bytes(cls, raw: bytes) -> "ServerSentEvent": """Create an event while retaining its original representation. Args: - raw: Complete SSE event bytes. + raw: Exactly one complete SSE block, without stream-start BOM handling. Returns: The parsed event representation. """ - body = raw - for separator in (b"\r\n\r\n", b"\n\n", b"\r\r"): - if body.endswith(separator): - body = body[: -len(separator)] - break - return cls(raw=raw, lines=tuple(body.replace(b"\r\n", b"\n").replace(b"\r", b"\n").split(b"\n"))) + framer = _SSEFramer(max_event_bytes=max(1, len(raw)), stream_start=False) + events = [*framer.feed(raw), *framer.finish()] + if len(events) != 1 or events[0].raw != raw: + raise ValueError("Expected exactly one complete SSE block.") + return events[0] @property def data(self) -> bytes | None: @@ -83,13 +86,72 @@ def event_type(self) -> bytes | None: return value -def _event_boundary(buffer: bytes) -> tuple[int, int] | None: - boundaries = [] - for separator in (b"\r\n\r\n", b"\n\n", b"\r\r"): - index = buffer.find(separator) - if index >= 0: - boundaries.append((index, len(separator))) - return min(boundaries) if boundaries else None +class _SSEFramer: + def __init__(self, *, max_event_bytes: int, stream_start: bool = True): + self.max_event_bytes = max_event_bytes + self.raw = bytearray() + self.lines: list[bytes] = [] + self.line_start = 0 + self.pending_cr = False + self.stream_start = stream_start + + def _append(self, value: int) -> None: + if len(self.raw) >= self.max_event_bytes: + raise ServerSentEventTooLarge + self.raw.append(value) + + def _field_line(self, line: bytes) -> bytes: + if self.stream_start: + self.stream_start = False + return line.removeprefix(b"\xef\xbb\xbf") + return line + + def _end_line(self, ending_length: int) -> ServerSentEvent | None: + line = self._field_line(bytes(self.raw[self.line_start : -ending_length])) + if line: + self.lines.append(line) + self.line_start = len(self.raw) + return None + event = ServerSentEvent(bytes(self.raw), tuple(self.lines)) + self.raw.clear() + self.lines.clear() + self.line_start = 0 + return event + + def feed(self, chunk: bytes) -> Iterator[ServerSentEvent]: + for value in chunk: + if self.pending_cr: + self.pending_cr = False + if value == 10: + self._append(value) + event = self._end_line(2) + if event is not None: + yield event + continue + event = self._end_line(1) + if event is not None: + yield event + self._append(value) + if value == 13: + self.pending_cr = True + elif value == 10: + event = self._end_line(1) + if event is not None: + yield event + + def finish(self) -> Iterator[ServerSentEvent]: + if self.pending_cr: + self.pending_cr = False + event = self._end_line(1) + if event is not None: + yield event + if self.raw: + tail = self._field_line(bytes(self.raw[self.line_start :])) + if any(line and not line.startswith(b":") for line in (*self.lines, tail)): + raise TruncatedServerSentEvent("The provider stream ends with an unfinished SSE field block.") + self.raw.clear() + self.lines.clear() + self.line_start = 0 async def iter_sse_events( @@ -97,7 +159,10 @@ async def iter_sse_events( *, max_event_bytes: int = 1024 * 1024, ) -> AsyncIterator[ServerSentEvent]: - """Yield complete SSE events from arbitrary HTTP byte chunks. + """Yield complete SSE blocks from arbitrary HTTP byte chunks. + + Preserve raw bytes, including an initial BOM ignored only for field parsing. + Drop unfinished comment tails; reject all other unfinished field blocks. Args: source: Raw upstream response byte stream. @@ -109,23 +174,15 @@ async def iter_sse_events( Raises: ValueError: If ``max_event_bytes`` is not positive. ServerSentEventTooLarge: If one event exceeds the limit. + TruncatedServerSentEvent: If the stream ends in an unfinished field block. """ - if max_event_bytes <= 0: + if type(max_event_bytes) is not int or max_event_bytes <= 0: raise ValueError("max_event_bytes must be positive.") - buffer = b"" + framer = _SSEFramer(max_event_bytes=max_event_bytes) async for chunk in source: - buffer += chunk - while boundary := _event_boundary(buffer): - index, separator_length = boundary - event_length = index + separator_length - if event_length > max_event_bytes: - raise ServerSentEventTooLarge - raw, buffer = buffer[:event_length], buffer[event_length:] - yield ServerSentEvent.from_bytes(raw) - if len(buffer) > max_event_bytes: - raise ServerSentEventTooLarge - - if buffer: - yield ServerSentEvent.from_bytes(buffer) + for event in framer.feed(chunk): + yield event + for event in framer.finish(): + yield event diff --git a/nemoguardrails/server/experimental/provider/stream.py b/nemoguardrails/server/experimental/provider/stream.py index 3bbf071c06..a175d6d604 100644 --- a/nemoguardrails/server/experimental/provider/stream.py +++ b/nemoguardrails/server/experimental/provider/stream.py @@ -153,7 +153,7 @@ def validate_end_of_stream(self) -> None: ... def encode_error(self, rendered_body: bytes) -> Sequence[bytes]: - """Frame a proxy error and its provider-native body for the stream.""" + """Frame a proxy error and its provider-native body as complete SSE blocks.""" ... @@ -174,7 +174,7 @@ def validate_end_of_stream(self) -> None: ... def encode_error(self, rendered_body: bytes) -> Sequence[bytes]: - """Frame a proxy error and its provider-native body for the stream.""" + """Frame a proxy error and its provider-native body as complete SSE blocks.""" ... diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index c45e5f5a12..1bfd540adb 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -351,3 +351,59 @@ def test_stream_policy_must_buffer_before_release(): with pytest.raises(UnsupportedStreamInspection, match="cannot release"): validate_streaming_policy(StreamBufferingPolicy(1, 0, True)) + + +@pytest.mark.asyncio +async def test_truncation_withholds_pending_bytes_and_error_starts_on_complete_boundary(): + approved = stream_event({"text": "approved"}) + incomplete = stream_event({"text": "hidden"})[:-2] + source = Source([approved, incomplete]) + hooks = Hooks() + + result = await collect(source, StaticChecker(), hooks) + + assert result.startswith(approved + b"data: ") + assert b"hidden" not in result + assert b'"code": "unsupported_stream"' in result + assert hooks.finished is False + assert source.closed is True + + +@pytest.mark.asyncio +async def test_end_validation_failure_withholds_pending_window(): + class RejectingHooks(Hooks): + def validate_end_of_stream(self): + raise UnsupportedProviderStream("truncated provider lifecycle") + + checker = StaticChecker() + source = Source([stream_event({"text": "hidden"})]) + + result = await collect(source, checker, RejectingHooks(), policy=StreamBufferingPolicy(2, 0)) + + assert b"hidden" not in result + assert b'"code": "unsupported_stream"' in result + assert checker.calls == [] + assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("role", "limits"), + [("assistant", {}), ("user", {"max_event_bytes": 0}), ("user", {"max_pending_bytes": -1})], +) +async def test_rejected_arguments_raise_and_close_source(role, limits): + source = Source([stream_event({"text": "answer"})]) + stream = guard_provider_stream( + source, + checker=StaticChecker(), + streaming_policy=StreamBufferingPolicy(1, 0), + input_message=GuardedMessage(role, "question"), + adapter=ClassifiedStreamAdapter(Classifier(), Hooks()), + render_outcome=render_outcome, + **limits, + ) + + with pytest.raises(ValueError): + await anext(stream) + + assert source.closed is True diff --git a/tests/server/experimental/test_sse.py b/tests/server/experimental/test_sse.py index 8e871a250d..0c54358d92 100644 --- a/tests/server/experimental/test_sse.py +++ b/tests/server/experimental/test_sse.py @@ -15,7 +15,12 @@ import pytest -from nemoguardrails.server.experimental.provider.sse import ServerSentEvent, ServerSentEventTooLarge, iter_sse_events +from nemoguardrails.server.experimental.provider.sse import ( + ServerSentEvent, + ServerSentEventTooLarge, + TruncatedServerSentEvent, + iter_sse_events, +) async def _chunks(*values): @@ -23,6 +28,11 @@ async def _chunks(*values): yield value +async def _byte_chunks(raw): + for byte in raw: + yield bytes([byte]) + + def test_sse_event_exposes_joined_data_and_last_event_type_without_changing_bytes(): raw = b"event: first\r\nevent: final\r\ndata: one\r\ndata:two\r\n\r\n" event = ServerSentEvent.from_bytes(raw) @@ -34,11 +44,11 @@ def test_sse_event_exposes_joined_data_and_last_event_type_without_changing_byte @pytest.mark.asyncio async def test_sse_parser_preserves_events_split_across_arbitrary_chunks(): - raw = b": metadata\n\ndata: first\n\ndata: second" + raw = b": metadata\n\ndata: first\n\ndata: second\n\n" events = [event async for event in iter_sse_events(_chunks(raw[:7], raw[7:19], raw[19:]), max_event_bytes=30)] - assert [event.raw for event in events] == [b": metadata\n\n", b"data: first\n\n", b"data: second"] + assert [event.raw for event in events] == [b": metadata\n\n", b"data: first\n\n", b"data: second\n\n"] @pytest.mark.asyncio @@ -47,3 +57,114 @@ async def test_sse_parser_bounds_complete_and_incomplete_events(): _ = [event async for event in iter_sse_events(_chunks(b"data: oversized\n\n"), max_event_bytes=8)] with pytest.raises(ServerSentEventTooLarge): _ = [event async for event in iter_sse_events(_chunks(b"data: unfinished"), max_event_bytes=8)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ending", [b"\n\n", b"\n\r", b"\n\r\n", b"\r\r", b"\r\r\n", b"\r\n\n", b"\r\n\r", b"\r\n\r\n"]) +async def test_all_event_line_endings_at_every_chunk_split(ending): + blocks = [b"data: one" + ending, b"data: two" + ending] + raw = b"".join(blocks) + + for split in range(len(raw) + 1): + events = [event async for event in iter_sse_events(_chunks(raw[:split], b"", raw[split:]))] + assert [event.raw for event in events] == blocks + assert [event.data for event in events] == [b"one", b"two"] + assert ServerSentEvent.from_bytes(blocks[0]) == events[0] + + +@pytest.mark.asyncio +async def test_bom_is_ignored_only_at_stream_start_and_preserved_in_raw(): + blocks = [b"\xef\xbb\xbfdata: one\n\n", b"\xef\xbb\xbfdata: ignored\n\n", b"data: \xef\xbb\xbfvalue\n\n"] + raw = b"".join(blocks) + + for split in range(len(raw) + 1): + events = [event async for event in iter_sse_events(_chunks(raw[:split], raw[split:]))] + assert [event.raw for event in events] == blocks + assert [event.data for event in events] == [b"one", None, b"\xef\xbb\xbfvalue"] + assert ServerSentEvent.from_bytes(blocks[0]).data is None + + +@pytest.mark.asyncio +async def test_comments_empty_blocks_and_unknown_fields_preserve_exact_bytes(): + blocks = [b"\n", b": keepalive\r\n\r", b"unknown: value\ndata\n\n"] + raw = b"".join(blocks) + + events = [event async for event in iter_sse_events(_chunks(*(bytes([byte]) for byte in raw)))] + + assert [event.raw for event in events] == blocks + assert [event.data for event in events] == [None, None, b""] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tail", + [ + b"data: partial", + b"data: partial\n", + b"data: partial\r", + b"event: future", + b"id: 1", + b"retry: 1", + b"unknown", + b" ", + b": comment\ndata: partial", + ], +) +async def test_unfinished_field_blocks_are_rejected(tail): + with pytest.raises(TruncatedServerSentEvent): + _ = [event async for event in iter_sse_events(_chunks(tail))] + with pytest.raises(ValueError): + ServerSentEvent.from_bytes(tail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tail", [b"", b": comment", b": comment\n", b": comment\r", b": one\n: two", b"\xef\xbb\xbf: comment"] +) +async def test_unfinished_comment_tails_are_discarded(tail): + events = [event async for event in iter_sse_events(_chunks(tail))] + + assert events == [] + + +@pytest.mark.asyncio +async def test_event_limit_counts_terminators_bom_and_unfinished_lines(): + block = b"\xef\xbb\xbfdata: one\r\n\r\n" + events = [event async for event in iter_sse_events(_chunks(block * 3), max_event_bytes=len(block))] + assert len(events) == 3 + for raw in (block, b"x" * len(block)): + with pytest.raises(ServerSentEventTooLarge): + _ = [ + event + async for event in iter_sse_events( + _chunks(*(bytes([byte]) for byte in raw)), max_event_bytes=len(block) - 1 + ) + ] + + +def test_from_bytes_rejects_multiple_blocks(): + with pytest.raises(ValueError, match="exactly one"): + ServerSentEvent.from_bytes(b"data: one\n\ndata: two\n\n") + + +@pytest.mark.asyncio +async def test_many_small_events_in_one_large_chunk(): + block = b"data: x\n\n" + count = 1024 * 1024 // len(block) + observed = 0 + + async for event in iter_sse_events(_chunks(block * count), max_event_bytes=len(block)): + assert event.raw == block + observed += 1 + + assert observed == count + + +@pytest.mark.asyncio +async def test_large_event_arriving_one_byte_per_chunk(): + raw = b"data: " + b"x" * (1024 * 1024 - 8) + b"\n\n" + + events = [event async for event in iter_sse_events(_byte_chunks(raw), max_event_bytes=len(raw))] + + assert len(events) == 1 + assert events[0].raw == raw From fbf35f3d6f25fb83dd661f14040c0aa7fd1db52f Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:28:24 +0200 Subject: [PATCH 03/16] fix(server): release safe stream keepalive prefixes promptly Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_stream.py | 18 +++-- .../experimental/test_guarded_stream.py | 71 +++++++++++++++++++ 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index 186062b2fa..07cc22ea23 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -134,14 +134,17 @@ def _take_checked_events( async def _text_windows( - source: AsyncIterator[str], + source: AsyncIterator[str | ServerSentEvent], *, chunk_size: int, context_size: int, -) -> AsyncIterator[_TextWindow]: +) -> AsyncIterator[_TextWindow | ServerSentEvent]: context: list[str] = [] pending: list[str] = [] async for text in source: + if isinstance(text, ServerSentEvent): + yield text + continue pending.append(text) if len(pending) < chunk_size: continue @@ -196,7 +199,8 @@ async def guard_provider_stream( Iteration closes ``source`` when it ends, including when arguments are rejected. A stream that is never iterated does not run that cleanup, so its - caller must close ``source``. + caller must close ``source``. Validated data-less keepalive prefixes are + released promptly; they never overtake pending payloads or terminal events. """ try: @@ -221,11 +225,14 @@ async def guard_provider_stream( observed_text_bytes = 0 snapshot_history_complete = True - async def text_deltas() -> AsyncIterator[str]: + async def text_deltas() -> AsyncIterator[str | ServerSentEvent]: nonlocal observed_text_bytes, pending_bytes, snapshot_history_complete async for event in iter_sse_events(_guard_upstream(source), max_event_bytes=max_event_bytes): classified = adapter.classify_event(event) validate_stream_event(adapter.contract, classified) + if not pending and event.data is None and classified.role is StreamEventRole.OPAQUE_METADATA: + yield event + continue if classified.role is StreamEventRole.TEXT_SNAPSHOT: if not snapshot_history_complete: raise StreamSnapshotHistoryTooLarge( @@ -258,6 +265,9 @@ async def text_deltas() -> AsyncIterator[str]: chunk_size=streaming_policy.chunk_size, context_size=streaming_policy.context_size, ): + if isinstance(window, ServerSentEvent): + yield window.raw + continue checked_events = _take_checked_events(pending, window.release_count) pending_bytes -= sum(len(event.raw) for event in checked_events) try: diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index 1bfd540adb..76c8dcd789 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import json import pytest @@ -208,6 +209,76 @@ async def test_zero_context_does_not_retain_the_previous_window(): assert [call.output_content for call in checker.calls] == ["one two ", "three four"] +class KeepaliveClassifier(Classifier): + """Recognize data-less transport events without inventing text.""" + + def classify_event(self, event): + if event.data is None: + return GuardedStreamEvent("metadata", StreamEventRole.OPAQUE_METADATA) + return super().classify_event(event) + + +@pytest.mark.asyncio +async def test_keepalive_prefix_is_released_before_reading_the_next_event(): + """A live idle stream must make progress without waiting for text or EOF.""" + checker = StaticChecker() + closed = [] + + async def source(): + try: + yield b": keepalive\n\n" + await asyncio.Event().wait() + finally: + closed.append(True) + + stream = guard_provider_stream( + source(), + checker=checker, + streaming_policy=StreamBufferingPolicy(2, 0), + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(KeepaliveClassifier(), Hooks()), + render_outcome=render_outcome, + ) + try: + assert await asyncio.wait_for(anext(stream), 1) == b": keepalive\n\n" + assert checker.calls == [] + finally: + await stream.aclose() + assert closed == [True] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefix", [b"", stream_event({"text": "pending"}), b"data: [END]\n\n"]) +async def test_keepalives_preserve_pending_payload_order_and_bounds(prefix): + """Only a transport-only prefix can bypass the pending payload window.""" + keepalives = b": alive\n\n" * 30 + source = Source([prefix + keepalives]) + checker = StaticChecker() + result = b"".join( + [ + part + async for part in guard_provider_stream( + source, + checker=checker, + streaming_policy=StreamBufferingPolicy(2, 0), + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(KeepaliveClassifier(), Hooks()), + render_outcome=render_outcome, + max_event_bytes=128, + max_pending_bytes=128, + ) + ] + ) + if prefix: + assert b"unsupported_stream" in result + assert b": alive" not in result + assert b"pending" not in result + else: + assert result == keepalives + assert checker.calls == [] + assert source.closed is True + + @pytest.mark.asyncio async def test_stream_block_suppresses_the_complete_pending_window(): first = stream_event({"text": "unsafe "}) From 9d4efc097c847e7321afa31eccd18b520dcd3d75 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:20:04 +0200 Subject: [PATCH 04/16] docs(server): state the single user input subject invariant Output checks receive one projected user input subject as context in both buffered and streaming execution. Document that this is a provider-neutral execution invariant rather than a restriction on provider conversation roles, and that relaxing it needs matching bindings, checker semantics, and execution tests, not only provider or profile declarations. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_content_checker.py | 9 ++++++++- nemoguardrails/server/experimental/_guarded_stream.py | 11 +++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/nemoguardrails/server/experimental/_content_checker.py b/nemoguardrails/server/experimental/_content_checker.py index e7c535fbad..425f2e0656 100644 --- a/nemoguardrails/server/experimental/_content_checker.py +++ b/nemoguardrails/server/experimental/_content_checker.py @@ -78,7 +78,14 @@ class InputContentCheck: @dataclass(frozen=True, slots=True) class OutputContentCheck: - """Carry generated text and its effective input context to one checker.""" + """Carry generated text and its effective input context to one checker. + + Buffered and streaming execution currently supply one projected user input + subject as context for assistant-output checks. This carrier does not + itself enforce the role. Supporting other roles or multiple input subjects + requires coordinated changes to bindings, checker semantics, and execution; + changing provider metadata alone does not extend that capability. + """ input_message: GuardedMessage output_content: str diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index 07cc22ea23..422dfc64a4 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -197,6 +197,17 @@ async def guard_provider_stream( ) -> AsyncIterator[bytes]: """Release provider SSE events only after required output checks pass. + The current execution model requires one projected user input subject as + context for assistant-output checks. This is a provider-neutral execution + invariant, not a restriction on roles in the provider's raw conversation. + Bindings supporting other input subjects require corresponding checker and + execution support. + + Change this invariant only when the framework implements a capability for + other input roles or multiple guarded input subjects, with matching bindings, + checker semantics, and buffered and streaming execution tests. Adding a + provider or changing a profile declaration alone does not relax it. + Iteration closes ``source`` when it ends, including when arguments are rejected. A stream that is never iterated does not run that cleanup, so its caller must close ``source``. Validated data-less keepalive prefixes are From ec0b1e32afeaf925760dc9c8f486ae6013052fa9 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 18:13:48 +0200 Subject: [PATCH 05/16] fix(server): reject line-leading BOMs after SSE stream start The SSE standard ignores one byte order mark at stream start. The parser treated a later line that starts with a BOM as an unknown field, so its event had no data and was released as opaque metadata without text checks. Some clients, including openai-node, strip a BOM at later line starts and read the line as a data field, so a provider could deliver text that output rails never inspected. Keep one stream-leading BOM. Reject any other line that starts with a BOM, including BOM-only lines and a second leading BOM, as an ambiguous event that fails the inspected stream closed. The check runs on complete lines, so it is independent of chunk boundaries and line endings. Uninspected relays are not parsed and are unaffected. Tests cover every chunk split, LF, CR, and CRLF endings, BOM-only lines, comments, unfinished tails, withheld events in the guarded stream, and a differential check that the OpenAI Python SDK decoder never yields data that this parser did not parse. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_stream.py | 2 + .../server/experimental/provider/sse.py | 12 +++- .../experimental/test_guarded_stream.py | 16 ++++++ tests/server/experimental/test_sse.py | 55 +++++++++++++++++-- 4 files changed, 80 insertions(+), 5 deletions(-) diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index 422dfc64a4..f858905c34 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -39,6 +39,7 @@ from nemoguardrails.server.experimental._guarded_operation import UnsupportedGuardedPayload from nemoguardrails.server.experimental._http_kernel import BufferedHttpResponse from nemoguardrails.server.experimental.provider.sse import ( + AmbiguousServerSentEvent, ServerSentEvent, ServerSentEventTooLarge, TruncatedServerSentEvent, @@ -308,6 +309,7 @@ async def text_deltas() -> AsyncIterator[str | ServerSentEvent]: for event in pending: yield event.event.raw except ( + AmbiguousServerSentEvent, PendingStreamBufferTooLarge, ServerSentEventTooLarge, TruncatedServerSentEvent, diff --git a/nemoguardrails/server/experimental/provider/sse.py b/nemoguardrails/server/experimental/provider/sse.py index 844b176544..e78bcc8a9d 100644 --- a/nemoguardrails/server/experimental/provider/sse.py +++ b/nemoguardrails/server/experimental/provider/sse.py @@ -31,6 +31,10 @@ class TruncatedServerSentEvent(ValueError): """Report an unfinished SSE block containing fields at end of stream.""" +class AmbiguousServerSentEvent(ValueError): + """Report SSE bytes that common clients interpret differently.""" + + @dataclass(frozen=True, slots=True) class ServerSentEvent: """Represent one complete SSE event. @@ -103,7 +107,12 @@ def _append(self, value: int) -> None: def _field_line(self, line: bytes) -> bytes: if self.stream_start: self.stream_start = False - return line.removeprefix(b"\xef\xbb\xbf") + line = line.removeprefix(b"\xef\xbb\xbf") + # The SSE standard ignores only one BOM at stream start. Some clients + # also strip a BOM at later line starts and read the line as a field, + # so such a line could carry data that this parser treats as unknown. + if line.startswith(b"\xef\xbb\xbf"): + raise AmbiguousServerSentEvent("A provider stream line starts with a byte order mark after stream start.") return line def _end_line(self, ending_length: int) -> ServerSentEvent | None: @@ -175,6 +184,7 @@ async def iter_sse_events( ValueError: If ``max_event_bytes`` is not positive. ServerSentEventTooLarge: If one event exceeds the limit. TruncatedServerSentEvent: If the stream ends in an unfinished field block. + AmbiguousServerSentEvent: If a line after stream start begins with a BOM. """ if type(max_event_bytes) is not int or max_event_bytes <= 0: diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index 76c8dcd789..2b73baf115 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -478,3 +478,19 @@ async def test_rejected_arguments_raise_and_close_source(role, limits): await anext(stream) assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("split", [0, 1, 2, 3]) +async def test_line_leading_bom_mid_stream_fails_closed_without_release(split): + """An ambiguous line is never forwarded, even when its BOM bytes are fragmented.""" + approved = stream_event({"text": "approved"}) + ambiguous = b"\xef\xbb\xbf" + stream_event({"text": "EVIL"}) + source = Source([approved, ambiguous[:split], ambiguous[split:]]) + + result = await collect(source, StaticChecker(), Hooks()) + + assert result.startswith(approved) + assert b"EVIL" not in result + assert b'"code": "unsupported_stream"' in result + assert source.closed is True diff --git a/tests/server/experimental/test_sse.py b/tests/server/experimental/test_sse.py index 0c54358d92..48e1c098b4 100644 --- a/tests/server/experimental/test_sse.py +++ b/tests/server/experimental/test_sse.py @@ -13,9 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio + import pytest from nemoguardrails.server.experimental.provider.sse import ( + AmbiguousServerSentEvent, ServerSentEvent, ServerSentEventTooLarge, TruncatedServerSentEvent, @@ -74,14 +77,35 @@ async def test_all_event_line_endings_at_every_chunk_split(ending): @pytest.mark.asyncio async def test_bom_is_ignored_only_at_stream_start_and_preserved_in_raw(): - blocks = [b"\xef\xbb\xbfdata: one\n\n", b"\xef\xbb\xbfdata: ignored\n\n", b"data: \xef\xbb\xbfvalue\n\n"] + blocks = [b"\xef\xbb\xbfdata: one\n\n", b"data: \xef\xbb\xbfvalue\n\n"] raw = b"".join(blocks) for split in range(len(raw) + 1): events = [event async for event in iter_sse_events(_chunks(raw[:split], raw[split:]))] assert [event.raw for event in events] == blocks - assert [event.data for event in events] == [b"one", None, b"\xef\xbb\xbfvalue"] - assert ServerSentEvent.from_bytes(blocks[0]).data is None + assert [event.data for event in events] == [b"one", b"\xef\xbb\xbfvalue"] + with pytest.raises(AmbiguousServerSentEvent): + ServerSentEvent.from_bytes(blocks[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raw", + [ + b"data: one\n\n\xef\xbb\xbfdata: two\n\n", + b"data: one\n\xef\xbb\xbf\ndata: two\n\n", + b"\xef\xbb\xbf\xef\xbb\xbfdata: one\n\n", + b"data: one\n\n\xef\xbb\xbf: comment\n\n", + b"data: one\r\r\xef\xbb\xbfdata: two\r\r", + b"data: one\r\n\r\n\xef\xbb\xbfdata: two\r\n\r\n", + b"data: one\n\n\xef\xbb\xbfdata: two", + ], +) +async def test_line_leading_bom_after_stream_start_is_rejected_at_every_split(raw): + """Clients disagree on such lines, so the parser never interprets them.""" + for split in range(len(raw) + 1): + with pytest.raises(AmbiguousServerSentEvent): + _ = [event async for event in iter_sse_events(_chunks(raw[:split], raw[split:]))] @pytest.mark.asyncio @@ -130,7 +154,7 @@ async def test_unfinished_comment_tails_are_discarded(tail): @pytest.mark.asyncio async def test_event_limit_counts_terminators_bom_and_unfinished_lines(): block = b"\xef\xbb\xbfdata: one\r\n\r\n" - events = [event async for event in iter_sse_events(_chunks(block * 3), max_event_bytes=len(block))] + events = [event async for event in iter_sse_events(_chunks(block + block[3:] * 2), max_event_bytes=len(block))] assert len(events) == 3 for raw in (block, b"x" * len(block)): with pytest.raises(ServerSentEventTooLarge): @@ -168,3 +192,26 @@ async def test_large_event_arriving_one_byte_per_chunk(): assert len(events) == 1 assert events[0].raw == raw + + +@pytest.mark.parametrize( + "raw", + [ + b"\xef\xbb\xbfdata: one\r\n\r\ndata: two\n\n", + b"data: a\ndata: b\r\rdata: c\r\n\r\n", + b": comment\n\nid: 1\nevent: delta\ndata: x\n\nretry: 5\n\n", + b"data:no-space\n\ndata: \xef\xbb\xbfvalue\n\nunknown: field\ndata\n\n", + b"Data: wrong-case\n\ndata: right\n\n", + ], +) +def test_client_never_sees_data_this_parser_did_not_parse(raw): + """The OpenAI Python SDK's events are a subsequence of the events parsed here.""" + sse_decoder = pytest.importorskip("openai._streaming").SSEDecoder + + async def parse(): + return [event.data for event in [event async for event in iter_sse_events(_chunks(raw))]] + + parsed = [data for data in asyncio.run(parse()) if data is not None] + client = [event.data.encode() for event in sse_decoder().iter_bytes(iter([raw])) if event.data] + remaining = iter(parsed) + assert all(any(data == candidate for candidate in remaining) for data in client) From ad9645f575ff04f7c948cf788b7db47272391d02 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 18:14:46 +0200 Subject: [PATCH 06/16] fix(server): reject unpaired surrogates in guarded stream text JSON can encode an unpaired UTF-16 surrogate as an escape, and the JSON parser also accepts its CESU-8 bytes. Such text cannot be re-encoded as UTF-8, so the guarded stream failed with a generic processing error while counting text bytes. Treat guarded or snapshot text with an unpaired surrogate as an unsupported provider stream: no output check runs, pending events are withheld, and the error is rendered as a stream shape failure. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_stream.py | 11 +++++++++++ tests/server/experimental/test_guarded_stream.py | 14 ++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index f858905c34..3128613925 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -134,6 +134,15 @@ def _take_checked_events( return tuple(event.event for event in checked) +def _is_unicode_text(text: str) -> bool: + """Return whether text is valid Unicode that can be checked and re-encoded as UTF-8.""" + try: + text.encode("utf-8") + except UnicodeEncodeError: + return False + return True + + async def _text_windows( source: AsyncIterator[str | ServerSentEvent], *, @@ -242,6 +251,8 @@ async def text_deltas() -> AsyncIterator[str | ServerSentEvent]: async for event in iter_sse_events(_guard_upstream(source), max_event_bytes=max_event_bytes): classified = adapter.classify_event(event) validate_stream_event(adapter.contract, classified) + if classified.text is not None and not _is_unicode_text(classified.text): + raise UnsupportedProviderStream("Provider stream text contains an unpaired surrogate.") if not pending and event.data is None and classified.role is StreamEventRole.OPAQUE_METADATA: yield event continue diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index 2b73baf115..a30a1e18fa 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -494,3 +494,17 @@ async def test_line_leading_bom_mid_stream_fails_closed_without_release(split): assert b"EVIL" not in result assert b'"code": "unsupported_stream"' in result assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event", [stream_event({"text": "a\ud800b"}), b'data: {"text":"a\xed\xa0\x80b"}\n\n'], ids=["escape", "bytes"] +) +async def test_unpaired_surrogate_text_is_rejected_as_unsupported(event): + """Guarded text must be valid Unicode; an unpaired surrogate is a shape error, not a crash.""" + checker = StaticChecker() + result = await collect(Source([event]), checker, Hooks()) + + assert b'"code": "unsupported_stream"' in result + assert b"\xed\xa0\x80" not in result + assert checker.calls == [] From e4f45eafbfd0e79491bcf200d62953c7efc36b64 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:07:30 +0200 Subject: [PATCH 07/16] test(server): require Python SDK SSE comparison Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- tests/server/experimental/test_sse.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/server/experimental/test_sse.py b/tests/server/experimental/test_sse.py index 48e1c098b4..a4d254d205 100644 --- a/tests/server/experimental/test_sse.py +++ b/tests/server/experimental/test_sse.py @@ -16,6 +16,7 @@ import asyncio import pytest +from openai._streaming import SSEDecoder from nemoguardrails.server.experimental.provider.sse import ( AmbiguousServerSentEvent, @@ -204,14 +205,13 @@ async def test_large_event_arriving_one_byte_per_chunk(): b"Data: wrong-case\n\ndata: right\n\n", ], ) -def test_client_never_sees_data_this_parser_did_not_parse(raw): - """The OpenAI Python SDK's events are a subsequence of the events parsed here.""" - sse_decoder = pytest.importorskip("openai._streaming").SSEDecoder +def test_python_sdk_never_sees_data_this_parser_did_not_parse(raw): + """Compare Python SDK parsing, separately from the Node BOM regression.""" async def parse(): return [event.data for event in [event async for event in iter_sse_events(_chunks(raw))]] parsed = [data for data in asyncio.run(parse()) if data is not None] - client = [event.data.encode() for event in sse_decoder().iter_bytes(iter([raw])) if event.data] + client = [event.data.encode() for event in SSEDecoder().iter_bytes(iter([raw])) if event.data] remaining = iter(parsed) assert all(any(data == candidate for candidate in remaining) for data in client) From 087623a7ab73f5c5a2bea85e2651ce0a45348db5 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 12 Oct 2026 11:06:36 +0200 Subject: [PATCH 08/16] fix(server): relay lone surrogates in guarded stream text A transparent proxy should not enforce Unicode validity that the provider and client do not. JSON escapes can produce lone surrogates, for example half of a pair split across deltas, and clients receive them as sent. The stream path now counts their bytes without failing, checks the text, and relays it unchanged. A checker that cannot handle such text still fails closed as a check failure. This replaces the earlier rejection. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_stream.py | 16 +++-------- .../experimental/test_guarded_stream.py | 28 +++++++++++++++---- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index 3128613925..3b1f2548f2 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -134,15 +134,6 @@ def _take_checked_events( return tuple(event.event for event in checked) -def _is_unicode_text(text: str) -> bool: - """Return whether text is valid Unicode that can be checked and re-encoded as UTF-8.""" - try: - text.encode("utf-8") - except UnicodeEncodeError: - return False - return True - - async def _text_windows( source: AsyncIterator[str | ServerSentEvent], *, @@ -251,8 +242,6 @@ async def text_deltas() -> AsyncIterator[str | ServerSentEvent]: async for event in iter_sse_events(_guard_upstream(source), max_event_bytes=max_event_bytes): classified = adapter.classify_event(event) validate_stream_event(adapter.contract, classified) - if classified.text is not None and not _is_unicode_text(classified.text): - raise UnsupportedProviderStream("Provider stream text contains an unpaired surrogate.") if not pending and event.data is None and classified.role is StreamEventRole.OPAQUE_METADATA: yield event continue @@ -274,7 +263,10 @@ async def text_deltas() -> AsyncIterator[str | ServerSentEvent]: ) if text: if snapshot_history_complete: - observed_text_bytes += len(text.encode()) + # JSON escapes can produce lone surrogates, for example half of a + # pair split across deltas. Clients receive them as sent, so the + # proxy relays them too and only needs a byte count here. + observed_text_bytes += len(text.encode("utf-8", "surrogatepass")) if observed_text_bytes > max_event_bytes: observed_text.clear() snapshot_history_complete = False diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index a30a1e18fa..c36b6e3c21 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -498,13 +498,29 @@ async def test_line_leading_bom_mid_stream_fails_closed_without_release(split): @pytest.mark.asyncio @pytest.mark.parametrize( - "event", [stream_event({"text": "a\ud800b"}), b'data: {"text":"a\xed\xa0\x80b"}\n\n'], ids=["escape", "bytes"] + ("events", "text"), + [ + ([stream_event({"text": "a\ud800b"})], "a\ud800b"), + ([b'data: {"text":"a\xed\xa0\x80b"}\n\n'], "a\ud800b"), + ([stream_event({"text": "\ud83d"}), stream_event({"text": "\ude00"})], "\ud83d\ude00"), + ], + ids=["escape", "bytes", "split-pair"], ) -async def test_unpaired_surrogate_text_is_rejected_as_unsupported(event): - """Guarded text must be valid Unicode; an unpaired surrogate is a shape error, not a crash.""" +async def test_lone_surrogate_text_is_checked_and_relayed_unchanged(events, text): + """Lone surrogates reach clients as sent, so the proxy checks and relays them too.""" checker = StaticChecker() + result = await collect(Source(events), checker, Hooks(), policy=StreamBufferingPolicy(len(events), 0)) + + assert result == b"".join(events) + assert [call.output_content for call in checker.calls] == [text] + + +@pytest.mark.asyncio +async def test_checker_that_cannot_encode_lone_surrogates_fails_closed(): + event = stream_event({"text": "a\ud800b"}) + checker = StaticChecker([UnicodeEncodeError("utf-8", "a\ud800b", 1, 2, "surrogates not allowed")]) + result = await collect(Source([event]), checker, Hooks()) - assert b'"code": "unsupported_stream"' in result - assert b"\xed\xa0\x80" not in result - assert checker.calls == [] + assert b'"code": "check_failed"' in result + assert b"ud800" not in result From ddac88c899c2c7c8360f324e172b5fc8dc165a8c Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 12 Oct 2026 11:07:15 +0200 Subject: [PATCH 09/16] test(server): compare SSE parsing with a Node-like BOM reader The Python SDK never strips a BOM, so comparing with it cannot catch the line-leading BOM differential. A small reader modelled on openai-node 7.31.0 strips a BOM at every line start; accepted streams must expose no data to it that this parser did not parse. Disabling the BOM rejection fails these cases. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- tests/server/experimental/test_sse.py | 45 +++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/server/experimental/test_sse.py b/tests/server/experimental/test_sse.py index a4d254d205..2aea9fc015 100644 --- a/tests/server/experimental/test_sse.py +++ b/tests/server/experimental/test_sse.py @@ -215,3 +215,48 @@ async def parse(): client = [event.data.encode() for event in SSEDecoder().iter_bytes(iter([raw])) if event.data] remaining = iter(parsed) assert all(any(data == candidate for candidate in remaining) for data in client) + + +def _node_like_events(raw: bytes) -> list[bytes]: + """Model openai-node 7.31.0, which strips a BOM at the start of every line. + + The preserved review evidence shows Node reading a later BOM-prefixed data + line that a stream-start-only parser ignores. This is a model, not Node. + """ + events, data = [], [] + for line in raw.replace(b"\r\n", b"\n").replace(b"\r", b"\n").split(b"\n"): + line = line.removeprefix(b"\xef\xbb\xbf") + if not line: + if data: + events.append(b"\n".join(data)) + data = [] + continue + name, _, value = line.partition(b":") + if name == b"data": + data.append(value.removeprefix(b" ")) + return events + + +@pytest.mark.parametrize( + "raw", + [ + b"data: safe\n\n\xef\xbb\xbfdata: EVIL\n\n", + b"data: safe\r\r\xef\xbb\xbfdata: EVIL\r\r", + b"data: safe\n\xef\xbb\xbfdata: EVIL\n\n", + b": note\n\xef\xbb\xbfdata: EVIL\n\n", + b"\xef\xbb\xbfdata: one\n\ndata: two\n\n", + ], +) +@pytest.mark.parametrize("split", range(4)) +def test_node_like_client_never_sees_data_this_parser_did_not_parse(raw, split): + """Guard the Node BOM differential: accepted streams expose no extra data.""" + + async def parse(): + return [event.data async for event in iter_sse_events(_chunks(raw[:split], raw[split:]))] + + try: + parsed = [data for data in asyncio.run(parse()) if data is not None] + except AmbiguousServerSentEvent: + return + remaining = iter(parsed) + assert all(any(data == candidate for candidate in remaining) for data in _node_like_events(raw)) From 10d19d35b90f44845a3fb9d37a8a684078d9761c Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:20:12 +0200 Subject: [PATCH 10/16] fix(server): inspect the original stream text The classifier read guarded text from the validated projection, while the runner relays the original event bytes. A validator that changed, dropped or reordered the selected text would let the checker approve one string while the client received another. Text is now resolved from the raw payload before validation and compared afterwards; any difference fails closed. Null, missing and empty text keep their existing meaning. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../provider/stream_classifier.py | 12 ++++--- .../experimental/test_stream_classifier.py | 34 ++++++++++++++++++- 2 files changed, 41 insertions(+), 5 deletions(-) diff --git a/nemoguardrails/server/experimental/provider/stream_classifier.py b/nemoguardrails/server/experimental/provider/stream_classifier.py index 09939dfc9c..eb039695d8 100644 --- a/nemoguardrails/server/experimental/provider/stream_classifier.py +++ b/nemoguardrails/server/experimental/provider/stream_classifier.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Classify provider stream events from generated semantic rules.""" +"""Validate declared event shapes without changing the inspected wire text.""" from collections.abc import Sequence from dataclasses import dataclass @@ -151,8 +151,10 @@ def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: shapes = ", ".join(repr(rule.shape) for rule in matching_rules) raise UnsupportedProviderStream(f"{self.subject} event ambiguously matches {shapes}.") rule = matching_rules[0] + text = _resolve_text(payload, rule.text_path) projection = rule.model.model_validate(payload) - text = _resolve_text(projection, rule.text_path) + if _resolve_text(projection, rule.text_path) != text: + raise UnsupportedProviderStream(f"{self.subject} validation changed the original guarded text.") if text is None and rule.missing_text_role is not None: return GuardedStreamEvent( rule.missing_text_shape or rule.shape, @@ -170,13 +172,15 @@ def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: raise UnsupportedProviderStream(str(error)) from error -def _resolve_text(projection: BaseModel, path: StreamTextPath | None) -> str | None: +def _resolve_text(projection: BaseModel | dict[str, Any], path: StreamTextPath | None) -> str | None: if path is None: return None value: object = projection for segment in path: if isinstance(segment, str) and isinstance(value, BaseModel): value = getattr(value, segment) + elif isinstance(segment, str) and isinstance(value, dict): + value = value.get(segment) elif isinstance(segment, int) and isinstance(value, list): if segment < 0: raise ValueError("Stream event text paths cannot use negative indexes.") @@ -195,7 +199,7 @@ def _resolve_text(projection: BaseModel, path: StreamTextPath | None) -> str | N def build_stream_classifier(definition: StreamClassifierDefinition) -> StreamClassifier: - """Validate generated event rules and return their runtime classifier.""" + """Validate declared event rules and return their runtime classifier.""" if not definition.subject: raise ValueError("A stream classifier definition must identify its provider event subject.") diff --git a/tests/server/experimental/test_stream_classifier.py b/tests/server/experimental/test_stream_classifier.py index c91ebb8b27..cec23eeaea 100644 --- a/tests/server/experimental/test_stream_classifier.py +++ b/tests/server/experimental/test_stream_classifier.py @@ -13,10 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +from dataclasses import replace from typing import Literal import pytest -from pydantic import BaseModel +from pydantic import BaseModel, model_validator from nemoguardrails.server.experimental.provider.sse import ServerSentEvent from nemoguardrails.server.experimental.provider.stream import ( @@ -208,3 +210,33 @@ def test_stream_discriminators_require_matching_types(expected, other): assert selector.select([incorrect, correct]) is correct assert rule.matches(correct) is True assert rule.matches(incorrect) is False + + +@pytest.mark.parametrize("mutation", ["text", "drop", "reorder"]) +def test_classifier_rejects_validation_that_changes_the_original_subject(mutation): + class TransformingChunk(Chunk): + @model_validator(mode="before") + @classmethod + def transform(cls, value): + if mutation == "text": + value["deltas"][0]["text"] = "safe" + elif mutation == "drop": + value["deltas"] = [] + else: + value["deltas"].reverse() + return value + + original = classifier() + modified = replace(original, rules=(replace(original.rules[1], model=TransformingChunk),)) + event = ServerSentEvent.from_bytes(b'data: {"kind":"chunk","deltas":[{"text":"unsafe"},{"text":"safe"}]}\n\n') + + with pytest.raises(UnsupportedProviderStream, match="original guarded text"): + modified.classify_event(event) + + +@pytest.mark.parametrize("value", [None, "", "hello"]) +def test_classifier_preserves_nullable_and_empty_original_text(value): + event = ServerSentEvent.from_bytes( + b"data: " + json.dumps({"kind": "chunk", "deltas": [{"text": value}]}).encode() + b"\n\n" + ) + assert classifier().classify_event(event).text == value From 164885b9108ee9eb7b69de0033be4f9cdccfcc63 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:20:42 +0200 Subject: [PATCH 11/16] fix(server): match stream discriminators by type and presence Rule construction compared discriminators with Python equality, so rules for true and 1 were rejected as duplicates although runtime matching treats them as different values. Duplicate detection now uses the same exact type and value as matching. Array selection also stops treating a missing member as an explicit null. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../experimental/provider/stream_classifier.py | 6 +++--- .../experimental/test_stream_classifier.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/nemoguardrails/server/experimental/provider/stream_classifier.py b/nemoguardrails/server/experimental/provider/stream_classifier.py index eb039695d8..9cc54d75d5 100644 --- a/nemoguardrails/server/experimental/provider/stream_classifier.py +++ b/nemoguardrails/server/experimental/provider/stream_classifier.py @@ -47,7 +47,7 @@ def select(self, items: Sequence[object]) -> object | None: item for item in items if all( - type(value := (item.get(name) if isinstance(item, dict) else getattr(item, name, object()))) + type(value := (item.get(name, object()) if isinstance(item, dict) else getattr(item, name, object()))) is type(expected) and value == expected for name, expected in self.match @@ -209,7 +209,7 @@ def build_stream_classifier(definition: StreamClassifierDefinition) -> StreamCla raise ValueError("A stream classifier subject must match its projection contract identifier.") if definition.require_event_type and definition.event_type_field is None: raise ValueError("A required SSE event field needs a data discriminator.") - signatures: set[tuple[tuple[tuple[str, Any], ...], frozenset[str]]] = set() + signatures: set[tuple[frozenset[tuple[str, type, StreamMatchValue]], frozenset[str]]] = set() def validate_declared_event(shape: str, role: StreamEventRole, text: str | None = None) -> None: try: @@ -226,7 +226,7 @@ def validate_declared_event(shape: str, role: StreamEventRole, text: str | None raise ValueError("A stream event rule must identify its provider shape.") if not rule.match and not rule.required_fields: raise ValueError(f"Stream event rule {rule.shape!r} must declare how it is selected.") - signature = (rule.match, rule.required_fields) + signature = (frozenset((name, type(value), value) for name, value in rule.match), rule.required_fields) if signature in signatures: raise ValueError(f"Stream event rule {rule.shape!r} duplicates another selection.") signatures.add(signature) diff --git a/tests/server/experimental/test_stream_classifier.py b/tests/server/experimental/test_stream_classifier.py index cec23eeaea..3d6b9f6bf7 100644 --- a/tests/server/experimental/test_stream_classifier.py +++ b/tests/server/experimental/test_stream_classifier.py @@ -240,3 +240,21 @@ def test_classifier_preserves_nullable_and_empty_original_text(value): b"data: " + json.dumps({"kind": "chunk", "deltas": [{"text": value}]}).encode() + b"\n\n" ) assert classifier().classify_event(event).text == value + + +def test_classifier_distinguishes_boolean_integer_and_float_selections(): + rules = tuple( + StreamEventRule( + shape="chunk:metadata", model=BaseModel, role=StreamEventRole.OPAQUE_METADATA, match=(("kind", value),) + ) + for value in (True, 1, 1.0) + ) + result = build_stream_classifier(StreamClassifierDefinition(subject="test.stream", contract=CONTRACT, rules=rules)) + for value in (True, 1, 1.0): + assert sum(rule.matches({"kind": value}) for rule in result.rules) == 1 + + +def test_array_selection_distinguishes_missing_from_null(): + selector = StreamArraySelection(match=(("kind", None),), cardinality="exactly_one") + explicit_null = {"kind": None} + assert selector.select([{}, explicit_null]) is explicit_null From 6fa1158005bd2e2a0925cf697d606525a401f2e4 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 12 Oct 2026 11:42:33 +0200 Subject: [PATCH 12/16] fix(server): bound stream cleanup and preserve primary outcomes The guarded stream closed its source without a time limit, so a provider body whose close stalled could hang an otherwise complete response. Iterator and resource cleanup now share one helper: each attempt is shielded and bounded, and timeouts or ordinary close failures are logged instead of replacing the stream's outcome. Close failures were previously raised; this is deliberate, because cleanup must not mask a block, an upstream failure or a cancellation. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_stream.py | 24 ++++----- .../server/experimental/_stream_cleanup.py | 50 +++++++++++++++++++ .../experimental/test_guarded_stream.py | 39 +++++++++++++++ 3 files changed, 99 insertions(+), 14 deletions(-) create mode 100644 nemoguardrails/server/experimental/_stream_cleanup.py diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index 3b1f2548f2..39e63261cd 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -15,7 +15,6 @@ """Apply buffered output checks to provider-native event streams.""" -import inspect from collections.abc import AsyncIterator, Callable from dataclasses import dataclass @@ -38,6 +37,7 @@ ) from nemoguardrails.server.experimental._guarded_operation import UnsupportedGuardedPayload from nemoguardrails.server.experimental._http_kernel import BufferedHttpResponse +from nemoguardrails.server.experimental._stream_cleanup import StreamClose, close_stream from nemoguardrails.server.experimental.provider.sse import ( AmbiguousServerSentEvent, ServerSentEvent, @@ -157,14 +157,6 @@ async def _text_windows( yield _TextWindow(tuple(context + pending), len(pending)) -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 _guard_upstream(source: AsyncIterator[bytes]) -> AsyncIterator[bytes]: try: async for chunk in source: @@ -195,6 +187,7 @@ async def guard_provider_stream( render_outcome: StreamOutcomeRenderer, max_event_bytes: int = 1024 * 1024, max_pending_bytes: int = 10 * 1024 * 1024, + close: StreamClose | None = None, ) -> AsyncIterator[bytes]: """Release provider SSE events only after required output checks pass. @@ -209,9 +202,12 @@ async def guard_provider_stream( checker semantics, and buffered and streaming execution tests. Adding a provider or changing a profile declaration alone does not relax it. - Iteration closes ``source`` when it ends, including when arguments are - rejected. A stream that is never iterated does not run that cleanup, so its - caller must close ``source``. Validated data-less keepalive prefixes are + Iteration attempts bounded, shielded cleanup of ``source`` and the optional + resource-owner ``close`` callback, including when arguments are rejected. + Iterator termination alone cannot recover interrupted generator cleanup; + callers owning external resources must supply their independent closer. + A stream never iterated does not run cleanup, so its caller retains ownership. + Validated data-less keepalive prefixes are released promptly; they never overtake pending payloads or terminal events. """ @@ -221,7 +217,7 @@ async def guard_provider_stream( if any(type(limit) is not int or limit <= 0 for limit in (max_event_bytes, max_pending_bytes)): raise ValueError("Stream byte limits must be positive.") except ValueError: - await _close_source(source) + await close_stream(source, close=close) raise try: @@ -340,4 +336,4 @@ async def text_deltas() -> AsyncIterator[str | ServerSentEvent]: ): yield chunk finally: - await _close_source(source) + await close_stream(source, close=close) diff --git a/nemoguardrails/server/experimental/_stream_cleanup.py b/nemoguardrails/server/experimental/_stream_cleanup.py new file mode 100644 index 0000000000..c216752fb9 --- /dev/null +++ b/nemoguardrails/server/experimental/_stream_cleanup.py @@ -0,0 +1,50 @@ +# 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. + +"""Best-effort, bounded cleanup independent of stream iteration state.""" + +import inspect +import logging +from collections.abc import Awaitable, Callable + +import anyio + +log = logging.getLogger(__name__) + +_STREAM_CLOSE_TIMEOUT_SECONDS = 5.0 +StreamClose = Callable[[], Awaitable[None]] + + +async def close_stream(source: object, *, close: StreamClose | None = None) -> None: + """Close an iterator and, independently, its resource owner. + + Each attempt has a shielded five-second budget. Timeouts and ordinary close + failures are logged without replacing the stream's outcome. Resource owners + must make closing idempotent and cancellation-cooperative. Iterator cleanup + alone cannot recover resources from an interrupted async-generator finally. + This does not forcibly stop blocking code or cancellation-suppressing code. + """ + for closer in (getattr(source, "aclose", None), close): + if not callable(closer): + continue + try: + with anyio.move_on_after(_STREAM_CLOSE_TIMEOUT_SECONDS, shield=True) as scope: + result = closer() + if inspect.isawaitable(result): + await result + if scope.cancelled_caught: + log.warning("Stream cleanup timed out; resources may remain open.") + except Exception: + log.warning("Stream cleanup failed; resources may remain open.") diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index c36b6e3c21..0513138cdb 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -16,8 +16,10 @@ import asyncio import json +import anyio import pytest +from nemoguardrails.server.experimental import _stream_cleanup from nemoguardrails.server.experimental._buffered_kernel import ( OperationBlocked, OperationCheckFailed, @@ -524,3 +526,40 @@ async def test_checker_that_cannot_encode_lone_surrogates_fails_closed(): assert b'"code": "check_failed"' in result assert b"ud800" not in result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["stall", "raise"]) +@pytest.mark.parametrize("decision", [ContentAllowed(), ContentBlocked("blocked")]) +async def test_cleanup_is_bounded_without_replacing_the_stream_outcome(monkeypatch, caplog, failure, decision): + monkeypatch.setattr(_stream_cleanup, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01) + + class FailingClose(Source): + async def aclose(self): + if failure == "stall": + await anyio.sleep_forever() + raise RuntimeError("private upstream detail") + + event = stream_event({"text": "answer"}) + with anyio.fail_after(1): + result = await collect(FailingClose([event]), StaticChecker([decision]), Hooks()) + + assert result == event if isinstance(decision, ContentAllowed) else b'"code": "blocked"' in result + assert "Stream cleanup" in caplog.text + assert "private upstream detail" not in caplog.text + + +@pytest.mark.asyncio +async def test_explicit_resource_owner_closes_after_iterator_cleanup_fails(): + closed = [] + + class FailingClose(Source): + async def aclose(self): + raise RuntimeError("iterator cleanup failed") + + async def close_owner(): + await anyio.sleep(0) + closed.append(True) + + await collect(FailingClose([]), StaticChecker(), Hooks(), close=close_owner) + assert closed == [True] From f86439c38ff0022c530e334d226873f3b49c68da Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 12 Oct 2026 11:58:19 +0200 Subject: [PATCH 13/16] fix(server): attempt resource cleanup even when iterator close is cancelled If closing the iterator raised cancellation, the separate resource closer was skipped and the upstream connection could stay open. The resource closer now runs in a finally block, so it is attempted however iterator cleanup ends. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_stream_cleanup.py | 11 +++++++++-- .../server/experimental/test_guarded_stream.py | 17 +++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/nemoguardrails/server/experimental/_stream_cleanup.py b/nemoguardrails/server/experimental/_stream_cleanup.py index c216752fb9..5785cc3aeb 100644 --- a/nemoguardrails/server/experimental/_stream_cleanup.py +++ b/nemoguardrails/server/experimental/_stream_cleanup.py @@ -36,9 +36,11 @@ async def close_stream(source: object, *, close: StreamClose | None = None) -> N alone cannot recover resources from an interrupted async-generator finally. This does not forcibly stop blocking code or cancellation-suppressing code. """ - for closer in (getattr(source, "aclose", None), close): + + async def attempt(closer: Callable[[], object] | None) -> None: + """Bound one closer without swallowing cancellation from its caller.""" if not callable(closer): - continue + return try: with anyio.move_on_after(_STREAM_CLOSE_TIMEOUT_SECONDS, shield=True) as scope: result = closer() @@ -48,3 +50,8 @@ async def close_stream(source: object, *, close: StreamClose | None = None) -> N log.warning("Stream cleanup timed out; resources may remain open.") except Exception: log.warning("Stream cleanup failed; resources may remain open.") + + try: + await attempt(getattr(source, "aclose", None)) + finally: + await attempt(close) diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index 0513138cdb..84c6a29f00 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -563,3 +563,20 @@ async def close_owner(): await collect(FailingClose([]), StaticChecker(), Hooks(), close=close_owner) assert closed == [True] + + +@pytest.mark.asyncio +async def test_iterator_close_cancellation_does_not_skip_the_resource_owner(): + closed = [] + + class CancelledClose(Source): + async def aclose(self): + raise asyncio.CancelledError + + async def close_owner(): + await anyio.sleep(0) + closed.append(True) + + with pytest.raises(asyncio.CancelledError): + await collect(CancelledClose([]), StaticChecker(), Hooks(), close=close_owner) + assert closed == [True] From 3a240dd00fafcc88f7278d2bddda8ede054f9b19 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:38:56 +0200 Subject: [PATCH 14/16] fix(server): log why guarded streams stop Stream failures reached the client as a native error event but left no record for operators. Each failure outcome is now logged once, where it is encoded: projection and framing rejections and unsupported modification requests as warnings with their stage and reason, upstream failures as warnings with the exception type, and checker or processing failures as errors with the exception attached. Blocks are ordinary checker decisions and are not logged here. Reasons describe shapes and never echo provider text. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_stream.py | 34 +++++++++++ .../experimental/test_guarded_stream.py | 61 +++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index 39e63261cd..13da1a447e 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -15,6 +15,7 @@ """Apply buffered output checks to provider-native event streams.""" +import logging from collections.abc import AsyncIterator, Callable from dataclasses import dataclass @@ -53,6 +54,8 @@ ) from nemoguardrails.server.experimental.provider.types import GuardedMessage +log = logging.getLogger(__name__) + class UnsupportedStreamInspection(ValueError): """Report an output policy that cannot safely release a guarded stream.""" @@ -167,12 +170,43 @@ async def _guard_upstream(source: AsyncIterator[bytes]) -> AsyncIterator[bytes]: raise _UpstreamStreamFailure(failure) from failure +def _log_outcome(outcome: StreamOutcome) -> None: + """Record why a guarded stream stopped, without logging provider text. + + Blocks are normal checker decisions and are not logged here. Failure + reasons come from projection and framing checks, which describe shapes + rather than echoing provider content. + """ + if isinstance(outcome, OperationBlocked): + return + if isinstance(outcome, OperationProjectionFailed): + log.warning("Guarded stream rejected at %s stage: %s", outcome.stage.value, outcome.failure) + elif isinstance(outcome, OperationModificationUnsupported): + log.warning( + "Guarded stream stopped at %s stage: the checker requested unsupported modification.", outcome.stage.value + ) + elif isinstance(outcome, OperationCheckFailed): + cause = outcome.failure.cause + log.error( + "Guarded stream check failed at %s stage: %s (%s)", + outcome.stage.value, + outcome.failure.message, + type(cause).__name__ if cause is not None else "no exception", + exc_info=cause, + ) + elif isinstance(outcome, StreamUpstreamFailed): + log.warning("Guarded stream upstream failed: %s", type(outcome.failure).__name__) + else: + log.error("Guarded stream processing failed: %s", type(outcome.failure).__name__, exc_info=outcome.failure) + + def _encoded_outcome( outcome: StreamOutcome, *, adapter: ProviderStreamAdapter, render_outcome: StreamOutcomeRenderer, ) -> tuple[bytes, ...]: + _log_outcome(outcome) response = render_outcome(outcome) return tuple(adapter.encode_error(response.body)) diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index 84c6a29f00..c87bca4505 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -580,3 +580,64 @@ async def close_owner(): with pytest.raises(asyncio.CancelledError): await collect(CancelledClose([]), StaticChecker(), Hooks(), close=close_owner) assert closed == [True] + + +class FailingEndHooks(Hooks): + def validate_end_of_stream(self): + raise RuntimeError("hook bug") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("source", "checker", "hooks", "level", "expected"), + [ + ( + lambda: Source([stream_event({"text": "secret"}) + b"\xef\xbb\xbfdata: x\n\n"]), + lambda: StaticChecker(), + Hooks, + "WARNING", + "rejected at output stage: A provider stream line starts with a byte order mark", + ), + ( + lambda: Source([stream_event({"text": "secret"})]), + lambda: StaticChecker([RuntimeError("checker down")]), + Hooks, + "ERROR", + "check failed at output stage: The output content check failed. (RuntimeError)", + ), + ( + lambda: Source([stream_event({"text": "secret"})], failure=ConnectionError("reset")), + lambda: StaticChecker(), + Hooks, + "WARNING", + "upstream failed: ConnectionError", + ), + ( + lambda: Source([stream_event({"text": "secret"})]), + lambda: StaticChecker(), + FailingEndHooks, + "ERROR", + "processing failed: RuntimeError", + ), + ], + ids=["projection", "check", "upstream", "processing"], +) +async def test_stream_failures_are_logged_without_provider_text(caplog, source, checker, hooks, level, expected): + """Operators see why a stream stopped; the logged reason never echoes guarded text.""" + caplog.set_level("WARNING", logger="nemoguardrails.server.experimental._guarded_stream") + + await collect(source(), checker(), hooks(), policy=StreamBufferingPolicy(2, 0)) + + records = [record for record in caplog.records if record.name.endswith("_guarded_stream")] + assert [record.levelname for record in records] == [level] + assert expected in records[0].getMessage() + assert "secret" not in caplog.text + + +@pytest.mark.asyncio +async def test_blocked_stream_is_not_logged_as_a_failure(caplog): + caplog.set_level("DEBUG", logger="nemoguardrails.server.experimental._guarded_stream") + + await collect(Source([stream_event({"text": "secret"})]), StaticChecker([ContentBlocked("blocked")]), Hooks()) + + assert not [record for record in caplog.records if record.name.endswith("_guarded_stream")] From 3d3b86c5d89cf1128a79c5a0750bc5e7e1ae1f2b Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:42:12 +0200 Subject: [PATCH 15/16] test(server): cover stream features Chat does not use yet The OpenAI Chat binding exercises only part of the shared stream kernel. Synthetic rules now cover data-less events, the SSE event-type discriminator, exactly-one array selection, dotted match keys, missing text handling, text paths that do not reach a string, snapshot history beyond the verification limit, non-byte provider chunks, and every classifier definition and contract consistency check. No runtime behavior changes. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../experimental/test_content_checker.py | 5 + .../experimental/test_guarded_stream.py | 20 +++ tests/server/experimental/test_sse.py | 6 + .../experimental/test_stream_classifier.py | 151 ++++++++++++++++++ 4 files changed, 182 insertions(+) diff --git a/tests/server/experimental/test_content_checker.py b/tests/server/experimental/test_content_checker.py index 339df5ec41..f52713d515 100644 --- a/tests/server/experimental/test_content_checker.py +++ b/tests/server/experimental/test_content_checker.py @@ -182,3 +182,8 @@ def test_stream_buffering_policy_rejects_invalid_values(buffering): def test_stream_buffering_requires_output_inspection(): with pytest.raises(ValueError, match="requires output inspection"): ContentInspectionPolicy(True, False, StreamBufferingPolicy(1, 0)) + + +def test_stream_buffering_must_be_a_buffering_policy(): + with pytest.raises(TypeError, match="StreamBufferingPolicy"): + ContentInspectionPolicy(True, True, {"chunk_size": 1}) diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index c87bca4505..a21108bd87 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -641,3 +641,23 @@ async def test_blocked_stream_is_not_logged_as_a_failure(caplog): await collect(Source([stream_event({"text": "secret"})]), StaticChecker([ContentBlocked("blocked")]), Hooks()) assert not [record for record in caplog.records if record.name.endswith("_guarded_stream")] + + +@pytest.mark.asyncio +async def test_non_byte_provider_chunks_are_upstream_failures(): + result = await collect(Source(["data: text\n\n"]), StaticChecker(), Hooks()) + + assert b'"code": "upstream_failed"' in result + + +@pytest.mark.asyncio +async def test_snapshot_after_unverifiable_history_fails_closed(): + """Once delta history exceeds the limit, a later snapshot cannot be verified and is rejected.""" + deltas = [stream_event({"text": "x" * 10}) for _ in range(5)] + checker = StaticChecker() + + result = await collect(Source([*deltas, stream_event({"snapshot": "x"})]), checker, Hooks(), max_event_bytes=40) + + assert result.startswith(b"".join(deltas)) + assert b'"code": "unsupported_stream"' in result + assert b'"snapshot"' not in result diff --git a/tests/server/experimental/test_sse.py b/tests/server/experimental/test_sse.py index 2aea9fc015..70277bccd2 100644 --- a/tests/server/experimental/test_sse.py +++ b/tests/server/experimental/test_sse.py @@ -260,3 +260,9 @@ async def parse(): return remaining = iter(parsed) assert all(any(data == candidate for candidate in remaining) for data in _node_like_events(raw)) + + +@pytest.mark.parametrize("limit", [0, -1, True, 1.5]) +def test_event_limit_must_be_a_positive_integer(limit): + with pytest.raises(ValueError, match="max_event_bytes must be positive"): + asyncio.run(anext(iter_sse_events(_chunks(b"data: x\n\n"), max_event_bytes=limit))) diff --git a/tests/server/experimental/test_stream_classifier.py b/tests/server/experimental/test_stream_classifier.py index 3d6b9f6bf7..aa82dfa7c5 100644 --- a/tests/server/experimental/test_stream_classifier.py +++ b/tests/server/experimental/test_stream_classifier.py @@ -258,3 +258,154 @@ def test_array_selection_distinguishes_missing_from_null(): selector = StreamArraySelection(match=(("kind", None),), cardinality="exactly_one") explicit_null = {"kind": None} assert selector.select([{}, explicit_null]) is explicit_null + + +def chunk_rule(**changes): + rule = StreamEventRule( + shape="chunk:text", + model=Chunk, + role=StreamEventRole.GUARDED_TEXT, + match=(("kind", "chunk"),), + text_path=("deltas", 0, "text"), + ) + return replace(rule, **changes) + + +def definition(**changes): + return replace( + StreamClassifierDefinition(subject="test.stream", contract=CONTRACT, rules=(chunk_rule(),)), **changes + ) + + +def event(raw): + return ServerSentEvent.from_bytes(raw) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"subject": ""}, "identify its provider event subject"), + ({"rules": ()}, "at least one event rule"), + ({"subject": "other.stream"}, "match its projection contract identifier"), + ({"require_event_type": True}, "needs a data discriminator"), + ({"non_data_shape": "chunk:text"}, "contract declares"), + ({"rules": (chunk_rule(shape=""),)}, "identify its provider shape"), + ({"rules": (chunk_rule(match=()),)}, "must declare how it is selected"), + ({"rules": (chunk_rule(), chunk_rule())}, "duplicates another selection"), + ({"rules": (chunk_rule(text_path=None),)}, "inconsistent text role and path"), + ( + { + "rules": ( + chunk_rule( + shape="chunk:metadata", + role=StreamEventRole.OPAQUE_METADATA, + text_path=None, + missing_text_role=StreamEventRole.OPAQUE_METADATA, + ), + ) + }, + "without a text path", + ), + ({"rules": (chunk_rule(missing_text_shape="chunk:metadata"),)}, "rename missing text without classifying"), + ({"rules": (chunk_rule(missing_text_value=""),)}, "supply missing text without classifying"), + ({"rules": (chunk_rule(missing_text_role=StreamEventRole.GUARDED_TEXT),)}, "classify missing text as content"), + ], +) +def test_classifier_definition_rejects_inconsistent_declarations(changes, message): + with pytest.raises(ValueError, match=message): + build_stream_classifier(definition(**changes)) + + +def test_data_less_events_need_a_declared_shape(): + data_less = event(b"event: ping\n\n") + with pytest.raises(UnsupportedProviderStream, match="must contain data"): + build_stream_classifier(definition()).classify_event(data_less) + + classified = build_stream_classifier(definition(non_data_shape="done")).classify_event(data_less) + assert classified == GuardedStreamEvent("done", StreamEventRole.OPAQUE_METADATA) + + +def test_event_type_field_must_agree_with_the_sse_event_name(): + typed = build_stream_classifier(definition(event_type_field="kind", require_event_type=True)) + chunk = b'data: {"kind":"chunk","deltas":[{"text":"hi"}]}\n\n' + + assert typed.classify_event(event(b"event: chunk\n" + chunk)).text == "hi" + with pytest.raises(UnsupportedProviderStream, match="SSE event field is required"): + typed.classify_event(event(chunk)) + with pytest.raises(UnsupportedProviderStream, match="must match data field 'kind'"): + typed.classify_event(event(b"event: other\n" + chunk)) + with pytest.raises(UnsupportedProviderStream, match="must declare string field 'kind'"): + typed.classify_event(event(b'event: chunk\ndata: {"kind":1,"deltas":[]}\n\n')) + + +@pytest.mark.parametrize( + ("raw", "message"), + [ + (b'data: {"kind":"chunk","deltas":[{"text":null}]}\n\n', "has no guarded text value"), + (b'data: {"kind":"chunk","deltas":[{"text":"hi"},5]}\n\n', "test.stream"), + (b"data: {not json}\n\n", "JSON"), + ], + ids=["missing-text", "schema", "json"], +) +def test_classifier_rejects_unusable_event_payloads(raw, message): + with pytest.raises(UnsupportedProviderStream, match=message): + build_stream_classifier(definition()).classify_event(event(raw)) + + +@pytest.mark.parametrize( + ("text_path", "message"), + [ + (("deltas", -1, "text"), "negative indexes"), + (("deltas", "text"), "does not match its projection"), + (("deltas",), "must resolve to a string"), + ], +) +def test_classifier_rejects_text_paths_that_do_not_reach_a_string(text_path, message): + classifier = build_stream_classifier(definition(rules=(chunk_rule(text_path=text_path),))) + with pytest.raises(UnsupportedProviderStream, match=message): + classifier.classify_event(event(b'data: {"kind":"chunk","deltas":[{"text":"hi"}]}\n\n')) + + +def test_exactly_one_array_selection_rejects_missing_or_repeated_items(): + selection = StreamArraySelection(match=(("kind", "message"),), cardinality="exactly_one") + with pytest.raises(ValueError, match="exactly one matching item"): + selection.select([{"kind": "other"}]) + with pytest.raises(ValueError, match="exactly one matching item"): + selection.select([{"kind": "message"}, {"kind": "message"}]) + + +def test_dotted_match_keys_do_not_match_through_scalars(): + rule = chunk_rule(match=(("kind.name", "chunk"),)) + assert rule.matches({"kind": {"name": "chunk"}}) is True + assert rule.matches({"kind": "chunk"}) is False + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"projection_id": " "}, "must have an identifier"), + ({"shapes": StreamShapeCoverage()}, "at least one shape"), + ], +) +def test_stream_projection_contract_requires_identity_and_shapes(kwargs, message): + values = {"projection_id": "p", "profile": StreamCapabilityProfile.SINGLE_TEXT_DELTA_V1, "shapes": CONTRACT.shapes} + with pytest.raises(ValueError, match=message): + StreamProjectionContract(**{**values, **kwargs}) + + +@pytest.mark.parametrize( + ("args", "message"), + [ + (("", StreamEventRole.OPAQUE_METADATA), "identify its provider shape"), + (("chunk:text", StreamEventRole.GUARDED_TEXT), "must carry text"), + (("done", StreamEventRole.OPAQUE_METADATA, "text"), "must carry text"), + ], +) +def test_guarded_stream_event_requires_shape_and_role_consistent_text(args, message): + with pytest.raises(ValueError, match=message): + GuardedStreamEvent(*args) + + +def test_stream_contract_rejects_undeclared_shapes(): + with pytest.raises(UnsupportedProviderStream, match="is not declared"): + validate_stream_event(CONTRACT, GuardedStreamEvent("future", StreamEventRole.OPAQUE_METADATA)) From 956753a410d62c782fc811ddf3b8634d43533ffa Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:44:16 +0200 Subject: [PATCH 16/16] fix(server): log stream failures by type only Failure messages and tracebacks can carry provider or checked text, such as a duplicate JSON member name or a checker error. Log only the outcome type, stage and exception type. Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> --- .../server/experimental/_guarded_stream.py | 32 ++++---------- .../experimental/test_guarded_stream.py | 43 ++++++++++--------- 2 files changed, 31 insertions(+), 44 deletions(-) diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py index 13da1a447e..bfd717d1b0 100644 --- a/nemoguardrails/server/experimental/_guarded_stream.py +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -171,33 +171,19 @@ async def _guard_upstream(source: AsyncIterator[bytes]) -> AsyncIterator[bytes]: def _log_outcome(outcome: StreamOutcome) -> None: - """Record why a guarded stream stopped, without logging provider text. + """Record why a guarded stream stopped, by type and stage only. - Blocks are normal checker decisions and are not logged here. Failure - reasons come from projection and framing checks, which describe shapes - rather than echoing provider content. + Failure messages and tracebacks can carry provider or checked text, so + they are never logged. Blocks are normal checker decisions and are not + logged here. """ if isinstance(outcome, OperationBlocked): return - if isinstance(outcome, OperationProjectionFailed): - log.warning("Guarded stream rejected at %s stage: %s", outcome.stage.value, outcome.failure) - elif isinstance(outcome, OperationModificationUnsupported): - log.warning( - "Guarded stream stopped at %s stage: the checker requested unsupported modification.", outcome.stage.value - ) - elif isinstance(outcome, OperationCheckFailed): - cause = outcome.failure.cause - log.error( - "Guarded stream check failed at %s stage: %s (%s)", - outcome.stage.value, - outcome.failure.message, - type(cause).__name__ if cause is not None else "no exception", - exc_info=cause, - ) - elif isinstance(outcome, StreamUpstreamFailed): - log.warning("Guarded stream upstream failed: %s", type(outcome.failure).__name__) - else: - log.error("Guarded stream processing failed: %s", type(outcome.failure).__name__, exc_info=outcome.failure) + failure = outcome.failure + if isinstance(failure, ContentCheckFailed): + failure = failure.cause + stage = "stream" if isinstance(outcome, StreamUpstreamFailed | StreamProcessingFailed) else outcome.stage.value + log.warning("Guarded stream stopped: %s at %s stage (%s)", type(outcome).__name__, stage, type(failure).__name__) def _encoded_outcome( diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py index a21108bd87..2bdb787782 100644 --- a/tests/server/experimental/test_guarded_stream.py +++ b/tests/server/experimental/test_guarded_stream.py @@ -582,55 +582,56 @@ async def close_owner(): assert closed == [True] +class SecretRejectingHooks(Hooks): + def observe_event(self, event, classified): + raise UnsupportedProviderStream("Duplicate JSON member 'secret' is not supported.") + + class FailingEndHooks(Hooks): def validate_end_of_stream(self): - raise RuntimeError("hook bug") + raise RuntimeError("secret") @pytest.mark.asyncio @pytest.mark.parametrize( - ("source", "checker", "hooks", "level", "expected"), + ("source", "checker", "hooks", "expected"), [ ( - lambda: Source([stream_event({"text": "secret"}) + b"\xef\xbb\xbfdata: x\n\n"]), + lambda: Source([stream_event({"text": "hi"})]), lambda: StaticChecker(), - Hooks, - "WARNING", - "rejected at output stage: A provider stream line starts with a byte order mark", + SecretRejectingHooks, + "OperationProjectionFailed at output stage (UnsupportedGuardedPayload)", ), ( - lambda: Source([stream_event({"text": "secret"})]), - lambda: StaticChecker([RuntimeError("checker down")]), + lambda: Source([stream_event({"text": "hi"})]), + lambda: StaticChecker([RuntimeError("secret")]), Hooks, - "ERROR", - "check failed at output stage: The output content check failed. (RuntimeError)", + "OperationCheckFailed at output stage (RuntimeError)", ), ( - lambda: Source([stream_event({"text": "secret"})], failure=ConnectionError("reset")), + lambda: Source([stream_event({"text": "hi"})], failure=ConnectionError("secret")), lambda: StaticChecker(), Hooks, - "WARNING", - "upstream failed: ConnectionError", + "StreamUpstreamFailed at stream stage (ConnectionError)", ), ( - lambda: Source([stream_event({"text": "secret"})]), + lambda: Source([stream_event({"text": "hi"})]), lambda: StaticChecker(), FailingEndHooks, - "ERROR", - "processing failed: RuntimeError", + "StreamProcessingFailed at stream stage (RuntimeError)", ), ], ids=["projection", "check", "upstream", "processing"], ) -async def test_stream_failures_are_logged_without_provider_text(caplog, source, checker, hooks, level, expected): - """Operators see why a stream stopped; the logged reason never echoes guarded text.""" - caplog.set_level("WARNING", logger="nemoguardrails.server.experimental._guarded_stream") +async def test_stream_failures_are_logged_by_type_only(caplog, source, checker, hooks, expected): + """Failure messages and tracebacks can carry private text, so only types reach the log.""" + caplog.set_level("DEBUG", logger="nemoguardrails.server.experimental._guarded_stream") await collect(source(), checker(), hooks(), policy=StreamBufferingPolicy(2, 0)) records = [record for record in caplog.records if record.name.endswith("_guarded_stream")] - assert [record.levelname for record in records] == [level] - assert expected in records[0].getMessage() + assert [record.getMessage() for record in records] == [f"Guarded stream stopped: {expected}"] + assert records[0].exc_info is None assert "secret" not in caplog.text