test(integration): audit cells for failure-path redaction

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-29 19:24:56 +00:00
parent be7734d726
commit 92f0f35f05
4 changed files with 1558 additions and 22 deletions

View file

@ -1,11 +1,17 @@
import asyncio
import json import json
import threading
import time
import uuid import uuid
from collections.abc import Iterator from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Final from typing import Final
import anthropic
import httpx import httpx
import openai
import pytest import pytest
import yaml import yaml
from integration._support.client import Gateway, JsonValue, eventually, gateway_from_environment, object_value from integration._support.client import Gateway, JsonValue, eventually, gateway_from_environment, object_value
@ -31,47 +37,193 @@ class Rig:
) )
def _provider(request: Request) -> Reply: def _prompt_text(body: dict[str, JsonValue]) -> str:
text: Final = json.loads(request.body)["messages"][-1]["content"] messages: Final = body.get("messages")
if isinstance(messages, list) and messages:
last: Final = messages[-1]
if isinstance(last, dict):
content: Final = last.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
return "".join(
str(part["text"])
for part in content
if isinstance(part, dict) and isinstance(part.get("text"), str)
)
input_value: Final = body.get("input")
if isinstance(input_value, str):
return input_value
if isinstance(input_value, list):
return json.dumps(input_value)
return json.dumps(body)[:200]
def _success_reply(target: str, text: str) -> Reply:
if target.endswith("/messages"):
return Reply(
body=json.dumps(
{
"id": "msg_" + uuid.uuid4().hex,
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5",
"content": [{"type": "text", "text": text}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 3, "output_tokens": 2},
}
).encode()
)
if target.endswith("/responses"):
return Reply(
body=json.dumps(
{
"id": "resp_" + uuid.uuid4().hex,
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"id": "msg_" + uuid.uuid4().hex,
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": text, "annotations": []}],
}
],
"usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
}
).encode()
)
return Reply( return Reply(
status=400,
body=json.dumps( body=json.dumps(
{"error": {"type": "invalid_request_error", "message": f"Unsupported content: {text}"}} {
).encode(), "id": "chatcmpl-" + uuid.uuid4().hex,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
}
).encode()
) )
@pytest.fixture(scope="module") def _provider(request: Request) -> Reply:
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: try:
root: Final = tmp_path_factory.mktemp("failure_redaction") body: Final = json.loads(request.body)
except json.JSONDecodeError:
return Reply(status=404, body=b"{}")
text: Final = _prompt_text(body)
if text.startswith("ok "):
return _success_reply(request.target, text)
status: Final = 401 if text.startswith("auth401 ") else 400
error: Final = {"type": "invalid_request_error", "message": f"Unsupported content: {text}"}
if request.target.endswith("/messages"):
return Reply(status=status, body=json.dumps({"type": "error", "error": error}).encode())
return Reply(status=status, body=json.dumps({"error": error}).encode())
_SINK_OUTAGE: Final = threading.Event()
_SINK_SLOW: Final = threading.Event()
def _sink(request: Request) -> Reply:
if _SINK_OUTAGE.is_set():
return Reply(status=503, body=b'{"error":"sink down"}')
if _SINK_SLOW.is_set():
time.sleep(2)
return Reply()
@contextmanager
def _booted_rig(
root: Path,
provider: Wire,
sink: Wire,
settings: Mapping[str, JsonValue],
) -> Iterator[Rig]:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update( config["general_settings"]["allow_client_side_credentials"] = True
{ config["litellm_settings"].update({"DEFAULT_FLUSH_INTERVAL_SECONDS": 1, **settings})
"callbacks": ["generic_api"],
"DEFAULT_FLUSH_INTERVAL_SECONDS": 1,
"turn_off_message_logging": True,
"standard_logging_payload_excluded_fields": ["hidden_params"],
}
)
path: Final = root / "failure_redaction.yaml" path: Final = root / "failure_redaction.yaml"
path.write_text(yaml.safe_dump(config)) path.write_text(yaml.safe_dump(config))
with ( with (
gateway_from_environment() as gateway, gateway_from_environment() as gateway,
wire_server(_provider) as provider, owned_proxy(gateway, root, {"GENERIC_LOGGER_ENDPOINT": sink.url}, config=path, workers=2) as proxy,
wire_server(lambda _: Reply()) as sink,
owned_proxy(gateway, root, {"GENERIC_LOGGER_ENDPOINT": sink.url}, config=path) as proxy,
): ):
yield Rig(proxy, provider, sink, []) # mutable-ok: sink drain consumes batches, later polls keep earlier ones yield Rig(proxy, provider, sink, []) # mutable-ok: sink drain consumes batches, later polls keep earlier ones
@pytest.fixture(scope="module")
def provider() -> Iterator[Wire]:
with wire_server(_provider) as wire:
yield wire
@pytest.fixture(scope="module")
def sink() -> Iterator[Wire]:
with wire_server(_sink) as wire:
yield wire
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory, provider: Wire, sink: Wire) -> Iterator[Rig]:
with _booted_rig(
tmp_path_factory.mktemp("failure_redaction"),
provider,
sink,
{
"callbacks": ["generic_api"],
"turn_off_message_logging": True,
"standard_logging_payload_excluded_fields": ["hidden_params"],
},
) as booted:
yield booted
@pytest.fixture(scope="module")
def rig_off(tmp_path_factory: pytest.TempPathFactory, provider: Wire, sink: Wire) -> Iterator[Rig]:
with _booted_rig(
tmp_path_factory.mktemp("failure_redaction_off"), provider, sink, {"callbacks": ["generic_api"]}
) as booted:
yield booted
@pytest.fixture(scope="module")
def rig_failure_callback(tmp_path_factory: pytest.TempPathFactory, provider: Wire, sink: Wire) -> Iterator[Rig]:
with _booted_rig(
tmp_path_factory.mktemp("failure_redaction_failure_cb"),
provider,
sink,
{"failure_callback": ["generic_api"], "turn_off_message_logging": True},
) as booted:
yield booted
@pytest.fixture(scope="module")
def rig_success_callback(tmp_path_factory: pytest.TempPathFactory, provider: Wire, sink: Wire) -> Iterator[Rig]:
with _booted_rig(
tmp_path_factory.mktemp("failure_redaction_success_cb"),
provider,
sink,
{"success_callback": ["generic_api"], "turn_off_message_logging": True},
) as booted:
yield booted
def _secret_prompt() -> str: def _secret_prompt() -> str:
return "confidential-prompt-" + uuid.uuid4().hex return "confidential-prompt-" + uuid.uuid4().hex
def _single_failure_event(rig: Rig, model: str) -> dict[str, JsonValue]: def _single_failure_event(rig: Rig, model: str) -> dict[str, JsonValue]:
events: Final = eventually(lambda: rig.failure_events(model), lambda values: len(values) >= 1, seconds=20) events: Final = eventually(
lambda: tuple(event for event in rig.failure_events(model) if event.get("status") == "failure"),
lambda values: len(values) >= 1,
seconds=30,
)
assert len(events) == 1, events assert len(events) == 1, events
assert events[0]["status"] == "failure", events[0]
return events[0] return events[0]
@ -155,6 +307,7 @@ def test_proxy_only_rejection_does_not_leak_the_prompt_to_callbacks(rig: Rig) ->
allowed: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key") allowed: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
denied: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key") denied: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[allowed]) key: Final = scenario.key(models=[allowed])
rig.provider.drain()
response: Final = _chat(rig, denied, [{"role": "user", "content": secret}], key=key) response: Final = _chat(rig, denied, [{"role": "user", "content": secret}], key=key)
assert response.status_code in (401, 403), response.text assert response.status_code in (401, 403), response.text
assert rig.provider.drain() == () assert rig.provider.drain() == ()
@ -162,3 +315,694 @@ def test_proxy_only_rejection_does_not_leak_the_prompt_to_callbacks(rig: Rig) ->
assert secret not in json.dumps(event), json.dumps(event) assert secret not in json.dumps(event), json.dumps(event)
assert "hidden_params" not in event, sorted(event) assert "hidden_params" not in event, sorted(event)
assert _error_information(event)["error_code"] == str(response.status_code), event assert _error_information(event)["error_code"] == str(response.status_code), event
def _openai(rig: Rig, key: str | None = None) -> openai.OpenAI:
return openai.OpenAI(
base_url=str(rig.proxy.client.base_url) + "/v1",
api_key=key or rig.proxy.key,
max_retries=0,
)
def _async_openai(rig: Rig, key: str | None = None) -> openai.AsyncOpenAI:
return openai.AsyncOpenAI(
base_url=str(rig.proxy.client.base_url) + "/v1",
api_key=key or rig.proxy.key,
max_retries=0,
)
def _anthropic(rig: Rig, key: str | None = None) -> anthropic.Anthropic:
return anthropic.Anthropic(base_url=str(rig.proxy.client.base_url), api_key=key or rig.proxy.key, max_retries=0)
def _async_anthropic(rig: Rig, key: str | None = None) -> anthropic.AsyncAnthropic:
return anthropic.AsyncAnthropic(
base_url=str(rig.proxy.client.base_url), api_key=key or rig.proxy.key, max_retries=0
)
def _call_id(response: httpx.Response) -> str:
call_id: Final = response.headers.get("x-litellm-call-id")
assert call_id, dict(response.headers)
return call_id
def _error_call_id(error: openai.APIStatusError | anthropic.APIStatusError) -> str:
call_id: Final = error.response.headers.get("x-litellm-call-id")
assert call_id, dict(error.response.headers)
return call_id
def _assert_failure_redacted(rig: Rig, model: str, secret: str, call_id: str) -> dict[str, JsonValue]:
event: Final = _single_failure_event(rig, model)
assert secret not in json.dumps(event), json.dumps(event)
error_information: Final = _error_information(event)
assert error_information["error_class"], error_information
assert error_information["llm_provider"], error_information
row: Final = _spend_row(call_id)
assert row["status"] == "failure", row
assert secret not in json.dumps(row, default=str), row
return event
def _assert_failure_raw(rig: Rig, model: str, secret: str, call_id: str | None = None) -> dict[str, JsonValue]:
event: Final = _single_failure_event(rig, model)
assert secret in json.dumps(event), json.dumps(event)
if call_id is not None:
row: Final = _spend_row(call_id)
assert secret in json.dumps(row, default=str), row
return event
def _denied_body(path: str, model: str, secret: str) -> dict[str, JsonValue]:
if path == "/v1/messages":
return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": secret}]}
if path == "/v1/responses":
return {"model": model, "input": secret}
return {"model": model, "messages": [{"role": "user", "content": secret}]}
# --- A. Endpoint x stream x client, provider 400 echoing the prompt -------------------
_SDK_CLIENTS: Final = ("openai_sync", "openai_async", "httpx")
def _chat_call(rig: Rig, client: str, model: str, secret: str, stream: bool) -> str:
messages: Final = [{"role": "user", "content": secret}]
if client == "httpx":
response: Final = _chat(rig, model, messages)
assert response.status_code == 400, response.text
return _call_id(response)
if client == "openai_sync":
with pytest.raises(openai.BadRequestError) as caught:
_openai(rig).chat.completions.create(model=model, messages=messages, stream=stream)
assert caught.value.status_code == 400, caught.value
return _error_call_id(caught.value)
async def fire() -> str:
with pytest.raises(openai.BadRequestError) as caught:
await _async_openai(rig).chat.completions.create(model=model, messages=messages, stream=stream)
return _error_call_id(caught.value)
return asyncio.run(fire())
@pytest.mark.parametrize("client", _SDK_CLIENTS, ids=lambda value: value)
def test_a1_chat_provider_error_redacted(rig: Rig, client: str) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
call_id: Final = _chat_call(rig, client, model, secret, stream=False)
assert any(secret.encode() in request.body for request in rig.provider.drain())
event: Final = _assert_failure_redacted(rig, model, secret, call_id)
assert object_value(event["error_information"])["error_class"] == "BadRequestError", event
@pytest.mark.parametrize("client", _SDK_CLIENTS, ids=lambda value: value)
def test_a2_chat_stream_provider_error_redacted(rig: Rig, client: str) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
call_id: Final = _chat_call(rig, client, model, secret, stream=True)
assert any(secret.encode() in request.body for request in rig.provider.drain())
_assert_failure_redacted(rig, model, secret, call_id)
def _consume_messages_stream(client: anthropic.Anthropic, model: str, messages: list[dict[str, str]]) -> None:
with client.messages.stream(model=model, max_tokens=16, messages=messages) as events:
for _ in events:
pass
async def _consume_messages_stream_async(
client: anthropic.AsyncAnthropic, model: str, messages: list[dict[str, str]]
) -> None:
async with client.messages.stream(model=model, max_tokens=16, messages=messages) as events:
async for _ in events:
pass
def _messages_call(rig: Rig, client: str, model: str, secret: str, stream: bool) -> str:
messages: Final = [{"role": "user", "content": secret}]
if client == "httpx":
response: Final = rig.proxy.request(
"POST", "/v1/messages", {"model": model, "max_tokens": 16, "messages": messages}
)
assert response.status_code == 400, response.text
return _call_id(response)
if client == "openai_sync":
anthropic_client: Final = _anthropic(rig)
if stream:
with pytest.raises(anthropic.BadRequestError) as caught:
_consume_messages_stream(anthropic_client, model, messages)
return _error_call_id(caught.value)
with pytest.raises(anthropic.BadRequestError) as caught:
anthropic_client.messages.create(model=model, max_tokens=16, messages=messages)
return _error_call_id(caught.value)
async def fire() -> str:
client_async: Final = _async_anthropic(rig)
if stream:
with pytest.raises(anthropic.BadRequestError) as caught:
await _consume_messages_stream_async(client_async, model, messages)
return _error_call_id(caught.value)
with pytest.raises(anthropic.BadRequestError) as caught:
await client_async.messages.create(model=model, max_tokens=16, messages=messages)
return _error_call_id(caught.value)
return asyncio.run(fire())
@pytest.mark.parametrize("client", _SDK_CLIENTS, ids=lambda value: value)
def test_a3_messages_provider_error_redacted(rig: Rig, client: str) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(
model="anthropic/claude-sonnet-5", api_base=rig.provider.url, api_key="synthetic-anthropic-key"
)
call_id: Final = _messages_call(rig, client, model, secret, stream=False)
assert any(secret.encode() in request.body for request in rig.provider.drain())
_assert_failure_redacted(rig, model, secret, call_id)
@pytest.mark.parametrize("client", _SDK_CLIENTS, ids=lambda value: value)
def test_a4_messages_stream_provider_error_redacted(rig: Rig, client: str) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(
model="anthropic/claude-sonnet-5", api_base=rig.provider.url, api_key="synthetic-anthropic-key"
)
call_id: Final = _messages_call(rig, client, model, secret, stream=True)
assert any(secret.encode() in request.body for request in rig.provider.drain())
_assert_failure_redacted(rig, model, secret, call_id)
def _responses_call(rig: Rig, client: str, model: str, secret: str, stream: bool) -> str:
if client == "httpx":
response: Final = rig.proxy.request(
"POST", "/v1/responses", {"model": model, "input": secret, "stream": stream}
)
assert response.status_code == 400, response.text
return _call_id(response)
if client == "openai_sync":
with pytest.raises(openai.BadRequestError) as caught:
_openai(rig).responses.create(model=model, input=secret, stream=stream)
return _error_call_id(caught.value)
async def fire() -> str:
with pytest.raises(openai.BadRequestError) as caught:
await _async_openai(rig).responses.create(model=model, input=secret, stream=stream)
return _error_call_id(caught.value)
return asyncio.run(fire())
@pytest.mark.parametrize("client", _SDK_CLIENTS, ids=lambda value: value)
def test_a5_responses_provider_error_redacted(rig: Rig, client: str) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
call_id: Final = _responses_call(rig, client, model, secret, stream=False)
assert any(secret.encode() in request.body for request in rig.provider.drain())
_assert_failure_redacted(rig, model, secret, call_id)
@pytest.mark.parametrize("client", _SDK_CLIENTS, ids=lambda value: value)
def test_a6_responses_stream_provider_error_redacted(rig: Rig, client: str) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
call_id: Final = _responses_call(rig, client, model, secret, stream=True)
assert any(secret.encode() in request.body for request in rig.provider.drain())
_assert_failure_redacted(rig, model, secret, call_id)
def test_a7_messages_transformation_error_redacted(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(
model="anthropic/claude-sonnet-5", api_base=rig.provider.url, api_key="synthetic-anthropic-key"
)
response: Final = rig.proxy.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 16,
"messages": [
{"role": "user", "content": secret},
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": 12345, "type": "function", "function": {"name": None, "arguments": "{}"}}
],
},
{"role": "tool", "tool_call_id": 12345, "content": secret},
],
},
)
assert response.status_code == 400, response.text
_assert_failure_redacted(rig, model, secret, _call_id(response))
@pytest.mark.parametrize(
"path", ("/v1/chat/completions", "/v1/messages", "/v1/responses"), ids=lambda v: v.split("/")[-1]
)
def test_a8_proxy_only_rejection_redacted(rig: Rig, path: str) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
allowed: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
denied: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[allowed])
rig.provider.drain()
response: Final = rig.proxy.request("POST", path, _denied_body(path, denied, secret), key=key)
assert response.status_code in (401, 403), response.text
assert rig.provider.drain() == ()
event: Final = _single_failure_event(rig, denied)
assert secret not in json.dumps(event), json.dumps(event)
assert _error_information(event)["error_code"] == str(response.status_code), event
def test_a9_unknown_model_failure_redacted(rig: Rig) -> None:
secret: Final = _secret_prompt()
model: Final = "unknown-" + uuid.uuid4().hex
rig.provider.drain()
response: Final = _chat(rig, model, [{"role": "user", "content": secret}])
assert response.status_code == 400, response.text
assert rig.provider.drain() == ()
event: Final = eventually(lambda: rig.failure_events(model), lambda values: len(values) >= 1, seconds=20)[0]
assert secret not in json.dumps(event), json.dumps(event)
row: Final = _spend_row(_call_id(response))
assert secret not in json.dumps(row, default=str), row
# --- B. Redaction source modes (YAML global off unless noted) -------------------------
@pytest.mark.parametrize(
("headers", "body"),
(
pytest.param(None, {}, id="b1_no_signal"),
pytest.param({"x-litellm-enable-message-redaction": "true"}, {}, id="b2_enable_header"),
pytest.param({"litellm-enable-message-redaction": "true"}, {}, id="b3_legacy_enable_header"),
pytest.param(None, {"turn_off_message_logging": True}, id="b4_body_param"),
pytest.param(None, {"metadata": {"turn_off_message_logging": True}}, id="b5_metadata_param"),
pytest.param(None, {"litellm_metadata": {"turn_off_message_logging": True}}, id="b5b_litellm_metadata"),
),
)
def test_b_request_level_opt_in_redacts(
rig_off: Rig, headers: Mapping[str, str] | None, body: Mapping[str, JsonValue]
) -> None:
secret: Final = _secret_prompt()
with rig_off.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig_off.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = rig_off.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}], **body},
headers=headers,
)
assert response.status_code == 400, response.text
assert any(secret.encode() in request.body for request in rig_off.provider.drain())
if headers is None and not body:
_assert_failure_raw(rig_off, model, secret, _call_id(response))
else:
_assert_failure_redacted(rig_off, model, secret, _call_id(response))
def test_b8_global_on_permitted_key_opt_out_keeps_raw(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model], metadata={"allow_client_message_redaction_opt_out": True})
response: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}]},
key=key,
headers={"litellm-disable-message-redaction": "true"},
)
assert response.status_code == 400, response.text
_assert_failure_raw(rig, model, secret, _call_id(response))
def test_b9_global_on_disable_header_without_permission_stays_redacted(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
response: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}]},
key=key,
headers={"litellm-disable-message-redaction": "true"},
)
assert response.status_code == 400, response.text
_assert_failure_redacted(rig, model, secret, _call_id(response))
def test_b10_global_on_team_permission_opt_out_keeps_raw(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
team: Final = scenario.team(metadata={"allow_client_message_redaction_opt_out": True})
key: Final = scenario.key(models=[model], team_id=team)
response: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}]},
key=key,
headers={"litellm-disable-message-redaction": "true"},
)
assert response.status_code == 400, response.text
_assert_failure_raw(rig, model, secret, _call_id(response))
# --- C. Callback registration modes (YAML global on) ----------------------------------
def test_c2_failure_callback_registration_redacts(rig_failure_callback: Rig) -> None:
rig: Final = rig_failure_callback
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = _chat(rig, model, [{"role": "user", "content": secret}])
assert response.status_code == 400, response.text
_assert_failure_redacted(rig, model, secret, _call_id(response))
def test_c3_success_callback_only_emits_no_failure_event(rig_success_callback: Rig) -> None:
rig: Final = rig_success_callback
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = _chat(rig, model, [{"role": "user", "content": secret}])
assert response.status_code == 400, response.text
events: Final = eventually(
lambda: rig.failure_events(model),
lambda values: len(values) >= 1,
seconds=8,
return_last_on_timeout=True,
)
assert events == (), events
def test_c4_excluded_fields_stripped_on_failure(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = _chat(rig, model, [{"role": "user", "content": secret}])
assert response.status_code == 400, response.text
event: Final = _assert_failure_redacted(rig, model, secret, _call_id(response))
assert "hidden_params" not in event, sorted(event)
# --- D. Cache-hit twin -----------------------------------------------------------------
@pytest.mark.timeout(280)
def test_d1_cache_hit_success_then_failure_redacted(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
rig.provider.drain()
ok_text: Final = "ok " + uuid.uuid4().hex
first: Final = _chat(rig, model, [{"role": "user", "content": ok_text}])
assert first.status_code == 200, first.text
second: Final = _chat(rig, model, [{"role": "user", "content": ok_text}])
assert second.status_code == 200, second.text
assert second.json()["id"] == first.json()["id"], "second call did not hit the response cache"
response: Final = _chat(rig, model, [{"role": "user", "content": secret}])
assert response.status_code == 400, response.text
_assert_failure_redacted(rig, model, secret, _call_id(response))
# --- E. Sad paths -----------------------------------------------------------------------
_HOSTILE_HEADERS: Final = ("", "0", "false", "1", "a,b", "h" * 5000)
@pytest.mark.parametrize("value", _HOSTILE_HEADERS, ids=lambda v: v[:8] or "empty")
@pytest.mark.timeout(280)
def test_e1_hostile_enable_header_unauthenticated_and_authenticated(rig_off: Rig, value: str) -> None:
secret: Final = _secret_prompt()
with rig_off.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig_off.provider.url + "/v1", api_key="synthetic-provider-key")
rejected: Final = rig_off.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}]},
key="sk-bogus",
headers={"x-litellm-enable-message-redaction": value},
)
assert rejected.status_code == 401, rejected.text
call_ids: list[str] = [] # mutable-ok: collect per-request call ids across the loop
for _ in range(2):
response: Final = rig_off.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}], "cache": {"no-cache": True}},
headers={"x-litellm-enable-message-redaction": value},
)
assert response.status_code == 400, response.text
call_ids.append(_call_id(response))
rig_off.provider.drain()
for call_id in call_ids:
events: Final = eventually(
lambda cid=call_id: tuple(
event
for event in rig_off.failure_events(model)
if event.get("litellm_call_id") == cid
and object_value(event.get("error_information")).get("error_class") != "KeyNotFoundError"
),
lambda values: len(values) == 1,
seconds=90,
)
if value:
assert secret not in json.dumps(events[0]), (value, json.dumps(events[0])[:400])
else:
assert secret in json.dumps(events[0]), (value, json.dumps(events[0])[:400])
@pytest.mark.parametrize("value", (1, [True], "", "v" * 5000, None), ids=lambda v: type(v).__name__)
def test_e2_odd_turn_off_message_logging_values_never_500(rig_off: Rig, value: JsonValue) -> None:
secret: Final = _secret_prompt()
with rig_off.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig_off.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = rig_off.proxy.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": secret}],
"turn_off_message_logging": value,
},
)
assert response.status_code == 400, response.text
event: Final = _single_failure_event(rig_off, model)
assert event["status"] == "failure", event
def test_e3_sink_rejections_keep_proxy_serving_and_later_events_still_redacted(rig: Rig) -> None:
secret: Final = _secret_prompt()
_SINK_OUTAGE.set()
try:
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
during: Final = _chat(rig, model, [{"role": "user", "content": "down " + secret}])
assert during.status_code == 400, during.text
ok: Final = _chat(rig, model, [{"role": "user", "content": "ok alive " + uuid.uuid4().hex}])
assert ok.status_code == 200, ok.text
rig.sink.drain()
finally:
_SINK_OUTAGE.clear()
with rig.proxy.scenario() as scenario:
model_two: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
after: Final = _chat(rig, model_two, [{"role": "user", "content": "after " + secret}])
assert after.status_code == 400, after.text
event: Final = _single_failure_event(rig, model_two)
assert secret not in json.dumps(event), json.dumps(event)
def test_e4_header_only_unknown_model_spend_row_never_carries_marker(rig_off: Rig) -> None:
secret: Final = _secret_prompt()
response: Final = rig_off.proxy.request(
"POST",
"/v1/chat/completions",
{"model": "unknown-" + uuid.uuid4().hex, "messages": [{"role": "user", "content": secret}]},
headers={"x-litellm-enable-message-redaction": "true"},
)
assert response.status_code == 400, response.text
row: Final = _spend_row(_call_id(response))
assert secret not in json.dumps(row, default=str), row
@pytest.mark.parametrize("callback_fixture", ("rig", "rig_failure_callback"), ids=("callbacks", "failure_callback"))
@pytest.mark.parametrize("text", ("auth401 ", ""), ids=("provider_401", "provider_400"))
def test_e5_provider_401_and_400_echo_redacted(
request: pytest.FixtureRequest, callback_fixture: str, text: str
) -> None:
rig: Final = request.getfixturevalue(callback_fixture)
secret: Final = text + _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = _chat(rig, model, [{"role": "user", "content": secret}])
assert response.status_code in (400, 401), response.text
_assert_failure_redacted(rig, model, secret, _call_id(response))
@pytest.mark.timeout(280)
def test_e6_malformed_callback_config_fails_boot_identically(
tmp_path_factory: pytest.TempPathFactory, provider: Wire, sink: Wire
) -> None:
with pytest.raises(AssertionError, match="exited before readiness"):
with _booted_rig(
tmp_path_factory.mktemp("failure_redaction_bad_cb"),
provider,
sink,
{
"callbacks": ["generic_api", "not_a_callback"],
"failure_callback": None,
"turn_off_message_logging": True,
},
):
raise AssertionError("malformed callback config must not boot")
def test_e7_spend_logs_detail_endpoint_redacted(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = _chat(rig, model, [{"role": "user", "content": secret}])
assert response.status_code == 400, response.text
call_id: Final = _call_id(response)
_assert_failure_redacted(rig, model, secret, call_id)
detail: Final = eventually(
lambda: rig.proxy.get(f"/spend/logs/ui/{call_id}"),
lambda value: bool(value),
seconds=30,
)
assert secret not in json.dumps(detail), json.dumps(detail)[:2000]
# --- F. Edge ----------------------------------------------------------------------------
_F1_SLOTS: Final = ("top", "metadata", "litellm_metadata")
def _redaction_body(model: str, secret: str, slot: str, value: JsonValue) -> dict[str, JsonValue]:
body: Final[dict[str, JsonValue]] = {
"model": model,
"messages": [{"role": "user", "content": secret}],
}
if slot == "top":
if value != "MISSING":
body["turn_off_message_logging"] = value
else:
body[slot] = {} if value == "MISSING" else {"turn_off_message_logging": value}
return body
@pytest.mark.parametrize("slot", _F1_SLOTS)
@pytest.mark.parametrize("value", ("MISSING", None, ""), ids=("missing", "null", "empty"))
def test_f1_turn_off_edge_values_global_on(rig: Rig, slot: str, value: JsonValue) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = rig.proxy.request("POST", "/v1/chat/completions", _redaction_body(model, secret, slot, value))
assert response.status_code == 400, response.text
_assert_failure_redacted(rig, model, secret, _call_id(response))
@pytest.mark.parametrize("slot", _F1_SLOTS)
@pytest.mark.parametrize("value", ("MISSING", None, ""), ids=("missing", "null", "empty"))
def test_f1_turn_off_edge_values_global_off(rig_off: Rig, slot: str, value: JsonValue) -> None:
secret: Final = _secret_prompt()
with rig_off.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig_off.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = rig_off.proxy.request(
"POST", "/v1/chat/completions", _redaction_body(model, secret, slot, value)
)
assert response.status_code == 400, response.text
_assert_failure_raw(rig_off, model, secret, _call_id(response))
def test_f2_permitted_opt_out_top_level_false_beats_metadata_true(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model], metadata={"allow_client_message_redaction_opt_out": True})
response: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": secret}],
"turn_off_message_logging": False,
"metadata": {"turn_off_message_logging": True},
},
key=key,
)
assert response.status_code == 400, response.text
_assert_failure_raw(rig, model, secret, _call_id(response))
def test_f3_team_opt_out_permission_flip_takes_effect(rig: Rig) -> None:
first_secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
team: Final = scenario.team(metadata={"allow_client_message_redaction_opt_out": True})
key: Final = scenario.key(models=[model], team_id=team)
allowed: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": first_secret}]},
key=key,
headers={"litellm-disable-message-redaction": "true"},
)
assert allowed.status_code == 400, allowed.text
_assert_failure_raw(rig, model, first_secret, _call_id(allowed))
rig.proxy.post("/team/update", {"team_id": team, "metadata": {}})
def now_redacted() -> tuple[bool, ...]:
secret: Final = _secret_prompt()
fired: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}]},
key=key,
headers={"litellm-disable-message-redaction": "true"},
)
assert fired.status_code == 400, fired.text
events: Final = eventually(
lambda: tuple(
event for event in rig.failure_events(model) if event.get("litellm_call_id") == _call_id(fired)
),
lambda values: len(values) == 1,
seconds=20,
)
return (secret not in json.dumps(events[0]),)
converged: Final = eventually(now_redacted, lambda values: values[0], seconds=120)
assert converged[0]
def test_f4_five_identical_failures_log_once_each(rig: Rig) -> None:
secret: Final = _secret_prompt()
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
call_ids: Final = tuple(_call_id(_chat(rig, model, [{"role": "user", "content": secret}])) for _ in range(5))
assert len(set(call_ids)) == 5
for call_id in call_ids:
eventually(
lambda cid=call_id: tuple(
event for event in rig.failure_events(model) if event.get("litellm_call_id") == cid
),
lambda values: len(values) == 1,
seconds=30,
)
assert _spend_row(call_id)["status"] == "failure", call_id

View file

@ -0,0 +1,289 @@
import json
import signal
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import httpx
import psutil
import pytest
import yaml
from integration._support.client import Gateway, JsonValue, eventually, gateway_from_environment, object_value
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
BURST: Final = 30
def _provider(request: Request) -> Reply:
try:
body: Final = json.loads(request.body)
except json.JSONDecodeError:
return Reply(status=404, body=b"{}")
text: Final = (
body["messages"][-1]["content"] if "messages" in body else str(body.get("input", json.dumps(body)[:200]))
)
return Reply(
status=400,
body=json.dumps(
{"error": {"type": "invalid_request_error", "message": f"Unsupported content: {text}"}}
).encode(),
)
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
process: OwnedProxy
provider: Wire
sink: Wire
model: str
outage: threading.Event
slow: threading.Event
batches: list[Request]
def failure_events(self) -> tuple[dict[str, JsonValue], ...]:
self.batches.extend(self.sink.drain()) # mutable-ok: drain consumes, polls keep earlier batches
return tuple(
object_value(event)
for batch in self.batches
for event in json.loads(batch.body)
if self.model in json.dumps(event)
)
def _bodies(model: str, marker: str) -> tuple[tuple[str, dict[str, JsonValue]], ...]:
chat: Final = tuple(
(
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"burst {marker} {index}"}],
**({"stream": True} if index % 2 else {}),
},
)
for index in range(BURST // 3 * 2)
)
messages: Final = tuple(
(
"/v1/messages",
{
"model": model,
"max_tokens": 16,
"messages": [{"role": "user", "content": f"burst {marker} m{index}"}],
},
)
for index in range(BURST // 6)
)
responses: Final = tuple(
("/v1/responses", {"model": model, "input": f"burst {marker} r{index}"})
for index in range(BURST - len(chat) - len(messages))
)
return chat + messages + responses
def _fire(rig: Rig, bodies: tuple[tuple[str, dict[str, JsonValue]], ...]) -> tuple[tuple[int, str | None], ...]:
def call(item: tuple[str, dict[str, JsonValue]]) -> tuple[int, str | None]:
try:
response: Final = rig.proxy.request("POST", item[0], item[1])
response.read()
return response.status_code, response.headers.get("x-litellm-call-id")
except httpx.HTTPError:
return -1, None
with ThreadPoolExecutor(max_workers=8) as pool:
return tuple(pool.map(call, bodies))
@pytest.mark.timeout(280)
def test_g1_sink_outage_mid_burst_lands_each_call_id_once_redacted(
tmp_path: Path,
) -> None:
marker: Final = uuid.uuid4().hex
outage: Final = threading.Event()
def sink(request: Request) -> Reply:
if outage.is_set():
return Reply(status=503, body=b'{"error":"sink down"}')
return Reply()
with wire_server(_provider) as provider, wire_server(sink) as endpoint:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update(
{"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, "turn_off_message_logging": True}
)
path: Final = tmp_path / "chaos.yaml"
path.write_text(yaml.safe_dump(config))
with (
gateway_from_environment() as gateway,
owned_proxy_process(
gateway, tmp_path, {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, config=path, workers=2
) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
rig: Final = Rig(owned.gateway, owned, provider, endpoint, model, outage, threading.Event(), [])
bodies: Final = _bodies(model, marker)
outage.set()
first_half: Final = _fire(rig, bodies[: BURST // 2])
outage.clear()
second_half: Final = _fire(rig, bodies[BURST // 2 :])
outcomes: Final = first_half + second_half
answered: Final = tuple(outcome for outcome in outcomes if outcome[0] >= 0)
assert all(status == 400 for status, _ in answered), outcomes
events: Final = eventually(rig.failure_events, lambda values: len(values) >= len(second_half), seconds=70)
seen: Final = tuple(str(event.get("litellm_call_id")) for event in events)
assert len(seen) == len(set(seen)), ("duplicate failure events", seen)
for event in events:
assert marker not in json.dumps(event), json.dumps(event)[:400]
@pytest.mark.timeout(280)
def test_g2_slow_sink_no_deadlock_no_duplicates(tmp_path: Path) -> None:
marker: Final = uuid.uuid4().hex
slow: Final = threading.Event()
def sink(request: Request) -> Reply:
if slow.is_set():
time.sleep(1)
return Reply()
with wire_server(_provider) as provider, wire_server(sink) as endpoint:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update(
{"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, "turn_off_message_logging": True}
)
path: Final = tmp_path / "chaos_slow.yaml"
path.write_text(yaml.safe_dump(config))
with (
gateway_from_environment() as gateway,
owned_proxy_process(
gateway, tmp_path, {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, config=path, workers=2
) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
rig: Final = Rig(owned.gateway, owned, provider, endpoint, model, threading.Event(), slow, [])
slow.set()
bodies: Final = _bodies(model, marker)[:6]
outcomes: Final = _fire(rig, bodies)
assert all(status == 400 for status, _ in outcomes), outcomes
call_ids: Final = tuple(cid for _, cid in outcomes if cid)
events: Final = eventually(rig.failure_events, lambda values: len(values) >= len(bodies), seconds=120)
landed: Final = tuple(str(event.get("litellm_call_id")) for event in events)
assert len(landed) == len(set(landed)), ("duplicate failure events", landed)
for event in events:
assert marker not in json.dumps(event), json.dumps(event)[:400]
assert set(call_ids) <= set(landed), (call_ids, landed)
@pytest.mark.timeout(280)
def test_g3_proxy_restart_mid_burst(tmp_path: Path) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(_provider) as provider, wire_server(lambda _: Reply()) as endpoint:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update(
{"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, "turn_off_message_logging": True}
)
path: Final = tmp_path / "chaos_restart.yaml"
path.write_text(yaml.safe_dump(config))
overrides: Final = {"GENERIC_LOGGER_ENDPOINT": endpoint.url}
with gateway_from_environment() as gateway:
bodies: Final = _bodies("restart-model", marker)
with owned_proxy_process(gateway, tmp_path, overrides, config=path, workers=2) as owned_one:
owned_one.gateway.post(
"/model/new",
{
"model_name": "restart-model",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": provider.url + "/v1",
"api_key": "synthetic-provider-key",
},
},
)
first: Final = _fire(
Rig(
owned_one.gateway,
owned_one,
provider,
endpoint,
"restart-model",
threading.Event(),
threading.Event(),
[],
),
bodies[: BURST // 2],
)
with owned_proxy_process(gateway, tmp_path, overrides, config=path, workers=2) as owned_two:
rig_two: Final = Rig(
owned_two.gateway,
owned_two,
provider,
endpoint,
"restart-model",
threading.Event(),
threading.Event(),
[],
)
def served() -> tuple[int, ...]:
probe: Final = rig_two.proxy.request("POST", *bodies[BURST // 2])
return (probe.status_code,)
eventually(served, lambda statuses: statuses[0] == 400, seconds=60)
second: Final = _fire(rig_two, bodies[BURST // 2 :])
answered: Final = first + second
assert all(status in (400, 500) for status, _ in answered), answered
assert any(status == 400 for status, _ in answered), answered
events: Final = tuple(object_value(event) for batch in endpoint.drain() for event in json.loads(batch.body))
call_ids: Final = tuple(str(event.get("litellm_call_id")) for event in events if event)
assert len(call_ids) == len(set(call_ids)), ("duplicate events after restart", call_ids)
for event in events:
assert marker not in json.dumps(event), json.dumps(event)[:400]
@pytest.mark.timeout(280)
def test_g4_worker_kill_keeps_serving_redacted(tmp_path: Path) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(_provider) as provider, wire_server(lambda _: Reply()) as endpoint:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update(
{"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, "turn_off_message_logging": True}
)
path: Final = tmp_path / "chaos_worker.yaml"
path.write_text(yaml.safe_dump(config))
with (
gateway_from_environment() as gateway,
owned_proxy_process(
gateway, tmp_path, {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, config=path, workers=2
) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
rig: Final = Rig(owned.gateway, owned, provider, endpoint, model, threading.Event(), threading.Event(), [])
bodies: Final = _bodies(model, marker)[:12]
children: Final = psutil.Process(owned.process.pid).children(recursive=True)
assert children, "no uvicorn worker children found"
with ThreadPoolExecutor(max_workers=8) as pool:
futures: Final = tuple(
pool.submit(lambda b: rig.proxy.request("POST", b[0], b[1]), body) for body in bodies
)
eventually(lambda: provider.received.qsize() >= 3, bool, seconds=30)
children[0].send_signal(signal.SIGKILL)
statuses: list[int] = [] # mutable-ok: collect per-request outcomes from concurrent futures
for future in futures:
try:
statuses.append(future.result().status_code)
except httpx.HTTPError:
statuses.append(-1)
assert all(status == 400 for status in statuses if status >= 0), statuses
events: Final = eventually(rig.failure_events, lambda values: len(values) >= 1, seconds=70)
landed: Final = tuple(str(event.get("litellm_call_id")) for event in events)
assert len(landed) == len(set(landed)), ("duplicate failure events", landed)
for event in events:
assert marker not in json.dumps(event), json.dumps(event)[:400]

View file

@ -0,0 +1,90 @@
import gzip
import json
import uuid
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, JsonValue, eventually, gateway_from_environment, object_value
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
def _body(batch: Request) -> bytes:
return gzip.decompress(batch.body) if batch.body[:2] == b"\x1f\x8b" else batch.body
def _provider(request: Request) -> Reply:
try:
body: Final = json.loads(request.body)
except json.JSONDecodeError:
return Reply(status=404, body=b"{}")
text: Final = body["messages"][-1]["content"]
return Reply(
status=400,
body=json.dumps(
{"error": {"type": "invalid_request_error", "message": f"Unsupported content: {text}"}}
).encode(),
)
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
provider: Wire
sink: Wire
def log_entries(self, model: str) -> tuple[dict[str, JsonValue], ...]:
return tuple(
object_value(entry)
for batch in self.sink.drain()
for entry in json.loads(_body(batch))
if model in json.dumps(entry)
)
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
root: Final = tmp_path_factory.mktemp("failure_redaction_datadog")
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update(
{"callbacks": ["datadog"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, "turn_off_message_logging": True}
)
path: Final = root / "datadog_failure.yaml"
path.write_text(yaml.safe_dump(config))
with (
wire_server(_provider) as provider,
wire_server(lambda _: Reply()) as sink,
gateway_from_environment() as gateway,
owned_proxy(
gateway,
root,
{"DD_API_KEY": "synthetic-dd-key", "DD_BASE_URL": sink.url, "DD_SITE": "localhost"},
config=path,
workers=2,
) as proxy,
):
yield Rig(proxy, provider, sink)
def test_c5_datadog_failure_log_redacted_keeps_status(rig: Rig) -> None:
secret: Final = "dd-secret-" + uuid.uuid4().hex
with rig.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}]},
)
assert response.status_code == 400, response.text
assert any(secret.encode() in request.body for request in rig.provider.drain())
entries: Final = eventually(lambda: rig.log_entries(model), lambda values: len(values) >= 1, seconds=30)
entry: Final = entries[0]
assert secret not in json.dumps(entry), json.dumps(entry)[:2000]
message: Final = object_value(json.loads(str(entry.get("message", "{}"))))
assert secret not in json.dumps(message), json.dumps(message)[:2000]
error_information: Final = object_value(message.get("error_information"))
assert error_information.get("error_class"), error_information

View file

@ -0,0 +1,313 @@
import json
import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from google.protobuf.json_format import MessageToDict
from integration._support.client import Gateway, JsonValue, eventually, gateway_from_environment, object_value
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
def _provider(request: Request) -> Reply:
try:
body: Final = json.loads(request.body)
except json.JSONDecodeError:
return Reply(status=404, body=b"{}")
text: Final = body["messages"][-1]["content"]
return Reply(
status=400,
body=json.dumps(
{"error": {"type": "invalid_request_error", "message": f"Unsupported content: {text}"}}
).encode(),
)
def _decode(body: bytes) -> dict[str, JsonValue]:
if body[:1] == b"{":
return json.loads(body)
request: Final = ExportTraceServiceRequest()
request.ParseFromString(body)
return object_value(MessageToDict(request))
@dataclass(frozen=True, slots=True)
class Spans:
wire: Wire
batches: list[Request]
def all(self) -> tuple[dict[str, JsonValue], ...]:
self.batches.extend(self.wire.drain())
return tuple(
span
for batch in self.batches
for resource in _decode(batch.body).get("resourceSpans", ())
for scope in resource.get("scopeSpans", ())
for span in scope.get("spans", ())
)
def named(self, name: str, model: str) -> tuple[dict[str, JsonValue], ...]:
return tuple(span for span in self.all() if span.get("name") == name and model in json.dumps(span))
def in_trace(self, name: str, trace_id: str) -> tuple[dict[str, JsonValue], ...]:
return tuple(span for span in self.all() if span.get("name") == name and span.get("traceId") == trace_id)
def _span_attributes(span: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
return {
str(attribute["key"]): object_value(attribute["value"]).get("stringValue")
or object_value(attribute["value"]).get("intValue")
for attribute in span.get("attributes", ())
if isinstance(attribute, dict)
}
def _exception_events(span: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
return tuple(
object_value(event)
for event in span.get("events", ())
if isinstance(event, dict) and event.get("name") == "exception"
)
def _exception_texts(span: Mapping[str, JsonValue]) -> str:
return json.dumps(_exception_events(span))
def _error_attribute(span: Mapping[str, JsonValue]) -> str:
attributes: Final = _span_attributes(span)
return str(attributes.get("error.message", ""))
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
provider: Wire
sink: Spans
@contextmanager
def _otel_rig(root: Path, provider: Wire, sink: Wire, v2: bool, global_on: bool) -> Iterator[Rig]:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
settings: Final[dict[str, JsonValue]] = {"callbacks": ["otel"]}
if global_on:
settings["turn_off_message_logging"] = True
config["litellm_settings"].update(settings)
if v2:
config["callback_settings"] = {
"otel": {"exporter": "http/json", "endpoint": sink.url, "mapper_names": ["genai"]}
}
path: Final = root / "otel_failure.yaml"
path.write_text(yaml.safe_dump(config))
env: Final = (
{"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"}
if v2
else {
"OTEL_EXPORTER": "http/json",
"OTEL_EXPORTER_OTLP_ENDPOINT": sink.url,
"OTEL_BSP_SCHEDULE_DELAY": "300",
}
)
with (
gateway_from_environment() as gateway,
owned_proxy(gateway, root, env, config=path, workers=2) as proxy,
):
yield Rig(proxy, provider, Spans(sink, [])) # mutable-ok: drain consumes, polls keep earlier batches
@pytest.fixture(scope="module")
def provider() -> Iterator[Wire]:
with wire_server(_provider) as wire:
yield wire
@pytest.fixture(scope="module")
def sink() -> Iterator[Wire]:
with wire_server(lambda _: Reply()) as wire:
yield wire
@pytest.fixture(scope="module")
def rig_v1_on(tmp_path_factory: pytest.TempPathFactory, provider: Wire, sink: Wire) -> Iterator[Rig]:
with _otel_rig(tmp_path_factory.mktemp("otel_v1_on"), provider, sink, v2=False, global_on=True) as booted:
yield booted
@pytest.fixture(scope="module")
def rig_v2_on(tmp_path_factory: pytest.TempPathFactory, provider: Wire, sink: Wire) -> Iterator[Rig]:
with _otel_rig(tmp_path_factory.mktemp("otel_v2_on"), provider, sink, v2=True, global_on=True) as booted:
yield booted
@pytest.fixture(scope="module")
def rig_v2_off(tmp_path_factory: pytest.TempPathFactory, provider: Wire, sink: Wire) -> Iterator[Rig]:
with _otel_rig(tmp_path_factory.mktemp("otel_v2_off"), provider, sink, v2=True, global_on=False) as booted:
yield booted
def _secret() -> str:
return "otel-secret-" + uuid.uuid4().hex
def _fail(rig: Rig, model: str, secret: str, **kwargs: JsonValue) -> httpx.Response:
headers: Final = kwargs.pop("headers", None)
key: Final = kwargs.pop("key", None)
return rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": secret}], **kwargs},
headers=headers if isinstance(headers, dict) else None,
key=key if isinstance(key, str) else None,
)
def _llm_span(rig: Rig, call_id: str) -> dict[str, JsonValue]:
def found() -> tuple[dict[str, JsonValue], ...]:
return tuple(
span
for span in rig.sink.all()
if call_id in json.dumps(span) and str(span.get("name", "")).startswith(("chat ", "litellm_request"))
)
return eventually(found, lambda values: len(values) >= 1, seconds=260)[0]
_SERVER_SPAN_NAMES: Final = ("Received Proxy Server Request", "POST /v1/chat/completions", "POST /v1/messages")
def _server_span(rig: Rig, call_id: str) -> dict[str, JsonValue]:
trace_id: Final = str(_llm_span(rig, call_id)["traceId"])
def found() -> tuple[dict[str, JsonValue], ...]:
return tuple(
span
for span in rig.sink.all()
if span.get("traceId") == trace_id and str(span.get("name", "")) in _SERVER_SPAN_NAMES
)
spans: Final = eventually(found, lambda values: len(values) >= 1, seconds=60, return_last_on_timeout=True)
assert spans, [(span.get("name"), span.get("traceId")) for span in rig.sink.all()]
return spans[0]
def _span_for(rig: Rig, name: str, model: str) -> dict[str, JsonValue]:
spans: Final = eventually(lambda: rig.sink.named(name, model), lambda values: len(values) >= 1, seconds=260)
return spans[0]
def _auth_exception_span(rig: Rig) -> dict[str, JsonValue]:
def found() -> tuple[dict[str, JsonValue], ...]:
return tuple(
span for span in rig.sink.all() if str(span.get("name", "")).startswith("auth") and _exception_events(span)
)
return eventually(found, lambda values: len(values) >= 1, seconds=260)[0]
# C6: OTEL v1 failure span redaction under global on
@pytest.mark.timeout(320)
def test_c6_v1_provider_error_spans_redacted(rig_v1_on: Rig) -> None:
secret: Final = _secret()
with rig_v1_on.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig_v1_on.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = _fail(rig_v1_on, model, secret)
assert response.status_code == 400, response.text
call_id: Final = response.headers["x-litellm-call-id"]
assert any(secret.encode() in request.body for request in rig_v1_on.provider.drain())
server: Final = _server_span(rig_v1_on, call_id)
assert _error_attribute(server) == "redacted-by-litellm", _span_attributes(server)
assert secret not in _exception_texts(server), _exception_texts(server)[:600]
request_span: Final = _llm_span(rig_v1_on, call_id)
assert _error_attribute(request_span) == "redacted-by-litellm", _span_attributes(request_span)
assert secret not in _exception_texts(request_span), _exception_texts(request_span)[:600]
# C7: OTEL v2 request + server spans redacted under global on
@pytest.mark.timeout(320)
def test_c7_v2_provider_error_spans_redacted(rig_v2_on: Rig) -> None:
secret: Final = _secret()
with rig_v2_on.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig_v2_on.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = _fail(rig_v2_on, model, secret)
assert response.status_code == 400, response.text
call_id: Final = response.headers["x-litellm-call-id"]
server: Final = _server_span(rig_v2_on, call_id)
assert _error_attribute(server) == "redacted-by-litellm", _span_attributes(server)
assert secret not in _exception_texts(server), _exception_texts(server)[:600]
# C8: v2 server span restamp keeps request opt-in under global off
@pytest.mark.timeout(320)
def test_c8_v2_header_opt_in_restamped_server_span_redacted(rig_v2_off: Rig) -> None:
secret: Final = _secret()
with rig_v2_off.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig_v2_off.provider.url + "/v1", api_key="synthetic-provider-key")
response: Final = _fail(rig_v2_off, model, secret, headers={"x-litellm-enable-message-redaction": "true"})
assert response.status_code == 400, response.text
call_id: Final = response.headers["x-litellm-call-id"]
server: Final = _server_span(rig_v2_off, call_id)
assert _error_attribute(server) == "redacted-by-litellm", _span_attributes(server)
assert secret not in _exception_texts(server), _exception_texts(server)[:600]
# C9: v2 server span keeps permitted opt-out raw under global on
@pytest.mark.timeout(320)
def test_c9_v2_permitted_opt_out_keeps_server_span_raw(rig_v2_on: Rig) -> None:
secret: Final = _secret()
with rig_v2_on.proxy.scenario() as scenario:
model: Final = scenario.model(api_base=rig_v2_on.provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model], metadata={"allow_client_message_redaction_opt_out": True})
response: Final = _fail(
rig_v2_on, model, secret, key=key, headers={"litellm-disable-message-redaction": "true"}
)
assert response.status_code == 400, response.text
call_id: Final = response.headers["x-litellm-call-id"]
server: Final = _server_span(rig_v2_on, call_id)
assert secret in _error_attribute(server), _span_attributes(server)
# C10: v2 auth phase span honors request opt-in under global off
@pytest.mark.timeout(320)
def test_c10_v2_auth_span_honors_header_opt_in(rig_v2_off: Rig) -> None:
with rig_v2_off.proxy.scenario() as scenario:
allowed: Final = scenario.model(api_base=rig_v2_off.provider.url + "/v1", api_key="synthetic-provider-key")
denied: Final = scenario.model(api_base=rig_v2_off.provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[allowed])
response: Final = rig_v2_off.proxy.request(
"POST",
"/v1/chat/completions",
{"model": denied, "messages": [{"role": "user", "content": "hi"}]},
key=key,
headers={"x-litellm-enable-message-redaction": "true"},
)
assert response.status_code in (400, 401, 403, 404), response.text
auth: Final = _auth_exception_span(rig_v2_off)
assert allowed not in _exception_texts(auth), _exception_texts(auth)[:600]
assert "redacted-by-litellm" in _exception_texts(auth), _exception_texts(auth)[:600]
# C11: v2 auth phase span redacts under global on even with unpermitted disable header
@pytest.mark.timeout(320)
def test_c11_v2_auth_span_redacts_under_global_on(rig_v2_on: Rig) -> None:
with rig_v2_on.proxy.scenario() as scenario:
allowed: Final = scenario.model(api_base=rig_v2_on.provider.url + "/v1", api_key="synthetic-provider-key")
denied: Final = scenario.model(api_base=rig_v2_on.provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[allowed])
response: Final = rig_v2_on.proxy.request(
"POST",
"/v1/chat/completions",
{"model": denied, "messages": [{"role": "user", "content": "hi"}]},
key=key,
headers={"litellm-disable-message-redaction": "true"},
)
assert response.status_code in (400, 401, 403, 404), response.text
auth: Final = _auth_exception_span(rig_v2_on)
assert allowed not in _exception_texts(auth), _exception_texts(auth)[:600]
assert "redacted-by-litellm" in _exception_texts(auth), _exception_texts(auth)[:600]