test(bedrock): audit the native GPT chat completions route on the integration rig

Adds the /audit cells for the runtime chat completions route: the scripted Bedrock runtime peer, the happy and fallback wire tests, the sad-path and regex worst-case tests, the chaos burst tests, the Messages adapter tests, and the Responses native-route tests. Tests only, no product diff.
This commit is contained in:
mateo-berri 2026-10-01 16:41:03 -07:00
parent 578f26edba
commit 4ee7ff4315
6 changed files with 1953 additions and 0 deletions

View file

@ -0,0 +1,276 @@
import json
import re
import threading
from collections.abc import Mapping
from multiprocessing.sharedctypes import Synchronized
from types import MappingProxyType
from typing import Final
from urllib.parse import unquote
from integration._support.upstream import _aws_event_frame
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
REASONING_EFFORTS: Final = frozenset(("none", "minimal", "low", "medium", "high", "xhigh"))
NATIVE_CHAT: Final = "/openai/v1/chat/completions"
NATIVE_RESPONSES: Final = "/openai/v1/responses"
PNG_1X1: Final = bytes.fromhex(
"89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489"
"0000000d49444154789c63f8cfc0f01f00050001ff89993d1d0000000049454e44ae426082"
)
USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"prompt_tokens": 9,
"completion_tokens": 5,
"total_tokens": 14,
"completion_tokens_details": {"reasoning_tokens": 3},
}
)
_STATUS: Final = re.compile(r"status=(\d{3})")
_CONVERSE: Final = re.compile(r"^/model/(.+)/converse$")
_CONVERSE_STREAM: Final = re.compile(r"^/model/(.+)/converse-stream$")
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_NO_MARKER: Final = "0" * 32
def marker_of(request: Request) -> str:
found: Final = MARKER.search(request.body.decode(errors="replace"))
return _NO_MARKER if found is None else found.group(1)
def body_of(request: Request) -> Mapping[str, JsonValue]:
try:
return _JSON_OBJECT.validate_json(request.body)
except ValueError:
return {}
def target_of(request: Request) -> str:
return unquote(request.target)
def answer(marker: str) -> str:
return f"answer marker-{marker}"
def reasoning_answer(marker: str) -> str:
return f"<reasoning>why marker-{marker}</reasoning> {answer(marker)}"
def _headers(marker: str) -> Mapping[str, str]:
return MappingProxyType({"x-amzn-requestid": marker})
def _json_reply(status: int, payload: Mapping[str, JsonValue], marker: str) -> Reply:
return Reply(status=status, body=json.dumps(payload).encode(), headers=_headers(marker))
def _error(status: int, message: str, marker: str) -> Reply:
return _json_reply(status, {"message": message}, marker)
def _effort_of(target: str, body: Mapping[str, JsonValue]) -> JsonValue:
if not _CONVERSE.match(target) and not _CONVERSE_STREAM.match(target):
return body.get("reasoning_effort")
fields: Final = body.get("additionalModelRequestFields")
reasoning: Final = fields.get("reasoning") if isinstance(fields, Mapping) else None
return reasoning.get("effort") if isinstance(reasoning, Mapping) else None
def forwarded_effort(request: Request) -> JsonValue:
return _effort_of(target_of(request), body_of(request))
def _sse(frames: tuple[Mapping[str, JsonValue], ...], pause: float) -> Reply:
return Reply(
content_type="text/event-stream",
chunks=(*(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames), b"data: [DONE]\n\n"),
pause_between_chunks=pause,
)
def _with_headers(reply: Reply, marker: str) -> Reply:
return Reply(
status=reply.status,
body=reply.body,
content_type=reply.content_type,
chunks=reply.chunks,
abort_after=reply.abort_after,
gate_after_first=reply.gate_after_first,
pause_between_chunks=reply.pause_between_chunks,
headers=_headers(marker),
)
def _content_deltas(model: str, marker: str) -> tuple[str, ...]:
if "gpt-oss" in model:
return ("<reason", "ing>why ", f"marker-{marker}", "</reas", "oning> answer ", f"marker-{marker}")
return ("answer ", f"marker-{marker}")
def _chat_text(model: str, marker: str) -> str:
return reasoning_answer(marker) if "gpt-oss" in model else answer(marker)
def _chat_reply(model: str, marker: str, stream: bool, pause: float) -> Reply:
identity: Final = f"chatcmpl-{marker}"
if not stream:
return _json_reply(
200,
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": _chat_text(model, marker)},
"finish_reason": "stop",
}
],
"usage": dict(USAGE),
},
marker,
)
deltas: Final = _content_deltas(model, marker)
frames: Final = tuple(
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [{"index": 0, "delta": {"role": "assistant", "content": delta}, "finish_reason": None}],
}
for delta in deltas
)
finish: Final[Mapping[str, JsonValue]] = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
"usage": dict(USAGE),
}
return _with_headers(_sse((*frames, finish), pause), marker)
def _responses_reply(model: str, marker: str, stream: bool, pause: float) -> Reply:
identity: Final = f"resp_upstream_{marker}"
item_id: Final = f"msg_{marker}"
response: Final[Mapping[str, JsonValue]] = {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": model,
"output": [
{
"type": "message",
"id": item_id,
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": answer(marker), "annotations": []}],
}
],
"usage": {"input_tokens": 30, "output_tokens": 5, "total_tokens": 35},
}
if not stream:
return _json_reply(200, response, marker)
events: Final[tuple[Mapping[str, JsonValue], ...]] = (
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
},
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": item_id,
"output_index": 0,
"content_index": 0,
"delta": answer(marker),
},
{"type": "response.completed", "sequence_number": 2, "response": response},
)
return Reply(
content_type="text/event-stream",
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
pause_between_chunks=pause,
headers=_headers(marker),
)
def _converse_reply(marker: str) -> Reply:
return _json_reply(
200,
{
"output": {"message": {"role": "assistant", "content": [{"text": answer(marker)}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14},
"metrics": {"latencyMs": 1},
},
marker,
)
def _converse_stream_reply(marker: str, pause: float) -> Reply:
events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = (
("messageStart", {"role": "assistant"}),
("contentBlockDelta", {"delta": {"text": "answer "}, "contentBlockIndex": 0}),
("contentBlockDelta", {"delta": {"text": f"marker-{marker}"}, "contentBlockIndex": 0}),
("contentBlockStop", {"contentBlockIndex": 0}),
("messageStop", {"stopReason": "end_turn"}),
("metadata", {"usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14}, "metrics": {"latencyMs": 1}}),
)
return Reply(
content_type=EVENT_STREAM,
chunks=tuple(_aws_event_frame(kind, payload, "sc", marker) for kind, payload in events),
pause_between_chunks=pause,
headers=_headers(marker),
)
def respond(request: Request, *, pause: float = 0.0) -> Reply:
target: Final = target_of(request)
marker: Final = marker_of(request)
if request.method == "GET":
if target == "/image.png":
return Reply(body=PNG_1X1, content_type="image/png", headers=_headers(marker))
return _error(404, f"no scripted object at {target}", marker)
scripted_status: Final = _STATUS.search(request.body.decode(errors="replace"))
if scripted_status is not None:
status: Final = int(scripted_status.group(1))
return _error(status, f"scripted {status}", marker)
body: Final = body_of(request)
effort: Final = _effort_of(target, body)
if effort is not None and (not isinstance(effort, str) or effort not in REASONING_EFFORTS):
return _error(400, f"Invalid reasoning effort: {json.dumps(effort)}", marker)
model: Final = str(body.get("model", ""))
stream: Final = body.get("stream") is True
if request.method == "POST" and target == NATIVE_CHAT:
return _chat_reply(model, marker, stream, pause)
if request.method == "POST" and target == NATIVE_RESPONSES:
return _responses_reply(model, marker, stream, pause)
if request.method == "POST" and _CONVERSE.match(target):
return _converse_reply(marker)
if request.method == "POST" and _CONVERSE_STREAM.match(target):
return _converse_stream_reply(marker, pause)
return _error(404, f"unknown bedrock route {request.method} {target}", marker)
def serve_peer(port: int, received: Synchronized[int], answer_first: int) -> None:
held: Final = threading.Event()
def respond_or_hold(request: Request) -> Reply:
with received.get_lock():
received.value += 1
ordinal: Final = received.value
if ordinal > answer_first:
held.wait()
return respond(request)
with wire_server(respond_or_hold, port=port):
threading.Event().wait()

View file

@ -0,0 +1,194 @@
import json
import uuid
from collections.abc import Mapping
from typing import Final
import anthropic
from integration._support.bedrock_runtime_peer import NATIVE_CHAT, answer, body_of, marker_of, respond, target_of
from integration._support.client import Gateway, Scenario, eventually
from integration._support.database import read_rows
from integration._support.wire import Request, Wire, wire_server
from pydantic import JsonValue
BEDROCK_MODEL: Final = "us.openai.gpt-5.6-sol"
TOKEN: Final = "synthetic-bedrock-bearer"
NO_CACHE: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}}
ANTHROPIC_VERSION: Final[Mapping[str, str]] = {"anthropic-version": "2023-06-01"}
def _question(marker: str) -> str:
return f"Question marker-{marker}"
def _deployment(scenario: Scenario, wire: Wire) -> str:
return scenario.model(
model=f"bedrock/{BEDROCK_MODEL}",
api_key=TOKEN,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
)
def _carrying(wire: Wire, marker: str) -> tuple[Request, ...]:
return tuple(request for request in wire.drain() if marker_of(request) == marker)
def _native_body(wire: Wire, marker: str) -> Mapping[str, JsonValue]:
received: Final = _carrying(wire, marker)
assert [(request.method, target_of(request)) for request in received] == [("POST", NATIVE_CHAT)]
assert received[0].headers["authorization"] == f"Bearer {TOKEN}", received[0].headers
return body_of(received[0])
def _native_request(marker: str, max_tokens: int, effort: str) -> Mapping[str, JsonValue]:
return {
"model": BEDROCK_MODEL,
"messages": [{"role": "user", "content": _question(marker)}],
"max_completion_tokens": max_tokens,
"reasoning_effort": effort,
}
def _spend_rows(identity: str, expected: int) -> list[dict[str, JsonValue]]:
return eventually(
lambda: read_rows(
"SELECT request_id, call_type, status, model_group, prompt_tokens, completion_tokens, cache_hit"
' FROM "LiteLLM_SpendLogs" WHERE starts_with(request_id, %s) ORDER BY "startTime"',
(identity,),
),
lambda found: len(found) == expected,
seconds=70,
)
def _success_row(identity: str, model: str, cache_hit: str = "None") -> dict[str, JsonValue]:
return {
"request_id": identity,
"call_type": "anthropic_messages",
"status": "success",
"model_group": model,
"prompt_tokens": 9,
"completion_tokens": 5,
"cache_hit": cache_hit,
}
def test_anthropic_sdk_thinking_budget_reaches_native_chat_completions_as_reasoning_effort(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
message: Final = client.messages.create(
model=model,
max_tokens=4096,
thinking={"type": "enabled", "budget_tokens": 2048},
messages=[{"role": "user", "content": _question(marker)}],
extra_body=NO_CACHE,
)
assert _native_body(wire, marker) == _native_request(marker, 4096, "medium")
assert message.id == f"chatcmpl-{marker}", message
assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", answer(marker))]
assert (message.usage.input_tokens, message.usage.output_tokens) == (9, 5), message
assert _spend_rows(message.id, 1) == [_success_row(message.id, model)]
def test_anthropic_sdk_stream_with_thinking_budget_is_served_by_native_chat_completions(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
stream: Final = client.messages.create(
model=model,
max_tokens=4096,
thinking={"type": "enabled", "budget_tokens": 2048},
messages=[{"role": "user", "content": _question(marker)}],
extra_body=NO_CACHE,
stream=True,
)
events: Final = list(stream)
assert _native_body(wire, marker) == {
**_native_request(marker, 4096, "medium"),
"stream": True,
"stream_options": {"include_usage": True},
}
assert events[0].type == "message_start" and events[-1].type == "message_stop", events
identity: Final = events[0].message.id
assert identity.startswith("msg_"), events
assert "".join(
event.delta.text
for event in events
if event.type == "content_block_delta" and event.delta.type == "text_delta"
) == answer(marker)
assert _spend_rows(identity, 1) == [_success_row(identity, model, cache_hit="False")]
def test_raw_thinking_summary_reaches_native_chat_completions_as_the_plain_effort(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048, "summary": "detailed"},
"messages": [{"role": "user", "content": _question(marker)}],
**NO_CACHE,
},
headers=ANTHROPIC_VERSION,
)
body: Final = _native_body(wire, marker)
assert body == _native_request(marker, 4096, "medium")
assert "summary" not in json.dumps(body), body
assert response.status_code == 200, response.text
assert response.json()["id"] == f"chatcmpl-{marker}", response.text
assert response.json()["content"] == [{"type": "text", "text": answer(marker)}], response.text
assert _spend_rows(f"chatcmpl-{marker}", 1) == [_success_row(f"chatcmpl-{marker}", model)]
async def test_async_anthropic_sdk_disabled_thinking_reaches_native_chat_completions_as_effort_none(
gateway: Gateway,
) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
client: Final = anthropic.AsyncAnthropic(
base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0
)
message: Final = await client.messages.create(
model=model,
max_tokens=64,
thinking={"type": "disabled"},
messages=[{"role": "user", "content": _question(marker)}],
extra_body=NO_CACHE,
)
assert _native_body(wire, marker) == _native_request(marker, 64, "none")
assert message.id == f"chatcmpl-{marker}", message
assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", answer(marker))]
assert _spend_rows(message.id, 1) == [_success_row(message.id, model)]
def test_identical_messages_requests_reach_the_peer_once_and_log_a_cache_hit_row(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
body: Final[dict[str, JsonValue]] = {
"model": model,
"max_tokens": 64,
"messages": [{"role": "user", "content": _question(marker)}],
}
first: Final = gateway.request("POST", "/v1/messages", body, headers=ANTHROPIC_VERSION)
assert first.status_code == 200, first.text
identity: Final = str(first.json()["id"])
assert first.json()["content"] == [{"type": "text", "text": answer(marker)}], first.text
second: Final = gateway.request("POST", "/v1/messages", body, headers=ANTHROPIC_VERSION)
assert second.status_code == 200, second.text
assert second.json()["id"] == identity, (first.text, second.text)
assert second.json()["content"] == [{"type": "text", "text": answer(marker)}], second.text
received: Final = _carrying(wire, marker)
assert [(request.method, marker_of(request)) for request in received] == [("POST", marker)], received
rows: Final = _spend_rows(identity, 2)
assert rows[0] == _success_row(identity, model), rows
assert str(rows[1]["request_id"]).startswith(identity + "_cache_hit"), rows
assert {**rows[1], "request_id": identity, "cache_hit": "None"} == _success_row(identity, model), rows

View file

@ -0,0 +1,117 @@
import base64
import uuid
from dataclasses import dataclass
from typing import Final
import openai
from integration._support.bedrock_runtime_peer import NATIVE_RESPONSES, answer, respond, target_of
from integration._support.client import Gateway, Scenario, eventually
from integration._support.database import read_rows
from integration._support.wire import Request, Wire, wire_server
from openai.types.responses import ResponseCompletedEvent, ResponseTextDeltaEvent
from pydantic import JsonValue, TypeAdapter
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with
GPT: Final = "us.openai.gpt-5.6-sol"
TOKEN: Final = "synthetic-bedrock-bearer"
SALT: Final = "sk-integration-salt"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
@dataclass(frozen=True, slots=True)
class _IssuedId:
issued: str
upstream: str
def _prompt(marker: str) -> str:
return f"synthetic responses request marker-{marker}"
def _deployment(scenario: Scenario, wire: Wire) -> str:
return scenario.model(
model=f"bedrock/{GPT}",
api_key=TOKEN,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
api_base=None,
)
def _issued_id(client_id: str) -> _IssuedId:
decrypted: Final = decrypt_if_encrypted_with(client_id.removeprefix("resp_"), SALT)
assert decrypted is not None, client_id
issued: Final = decrypted.split(";")[0].split("response_id:")[-1]
decoded: Final = base64.b64decode(issued.removeprefix("resp_")).decode()
return _IssuedId(issued, decoded.split(";")[-1].removeprefix("response_id:"))
def _native_request(wire: Wire) -> Request:
received: Final = wire.drain()
assert [(request.method, target_of(request)) for request in received] == [("POST", NATIVE_RESPONSES)], received
assert received[0].headers["authorization"] == f"Bearer {TOKEN}", dict(received[0].headers)
return received[0]
def _body(request: Request) -> dict[str, JsonValue]:
return _JSON_OBJECT.validate_json(request.body)
def _spend_row(identity: str) -> dict[str, JsonValue]:
rows: Final = eventually(
lambda: read_rows(
'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(identity,),
),
lambda found: len(found) == 1,
seconds=70,
)
return rows[0]
def _success_row(model: str) -> dict[str, JsonValue]:
return {"model_group": model, "status": "success", "prompt_tokens": 30, "completion_tokens": 5}
def test_openai_sdk_responses_request_is_served_by_the_native_responses_route(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
client: Final = openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0)
raw: Final = client.responses.with_raw_response.create(
model=model, input=_prompt(marker), extra_body={"cache": {"no-cache": True}}
)
response: Final = raw.parse()
assert response.output_text == answer(marker), raw.text
assert response.usage is not None and (response.usage.input_tokens, response.usage.output_tokens) == (30, 5)
assert _issued_id(response.id).upstream == f"resp_upstream_{marker}", response.id
request: Final = _native_request(wire)
assert _body(request) == {"model": GPT, "input": _prompt(marker)}, request.body
assert _spend_row(response.id) == _success_row(model)
async def test_async_openai_sdk_responses_stream_is_served_by_the_native_responses_route(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
client: Final = openai.AsyncOpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0)
stream: Final = await client.responses.create(
model=model, input=_prompt(marker), stream=True, extra_body={"cache": {"no-cache": True}}
)
events: Final = [event async for event in stream]
assert [event.type for event in events] == [
"response.created",
"response.output_text.delta",
"response.completed",
], events
deltas: Final = "".join(event.delta for event in events if isinstance(event, ResponseTextDeltaEvent))
assert deltas == answer(marker), events
completed: Final = events[-1]
assert isinstance(completed, ResponseCompletedEvent), completed
assert completed.response.output_text == answer(marker), completed
issued: Final = _issued_id(completed.response.id)
assert issued.upstream == f"resp_upstream_{marker}", completed.response.id
request: Final = _native_request(wire)
assert _body(request) == {"model": GPT, "input": _prompt(marker), "stream": True}, request.body
assert _spend_row(issued.issued) == _success_row(model)

View file

@ -0,0 +1,402 @@
import asyncio
import base64
import binascii
import multiprocessing
import os
import re
import signal
import socket
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from multiprocessing.process import BaseProcess
from multiprocessing.sharedctypes import Synchronized
from pathlib import Path
from queue import SimpleQueue
from types import MappingProxyType
from typing import Final, Literal
from urllib.parse import urlsplit, urlunsplit
import httpx
import psutil
import pytest
import yaml
from integration._support.bedrock_runtime_peer import MARKER, marker_of, respond, serve_peer
from integration._support.client import Gateway, Scenario, eventually, object_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
BEDROCK_MODEL: Final = "us.openai.gpt-5.6-sol"
TOKEN: Final = "synthetic-bedrock-bearer"
_CONFIG_MODEL: Final = "bedrock-gpt-chat-completions-chaos"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_ENDPOINTS: Final[tuple["Endpoint", ...]] = ("chat", "messages", "responses")
Endpoint = Literal["chat", "messages", "responses"]
@dataclass(frozen=True, slots=True)
class _Call:
endpoint: Endpoint
stream: bool
marker: str
@dataclass(frozen=True, slots=True)
class _Served:
call: _Call
status: int
text: str
call_id: str | None
@dataclass(frozen=True, slots=True)
class _ChildPeer:
process: BaseProcess
received: Synchronized[int]
url: str
def _path(endpoint: Endpoint) -> str:
match endpoint:
case "chat":
return "/v1/chat/completions"
case "messages":
return "/v1/messages"
case "responses":
return "/v1/responses"
def _terminal(endpoint: Endpoint) -> str:
match endpoint:
case "chat":
return "data: [DONE]"
case "messages":
return "event: message_stop"
case "responses":
return '"type":"response.completed"'
def _body(model: str, call: _Call) -> dict[str, JsonValue]:
question: Final = f"Question marker-{call.marker}"
common: Final[dict[str, JsonValue]] = {"model": model, "stream": call.stream, "cache": {"no-cache": True}}
match call.endpoint:
case "chat":
return {**common, "messages": [{"role": "user", "content": question}]}
case "messages":
return {**common, "max_tokens": 64, "messages": [{"role": "user", "content": question}]}
case "responses":
return {**common, "input": question}
def _deployment(scenario: Scenario, endpoint: str) -> str:
return scenario.model(
model=f"bedrock/{BEDROCK_MODEL}",
api_key=TOKEN,
api_base=None,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=endpoint,
)
def _frames(text: str) -> tuple[dict[str, JsonValue], ...]:
return tuple(
_JSON_OBJECT.validate_json(line[6:])
for line in text.splitlines()
if line.startswith("data: ") and line != "data: [DONE]"
)
def _frame_id(frame: Mapping[str, JsonValue]) -> str | None:
if frame.get("type") == "message_start":
return str(object_value(frame["message"])["id"])
response: Final = frame.get("response")
if isinstance(response, dict) and "id" in response:
return str(response["id"])
identity: Final = frame.get("id")
return identity if isinstance(identity, str) else None
def _response_id(served: _Served) -> str:
if not served.call.stream:
return str(_JSON_OBJECT.validate_json(served.text)["id"])
ids: Final = tuple(identity for identity in map(_frame_id, _frames(served.text)) if identity is not None)
assert ids, served.text
return ids[0]
def _assert_answered_with_its_own_marker(served: _Served) -> None:
assert served.status == 200, served.text
assert set(MARKER.findall(served.text)) == {served.call.marker}, served.text
if served.call.stream:
assert _terminal(served.call.endpoint) in served.text, served.text
def _spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]:
return eventually(
lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)),
lambda found: len(found) >= expected,
seconds=60,
)
def _rows_by_status(rows: list[dict[str, JsonValue]], status: str) -> list[str]:
return sorted(str(row["request_id"]) for row in rows if row["status"] == status)
def _upstream_id_inside(row_id: str) -> str | None:
try:
payload: Final = base64.b64decode(row_id.removeprefix("resp_"), validate=True).decode()
except (binascii.Error, UnicodeDecodeError):
return None
return payload.rsplit("response_id:", 1)[1] if "response_id:" in payload else None
# TODO: a Bedrock non-stream /v1/responses spend row can carry the pre-encryption resp_<base64> id instead of the
# ciphertext the caller received, because the spend row id is read from response_obj["id"] before the
# ResponsesIDSecurity hook rewrites it in place; such a row is matched by the upstream id inside that payload until
# that ordering is fixed on main
def _row_belongs_to(row_id: str, served: _Served) -> bool:
if row_id == _response_id(served):
return True
return served.call.endpoint == "responses" and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}"
def _assert_each_success_landed_once(rows: list[dict[str, JsonValue]], served: tuple[_Served, ...]) -> None:
success_ids: Final = _rows_by_status(rows, "success")
assert len(success_ids) == len(served), rows
for item in served:
owned: Final = [row_id for row_id in success_ids if _row_belongs_to(row_id, item)]
assert len(owned) == 1, (item.call, owned, success_ids)
async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served:
async with client.stream(
"POST",
_path(call.endpoint),
json=_body(model, call),
headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"},
) as response:
raw: Final = await response.aread()
return _Served(
call=call, status=response.status_code, text=raw.decode(), call_id=response.headers.get("x-litellm-call-id")
)
async def _burst(
base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False
) -> tuple[_Served, ...]:
async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
results: Final = await asyncio.gather(
*(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors
)
for result in results:
assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
return tuple(result for result in results if isinstance(result, _Served))
def _calls(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool]) -> tuple[_Call, ...]:
return tuple(
_Call(endpoint=endpoints[index % len(endpoints)], stream=stream(index), marker=uuid.uuid4().hex)
for index in range(count)
)
def _free_port() -> int:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
return reserve.getsockname()[1]
def _accepts_connections(port: int) -> bool:
try:
with socket.create_connection(("127.0.0.1", port), timeout=0.2):
return True
except OSError:
return False
@contextmanager
def _child_peer(port: int, answer_first: int) -> Iterator[_ChildPeer]:
received: Final = multiprocessing.Value("i", 0)
process: Final = multiprocessing.get_context("spawn").Process(
target=serve_peer, args=(port, received, answer_first), daemon=True
)
process.start()
try:
eventually(lambda: _accepts_connections(port), bool, seconds=30)
yield _ChildPeer(process=process, received=received, url=f"http://127.0.0.1:{port}")
finally:
process.kill()
process.join(timeout=10)
assert not process.is_alive(), "Owned peer survived cleanup"
async def test_burst_across_every_endpoint_lands_each_response_id_once(gateway: Gateway) -> None:
calls: Final = _calls(36, _ENDPOINTS, lambda index: index % 2 == 0)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire.url)
served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
assert len(served) == 36
for item in served:
_assert_answered_with_its_own_marker(item)
ids: Final = sorted(_response_id(item) for item in served)
assert len(set(ids)) == 36, ids
assert sorted(marker_of(request) for request in wire.drain()) == sorted(call.marker for call in calls)
rows: Final = _spend_rows(model, 36)
_assert_each_success_landed_once(rows, served)
assert len(rows) == 36, rows
@pytest.mark.timeout(180)
async def test_peer_killed_mid_burst_fails_only_the_held_calls_and_a_restarted_peer_serves_again(
gateway: Gateway,
) -> None:
calls: Final = _calls(12, _ENDPOINTS, lambda index: index % 2 == 0)
recovery: Final = _calls(6, _ENDPOINTS, lambda index: index % 2 == 1)
port: Final = _free_port()
with gateway.scenario() as scenario:
model: Final = _deployment(scenario, f"http://127.0.0.1:{port}")
with _child_peer(port, answer_first=6) as peer:
burst: Final = asyncio.create_task(_burst(str(gateway.client.base_url), gateway.key, model, calls))
await asyncio.to_thread(eventually, lambda: peer.received.value, lambda count: count == 12, 60)
peer.process.kill()
peer.process.join(timeout=10)
served: Final = await burst
succeeded: Final = tuple(item for item in served if item.status == 200)
failed: Final = tuple(item for item in served if item.status != 200)
assert (len(succeeded), len(failed)) == (6, 6), [(item.call.marker, item.status) for item in served]
for item in succeeded:
_assert_answered_with_its_own_marker(item)
assert {item.status for item in failed} == {503}, [
(item.call.endpoint, item.call.stream, item.status, item.text) for item in failed
]
for item in failed:
assert "ServiceUnavailableError: BedrockException - Server disconnected" in item.text, item.text
assert "marker-" not in item.text and item.call_id is not None, item.text
with _child_peer(port, answer_first=10**6) as revived:
recovered: Final = await _burst(str(gateway.client.base_url), gateway.key, model, recovery)
assert revived.received.value == 6, revived.received.value
for item in recovered:
_assert_answered_with_its_own_marker(item)
rows: Final = _spend_rows(model, 18)
_assert_each_success_landed_once(rows, (*succeeded, *recovered))
assert _rows_by_status(rows, "failure") == sorted(str(item.call_id) for item in failed), rows
assert len(rows) == 18, rows
async def test_slow_peer_streams_are_forwarded_once_and_terminated(gateway: Gateway) -> None:
calls: Final = _calls(10, ("chat",), lambda _: True)
with wire_server(lambda request: respond(request, pause=0.3)) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire.url)
served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
assert len(served) == 10
for item in served:
_assert_answered_with_its_own_marker(item)
assert sorted(marker_of(request) for request in wire.drain()) == sorted(call.marker for call in calls)
ids: Final = sorted(_response_id(item) for item in served)
rows: Final = _spend_rows(model, 10)
assert _rows_by_status(rows, "success") == ids, rows
assert len(rows) == 10, rows
def _chaos_config(wire: Wire, tmp_path: Path) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["model_list"] = [
{
"model_name": _CONFIG_MODEL,
"litellm_params": {
"model": f"bedrock/{BEDROCK_MODEL}",
"api_key": TOKEN,
"aws_region_name": "us-east-1",
"aws_bedrock_runtime_endpoint": wire.url,
},
}
]
path: Final = tmp_path / "bedrock-gpt-chat-completions-chaos.yaml"
path.write_text(yaml.safe_dump(config))
return path
def _pooled_database_url() -> str:
parts: Final = urlsplit(os.environ["DATABASE_URL"])
query: Final = "&".join(part for part in (parts.query, "connection_limit=5") if part)
return urlunsplit(parts._replace(query=query))
def _open_upstream_connections(pid: int, upstream: str) -> int:
port: Final = urlsplit(upstream).port
return sum(
1
for connection in psutil.Process(pid).net_connections(kind="tcp")
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
)
def _landed_once(ids: tuple[str, ...]) -> list[dict[str, JsonValue]]:
return eventually(
lambda: read_rows(
'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)',
(list(ids),), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array
),
lambda found: len(found) >= len(ids),
seconds=60,
)
@pytest.mark.timeout(180)
async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving(gateway: Gateway, tmp_path: Path) -> None:
calls: Final = _calls(20, ("chat",), lambda _: False)
release: Final = threading.Event()
held_markers: Final[SimpleQueue[str]] = SimpleQueue()
def held(request: Request) -> Reply:
held_markers.put(marker_of(request))
assert release.wait(timeout=60), "The burst was never released"
return respond(request)
with wire_server(held) as wire:
path: Final = _chaos_config(wire, tmp_path)
overrides: Final = {"DATABASE_URL": _pooled_database_url()}
with owned_proxy_process(gateway, tmp_path, overrides, config=path, workers=2) as owned:
candidate: Final = owned.gateway
workers: Final = eventually(
lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
lambda pids: len(pids) == 2,
seconds=30,
)
burst: Final = asyncio.create_task(
_burst(
str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True
)
)
await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60)
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers})
assert sum(held_by.values()) == 20, held_by
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
victim: Final = psutil.Process(victim_pid)
victim.suspend()
victim.send_signal(signal.SIGKILL)
release.set()
served: Final = await burst
assert held_by[survivor_pid] >= 10, held_by
assert len(served) == held_by[survivor_pid], (held_by, len(served))
for item in served:
_assert_answered_with_its_own_marker(item)
follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex)
(answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,))
_assert_answered_with_its_own_marker(answered)
received: Final = wire.drain()
assert {request.method for request in received} == {"POST"}, received
assert sorted(marker_of(request) for request in received) == sorted(
call.marker for call in (*calls, follow_up)
)
ids: Final = tuple(sorted(_response_id(item) for item in (*served, answered)))
rows: Final = _landed_once(ids)
assert _rows_by_status(rows, "success") == list(ids), rows
assert len(rows) == len(ids), rows

View file

@ -0,0 +1,415 @@
import json
import os
import time
import uuid
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from hashlib import sha256
from pathlib import Path
from types import MappingProxyType
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import httpx
import pytest
import yaml
from integration._support.bedrock_runtime_peer import answer, forwarded_effort, marker_of, respond, target_of
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy_process
from integration._support.wire import Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
GPT: Final = "us.openai.gpt-5.6-sol"
TOKEN: Final = "synthetic-bedrock-bearer"
BAD_KEY: Final = "sk-synthetic-bad-key"
NATIVE_TARGET: Final = "/openai/v1/chat/completions"
CONVERSE_TARGET: Final = f"/model/{GPT}/converse"
LONG_VERSION_GPT: Final = "openai.gpt-" + "1" * 30000
PNG_DATA_URL: Final = (
"data:image/png;base64,"
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg=="
)
GPT_DEPLOYMENT: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"model": f"bedrock/{GPT}", "api_key": TOKEN, "aws_region_name": "us-east-1"}
)
_ALLOWLISTED_MODEL: Final = "bedrock-gpt-image-allowlist"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _prompt(marker: str) -> str:
return f"synthetic sad request marker-{marker}"
def _messages(marker: str) -> list[dict[str, JsonValue]]:
return [{"role": "user", "content": _prompt(marker)}]
def _image_messages(marker: str, url: str) -> list[dict[str, JsonValue]]:
return [
{
"role": "user",
"content": [{"type": "text", "text": _prompt(marker)}, {"type": "image_url", "image_url": {"url": url}}],
}
]
def _deployment(scenario: Scenario, wire: Wire, **overrides: JsonValue) -> str:
return scenario.model(**{**GPT_DEPLOYMENT, "aws_bedrock_runtime_endpoint": wire.url, **overrides})
def _chat(gateway: Gateway, model: str, marker: str, *, key: str | None = None, **params: JsonValue) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": _messages(marker), "cache": {"no-cache": True}, **params},
key=key,
)
def _payload(response: httpx.Response) -> dict[str, JsonValue]:
assert response.status_code == 200, response.text
return _JSON_OBJECT.validate_json(response.content)
def _content(response: httpx.Response) -> JsonValue:
choices: Final = _payload(response)["choices"]
assert isinstance(choices, list), response.text
return object_value(object_value(choices[0])["message"])["content"]
def _error_message(response: httpx.Response) -> str:
return string_value(object_value(_JSON_OBJECT.validate_json(response.content)["error"])["message"])
def _call_id(response: httpx.Response) -> str:
return response.headers["x-litellm-call-id"]
def _body(request: Request) -> dict[str, JsonValue]:
return _JSON_OBJECT.validate_json(request.body)
def _routes(received: tuple[Request, ...]) -> list[tuple[str, str]]:
return [(request.method, target_of(request)) for request in received]
def _only_request(wire: Wire, marker: str) -> Request:
received: Final = wire.drain()
assert len(received) == 1, _routes(received)
assert marker_of(received[0]) == marker, received[0].body
return received[0]
def _spend_rows(identity: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT request_id, model_group, status, cache_hit, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(identity,),
)
def _spend_row(identity: str) -> dict[str, JsonValue]:
return eventually(lambda: _spend_rows(identity), lambda found: len(found) == 1, seconds=70)[0]
def _assert_row(identity: str, model: str, status: str) -> None:
row: Final = _spend_row(identity)
assert (row["model_group"], row["status"]) == (model, status), row
def _timed_liveliness(gateway: Gateway) -> tuple[int, float]:
started: Final = time.monotonic()
response: Final = gateway.request("GET", "/health/liveliness")
return response.status_code, time.monotonic() - started
def _pooled_database_url(url: str) -> str:
parts: Final = urlsplit(url)
query: Final = "&".join(part for part in (parts.query, "connection_limit=5") if part)
return urlunsplit(parts._replace(query=query))
def _allowlist_config(wire: Wire, tmp_path: Path) -> Path:
config: Final = _JSON_OBJECT.validate_python(
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
)
path: Final = tmp_path / "bedrock-gpt-image-allowlist.yaml"
path.write_text(
yaml.safe_dump(
{
**config,
"model_list": [
{
"model_name": _ALLOWLISTED_MODEL,
"litellm_params": {**GPT_DEPLOYMENT, "aws_bedrock_runtime_endpoint": wire.url},
}
],
"general_settings": {
**object_value(config["general_settings"]),
"user_url_allowed_hosts": ["127.0.0.1"],
},
}
)
)
return path
def test_remote_image_url_on_the_shared_proxy_is_rejected_before_any_fetch(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": _image_messages(marker, f"{wire.url}/image.png"), "cache": {"no-cache": True}},
)
assert response.status_code == 400, response.text
message: Final = _error_message(response)
assert "Unable to fetch image from URL" in message and "user_url_allowed_hosts" in message, response.text
_assert_row(_call_id(response), model, "failure")
assert _routes(wire.drain()) == []
@pytest.mark.timeout(180)
def test_allowlisted_remote_image_is_inlined_for_the_native_route(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = uuid.uuid4().hex
missing_marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire:
path: Final = _allowlist_config(wire, tmp_path)
overrides: Final = {"DATABASE_URL": _pooled_database_url(os.environ["DATABASE_URL"])}
with owned_proxy_process(gateway, tmp_path, overrides, config=path) as owned:
candidate: Final = owned.gateway
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{
"model": _ALLOWLISTED_MODEL,
"messages": _image_messages(marker, f"{wire.url}/image.png"),
"cache": {"no-cache": True},
},
)
assert _content(response) == answer(marker), response.text
received: Final = wire.drain()
assert _routes(received) == [("GET", "/image.png"), ("POST", NATIVE_TARGET)], received
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(received[1]) == {
"model": GPT,
"messages": _image_messages(marker, PNG_DATA_URL),
"stream": False,
}, received[1].body
_assert_row(f"chatcmpl-{marker}", _ALLOWLISTED_MODEL, "success")
missing: Final = candidate.request(
"POST",
"/v1/chat/completions",
{
"model": _ALLOWLISTED_MODEL,
"messages": _image_messages(missing_marker, f"{wire.url}/missing.png"),
"cache": {"no-cache": True},
},
)
assert missing.status_code == 400, missing.text
assert "Unable to fetch image from URL. Status code: 404" in _error_message(missing), missing.text
_assert_row(_call_id(missing), _ALLOWLISTED_MODEL, "failure")
assert _routes(wire.drain()) == [("GET", "/missing.png")]
def test_response_cache_twin_serves_the_second_request_without_a_second_wire_call(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
body: Final[dict[str, JsonValue]] = {"model": model, "messages": _messages(marker)}
first: Final = gateway.request("POST", "/v1/chat/completions", body)
second: Final = gateway.request("POST", "/v1/chat/completions", body)
identity: Final = string_value(_payload(first)["id"])
assert _content(first) == answer(marker), first.text
assert _payload(second)["id"] == identity, (first.text, second.text)
assert _content(second) == answer(marker), second.text
_only_request(wire, marker)
rows: Final = eventually(
lambda: read_rows(
'SELECT request_id, cache_hit, spend FROM "LiteLLM_SpendLogs" WHERE starts_with(request_id, %s)'
" ORDER BY request_id",
(identity,),
),
lambda found: len(found) == 2,
seconds=70,
)
assert [(row["request_id"] == identity, row["cache_hit"]) for row in rows] == [(True, "None"), (False, "True")]
assert string_value(rows[1]["request_id"]).startswith(f"{identity}_cache_hit"), rows
assert rows[1]["spend"] == 0.0, rows
assert isinstance(rows[0]["spend"], float) and rows[0]["spend"] > 0.0, rows
def test_model_group_info_lists_the_native_supported_params(gateway: Gateway) -> None:
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
groups: Final = gateway.get("/model_group/info", {"model_group": model})["data"]
assert isinstance(groups, list) and len(groups) == 1, groups
group: Final = object_value(groups[0])
assert group["model_group"] == model, group
params: Final = group["supported_openai_params"]
assert isinstance(params, list), group
assert {"reasoning_effort", "logprobs", "top_logprobs"} <= set(params) and "n" not in params, params
assert _routes(wire.drain()) == []
def test_thirty_thousand_digit_version_is_classified_quickly_and_served_by_converse(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=1) as pool:
model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_VERSION_GPT}")
liveliness: Final = pool.submit(_timed_liveliness, gateway)
started: Final = time.monotonic()
response: Final = _chat(gateway, model, marker)
elapsed: Final = time.monotonic() - started
health_status, health_elapsed = liveliness.result()
assert _content(response) == answer(marker), response.text
assert elapsed < 10, elapsed
assert (health_status, health_elapsed < 2) == (200, True), (health_status, health_elapsed)
request: Final = _only_request(wire, marker)
assert (request.method, target_of(request)) == ("POST", f"/model/{LONG_VERSION_GPT}/converse"), request.target
_assert_row(string_value(_payload(response)["id"]), model, "success")
def test_bad_key_on_the_long_version_model_is_refused_before_any_route(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
control_marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_VERSION_GPT}")
started: Final = time.monotonic()
refused: Final = _chat(gateway, model, marker, key=BAD_KEY)
elapsed: Final = time.monotonic() - started
assert refused.status_code == 401, refused.text
assert elapsed < 2, elapsed
assert "Authentication Error" in _error_message(refused), refused.text
refused_rows: Final = eventually(
lambda: read_rows(
"SELECT request_id, status, spend, metadata->'error_information'->>'error_code' AS error_code"
' FROM "LiteLLM_SpendLogs" WHERE model_group=%s AND api_key=%s',
(model, sha256(BAD_KEY.encode()).hexdigest()),
),
lambda found: len(found) == 1,
seconds=70,
)
assert (refused_rows[0]["status"], refused_rows[0]["spend"], refused_rows[0]["error_code"]) == (
"failure",
0.0,
"401",
), refused_rows
control: Final = _chat(gateway, model, control_marker)
control_id: Final = string_value(_payload(control)["id"])
_assert_row(control_id, model, "success")
landed: Final = read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,))
assert {row["request_id"] for row in landed} == {control_id, refused_rows[0]["request_id"]}, landed
received: Final = wire.drain()
assert [marker_of(request) for request in received] == [control_marker], _routes(received)
@pytest.mark.parametrize(
"effort",
[
pytest.param(7, id="int"),
pytest.param(["high"], id="list"),
pytest.param("", id="empty"),
pytest.param("x" * 5120, id="five_kb"),
],
)
def test_invalid_reasoning_effort_reaches_the_peer_and_its_400_reaches_the_caller(
gateway: Gateway, effort: JsonValue
) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, reasoning_effort=effort)
assert response.status_code == 400, response.text
peer_error: Final = json.dumps({"message": f"Invalid reasoning effort: {json.dumps(effort)}"})
assert f"BedrockException - {peer_error}" in _error_message(response), response.text
request: Final = _only_request(wire, marker)
assert forwarded_effort(request) == effort, request.body
_assert_row(_call_id(response), model, "failure")
def test_duplicated_reasoning_effort_key_lets_the_last_value_win(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
prefix: Final = json.dumps({"model": model, "messages": _messages(marker), "cache": {"no-cache": True}})[:-1]
response: Final = gateway.client.post(
"/v1/chat/completions",
content=f'{prefix}, "reasoning_effort": "low", "reasoning_effort": "high"}}'.encode(),
headers={"Authorization": f"Bearer {gateway.key}", "content-type": "application/json"},
)
assert _content(response) == answer(marker), response.text
request: Final = _only_request(wire, marker)
assert forwarded_effort(request) == "high", request.body
_assert_row(string_value(_payload(response)["id"]), model, "success")
def test_string_temperature_is_refused_before_any_wire_request(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, temperature="0.2")
assert response.status_code == 400, response.text
message: Final = _error_message(response)
assert message.startswith("litellm.UnsupportedParamsError") and "['temperature']" in message, response.text
_assert_row(_call_id(response), model, "failure")
assert _routes(wire.drain()) == []
@pytest.mark.parametrize(
("scripted", "expected"),
[pytest.param(401, 401, id="401"), pytest.param(429, 429, id="429"), pytest.param(500, 503, id="500")],
)
def test_peer_error_status_reaches_the_caller_and_unrelated_deployments_keep_serving(
gateway: Gateway, scripted: int, expected: int
) -> None:
marker: Final = uuid.uuid4().hex
control_marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
unrelated: Final = scenario.model()
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"status={scripted} marker-{marker}"}],
"cache": {"no-cache": True},
},
)
assert response.status_code == expected, response.text
assert f'BedrockException - {{"message": "scripted {scripted}"}}' in _error_message(response), response.text
_only_request(wire, marker)
_assert_row(_call_id(response), model, "failure")
control: Final = _chat(gateway, unrelated, control_marker)
assert control.status_code == 200, control.text
_assert_row(string_value(_payload(control)["id"]), unrelated, "success")
assert _routes(wire.drain()) == []
@pytest.mark.parametrize(
"params", [pytest.param({"reasoning_effort": None}, id="null"), pytest.param({}, id="missing")]
)
def test_absent_reasoning_effort_is_forwarded_as_absent_on_every_repeat(
gateway: Gateway, params: dict[str, JsonValue]
) -> None:
markers: Final = tuple(uuid.uuid4().hex for _ in range(3))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
responses: Final = tuple(_chat(gateway, model, marker, **params) for marker in markers)
assert [_content(response) for response in responses] == [answer(marker) for marker in markers]
ids: Final = tuple(string_value(_payload(response)["id"]) for response in responses)
assert len(set(ids)) == 3, ids
received: Final = wire.drain()
assert [marker_of(request) for request in received] == list(markers), _routes(received)
assert [forwarded_effort(request) for request in received] == [None, None, None], [_body(r) for r in received]
rows: Final = eventually(
lambda: read_rows(
'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id IN (%s, %s, %s)', ids
),
lambda found: len(found) == 3,
seconds=70,
)
assert {(string_value(row["request_id"]), row["status"]) for row in rows} == {
(identity, "success") for identity in ids
}, rows

View file

@ -0,0 +1,549 @@
import json
import uuid
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from urllib.parse import quote
import httpx
import openai
import pytest
from integration._support.bedrock_runtime_peer import answer, respond, target_of
from integration._support.client import Gateway, Scenario, eventually
from integration._support.database import read_rows
from integration._support.sigv4 import signature
from integration._support.wire import Request, Wire, wire_server
from openai.types.chat import ChatCompletionChunk, ChatCompletionMessageParam
from openai.types.chat.chat_completion_chunk import ChoiceDelta
from pydantic import JsonValue, TypeAdapter
GPT: Final = "us.openai.gpt-5.6-sol"
GLOBAL_GPT: Final = "global.openai.gpt-5.6-sol"
GPT_OSS: Final = "openai.gpt-oss-120b-1:0"
TOKEN: Final = "synthetic-bedrock-bearer"
ACCESS_KEY: Final = "AKIASYNTHETICKEY0001"
SECRET_KEY: Final = "synthetic-secret-key-for-testing"
PROFILE_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/a1b2c3d4e5f6"
NATIVE_TARGET: Final = "/openai/v1/chat/completions"
CONVERSE_TARGET: Final = f"/model/{GPT}/converse"
GPT_DEPLOYMENT: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"model": f"bedrock/{GPT}", "api_key": TOKEN, "aws_region_name": "us-east-1"}
)
GUARDRAIL: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"guardrailIdentifier": "gr-synthetic", "guardrailVersion": "1"}
)
TOOL_PARAMETERS: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}
)
TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"type": "function",
"function": {
"name": "lookup_invoice",
"description": "Look up an invoice",
"parameters": dict(TOOL_PARAMETERS),
},
}
)
CONVERSE_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"toolSpec": {
"inputSchema": {"json": dict(TOOL_PARAMETERS)},
"name": "lookup_invoice",
"description": "Look up an invoice",
}
}
)
JSON_SCHEMA: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"type": "json_schema",
"json_schema": {
"name": "verdict",
"strict": True,
"schema": {
"type": "object",
"properties": {"ok": {"type": "boolean"}},
"required": ["ok"],
"additionalProperties": False,
},
},
}
)
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_OBSERVATIONS: Final = TypeAdapter(list[dict[str, JsonValue]])
def _prompt(marker: str) -> str:
return f"synthetic native request marker-{marker}"
def _messages(marker: str) -> list[JsonValue]:
return [{"role": "user", "content": _prompt(marker)}]
def _sdk_messages(marker: str) -> list[ChatCompletionMessageParam]:
return [{"role": "user", "content": _prompt(marker)}]
def _converse_messages(marker: str) -> list[JsonValue]:
return [{"role": "user", "content": [{"text": _prompt(marker)}]}]
def _native_body(model: str, marker: str, **params: JsonValue) -> dict[str, JsonValue]:
return {"model": model, "messages": _messages(marker), "stream": False, **params}
def _streamed_native_body(model: str, marker: str) -> dict[str, JsonValue]:
return _native_body(model, marker, stream=True, stream_options={"include_usage": True})
def _deployment(scenario: Scenario, wire: Wire, **overrides: JsonValue) -> str:
return scenario.model(model_info=None, **{**GPT_DEPLOYMENT, "aws_bedrock_runtime_endpoint": wire.url, **overrides})
def _openai_client(gateway: Gateway) -> openai.OpenAI:
return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
def _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI:
return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
def _chat(gateway: Gateway, model: str, marker: str, **params: JsonValue) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": _messages(marker), "cache": {"no-cache": True}, **params},
)
def _payload(response: httpx.Response) -> dict[str, JsonValue]:
assert response.status_code == 200, response.text
return _JSON_OBJECT.validate_json(response.content)
def _only_request(wire: Wire) -> Request:
received: Final = wire.drain()
assert len(received) == 1, [(request.method, target_of(request)) for request in received]
return received[0]
def _body(request: Request) -> dict[str, JsonValue]:
return _JSON_OBJECT.validate_json(request.body)
def _native_request(wire: Wire) -> Request:
request: Final = _only_request(wire)
assert (request.method, target_of(request)) == ("POST", NATIVE_TARGET), request.target
assert request.headers["authorization"] == f"Bearer {TOKEN}", dict(request.headers)
return request
def _converse_request(wire: Wire, target: str = CONVERSE_TARGET) -> Request:
request: Final = _only_request(wire)
assert (request.method, target_of(request)) == ("POST", target), request.target
assert request.headers["authorization"] == f"Bearer {TOKEN}", dict(request.headers)
return request
def _spend_row(identity: str) -> dict[str, JsonValue]:
rows: Final = eventually(
lambda: read_rows(
'SELECT model_group, status, prompt_tokens, completion_tokens, api_base FROM "LiteLLM_SpendLogs"'
" WHERE request_id=%s",
(identity,),
),
lambda found: len(found) == 1,
seconds=70,
)
return rows[0]
def _success_row(model: str, api_base: str) -> dict[str, JsonValue]:
return {"model_group": model, "status": "success", "prompt_tokens": 9, "completion_tokens": 5, "api_base": api_base}
def _delta_text(delta: ChoiceDelta, field: str) -> str:
value: Final = delta.model_dump().get(field)
return value if isinstance(value, str) else ""
def _chunk_text(chunk: ChatCompletionChunk, field: str) -> str:
return "".join(_delta_text(choice.delta, field) for choice in chunk.choices)
def _joined(chunks: Sequence[ChatCompletionChunk], field: str) -> str:
return "".join(_chunk_text(chunk, field) for chunk in chunks)
def _upstream_requests_mentioning(gateway: Gateway, marker: str) -> list[dict[str, JsonValue]]:
observed: Final = httpx.get(f"{gateway.upstream_url}/__observations", trust_env=False, timeout=15)
observed.raise_for_status()
requests: Final = _OBSERVATIONS.validate_python(_JSON_OBJECT.validate_json(observed.content)["requests"])
return [request for request in requests if marker in json.dumps(request["body"])]
def _authorization_field(part: str) -> tuple[str, str]:
name, _, value = part.partition("=")
return name, value
def _assert_sigv4_signed(request: Request, path: str) -> None:
authorization: Final = request.headers["authorization"]
assert authorization.startswith("AWS4-HMAC-SHA256 "), dict(request.headers)
fields: Final = dict(
_authorization_field(part) for part in authorization.removeprefix("AWS4-HMAC-SHA256 ").split(", ")
)
access_key, scope = fields["Credential"].split("/", 1)
assert access_key == ACCESS_KEY, authorization
assert scope == f"{request.headers['x-amz-date'][:8]}/us-east-1/bedrock/aws4_request", authorization
assert {"host", "x-amz-date"}.issubset(fields["SignedHeaders"].split(";")), authorization
expected: Final = signature("POST", path, request.headers, fields["SignedHeaders"], request.body, SECRET_KEY, scope)
assert fields["Signature"] == expected[1], authorization
def test_openai_sdk_reasoning_request_is_served_by_native_chat_completions(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
raw: Final = _openai_client(gateway).chat.completions.with_raw_response.create(
model=model,
messages=_sdk_messages(marker),
reasoning_effort="high",
max_tokens=16,
extra_body={"cache": {"no-cache": True}},
)
completion: Final = raw.parse()
assert completion.id == f"chatcmpl-{marker}", raw.text
assert completion.choices[0].message.content == answer(marker), raw.text
assert completion.usage is not None and completion.usage.model_dump(exclude_none=True) == {
"prompt_tokens": 9,
"completion_tokens": 5,
"total_tokens": 14,
"completion_tokens_details": {"reasoning_tokens": 3},
}, raw.text
assert raw.headers["llm_provider-x-amzn-requestid"] == marker, dict(raw.headers)
request: Final = _native_request(wire)
assert _body(request) == _native_body(GPT, marker, max_completion_tokens=16, reasoning_effort="high")
assert _spend_row(completion.id) == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
async def test_async_openai_sdk_stream_keeps_the_upstream_id_and_usage(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
identity: Final = f"chatcmpl-{marker}"
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
stream: Final = await _async_openai_client(gateway).chat.completions.create(
model=model,
messages=_sdk_messages(marker),
stream=True,
stream_options={"include_usage": True},
extra_body={"cache": {"no-cache": True}},
)
chunks: Final = [chunk async for chunk in stream]
assert {chunk.id for chunk in chunks} == {identity}, chunks
assert _joined(chunks, "content") == answer(marker), chunks
usage: Final = chunks[-1].usage
assert usage is not None and (usage.prompt_tokens, usage.completion_tokens) == (9, 5), chunks[-1]
assert usage.completion_tokens_details is not None and usage.completion_tokens_details.reasoning_tokens == 3
assert all(chunk.usage is None for chunk in chunks[:-1]), chunks
assert _body(_native_request(wire)) == _streamed_native_body(GPT, marker)
assert _spend_row(identity) == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_temperature_is_forwarded_natively_when_reasoning_is_off(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, temperature=0.2, reasoning_effort="none")
payload: Final = _payload(response)
assert payload["id"] == f"chatcmpl-{marker}", response.text
assert _body(_native_request(wire)) == _native_body(GPT, marker, temperature=0.2, reasoning_effort="none")
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_temperature_while_reasoning_is_refused_before_any_wire_request(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, temperature=0.2, reasoning_effort="high")
assert response.status_code == 400, response.text
assert "UnsupportedParamsError" in response.text and "'temperature'" in response.text, response.text
assert wire.drain() == (), response.text
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (row["status"], row["model_group"], row["prompt_tokens"]) == ("failure", model, 0), row
assert "while reasoning is active" in response.text, response.text
def test_drop_params_deployment_drops_temperature_while_reasoning(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, drop_params=True)
response: Final = _chat(gateway, model, marker, temperature=0.2, reasoning_effort="high")
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(_native_request(wire)) == _native_body(GPT, marker, reasoning_effort="high")
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_guardrail_config_keeps_converse(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, guardrailConfig=dict(GUARDRAIL))
payload: Final = _payload(response)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"content": answer(marker), "role": "assistant"}}
], response.text
assert response.headers["llm_provider-x-amzn-requestid"] == marker, dict(response.headers)
body: Final = _body(_converse_request(wire))
assert body["guardrailConfig"] == GUARDRAIL, body
assert body["messages"] == [
{"role": "user", "content": [{"guardContent": {"text": {"text": _prompt(marker)}}}]}
], body
assert _spend_row(str(payload["id"])) == _success_row(model, f"{wire.url}{CONVERSE_TARGET}")
def test_converse_prefix_pins_the_model_to_converse(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, model=f"bedrock/converse/{GPT}")
response: Final = _chat(gateway, model, marker, reasoning_effort="high")
payload: Final = _payload(response)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"content": answer(marker), "role": "assistant"}}
], response.text
body: Final = _body(_converse_request(wire))
assert body["messages"] == _converse_messages(marker), body
assert body["additionalModelRequestFields"] == {"reasoning": {"effort": "high"}}, body
assert _spend_row(str(payload["id"])) == _success_row(model, f"{wire.url}{CONVERSE_TARGET}")
def test_application_inference_profile_arn_keeps_converse(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, model=f"bedrock/{PROFILE_ARN}")
response: Final = _chat(gateway, model, marker)
payload: Final = _payload(response)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"content": answer(marker), "role": "assistant"}}
], response.text
request: Final = _converse_request(wire, f"/model/{PROFILE_ARN}/converse")
assert request.target == f"/model/{quote(PROFILE_ARN, safe='')}/converse", request.target
assert _body(request)["messages"] == _converse_messages(marker), request.body
assert _spend_row(str(payload["id"])) == _success_row(
model, f"{wire.url}/model/{quote(PROFILE_ARN, safe='')}/converse"
)
def test_model_id_application_inference_profile_keeps_converse_at_the_profile_url(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, model_id=PROFILE_ARN)
response: Final = _chat(gateway, model, marker)
payload: Final = _payload(response)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"content": answer(marker), "role": "assistant"}}
], response.text
request: Final = _converse_request(wire, f"/model/{PROFILE_ARN}/converse")
assert request.target == f"/model/{quote(PROFILE_ARN, safe='')}/converse", request.target
body: Final = _body(request)
assert body["messages"] == _converse_messages(marker), request.body
assert "model_id" not in body and "model" not in body, request.body
assert _spend_row(str(payload["id"])) == _success_row(
model, f"{wire.url}/model/{quote(PROFILE_ARN, safe='')}/converse"
)
def test_stop_sequences_keep_converse(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, stop=["END"])
payload: Final = _payload(response)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"content": answer(marker), "role": "assistant"}}
], response.text
body: Final = _body(_converse_request(wire))
assert body["messages"] == _converse_messages(marker), body
assert body["inferenceConfig"] == {"stopSequences": ["END"]}, body
assert _spend_row(str(payload["id"])) == _success_row(model, f"{wire.url}{CONVERSE_TARGET}")
def test_json_object_response_format_keeps_converse(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, response_format={"type": "json_object"})
payload: Final = _payload(response)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"content": answer(marker), "role": "assistant"}}
], response.text
assert _body(_converse_request(wire))["messages"] == _converse_messages(marker), response.text
assert _spend_row(str(payload["id"])) == _success_row(model, f"{wire.url}{CONVERSE_TARGET}")
def test_json_schema_response_format_is_forwarded_natively(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, response_format=dict(JSON_SCHEMA))
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(_native_request(wire)) == _native_body(GPT, marker, response_format=dict(JSON_SCHEMA))
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_tools_while_reasoning_keep_converse(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, tools=[dict(TOOL)], reasoning_effort="high")
payload: Final = _payload(response)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"content": answer(marker), "role": "assistant"}}
], response.text
body: Final = _body(_converse_request(wire))
assert body["toolConfig"] == {"tools": [CONVERSE_TOOL]}, body
assert body["additionalModelRequestFields"] == {"reasoning": {"effort": "high"}}, body
assert _spend_row(str(payload["id"])) == _success_row(model, f"{wire.url}{CONVERSE_TARGET}")
def test_tools_with_reasoning_off_are_forwarded_natively(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, tools=[dict(TOOL)], reasoning_effort="none")
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(_native_request(wire)) == _native_body(GPT, marker, tools=[dict(TOOL)], reasoning_effort="none")
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_empty_tools_list_while_reasoning_stays_native(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker, tools=[], reasoning_effort="high")
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(_native_request(wire)) == _native_body(GPT, marker, tools=[], reasoning_effort="high")
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_chat_completions_prefix_splits_gpt_oss_reasoning_tag(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, model=f"bedrock/chat_completions/{GPT_OSS}")
raw: Final = _openai_client(gateway).chat.completions.with_raw_response.create(
model=model, messages=_sdk_messages(marker), extra_body={"cache": {"no-cache": True}}
)
completion: Final = raw.parse()
assert completion.id == f"chatcmpl-{marker}", raw.text
message: Final = completion.choices[0].message
assert message.content == answer(marker), raw.text
assert (message.model_extra or {}).get("reasoning_content") == f"why marker-{marker}", raw.text
assert _body(_native_request(wire)) == _native_body(GPT_OSS, marker)
assert _spend_row(completion.id) == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_chat_completions_prefix_splits_gpt_oss_reasoning_tag_across_stream_deltas(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
identity: Final = f"chatcmpl-{marker}"
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, model=f"bedrock/chat_completions/{GPT_OSS}")
stream: Final = _openai_client(gateway).chat.completions.create(
model=model,
messages=_sdk_messages(marker),
stream=True,
stream_options={"include_usage": True},
extra_body={"cache": {"no-cache": True}},
)
chunks: Final = list(stream)
assert {chunk.id for chunk in chunks} == {identity}, chunks
assert _joined(chunks, "reasoning_content") == f"why marker-{marker}", chunks
assert _joined(chunks, "content") == answer(marker), chunks
assert _body(_native_request(wire)) == _streamed_native_body(GPT_OSS, marker)
assert _spend_row(identity) == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_region_path_model_is_served_natively_without_the_region(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/us-west-2/{GLOBAL_GPT}", api_key=TOKEN, aws_bedrock_runtime_endpoint=wire.url
)
response: Final = _chat(gateway, model, marker)
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(_native_request(wire)) == _native_body(GLOBAL_GPT, marker)
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_sigv4_deployment_signs_the_native_request(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/{GPT}",
api_key=None,
aws_access_key_id=ACCESS_KEY,
aws_secret_access_key=SECRET_KEY,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
)
response: Final = _chat(gateway, model, marker)
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
request: Final = _only_request(wire)
assert (request.method, target_of(request)) == ("POST", NATIVE_TARGET), request.target
_assert_sigv4_signed(request, NATIVE_TARGET)
assert _body(request) == _native_body(GPT, marker)
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_blank_api_key_on_a_sigv4_deployment_is_signed_not_sent_as_an_empty_bearer(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/{GPT}",
api_key="",
aws_access_key_id=ACCESS_KEY,
aws_secret_access_key=SECRET_KEY,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
)
response: Final = _chat(gateway, model, marker)
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
request: Final = _only_request(wire)
assert (request.method, target_of(request)) == ("POST", NATIVE_TARGET), request.target
_assert_sigv4_signed(request, NATIVE_TARGET)
assert _body(request) == _native_body(GPT, marker)
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_runtime_endpoint_without_api_base_is_used_natively(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, api_base=None)
response: Final = _chat(gateway, model, marker)
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(_native_request(wire)) == _native_body(GPT, marker)
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
def test_runtime_endpoint_wins_over_an_unrelated_api_base(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
response: Final = _chat(gateway, model, marker)
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(_native_request(wire)) == _native_body(GPT, marker)
assert _upstream_requests_mentioning(gateway, marker) == [], response.text
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")
@pytest.mark.parametrize("suffix", ["/openai/v1", "/openai/v1/chat/completions"])
def test_api_base_already_naming_the_native_path_is_not_doubled(gateway: Gateway, suffix: str) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model_info=None, **{**GPT_DEPLOYMENT, "api_base": f"{wire.url}{suffix}"})
response: Final = _chat(gateway, model, marker)
request: Final = _only_request(wire)
assert (request.method, request.target) == ("POST", NATIVE_TARGET), response.text
assert _payload(response)["id"] == f"chatcmpl-{marker}", response.text
assert _body(request) == _native_body(GPT, marker)
assert _spend_row(f"chatcmpl-{marker}") == _success_row(model, f"{wire.url}{NATIVE_TARGET}")