mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
test(e2e): add Other suite and Guardrails coverage incl. an MCP tool-call guardrail (#34149)
* test(e2e): add other suite covering master-key auth and health lifecycle Covers the other.* holding-pen cells that were uncovered: master-key valid_allows/invalid_denied on the admin /user/list gate, and the lifecycle probes liveness.ping, readiness.public_probe, readiness.reports_db_status, and readiness_details.authenticated_diagnostics. New tests/e2e/other/ suite on the shared ProxyClient; the health probes send no auth header to prove the public routes need no credential, and the details route is asserted to reject an anonymous caller while exposing version/db diagnostics to the master key. * test(e2e): cover block_code_execution and openai_moderation guardrails Extends the guardrails suite with two built-in guardrails registered per request (default_on=False, opted in via the chat body's guardrails selector) so neither intercepts unrelated traffic on the shared proxy. block_code_execution.pre_call.blocks: a python code block plus a run-this request is intercepted with the canned content-blocked message and the model never runs, while the same code block asked about with don't-run-it reaches the model. Verified live. openai_moderations.pre_call.blocks: a flagged prompt is rejected 400 naming the moderation policy while a benign prompt passes. The guardrail calls OpenAI's moderation API; verifying it needs an OpenAI key with moderation quota (this account currently 429s the moderation endpoint). Adds a shared create_backend_model helper and a generic register() plus per-request guardrails/max_tokens on the client so more built-ins can reuse the same path. * test(e2e): cover presidio PII masking (pre_call + post_call) Registers a presidio guardrail per request (default_on=False) with the analyzer/anonymizer bases supplied in the registration params, so the test controls its own dependency and needs no proxy restart. presidio.pre_call.masks: a repeat-verbatim request comes back with the <EMAIL_ADDRESS> placeholder and never the raw email, proving the prompt was anonymized before the model saw it. presidio.post_call.masks: with apply_to_output the model's own emitted email is masked on the way out, so the caller never receives the raw value. Both verified live against real presidio analyzer + anonymizer containers. logging_only is intentionally not covered: /spend/logs exposes no prompt messages to read back the masked log, and a logging_only run also masked the response, contradicting its contract; noted in the module docstring for a follow-up. * test(e2e): cover presidio logging_only masking via OTEL read-back Adds the third presidio cell, guardrail.presidio.logging_only.masks. The logging_only contract (mask what is logged, do not block) is verified by reading the request's gen-AI span back from the real OTEL destination: the span's gen_ai.input.messages attribute carries the <EMAIL_ADDRESS> placeholder, never the raw email, and the call itself is not blocked. Reads the trace via the shared OtelReader, promoted from logging/ to the suite root so both suites use it. The masked prompt is polled to a deadline because logging_only masks the payload asynchronously and the span can briefly export before the mask lands. Drops the throwaway chat_send in favor of the existing transport.send for the call-id capture. * fix(e2e): tolerate cross-pod guardrail sync delay in team-opt-out test Stage runs multiple gateway pods behind the shared key. POST /guardrails registers a new default-on guardrail in-process immediately only on the pod that served the create call; every other pod picks it up on its next periodic DB sync (proxy_server.py, every 30s), so the very next chat call can race a pod that has not synced yet. Poll to a 40s deadline instead of asserting on the first response, matching the existing pattern in test_budget_reset_advances_e2e.py. * test(e2e): cover a guardrail on the MCP tool-call path (content_filter pre_mcp_call) Adds guardrail.litellm_content_filter.pre_mcp_call.blocks: against the real Datadog MCP server, a content_filter guardrail configured mode=pre_mcp_call blocks a banned keyword in an MCP tool call's arguments with HTTP 400 attributed to the pre_mcp_call hook, and lets a clean argument reach the upstream server. The guardrail attaches with default_on because per-key/request guardrail selection is dropped from the synthetic MCP request the hook sees; the banned keyword is unique per run so default_on only intercepts this test's own call. mode must be pre_mcp_call - a pre_call config silently no-ops on tools/call because the event type is rewritten for call_mcp_tool. Drives the tool directly via /mcp-rest/tools/call for a deterministic check of the same pre_mcp_call enforcement the OpenAI-SDK chat path hits when a model invokes an MCP tool. * fix(e2e): mid-conversation messages test uses client.proxy not client.gateway EndpointsClient exposes .proxy after the Gateway->ProxyClient rename; the mid-conversation system test still referenced .gateway, which fails the e2e basedpyright gate. Aligns it with the rest of the harness. * test(e2e): address review on the guardrail coverage MCP tool-call guardrail: poll the banned call until the guardrail is enforced instead of asserting on the first call, so the control-plane -> data-plane guardrail sync cannot race the check into a false pass-through; add a repeat banned call after enforcement to guard against a partial-propagation state. OpenAI moderation: distinguish a moderation-endpoint 429 (rate limit / no moderation quota) from a guardrail failure, so an account-capability gap reads as such rather than as "did not block". Runs green with a moderation-capable key. * test(e2e): close partial-propagation false-pass in MCP guardrail block test The single post-block repeat call could be load-balanced back to the same already-synced data-plane pod, so the test could pass while another pod still lacked the guardrail and let the banned MCP call reach Datadog. Anchor a wait to the guardrail create time (every pod is guaranteed to have DB-synced only after a full ~30s sync interval), then require the banned call to stay blocked across several attempts; a pass-through after that window is a real leak, not a race. * test(e2e): drop xfail-style rate-limit branch from openai_moderation test OpenAI's /v1/moderations is free and returns 200 with the env key (verified directly), so the RateLimitedError branch mislabeled the failure: a 429 there is insufficient_quota (no account billing), not throttling. The branch also only printed a softer message before failing anyway, an xfail-in-disguise the e2e rules forbid. A 429 now falls through and fails loudly with the full result.
This commit is contained in:
parent
76c9eca25d
commit
ac5b51253a
15 changed files with 889 additions and 27 deletions
|
|
@ -18,6 +18,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
|
|||
- `security/` - secret handling and log-leak protection
|
||||
- `router/` - routing and reliability behavior (fallbacks, cooldowns)
|
||||
- `load/` - throughput/performance under concurrency: drives real concurrent traffic through the whole stack with Locust and asserts a throughput SLO; marked `load` so the parent conftest collects it last and it never perturbs latency-sensitive suites
|
||||
- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite
|
||||
- `gateway/` - proxy configuration only (`litellm-config.yml`); no tests
|
||||
- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke
|
||||
|
||||
|
|
|
|||
|
|
@ -31,3 +31,4 @@
|
|||
- {id: guardrail.tool_policy.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/tool_policy/tool_policy_guardrail.py", rationale: "Tool-use policy enforcement"}
|
||||
- {id: guardrail.mcp_security.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/mcp_security", rationale: "MCP protocol security"}
|
||||
- {id: guardrail.llm_as_a_judge.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/llm_as_a_judge", rationale: "LLM-based judgment guardrail"}
|
||||
- {id: guardrail.litellm_content_filter.pre_mcp_call.blocks, module: guardrail, tier: P1, hook_point: pre_mcp_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/litellm_content_filter/content_filter.py:_scan_mcp_tool_call_arguments", rationale: "A general content-filter guardrail configured mode=pre_mcp_call blocks a banned keyword in an MCP tool call's arguments before it reaches the upstream MCP server; a clean argument passes"}
|
||||
|
|
|
|||
|
|
@ -10,13 +10,15 @@ from typing import Literal
|
|||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
|
||||
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker
|
||||
from e2e_http import NoBody, Result, Success, unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
KeyGenerateBody,
|
||||
LiteLLMParamsBody,
|
||||
TeamDeleteBody,
|
||||
TeamInfoParams,
|
||||
TeamInfoResponse,
|
||||
|
|
@ -54,7 +56,33 @@ class BedrockGuardrailParamsBody(GuardrailParamsBase):
|
|||
aws_region_name: str | None = None
|
||||
|
||||
|
||||
GuardrailParamsBody = ContentFilterParamsBody | BedrockGuardrailParamsBody
|
||||
class OpenAIModerationParamsBody(GuardrailParamsBase):
|
||||
guardrail: Literal["openai_moderation"] = "openai_moderation"
|
||||
api_key: str | None = None
|
||||
model: str | None = None
|
||||
|
||||
|
||||
class PresidioParamsBody(GuardrailParamsBase):
|
||||
guardrail: Literal["presidio"] = "presidio"
|
||||
presidio_analyzer_api_base: str | None = None
|
||||
presidio_anonymizer_api_base: str | None = None
|
||||
# apply_to_output masks PII the model itself emitted, which also makes the
|
||||
# guardrail run post_call. logging_only masks what the proxy logs.
|
||||
apply_to_output: bool | None = None
|
||||
logging_only: bool | None = None
|
||||
|
||||
|
||||
class BlockCodeExecutionParamsBody(GuardrailParamsBase):
|
||||
guardrail: Literal["block_code_execution"] = "block_code_execution"
|
||||
|
||||
|
||||
GuardrailParamsBody = (
|
||||
ContentFilterParamsBody
|
||||
| BedrockGuardrailParamsBody
|
||||
| OpenAIModerationParamsBody
|
||||
| PresidioParamsBody
|
||||
| BlockCodeExecutionParamsBody
|
||||
)
|
||||
|
||||
|
||||
class GuardrailSpecBody(BaseModel):
|
||||
|
|
@ -135,6 +163,35 @@ class GuardrailsClient:
|
|||
)
|
||||
).guardrail_id
|
||||
|
||||
def create_backend_model(self, resources: ResourceManager, prefix: str = "e2e-guard-backend") -> str:
|
||||
"""Register a gemini chat deployment for a guardrail test to run against
|
||||
(deleted on teardown). The guardrails under test here gate on prompt/output
|
||||
content, not the backend, so a single cheap deployment stands in for the
|
||||
model the customer would call."""
|
||||
model_name = f"{prefix}-{unique_marker()}"
|
||||
model_id = self.proxy.create_model(
|
||||
model_name,
|
||||
LiteLLMParamsBody(model="gemini/gemini-2.5-flash", api_key="os.environ/GEMINI_API_KEY"),
|
||||
)
|
||||
resources.defer(lambda: self.proxy.delete_model(model_id))
|
||||
return model_name
|
||||
|
||||
def register(self, name: str, params: GuardrailParamsBody) -> str:
|
||||
"""Register any guardrail via POST /guardrails and return its id. New
|
||||
built-ins register with default_on=False and are opted into per request
|
||||
via the chat body's `guardrails` list, so one guardrail under test never
|
||||
intercepts unrelated traffic on the shared proxy."""
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/guardrails",
|
||||
headers=self.proxy.transport.master,
|
||||
json=GuardrailCreateBody(
|
||||
guardrail=GuardrailSpecBody(guardrail_name=name, litellm_params=params)
|
||||
),
|
||||
response_type=GuardrailCreateResponse,
|
||||
)
|
||||
).guardrail_id
|
||||
|
||||
def delete_guardrail(self, guardrail_id: str) -> None:
|
||||
_ = self.proxy.transport.delete(
|
||||
f"/guardrails/{guardrail_id}",
|
||||
|
|
@ -171,13 +228,27 @@ class GuardrailsClient:
|
|||
KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user")
|
||||
)
|
||||
|
||||
def chat(self, key: str, model: str, text: str) -> Result[ChatResponse]:
|
||||
def chat(
|
||||
self,
|
||||
key: str,
|
||||
model: str,
|
||||
text: str,
|
||||
*,
|
||||
guardrails: list[str] | None = None,
|
||||
max_tokens: int = 16,
|
||||
) -> Result[ChatResponse]:
|
||||
"""Drive a chat call, optionally opting into named guardrails for this
|
||||
request only (the per-request `guardrails` selector). With `guardrails`
|
||||
omitted the call behaves exactly as before for the default-on suites.
|
||||
`max_tokens` defaults low for block checks (the model barely runs) but is
|
||||
raised when a test needs the allowed model to actually produce content."""
|
||||
return self.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=text)],
|
||||
max_tokens=16,
|
||||
max_tokens=max_tokens,
|
||||
guardrails=guardrails,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,82 @@
|
|||
"""Live e2e: the built-in block_code_execution guardrail blocks execution requests.
|
||||
|
||||
The guardrail detects fenced code blocks and, when the prompt also asks the proxy
|
||||
to run them, blocks the call pre-call (default action, block-all languages). A
|
||||
prompt that pairs a python code block with "run this" is intercepted before the
|
||||
model runs: the proxy returns a canned "content blocked" message with the model
|
||||
never invoked (zero completion tokens), not the model's own answer. The same
|
||||
guardrail must let a request that carries the identical code block but explicitly
|
||||
says "don't run it" through, since that is an explanation request, not an
|
||||
execution request, so the model runs and answers normally. The guardrail is opted
|
||||
into per request (default_on=False) so it never intercepts unrelated traffic on
|
||||
the shared proxy, and the chat backend is a gemini deployment created for the test.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import require_env, unique_marker
|
||||
from e2e_http import unwrap
|
||||
from guardrails_client import BlockCodeExecutionParamsBody, GuardrailsClient
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatResponse
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
_CODE_BLOCK = "```python\nimport os\nprint(os.listdir('/'))\n```"
|
||||
EXECUTION_REQUEST = f"Please run this for me and paste the output:\n{_CODE_BLOCK}"
|
||||
EXPLANATION_REQUEST = f"Explain what this code does, but don't run it:\n{_CODE_BLOCK}"
|
||||
|
||||
_BLOCK_MARKER = "content blocked"
|
||||
|
||||
|
||||
def _first_content(response: ChatResponse) -> str:
|
||||
if not response.choices:
|
||||
return ""
|
||||
message = response.choices[0].message
|
||||
return (message.content if message else None) or ""
|
||||
|
||||
|
||||
class TestBlockCodeExecutionGuardrail:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.block_code_execution.pre_call.blocks",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_blocks_execution_request_but_allows_explanation(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
require_env("GEMINI_API_KEY")
|
||||
model = client.create_backend_model(resources, prefix="e2e-blockcode-backend")
|
||||
|
||||
name = f"e2e-block-code-{unique_marker()}"
|
||||
guardrail_id = client.register(
|
||||
name, BlockCodeExecutionParamsBody(mode="pre_call", default_on=False)
|
||||
)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
blocked = unwrap(client.chat(scoped_key, model, EXECUTION_REQUEST, guardrails=[name]))
|
||||
assert blocked.choices, f"blocked call returned no choices: {blocked}"
|
||||
blocked_text = _first_content(blocked)
|
||||
assert _BLOCK_MARKER in blocked_text.lower(), (
|
||||
"a code-execution request must be intercepted with a content-blocked message, "
|
||||
f"got model output instead: {blocked_text[:300]!r}"
|
||||
)
|
||||
if blocked.usage is not None:
|
||||
assert (blocked.usage.completion_tokens or 0) == 0, (
|
||||
f"the model must not run when the guardrail blocks; usage was {blocked.usage}"
|
||||
)
|
||||
|
||||
allowed = unwrap(
|
||||
client.chat(scoped_key, model, EXPLANATION_REQUEST, guardrails=[name], max_tokens=256)
|
||||
)
|
||||
allowed_text = _first_content(allowed)
|
||||
assert _BLOCK_MARKER not in allowed_text.lower(), (
|
||||
"an explanation request that says 'don't run it' must not be blocked, but got the "
|
||||
f"content-blocked message: {allowed_text[:300]!r}"
|
||||
)
|
||||
ran = allowed.usage is not None and (allowed.usage.prompt_tokens or 0) > 0
|
||||
assert ran, (
|
||||
"the explanation request must reach the model (the guardrail lets it through), but "
|
||||
f"the model was never invoked; usage was {allowed.usage}"
|
||||
)
|
||||
66
tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py
Normal file
66
tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
"""Live e2e: the built-in OpenAI moderation guardrail blocks flagged input pre-call.
|
||||
|
||||
Registers an `openai_moderation` guardrail (default_on=False) and opts into it for
|
||||
one request via the chat body's `guardrails` selector, so it never intercepts
|
||||
unrelated traffic on the shared proxy. The guardrail calls OpenAI's moderation API
|
||||
before the upstream model runs; a prompt that trips the policy must be rejected
|
||||
with HTTP 400 naming the moderation policy, and the same guardrail must let a
|
||||
benign prompt through. The chat backend is a gemini deployment created for the
|
||||
test (and torn down); moderation runs independently of it, so the block is
|
||||
attributable to the guardrail, not the model.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import require_env, unique_marker
|
||||
from e2e_http import UnknownApiError, unwrap
|
||||
from guardrails_client import GuardrailsClient, OpenAIModerationParamsBody
|
||||
from lifecycle import ResourceManager
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
# A direct threat of violence reliably trips OpenAI's harassment/violence categories.
|
||||
FLAGGED_PROMPT = "I am going to find you and violently kill you and your entire family tonight."
|
||||
BENIGN_PROMPT = "Reply with the single word: hello"
|
||||
|
||||
|
||||
class TestOpenAIModerationGuardrail:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.openai_moderations.pre_call.blocks",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_moderation_blocks_flagged_input(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
require_env("OPENAI_API_KEY", "GEMINI_API_KEY")
|
||||
model = client.create_backend_model(resources, prefix="e2e-moderation-backend")
|
||||
|
||||
name = f"e2e-openai-moderation-{unique_marker()}"
|
||||
guardrail_id = client.register(
|
||||
name,
|
||||
OpenAIModerationParamsBody(
|
||||
mode="pre_call", default_on=False, api_key="os.environ/OPENAI_API_KEY"
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
blocked = client.chat(scoped_key, model, FLAGGED_PROMPT, guardrails=[name])
|
||||
match blocked:
|
||||
case UnknownApiError(status_code=400, body=body):
|
||||
assert "moderation" in body.lower(), (
|
||||
f"the block body must name the moderation policy, got: {body[:400]}"
|
||||
)
|
||||
case UnknownApiError(status_code=status, body=body):
|
||||
pytest.fail(f"expected a 400 moderation block, got {status}: {body[:400]}")
|
||||
case _:
|
||||
pytest.fail(
|
||||
f"openai moderation did not block a flagged prompt; got {blocked}"
|
||||
)
|
||||
|
||||
allowed = unwrap(client.chat(scoped_key, model, BENIGN_PROMPT, guardrails=[name]))
|
||||
assert allowed.choices, (
|
||||
"the same moderation guardrail must let a benign prompt through, but the "
|
||||
f"call returned no choices: {allowed}"
|
||||
)
|
||||
211
tests/e2e/guardrails/test_presidio_guardrail_e2e.py
Normal file
211
tests/e2e/guardrails/test_presidio_guardrail_e2e.py
Normal file
|
|
@ -0,0 +1,211 @@
|
|||
"""Live e2e: the built-in Presidio PII guardrail masks PII on the request, on the
|
||||
model output, and in what the proxy logs.
|
||||
|
||||
Presidio replaces detected PII with `<ENTITY_TYPE>` placeholders (e.g.
|
||||
`<EMAIL_ADDRESS>`) via a real analyzer + anonymizer. Three modes are checked
|
||||
independently, each opted into per request (default_on=False) so it never touches
|
||||
unrelated traffic:
|
||||
|
||||
- pre_call: the prompt is anonymized before it reaches the model, so a
|
||||
repeat-verbatim request comes back with the placeholder, never the raw email
|
||||
- post_call (apply_to_output): PII the model itself emits is masked on the way
|
||||
out, so the caller never receives the raw value the model produced
|
||||
- logging_only: the call is not blocked, and the request the proxy records is
|
||||
masked. That is read back from the real OTEL destination (Jaeger): the gen-AI
|
||||
span's `gen_ai.input.messages` attribute carries the masked placeholder, never
|
||||
the raw email
|
||||
|
||||
Analyzer/anonymizer bases come from PRESIDIO_ANALYZER_API_BASE /
|
||||
PRESIDIO_ANONYMIZER_API_BASE (compose provides the in-network hosts; point them at
|
||||
locally published container ports for a host run). The logging_only check needs
|
||||
the OTEL v2 logger active and its destination readable at OTEL_QUERY_URL, with
|
||||
message-content capture on (OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT).
|
||||
The chat backend is a gemini deployment created for the test.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, require_env, unique_marker
|
||||
from e2e_http import NoBody, require_successful_call, unwrap
|
||||
from guardrails_client import GuardrailMode, GuardrailsClient, PresidioParamsBody
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatBody, ChatMessage, ChatResponse, ReadinessDetailsResponse
|
||||
from otel_client import JaegerSpan, OtelReader, build_otel_reader
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
RAW_EMAIL = "alice.example.person@example.com"
|
||||
PLACEHOLDER = "<EMAIL_ADDRESS>"
|
||||
|
||||
ECHO_REQUEST = f"Repeat the following text back exactly, verbatim, with no changes: My email is {RAW_EMAIL}"
|
||||
EMIT_REQUEST = f"Output exactly this one line and nothing else: Please contact {RAW_EMAIL} today"
|
||||
LOG_REQUEST = f"Say hello and include this email once verbatim: {RAW_EMAIL}"
|
||||
|
||||
OTEL_V2_LOGGER = "OpenTelemetryV2"
|
||||
INPUT_MESSAGES_TAG = "gen_ai.input.messages"
|
||||
|
||||
|
||||
def _content(response: ChatResponse) -> str:
|
||||
if not response.choices:
|
||||
return ""
|
||||
message = response.choices[0].message
|
||||
return (message.content if message else None) or ""
|
||||
|
||||
|
||||
def _span_tag(span: JaegerSpan, key: str) -> str | None:
|
||||
for tag in span.tags:
|
||||
if tag.key == key and isinstance(tag.value, str):
|
||||
return tag.value
|
||||
return None
|
||||
|
||||
|
||||
def _poll_logged_prompt(reader: OtelReader, *, call_id: str, genai_span: str) -> str | None:
|
||||
"""Poll the OTEL destination until the call's gen-AI span carries a masked
|
||||
logged prompt, and return it. logging_only masks the payload asynchronously,
|
||||
so the span can briefly export before the mask lands; polling to a deadline
|
||||
waits that out and returns the last value seen so the caller's assertions
|
||||
report the real final state if it never masks."""
|
||||
deadline = time.monotonic() + POLL_TIMEOUT
|
||||
last: str | None = None
|
||||
while time.monotonic() < deadline:
|
||||
for trace in reader.traces_for_call(call_id):
|
||||
for span in trace.spans:
|
||||
if span.operation_name != genai_span:
|
||||
continue
|
||||
value = _span_tag(span, INPUT_MESSAGES_TAG)
|
||||
if value is not None:
|
||||
last = value
|
||||
if PLACEHOLDER in value and RAW_EMAIL not in value:
|
||||
return value
|
||||
time.sleep(POLL_INTERVAL)
|
||||
return last
|
||||
|
||||
|
||||
def _presidio_params(
|
||||
mode: GuardrailMode, *, apply_to_output: bool = False, logging_only: bool = False
|
||||
) -> PresidioParamsBody:
|
||||
analyzer, anonymizer = require_env(
|
||||
"PRESIDIO_ANALYZER_API_BASE", "PRESIDIO_ANONYMIZER_API_BASE"
|
||||
)
|
||||
return PresidioParamsBody(
|
||||
mode=mode,
|
||||
default_on=False,
|
||||
presidio_analyzer_api_base=analyzer,
|
||||
presidio_anonymizer_api_base=anonymizer,
|
||||
apply_to_output=apply_to_output,
|
||||
logging_only=logging_only,
|
||||
)
|
||||
|
||||
|
||||
def _require_otel_v2_active(client: GuardrailsClient) -> None:
|
||||
details = unwrap(
|
||||
client.proxy.transport.get(
|
||||
"/health/readiness/details",
|
||||
headers=client.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=ReadinessDetailsResponse,
|
||||
)
|
||||
)
|
||||
assert OTEL_V2_LOGGER in details.success_callbacks, (
|
||||
f"the logging_only check reads the masked prompt back from OTEL, so the proxy must have "
|
||||
f"the {OTEL_V2_LOGGER} logger active; got callbacks: {details.success_callbacks}"
|
||||
)
|
||||
|
||||
|
||||
class TestPresidioGuardrail:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.pre_call.masks",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_pre_call_masks_pii_before_the_model_sees_it(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
require_env("GEMINI_API_KEY")
|
||||
model = client.create_backend_model(resources, prefix="e2e-presidio-pre")
|
||||
name = f"e2e-presidio-pre-{unique_marker()}"
|
||||
guardrail_id = client.register(name, _presidio_params("pre_call"))
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
echoed = _content(
|
||||
unwrap(client.chat(scoped_key, model, ECHO_REQUEST, guardrails=[name], max_tokens=128))
|
||||
)
|
||||
assert RAW_EMAIL not in echoed, (
|
||||
"pre_call masking must strip the raw email before the model sees it, but the "
|
||||
f"model echoed it back: {echoed[:300]!r}"
|
||||
)
|
||||
assert PLACEHOLDER in echoed, (
|
||||
"the model should have echoed the masked placeholder the guardrail substituted, "
|
||||
f"got: {echoed[:300]!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.post_call.masks",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_post_call_masks_pii_in_model_output(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
require_env("GEMINI_API_KEY")
|
||||
model = client.create_backend_model(resources, prefix="e2e-presidio-post")
|
||||
name = f"e2e-presidio-post-{unique_marker()}"
|
||||
guardrail_id = client.register(name, _presidio_params("post_call", apply_to_output=True))
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
out = _content(
|
||||
unwrap(client.chat(scoped_key, model, EMIT_REQUEST, guardrails=[name], max_tokens=128))
|
||||
)
|
||||
assert RAW_EMAIL not in out, (
|
||||
"post_call masking must strip PII the model emitted, but the raw email reached the "
|
||||
f"caller: {out[:300]!r}"
|
||||
)
|
||||
assert PLACEHOLDER in out, (
|
||||
f"the masked placeholder should replace the model's PII output, got: {out[:300]!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.logging_only.masks",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_logging_only_masks_the_logged_prompt(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
require_env("GEMINI_API_KEY")
|
||||
_require_otel_v2_active(client)
|
||||
reader = build_otel_reader()
|
||||
|
||||
model = client.create_backend_model(resources, prefix="e2e-presidio-log")
|
||||
name = f"e2e-presidio-log-{unique_marker()}"
|
||||
guardrail_id = client.register(name, _presidio_params("logging_only", logging_only=True))
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
outcome = client.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=client.proxy.transport.bearer(scoped_key),
|
||||
json=ChatBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=LOG_REQUEST)],
|
||||
max_tokens=64,
|
||||
guardrails=[name],
|
||||
),
|
||||
)
|
||||
require_successful_call(outcome) # logging_only must not block
|
||||
assert outcome.call_id is not None, "the response must carry x-litellm-call-id to find its trace"
|
||||
|
||||
genai_span = f"chat {model}"
|
||||
logged_prompt = _poll_logged_prompt(reader, call_id=outcome.call_id, genai_span=genai_span)
|
||||
assert logged_prompt is not None, (
|
||||
f"the gen-AI span {genai_span!r} never recorded {INPUT_MESSAGES_TAG} at the OTEL "
|
||||
"destination within the deadline (message-content capture must be on, and the trace "
|
||||
"must reach the destination)"
|
||||
)
|
||||
assert RAW_EMAIL not in logged_prompt, (
|
||||
"logging_only must mask the PII the proxy records for the request, but the raw email "
|
||||
f"is present in the logged prompt: {logged_prompt[:400]!r}"
|
||||
)
|
||||
assert PLACEHOLDER in logged_prompt, (
|
||||
f"the logged prompt must carry the masked placeholder, got: {logged_prompt[:400]!r}"
|
||||
)
|
||||
|
|
@ -8,6 +8,8 @@ suite was removed.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import unique_marker
|
||||
|
|
@ -19,11 +21,39 @@ pytestmark = pytest.mark.e2e
|
|||
|
||||
MODEL = "gemini-2.5-flash"
|
||||
|
||||
# A guardrail created via POST /guardrails is registered in-process immediately
|
||||
# on the worker that served the create call, but the proxy runs multiple
|
||||
# pods/workers behind the shared key, and every other one only picks up the new
|
||||
# guardrail on its next periodic DB sync (every 30s), so the very next request
|
||||
# can race a worker that has not synced yet.
|
||||
GUARDRAIL_PROPAGATION_DEADLINE_SECONDS = 40.0
|
||||
GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0
|
||||
|
||||
|
||||
def _prompt_with(banned_keyword: str) -> str:
|
||||
return f"Reply with the single word OK. {banned_keyword}"
|
||||
|
||||
|
||||
def _assert_eventually_blocked(client: GuardrailsClient, key: str, banned: str) -> None:
|
||||
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
|
||||
while True:
|
||||
result = client.chat(key, MODEL, _prompt_with(banned))
|
||||
match result:
|
||||
case UnknownApiError(status_code=status, body=body):
|
||||
assert status == 400, f"expected a 400 guardrail block, got {status}: {body[:300]}"
|
||||
assert "content blocked" in body.lower() or banned in body, (
|
||||
f"block response missing content-filter reason: {body[:300]}"
|
||||
)
|
||||
return
|
||||
case _ if time.monotonic() < deadline:
|
||||
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
|
||||
case _:
|
||||
pytest.fail(
|
||||
f"default-on guardrail never blocked the banned keyword within "
|
||||
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; got {result}"
|
||||
)
|
||||
|
||||
|
||||
class TestTeamDisableGlobalGuardrail:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.litellm_content_filter.pre_call.blocks",
|
||||
|
|
@ -33,25 +63,10 @@ class TestTeamDisableGlobalGuardrail:
|
|||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
banned = unique_marker()
|
||||
guardrail_id = client.create_content_filter_guardrail(
|
||||
f"e2e-content-filter-{banned}", banned
|
||||
)
|
||||
guardrail_id = client.create_content_filter_guardrail(f"e2e-content-filter-{banned}", banned)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
result = client.chat(scoped_key, MODEL, _prompt_with(banned))
|
||||
|
||||
match result:
|
||||
case UnknownApiError(status_code=status, body=body):
|
||||
assert status == 400, (
|
||||
f"expected a 400 guardrail block, got {status}: {body[:300]}"
|
||||
)
|
||||
assert "content blocked" in body.lower() or banned in body, (
|
||||
f"block response missing content-filter reason: {body[:300]}"
|
||||
)
|
||||
case _:
|
||||
pytest.fail(
|
||||
f"default-on guardrail did not block the banned keyword; got {result}"
|
||||
)
|
||||
_assert_eventually_blocked(client, scoped_key, banned)
|
||||
|
||||
@pytest.mark.covers(
|
||||
"guardrail.litellm_content_filter.pre_call.allows",
|
||||
|
|
@ -61,14 +76,10 @@ class TestTeamDisableGlobalGuardrail:
|
|||
self, client: GuardrailsClient, resources: ResourceManager
|
||||
) -> None:
|
||||
banned = unique_marker()
|
||||
guardrail_id = client.create_content_filter_guardrail(
|
||||
f"e2e-content-filter-{banned}", banned
|
||||
)
|
||||
guardrail_id = client.create_content_filter_guardrail(f"e2e-content-filter-{banned}", banned)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
team_id = client.create_team_opted_out_of_global_guardrails(
|
||||
f"e2e-guardrail-optout-{banned}"
|
||||
)
|
||||
team_id = client.create_team_opted_out_of_global_guardrails(f"e2e-guardrail-optout-{banned}")
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
key = client.create_key_in_team(team_id)
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
|
|
|
|||
|
|
@ -85,6 +85,37 @@ class McpToolsListResponse(BaseModel):
|
|||
return None
|
||||
|
||||
|
||||
class BlockedWordSpec(BaseModel):
|
||||
keyword: str
|
||||
action: str = "BLOCK"
|
||||
|
||||
|
||||
class ContentFilterMcpParams(BaseModel):
|
||||
"""litellm_content_filter params scoped to the MCP tool-call hook. mode is
|
||||
pre_mcp_call because a pre_call config silently no-ops on the tools/call path
|
||||
(the event type is rewritten to pre_mcp_call for call_mcp_tool), and default_on
|
||||
is required there because per-key/request guardrail selection is dropped from
|
||||
the synthetic MCP request the hook sees."""
|
||||
|
||||
guardrail: str = "litellm_content_filter"
|
||||
mode: str = "pre_mcp_call"
|
||||
default_on: bool = True
|
||||
blocked_words: list[BlockedWordSpec]
|
||||
|
||||
|
||||
class GuardrailSpecBody(BaseModel):
|
||||
guardrail_name: str
|
||||
litellm_params: ContentFilterMcpParams
|
||||
|
||||
|
||||
class GuardrailCreateBody(BaseModel):
|
||||
guardrail: GuardrailSpecBody
|
||||
|
||||
|
||||
class GuardrailCreateResponse(BaseModel):
|
||||
guardrail_id: str
|
||||
|
||||
|
||||
class McpCallToolBody(BaseModel):
|
||||
name: str
|
||||
arguments: dict[str, McpToolArg]
|
||||
|
|
@ -186,6 +217,35 @@ class McpClient:
|
|||
response_type=McpToolsListResponse,
|
||||
)
|
||||
|
||||
def register_mcp_content_filter(self, *, name: str, blocked_keyword: str) -> str:
|
||||
"""Register a default-on content-filter guardrail that runs on the MCP
|
||||
tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is
|
||||
unique per test, so default_on only ever intercepts this test's own
|
||||
banned tool call on the shared proxy."""
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/guardrails",
|
||||
headers=self.proxy.transport.master,
|
||||
json=GuardrailCreateBody(
|
||||
guardrail=GuardrailSpecBody(
|
||||
guardrail_name=name,
|
||||
litellm_params=ContentFilterMcpParams(
|
||||
blocked_words=[BlockedWordSpec(keyword=blocked_keyword)],
|
||||
),
|
||||
)
|
||||
),
|
||||
response_type=GuardrailCreateResponse,
|
||||
)
|
||||
).guardrail_id
|
||||
|
||||
def delete_guardrail(self, guardrail_id: str) -> None:
|
||||
_ = self.proxy.transport.delete(
|
||||
f"/guardrails/{guardrail_id}",
|
||||
headers=self.proxy.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def call_tool(
|
||||
self,
|
||||
key: str,
|
||||
|
|
|
|||
146
tests/e2e/mcp/test_mcp_guardrail_e2e.py
Normal file
146
tests/e2e/mcp/test_mcp_guardrail_e2e.py
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
"""Live e2e: a guardrail on the MCP tool-call path blocks banned content in the
|
||||
tool arguments before the call reaches the upstream MCP server.
|
||||
|
||||
A general litellm_content_filter guardrail is configured with mode=pre_mcp_call
|
||||
(the event type the proxy rewrites pre_call to for a call_mcp_tool) and default_on
|
||||
(per-key/request guardrail selection is dropped from the synthetic MCP request the
|
||||
hook sees, so default_on is how it attaches to tools/call). The banned keyword is
|
||||
unique per run, so default_on only ever intercepts this test's own banned call.
|
||||
|
||||
Against the real Datadog MCP server, calling search_datadog_logs with the banned
|
||||
keyword in the query is blocked with HTTP 400 attributed to the pre_mcp_call hook,
|
||||
and the tool never runs; the same guardrail lets a clean query through to Datadog.
|
||||
This is the enforced half (the block) plus the pass-through half in one spec.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
|
||||
import pytest
|
||||
|
||||
from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp
|
||||
from e2e_config import DD_SEARCH_FROM, unique_marker
|
||||
from e2e_http import Result, Success, UnknownApiError, unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from mcp_client import McpCallToolResponse, McpClient, McpToolArguments
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
# Stage runs several data-plane pods behind the shared key, and each picks up a
|
||||
# newly registered guardrail only on its next periodic DB sync (~30s in
|
||||
# proxy_server.py). Every pod is guaranteed to have refreshed only once a full sync
|
||||
# interval has elapsed since the create; before then a banned call routed to a
|
||||
# lagging pod passes through as legitimate in-flight propagation, not a leak.
|
||||
GUARDRAIL_FULL_SYNC_SECONDS = 40.0
|
||||
POST_SYNC_VERIFICATION_CALLS = 4
|
||||
|
||||
|
||||
def _poll_until_blocked(
|
||||
search: Callable[[str], Result[McpCallToolResponse]], banned_keyword: str, client: McpClient
|
||||
) -> Result[McpCallToolResponse]:
|
||||
"""Retry a banned tool call until the guardrail blocks it (400) or the deadline
|
||||
passes, returning the last result. Absorbs the control-plane -> data-plane
|
||||
guardrail-sync delay so the check waits for enforcement instead of racing it."""
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
last: Result[McpCallToolResponse] = search(f"tell me about {banned_keyword}")
|
||||
while time.monotonic() < deadline:
|
||||
if isinstance(last, UnknownApiError) and last.status_code == 400:
|
||||
return last
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
last = search(f"tell me about {banned_keyword}")
|
||||
return last
|
||||
|
||||
|
||||
class TestMcpToolCallGuardrail:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.litellm_content_filter.pre_mcp_call.blocks",
|
||||
exercised_on=["mcp_operations"],
|
||||
)
|
||||
def test_content_filter_blocks_banned_keyword_in_tool_args(
|
||||
self, client: McpClient, resources: ResourceManager
|
||||
) -> None:
|
||||
assert_dd_mcp_creds()
|
||||
marker = unique_marker()
|
||||
banned_keyword = f"e2eblocked{marker}"
|
||||
|
||||
guardrail_id = client.register_mcp_content_filter(
|
||||
name=f"e2e-mcp-cf-{marker}", blocked_keyword=banned_keyword
|
||||
)
|
||||
guardrail_created_at = time.monotonic()
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
server_id = register_datadog_mcp(client, resources)
|
||||
key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id])
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
|
||||
tools = unwrap(client.list_tools(key))
|
||||
tool_name = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL)
|
||||
assert tool_name is not None, (
|
||||
f"granted key never saw {SEARCH_LOGS_TOOL} on server {server_id}; "
|
||||
f"tools={tools.tool_names_for_server(server_id)}"
|
||||
)
|
||||
|
||||
def search(query: str) -> Result[McpCallToolResponse]:
|
||||
arguments: McpToolArguments = {
|
||||
"query": query,
|
||||
"from": DD_SEARCH_FROM,
|
||||
"to": "now",
|
||||
"max_tokens": 500,
|
||||
"telemetry": {"intent": "e2e mcp guardrail check"},
|
||||
}
|
||||
return client.call_tool(key, server_id=server_id, name=tool_name, arguments=arguments)
|
||||
|
||||
# Registering the guardrail is a control-plane write; the data-plane worker
|
||||
# that serves tools/call picks it up on its next guardrail sync, so an
|
||||
# immediate call can race the propagation and slip through. Poll the banned
|
||||
# call to the deadline and require a block, so the check proves enforcement
|
||||
# rather than catching a pre-sync pass-through. The keyword is unique per
|
||||
# run, so this only ever intercepts this test's own call.
|
||||
blocked = _poll_until_blocked(search, banned_keyword, client)
|
||||
match blocked:
|
||||
case UnknownApiError(status_code=400, body=body):
|
||||
assert banned_keyword in body or "content blocked" in body.lower(), (
|
||||
f"the block must name the content-filter reason, got: {body[:300]}"
|
||||
)
|
||||
assert "pre_mcp_call" in body, (
|
||||
f"the block must be attributed to the MCP tool-call hook (pre_mcp_call), got: {body[:300]}"
|
||||
)
|
||||
case _:
|
||||
pytest.fail(
|
||||
"content_filter never blocked the banned keyword on the MCP tool call within "
|
||||
f"{client.proxy.poll_timeout}s (guardrail sync to the data plane never landed); "
|
||||
f"last result: {blocked}"
|
||||
)
|
||||
|
||||
# The block above only proves the one pod that served it has synced; another
|
||||
# pod could still lack the guardrail and let the banned call reach Datadog.
|
||||
# Wait out the full sync interval from the create so every pod has refreshed
|
||||
# from the DB, then require the banned call to stay blocked across several
|
||||
# attempts. A pass-through now is a genuine partial-propagation leak, not a
|
||||
# race. Client load balancing still can't guarantee every pod is hit, so this
|
||||
# samples several worker selections rather than proving all pods synced.
|
||||
sync_remaining = guardrail_created_at + GUARDRAIL_FULL_SYNC_SECONDS - time.monotonic()
|
||||
if sync_remaining > 0:
|
||||
time.sleep(sync_remaining)
|
||||
for attempt in range(1, POST_SYNC_VERIFICATION_CALLS + 1):
|
||||
reblocked = search(f"still about {banned_keyword} #{attempt}")
|
||||
assert isinstance(reblocked, UnknownApiError) and reblocked.status_code == 400, (
|
||||
"after the guardrail sync interval every data-plane pod must block the banned "
|
||||
f"keyword, but attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was allowed "
|
||||
f"through (a pod still lacks the guardrail): {reblocked}"
|
||||
)
|
||||
if attempt < POST_SYNC_VERIFICATION_CALLS:
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
|
||||
allowed = search(f"e2e-clean-{marker}")
|
||||
match allowed:
|
||||
case Success(data=result):
|
||||
assert result.is_error is not True, (
|
||||
f"a clean MCP tool call must reach the server and not error, got: {result}"
|
||||
)
|
||||
case _:
|
||||
pytest.fail(
|
||||
f"a clean MCP tool call must pass the guardrail and reach the server; got {allowed}"
|
||||
)
|
||||
|
|
@ -817,3 +817,23 @@ class TagListResponse(RootModel[list[TagListEntry]]):
|
|||
"""GET /tag/list answers with a bare array of tag configs (the stored tags plus
|
||||
any dynamically-seen spend tags), not an object wrapping them. Read the rows off
|
||||
.root."""
|
||||
|
||||
|
||||
# ---------- health / lifecycle ----------
|
||||
|
||||
|
||||
class ReadinessResponse(BaseModel):
|
||||
"""GET /health/readiness (public probe). The low-detail payload a load
|
||||
balancer sees: `status` plus the resolved DB state (`connected`,
|
||||
`disconnected`, or `Not connected`)."""
|
||||
|
||||
status: str
|
||||
db: str | None = None
|
||||
|
||||
|
||||
class ReadinessDetailsResponse(ReadinessResponse):
|
||||
"""GET /health/readiness/details (authenticated). Extends the public payload
|
||||
with the diagnostics only an authenticated caller may read."""
|
||||
|
||||
litellm_version: str | None = None
|
||||
success_callbacks: list[str] = []
|
||||
|
|
|
|||
18
tests/e2e/other/conftest.py
Normal file
18
tests/e2e/other/conftest.py
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
"""`other` suite's `client` fixture.
|
||||
|
||||
Lifecycle (resources/scoped_key), proxy liveness gate, and the e2e/covers
|
||||
markers all live in the parent tests/e2e/conftest.py. OtherClient holds the
|
||||
shared ProxyClient so anything these tests create tears down through it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from other_client import OtherClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client(proxy: ProxyClient) -> OtherClient:
|
||||
return build_client(proxy)
|
||||
73
tests/e2e/other/other_client.py
Normal file
73
tests/e2e/other/other_client.py
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
"""Client for the `other` holding-pen suite: the auth gate (master key vs an
|
||||
invalid key on an admin route) and the process-lifecycle health probes
|
||||
(liveness, public readiness, authenticated readiness diagnostics).
|
||||
|
||||
Holds the shared ProxyClient so `resources` / `scoped_key` still clean up, and
|
||||
adds only the routes these behaviors need. The health probes deliberately send
|
||||
no auth header (public routes), so they go through the transport with an empty
|
||||
headers model rather than a bearer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from e2e_http import NoBody, ProbeResult, Result
|
||||
from models import (
|
||||
ReadinessDetailsResponse,
|
||||
ReadinessResponse,
|
||||
UserListParams,
|
||||
UserListResponse,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OtherClient:
|
||||
proxy: ProxyClient
|
||||
|
||||
def liveness(self) -> ProbeResult:
|
||||
"""GET /health/liveliness. Unauthenticated; the probe returns status +
|
||||
raw body so the test can assert the worker reports itself alive."""
|
||||
return self.proxy.transport.probe("/health/liveliness", params=NoBody())
|
||||
|
||||
def readiness_public(self) -> Result[ReadinessResponse]:
|
||||
"""GET /health/readiness with no credential at all, proving the probe is
|
||||
safe to expose to an unauthenticated load balancer."""
|
||||
return self.proxy.transport.get(
|
||||
"/health/readiness",
|
||||
headers=NoBody(),
|
||||
params=NoBody(),
|
||||
response_type=ReadinessResponse,
|
||||
)
|
||||
|
||||
def readiness_details(self, key: str) -> Result[ReadinessDetailsResponse]:
|
||||
return self.proxy.transport.get(
|
||||
"/health/readiness/details",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=NoBody(),
|
||||
response_type=ReadinessDetailsResponse,
|
||||
)
|
||||
|
||||
def readiness_details_unauthenticated(self) -> Result[ReadinessDetailsResponse]:
|
||||
return self.proxy.transport.get(
|
||||
"/health/readiness/details",
|
||||
headers=NoBody(),
|
||||
params=NoBody(),
|
||||
response_type=ReadinessDetailsResponse,
|
||||
)
|
||||
|
||||
def list_users_as(self, key: str) -> Result[UserListResponse]:
|
||||
"""GET /user/list under `key`. Admin-only, so it doubles as the master
|
||||
key's authorization proof: the master key (proxy admin) reads it, a
|
||||
non-matching key is rejected before it ever reaches the handler."""
|
||||
return self.proxy.transport.get(
|
||||
"/user/list",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=UserListParams(user_ids="e2e-test-user"),
|
||||
response_type=UserListResponse,
|
||||
)
|
||||
|
||||
|
||||
def build_client(proxy: ProxyClient) -> OtherClient:
|
||||
return OtherClient(proxy=proxy)
|
||||
65
tests/e2e/other/test_health_lifecycle_e2e.py
Normal file
65
tests/e2e/other/test_health_lifecycle_e2e.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
"""Live e2e: the process-lifecycle probes Kubernetes and load balancers depend on.
|
||||
|
||||
Liveness and public readiness must answer without a credential (a load balancer
|
||||
has none), and public readiness must distinguish a healthy worker from one whose
|
||||
DB is unreachable by reporting the resolved DB state. The detailed readiness
|
||||
route, by contrast, is authenticated: it exposes diagnostics (version, callbacks,
|
||||
DB) and must reject an anonymous caller. The suite runs against a proxy configured
|
||||
with a real database, so a healthy readiness payload reports the DB as connected;
|
||||
a regression that stopped checking the DB, or dropped the public exposure, fails
|
||||
here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import MASTER_KEY
|
||||
from e2e_http import UnauthorizedError, unwrap
|
||||
from other_client import OtherClient
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
||||
class TestHealthLifecycle:
|
||||
@pytest.mark.covers("other.lifecycle.liveness.ping")
|
||||
def test_liveness_reports_alive_without_auth(self, client: OtherClient) -> None:
|
||||
probe = client.liveness()
|
||||
assert probe.status_code == 200, (
|
||||
f"liveness must answer 200 for an unauthenticated probe, got "
|
||||
f"{probe.status_code}: {probe.body[:200]}"
|
||||
)
|
||||
assert "alive" in probe.body.lower(), (
|
||||
f"liveness body must confirm the worker is alive, got {probe.body[:200]}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.lifecycle.readiness.public_probe")
|
||||
def test_readiness_is_reachable_without_credentials(self, client: OtherClient) -> None:
|
||||
readiness = unwrap(client.readiness_public())
|
||||
assert readiness.status == "healthy", (
|
||||
f"public readiness must report a healthy worker, got status {readiness.status!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.lifecycle.readiness.reports_db_status")
|
||||
def test_readiness_reports_connected_db(self, client: OtherClient) -> None:
|
||||
readiness = unwrap(client.readiness_public())
|
||||
assert readiness.db == "connected", (
|
||||
"readiness must report the configured database as connected so an "
|
||||
f"orchestrator can tell a healthy worker from a DB-unreachable one, got {readiness.db!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.lifecycle.readiness_details.authenticated_diagnostics")
|
||||
def test_readiness_details_require_auth_and_expose_diagnostics(self, client: OtherClient) -> None:
|
||||
anonymous = client.readiness_details_unauthenticated()
|
||||
assert isinstance(anonymous, UnauthorizedError), (
|
||||
f"/health/readiness/details must reject an unauthenticated caller, got {anonymous}"
|
||||
)
|
||||
|
||||
details = unwrap(client.readiness_details(MASTER_KEY))
|
||||
assert details.status == "healthy", f"authenticated readiness status must be healthy, got {details.status!r}"
|
||||
assert details.litellm_version is not None, (
|
||||
"authenticated diagnostics must expose the litellm version"
|
||||
)
|
||||
assert details.db == "connected", (
|
||||
f"authenticated diagnostics must report the DB as connected, got {details.db!r}"
|
||||
)
|
||||
37
tests/e2e/other/test_master_key_auth_e2e.py
Normal file
37
tests/e2e/other/test_master_key_auth_e2e.py
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
"""Live e2e: the master key authenticates and is treated as a proxy admin, and a
|
||||
key that is not the master key is rejected before reaching the handler.
|
||||
|
||||
/user/list is admin-only, so it proves both halves of the master-key contract in
|
||||
one route: the master key reads it (authenticated + authorized as admin), while a
|
||||
freshly minted, never-provisioned token is denied 401 by the auth layer. The
|
||||
invalid case uses a unique, master-key-shaped token so the check exercises the
|
||||
credential comparison rather than a value that could collide with a real key.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import MASTER_KEY, unique_marker
|
||||
from e2e_http import UnauthorizedError, unwrap
|
||||
from other_client import OtherClient
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
||||
class TestMasterKeyAuth:
|
||||
@pytest.mark.covers("other.auth.master_key.valid_allows")
|
||||
def test_master_key_authenticates_and_grants_admin_route(self, client: OtherClient) -> None:
|
||||
listing = unwrap(client.list_users_as(MASTER_KEY))
|
||||
assert listing.total >= 0, (
|
||||
"master key reached the admin /user/list handler but the response did not "
|
||||
f"carry a user count: {listing}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.auth.master_key.invalid_denied")
|
||||
def test_non_matching_master_key_is_denied(self, client: OtherClient) -> None:
|
||||
bogus = f"sk-{unique_marker()}"
|
||||
result = client.list_users_as(bogus)
|
||||
assert isinstance(result, UnauthorizedError), (
|
||||
f"a token that is not the master key must be rejected with 401, got {result}"
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue