diff --git a/tests/integration/observability/test_guardrail_pipeline_live_passthrough.py b/tests/integration/observability/test_guardrail_pipeline_live_passthrough.py new file mode 100644 index 00000000000..3589b72679c --- /dev/null +++ b/tests/integration/observability/test_guardrail_pipeline_live_passthrough.py @@ -0,0 +1,654 @@ +import base64 +import contextlib +import json +import re +import socket +import socketserver +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, held, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.upstream import aws_event_stream_frame +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +SECRET_WORD: Final = "pipelinesecretword" +REWRITTEN: Final = "[PIPELINE_REDACTED]" +TOKEN: Final = re.compile(rb"plive-[0-9a-f]{32}-\d+") +REGION: Final = "us-west-2" +VERTEX_HOST: Final = "us-central1-aiplatform.googleapis.com" +AGENT_RUNTIME_HOST: Final = f"bedrock-agent-runtime.{REGION}.amazonaws.com" +TUNNEL_HOSTS: Final = frozenset((VERTEX_HOST, AGENT_RUNTIME_HOST)) +BEDROCK_MODEL: Final = "anthropic.claude-sonnet-4-5-20250929-v1:0" +VERTEX_PROJECT: Final = "plive-pass-through" +SSE: Final = "text/event-stream" +EVENT_STREAM: Final = "application/vnd.amazon.eventstream" + + +def _sse(payload: Mapping[str, JsonValue], event: str | None = None) -> bytes: + prefix: Final = b"" if event is None else f"event: {event}\n".encode() + return prefix + b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _chat_frames(text: str) -> tuple[bytes, ...]: + def frame(delta: Mapping[str, JsonValue], finish: str | None) -> bytes: + return _sse( + { + "id": "chatcmpl-plive", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish}], + } + ) + + return ( + frame({"role": "assistant", "content": text}, None), + frame({"content": " tail"}, "stop"), + b"data: [DONE]\n\n", + ) + + +def _anthropic_frames(text: str) -> tuple[bytes, ...]: + message: Final = { + "id": "msg_plive", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 0}, + } + return ( + _sse({"type": "message_start", "message": message}, "message_start"), + _sse( + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + "content_block_start", + ), + _sse( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + "content_block_delta", + ), + _sse( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " tail"}}, + "content_block_delta", + ), + _sse({"type": "content_block_stop", "index": 0}, "content_block_stop"), + _sse( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + "message_delta", + ), + _sse({"type": "message_stop"}, "message_stop"), + ) + + +def _gemini_frames(text: str) -> tuple[bytes, ...]: + def frame(piece: str, finish: str | None) -> bytes: + candidate: Final = {"content": {"role": "model", "parts": [{"text": piece}]}, "index": 0} + return _sse( + { + "candidates": [candidate if finish is None else {**candidate, "finishReason": finish}], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 4, "totalTokenCount": 9}, + "modelVersion": "gemini-2.5-flash", + } + ) + + return (frame(text, None), frame(" tail", "STOP")) + + +def _watsonx_frames(text: str) -> tuple[bytes, ...]: + def frame(piece: str, reason: str) -> bytes: + return _sse( + { + "model_id": "ibm/granite-3-8b-instruct", + "results": [ + { + "generated_text": piece, + "generated_token_count": 2, + "input_token_count": 5, + "stop_reason": reason, + } + ], + }, + "message", + ) + + return (frame(text, "not_finished"), frame(" tail", "eos_token")) + + +def _event(kind: str, payload: Mapping[str, JsonValue]) -> bytes: + return aws_event_stream_frame( + {":event-type": kind, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload, separators=(",", ":")).encode(), + ) + + +def _converse_frames(text: str) -> tuple[bytes, ...]: + return ( + _event("messageStart", {"role": "assistant"}), + _event("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}), + _event("contentBlockDelta", {"delta": {"text": " tail"}, "contentBlockIndex": 0}), + _event("messageStop", {"stopReason": "end_turn"}), + _event("metadata", {"usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9}}), + ) + + +def _agent_frames(text: str) -> tuple[bytes, ...]: + return tuple(_event("chunk", {"bytes": base64.b64encode(piece.encode()).decode()}) for piece in (text, " tail")) + + +def _frames(target: str, text: str) -> tuple[tuple[bytes, ...], str]: + path: Final = target.split("?", 1)[0] + if path.endswith("/converse-stream"): + return _converse_frames(text), EVENT_STREAM + if path.startswith("/agents/"): + return _agent_frames(text), EVENT_STREAM + if path.endswith(":streamGenerateContent"): + return _gemini_frames(text), SSE + if path.endswith("/v1/messages"): + return _anthropic_frames(text), SSE + if path.endswith("/generation_stream"): + return _watsonx_frames(text), SSE + return _chat_frames(text), SSE + + +def _token_of(body: bytes) -> str: + found: Final = TOKEN.search(body) + assert found is not None, body + return found.group().decode() + + +def _provider_text(token: str) -> str: + return f"{SECRET_WORD} {token}" + + +def _expected(target: str, token: str) -> bytes: + chunks, _content_type = _frames(target, _provider_text(token)) + return b"".join(chunks) + + +class _Gates: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.held: Mapping[str, threading.Event] = MappingProxyType({}) + + def hold(self, token: str) -> threading.Event: + gate: Final = threading.Event() + with self._lock: + self.held = MappingProxyType({**self.held, token: gate}) + return gate + + +class _Received: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.pieces: tuple[bytes, ...] = () + + def add(self, piece: bytes) -> None: + with self._lock: + self.pieces = (*self.pieces, piece) + + @property + def content(self) -> bytes: + return b"".join(self.pieces) + + +def _replying(gates: _Gates) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + if request.target.startswith("/_oauth/token"): + return Reply( + body=b'{"access_token": "synthetic-vertex-access-token", "expires_in": 3600, "token_type": "Bearer"}' + ) + token: Final = _token_of(request.body) + chunks, content_type = _frames(request.target, _provider_text(token)) + return Reply(chunks=chunks, content_type=content_type, gate_after_first=gates.held.get(token)) + + return reply + + +def _rewrite(request: Request) -> Reply: + texts: Final = object_value(json.loads(request.body)).get("texts") or [] + assert isinstance(texts, list), request.body + rewritten: Final = [str(text).replace(SECRET_WORD, REWRITTEN) for text in texts] + return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "texts": rewritten}).encode()) + + +class _Log: + def __init__(self, wire: Wire) -> None: + self.wire: Final = wire + self._lock: Final = threading.Lock() + self.seen: tuple[Request, ...] = () + + def matching(self, token: str) -> tuple[Request, ...]: + with self._lock: + self.seen = (*self.seen, *self.wire.drain()) + return tuple(request for request in self.seen if token.encode() in request.body) + + +def _pipe(source: socket.socket, sink: socket.socket) -> None: + with contextlib.suppress(OSError): + for chunk in iter(lambda: source.recv(65536), b""): + sink.sendall(chunk) + with contextlib.suppress(OSError): + sink.shutdown(socket.SHUT_WR) + + +@contextlib.contextmanager +def _connect_tunnel(port: int) -> Iterator[str]: + class Handler(socketserver.StreamRequestHandler): + rbufsize = 0 + request: socket.socket + + def handle(self) -> None: + line: Final = self.rfile.readline().decode().split() + authority: Final = line[1] if len(line) > 1 else "" + while self.rfile.readline() not in (b"\r\n", b""): + pass + if authority.rsplit(":", 1)[0] not in TUNNEL_HOSTS: + self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n") + return + self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n") + self.request.settimeout(15) + with socket.create_connection(("127.0.0.1", port), timeout=15) as upstream: + outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream)) + outbound.start() + _pipe(upstream, self.request) + outbound.join(timeout=16) + + with socketserver.ThreadingTCPServer(("127.0.0.1", 0), Handler) as server: + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + thread.join(timeout=6) + + +@dataclass(frozen=True, slots=True) +class _Models: + azure: str + bedrock: str + nim: str + vllm: str + gigachat: str + + +@dataclass(frozen=True, slots=True) +class _PassRig: + owned: OwnedProxy + upstream: _Log + tunneled: _Log + peer: _Log + rail: str + models: _Models + gates: _Gates + + +def _model_list(models: _Models, upstream: str) -> list[JsonValue]: + return [ + { + "model_name": models.azure, + "litellm_params": { + "model": f"azure/{models.azure}", + "api_base": upstream, + "api_key": "synthetic-azure-key", + "api_version": "2024-10-21", + }, + }, + { + "model_name": models.bedrock, + "litellm_params": { + "model": f"bedrock/{BEDROCK_MODEL}", + "aws_bedrock_runtime_endpoint": upstream, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": REGION, + }, + }, + { + "model_name": models.nim, + "litellm_params": { + "model": "nvidia_nim/meta/llama-3.1-8b-instruct", + "api_base": f"{upstream}/v1", + "api_key": "synthetic-nim-key", + }, + }, + { + "model_name": models.vllm, + "litellm_params": {"model": "hosted_vllm/meta-llama/Llama-3.1-8B-Instruct", "api_base": upstream}, + }, + { + "model_name": models.gigachat, + "litellm_params": {"model": "gigachat/GigaChat-2-Max", "api_base": upstream, "api_key": "synthetic"}, + }, + ] + + +def _environment(upstream: str, credentials: Path) -> dict[str, JsonValue]: + return { + "ANTHROPIC_API_BASE": upstream, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + "GEMINI_API_BASE": upstream, + "GEMINI_API_KEY": "synthetic-gemini-key", + "OPENAI_API_BASE": upstream, + "OPENAI_API_KEY": "synthetic-openai-key", + "AZURE_API_BASE": upstream, + "AZURE_API_KEY": "synthetic-azure-key", + "AWS_BEDROCK_RUNTIME_ENDPOINT": upstream, + "AWS_ACCESS_KEY_ID": "AKIASCRIPTEDPROVIDER", + "AWS_SECRET_ACCESS_KEY": "scripted-secret", + "AWS_REGION_NAME": REGION, + "COHERE_API_BASE": upstream, + "COHERE_API_KEY": "synthetic-cohere-key", + "MISTRAL_API_BASE": upstream, + "MISTRAL_API_KEY": "synthetic-mistral-key", + "OPENROUTER_API_BASE": f"{upstream}/api/v1", + "OPENROUTER_API_KEY": "synthetic-openrouter-key", + "VLLM_API_BASE": upstream, + "WATSONX_API_BASE": upstream, + "WATSONX_TOKEN": "synthetic-watsonx-token", + "GIGACHAT_API_BASE": upstream, + "GIGACHAT_ACCESS_TOKEN": "synthetic-gigachat-token", + "TINYFISH_AGENT_API_BASE": upstream, + "TINYFISH_API_KEY": "synthetic-tinyfish-key", + "DEFAULT_VERTEXAI_PROJECT": VERTEX_PROJECT, + "DEFAULT_VERTEXAI_LOCATION": "us-central1", + "DEFAULT_GOOGLE_APPLICATION_CREDENTIALS": str(credentials), + } + + +def _config(rail: str, peer: str, upstream: str, models: _Models, credentials: Path) -> dict[str, JsonValue]: + base: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + return { + **base, + "environment_variables": { + **object_value(base.get("environment_variables") or {}), + **_environment(upstream, credentials), + }, + "general_settings": { + **object_value(base["general_settings"]), + "pass_through_endpoints": [ + {"path": "/plive-config-route", "target": f"{upstream}/configured", "auth": True} + ], + }, + "guardrails": [ + { + "guardrail_name": rail, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": f"{peer}/{rail}", + "api_key": "synthetic-guardrail-key", + "streaming_buffer_until_moderated": False, + }, + } + ], + "model_list": _model_list(models, upstream), + "policies": { + f"{rail}-policy": { + "guardrails": {"add": [rail]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": rail, "on_pass": "allow", "on_fail": "next"}], + }, + } + }, + "policy_attachments": [{"policy": f"{rail}-policy", "scope": "*"}], + } + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_PassRig]: + directory: Final = tmp_path_factory.mktemp("pipeline-live-passthrough") + suffix: Final = uuid.uuid4().hex[:8] + rail: Final = f"plive-pt-{suffix}" + models: Final = _Models(*(f"plive-pt-{kind}-{suffix}" for kind in ("azure", "bedrock", "nim", "vllm", "gigachat"))) + certificate, key = write_self_signed_cert(directory, tuple(TUNNEL_HOSTS)) + gates: Final = _Gates() + with ( + gateway_from_environment() as root, + wire_server(_replying(gates)) as upstream, + wire_server(_replying(gates), tls=server_context(certificate, key)) as tunneled, + wire_server(_rewrite) as peer, + _connect_tunnel(int(tunneled.url.rsplit(":", 1)[1])) as tunnel, + ): + credentials: Final = directory / "vertex-service-account.json" + credentials.write_text(service_account_json(VERTEX_PROJECT, upstream.url)) + config: Final = directory / "pipeline-live-passthrough.yaml" + config.write_text(yaml.safe_dump(_config(rail, peer.url, upstream.url, models, credentials))) + overrides: Final = { + "HTTPS_PROXY": tunnel, + "https_proxy": tunnel, + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(certificate), + } + with owned_proxy_process(root, directory, overrides, config=config, workers=2) as owned: + yield _PassRig(owned, _Log(upstream), _Log(tunneled), _Log(peer), rail, models, gates) + + +@dataclass(frozen=True, slots=True) +class _Route: + path: str + body: Mapping[str, JsonValue] + upstream_target: str + tunneled: bool = False + headers: Mapping[str, str] = MappingProxyType({}) + + +def _chat_body(prompt: str, model: str) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": True} + + +def _gemini_body(prompt: str) -> dict[str, JsonValue]: + return {"contents": [{"role": "user", "parts": [{"text": prompt}]}]} + + +VERTEX_TARGET: Final = ( + f"/v1/projects/{VERTEX_PROJECT}/locations/us-central1" + "/publishers/google/models/gemini-2.5-flash:streamGenerateContent" +) + + +def _route(row: str, prompt: str, models: _Models, key: str) -> _Route: + converse: Final = {"messages": [{"role": "user", "content": [{"text": prompt}]}]} + azure_path: Final = "/openai/deployments/plive-direct/chat/completions" + azure_target: Final = f"{azure_path}?api-version=2024-10-21" + match row: + case "G1-anthropic": + body: Final = {**_chat_body(prompt, "claude-sonnet-4-5"), "max_tokens": 32} + return _Route("/anthropic/v1/messages", body, "/v1/messages") + case "G2-gemini": + return _Route( + "/gemini/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse", + _gemini_body(prompt), + "/v1beta/models/gemini-2.5-flash:streamGenerateContent", + headers=MappingProxyType({"x-goog-api-key": key}), + ) + case "G3-openai": + return _Route("/openai/v1/chat/completions", _chat_body(prompt, "gpt-4o-mini"), "/v1/chat/completions") + case "G4-openai-passthrough": + return _Route( + "/openai_passthrough/v1/chat/completions", _chat_body(prompt, "gpt-4o-mini"), "/v1/chat/completions" + ) + case "G5-azure-direct": + return _Route(f"/azure{azure_target}", _chat_body(prompt, "plive-direct"), azure_path) + case "G6-azure-router-model": + return _Route( + f"/azure/openai/deployments/{models.azure}/chat/completions?api-version=2024-10-21", + {"messages": [{"role": "user", "content": prompt}], "stream": True}, + f"/openai/deployments/{models.azure}/chat/completions", + ) + case "G7-azure-ai-direct": + return _Route(f"/azure_ai{azure_target}", _chat_body(prompt, "plive-direct"), azure_path) + case "G9-bedrock-router-model": + return _Route( + f"/bedrock/model/{models.bedrock}/converse-stream", converse, f"/model/{BEDROCK_MODEL}/converse-stream" + ) + case "G10-bedrock-agent-runtime": + agent_path: Final = f"/agents/PLIVEAGENT/agentAliases/PLIVEALIAS/sessions/{uuid.uuid4().hex}/text" + return _Route(f"/bedrock{agent_path}", {"inputText": prompt}, agent_path, tunneled=True) + case "G11-vertex-ai": + return _Route(f"/vertex_ai{VERTEX_TARGET}?alt=sse", _gemini_body(prompt), VERTEX_TARGET, tunneled=True) + case "G12-vertex-ai-dash-alias": + return _Route(f"/vertex-ai{VERTEX_TARGET}?alt=sse", _gemini_body(prompt), VERTEX_TARGET, tunneled=True) + case "G13-cohere": + return _Route("/cohere/v2/chat", _chat_body(prompt, "command-a-03-2025"), "/v2/chat") + case "G14-mistral": + return _Route( + "/mistral/v1/chat/completions", _chat_body(prompt, "mistral-large-latest"), "/v1/chat/completions" + ) + case "G15-nvidia-nim-router-model": + return _Route( + f"/nvidia_nim/{models.nim}/v1/chat/completions", + _chat_body(prompt, "meta/llama-3.1-8b-instruct"), + "/v1/chat/completions", + ) + case "G16-openrouter": + return _Route( + "/openrouter/v1/chat/completions", _chat_body(prompt, "openai/gpt-4o-mini"), "/api/v1/chat/completions" + ) + case "G17-vllm-router-model": + return _Route("/vllm/v1/chat/completions", _chat_body(prompt, models.vllm), "/v1/chat/completions") + case "G18-vllm-direct": + return _Route( + "/vllm/v1/chat/completions", + _chat_body(prompt, "meta-llama/Llama-3.1-8B-Instruct"), + "/v1/chat/completions", + ) + case "G19-watsonx": + return _Route( + "/watsonx/ml/v1/text/generation_stream?version=2024-03-13", + {"input": prompt, "model_id": "ibm/granite-3-8b-instruct", "project_id": "plive-project"}, + "/ml/v1/text/generation_stream", + ) + case "G20-gigachat-direct": + return _Route("/gigachat/chat/completions", _chat_body(prompt, "GigaChat-2-Max"), "/chat/completions") + case "G21-gigachat-router-model": + return _Route("/gigachat/chat/completions", _chat_body(prompt, models.gigachat), "/chat/completions") + case "G22-tinyfish-run-sse": + return _Route( + "/tinyfish/v1/automation/run-sse", + {"url": "https://example.com", "goal": prompt}, + "/v1/automation/run-sse", + ) + case "G23-config-pass-through-endpoint": + return _Route("/plive-config-route", _chat_body(prompt, "gpt-4o-mini"), "/configured") + raise AssertionError(row) + + +ROWS: Final = ( + "G1-anthropic", + "G2-gemini", + "G3-openai", + "G4-openai-passthrough", + "G5-azure-direct", + "G6-azure-router-model", + "G7-azure-ai-direct", + "G10-bedrock-agent-runtime", + "G11-vertex-ai", + "G12-vertex-ai-dash-alias", + "G13-cohere", + "G14-mistral", + "G15-nvidia-nim-router-model", + "G16-openrouter", + "G17-vllm-router-model", + "G18-vllm-direct", + "G19-watsonx", + "G20-gigachat-direct", + "G21-gigachat-router-model", + "G22-tinyfish-run-sse", + "G23-config-pass-through-endpoint", +) + + +@dataclass(frozen=True, slots=True) +class _Sent: + token: str + route: _Route + released_while_held: bytes + content: bytes + + +def _stream(gateway: Gateway, route: _Route, received: _Received) -> int: + headers: Final = {"Authorization": f"Bearer {gateway.key}", **route.headers} + with gateway.client.stream("POST", route.path, json=dict(route.body), headers=headers) as response: + for piece in response.iter_bytes(): + received.add(piece) + return response.status_code + + +def _prepared(rig: _PassRig, row: str) -> tuple[str, _Route]: + token: Final = f"plive-{uuid.uuid4().hex}-0" + return token, _route(row, f"synthetic prompt {token}", rig.models, rig.owned.gateway.key) + + +def _vendor_log(rig: _PassRig, route: _Route) -> _Log: + return rig.tunneled if route.tunneled else rig.upstream + + +def _assert_reached_the_vendor(rig: _PassRig, row: str, sent: tuple[str, _Route], status: int, body: bytes) -> None: + token, route = sent + targets: Final = tuple(request.target.split("?", 1)[0] for request in _vendor_log(rig, route).matching(token)) + assert status == 200, (row, status, body, targets) + assert targets == (route.upstream_target,), targets + + +def _send(rig: _PassRig, row: str) -> _Sent: + token, route = _prepared(rig, row) + first_chunk: Final = _frames(route.upstream_target, _provider_text(token))[0][0] + gate: Final = rig.gates.hold(token) + received: Final = _Received() + with ThreadPoolExecutor(max_workers=1) as pool: + future: Final = pool.submit(_stream, rig.owned.gateway, route, received) + try: + eventually(lambda: _vendor_log(rig, route).matching(token), lambda requests: len(requests) >= 1, seconds=30) + early: Final = eventually( + lambda: received.content, + lambda content: content.startswith(first_chunk), + seconds=8, + return_last_on_timeout=True, + ) + finally: + gate.set() + status: Final = future.result(timeout=60) + _assert_reached_the_vendor(rig, row, (token, route), status, received.content) + return _Sent(token, route, early, received.content) + + +@pytest.mark.parametrize("row", ROWS) +def test_g_a_live_pipeline_leaves_every_pass_through_stream_live_unscanned_and_byte_identical( + rig: _PassRig, row: str +) -> None: + sent: Final = _send(rig, row) + first_chunk: Final = _frames(sent.route.upstream_target, _provider_text(sent.token))[0][0] + observed: Final = ( + sent.released_while_held.startswith(first_chunk), + sent.content == _expected(sent.route.upstream_target, sent.token), + ) + assert observed == (True, True), (row, sent) + held(lambda: rig.peer.matching(sent.token), lambda scans: not scans, holding=3) + + +def test_g9_a_bedrock_router_model_converse_stream_keeps_buffering_for_the_pipeline_rewrite(rig: _PassRig) -> None: + token, route = _prepared(rig, "G9-bedrock-router-model") + response: Final = rig.owned.gateway.request("POST", route.path, route.body, headers=route.headers) + _assert_reached_the_vendor(rig, "G9-bedrock-router-model", (token, route), response.status_code, response.content) + content: Final = response.content + assert f"{REWRITTEN} {token}".encode() in content and SECRET_WORD.encode() not in content, content + held(lambda: len(rig.peer.matching(token)), lambda scans: scans == 1, holding=3) diff --git a/tests/integration/observability/test_guardrail_pipeline_live_streaming.py b/tests/integration/observability/test_guardrail_pipeline_live_streaming.py new file mode 100644 index 00000000000..6b4cdd43b00 --- /dev/null +++ b/tests/integration/observability/test_guardrail_pipeline_live_streaming.py @@ -0,0 +1,1257 @@ +import asyncio +import json +import re +import signal +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, replace +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import anthropic +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, held, list_value, object_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, group_members, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from openai import AsyncOpenAI, OpenAI +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +SECRET_WORD: Final = "pipelinesecretword" +REWRITTEN: Final = "[PIPELINE_REDACTED]" +KEYWORD_MASK: Final = "[KEYWORD_REDACTED]" +TOKEN: Final = re.compile(rb"plive-[0-9a-f]{32}-\d+") +HOLD: Final = "holdme" +DROP: Final = "dropme" +PROVIDER_FAILURE: Final = "providerfails" +GENERIC_PATH: Final = "/beta/litellm_basic_guardrail_api" + +Endpoint: TypeAlias = Literal["chat", "responses", "messages"] +Outcome: TypeAlias = Literal["rewritten", "provider", "keyword_masked"] +ModelKind: TypeAlias = Literal["openai", "gemini", "text"] +Sink: TypeAlias = Callable[[str], None] + +PATHS: Final[Mapping[Endpoint, str]] = MappingProxyType( + {"chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"} +) + + +def _token(index: int = 0) -> str: + return f"plive-{uuid.uuid4().hex}-{index}" + + +def _provider_text(token: str) -> str: + return f"{SECRET_WORD} {token}" + + +def _prompt(token: str, *markers: str) -> str: + return " ".join(("synthetic prompt", token, *markers)) + + +def _token_of(body: bytes) -> str: + found: Final = TOKEN.search(body) + assert found is not None, body + return found.group().decode() + + +class _Log: + def __init__(self, wire: Wire) -> None: + self.wire: Final = wire + self._lock: Final = threading.Lock() + self.seen: tuple[Request, ...] = () + + def matching(self, needle: str) -> tuple[Request, ...]: + with self._lock: + self.seen = (*self.seen, *self.wire.drain()) + return tuple(request for request in self.seen if needle.encode() in request.body) + + +class _Gates: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.held: Mapping[str, threading.Event] = MappingProxyType({}) + + def hold(self, token: str) -> threading.Event: + gate: Final = threading.Event() + with self._lock: + self.held = MappingProxyType({**self.held, token: gate}) + return gate + + def release_all(self) -> None: + for gate in self.held.values(): + gate.set() + + +class _Received: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.pieces: tuple[str, ...] = () + + def add(self, piece: str) -> None: + with self._lock: + self.pieces = (*self.pieces, piece) + + @property + def text(self) -> str: + return "".join(self.pieces) + + +def _sse(payload: Mapping[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _chat_stream(text: str) -> tuple[bytes, ...]: + identity: Final = "chatcmpl-" + uuid.uuid4().hex + + def frame(delta: Mapping[str, JsonValue], finish: str | None) -> bytes: + return _sse( + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish}], + } + ) + + return ( + frame({"role": "assistant", "content": text}, None), + frame({"content": " tail"}, "stop"), + b"data: [DONE]\n\n", + ) + + +def _completion_stream(text: str) -> tuple[bytes, ...]: + identity: Final = "cmpl-" + uuid.uuid4().hex + + def frame(piece: str, finish: str | None) -> bytes: + return _sse( + { + "id": identity, + "object": "text_completion", + "created": 1, + "model": "gpt-3.5-turbo-instruct", + "choices": [{"text": piece, "index": 0, "logprobs": None, "finish_reason": finish}], + } + ) + + return (frame(text, None), frame(" tail", "stop"), b"data: [DONE]\n\n") + + +def _responses_object(identity: str, text: str, status: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "response", + "created_at": 1, + "status": status, + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_" + identity, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ], + "usage": { + "input_tokens": 11, + "output_tokens": 4, + "total_tokens": 15, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +def _responses_stream(text: str) -> tuple[bytes, ...]: + identity: Final = "resp_" + uuid.uuid4().hex + completed: Final = _responses_object(identity, text + " tail", "completed") + events: Final = ( + {"type": "response.created", "response": {**completed, "status": "in_progress", "output": [], "usage": None}}, + { + "type": "response.output_text.delta", + "item_id": "msg_" + identity, + "output_index": 0, + "content_index": 0, + "delta": text, + }, + { + "type": "response.output_text.delta", + "item_id": "msg_" + identity, + "output_index": 0, + "content_index": 0, + "delta": " tail", + }, + {"type": "response.completed", "response": completed}, + ) + encoded: Final = tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) + return (encoded[0] + encoded[1], encoded[2] + encoded[3]) + + +def _gemini_stream(text: str) -> tuple[bytes, ...]: + def frame(piece: str, finish: str | None) -> bytes: + candidate: Final = { + "content": {"parts": [{"text": piece}], "role": "model"}, + "index": 0, + **({"finishReason": finish} if finish else {}), + } + payload: Final = { + "candidates": [candidate], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, + "modelVersion": "gemini-2.5-flash", + } + return b"data: " + json.dumps(payload).encode() + b"\r\n\r\n" + + return (frame(text, None), frame(" tail", "STOP")) + + +def _chat_body(text: str) -> bytes: + return json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": text + " tail"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13}, + } + ).encode() + + +def _stream_frames(path: str, text: str) -> tuple[bytes, ...]: + if "streamGenerateContent" in path: + return _gemini_stream(text) + if path.endswith("/responses"): + return _responses_stream(text) + if path.endswith("/chat/completions"): + return _chat_stream(text) + return _completion_stream(text) + + +def _provider(gates: _Gates) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + token: Final = _token_of(request.body) + text: Final = _provider_text(token) + path: Final = request.target.split("?")[0] + if PROVIDER_FAILURE.encode() in request.body: + return Reply( + status=500, + body=json.dumps({"error": {"message": "synthetic provider failure", "type": "server_error"}}).encode(), + ) + streamed: Final = ( + "streamGenerateContent" in path or object_value(json.loads(request.body)).get("stream") is True + ) + if not streamed: + return Reply( + body=json.dumps(_responses_object("resp_" + uuid.uuid4().hex, text + " tail", "completed")).encode() + if path.endswith("/responses") + else _chat_body(text) + ) + return Reply( + content_type="text/event-stream", + chunks=_stream_frames(path, text), + gate_after_first=gates.held.get(token) if HOLD.encode() in request.body else None, + abort_after=1 if DROP.encode() in request.body else None, + ) + + return respond + + +def _rewritten(text: str) -> str: + return text.replace(SECRET_WORD, REWRITTEN) + + +def _generic_reply(behavior: str, body: Mapping[str, JsonValue]) -> Reply: + texts: Final = body.get("texts") or [] + assert isinstance(texts, list), body + if behavior == "rw": + return Reply( + body=json.dumps( + {"action": "GUARDRAIL_INTERVENED", "texts": [_rewritten(str(text)) for text in texts]} + ).encode() + ) + if behavior == "fail500": + return Reply(status=500, body=b'{"error": "synthetic guardrail failure"}') + if behavior == "malformed": + return Reply(body=b"synthetic malformed guardrail body", content_type="text/html") + if behavior == "blocked": + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode()) + if behavior == "outage" and int(_token_of(json.dumps(body).encode()).rsplit("-", 1)[1]) % 2 == 0: + return Reply(status=503, body=b'{"error": "synthetic guardrail outage"}') + return Reply(body=b'{"action": "NONE"}') + + +def _crowdstrike_reply(body: Mapping[str, JsonValue]) -> Reply: + guard_input: Final = object_value(body["guard_input"]) + messages: Final = guard_input.get("messages") or [] + assert isinstance(messages, list), body + redacted: Final = [ + {**object_value(message), "content": _rewritten(str(object_value(message).get("content", "")))} + for message in messages + ] + return Reply( + body=json.dumps( + {"result": {"blocked": False, "transformed": True, "guard_output": {"messages": redacted}}} + ).encode() + ) + + +def _bedrock_reply(body: Mapping[str, JsonValue]) -> Reply: + content: Final = body.get("content") or [] + assert isinstance(content, list), body + texts: Final = tuple(str(object_value(object_value(item)["text"])["text"]) for item in content) + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": _rewritten(text)} for text in texts], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [{"type": "NAME", "match": SECRET_WORD, "action": "ANONYMIZED"}] + } + } + ], + } + ).encode() + ) + + +def _peer(request: Request) -> Reply: + behavior: Final = request.target.split("/")[2] + body: Final = object_value(json.loads(request.body)) + if behavior == "crowdstrike": + assert request.target.endswith("/v1/guard_chat_completions"), request.target + return _crowdstrike_reply(body) + if behavior == "bedrock": + assert request.target.endswith("/version/DRAFT/apply"), request.target + return _bedrock_reply(body) + if behavior == "straiker": + assert request.target.endswith("/api/v1/detect/webhook"), request.target + return Reply(body=b'{"action": "NONE"}') + assert request.target.endswith(GENERIC_PATH), request.target + return _generic_reply(behavior, body) + + +def _flag(value: JsonValue) -> Mapping[str, JsonValue]: + return MappingProxyType({"streaming_buffer_until_moderated": value}) + + +FLAG_FALSE: Final = _flag(False) +FLAG_TRUE: Final = _flag(True) +UNSET: Final[Mapping[str, JsonValue]] = MappingProxyType({}) + + +def _generic( + name: str, peer: str, behavior: str, params: Mapping[str, JsonValue], mode: str = "post_call" +) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": f"{peer}/{name}/{behavior}", + "api_key": "synthetic-guardrail-key", + **params, + }, + } + + +def _crowdstrike(name: str, peer: str, params: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "crowdstrike_aidr", + "mode": "post_call", + "default_on": False, + "api_base": f"{peer}/{name}/crowdstrike", + "api_key": "synthetic-crowdstrike-key", + **params, + }, + } + + +def _bedrock(name: str, peer: str, params: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "bedrock", + "mode": "post_call", + "default_on": False, + "guardrailIdentifier": "synthetic-guardrail", + "guardrailVersion": "DRAFT", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIASYNTHETICGUARDRAIL", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_bedrock_runtime_endpoint": f"{peer}/{name}/bedrock", + **params, + }, + } + + +def _straiker(name: str, peer: str) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "straiker", + "mode": "post_call", + "default_on": False, + "api_base": f"{peer}/{name}/straiker", + "api_key": "synthetic-straiker-key", + "default_app": "synthetic-app", + **FLAG_FALSE, + }, + } + + +def _content_filter(name: str) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "blocked_words": [{"keyword": SECRET_WORD, "action": "MASK"}], + **FLAG_FALSE, + }, + } + + +def _step(guardrail: str, **actions: str) -> dict[str, JsonValue]: + return {"guardrail": guardrail, "on_pass": "allow", "on_fail": "next", **actions} + + +def _policy(steps: tuple[dict[str, JsonValue], ...]) -> dict[str, JsonValue]: + return { + "guardrails": {"add": list(dict.fromkeys(str(step["guardrail"]) for step in steps))}, + "pipeline": {"mode": "post_call", "steps": list(steps)}, + } + + +@dataclass(frozen=True, slots=True) +class _RowSpec: + rails: tuple[dict[str, JsonValue], ...] + pipelines: tuple[tuple[dict[str, JsonValue], ...], ...] + model: ModelKind = "openai" + + +def _single(name: str, rail: dict[str, JsonValue], **actions: str) -> _RowSpec: + return _RowSpec((rail,), ((_step(name, **actions),),)) + + +def _row_specs(peer: str, suffix: str) -> Mapping[str, _RowSpec]: + def rail(row: str, index: int = 0) -> str: + return f"plive-{row.lower()}-{index}-{suffix}" + + def generic(row: str, behavior: str, params: Mapping[str, JsonValue], **actions: str) -> _RowSpec: + return _single(rail(row), _generic(rail(row), peer, behavior, params), **actions) + + def unattached(row: str, params: Mapping[str, JsonValue], mode: str = "post_call") -> _RowSpec: + return _RowSpec((_generic(rail(row), peer, "allow", params, mode),), ()) + + hold_rows: Final = {row: generic(row, "allow", FLAG_FALSE) for row in ("A1", "A2", "A3", "A4", "A5", "A6")} + return MappingProxyType( + { + **hold_rows, + "A7": replace(generic("A7", "allow", FLAG_FALSE), model="gemini"), + "A8": replace(generic("A8", "allow", FLAG_FALSE), model="text"), + "B1": generic("B1", "rw", FLAG_FALSE), + "B2": generic("B2", "rw", FLAG_TRUE), + "B3": generic("B3", "rw", UNSET), + "B4": generic("B4", "rw", _flag(None)), + "B5": generic("B5", "rw", FLAG_FALSE, on_fail="block"), + "B6": generic("B6", "rw", FLAG_FALSE, on_error="block"), + "B7": _RowSpec( + (_generic(rail("B7", 0), peer, "rw", FLAG_FALSE), _generic(rail("B7", 1), peer, "rw", FLAG_TRUE)), + ((_step(rail("B7", 0), on_pass="next"), _step(rail("B7", 1))),), + ), + "B8": _RowSpec( + (_generic(rail("B8", 0), peer, "rw", FLAG_FALSE), _generic(rail("B8", 1), peer, "rw", FLAG_TRUE)), + ((_step(rail("B8", 0)),), (_step(rail("B8", 1)),)), + ), + "B10": generic("B10", "rw", FLAG_FALSE), + "B11": generic("B11", "rw", FLAG_FALSE), + "B12": generic("B12", "rw", FLAG_FALSE), + "B13": generic("B13", "rw", FLAG_FALSE), + "B14": generic("B14", "rw", FLAG_FALSE), + "B15": generic("B15", "rw", {"optional_params": dict(FLAG_FALSE)}), + "B16": generic("B16", "rw", _flag("")), + "B17": generic("B17", "rw", _flag("false")), + "B18": generic("B18", "rw", _flag(0)), + "B19": generic("B19", "rw", FLAG_FALSE, on_fail="modify_response"), + "C1": _single(rail("C1"), _crowdstrike(rail("C1"), peer, UNSET)), + "C2": _single(rail("C2"), _crowdstrike(rail("C2"), peer, FLAG_FALSE)), + "C3": _single(rail("C3"), _crowdstrike(rail("C3"), peer, FLAG_TRUE)), + "C4": _single(rail("C4"), _bedrock(rail("C4"), peer, UNSET)), + "C5": _single(rail("C5"), _bedrock(rail("C5"), peer, FLAG_FALSE)), + "C6": _single(rail("C6"), _bedrock(rail("C6"), peer, {**FLAG_FALSE, "mask_response_content": True})), + "C7": _single(rail("C7"), _content_filter(rail("C7"))), + "C8": _single(rail("C8"), _straiker(rail("C8"), peer)), + "D1": generic("D1", "fail500", FLAG_FALSE), + "D2": generic("D2", "malformed", FLAG_FALSE), + "D3": generic("D3", "blocked", FLAG_FALSE), + "D4-chat": generic("D4-chat", "allow", FLAG_FALSE), + "D4-responses": generic("D4-responses", "allow", FLAG_FALSE), + "D4-messages": generic("D4-messages", "allow", FLAG_FALSE), + "E1": generic("E1", "allow", FLAG_FALSE), + "F1": generic("F1", "outage", FLAG_FALSE), + "H2-chat": unattached("H2-chat", UNSET, "pre_call"), + "H2-responses": unattached("H2-responses", UNSET, "pre_call"), + "H2-messages": unattached("H2-messages", UNSET, "pre_call"), + "H3": _RowSpec((), ()), + "I1": unattached("I1", FLAG_TRUE), + "I2": unattached("I2", FLAG_FALSE), + "I3": unattached("I3", UNSET), + } + ) + + +def _deployment(kind: ModelKind, provider: str) -> dict[str, JsonValue]: + if kind == "gemini": + return {"model": "gemini/gemini-2.5-flash", "api_base": provider, "api_key": "synthetic-gemini-key"} + if kind == "text": + return { + "model": "text-completion-openai/gpt-3.5-turbo-instruct", + "api_base": f"{provider}/v1", + "api_key": "synthetic-provider-key", + } + return {"model": "openai/gpt-4o-mini", "api_base": f"{provider}/v1", "api_key": "synthetic-provider-key"} + + +def _row_policies( + specs: Mapping[str, _RowSpec], models: Mapping[str, str] +) -> Iterator[tuple[str, str, dict[str, JsonValue]]]: + for row, spec in specs.items(): + for index, steps in enumerate(spec.pipelines): + yield f"{models[row]}-p{index}", models[row], _policy(steps) + + +def _row_rails(specs: Mapping[str, _RowSpec]) -> Iterator[dict[str, JsonValue]]: + for spec in specs.values(): + yield from spec.rails + + +def _base_config() -> dict[str, JsonValue]: + return object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + + +def _matrix_config(specs: Mapping[str, _RowSpec], models: Mapping[str, str], provider: str) -> dict[str, JsonValue]: + policies: Final = tuple(_row_policies(specs, models)) + return { + **_base_config(), + "guardrails": list(_row_rails(specs)), + "model_list": [ + {"model_name": models[row], "litellm_params": _deployment(spec.model, provider)} + for row, spec in specs.items() + ], + "policies": {name: policy for name, _model, policy in policies}, + "policy_attachments": [{"policy": name, "models": [model]} for name, model, _policy_body in policies], + } + + +@dataclass(frozen=True, slots=True) +class _Stage: + gateway: Gateway + gates: _Gates + provider: _Log + + def upstream(self, token: str) -> tuple[Request, ...]: + return self.provider.matching(token) + + +@dataclass(frozen=True, slots=True) +class _Rig: + stage: _Stage + peer: _Log + models: Mapping[str, str] + rails: Mapping[str, tuple[str, ...]] + + @property + def gateway(self) -> Gateway: + return self.stage.gateway + + def scans(self, row: str, token: str, index: int = 0) -> tuple[Request, ...]: + prefix: Final = f"/{self.rails[row][index]}/" + return tuple(request for request in self.peer.matching(token) if request.target.startswith(prefix)) + + def upstream(self, token: str) -> tuple[Request, ...]: + return self.stage.upstream(token) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("pipeline-live-streaming") + gates: Final = _Gates() + suffix: Final = uuid.uuid4().hex[:8] + try: + with ( + gateway_from_environment() as root, + wire_server(_provider(gates)) as provider, + wire_server(_peer) as peer, + ): + specs: Final = _row_specs(peer.url, suffix) + models: Final = MappingProxyType({row: f"plive-{row.lower()}-{suffix}" for row in specs}) + config: Final = directory / "pipeline-live-streaming.yaml" + config.write_text(yaml.safe_dump(_matrix_config(specs, models, provider.url))) + rails: Final = MappingProxyType( + {row: tuple(str(rail["guardrail_name"]) for rail in spec.rails) for row, spec in specs.items()} + ) + with owned_proxy_process(root, directory, {}, config=config, workers=2) as owned: + yield _Rig(_Stage(owned.gateway, gates, _Log(provider)), _Log(peer), models, rails) + finally: + gates.release_all() + + +def _request_body(endpoint: Endpoint, model: str, prompt: str, stream: bool) -> dict[str, JsonValue]: + if endpoint == "chat": + return {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": stream} + if endpoint == "responses": + return {"model": model, "input": prompt, "stream": stream} + return {"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": prompt}], "stream": stream} + + +def _call( + gateway: Gateway, + endpoint: Endpoint, + model: str, + prompt: str, + *, + stream: bool = True, + extra: Mapping[str, JsonValue] = UNSET, +) -> httpx.Response: + return gateway.request("POST", PATHS[endpoint], {**_request_body(endpoint, model, prompt, stream), **extra}) + + +def _parts_text(parts: JsonValue) -> str: + return "".join(str(object_value(part).get("text") or "") for part in list_value(parts or [])) + + +def _choice_text(choice: JsonValue) -> str: + fields: Final = object_value(choice) + holder: Final = fields.get("delta") or fields.get("message") + if holder is None: + return str(fields.get("text") or "") + return str(object_value(holder).get("content") or "") + + +def _candidate_text(candidate: JsonValue) -> str: + return _parts_text(object_value(object_value(candidate).get("content") or {}).get("parts")) + + +def _event_text(event: Mapping[str, JsonValue]) -> str: + kind: Final = event.get("type") + if kind == "response.output_text.delta": + return str(event["delta"]) + if kind == "content_block_delta": + return str(object_value(event["delta"]).get("text") or "") + if kind == "message": + return _parts_text(event.get("content")) + if event.get("object") == "response": + return "".join(_parts_text(object_value(item).get("content")) for item in list_value(event.get("output") or [])) + choices: Final = "".join(_choice_text(choice) for choice in list_value(event.get("choices") or [])) + return choices + "".join(_candidate_text(candidate) for candidate in list_value(event.get("candidates") or [])) + + +def _line_text(line: str) -> str: + payload: Final = line.removeprefix("data:").strip() + if not line.startswith("data:") or not payload.startswith("{"): + return "" + return _event_text(object_value(json.loads(payload))) + + +def _text(response: httpx.Response) -> str: + if response.headers.get("content-type", "").startswith("text/event-stream"): + return "".join(_line_text(line) for line in response.text.splitlines()) + return _event_text(object_value(response.json())) + + +def _assert_outcome(text: str, token: str, outcome: Outcome) -> None: + if outcome == "provider": + assert f"{_provider_text(token)} tail" in text and REWRITTEN not in text, text + return + mask: Final = REWRITTEN if outcome == "rewritten" else KEYWORD_MASK + assert f"{mask} {token}" in text and SECRET_WORD not in text, text + + +@dataclass(frozen=True, slots=True) +class _SpendRow: + status: str + cache_hit: bool + verdicts: Mapping[str, tuple[str, ...]] + + +def _verdicts(metadata: JsonValue, rail: str) -> tuple[str, ...]: + entries: Final = object_value(metadata).get("guardrail_information") or [] + return tuple( + str(entry["guardrail_status"]) + for entry in (object_value(value) for value in list_value(entries)) + if entry.get("guardrail_name") == rail + ) + + +def _read_spend_rows(model: str) -> tuple[Mapping[str, JsonValue], ...]: + return tuple( + read_rows( + 'SELECT status, cache_hit, metadata FROM "LiteLLM_SpendLogs" WHERE model_group=%s ORDER BY "startTime"', + (model,), + ) + ) + + +def _spend_rows(model: str, rails: tuple[str, ...], count: int) -> tuple[_SpendRow, ...]: + found: Final = ( + held(lambda: _read_spend_rows(model), lambda rows: not rows, holding=10) + if count == 0 + else eventually( + lambda: _read_spend_rows(model), + lambda rows: len(rows) == count, + seconds=70, + return_last_on_timeout=True, + ) + ) + return tuple( + _SpendRow( + str(row["status"]), + str(row["cache_hit"]) == "True", + MappingProxyType({rail: _verdicts(row["metadata"], rail) for rail in rails}), + ) + for row in found + ) + + +def _assert_cached_twin_rows( + model: str, rails: tuple[str, ...], verdict: str, *, cache_hit_verdict_recorded: bool +) -> None: + rows: Final = _spend_rows(model, rails, 2) + expected: Final = {rail: (verdict,) for rail in rails} + assert len(rows) == 2 and rows[0].verdicts == expected, rows + assert [spend.cache_hit for spend in rows] == [False, True], rows + assert not cache_hit_verdict_recorded or rows[1].verdicts == expected, rows + + +def _assert_twin( + rig: _Rig, + row: str, + endpoint: Endpoint, + outcome: Outcome, + *, + stream: bool = True, + verdict: str = "success", + scans_per_request: int = 1, +) -> None: + token: Final = _token() + model: Final = rig.models[row] + first: Final = _call(rig.gateway, endpoint, model, _prompt(token), stream=stream) + second: Final = _call(rig.gateway, endpoint, model, _prompt(token), stream=stream) + assert first.status_code == second.status_code == 200, (first.text, second.text) + _assert_outcome(_text(first), token, outcome) + _assert_outcome(_text(second), token, outcome) + assert len(rig.upstream(token)) == 1, "the identical second request must be served from the cache" + expected_scans: Final = (2 * scans_per_request,) * len(rig.rails[row]) + scans: Final = eventually( + lambda: tuple(len(rig.scans(row, token, index)) for index in range(len(rig.rails[row]))), + lambda counts: counts == expected_scans, + seconds=10, + return_last_on_timeout=True, + ) + assert scans == expected_scans, scans + _assert_cached_twin_rows(model, rig.rails[row], verdict, cache_hit_verdict_recorded=stream) + + +Reader: TypeAlias = Callable[[Gateway, str, str, Sink], None] + + +def _read_lines(gateway: Gateway, path: str, body: Mapping[str, JsonValue], sink: Sink) -> None: + with gateway.client.stream( + "POST", path, json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + assert response.status_code == 200, response.read() + for line in response.iter_lines(): + sink(_line_text(line)) + + +def _chat_httpx(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + _read_lines(gateway, PATHS["chat"], _request_body("chat", model, prompt, True), sink) + + +def _responses_httpx(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + _read_lines(gateway, PATHS["responses"], _request_body("responses", model, prompt, True), sink) + + +def _messages_httpx(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + _read_lines(gateway, PATHS["messages"], _request_body("messages", model, prompt, True), sink) + + +def _gemini_httpx(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + body: Final = {"contents": [{"role": "user", "parts": [{"text": prompt}]}]} + _read_lines(gateway, f"/v1beta/models/{model}:streamGenerateContent?alt=sse", body, sink) + + +def _completions_httpx(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + _read_lines(gateway, "/v1/completions", {"model": model, "prompt": prompt, "stream": True}, sink) + + +def _opting_in(guardrails: tuple[str, ...]) -> Reader: + def read(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + body: Final = {**_request_body("chat", model, prompt, True), "guardrails": list(guardrails)} + _read_lines(gateway, PATHS["chat"], body, sink) + + return read + + +def _raw_reader(endpoint: Endpoint, raw: _Received) -> Reader: + def read(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + with gateway.client.stream( + "POST", + PATHS[endpoint], + json=_request_body(endpoint, model, prompt, True), + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + raw.add(f"status {response.status_code}\n") + for line in response.iter_lines(): + raw.add(line + "\n") + sink(_line_text(line)) + + return read + + +def _chat_async_openai_sdk(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + async def read() -> None: + client: Final = AsyncOpenAI( + base_url=str(gateway.client.base_url) + "/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=30), + ) + async with client: + stream: Final = await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": prompt}], stream=True + ) + async for chunk in stream: + sink("".join(choice.delta.content or "" for choice in chunk.choices)) + + asyncio.run(read()) + + +def _responses_openai_sdk(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + client: Final = OpenAI( + base_url=str(gateway.client.base_url) + "/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=30), + ) + with client: + for event in client.responses.create(model=model, input=prompt, stream=True): + if event.type == "response.output_text.delta": + sink(event.delta) + + +def _messages_anthropic_sdk(gateway: Gateway, model: str, prompt: str, sink: Sink) -> None: + client: Final = anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=30), + ) + with client: + for event in client.messages.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": prompt}], stream=True + ): + if event.type == "content_block_delta" and event.delta.type == "text_delta": + sink(event.delta.text) + + +@dataclass(frozen=True, slots=True) +class _HeldStream: + released_while_held: str + received: str + + +def _released(stage: _Stage, reader: Reader, model: str, token: str, *markers: str) -> _HeldStream: + gate: Final = stage.gates.hold(token) + received: Final = _Received() + with ThreadPoolExecutor(max_workers=1) as pool: + future: Final = pool.submit(reader, stage.gateway, model, _prompt(token, HOLD, *markers), received.add) + try: + eventually(lambda: stage.upstream(token), lambda requests: len(requests) >= 1, seconds=30) + early: Final = eventually( + lambda: received.text, + lambda text: _provider_text(token) in text, + seconds=8, + return_last_on_timeout=True, + ) + finally: + gate.set() + future.result(timeout=30) + return _HeldStream(early, received.text) + + +def _withheld(stage: _Stage, reader: Reader, model: str, token: str) -> str: + gate: Final = stage.gates.hold(token) + received: Final = _Received() + with ThreadPoolExecutor(max_workers=1) as pool: + future: Final = pool.submit(reader, stage.gateway, model, _prompt(token, HOLD), received.add) + try: + eventually(lambda: stage.upstream(token), lambda requests: len(requests) == 1, seconds=15) + held(lambda: received.text, lambda text: _provider_text(token) not in text, holding=3) + finally: + gate.set() + future.result(timeout=30) + return received.text + + +def _replayed(stage: _Stage, reader: Reader, model: str, token: str) -> str: + received: Final = _Received() + reader(stage.gateway, model, _prompt(token, HOLD), received.add) + return received.text + + +@dataclass(frozen=True, slots=True) +class _LiveRelease: + released_while_held: bool + whole_received: bool + upstream_requests: int + scans: int + replay_whole: bool + upstream_requests_after_replay: int + scans_after_replay: int + cache_hits: tuple[bool, ...] + + +def _scan_count(rig: _Rig, row: str, token: str, expected: int) -> int: + return eventually( + lambda: len(rig.scans(row, token)), + lambda count: count == expected, + seconds=10, + return_last_on_timeout=True, + ) + + +LIVE_RELEASE: Final = _LiveRelease(True, True, 1, 1, True, 1, 2, (False, True)) +RECORDED: Final = frozenset({("success",)}) +NEVER_SCANNED: Final = frozenset({()}) +LOGGED_BEFORE_THE_SCAN_ENDS: Final = frozenset({(), ("success",)}) + + +@pytest.mark.parametrize( + ("row", "reader", "expected", "verdicts"), + ( + pytest.param("A1", _chat_httpx, LIVE_RELEASE, RECORDED, id="A1-chat-httpx"), + pytest.param("A2", _chat_async_openai_sdk, LIVE_RELEASE, RECORDED, id="A2-chat-async-openai-sdk"), + pytest.param("A3", _responses_httpx, LIVE_RELEASE, RECORDED, id="A3-responses-httpx"), + pytest.param("A4", _responses_openai_sdk, LIVE_RELEASE, RECORDED, id="A4-responses-openai-sdk"), + pytest.param("A5", _messages_httpx, LIVE_RELEASE, RECORDED, id="A5-messages-httpx"), + pytest.param("A6", _messages_anthropic_sdk, LIVE_RELEASE, RECORDED, id="A6-messages-anthropic-sdk"), + pytest.param( + "A7", + _gemini_httpx, + _LiveRelease(True, True, 1, 1, True, 2, 2, (False, False)), + LOGGED_BEFORE_THE_SCAN_ENDS, + id="A7-native-gemini-stream-generate-content", + ), + pytest.param( + "A8", + _completions_httpx, + _LiveRelease(True, True, 1, 0, True, 1, 0, (False, True)), + NEVER_SCANNED, + id="A8-text-completions", + ), + ), +) +def test_a_live_detect_only_pipeline_releases_the_first_chunk_while_the_provider_holds_the_rest( + rig: _Rig, row: str, reader: Reader, expected: _LiveRelease, verdicts: frozenset[tuple[str, ...]] +) -> None: + token: Final = _token() + whole: Final = f"{_provider_text(token)} tail" + model: Final = rig.models[row] + rail: Final = rig.rails[row][0] + stream: Final = _released(rig.stage, reader, model, token) + upstream_requests: Final = len(rig.upstream(token)) + scans: Final = _scan_count(rig, row, token, expected.scans) + replay: Final = _replayed(rig.stage, reader, model, token) + upstream_requests_after_replay: Final = len(rig.upstream(token)) + scans_after_replay: Final = _scan_count(rig, row, token, expected.scans_after_replay) + rows: Final = _spend_rows(model, rig.rails[row], 2) + observed: Final = _LiveRelease( + _provider_text(token) in stream.released_while_held, + whole in stream.received, + upstream_requests, + scans, + whole in replay, + upstream_requests_after_replay, + scans_after_replay, + tuple(spend.cache_hit for spend in rows), + ) + assert observed == expected, (stream, replay, rows) + assert all(spend.verdicts[rail] in verdicts for spend in rows), rows + + +@pytest.mark.parametrize( + ("row", "endpoint", "outcome", "stream"), + ( + pytest.param("B1", "chat", "provider", True, id="B1-flag-false"), + pytest.param("B2", "chat", "rewritten", True, id="B2-flag-true"), + pytest.param("B3", "chat", "rewritten", True, id="B3-flag-unset"), + pytest.param("B4", "chat", "rewritten", True, id="B4-flag-null"), + pytest.param("B5", "chat", "rewritten", True, id="B5-on-fail-block"), + pytest.param("B6", "chat", "rewritten", True, id="B6-on-error-block"), + pytest.param("B7", "chat", "rewritten", True, id="B7-one-buffering-step-gates-the-pipeline"), + pytest.param("B8", "chat", "rewritten", True, id="B8-one-buffering-policy-gates-the-stream"), + pytest.param("B10", "chat", "rewritten", False, id="B10-non-streaming-chat"), + pytest.param("B11", "responses", "rewritten", False, id="B11-non-streaming-responses"), + pytest.param("B12", "messages", "rewritten", False, id="B12-non-streaming-messages"), + pytest.param("B13", "responses", "provider", True, id="B13-responses-stream"), + pytest.param("B14", "messages", "provider", True, id="B14-messages-stream"), + pytest.param("B15", "chat", "provider", True, id="B15-flag-false-in-optional-params"), + pytest.param("B16", "chat", "provider", True, id="B16-flag-empty-string"), + pytest.param("B17", "chat", "rewritten", True, id="B17-flag-false-string"), + pytest.param("B18", "chat", "provider", True, id="B18-flag-zero"), + pytest.param("B19", "chat", "rewritten", True, id="B19-on-fail-modify-response"), + pytest.param("C1", "chat", "provider", True, id="C1-crowdstrike-flag-unset"), + pytest.param("C2", "chat", "provider", True, id="C2-crowdstrike-flag-false"), + pytest.param("C3", "chat", "rewritten", True, id="C3-crowdstrike-flag-true"), + pytest.param("C4", "chat", "rewritten", True, id="C4-bedrock-flag-unset"), + pytest.param("C5", "chat", "provider", True, id="C5-bedrock-flag-false"), + pytest.param("C6", "chat", "rewritten", True, id="C6-bedrock-masking-keeps-buffering"), + ), +) +def test_b_a_pipeline_rewrite_reaches_the_client_only_while_the_stream_buffers( + rig: _Rig, row: str, endpoint: Endpoint, outcome: Outcome, stream: bool +) -> None: + _assert_twin(rig, row, endpoint, outcome, stream=stream) + + +def test_c7_an_in_process_masking_filter_keeps_masking_a_pipeline_stream(rig: _Rig) -> None: + _assert_twin(rig, "C7", "chat", "keyword_masked", scans_per_request=0) + + +def test_c8_straiker_keeps_buffering_when_the_flag_is_set_false(rig: _Rig) -> None: + token: Final = _token() + received: Final = _withheld(rig.stage, _chat_httpx, rig.models["C8"], token) + assert f"{_provider_text(token)} tail" in received, received + eventually(lambda: rig.scans("C8", token), lambda scans: len(scans) == 1, seconds=10) + + +@pytest.mark.parametrize( + ("row", "verdict"), + ( + pytest.param("D1", "guardrail_failed_to_respond", id="D1-guardrail-500"), + pytest.param("D2", "guardrail_failed_to_respond", id="D2-guardrail-malformed-body"), + pytest.param("D3", "guardrail_intervened", id="D3-guardrail-blocked-verdict-on-fail-next"), + ), +) +def test_d_a_failing_scan_on_a_live_pipeline_still_delivers_the_stream_and_records_the_verdict( + rig: _Rig, row: str, verdict: str +) -> None: + _assert_twin(rig, row, "chat", "provider", verdict=verdict) + + +@dataclass(frozen=True, slots=True) +class _CutStream: + released_while_held: bool + status_line: str + error_after_text: bool + upstream_requests: int + scans: int + rows: tuple[tuple[str, tuple[str, ...]], ...] + + +CUT_STREAM: Final = _CutStream(True, "status 200", True, 1, 1, (("failure", ("success",)),)) + + +@pytest.mark.parametrize( + ("endpoint", "expected"), + ( + pytest.param("chat", CUT_STREAM, id="chat"), + pytest.param("responses", CUT_STREAM, id="responses"), + pytest.param("messages", _CutStream(True, "status 200", True, 1, 1, ()), id="messages"), + ), +) +def test_d4_a_provider_drop_after_a_released_chunk_ends_the_stream_with_an_error_and_records_the_scan( + rig: _Rig, endpoint: Endpoint, expected: _CutStream +) -> None: + row: Final = f"D4-{endpoint}" + token: Final = _token() + raw: Final = _Received() + stream: Final = _released(rig.stage, _raw_reader(endpoint, raw), rig.models[row], token, DROP) + rows: Final = _spend_rows(rig.models[row], rig.rails[row], len(expected.rows)) + observed: Final = _CutStream( + _provider_text(token) in stream.released_while_held, + raw.text.split("\n", 1)[0], + "error" in raw.text.split(token, 1)[-1], + len(rig.upstream(token)), + _scan_count(rig, row, token, expected.scans), + tuple((spend.status, spend.verdicts[rig.rails[row][0]]) for spend in rows), + ) + assert observed == expected, (stream, raw.text, rows) + + +def test_e1_repeated_identical_live_streams_each_get_one_upstream_call_scan_and_spend_row(rig: _Rig) -> None: + token: Final = _token() + responses: Final = tuple( + _call(rig.gateway, "chat", rig.models["E1"], _prompt(token), extra={"cache": {"no-cache": True}}) + for _ in range(5) + ) + assert all(response.status_code == 200 for response in responses), [response.text for response in responses] + assert all(f"{_provider_text(token)} tail" in _text(response) for response in responses) + assert len(rig.upstream(token)) == 5 + eventually(lambda: rig.scans("E1", token), lambda scans: len(scans) == 5, seconds=10) + rows: Final = _spend_rows(rig.models["E1"], rig.rails["E1"], 5) + assert [spend.verdicts for spend in rows] == [{rig.rails["E1"][0]: ("success",)}] * 5, rows + + +def test_f1_concurrent_live_streams_during_a_guardrail_outage_each_scan_once(rig: _Rig) -> None: + base: Final = uuid.uuid4().hex + tokens: Final = tuple(f"plive-{base}-{index}" for index in range(30)) + endpoints: Final[tuple[Endpoint, ...]] = ("chat", "responses", "messages") + with ThreadPoolExecutor(max_workers=30) as pool: + responses: Final = tuple( + pool.map( + lambda index: _call(rig.gateway, endpoints[index % 3], rig.models["F1"], _prompt(tokens[index])), + range(30), + ) + ) + assert all(response.status_code == 200 for response in responses), [response.text for response in responses] + assert all( + f"{_provider_text(token)} tail" in _text(response) for token, response in zip(tokens, responses, strict=True) + ) + scanned: Final = eventually( + lambda: tuple(len(rig.scans("F1", token + " ")) for token in tokens), + lambda counts: all(count >= 1 for count in counts), + seconds=20, + ) + assert scanned == (1,) * 30, scanned + rows: Final = _spend_rows(rig.models["F1"], rig.rails["F1"], 30) + assert sorted(spend.verdicts[rig.rails["F1"][0]] for spend in rows) == sorted( + ("guardrail_failed_to_respond",) if index % 2 == 0 else ("success",) for index in range(30) + ), rows + + +@pytest.mark.parametrize("endpoint", ("chat", "responses", "messages")) +def test_h2_a_provider_failure_row_keeps_the_pre_call_guardrail_information(rig: _Rig, endpoint: Endpoint) -> None: + row: Final = f"H2-{endpoint}" + token: Final = _token() + response: Final = _call( + rig.gateway, + endpoint, + rig.models[row], + _prompt(token, PROVIDER_FAILURE), + stream=False, + extra={"guardrails": list(rig.rails[row])}, + ) + assert response.status_code == 500, response.text + rows: Final = _spend_rows(rig.models[row], rig.rails[row], 1) + assert rows[0].status == "failure" and rows[0].verdicts == {rig.rails[row][0]: ("success",)}, rows + + +def test_h3_a_provider_failure_without_a_guardrail_records_no_guardrail_information(rig: _Rig) -> None: + token: Final = _token() + response: Final = _call(rig.gateway, "chat", rig.models["H3"], _prompt(token, PROVIDER_FAILURE), stream=False) + assert response.status_code == 500, response.text + found: Final = eventually( + lambda: read_rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (rig.models["H3"],)), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert found[0]["status"] == "failure", found + assert not object_value(found[0]["metadata"]).get("guardrail_information"), found + + +def test_i1_a_request_opted_guardrail_with_the_flag_set_true_buffers_without_a_pipeline(rig: _Rig) -> None: + token: Final = _token() + received: Final = _withheld(rig.stage, _opting_in(rig.rails["I1"]), rig.models["I1"], token) + assert f"{_provider_text(token)} tail" in received, received + assert _scan_count(rig, "I1", token, 1) == 1 + + +@pytest.mark.parametrize("row", ("I2", "I3"), ids=("I2-flag-false", "I3-flag-unset")) +def test_i_a_request_opted_guardrail_without_the_flag_set_true_streams_live_without_a_pipeline( + rig: _Rig, row: str +) -> None: + token: Final = _token() + stream: Final = _released(rig.stage, _opting_in(rig.rails[row]), rig.models[row], token) + assert _provider_text(token) in stream.released_while_held, stream + assert f"{_provider_text(token)} tail" in stream.received, stream + assert _scan_count(rig, row, token, 1) == 1 + + +def _single_rail_config( + directory: Path, name: str, provider: str, guardrails: tuple[dict[str, JsonValue], ...] +) -> Path: + model: Final = f"{name}-model" + config: Final = { + **_base_config(), + "guardrails": list(guardrails), + "model_list": [{"model_name": model, "litellm_params": _deployment("openai", provider)}], + "policies": {f"{name}-policy": _policy((_step(name),))}, + "policy_attachments": [{"policy": f"{name}-policy", "models": [model]}], + } + path: Final = directory / f"{name}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def test_b9_a_pipeline_step_whose_guardrail_is_not_loaded_does_not_hold_the_stream( + gateway: Gateway, tmp_path: Path +) -> None: + gates: Final = _Gates() + name: Final = f"plive-b9-missing-{uuid.uuid4().hex[:8]}" + with wire_server(_provider(gates)) as provider: + config: Final = _single_rail_config(tmp_path, name, provider.url, ()) + try: + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + token: Final = _token() + stage: Final = _Stage(owned.gateway, gates, _Log(provider)) + stream: Final = _released(stage, _chat_httpx, f"{name}-model", token) + assert _provider_text(token) in stream.released_while_held, stream + assert f"{_provider_text(token)} tail" in stream.received, stream + finally: + gates.release_all() + + +def test_f2_live_pipeline_streams_keep_scanning_after_a_worker_is_killed(gateway: Gateway, tmp_path: Path) -> None: + gates: Final = _Gates() + name: Final = f"plive-f2-{uuid.uuid4().hex[:8]}" + with wire_server(_provider(gates)) as provider, wire_server(_peer) as peer: + config: Final = _single_rail_config( + tmp_path, name, provider.url, (_generic(name, peer.url, "allow", FLAG_FALSE),) + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + members: Final = tuple( + member for member in group_members(owned.process.pid) if member.pid != owned.process.pid + ) + workers: Final = tuple( + member for member in members if any("spawn_main" in part for part in member.cmdline()) + ) + assert len(workers) >= 2, members + workers[0].send_signal(signal.SIGKILL) + psutil.wait_procs((workers[0],), timeout=10) + base: Final = uuid.uuid4().hex + tokens: Final = tuple(f"plive-{base}-{index}" for index in range(8)) + with ThreadPoolExecutor(max_workers=8) as pool: + responses: Final = tuple( + pool.map(lambda token: _call(owned.gateway, "chat", f"{name}-model", _prompt(token)), tokens) + ) + assert all(response.status_code == 200 for response in responses), [r.text for r in responses] + assert all(f"{_provider_text(token)} tail" in _text(r) for token, r in zip(tokens, responses, strict=True)) + peer_log: Final = _Log(peer) + eventually( + lambda: tuple(len(peer_log.matching(token + " ")) for token in tokens), + lambda counts: counts == (1,) * 8, + seconds=20, + ) + rows: Final = _spend_rows(f"{name}-model", (name,), 8) + assert [spend.verdicts for spend in rows] == [{name: ("success",)}] * 8, rows + assert owned.process.poll() is None