litellm/tests/integration/observability/test_guardrail_stream_scope.py
devin-ai-integration[bot] 85a3869dfd
feat(guardrails): per-mode stream_scope with bedrock stream and pass-through fixes (#43801)
* feat(guardrails): run each mode only on streaming, non-streaming, or both

Add stream_scope so a rail can target streaming inference, non-streaming inference, or both per pre, during, and post mode

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(guardrails): honor stream_scope in pipelines and dashboard types

Pipeline steps skipped the stream_scope filter, direct construction ignored mixed-case maps, and schema.d.ts was stale.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(guardrails): skip unmatched stream_scope steps instead of allowing

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(ui): keep stored stream_scope keys for modes not on screen

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(ci): format stream_scope helpers and update fork MCP unit tests

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(mcp): keep hang cancellation tests from timing out during setup

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(guardrails): honor path-defined streaming for stream_scope

Passthrough routes like Gemini streamGenerateContent decide streaming from the URL, so stamp that onto hook data before guardrails run.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(rust): copy AnthropicModelCapabilities instead of cloning

Clippy treats clone-on-Copy as an error, which failed rust-lint on the messages request tests.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(guardrails): trust only server stream classification

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(guardrails): keep streaming marker through deepcopy

scan_raw_request snapshots copy each field, so a plain object() marker would lose identity and skip streaming-only rails.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(cost-map): drop duplicate perceptron-mk1.5 row

Two main cost-map PRs both added the OpenRouter model, so the merge left a second key that CI rejects.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(rust): expect native transcription 429 as RustUpstreamError

HTTP status errors from native routes map through route_error_to_pyerr, so the wheel SIGINT child was dying on an outdated RuntimeError check and never reached the hang probe.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(tests): follow Google Interactions OpenAPI without hardcoded names

The live spec dropped CreateModelInteractionParams and renamed the item path to {interactionsId}. Misc CI failed because the compliance tests still looked those names up as literals.

* test(guardrails): reproduce stream_scope bedrock and passthrough field gaps

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): cover invalid stored stream_scope reads

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): classify bedrock stream actions and keep caller is_streaming_request

Co-authored-by: Shivi Jain <mobile.350017@gmail.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(guardrails): make stream_scope_allows public and drop mutable builds

Co-authored-by: Shivi Jain <mobile.350017@gmail.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): classify pass-through and Bedrock stream scopes

Co-authored-by: Shivi Jain <mobile.350017@gmail.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): harden stream classification and validation

Co-authored-by: Shivi Jain <mobile.350017@gmail.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): add stream scope integration audit

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): avoid mutating passthrough custom body

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): run stream scope audit without enterprise license

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): keep stored scope restart cell on one worker

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): scope stream scope audit sink assertions to the rail under test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): assert logging_only scope absence behind an ordered barrier rail

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): assert one logging_only scan per phase after the barrier rail

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): set request scope on pass-through stream fixtures

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): tolerate invalid YAML stream_scope in v1 guardrails list

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* chore: merge main into litellm_guardrail_stream_scope_fixes

Update pass-through pre-call test callbacks for main's endpoint_type argument

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): add request paths to pass-through fixtures

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): cover websocket pass-through stream scope and LIT-9050 outage spend rows

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* chore(lint): remove unused type discipline suppressions

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): script the vertex live upstream in the websocket stream scope test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): keep stream marker json-serializable and restore pass-through helper names

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(lint): allow required Bedrock action re-export

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): strip the stream marker from pass-through payloads

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): drop stream marker by value in scans and snapshots, plain-tuple stream scope state

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(ui): render tag-scoped guardrail modes read-only in the custom code editor

A guardrail whose litellm_params.mode is the tag-scoped dict {tags, default}
crashed the Custom Code editor on open: normalizeMode wrapped the dict into
the mode array and StreamScopeFields rendered it as a React child (error #31,
whole dashboard unmounted). Treat a non-string non-array mode as no editable
modes, show formatGuardrailMode(mode) in a disabled input (read-only, matching
the guardrail info view), and keep mode/stream_scope out of the update payload
for such guardrails.

* fix(ui): resolve merge fallout in guardrails components

Deduplicate toModeArray import after the merge, and move the read-only
guardrail details block back into GuardrailReadOnlyDetails (now rendering
the shared mode/logging-only rows plus the stream-scope detail) so
guardrail_info.tsx stays under the 800-line lint budget. Guardrails UI
suite: 298 passed.

* fix(ui): drop duplicate toModeArray import reintroduced by merge

* chore(pass-through): document the deliberate in-place marker strip as mutable-ok

The clear/update on _parsed_body is load-bearing: rebinding to a fresh
mapping instead breaks 75 pass-through tests because the marker-free body
must propagate through the caller's request dict so downstream guardrail
scans and snapshots never observe the server streaming marker.

* fix(guardrails): move stream_scope after timeout in CustomGuardrail init

Inserting stream_scope before the existing timeout parameter shifted the
positional slot of timeout, so positional callers constructed with their
timeout bound to stream_scope (ValueError) and timeout silently None.
Restores the base parameter order; keyword callers are unaffected.

* fix(guardrails): typing pass for the lint gates

stream_scope leaves the declared constructor parameters (restoring the
base positional surface; keyword construction unchanged), unknown config
values crossing the new stream-scope code paths get typed locals or
cast-ok boundaries, and the passthrough payload literals are annotated.
All three lint gates pass against current main; the guardrail suites are
unchanged (236+5019 passing; one known anyio-driver failure pre-existing
on base).

* style: sort cast imports for the ruff gate

* chore(pass-through): drop unused BEDROCK_STREAMING_ACTIONS re-export

The streaming check now uses is_bedrock_streaming_endpoint; nothing in
the repo imports the name from this module.

---------

Co-authored-by: Shivi Jain <mobile.350017@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: gabriele <gabriele@berri.ai>
2026-10-08 13:46:53 -07:00

856 lines
36 KiB
Python

from __future__ import annotations
import base64
import json
import os
import uuid
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
import httpx
import pytest
import yaml
from integration._support.client import Gateway, JsonValue, Scenario, eventually
from integration._support.database import read_rows, write_rows
from integration._support.process import owned_proxy_process
from integration._support.upstream import (
_aws_event_frame, # pyright: ignore[reportPrivateUsage] # project Bedrock event-stream encoder
)
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import TypeAdapter
import litellm
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
EndpointScope: TypeAlias = Literal["streaming", "non_streaming"]
ModelKind: TypeAlias = Literal["router", "direct"]
StreamingAction: TypeAlias = Literal["converse-stream", "invoke-with-response-stream"]
NonStreamingAction: TypeAlias = Literal["converse", "invoke"]
BedrockAction: TypeAlias = StreamingAction | NonStreamingAction
BEDROCK_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
BEDROCK_FALSE_POSITIVE_MODEL_ID: Final = "anthropic.claude-converse-stream-test-v1:0"
BEDROCK_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api"
SERVER_STREAMING_CLASSIFICATION_KEY: Final = "litellm_server_streaming_classification"
STREAMING_ACTIONS: Final[tuple[StreamingAction, ...]] = (
"converse-stream",
"invoke-with-response-stream",
)
NON_STREAMING_ACTIONS: Final[tuple[NonStreamingAction, ...]] = ("converse", "invoke")
SCOPES: Final[tuple[EndpointScope, ...]] = ("streaming", "non_streaming")
HOSTILE_CLASSIFICATION_CASES: Final = (
pytest.param("is_streaming_request", True, id="boolean-marker"),
pytest.param("is_streaming_request", "litellm-server-streaming", id="server-marker-string"),
pytest.param("litellm_server_streaming_classification", True, id="classification-field"),
)
WORKTREE: Final = Path(__file__).resolve().parents[3]
LITELLM_PATH: Final = Path(litellm.__file__).resolve()
assert LITELLM_PATH.is_relative_to(WORKTREE), (LITELLM_PATH, WORKTREE)
print(f"stream_scope repro litellm import: {LITELLM_PATH}") # noqa: T201 # required worktree evidence
def _json(value: object) -> bytes:
return json.dumps(value, separators=(",", ":")).encode()
def _strings(value: JsonValue) -> tuple[str, ...]:
if isinstance(value, str):
return (value,)
if isinstance(value, list):
return tuple(chain.from_iterable(_strings(item) for item in value))
if isinstance(value, dict):
return tuple(chain.from_iterable(_strings(item) for item in value.values()))
return ()
def _key_names(value: JsonValue) -> tuple[str, ...]:
if isinstance(value, dict):
return tuple(value) + tuple(chain.from_iterable(_key_names(item) for item in value.values()))
if isinstance(value, list):
return tuple(chain.from_iterable(_key_names(item) for item in value))
return ()
def _marker(body: JsonValue) -> str:
return next(value for value in _strings(body) if value.startswith("scope-"))
def _bedrock_converse_stream(marker: str) -> bytes:
return b"".join(
_aws_event_frame(event_type, payload, marker, marker)
for event_type, payload in (
("messageStart", {"role": "assistant"}),
(
"contentBlockDelta",
{"delta": {"text": f"scripted Bedrock reply {marker}"}, "contentBlockIndex": 0},
),
("messageStop", {"stopReason": "end_turn"}),
("metadata", {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}}),
)
)
def _invoke_chunk(payload: Mapping[str, JsonValue], marker: str) -> bytes:
encoded: Final = base64.b64encode(_json(payload)).decode()
return _aws_event_frame("chunk", {"bytes": encoded}, marker, marker)
def _bedrock_invoke_stream(marker: str) -> bytes:
events: Final = (
{
"type": "message_start",
"message": {
"id": f"msg-{marker}",
"type": "message",
"role": "assistant",
"model": BEDROCK_MODEL_ID,
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 11, "output_tokens": 0},
},
},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": f"scripted Bedrock reply {marker}"},
},
{"type": "content_block_stop", "index": 0},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"input_tokens": 11, "output_tokens": 4},
},
{"type": "message_stop"},
)
return b"".join(_invoke_chunk(event, marker) for event in events)
def _provider(request: Request) -> Reply:
if not request.body:
return Reply(status=400, body=_json({"error": "empty request body"}))
body: Final = JSON_OBJECT.validate_json(request.body)
marker: Final = _marker(body)
target: Final = request.target.split("?", 1)[0]
if target.startswith("/passthrough"):
if body.get("stream") is True:
streamed_response: Final = _json({"received": body})
return Reply(
content_type="text/event-stream",
chunks=(b"data: " + streamed_response + b"\n\n", b"data: [DONE]\n\n"),
)
return Reply(body=_json({"received": body}))
if target.endswith("/converse-stream"):
return Reply(body=_bedrock_converse_stream(marker), content_type=BEDROCK_EVENT_STREAM)
if target.endswith("/invoke-with-response-stream"):
return Reply(body=_bedrock_invoke_stream(marker), content_type=BEDROCK_EVENT_STREAM)
if target.endswith("/converse") or target.endswith("/invoke"):
return Reply(body=_json({"output": f"scripted Bedrock reply {marker}"}))
if target == "/v1/chat/completions":
if body.get("stream") is True:
chunk: Final = {
"id": f"chatcmpl-{marker}",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": f"scripted chat reply {marker}"},
"finish_reason": None,
}
],
}
final_chunk: Final = {
**chunk,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}
return Reply(
content_type="text/event-stream",
chunks=(
b"data: " + _json(chunk) + b"\n\n",
b"data: " + _json(final_chunk) + b"\n\n",
b"data: [DONE]\n\n",
),
)
return Reply(
body=_json(
{
"id": f"chatcmpl-{marker}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": f"scripted chat reply {marker}"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
)
)
return Reply(status=404, body=_json({"error": f"unexpected upstream path: {target}"}))
def _sink(request: Request) -> Reply:
assert request.target.endswith(GUARDRAIL_PATH), request.target
assert b"scope-" in request.body, request.body.decode()
return Reply(body=_json({"action": "NONE"}))
def _rail(name: str, sink: Wire, scope: EndpointScope) -> dict[str, JsonValue]:
return {
"guardrail_name": name,
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"default_on": False,
"stream_scope": scope,
"api_base": f"{sink.url}/{name}",
"api_key": "synthetic-guardrail-key",
},
}
def _chat_proxy_config(provider_url: str, guardrails: list[dict[str, JsonValue]]) -> dict[str, object]:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
return {
**config,
"guardrails": guardrails,
"model_list": [
{
"model_name": "scope-invalid-config-chat",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": f"{provider_url}/v1",
"api_key": "synthetic-provider-key",
},
}
],
}
@dataclass(frozen=True, slots=True)
class ReproRig:
candidate: Gateway
direct_candidate: Gateway
scenario: Scenario
models: Mapping[str, str]
rails: Mapping[str, str]
provider: Wire
sink: Wire
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ReproRig]:
with httpx.Client(
base_url=os.environ["INTEGRATION_PROXY_URL"],
timeout=30,
trust_env=False,
) as root_client:
root_gateway: Final = Gateway(
root_client,
os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"),
os.environ["INTEGRATION_UPSTREAM_URL"],
)
directory: Final = tmp_path_factory.mktemp("guardrail-stream-scope-repro")
with wire_server(_provider) as provider, wire_server(_sink) as sink:
rails: Final = MappingProxyType(
{
"bedrock_streaming": "bedrock_streaming",
"bedrock_non_streaming": "bedrock_non_streaming",
"chat_streaming": "chat_streaming",
"passthrough_streaming": "passthrough_streaming",
"passthrough_non_streaming": "passthrough_non_streaming",
"passthrough_spoof_streaming": "passthrough_spoof_streaming",
}
)
models: Final = MappingProxyType(
{
"chat": "scope-chat",
"bedrock_router": "scope-bedrock-router",
"bedrock_false_positive_router": "scope-bedrock-converse-stream-model",
}
)
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["guardrails"] = [
_rail(rails["bedrock_streaming"], sink, "streaming"),
_rail(rails["bedrock_non_streaming"], sink, "non_streaming"),
_rail(rails["chat_streaming"], sink, "streaming"),
_rail(rails["passthrough_streaming"], sink, "streaming"),
_rail(rails["passthrough_non_streaming"], sink, "non_streaming"),
_rail(rails["passthrough_spoof_streaming"], sink, "streaming"),
]
config["model_list"] = [
{
"model_name": models["chat"],
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": f"{provider.url}/v1",
"api_key": "synthetic-provider-key",
},
},
{
"model_name": models["bedrock_router"],
"litellm_params": {
"model": f"bedrock/{BEDROCK_MODEL_ID}",
"api_base": provider.url,
"aws_access_key_id": "AKIASYNTHETICSTREAMSCOPE",
"aws_secret_access_key": "synthetic-bedrock-secret",
"aws_region_name": "us-east-1",
},
},
{
"model_name": models["bedrock_false_positive_router"],
"litellm_params": {
"model": f"bedrock/{BEDROCK_FALSE_POSITIVE_MODEL_ID}",
"api_base": provider.url,
"aws_access_key_id": "AKIASYNTHETICSTREAMSCOPE",
"aws_secret_access_key": "synthetic-bedrock-secret",
"aws_region_name": "us-east-1",
},
},
]
config["environment_variables"] = {
"AWS_BEDROCK_RUNTIME_ENDPOINT": provider.url,
"AWS_ACCESS_KEY_ID": "AKIASYNTHETICSTREAMSCOPE",
"AWS_SECRET_ACCESS_KEY": "synthetic-bedrock-secret",
"AWS_REGION": "us-east-1",
"AWS_REGION_NAME": "us-east-1",
}
config["general_settings"]["pass_through_endpoints"] = [
{
"path": "/pt-forward",
"target": f"{provider.url}/passthrough",
"include_subpath": True,
},
{
"path": "/pt-spoof",
"target": f"{provider.url}/passthrough",
"include_subpath": True,
"guardrails": {rails["passthrough_spoof_streaming"]: None},
},
{
"path": "/pt-scope",
"target": f"{provider.url}/passthrough",
"include_subpath": True,
"guardrails": {
rails["passthrough_streaming"]: None,
rails["passthrough_non_streaming"]: None,
},
},
]
config_path: Final = directory / "stream-scope-repro.yaml"
config_path.write_text(yaml.safe_dump(config))
with owned_proxy_process(root_gateway, directory, {}, config=config_path, workers=1) as owned:
with owned.gateway.scenario() as scenario:
direct_config: Final = {
**config,
"model_list": [
*config["model_list"],
{
"model_name": f"scope-unused-{uuid.uuid4().hex}*",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": f"{provider.url}/v1",
"api_key": "synthetic-provider-key",
},
},
],
}
direct_config_path: Final = directory / "stream-scope-direct.yaml"
direct_config_path.write_text(yaml.safe_dump(direct_config))
with owned_proxy_process(
root_gateway,
directory,
{},
config=direct_config_path,
workers=1,
) as direct:
yield ReproRig(
owned.gateway,
direct.gateway,
scenario,
models,
rails,
provider,
sink,
)
def _bedrock_request(
action: BedrockAction,
model_path: str,
marker: str,
) -> tuple[str, dict[str, JsonValue]]:
if action in ("converse", "converse-stream"):
body: Final = {
"messages": [{"role": "user", "content": [{"text": marker}]}],
"inferenceConfig": {"maxTokens": 16},
}
else:
body = {
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 16,
"messages": [{"role": "user", "content": [{"type": "text", "text": marker}]}],
}
return f"/bedrock/model/{model_path}/{action}", body
def _matching_requests(wire: Wire, marker: str) -> tuple[Request, ...]:
return tuple(request for request in wire.drain() if marker.encode() in request.body)
def _rail_scans(rows: Sequence[Request], rail_name: str, marker: str) -> tuple[Request, ...]:
return tuple(
request for request in rows if request.target.startswith(f"/{rail_name}/") and marker.encode() in request.body
)
def _chat_request_with_scans(
gateway: Gateway,
sink: Wire,
model: str,
marker: str,
streamed: bool,
rail_name: str,
) -> tuple[httpx.Response, tuple[Request, ...]]:
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": marker}],
"stream": streamed,
"guardrails": [rail_name],
},
)
return response, _rail_scans(sink.drain(), rail_name, marker)
def _spend_row_for_call(call_id: str, content: bytes) -> dict[str, JsonValue]:
rows: Final = eventually(
lambda: read_rows(
"SELECT request_id, litellm_call_id, spend, prompt_tokens, completion_tokens, metadata "
'FROM "LiteLLM_SpendLogs" WHERE request_id=%s OR litellm_call_id=%s',
(call_id, call_id),
),
lambda values: len(values) >= 1,
seconds=70,
)
assert len(rows) == 1, (call_id, rows, content)
assert call_id in (rows[0]["request_id"], rows[0]["litellm_call_id"]), content
return rows[0]
@pytest.mark.parametrize("model_kind", ("router", "direct"))
@pytest.mark.parametrize("action", STREAMING_ACTIONS)
@pytest.mark.parametrize("scope", SCOPES)
def test_bedrock_streaming_actions_run_streaming_scoped_rails(
rig: ReproRig,
model_kind: ModelKind,
action: StreamingAction,
scope: EndpointScope,
) -> None:
marker: Final = f"scope-bedrock-stream-{uuid.uuid4().hex}"
model_path: Final = rig.models["bedrock_router"] if model_kind == "router" else BEDROCK_MODEL_ID
path, body = _bedrock_request(action, model_path, marker)
expected_body: Final = (
_bedrock_converse_stream(marker) if action == "converse-stream" else _bedrock_invoke_stream(marker)
)
call_id: Final = f"stream-scope-{uuid.uuid4().hex}"
key: Final = rig.scenario.key(guardrails=[rig.rails[f"bedrock_{scope}"]])
candidate: Final = rig.candidate if model_kind == "router" else rig.direct_candidate
response: Final = candidate.request(
"POST",
path,
body,
key=key,
headers={"x-litellm-call-id": call_id},
)
assert response.status_code == 200, response.content
assert response.headers.get("content-type") == BEDROCK_EVENT_STREAM, dict(response.headers)
assert response.content == expected_body, response.content
provider_rows: Final = _matching_requests(rig.provider, marker)
assert len(provider_rows) == 1, (marker, provider_rows, response.content)
provider_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body)
key_names: Final = _key_names(provider_body)
assert "is_streaming_request" not in key_names, provider_body
assert not tuple(name for name in key_names if name.startswith("litellm_")), provider_body
_spend_row_for_call(call_id, response.content)
sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails[f"bedrock_{scope}"], marker)
expected_scans: Final = int(scope == "streaming")
assert len(sink_rows) == expected_scans, (marker, model_kind, action, scope, sink_rows, response.content)
@pytest.mark.parametrize("model_kind", ("router", "direct"))
@pytest.mark.parametrize("action", NON_STREAMING_ACTIONS)
@pytest.mark.parametrize("scope", SCOPES)
def test_bedrock_non_streaming_actions_run_non_streaming_scoped_rails(
rig: ReproRig,
model_kind: ModelKind,
action: NonStreamingAction,
scope: EndpointScope,
) -> None:
marker: Final = f"scope-bedrock-nonstream-{uuid.uuid4().hex}"
model_path: Final = rig.models["bedrock_router"] if model_kind == "router" else BEDROCK_MODEL_ID
path, body = _bedrock_request(action, model_path, marker)
key: Final = rig.scenario.key(guardrails=[rig.rails[f"bedrock_{scope}"]])
candidate: Final = rig.candidate if model_kind == "router" else rig.direct_candidate
response: Final = candidate.request("POST", path, body, key=key)
assert response.status_code == 200, response.text
assert marker in response.text, response.text
provider_rows: Final = _matching_requests(rig.provider, marker)
assert len(provider_rows) == 1, (marker, provider_rows, response.text)
provider_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body)
key_names: Final = _key_names(provider_body)
assert "is_streaming_request" not in key_names, provider_body
assert not tuple(name for name in key_names if name.startswith("litellm_")), provider_body
sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails[f"bedrock_{scope}"], marker)
expected_scans: Final = int(scope == "non_streaming")
assert len(sink_rows) == expected_scans, (marker, model_kind, action, scope, sink_rows, response.text)
@pytest.mark.parametrize("model_kind", ("router", "direct"))
def test_bedrock_model_id_streaming_action_text_on_converse_is_non_streaming(
rig: ReproRig,
model_kind: ModelKind,
) -> None:
marker: Final = f"scope-bedrock-converse-model-{uuid.uuid4().hex}"
model_path: Final = (
rig.models["bedrock_false_positive_router"] if model_kind == "router" else BEDROCK_FALSE_POSITIVE_MODEL_ID
)
path, body = _bedrock_request("converse", model_path, marker)
call_id: Final = f"stream-scope-bedrock-{uuid.uuid4().hex}"
key: Final = rig.scenario.key(
guardrails=[
rig.rails["bedrock_streaming"],
rig.rails["bedrock_non_streaming"],
]
)
candidate: Final = rig.candidate if model_kind == "router" else rig.direct_candidate
response: Final = candidate.request(
"POST",
path,
body,
key=key,
headers={"x-litellm-call-id": call_id},
)
assert response.status_code == 200, response.text
assert marker in response.text, response.text
provider_rows: Final = _matching_requests(rig.provider, marker)
assert len(provider_rows) == 1, (marker, model_kind, provider_rows, response.text)
sink_rows: Final = rig.sink.drain()
streaming_rows: Final = _rail_scans(sink_rows, rig.rails["bedrock_streaming"], marker)
non_streaming_rows: Final = _rail_scans(sink_rows, rig.rails["bedrock_non_streaming"], marker)
assert streaming_rows == (), (marker, model_kind, streaming_rows, response.text)
assert len(non_streaming_rows) == 1, (marker, model_kind, non_streaming_rows, response.text)
@pytest.mark.parametrize("streamed", (False, True), ids=("stream-absent", "stream-true"))
def test_configured_passthrough_forwards_caller_is_streaming_request_field(
rig: ReproRig,
streamed: bool,
) -> None:
marker: Final = f"scope-passthrough-forward-{uuid.uuid4().hex}"
caller_value: Final = f"caller-{uuid.uuid4().hex}"
body: Final = {
"marker": marker,
"is_streaming_request": caller_value,
**({"stream": True} if streamed else {}),
}
response: Final = rig.candidate.request("POST", "/pt-forward", body)
assert response.status_code == 200, response.text
provider_rows: Final = _matching_requests(rig.provider, marker)
assert len(provider_rows) == 1, (marker, provider_rows, response.text)
upstream_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body)
assert upstream_body.get("is_streaming_request") == caller_value, (
marker,
caller_value,
upstream_body,
response.text,
)
assert upstream_body == body, (marker, body, upstream_body, response.text)
if streamed:
assert response.headers.get("content-type", "").lower().startswith("text/event-stream"), dict(response.headers)
event_body: Final = response.text.removeprefix("data: ").split("\n", maxsplit=1)[0]
response_body: Final = JSON_OBJECT.validate_json(event_body)
else:
response_body = JSON_OBJECT.validate_json(response.content)
assert response_body == {"received": body}, response.text
@pytest.mark.parametrize(("hostile_field", "hostile_value"), HOSTILE_CLASSIFICATION_CASES)
def test_client_cannot_spoof_server_stream_classification(
rig: ReproRig,
hostile_field: str,
hostile_value: JsonValue,
) -> None:
chat_marker: Final = f"scope-chat-spoof-{uuid.uuid4().hex}"
chat_body: Final = {
"model": rig.models["chat"],
"messages": [{"role": "user", "content": chat_marker}],
"stream": False,
hostile_field: hostile_value,
}
chat_key: Final = rig.scenario.key(guardrails=[rig.rails["chat_streaming"]])
chat_response: Final = rig.candidate.request(
"POST",
"/v1/chat/completions",
chat_body,
key=chat_key,
)
chat_provider_rows: Final = _matching_requests(rig.provider, chat_marker)
chat_upstream_body: Final = chat_provider_rows[0].body.decode() if chat_provider_rows else "<no upstream request>"
print( # noqa: T201 # required chat hostile-body observation
f"chat hostile {hostile_field}={hostile_value!r}: "
f"status={chat_response.status_code}, response={chat_response.text!r}, upstream={chat_upstream_body}"
)
assert chat_response.status_code == 200, chat_response.text
assert len(chat_provider_rows) == 1, (chat_marker, chat_provider_rows, chat_response.text)
assert JSON_OBJECT.validate_json(chat_provider_rows[0].body) == {
"messages": [{"role": "user", "content": chat_marker}],
"model": "gpt-4o-mini",
hostile_field: hostile_value,
}, (hostile_field, hostile_value, chat_upstream_body)
assert _rail_scans(rig.sink.drain(), rig.rails["chat_streaming"], chat_marker) == (), (
chat_marker,
hostile_field,
hostile_value,
chat_response.text,
)
@pytest.mark.parametrize(("hostile_field", "hostile_value"), HOSTILE_CLASSIFICATION_CASES)
def test_configured_passthrough_cannot_spoof_server_stream_classification(
rig: ReproRig,
hostile_field: str,
hostile_value: JsonValue,
) -> None:
passthrough_marker: Final = f"scope-passthrough-spoof-{uuid.uuid4().hex}"
passthrough_body: Final = {
"marker": passthrough_marker,
"stream": False,
hostile_field: hostile_value,
}
passthrough_response: Final = rig.candidate.request("POST", "/pt-spoof", passthrough_body)
assert passthrough_response.status_code == 200, passthrough_response.text
passthrough_provider_rows: Final = _matching_requests(rig.provider, passthrough_marker)
assert len(passthrough_provider_rows) == 1, (
passthrough_marker,
passthrough_provider_rows,
passthrough_response.text,
)
assert _rail_scans(rig.sink.drain(), rig.rails["passthrough_spoof_streaming"], passthrough_marker) == (), (
passthrough_marker,
hostile_field,
hostile_value,
passthrough_response.text,
)
@pytest.mark.parametrize(
("streamed", "request_fields"),
((True, {"stream": True}), (False, {"stream": False}), (False, {})),
ids=("stream-true", "stream-false", "stream-absent"),
)
def test_passthrough_scope_follows_proxy_stream_decision(
rig: ReproRig,
streamed: bool,
request_fields: dict[str, JsonValue],
) -> None:
marker: Final = f"scope-passthrough-scope-{uuid.uuid4().hex}"
body: Final = {"marker": marker, **request_fields}
response: Final = rig.candidate.request("POST", "/pt-scope", body)
assert response.status_code == 200, response.text
if streamed:
expected_frame: Final = f"data: {_json({'received': body}).decode()}\n\ndata: [DONE]\n\n"
assert response.headers.get("content-type", "").lower().startswith("text/event-stream"), dict(response.headers)
assert response.text == expected_frame, response.text
assert response.headers.get("transfer-encoding", "").lower() == "chunked", dict(response.headers)
else:
response_body: Final = JSON_OBJECT.validate_json(response.content)
assert response_body == {"received": body}, response.text
assert "content-length" in response.headers, dict(response.headers)
provider_rows: Final = _matching_requests(rig.provider, marker)
assert len(provider_rows) == 1, (marker, provider_rows, response.text)
upstream_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body)
assert upstream_body == body, (marker, body, upstream_body, response.text)
assert SERVER_STREAMING_CLASSIFICATION_KEY not in upstream_body, upstream_body
sink_rows: Final = rig.sink.drain()
streaming_rows: Final = _rail_scans(sink_rows, rig.rails["passthrough_streaming"], marker)
non_streaming_rows: Final = _rail_scans(sink_rows, rig.rails["passthrough_non_streaming"], marker)
assert len(streaming_rows) == int(streamed), (marker, streamed, streaming_rows, response.text)
assert len(non_streaming_rows) == int(not streamed), (marker, streamed, non_streaming_rows, response.text)
def test_invalid_yaml_stream_scope_keeps_rail_running_on_both_shapes(rig: ReproRig, tmp_path: Path) -> None:
name: Final = f"scope-invalid-yaml-{uuid.uuid4().hex}"
invalid_rail: Final = {
"guardrail_name": name,
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"default_on": False,
"stream_scope": "sometimes",
"api_base": f"{rig.sink.url}/{name}",
"api_key": "synthetic-guardrail-key",
},
}
config: Final = _chat_proxy_config(rig.provider.url, [invalid_rail])
config_path: Final = tmp_path / "invalid-stream-scope.yaml"
config_path.write_text(yaml.safe_dump(config))
with owned_proxy_process(rig.candidate, tmp_path, {}, config=config_path, workers=1) as owned:
markers: Final = tuple(f"scope-invalid-yaml-{int(streamed)}-{uuid.uuid4().hex}" for streamed in (False, True))
observations: Final = tuple(
_chat_request_with_scans(
owned.gateway,
rig.sink,
"scope-invalid-config-chat",
marker,
streamed,
name,
)
for streamed, marker in zip((False, True), markers)
)
assert tuple(response.status_code for response, _ in observations) == (200, 200), tuple(
response.text for response, _ in observations
)
assert tuple(len(sink_rows) for _, sink_rows in observations) == (1, 1), (markers, observations)
listed: Final = owned.gateway.request("GET", "/guardrails/list")
assert listed.status_code == 200, listed.text
list_payload: Final = JSON_OBJECT.validate_json(listed.content)
listed_guardrails: Final = list_payload.get("guardrails")
assert isinstance(listed_guardrails, list), list_payload
listed_rail: Final = next(
(row for row in listed_guardrails if isinstance(row, dict) and row.get("guardrail_name") == name),
None,
)
assert isinstance(listed_rail, dict), list_payload
listed_params: Final = listed_rail.get("litellm_params")
assert isinstance(listed_params, dict), listed_rail
assert listed_params.get("stream_scope") is None, listed_rail
def test_persisted_invalid_stream_scope_row_stays_readable_and_enforced(
rig: ReproRig,
tmp_path: Path,
) -> None:
guardrail_id: Final = str(uuid.uuid4())
guardrail_name: Final = f"scope-invalid-persisted-{uuid.uuid4().hex}"
params: Final = {
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"default_on": False,
"api_base": f"{rig.sink.url}/{guardrail_name}",
"api_key": "synthetic-guardrail-key",
"stream_scope": "sometimes",
}
database_url: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL") or os.environ["DATABASE_URL"]
write_rows(
'INSERT INTO "LiteLLM_GuardrailsTable" '
'("guardrail_id", "guardrail_name", "litellm_params", "created_at", "updated_at") '
"VALUES (%s, %s, %s::jsonb, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)",
(guardrail_id, guardrail_name, json.dumps(params)),
database_url=database_url,
)
try:
config: Final = _chat_proxy_config(rig.provider.url, [])
config_path: Final = tmp_path / "persisted-invalid-stream-scope.yaml"
config_path.write_text(yaml.safe_dump(config))
with owned_proxy_process(rig.candidate, tmp_path, {}, config=config_path, workers=1) as owned:
info: Final = eventually(
lambda: owned.gateway.request("GET", f"/guardrails/{guardrail_id}/info"),
lambda response: response.status_code != 404,
seconds=20,
)
listed: Final = owned.gateway.request("GET", "/v2/guardrails/list")
list_payload: Final = JSON_OBJECT.validate_json(listed.content) if listed.status_code == 200 else {}
listed_guardrails: Final = list_payload.get("guardrails")
includes_row: Final = isinstance(listed_guardrails, list) and any(
isinstance(row, dict) and row.get("guardrail_id") == guardrail_id for row in listed_guardrails
)
markers: Final = (
f"scope-invalid-persisted-0-{uuid.uuid4().hex}",
f"scope-invalid-persisted-1-{uuid.uuid4().hex}",
)
observations: Final = tuple(
_chat_request_with_scans(
owned.gateway,
rig.sink,
"scope-invalid-config-chat",
marker,
streamed,
guardrail_name,
)
for streamed, marker in zip((False, True), markers)
)
assert (
info.status_code == 200
and listed.status_code == 200
and includes_row
and tuple(response.status_code for response, _ in observations) == (200, 200)
and tuple(len(sink_rows) for _, sink_rows in observations) == (1, 1)
), {
"info": (info.status_code, info.text),
"list": (listed.status_code, listed.text),
"includes_row": includes_row,
"responses": tuple((response.status_code, response.text) for response, _ in observations),
"scan_counts": tuple(len(sink_rows) for _, sink_rows in observations),
}
finally:
write_rows(
'DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id=%s',
(guardrail_id,),
database_url=database_url,
)
def test_management_rejects_invalid_stream_scope(rig: ReproRig) -> None:
name: Final = f"scope-invalid-management-{uuid.uuid4().hex}"
invalid_params: Final = {
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"default_on": False,
"api_base": f"{rig.sink.url}/{name}",
"api_key": "synthetic-guardrail-key",
"stream_scope": "sometimes",
}
created_invalid: Final = rig.candidate.request(
"POST",
"/guardrails",
{"guardrail": {"guardrail_name": name, "litellm_params": invalid_params}},
)
assert created_invalid.status_code == 422, created_invalid.text
valid_params: Final = {**invalid_params, "stream_scope": "both"}
created: Final = rig.candidate.request(
"POST",
"/guardrails",
{"guardrail": {"guardrail_name": name, "litellm_params": valid_params}},
)
assert created.status_code == 200, created.text
guardrail_id: Final = str(created.json()["guardrail_id"])
try:
put_response: Final = rig.candidate.request(
"PUT",
f"/guardrails/{guardrail_id}",
{"guardrail": {"guardrail_name": name, "litellm_params": invalid_params}},
)
patch_response: Final = rig.candidate.request(
"PATCH",
f"/guardrails/{guardrail_id}",
{"litellm_params": {"stream_scope": "sometimes"}},
)
assert put_response.status_code == 422, put_response.text
assert patch_response.status_code == 422, patch_response.text
finally:
deleted: Final = rig.candidate.request("DELETE", f"/guardrails/{guardrail_id}")
assert deleted.status_code == 200, deleted.text