mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
7efce7f8e0
* test(integration): give chaos and timing tests a harness API and gate raw signals Chaos tests killed and suspended proxy workers with raw os.kill and psutil calls and bounded elapsed time with fixed literals, so they raced the uvicorn worker healthcheck and failed on slow CI boxes OwnedProxy now offers kill_worker and replacing_worker, which return only once a replacement answers readiness with its own pid, and suspended, which is refused unless the proxy was launched with a lenient worker healthcheck. timing.py bounds durations against a measured control or the injected delay. run.py measures the proxy import time once and derives readiness deadlines and the per-test timeout from it, writing boot-calibration.json check_test_quality.py adds TQ010 (raw process signals) and TQ011 (elapsed time asserted under a numeric literal) for tests/integration outside _support, and every call site is migrated so both report zero with no suppressions * test(integration): steer the assistants wire burst onto both workers The mixed-status burst assumed the OS would spread 48 connections across both workers and failed when one worker took them all. Each half is now sent while the other worker is suspended, and the test asserts an exact 24 and 24 split * test(integration): restore product-enforced timing bounds lost in the migration Some TQ011 migrations dropped or loosened a bound whose time is the behaviour under test. The websocket session cap only kept its floor, so a cap closing at 300 s would have passed, and several timeouts were bounded by the injected stall instead of the configured product value assert_enforced_within bounds an elapsed time by the configured value plus explicit slack, with an optional floor. The websocket cap, the MCP validation timeout, the guardrail litellm_params timeout, the prisma command timeouts, the baseline projection lock timeout and the key metadata statement timeout use it. The device code refusal keeps its 15 s margin, and the langfuse retried prompt fetch is now bounded by a cold fetch of a steady prompt as well as by the ignored retry-after * fix(tests): wait for the supervisor to reap killed workers and fix real-kill migrations The lenient worker healthcheck was the whole case budget, so uvicorn's supervisor sat in a ping on the killed worker for up to that long before respawning it. It now uses the graceful stop budget, and a worker kill returns only after the victim is reaped, so pooled connections no longer race the dead worker's socket. Several tests on main picked the multiprocessing resource tracker as their victim and never killed a worker. Now that they kill a real one, presidio counts streams cut by the death, delete_model waits for the config reload that removes a deleted model from a replacement worker, and the wedged coordination Redis test kills before the wedge and waits for every parked publish to fail before resuming Redis. The redis cluster test imports the harness through the tests package so it collects in the unit job, and the per-model fallback rig exposes the owned proxy it now needs. * test(observability): only require S3 deliveries the killed worker could not have buffered A real SIGKILL drops the victim's unflushed batch, so requests it already answered can't be expected to land. The test now requires every request served after the replacement to land, some of the surviving worker's burst share to land, and nothing unsent to land --------- Co-authored-by: Yuneng Jiang <yuneng@berri.ai>
1150 lines
51 KiB
Python
1150 lines
51 KiB
Python
"""Straiker guardrail on both platform APIs, driven through a real proxy.
|
|
|
|
The Straiker platform is the only double: an owned HTTP sink that speaks the v1 webhook and the v3
|
|
detect wire protocols and records every request. The provider is a second owned sink. The proxy,
|
|
its guardrail registry, Postgres and Redis run for real with two workers.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import itertools
|
|
import json
|
|
import socket
|
|
import threading
|
|
import uuid
|
|
from collections.abc import Callable, Iterator
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import anthropic
|
|
import httpx
|
|
import openai
|
|
import psutil
|
|
import pytest
|
|
import yaml
|
|
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
|
|
from integration._support.database import read_rows
|
|
from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy, owned_proxy_process
|
|
from integration._support.wire import Reply, Request, wire_server
|
|
|
|
V3_KEY: Final = "sk_agt_synthetic_integration_key"
|
|
V1_KEY: Final = "synthetic-v1-collection-key"
|
|
V3_PATH: Final = "/api/v3/detect"
|
|
V1_PATH: Final = "/api/v1/detect/webhook"
|
|
BLOCK_MARK: Final = "SYNTHETIC-INJECTION"
|
|
STRAY_V3_BLOCK_MARK: Final = "SYNTHETIC-STRAY-VERSION-BLOCK"
|
|
KILL_MARK: Final = "SYNTHETIC-KILLSWITCH"
|
|
DENY_MARK: Final = "SYNTHETIC-DENY"
|
|
SINK_500_MARK: Final = "SYNTHETIC-SINK-500"
|
|
SINK_401_MARK: Final = "SYNTHETIC-SINK-401"
|
|
SINK_GARBAGE_MARK: Final = "SYNTHETIC-SINK-GARBAGE"
|
|
LOG_BLOCK_MARK: Final = "SYNTHETIC-LOG-ONLY-BLOCK"
|
|
OPEN_500_MARK: Final = "SYNTHETIC-OPEN-500"
|
|
V1_500_MARK: Final = "SYNTHETIC-V1-500"
|
|
V1_BLOCK_MARK: Final = "SYNTHETIC-V1-BLOCK"
|
|
AUDIT_AGENT: Final = "audit-agent"
|
|
POST_AGENT: Final = "post-agent"
|
|
LOG_AGENT: Final = "log-agent"
|
|
OPEN_AGENT: Final = "open-agent"
|
|
BLOCK_MESSAGE: Final = "Straiker blocked this turn: prompt-injection"
|
|
DENY_MESSAGE: Final = "Straiker denied this turn"
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Seen:
|
|
target: str
|
|
headers: dict[str, str]
|
|
body: dict[str, object]
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class Sink:
|
|
"""Owned Straiker platform double on a fixed port so a test can stop and restart it."""
|
|
|
|
port: int
|
|
seen: list[Seen] = field(default_factory=list)
|
|
lock: threading.Lock = field(default_factory=threading.Lock)
|
|
server: ThreadingHTTPServer | None = None
|
|
thread: threading.Thread | None = None
|
|
|
|
@property
|
|
def url(self) -> str:
|
|
return f"http://127.0.0.1:{self.port}"
|
|
|
|
def start(self) -> None:
|
|
sink: Final = self
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
protocol_version = "HTTP/1.1"
|
|
|
|
def do_POST(self) -> None:
|
|
raw: Final = self.rfile.read(int(self.headers.get("content-length", "0")))
|
|
body: Final = json.loads(raw)
|
|
seen: Final = Seen(self.path, {k.lower(): v for k, v in self.headers.items()}, body)
|
|
with sink.lock:
|
|
sink.seen.append(seen)
|
|
status, payload = _verdict(seen, raw.decode())
|
|
self.send_response(status)
|
|
self.send_header("content-type", "application/json")
|
|
self.send_header("content-length", str(len(payload)))
|
|
self.send_header("connection", "close")
|
|
self.end_headers()
|
|
self.wfile.write(payload)
|
|
|
|
def log_message(self, format: str, *args: object) -> None:
|
|
pass
|
|
|
|
class Server(ThreadingHTTPServer):
|
|
allow_reuse_address = True
|
|
daemon_threads = True
|
|
|
|
self.server = Server(("127.0.0.1", self.port), Handler)
|
|
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
|
|
def stop(self) -> None:
|
|
assert self.server is not None and self.thread is not None
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.thread.join(timeout=5)
|
|
self.server = None
|
|
self.thread = None
|
|
|
|
def drain(self) -> tuple[Seen, ...]:
|
|
with self.lock:
|
|
taken: Final = tuple(self.seen)
|
|
self.seen.clear()
|
|
return taken
|
|
|
|
def for_marker(self, marker: str) -> tuple[Seen, ...]:
|
|
with self.lock:
|
|
return tuple(s for s in self.seen if marker in json.dumps(s.body))
|
|
|
|
|
|
def _verdict(seen: Seen, text: str) -> tuple[int, bytes]:
|
|
agent: Final = seen.headers.get("x-s6r-agent")
|
|
if (
|
|
SINK_500_MARK in text
|
|
or (OPEN_500_MARK in text and agent == OPEN_AGENT)
|
|
or (V1_500_MARK in text and seen.target == V1_PATH)
|
|
):
|
|
return 500, b'{"error":"synthetic outage"}'
|
|
if SINK_401_MARK in text:
|
|
return 401, b'{"error":"synthetic bad key"}'
|
|
if SINK_GARBAGE_MARK in text:
|
|
return 200, b"<html>not json</html>"
|
|
if seen.target == V1_PATH:
|
|
if BLOCK_MARK in text or V1_BLOCK_MARK in text:
|
|
return 200, json.dumps({"action": "BLOCKED", "blocked_reason": BLOCK_MESSAGE}).encode()
|
|
return 200, json.dumps({"action": "NONE"}).encode()
|
|
assert seen.target == V3_PATH, seen.target
|
|
turn: Final = "turn-" + hashlib.sha256(text.encode()).hexdigest()[:12]
|
|
if (
|
|
BLOCK_MARK in text
|
|
or (STRAY_V3_BLOCK_MARK in text and agent is None)
|
|
or (LOG_BLOCK_MARK in text and agent == LOG_AGENT)
|
|
):
|
|
return 200, json.dumps(
|
|
{
|
|
"hookSpecificOutput": {"permissionDecision": "block"},
|
|
"straiker": {
|
|
"action": "block",
|
|
"blocked_by": ["prompt-injection"],
|
|
"block_message": BLOCK_MESSAGE,
|
|
"turn_id": turn,
|
|
},
|
|
}
|
|
).encode()
|
|
if DENY_MARK in text:
|
|
return 200, json.dumps({"action": "deny", "deny_reason": DENY_MESSAGE, "turn_id": turn}).encode()
|
|
if KILL_MARK in text:
|
|
return 200, json.dumps(
|
|
{"straiker": {"action": "block", "block_message": BLOCK_MESSAGE, "turn_id": turn}}
|
|
).encode()
|
|
return 200, json.dumps(
|
|
{"hookSpecificOutput": {"permissionDecision": "allow"}, "straiker": {"action": "allow", "turn_id": turn}}
|
|
).encode()
|
|
|
|
|
|
def _marker_in(body: bytes) -> str:
|
|
text: Final = body.decode()
|
|
start: Final = text.find("mark-")
|
|
return text[start : start + 37] if start >= 0 else "mark-" + uuid.uuid4().hex
|
|
|
|
|
|
def _chat_body(marker: str, answer: str) -> bytes:
|
|
return json.dumps(
|
|
{
|
|
"id": "chatcmpl-" + marker,
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"}],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
|
}
|
|
).encode()
|
|
|
|
|
|
def _chat_chunks(marker: str, answer: str) -> tuple[bytes, ...]:
|
|
def chunk(delta: dict[str, object], finish: str | None) -> bytes:
|
|
return (
|
|
"data: "
|
|
+ json.dumps(
|
|
{
|
|
"id": "chatcmpl-" + marker,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
|
|
}
|
|
)
|
|
+ "\n\n"
|
|
).encode()
|
|
|
|
return (
|
|
chunk({"role": "assistant", "content": answer[:3]}, None),
|
|
chunk({"content": answer[3:]}, "stop"),
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
|
|
|
|
def _messages_body(marker: str, answer: str) -> bytes:
|
|
return json.dumps(
|
|
{
|
|
"id": "msg_" + marker,
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"model": "claude-sonnet-4-5-20250929",
|
|
"content": [{"type": "text", "text": answer}],
|
|
"stop_reason": "end_turn",
|
|
"stop_sequence": None,
|
|
"usage": {"input_tokens": 10, "output_tokens": 5},
|
|
}
|
|
).encode()
|
|
|
|
|
|
def _messages_chunks(marker: str, answer: str) -> tuple[bytes, ...]:
|
|
def event(name: str, payload: dict[str, object]) -> bytes:
|
|
return f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode()
|
|
|
|
return (
|
|
event(
|
|
"message_start",
|
|
{
|
|
"type": "message_start",
|
|
"message": {
|
|
"id": "msg_" + marker,
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"model": "claude-sonnet-4-5-20250929",
|
|
"content": [],
|
|
"stop_reason": None,
|
|
"stop_sequence": None,
|
|
"usage": {"input_tokens": 10, "output_tokens": 1},
|
|
},
|
|
},
|
|
),
|
|
event(
|
|
"content_block_start",
|
|
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
|
),
|
|
event(
|
|
"content_block_delta",
|
|
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": answer}},
|
|
),
|
|
event("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
|
event(
|
|
"message_delta",
|
|
{
|
|
"type": "message_delta",
|
|
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
|
"usage": {"output_tokens": 5},
|
|
},
|
|
),
|
|
event("message_stop", {"type": "message_stop"}),
|
|
)
|
|
|
|
|
|
def _responses_body(marker: str, answer: str) -> bytes:
|
|
return json.dumps(
|
|
{
|
|
"id": "resp_" + marker,
|
|
"object": "response",
|
|
"created_at": 1,
|
|
"status": "completed",
|
|
"model": "gpt-4o-mini",
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"id": "msgo_" + marker,
|
|
"status": "completed",
|
|
"role": "assistant",
|
|
"content": [{"type": "output_text", "text": answer, "annotations": []}],
|
|
}
|
|
],
|
|
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
|
}
|
|
).encode()
|
|
|
|
|
|
def _completion_body(marker: str, answer: str) -> bytes:
|
|
return json.dumps(
|
|
{
|
|
"id": "cmpl-" + marker,
|
|
"object": "text_completion",
|
|
"created": 1,
|
|
"model": "gpt-3.5-turbo-instruct",
|
|
"choices": [{"index": 0, "text": answer, "finish_reason": "stop", "logprobs": None}],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
|
}
|
|
).encode()
|
|
|
|
|
|
_PROVIDER_CALLS: Final = itertools.count(1)
|
|
|
|
|
|
def _provider(request: Request) -> Reply:
|
|
if not request.body:
|
|
return Reply(status=404, body=b'{"error":"synthetic provider: no body"}')
|
|
marker: Final = _marker_in(request.body)
|
|
ident: Final = f"{marker}-{next(_PROVIDER_CALLS)}"
|
|
body: Final = json.loads(request.body)
|
|
answer: Final = "synthetic answer " + marker + (" " + BLOCK_MARK if "ANSWER-BLOCK" in request.body.decode() else "")
|
|
streaming: Final = bool(body.get("stream"))
|
|
if request.target.endswith("/v1/messages"):
|
|
return (
|
|
Reply(chunks=_messages_chunks(ident, answer), content_type="text/event-stream")
|
|
if streaming
|
|
else Reply(body=_messages_body(ident, answer))
|
|
)
|
|
if request.target.endswith("/v1/responses"):
|
|
return Reply(body=_responses_body(ident, answer))
|
|
if request.target.endswith("/v1/completions"):
|
|
return Reply(body=_completion_body(ident, answer))
|
|
assert request.target.endswith("/v1/chat/completions"), request.target
|
|
return (
|
|
Reply(chunks=_chat_chunks(ident, answer), content_type="text/event-stream")
|
|
if streaming
|
|
else Reply(body=_chat_body(ident, answer))
|
|
)
|
|
|
|
|
|
def _guardrail(name: str, key: str, url: str, mode: str, default_on: bool, **params: object) -> dict[str, object]:
|
|
return {
|
|
"guardrail_name": name,
|
|
"litellm_params": {
|
|
"guardrail": "straiker",
|
|
"mode": mode,
|
|
"default_on": default_on,
|
|
"api_key": key,
|
|
"api_base": url,
|
|
"max_retries": 0,
|
|
**params,
|
|
},
|
|
}
|
|
|
|
|
|
def _rig_config(sink_url: str, root: Path) -> Path:
|
|
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
config["litellm_settings"]["cache"] = False
|
|
config["guardrails"] = [
|
|
_guardrail("straiker-v3", V3_KEY, sink_url, "pre_call", True, agent_ref=AUDIT_AGENT),
|
|
_guardrail("straiker-v3-post", V3_KEY, sink_url, "post_call", False, agent_ref=POST_AGENT),
|
|
_guardrail("straiker-v3-log", V3_KEY, sink_url, "logging_only", False, agent_ref=LOG_AGENT),
|
|
_guardrail("straiker-v3-open", V3_KEY, sink_url, "pre_call", False, fail_on_error=False, agent_ref=OPEN_AGENT),
|
|
_guardrail(
|
|
"straiker-v3-hint",
|
|
V3_KEY,
|
|
sink_url,
|
|
"pre_call",
|
|
False,
|
|
client="named-client",
|
|
format_hint="anthropic.messages",
|
|
),
|
|
_guardrail("straiker-v3-as-v1", V3_KEY, sink_url, "pre_call", False, api_version="v1"),
|
|
_guardrail("straiker-v3-stray-version", V3_KEY, sink_url, "pre_call", False, api_version="2024-09-01"),
|
|
_guardrail("straiker-v1", V1_KEY, sink_url, "pre_call", False),
|
|
_guardrail("straiker-v1-empty-version", V1_KEY, sink_url, "pre_call", False, api_version=""),
|
|
_guardrail("straiker-v1-post", V1_KEY, sink_url, "post_call", False),
|
|
]
|
|
path: Final = root / "straiker.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
return path
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Rig:
|
|
proxy: Gateway
|
|
owned: OwnedProxy
|
|
sink: Sink
|
|
provider_url: str
|
|
provider_drain: Callable[[], tuple[Request, ...]]
|
|
chat_model: str
|
|
anthropic_model: str
|
|
completion_model: str
|
|
|
|
def marker(self) -> str:
|
|
return "mark-" + uuid.uuid4().hex
|
|
|
|
def _base(self) -> str:
|
|
return str(self.proxy.client.base_url).rstrip("/")
|
|
|
|
def openai(self, key: str | None = None) -> openai.OpenAI:
|
|
return openai.OpenAI(base_url=self._base() + "/v1", api_key=key or self.proxy.key, max_retries=0)
|
|
|
|
def async_openai(self, key: str | None = None) -> openai.AsyncOpenAI:
|
|
return openai.AsyncOpenAI(base_url=self._base() + "/v1", api_key=key or self.proxy.key, max_retries=0)
|
|
|
|
def anthropic(self) -> anthropic.Anthropic:
|
|
return anthropic.Anthropic(base_url=self._base(), api_key=self.proxy.key, max_retries=0)
|
|
|
|
def async_anthropic(self) -> anthropic.AsyncAnthropic:
|
|
return anthropic.AsyncAnthropic(base_url=self._base(), api_key=self.proxy.key, max_retries=0)
|
|
|
|
def sink_calls(self, marker: str) -> tuple[Seen, ...]:
|
|
return self.sink.for_marker(marker)
|
|
|
|
def provider_calls(self, marker: str, requests: tuple[Request, ...]) -> tuple[Request, ...]:
|
|
return tuple(r for r in requests if marker.encode() in r.body)
|
|
|
|
def spend_row(self, request_id: str) -> dict[str, object]:
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT request_id, model, call_type, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s',
|
|
(request_id,),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
return rows[0]
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
|
root: Final = tmp_path_factory.mktemp("straiker")
|
|
with socket.socket() as reserve:
|
|
reserve.bind(("127.0.0.1", 0))
|
|
port: Final = reserve.getsockname()[1]
|
|
sink: Final = Sink(port)
|
|
sink.start()
|
|
with gateway_from_environment() as gateway, wire_server(_provider) as provider:
|
|
config: Final = _rig_config(sink.url, root)
|
|
with (
|
|
owned_proxy_process(
|
|
gateway, root, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT
|
|
) as owned,
|
|
owned.gateway.scenario() as scenario,
|
|
):
|
|
chat: Final = scenario.model(
|
|
model="openai/gpt-4o-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key"
|
|
)
|
|
claude: Final = scenario.model(
|
|
model="anthropic/claude-sonnet-4-5-20250929", api_base=provider.url, api_key="synthetic-anthropic-key"
|
|
)
|
|
completion: Final = scenario.model(
|
|
model="text-completion-openai/gpt-3.5-turbo-instruct",
|
|
api_base=provider.url + "/v1",
|
|
api_key="synthetic-openai-key",
|
|
)
|
|
yield Rig(owned.gateway, owned, sink, provider.url, provider.drain, chat, claude, completion)
|
|
if sink.server is not None:
|
|
sink.stop()
|
|
|
|
|
|
def _messages(text: str, system: str | None = None) -> list[dict[str, object]]:
|
|
return ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": text}]
|
|
|
|
|
|
def _chat(
|
|
rig: Rig, text: str, *, key: str | None = None, headers: dict[str, str] | None = None, **extra: object
|
|
) -> httpx.Response:
|
|
return rig.proxy.client.post(
|
|
"/v1/chat/completions",
|
|
json={"model": rig.chat_model, "messages": _messages(text), **extra},
|
|
headers={"Authorization": f"Bearer {key or rig.proxy.key}", **(headers or {})},
|
|
)
|
|
|
|
|
|
def _v3_request_calls(rig: Rig, marker: str, agent: str | None = AUDIT_AGENT) -> tuple[Seen, ...]:
|
|
return tuple(
|
|
s
|
|
for s in rig.sink_calls(marker)
|
|
if s.target == V3_PATH and "straiker_phase" not in s.body and s.headers.get("x-s6r-agent") == agent
|
|
)
|
|
|
|
|
|
def _v3_response_calls(rig: Rig, marker: str, agent: str | None = POST_AGENT) -> tuple[Seen, ...]:
|
|
return tuple(
|
|
s
|
|
for s in rig.sink_calls(marker)
|
|
if s.target == V3_PATH
|
|
and s.body.get("straiker_phase") == "response-sync"
|
|
and s.headers.get("x-s6r-agent") == agent
|
|
)
|
|
|
|
|
|
def _v1_calls(rig: Rig, marker: str, key: str) -> tuple[Seen, ...]:
|
|
return tuple(
|
|
s for s in rig.sink_calls(marker) if s.target == V1_PATH and s.headers.get("authorization") == "Bearer " + key
|
|
)
|
|
|
|
|
|
# H1: default_on v3 pre_call, OpenAI SDK sync, non-streaming
|
|
def test_v3_pre_call_allow_relays_provider_body_and_key_identity(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
with rig.proxy.scenario() as scenario:
|
|
key: Final = scenario.key(key_alias="alias-" + marker, metadata={"user_api_key_user_email": "n/a"})
|
|
response: Final = rig.openai(key).chat.completions.create(
|
|
model=rig.chat_model,
|
|
messages=[{"role": "user", "content": "hello " + marker}],
|
|
temperature=0.2,
|
|
user="end-" + marker,
|
|
)
|
|
assert response.id.startswith("chatcmpl-" + marker), response.id
|
|
assert response.choices[0].message.content == "synthetic answer " + marker
|
|
calls: Final = _v3_request_calls(rig, marker)
|
|
assert len(calls) == 1, calls
|
|
sent: Final = calls[0]
|
|
assert sent.headers["authorization"] == "Bearer " + V3_KEY
|
|
assert "x-straiker-webhook-format" not in sent.headers
|
|
assert sent.headers["x-s6r-agent"] == "audit-agent"
|
|
assert sent.body["messages"] == [{"role": "user", "content": "hello " + marker}]
|
|
assert sent.body["temperature"] == 0.2
|
|
assert sent.body["model"] == rig.chat_model
|
|
assert "api_key" not in sent.body and "synthetic-openai-key" not in json.dumps(sent.body)
|
|
assert object_value(sent.body["metadata"])["user_api_key_alias"] == "alias-" + marker
|
|
assert sent.body.get("session_id", "").startswith("litellm-")
|
|
upstream: Final = rig.provider_calls(marker, rig.provider_drain())
|
|
assert len(upstream) == 1 and upstream[0].target == "/v1/chat/completions"
|
|
row: Final = rig.spend_row(response.id)
|
|
assert row["model"] == "openai/gpt-4o-mini", row
|
|
|
|
|
|
# H2: v3 block verdict on the request phase blocks with the platform's message
|
|
def test_v3_block_verdict_returns_400_with_block_message_and_no_provider_call(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{BLOCK_MARK} {marker}")
|
|
assert response.status_code == 400, response.text
|
|
assert response.json()["error"]["message"] == BLOCK_MESSAGE, response.text
|
|
assert len(_v3_request_calls(rig, marker)) == 1
|
|
assert rig.provider_calls(marker, rig.provider_drain()) == ()
|
|
|
|
|
|
# H3: a resend of a blocked conversation is blocked by the process that saw the block, without a second detect call
|
|
def test_v3_blocked_conversation_replays_block_without_asking_again(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
session: Final = {"x-claude-code-session-id": "session-" + marker}
|
|
first: Final = _chat(rig, f"{BLOCK_MARK} {marker}", headers=session)
|
|
assert first.status_code == 400, first.text
|
|
baseline: Final = len(_v3_request_calls(rig, marker))
|
|
assert baseline == 1
|
|
outcomes: Final = tuple(_chat(rig, f"{BLOCK_MARK} {marker}", headers=session) for _ in range(6))
|
|
assert all(r.status_code == 400 and r.json()["error"]["message"] == BLOCK_MESSAGE for r in outcomes), [
|
|
r.text for r in outcomes
|
|
]
|
|
later: Final = len(_v3_request_calls(rig, marker))
|
|
# Two workers: only the worker that saw the block replays from memory, the other asks Straiker once
|
|
assert baseline <= later <= 2, later
|
|
grown: Final = rig.proxy.client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": rig.chat_model,
|
|
"messages": _messages(f"{BLOCK_MARK} {marker}")
|
|
+ [{"role": "assistant", "content": "x"}, {"role": "user", "content": "more"}],
|
|
},
|
|
headers={"Authorization": f"Bearer {rig.proxy.key}", **session},
|
|
)
|
|
assert grown.status_code == 400, grown.text
|
|
assert rig.provider_calls(marker, rig.provider_drain()) == ()
|
|
|
|
|
|
# H4: a kill-switch block (no blocked_by) blocks but is not remembered, so Straiker is asked every time
|
|
def test_v3_killswitch_block_is_not_remembered(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
session: Final = {"x-claude-code-session-id": "session-" + marker}
|
|
outcomes: Final = tuple(_chat(rig, f"{KILL_MARK} {marker}", headers=session) for _ in range(3))
|
|
assert all(r.status_code == 400 and r.json()["error"]["message"] == BLOCK_MESSAGE for r in outcomes)
|
|
assert len(_v3_request_calls(rig, marker)) == 3
|
|
|
|
|
|
# H4b: a deny decision on the flat envelope also blocks, with the deny_reason
|
|
def test_v3_flat_deny_decision_blocks_with_deny_reason(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{DENY_MARK} {marker}")
|
|
assert response.status_code == 400, response.text
|
|
assert response.json()["error"]["message"] == DENY_MESSAGE
|
|
|
|
|
|
# H5: post_call non-streaming, selected per request, async OpenAI SDK
|
|
@pytest.mark.asyncio
|
|
async def test_v3_post_call_sends_response_phase_with_answer(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = await rig.async_openai().chat.completions.create(
|
|
model=rig.chat_model,
|
|
messages=[{"role": "user", "content": "post " + marker}],
|
|
extra_body={"guardrails": ["straiker-v3-post"]},
|
|
)
|
|
assert response.id.startswith("chatcmpl-" + marker), response.id
|
|
calls: Final = eventually(lambda: _v3_response_calls(rig, marker), lambda c: len(c) == 1)
|
|
phase: Final = calls[0].body
|
|
assert phase["model"] == "gpt-4o-mini", "the deployment's model, not the alias"
|
|
assert object_value(phase["request"])["messages"] == [{"role": "user", "content": "post " + marker}]
|
|
assert json.loads(str(phase["sse"]))["id"].startswith("chatcmpl-" + marker)
|
|
assert json.loads(str(phase["sse"]))["choices"][0]["message"]["content"] == "synthetic answer " + marker
|
|
assert len(_v3_request_calls(rig, marker)) == 1, "the default_on pre_call route still runs beside the selected one"
|
|
assert rig.spend_row(response.id)["request_id"] == response.id
|
|
|
|
|
|
# H5b: post_call block replaces the answer with the block message as a 200
|
|
def test_v3_post_call_block_replaces_answer(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, "ANSWER-BLOCK " + marker, guardrails=["straiker-v3-post"])
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["choices"][0]["message"]["content"] == BLOCK_MESSAGE, response.text
|
|
assert len(_v3_response_calls(rig, marker)) == 1
|
|
|
|
|
|
# H6: post_call streaming, OpenAI SDK sync; the stream is consumed to the end before the phase is sent
|
|
def test_v3_post_call_streaming_sends_assembled_answer(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
stream: Final = rig.openai().chat.completions.create(
|
|
model=rig.chat_model,
|
|
messages=[{"role": "user", "content": "stream " + marker}],
|
|
stream=True,
|
|
extra_body={"guardrails": ["straiker-v3-post"]},
|
|
)
|
|
chunks: Final = list(stream)
|
|
assert chunks and all(c.id.startswith("chatcmpl-" + marker) for c in chunks)
|
|
text: Final = "".join(c.choices[0].delta.content or "" for c in chunks if c.choices)
|
|
assert text == "synthetic answer " + marker
|
|
calls: Final = eventually(lambda: _v3_response_calls(rig, marker), lambda c: len(c) == 1)
|
|
sse: Final = json.loads(str(calls[0].body["sse"]))
|
|
assert "synthetic answer " + marker in json.dumps(sse)
|
|
assert object_value(calls[0].body["request"])["stream"] is True
|
|
|
|
|
|
# H7: Anthropic Messages sync, pre_call, session header and recognised client
|
|
def test_v3_anthropic_messages_relays_system_and_routing_headers(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
client: Final = rig.anthropic().with_options(
|
|
default_headers={"x-claude-code-session-id": "cc-" + marker, "User-Agent": "claude-cli/2.0.0 (external, cli)"}
|
|
)
|
|
response: Final = client.messages.create(
|
|
model=rig.anthropic_model,
|
|
max_tokens=16,
|
|
system="synthetic system " + marker,
|
|
messages=[{"role": "user", "content": "anthropic " + marker}],
|
|
)
|
|
assert response.id.startswith("msg_" + marker), response.id
|
|
assert response.content[0].text == "synthetic answer " + marker
|
|
calls: Final = _v3_request_calls(rig, marker)
|
|
assert len(calls) == 1, calls
|
|
sent: Final = calls[0]
|
|
assert sent.headers["x-claude-code-session-id"] == "cc-" + marker
|
|
assert sent.headers["x-s6r-client"] == "claude"
|
|
assert sent.headers["x-s6r-agent"] == "audit-agent", "YAML agent_ref wins over the User-Agent derived agent"
|
|
assert sent.body["session_id"] == "cc-" + marker
|
|
assert sent.body["system"] == "synthetic system " + marker
|
|
assert sent.body["messages"] == [{"role": "user", "content": "anthropic " + marker}]
|
|
assert sent.body["max_tokens"] == 16
|
|
upstream: Final = rig.provider_calls(marker, rig.provider_drain())
|
|
assert len(upstream) == 1 and upstream[0].target == "/v1/messages"
|
|
assert rig.spend_row(response.id)["call_type"] == "anthropic_messages"
|
|
|
|
|
|
# H8: Anthropic Messages streaming, async SDK, post_call: the answer is scored in Messages shape
|
|
@pytest.mark.asyncio
|
|
async def test_v3_anthropic_streaming_post_call_scores_messages_shaped_answer(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
client: Final = rig.async_anthropic()
|
|
async with client.messages.stream(
|
|
model=rig.anthropic_model,
|
|
max_tokens=16,
|
|
messages=[{"role": "user", "content": "astream " + marker}],
|
|
extra_body={"guardrails": ["straiker-v3-post"]},
|
|
) as stream:
|
|
final: Final = await stream.get_final_message()
|
|
assert final.id.startswith("msg_" + marker), final.id
|
|
assert final.content[0].text == "synthetic answer " + marker
|
|
calls: Final = eventually(lambda: _v3_response_calls(rig, marker), lambda c: len(c) == 1)
|
|
sse: Final = json.loads(str(calls[0].body["sse"]))
|
|
assert sse.get("type") == "message", sse
|
|
assert sse["content"][0]["text"] == "synthetic answer " + marker
|
|
|
|
|
|
# H9: Responses API, raw httpx, pre_call relays `input`, `instructions`, and the answer on post_call
|
|
def test_v3_responses_api_relays_input_and_answer(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = rig.proxy.client.post(
|
|
"/v1/responses",
|
|
json={
|
|
"model": rig.chat_model,
|
|
"input": "responses " + marker,
|
|
"instructions": "be brief",
|
|
"guardrails": ["straiker-v3", "straiker-v3-post"],
|
|
},
|
|
headers={"Authorization": f"Bearer {rig.proxy.key}"},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["id"].startswith("resp_"), response.text
|
|
assert "synthetic answer " + marker in response.text
|
|
pre: Final = _v3_request_calls(rig, marker)
|
|
assert len(pre) == 1 and pre[0].body["input"] == "responses " + marker and pre[0].body["instructions"] == "be brief"
|
|
post: Final = eventually(lambda: _v3_response_calls(rig, marker), lambda c: len(c) == 1)
|
|
assert "synthetic answer " + marker in str(post[0].body["sse"])
|
|
upstream: Final = rig.provider_calls(marker, rig.provider_drain())
|
|
assert len(upstream) == 1 and upstream[0].target == "/v1/responses"
|
|
|
|
|
|
# H10/H11: Completions API prompt becomes messages on the request phase; the answer is sent as a chat completion
|
|
def test_v3_completions_prompt_is_relayed_as_messages_and_answer_as_chat(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = rig.proxy.client.post(
|
|
"/v1/completions",
|
|
json={
|
|
"model": rig.completion_model,
|
|
"prompt": "complete " + marker,
|
|
"guardrails": ["straiker-v3", "straiker-v3-post"],
|
|
},
|
|
headers={"Authorization": f"Bearer {rig.proxy.key}"},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["id"].startswith("cmpl-" + marker), response.json()["id"]
|
|
pre: Final = _v3_request_calls(rig, marker)
|
|
assert len(pre) == 1, pre
|
|
assert pre[0].body["messages"] == [{"role": "user", "content": "complete " + marker}]
|
|
assert "prompt" not in pre[0].body
|
|
post: Final = eventually(lambda: _v3_response_calls(rig, marker), lambda c: len(c) == 1)
|
|
sse: Final = json.loads(str(post[0].body["sse"]))
|
|
assert sse["object"] == "chat.completion", sse
|
|
assert sse["choices"][0]["message"]["content"] == "synthetic answer " + marker
|
|
|
|
|
|
# H12: tool and MCP server credentials are redacted one level deep; a schema property named headers is kept
|
|
def test_v3_redacts_tool_credentials_but_keeps_schema_properties(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
tools: Final = [
|
|
{
|
|
"type": "function",
|
|
"authorization": "Bearer synthetic-tool-secret",
|
|
"function": {
|
|
"name": "lookup",
|
|
"parameters": {"type": "object", "properties": {"headers": {"type": "string"}}},
|
|
},
|
|
}
|
|
]
|
|
response: Final = _chat(
|
|
rig,
|
|
"tools " + marker,
|
|
tools=tools,
|
|
mcp_servers=[{"url": "http://mcp", "authorization_token": "synthetic-mcp-secret"}],
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
sent: Final = _v3_request_calls(rig, marker)[0].body
|
|
assert sent["tools"][0]["authorization"] == "[redacted]" # pyright: ignore[reportIndexIssue] # sink body is loose JSON
|
|
assert sent["tools"][0]["function"]["parameters"]["properties"]["headers"] == {"type": "string"} # pyright: ignore[reportIndexIssue] # sink body is loose JSON
|
|
assert sent["mcp_servers"][0]["authorization_token"] == "[redacted]" # pyright: ignore[reportIndexIssue] # sink body is loose JSON
|
|
assert "synthetic-tool-secret" not in json.dumps(sent) and "synthetic-mcp-secret" not in json.dumps(sent)
|
|
|
|
|
|
# U1: a v1 collection key still speaks the v1 webhook with the litellm envelope
|
|
def test_v1_key_keeps_webhook_envelope_and_format_header(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, "v1 " + marker, guardrails=["straiker-v1"])
|
|
assert response.status_code == 200, response.text
|
|
calls: Final = _v1_calls(rig, marker, V1_KEY)
|
|
assert len(calls) == 1, rig.sink_calls(marker)
|
|
assert calls[0].headers["x-straiker-webhook-format"] == "litellm"
|
|
assert calls[0].body["schema_version"] and object_value(calls[0].body["event"])["type"]
|
|
assert "v1 " + marker in json.dumps(object_value(calls[0].body["request"]))
|
|
assert len(_v3_request_calls(rig, marker)) == 1, "the default_on v3 route runs beside it"
|
|
|
|
|
|
# U2: v1 block verdict still blocks
|
|
def test_v1_block_verdict_still_blocks(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{V1_BLOCK_MARK} {marker}", guardrails=["straiker-v1"])
|
|
assert response.status_code == 400, response.text
|
|
assert response.json()["error"]["message"] == BLOCK_MESSAGE
|
|
assert len(_v1_calls(rig, marker, V1_KEY)) == 1, rig.sink_calls(marker)
|
|
assert rig.provider_calls(marker, rig.provider_drain()) == ()
|
|
|
|
|
|
# U3: v1 post_call still receives the response envelope
|
|
def test_v1_post_call_sends_response_envelope(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, "v1post " + marker, guardrails=["straiker-v1-post"])
|
|
assert response.status_code == 200, response.text
|
|
calls: Final = eventually(lambda: _v1_calls(rig, marker, V1_KEY), lambda c: len(c) == 1)
|
|
assert "synthetic answer " + marker in json.dumps(calls[0].body.get("response"))
|
|
|
|
|
|
# E: a saved api_version v1 with an sk_agt_ key routes to v3; the key prefix decides, not the saved version
|
|
def test_saved_api_version_v1_with_v3_key_routes_to_v3_not_the_v1_webhook(rig: Rig) -> None:
|
|
allowed_marker: Final = rig.marker()
|
|
allowed: Final = _chat(rig, "saved v1 " + allowed_marker, guardrails=["straiker-v3-as-v1"])
|
|
assert allowed.status_code == 200, allowed.text
|
|
assert len(_v3_request_calls(rig, allowed_marker, agent=None)) == 1, rig.sink_calls(allowed_marker)
|
|
assert _v1_calls(rig, allowed_marker, V3_KEY) == ()
|
|
assert len(rig.provider_calls(allowed_marker, rig.provider_drain())) == 1
|
|
|
|
blocked_marker: Final = rig.marker()
|
|
blocked: Final = _chat(rig, f"{STRAY_V3_BLOCK_MARK} {blocked_marker}", guardrails=["straiker-v3-as-v1"])
|
|
assert blocked.status_code == 400, blocked.text
|
|
assert blocked.json()["error"]["message"] == BLOCK_MESSAGE, blocked.text
|
|
assert len(_v3_request_calls(rig, blocked_marker, agent=None)) == 1, rig.sink_calls(blocked_marker)
|
|
assert _v1_calls(rig, blocked_marker, V3_KEY) == ()
|
|
assert rig.provider_calls(blocked_marker, rig.provider_drain()) == ()
|
|
|
|
|
|
def test_stray_api_version_with_v3_key_still_enforces_on_v3(rig: Rig) -> None:
|
|
allowed_marker: Final = rig.marker()
|
|
allowed: Final = _chat(rig, "stray version " + allowed_marker, guardrails=["straiker-v3-stray-version"])
|
|
assert allowed.status_code == 200, allowed.text
|
|
assert len(_v3_request_calls(rig, allowed_marker, agent=None)) == 1
|
|
assert len(rig.provider_calls(allowed_marker, rig.provider_drain())) == 1
|
|
|
|
blocked_marker: Final = rig.marker()
|
|
blocked: Final = _chat(rig, f"{STRAY_V3_BLOCK_MARK} {blocked_marker}", guardrails=["straiker-v3-stray-version"])
|
|
assert blocked.status_code == 400, blocked.text
|
|
assert blocked.json()["error"]["message"] == BLOCK_MESSAGE, blocked.text
|
|
assert len(_v3_request_calls(rig, blocked_marker, agent=None)) == 1
|
|
assert rig.provider_calls(blocked_marker, rig.provider_drain()) == ()
|
|
|
|
|
|
def test_empty_api_version_with_v1_key_still_enforces_on_v1(rig: Rig) -> None:
|
|
allowed_marker: Final = rig.marker()
|
|
allowed: Final = _chat(rig, "empty version " + allowed_marker, guardrails=["straiker-v1-empty-version"])
|
|
assert allowed.status_code == 200, allowed.text
|
|
assert len(_v1_calls(rig, allowed_marker, V1_KEY)) == 1
|
|
assert len(rig.provider_calls(allowed_marker, rig.provider_drain())) == 1
|
|
|
|
blocked_marker: Final = rig.marker()
|
|
blocked: Final = _chat(rig, f"{V1_BLOCK_MARK} {blocked_marker}", guardrails=["straiker-v1-empty-version"])
|
|
assert blocked.status_code == 400, blocked.text
|
|
assert blocked.json()["error"]["message"] == BLOCK_MESSAGE, blocked.text
|
|
assert len(_v1_calls(rig, blocked_marker, V1_KEY)) == 1
|
|
assert rig.provider_calls(blocked_marker, rig.provider_drain()) == ()
|
|
|
|
|
|
# E: configured client and format_hint ride as headers; request header for agent fills in when YAML has none
|
|
def test_v3_client_and_format_hint_headers_and_request_agent_header(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(
|
|
rig, "hint " + marker, guardrails=["straiker-v3-hint"], headers={"x-s6r-agent": "caller-agent"}
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
hinted: Final = tuple(s for s in _v3_request_calls(rig, marker, agent="caller-agent"))
|
|
assert len(hinted) == 1, rig.sink_calls(marker)
|
|
sent: Final = hinted[0]
|
|
assert sent.headers["x-s6r-client"] == "named-client"
|
|
assert sent.headers["x-s6r-format"] == "anthropic.messages"
|
|
assert sent.headers["x-s6r-agent"] == "caller-agent"
|
|
|
|
|
|
# E: identity precedence: the key's user email wins over an end user in the body
|
|
def test_v3_user_prefers_key_email_over_body_user(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
with rig.proxy.scenario() as scenario:
|
|
user: Final = scenario.user(user_email=f"{marker}@example.test")
|
|
key: Final = scenario.key(user_id=user)
|
|
response: Final = _chat(rig, "identity " + marker, key=key, user="body-user-" + marker)
|
|
assert response.status_code == 200, response.text
|
|
sent: Final = _v3_request_calls(rig, marker)[0].body
|
|
meta: Final = object_value(sent["original"])
|
|
assert object_value(object_value(object_value(meta["processed"])["Meta"]))["user"] == f"{marker}@example.test"
|
|
assert object_value(sent["metadata"])["user_api_key_user_email"] == f"{marker}@example.test"
|
|
assert sent["user"] == "body-user-" + marker
|
|
|
|
|
|
# E: logging_only observes the turn but never blocks
|
|
def test_v3_logging_only_observes_block_verdict_without_blocking(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{LOG_BLOCK_MARK} {marker}")
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["id"].startswith("chatcmpl-" + marker), response.json()["id"]
|
|
calls: Final = eventually(lambda: _v3_request_calls(rig, marker, agent=LOG_AGENT), lambda c: len(c) >= 1)
|
|
assert calls[0].headers["x-s6r-agent"] == LOG_AGENT
|
|
assert len(rig.provider_calls(marker, rig.provider_drain())) == 1
|
|
row: Final = rig.spend_row(response.json()["id"])
|
|
assert row["request_id"] == response.json()["id"]
|
|
|
|
|
|
# E: the same identical allowed request three times yields three detect calls and three spend rows
|
|
def test_v3_repeated_allowed_request_is_scored_and_logged_each_time(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
responses: Final = tuple(_chat(rig, "repeat " + marker) for _ in range(3))
|
|
assert all(r.status_code == 200 for r in responses), [r.text for r in responses]
|
|
ids: Final = {r.json()["id"] for r in responses}
|
|
assert len(ids) == 3 and all(i.startswith("chatcmpl-" + marker) for i in ids), ids
|
|
assert len(_v3_request_calls(rig, marker)) == 3
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', ("chatcmpl-" + marker + "%",)
|
|
),
|
|
lambda values: len(values) == 3,
|
|
seconds=70,
|
|
)
|
|
assert {str(r["request_id"]) for r in rows} == ids
|
|
|
|
|
|
# S1: platform answers 500: fail closed with the reason in the body, no provider call
|
|
def test_v3_sink_500_fails_closed_with_reason(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{SINK_500_MARK} {marker}")
|
|
assert response.status_code == 400, response.text
|
|
assert "Straiker detection unavailable" in response.json()["error"]["message"], response.text
|
|
assert rig.provider_calls(marker, rig.provider_drain()) == ()
|
|
|
|
|
|
# S1a: the v1 webhook route fails the same way when the platform answers 500
|
|
def test_v1_sink_500_fails_closed_with_reason(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{V1_500_MARK} {marker}", guardrails=["straiker-v1"])
|
|
assert response.status_code == 400, response.text
|
|
assert "Straiker detection unavailable" in response.json()["error"]["message"], response.text
|
|
assert len(_v1_calls(rig, marker, V1_KEY)) == 1
|
|
assert rig.provider_calls(marker, rig.provider_drain()) == ()
|
|
|
|
|
|
# S1b: fail_on_error false lets the request through on a 500
|
|
def test_v3_fail_open_guardrail_passes_on_sink_500(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{OPEN_500_MARK} {marker}", guardrails=["straiker-v3-open"])
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["id"].startswith("chatcmpl-" + marker), response.json()["id"]
|
|
assert len(_v3_request_calls(rig, marker, agent=OPEN_AGENT)) == 1
|
|
|
|
|
|
# S2: platform rejects the key: 401 is not retried and fails closed
|
|
def test_v3_sink_401_fails_closed_once(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{SINK_401_MARK} {marker}")
|
|
assert response.status_code == 400, response.text
|
|
assert "401" in response.json()["error"]["message"], response.text
|
|
assert len(_v3_request_calls(rig, marker)) == 1
|
|
|
|
|
|
# S3: platform answers non JSON: fail closed, caller sees the parse failure
|
|
def test_v3_sink_garbage_fails_closed(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, f"{SINK_GARBAGE_MARK} {marker}")
|
|
assert response.status_code == 400, response.text
|
|
assert "Straiker detection unavailable" in response.json()["error"]["message"]
|
|
|
|
|
|
# S4: unauthenticated request never reaches the platform
|
|
def test_unauthenticated_request_does_not_reach_platform(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = _chat(rig, "anon " + marker, key="sk-not-a-real-key")
|
|
assert response.status_code == 401, response.text
|
|
assert rig.sink_calls(marker) == ()
|
|
|
|
|
|
# S5: unknown model: the guardrail still runs, then the router error reaches the caller
|
|
def test_unknown_model_error_reaches_caller_after_detect(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
response: Final = rig.proxy.client.post(
|
|
"/v1/chat/completions",
|
|
json={"model": "no-such-model-" + marker, "messages": _messages("unknown " + marker)},
|
|
headers={"Authorization": f"Bearer {rig.proxy.key}"},
|
|
)
|
|
assert response.status_code in (400, 401, 404), response.text
|
|
assert "no-such-model-" + marker in response.text
|
|
assert len(_v3_request_calls(rig, marker)) == 1, rig.sink_calls(marker)
|
|
assert rig.provider_calls(marker, rig.provider_drain()) == ()
|
|
|
|
|
|
# S6: odd shapes in the routing header and a 5 KB prompt are relayed verbatim, not crashed on
|
|
def test_v3_oversized_prompt_and_odd_header_values_are_relayed(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
big: Final = "x" * 5000 + " " + marker
|
|
response: Final = _chat(rig, big, headers={"x-claude-code-session-id": "", "x-s6r-agent": "1"})
|
|
assert response.status_code == 200, response.text
|
|
sent: Final = _v3_request_calls(rig, marker)[0]
|
|
assert sent.body["messages"] == [{"role": "user", "content": big}]
|
|
assert sent.headers["x-s6r-agent"] == "audit-agent"
|
|
assert "x-claude-code-session-id" not in sent.headers
|
|
assert sent.body.get("session_id", "").startswith("litellm-")
|
|
|
|
|
|
# S7: a guardrail with a malformed format_hint is rejected at /guardrails/apply_guardrail time, not at boot
|
|
def test_malformed_format_hint_config_is_rejected_by_guardrail_management(rig: Rig) -> None:
|
|
response: Final = rig.proxy.client.post(
|
|
"/guardrails",
|
|
json={
|
|
"guardrail": {
|
|
"guardrail_name": "straiker-bad-" + uuid.uuid4().hex,
|
|
"litellm_params": {
|
|
"guardrail": "straiker",
|
|
"mode": "pre_call",
|
|
"api_key": V3_KEY,
|
|
"api_base": rig.sink.url,
|
|
"format_hint": "bogus",
|
|
},
|
|
}
|
|
},
|
|
headers={"Authorization": f"Bearer {rig.proxy.key}"},
|
|
)
|
|
assert response.status_code in (400, 422, 500), response.text
|
|
assert "format_hint" in response.text or "bogus" in response.text, response.text
|
|
healthy: Final = _chat(rig, "still-fine " + uuid.uuid4().hex)
|
|
assert healthy.status_code == 200, healthy.text
|
|
|
|
|
|
# S8: /key/health reports the key without touching the platform
|
|
def test_key_health_does_not_call_platform(rig: Rig) -> None:
|
|
marker: Final = rig.marker()
|
|
with rig.proxy.scenario() as scenario:
|
|
key: Final = scenario.key(key_alias="health-" + marker)
|
|
response: Final = rig.proxy.client.post("/key/health", headers={"Authorization": f"Bearer {key}"})
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["key"] == "healthy"
|
|
assert rig.sink_calls(marker) == ()
|
|
|
|
|
|
# C1: 30 request mixed burst while the platform sink is down mid burst, then recovers; every allowed id lands once
|
|
def test_burst_with_platform_outage_recovers_without_duplicate_spend(rig: Rig) -> None:
|
|
burst: Final = 30
|
|
markers: Final = tuple(rig.marker() for _ in range(burst))
|
|
down: Final = threading.Event()
|
|
up: Final = threading.Event()
|
|
|
|
def call(index: int) -> tuple[int, int, str]:
|
|
if index == 8:
|
|
rig.sink.stop()
|
|
down.set()
|
|
if index == 20:
|
|
assert down.wait(10)
|
|
rig.sink.start()
|
|
up.set()
|
|
marker: Final = markers[index]
|
|
if index % 3 == 0:
|
|
response: Final = rig.proxy.client.post(
|
|
"/v1/messages",
|
|
json={
|
|
"model": rig.anthropic_model,
|
|
"max_tokens": 8,
|
|
"messages": [{"role": "user", "content": "burst " + marker}],
|
|
},
|
|
headers={"Authorization": f"Bearer {rig.proxy.key}"},
|
|
)
|
|
return index, response.status_code, response.text
|
|
streaming: Final = index % 2 == 1
|
|
response = _chat(rig, "burst " + marker, stream=streaming)
|
|
return index, response.status_code, response.text
|
|
|
|
with ThreadPoolExecutor(max_workers=6) as pool:
|
|
results: Final = sorted(pool.map(call, range(burst)))
|
|
assert up.is_set()
|
|
statuses: Final = {index: status for index, status, _ in results}
|
|
assert all(status in (200, 400) for status in statuses.values()), results
|
|
failed: Final = tuple(index for index, status, text in results if status == 400)
|
|
assert failed, "the outage must be visible to at least one caller"
|
|
assert all("Straiker detection unavailable" in text for index, status, text in results if status == 400), results
|
|
for index, status, text in results:
|
|
if status != 200 or (index % 3 != 0 and index % 2 == 1):
|
|
continue
|
|
marker = markers[index]
|
|
expected: Final = ("msg_" if index % 3 == 0 else "chatcmpl-") + marker + "%"
|
|
rows: Final = eventually(
|
|
lambda like=expected: read_rows(
|
|
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', (like,)
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert len(rows) == 1, rows
|
|
provider_seen: Final = rig.provider_drain()
|
|
for index, status, _ in results:
|
|
if status == 400:
|
|
assert rig.provider_calls(markers[index], provider_seen) == (), (
|
|
"a failed-closed turn must not reach the provider"
|
|
)
|
|
recovered: Final = _chat(rig, "after-outage " + rig.marker())
|
|
assert recovered.status_code == 200, recovered.text
|
|
|
|
|
|
# C2: one proxy worker is killed during a burst; the other keeps serving and detect still runs for each call
|
|
def test_an_unreaped_worker_does_not_hide_the_surviving_proxy_worker(rig: Rig) -> None:
|
|
workers: Final = tuple(sorted(eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2)))
|
|
victim: Final = workers[0]
|
|
survivor: Final = workers[1]
|
|
marker: Final = rig.marker()
|
|
with rig.owned.replacing_worker(victim) as kill, rig.owned.suspended(rig.owned.process.pid):
|
|
kill()
|
|
assert eventually(psutil.Process(victim).status, lambda status: status == psutil.STATUS_ZOMBIE) == (
|
|
psutil.STATUS_ZOMBIE
|
|
)
|
|
assert rig.owned.worker_pids() == frozenset((survivor,))
|
|
with httpx.Client(base_url=rig._base(), timeout=15, trust_env=False) as client:
|
|
response: Final = client.post(
|
|
"/v1/chat/completions",
|
|
json={"model": rig.chat_model, "messages": _messages(marker)},
|
|
headers={"Authorization": f"Bearer {rig.proxy.key}"},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert len(_v3_request_calls(rig, marker)) == 1
|
|
assert len(rig.owned.worker_pids() - frozenset((victim,))) == 2
|
|
|
|
|
|
def test_burst_survives_one_worker_kill(rig: Rig) -> None:
|
|
workers: Final = tuple(sorted(eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2)))
|
|
markers: Final = tuple(rig.marker() for _ in range(24))
|
|
|
|
def fresh_chat(text: str) -> tuple[int, str]:
|
|
with httpx.Client(base_url=rig._base(), timeout=15, trust_env=False) as fresh:
|
|
try:
|
|
response: Final = fresh.post(
|
|
"/v1/chat/completions",
|
|
json={"model": rig.chat_model, "messages": _messages(text)},
|
|
headers={"Authorization": f"Bearer {rig.proxy.key}"},
|
|
)
|
|
except httpx.TransportError as error:
|
|
return 0, repr(error)
|
|
return response.status_code, response.text
|
|
|
|
with rig.owned.replacing_worker(workers[0]) as kill:
|
|
|
|
def call(index: int) -> tuple[int, str]:
|
|
if index == 6:
|
|
kill()
|
|
return fresh_chat("kill " + markers[index])
|
|
|
|
with ThreadPoolExecutor(max_workers=4) as pool:
|
|
results: Final = tuple(pool.map(call, range(24)))
|
|
ok: Final = tuple(i for i, (status, _) in enumerate(results) if status == 200)
|
|
assert len(ok) >= 20, results
|
|
for index in ok:
|
|
assert len(_v3_request_calls(rig, markers[index])) >= 1, markers[index]
|
|
after: Final = fresh_chat("after-kill " + rig.marker())
|
|
assert after[0] == 200, after
|
|
|
|
|
|
# C3: proxy restart between a blocked turn and its replay: the memory is per process and empties, so Straiker is asked again
|
|
def test_proxy_restart_forgets_blocked_turns_and_asks_platform_again(tmp_path: Path, rig: Rig) -> None:
|
|
with gateway_from_environment() as gateway:
|
|
config: Final = _rig_config(rig.sink.url, tmp_path)
|
|
marker: Final = rig.marker()
|
|
session: Final = {"x-claude-code-session-id": "restart-" + marker}
|
|
body: Final = {"model": rig.chat_model, "messages": _messages(f"{BLOCK_MARK} {marker}")}
|
|
with owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as first:
|
|
blocked: Final = first.client.post(
|
|
"/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {first.key}", **session}
|
|
)
|
|
assert blocked.status_code == 400, blocked.text
|
|
replayed: Final = first.client.post(
|
|
"/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {first.key}", **session}
|
|
)
|
|
assert replayed.status_code == 400, replayed.text
|
|
assert len(_v3_request_calls(rig, marker)) == 1, "one worker replays from memory"
|
|
with owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as second:
|
|
again: Final = second.client.post(
|
|
"/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {second.key}", **session}
|
|
)
|
|
assert again.status_code == 400, again.text
|
|
assert len(_v3_request_calls(rig, marker)) == 2, "a restarted process has no memory and asks once more"
|