import hashlib import json import os import re from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from itertools import chain from pathlib import Path from types import MappingProxyType from typing import Final, TypeAlias import yaml from integration._support.anthropic_thinking import ( JSON_LIST, JSON_OBJECT, Event, marker_of, stream_reply, streams, ) from integration._support.client import eventually from integration._support.wire import Reply, Request, Wire from pydantic import JsonValue, TypeAdapter from redis import Redis DOCS_REFUSAL_URL: Final = "https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback#refusal-response" DOCS_STREAMING_REFUSAL_URL: Final = ( "https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals" ) # Copied verbatim from DOCS_REFUSAL_URL and DOCS_STREAMING_REFUSAL_URL on 2026-10-12. DOCS_REFUSAL: Final[Event] = { "id": "msg_01XFUDYJgAACzvnptvVoYEL", "type": "message", "role": "assistant", "model": "claude-fable-5", "content": [], "stop_reason": "refusal", "stop_details": { "type": "refusal", "category": "cyber", "explanation": "This request was declined because it could enable cyber harm.", }, "usage": {"input_tokens": 412, "output_tokens": 0}, } DOCS_PARTIAL_REFUSAL: Final[Event] = { "role": "assistant", "content": [{"type": "text", "text": "Hello.."}], "stop_reason": "refusal", "stop_details": { "type": "refusal", "category": "cyber", "explanation": "This request was declined because it could enable cyber harm.", }, } # Captured from api.anthropic.com (claude-sonnet-5-5, prompt "Reply with the single word: ready") on 2026-10-13. CAPTURED_TURN: Final[Event] = { "model": "claude-sonnet-5-5", "id": "msg_011CfyXFGze1cDkk8o7gDRRS", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "ready"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "usage": { "input_tokens": 18, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, "output_tokens": 4, "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", }, "diagnostics": None, } CAPTURED_EVENTS: Final[tuple[Event, ...]] = ( { "type": "message_start", "message": { "model": "claude-sonnet-5-5", "id": "msg_011CfyXFLwXUxsp7nNAm3gcu", "type": "message", "role": "assistant", "content": [], "container": None, "stop_reason": None, "stop_sequence": None, "stop_details": None, "usage": { "input_tokens": 18, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, "output_tokens": 4, "service_tier": "standard", "inference_geo": "global", }, "diagnostics": None, }, }, {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, {"type": "ping"}, {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ready"}}, {"type": "content_block_stop", "index": 0}, { "type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, "usage": { "input_tokens": 18, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 4, "output_tokens_details": {"thinking_tokens": 0}, }, }, {"type": "message_stop"}, ) DOCS_STOP_DETAILS: Final = DOCS_REFUSAL["stop_details"] REFUSAL_FIELDS: Final[Event] = {"stop_reason": "refusal", "stop_details": DOCS_STOP_DETAILS} STOP_FIELDS: Final = ("stop_reason", "stop_details") LONG_STOP_DETAILS: Final = "d" * 5120 MALFORMED_STOP_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType( {"int": 7, "list": ["refusal"], "empty": "", "long": LONG_STOP_DETAILS} ) _FIXTURE: Final = re.compile(r"refusal-fixture-([a-z]+)") Location: TypeAlias = tuple[str | int, ...] @dataclass(frozen=True, slots=True) class Fixture: status: int body: bytes events: tuple[Event, ...] = () def _text_block(index: int, block: JsonValue) -> tuple[Event, ...]: text: Final = JSON_OBJECT.validate_python(block)["text"] return ( {"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}}, {"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}}, {"type": "content_block_stop", "index": index}, ) def refusal_events(message: Mapping[str, JsonValue]) -> tuple[Event, ...]: usage: Final = JSON_OBJECT.validate_python(message["usage"]) start: Final[Event] = { "type": "message_start", "message": { "id": message["id"], "type": "message", "role": "assistant", "model": message["model"], "content": [], "stop_reason": None, "stop_sequence": None, "stop_details": None, "usage": usage, }, } blocks: Final = chain.from_iterable( _text_block(index, block) for index, block in enumerate(JSON_LIST.validate_python(message["content"])) ) delta: Final[Event] = { "type": "message_delta", "delta": { "stop_reason": message["stop_reason"], "stop_sequence": None, "stop_details": message["stop_details"], }, "usage": {name: usage[name] for name in ("output_tokens", "output_tokens_details") if name in usage}, } return (start, *blocks, delta, {"type": "message_stop"}) def _without_stop_details(value: Mapping[str, JsonValue]) -> Event: return {name: item for name, item in value.items() if name != "stop_details"} def _captured_event(event: Event, stop_details: JsonValue, *, keep: bool) -> Event: if event["type"] == "message_start": message: Final = JSON_OBJECT.validate_python(event["message"]) return {**event, "message": message if keep else _without_stop_details(message)} if event["type"] != "message_delta": return event delta: Final = JSON_OBJECT.validate_python(event["delta"]) replaced: Final = {**delta, "stop_details": stop_details} if keep else _without_stop_details(delta) return {**event, "delta": replaced} def _captured(stop_details: JsonValue, *, keep: bool = True) -> Fixture: body: Final = {**CAPTURED_TURN, "stop_details": stop_details} if keep else _without_stop_details(CAPTURED_TURN) return Fixture( 200, json.dumps(body).encode(), tuple(_captured_event(event, stop_details, keep=keep) for event in CAPTURED_EVENTS), ) def _duplicated_stop_details() -> bytes: pair: Final = f'"stop_details": {json.dumps(DOCS_STOP_DETAILS)}' return json.dumps(DOCS_REFUSAL).replace(pair, f"{pair}, {pair}").encode() def _anthropic_error(kind: str, message: str) -> bytes: return json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode() _PARTIAL: Final[Event] = {**DOCS_REFUSAL, **DOCS_PARTIAL_REFUSAL} _REASONED: Final[Event] = { **_PARTIAL, "usage": {"input_tokens": 412, "output_tokens": 10, "output_tokens_details": {"thinking_tokens": 4}}, } JSON_TOOL_NAME: Final = "json_tool_call" JSON_ANSWER_SCHEMA: Final[Event] = { "type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"], "additionalProperties": False, } JSON_ANSWER_FORMAT: Final[Event] = { "type": "json_schema", "json_schema": {"name": "answer", "schema": JSON_ANSWER_SCHEMA, "strict": True}, } DOCS_PARTIAL_TEXT: Final = "Hello.." _DECLARED_INPUT: Final[Event] = {"answer": DOCS_PARTIAL_TEXT} _ENVELOPED_INPUT: Final[Event] = {"parameters": _DECLARED_INPUT} DECLARED_PARTIAL_JSON: Final = '{"answer": "Hello..' ENVELOPED_PARTIAL_JSON: Final = '{"parameters": {"answer": "Hello..' _JSON_TOOL_USE_ID: Final = "toolu_01RefusalFixtureJsonTool" def _json_tool_refusal(tool_input: Event, partial_json: str) -> Fixture: message: Final[Event] = { **DOCS_REFUSAL, "content": [{"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": tool_input}], } start, delta, stop = refusal_events(DOCS_REFUSAL) tool_block: Final[tuple[Event, ...]] = ( { "type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": {}}, }, { "type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": partial_json}, }, {"type": "content_block_stop", "index": 0}, ) return Fixture(200, json.dumps(message).encode(), (start, *tool_block, delta, stop)) FIXTURES: Final[Mapping[str, Fixture]] = MappingProxyType( { "refusal": Fixture(200, json.dumps(DOCS_REFUSAL).encode(), refusal_events(DOCS_REFUSAL)), "partial": Fixture(200, json.dumps(_PARTIAL).encode(), refusal_events(_PARTIAL)), "reasoned": Fixture(200, json.dumps(_REASONED).encode(), refusal_events(_REASONED)), "jsontool": _json_tool_refusal(_DECLARED_INPUT, DECLARED_PARTIAL_JSON), "jsonheld": _json_tool_refusal(_ENVELOPED_INPUT, ENVELOPED_PARTIAL_JSON), "duplicate": Fixture(200, _duplicated_stop_details()), "null": _captured(None), "missing": _captured(None, keep=False), **{name: _captured(value) for name, value in MALFORMED_STOP_DETAILS.items()}, "nondict": Fixture(200, b"[]"), "unauthorized": Fixture(401, _anthropic_error("authentication_error", "invalid x-api-key")), "invalid": Fixture(400, _anthropic_error("invalid_request_error", "max_tokens: Field required")), } ) def refusal_prompt(fixture: str, marker: str) -> str: return f"refusal-fixture-{fixture} marker-{marker}" CITED_FIELDS: Final[Event] = {"citations": ["https://example.com/source"]} _PLAIN_DELTA: Final[Event] = {"role": "assistant", "content": "ready"} _CITED_DELTA: Final[Event] = {**_PLAIN_DELTA, "provider_specific_fields": CITED_FIELDS} _OPENAI_USAGE: Final[Event] = { "prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12, "completion_tokens_details": {"reasoning_tokens": 4, "text_tokens": 3}, } UNVALIDATED_USAGE_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType( { "text": {"reasoning_tokens": "abc"}, "fraction": {"reasoning_tokens": 1.5}, "object": {"reasoning_tokens": {"a": 1}}, } ) @dataclass(frozen=True, slots=True) class OpenAIFixture: first_delta: Event usage: tuple[Event, ...] = () _OPENAI_FIXTURES: Final[Mapping[str, OpenAIFixture]] = MappingProxyType( { "plain": OpenAIFixture(_PLAIN_DELTA), "cited": OpenAIFixture(_CITED_DELTA), "plainusage": OpenAIFixture(_PLAIN_DELTA, (_OPENAI_USAGE,)), "citedusage": OpenAIFixture(_CITED_DELTA, (_OPENAI_USAGE,)), } ) def openai_id(marker: str) -> str: return f"chatcmpl-{marker}" def _openai_frame(marker: str, body: Event) -> bytes: chunk: Final[Event] = { "id": openai_id(marker), "object": "chat.completion.chunk", "created": 1760000000, "model": "gpt-4o-mini", **body, } return b"data: " + json.dumps(chunk).encode() + b"\n\n" def _openai_choice(delta: Event, finish_reason: str | None) -> Event: return {"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}]} def _openai_stream(name: str, request: Request) -> Reply: assert streams(request), request.body marker: Final = marker_of(request) fixture: Final = _OPENAI_FIXTURES[name] return Reply( content_type="text/event-stream", chunks=( _openai_frame(marker, _openai_choice(fixture.first_delta, None)), _openai_frame(marker, _openai_choice({"content": " now"}, "stop")), *(_openai_frame(marker, {"choices": [], "usage": usage}) for usage in fixture.usage), b"data: [DONE]\n\n", ), ) def respond(request: Request) -> Reply: found: Final[Sequence[str]] = _FIXTURE.findall(request.body.decode()) assert found, request.body if request.target.endswith("/chat/completions"): return _openai_stream(found[-1], request) fixture: Final = FIXTURES[found[-1]] if fixture.status != 200 or not fixture.events or not streams(request): return Reply(status=fixture.status, body=fixture.body) return stream_reply(request, fixture.events) def _stop_fields(fields: JsonValue) -> Event: if not isinstance(fields, dict): return {} return {name: fields[name] for name in STOP_FIELDS if name in fields} def relayed(value: JsonValue, location: Location = ()) -> tuple[tuple[Location, Event], ...]: if isinstance(value, list): return tuple(chain.from_iterable(relayed(item, (*location, index)) for index, item in enumerate(value))) if not isinstance(value, dict): return () found: Final = _stop_fields(value.get("provider_specific_fields")) here: Final = (((*location, "provider_specific_fields"), found),) if found else () return (*here, *chain.from_iterable(relayed(item, (*location, name)) for name, item in value.items())) _CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") _STORED: Final = TypeAdapter(bytes | None) STALE_ID_SUFFIX: Final = "-stale" def _stored(store: Redis, key: str) -> bytes: if not _CACHE_KEY.fullmatch(key) or store.type(key) != b"string": return b"" return _STORED.validate_python(store.get(key)) or b"" def _entry_key(store: Redis, needle: str) -> str | None: keys: Final = (raw.decode() for raw in store.scan_iter(count=1000)) return next((key for key in keys if needle.encode() in _stored(store, key)), None) def _with_message_fields(response: Event, fields: JsonValue) -> Event: choice: Final = JSON_OBJECT.validate_python(JSON_LIST.validate_python(response["choices"])[0]) message: Final = JSON_OBJECT.validate_python(choice["message"]) return {**response, "choices": [{**choice, "message": {**message, "provider_specific_fields": fields}}]} def _with_usage_details(response: Event, details: JsonValue) -> Event: usage: Final = JSON_OBJECT.validate_python(response["usage"]) return {**response, "usage": {**usage, "completion_tokens_details": details}} def _rewrite_cached_response(needle: str, rewrite: Callable[[Event], Event]) -> str: store: Final = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) key: Final = eventually(lambda: _entry_key(store, needle), lambda found: found is not None) assert key is not None, needle entry: Final = JSON_OBJECT.validate_json(_stored(store, key)) stored: Final = entry["response"] encoded: Final = isinstance(stored, str) response: Final = ( JSON_OBJECT.validate_json(stored) if isinstance(stored, str) else JSON_OBJECT.validate_python(stored) ) rewritten: Final = {**rewrite(response), "id": f"{response['id']}{STALE_ID_SUFFIX}"} store.set(key, json.dumps({**entry, "response": json.dumps(rewritten) if encoded else rewritten}), ex=60) return str(rewritten["id"]) def rewrite_cached_message_fields(needle: str, fields: JsonValue) -> str: return _rewrite_cached_response(needle, lambda response: _with_message_fields(response, fields)) def rewrite_cached_usage_details(needle: str, details: JsonValue) -> str: return _rewrite_cached_response(needle, lambda response: _with_usage_details(response, details)) def key_hash(key: str) -> str: return hashlib.sha256(key.encode()).hexdigest() def _key_hash_of(event: Mapping[str, JsonValue]) -> JsonValue: return JSON_OBJECT.validate_python(event.get("metadata") or {}).get("user_api_key_hash") class Ledger: def __init__(self, wire: Wire) -> None: self.wire: Final = wire self.batches: tuple[Request, ...] = () def events(self) -> tuple[Event, ...]: self.batches = (*self.batches, *self.wire.drain()) decoded: Final = chain.from_iterable(JSON_LIST.validate_json(batch.body) for batch in self.batches) return tuple(JSON_OBJECT.validate_python(event) for event in decoded) def for_key(self, key: str) -> tuple[Event, ...]: hashed: Final = key_hash(key) return tuple(event for event in self.events() if _key_hash_of(event) == hashed) def await_key(self, key: str, count: int = 1, seconds: float = 30) -> tuple[Event, ...]: return eventually(lambda: self.for_key(key), lambda found: len(found) >= count, seconds=seconds) _HIDDEN_PARAMS_PROBE: Final = """ import httpx from litellm.integrations.custom_logger import CustomLogger class HiddenParamsProbe(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): payload = kwargs.get("standard_logging_object") or {{}} hidden = getattr(response_obj, "_hidden_params", None) or {{}} event = {{ "metadata": {{"user_api_key_hash": (payload.get("metadata") or {{}}).get("user_api_key_hash")}}, "response_id": getattr(response_obj, "id", None), "received_model_id": hidden.get("received_model_id"), }} async with httpx.AsyncClient(trust_env=False) as client: await client.post({url!r}, json=[event]) instance = HiddenParamsProbe() """ def hidden_params_probe(directory: Path, url: str) -> str: (directory / "hidden_params_probe.py").write_text(_HIDDEN_PARAMS_PROBE.format(url=url)) return "hidden_params_probe.instance" def sink_config( directory: Path, sinks: Mapping[str, str], *, retries: int = 0, extra_callbacks: Sequence[str] = (), store_prompts: bool = False, ) -> Path: shared: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) settings: Final = { **JSON_OBJECT.validate_python(shared["litellm_settings"]), "callbacks": [*sinks, *extra_callbacks], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, } callback_settings: Final = { name: { "callback_type": "generic_api", "endpoint": url, "headers": {"x-refusal-sink": name}, "max_retries": retries, "retry_delay": 0.5, } for name, url in sinks.items() } general_settings: Final = { **JSON_OBJECT.validate_python(shared["general_settings"]), "store_prompts_in_spend_logs": store_prompts, } path: Final = directory / "anthropic-refusal-sinks.yaml" path.write_text( yaml.safe_dump( { **shared, "general_settings": general_settings, "litellm_settings": settings, "callback_settings": callback_settings, } ) ) return path def frames(text: str) -> tuple[Event, ...]: return tuple( JSON_OBJECT.validate_json(line.removeprefix("data: ")) for line in text.splitlines() if line.startswith("data: {") ) def frame_types(events: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: return tuple(event.get("type") for event in events)