scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-20 05:04:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 155

Merge remote-tracking branch 'origin/main' into litellm_ws_spend_log_per_turn

Browse source 
# Conflicts:
#	tests/unit/responses/test_responses_websocket_all_providers.py
This commit is contained in:
mateo-berri 2026-10-15 19:13:33 -07:00
commit 10d6748a90
169 changed files with 6490 additions and 1060 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
.circleci/config.yml
View file

@ -1293,6 +1293,19 @@ jobs:
uv venv --python 3.12 .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
- when:
condition:
equal: [sdk, << parameters.suite >>]
steps:
- run:
name: Install redis-server 7.2.16
command: |
mkdir -p /tmp/redis-bin
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
docker cp "${cid}:/usr/local/bin/redis-server" /tmp/redis-bin/redis-server
docker rm "$cid"
/tmp/redis-bin/redis-server --version | grep -q 'v=7.2.16'
echo 'export PATH="/tmp/redis-bin:$PATH"' >> "$BASH_ENV"
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"

13
.github/workflows/test-unit.yml vendored
View file

@ -565,7 +565,6 @@ jobs:
- shard: mcp-elicitation
test-path: |-
--override-ini=pythonpath=tests
tests/unit/experimental_mcp_client/test_mcp_client.py
tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
tests/unit/proxy/_experimental/mcp_server/test_interactions.py
@ -580,17 +579,6 @@ jobs:
job-timeout-minutes: 60
dist: loadscope
- shard: mcp-interactions
test-path: |-
--override-ini=pythonpath=tests
tests/integration/mcp/test_interactions.py
python-version: "3.12"
workers: 4
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60
dist: load
- shard: proxy-infra
test-path: |-
tests/unit/proxy/db
@ -1463,7 +1451,6 @@ jobs:
tests/unit/caching/test_redis_connection_pool.py \
tests/unit/caching/test_redis_cluster_cache.py \
tests/unit/caching/test_evicted_client_closer.py \
tests/integration/sdk/test_redis_cluster_iam_auth.py \
--tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \
--cov=./litellm --cov-report=xml:coverage.xml

9
ci_cd/generate_model_prices_schema.py
View file

@ -189,6 +189,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = {
**NONNEG_INTEGER,
"description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
},
"token_count_framing": {
**NONNEG_INTEGER,
"description": "Tokens added once per request to a local count estimate, for the provider's chat framing.",
},
"tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
"rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
}
@ -283,6 +287,11 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
},
"comment": STRING,
"token_count_fallback_model": {
"type": "string",
"minLength": 1,
"description": "Model that counts tokens for this one when the provider cannot count it.",
},
"audio_transcription_config": STRING,
"vertex_ai_audio_api": {
"type": "string",

81
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json
View file

@ -6563,6 +6563,87 @@
],
"title": "litellm_project_model_rate_limit_used_metric",
"type": "timeseries"
},
{
"collapsed": false,
"gridPos": {
"h": 1,
"w": 24,
"x": 0,
"y": 447
},
"id": 115,
"panels": [],
"title": "Agent tracing export (needs LITELLM_LENS_URL)",
"type": "row"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"description": "Request records the Lens exporter delivered to the Lens service against the records it dropped before the service accepted them",
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"drawStyle": "line",
"fillOpacity": 10,
"lineWidth": 1,
"showPoints": "never",
"spanNulls": false
},
"unit": "short"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 24,
"x": 0,
"y": 448
},
"id": 116,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_lens_export_rows_written_total[$__rate_interval]))",
"legendFormat": "written",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_lens_export_rows_dropped_total[$__rate_interval]))",
"legendFormat": "dropped",
"range": true,
"refId": "B"
}
],
"title": "litellm_lens_export_rows_written and litellm_lens_export_rows_dropped rate",
"type": "timeseries"
}
],
"preload": false,

4
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md
View file

@ -1,10 +1,10 @@
# LiteLLM All Prometheus Metrics dashboard
Every `litellm_*` metric family the proxy can expose on `/metrics` (141 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
Every `litellm_*` metric family the proxy can expose on `/metrics` (143 families across 98 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes), project model rate limits, and the agent tracing export to Lens. Panel titles are the metric names so you can grep the JSON for the metric you care about
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
The first eleven rows and the project model rate limits row need only `callbacks: ["prometheus"]`. The other rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, the agent tracing export row needs `LITELLM_LENS_URL`, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
## Pre-requisites

37
litellm-rust/crates/inference-messages/src/prepare.rs
View file

@ -124,7 +124,7 @@ fn prepare_provider_request(
url,
body: transformed,
environment,
timeout,
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
api_key: api_key.map(SecretValue::new),
})
}
@ -253,6 +253,41 @@ mod tests {
);
}
#[rstest]
#[case::nothing_configured(None, None, None)]
#[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))]
#[case::call_timeout_wins(
Some(Duration::from_secs(30)),
Some(1800.0),
Some(Duration::from_secs(30))
)]
#[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))]
#[case::unusable_setting_ignored(None, Some(-1.0), None)]
#[case::zero_setting_ignored(None, Some(0.0), None)]
fn timeout_falls_back_to_the_configured_request_timeout(
shaping: MessagesShaping,
#[case] call_timeout: Option<Duration>,
#[case] request_timeout: Option<f64>,
#[case] expected: Option<Duration>,
) {
let prepared = prepare(MessagesCall {
body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: call_timeout,
shaping: MessagesShaping {
settings: MessagesSettings { request_timeout, ..shaping.settings },
..shaping
},
})
.unwrap();
assert_eq!(prepared.timeout, expected);
}
fn prepared_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
prepare(MessagesCall {
body: body(fields),

14
litellm-rust/crates/inference-messages/src/types.rs
View file

@ -60,6 +60,16 @@ pub struct MessagesSettings {
pub reasoning_auto_summary: bool,
#[serde(default)]
pub additional_drop_params: Vec<String>,
#[serde(default)]
pub request_timeout: Option<f64>,
}
impl MessagesSettings {
pub fn request_timeout(&self) -> Option<Duration> {
self.request_timeout
.and_then(|secs| Duration::try_from_secs_f64(secs).ok())
.filter(|timeout| !timeout.is_zero())
}
}
#[cfg(test)]
@ -120,13 +130,15 @@ mod tests {
},
"drop_params": true,
"reasoning_auto_summary": true,
"additional_drop_params": ["metadata.user_id", "thinking"]
"additional_drop_params": ["metadata.user_id", "thinking"],
"request_timeout": 1800.0
}),
MessagesShaping {
settings: MessagesSettings {
drop_params: true,
reasoning_auto_summary: true,
additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()],
request_timeout: Some(1800.0),
},
capabilities: MessagesModelCapabilities {
supports_reasoning: true,

2
litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs
View file

@ -23,7 +23,7 @@ fn run_chat_completions(
) -> PyResult<Py<PyAny>> {
let host = InferenceHost::new(
call.view()?,
"litellm.rust_bridge.chat_completions.route_host",
"litellm.rust_bridge.host.routes.chat_completions",
);
run_inference::<ChatCompletionsRoute, _>(
py,

10
litellm-rust/crates/python-bridge/src/routes/inference.rs
View file

@ -21,10 +21,7 @@ use serde_json::{Map, Value};
use super::{NativeCall, RequestView};
use crate::{
errors::{RustUpstreamError, route_error_to_pyerr},
marshal::{
RouteOptions, optional_timeout, project_optional_fields, public_response,
python_timeout_seconds,
},
marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds},
};
pub(super) struct InferenceHost {
@ -115,10 +112,7 @@ impl InferenceHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<Map<String, Value>> {
let names: Vec<String> = py.import(self.module)?.getattr("PARAMETERS")?.extract()?;
project_optional_fields(names.iter().map(String::as_str), |name| {
self.argument(py, arguments, name)
})
self.view.body_params(py, arguments, self.module)
}
pub fn response(&self, py: Python<'_>, response: &impl Serialize) -> PyResult<Py<PyAny>> {

29
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -26,34 +26,9 @@ use crate::{
routes::RequestView,
};
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host";
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages";
const REQUEST_ERROR_MARKER: &str = "messages_request_error";
const BODY_FIELDS: [&str; 22] = [
"max_tokens",
"metadata",
"stop_sequences",
"stream",
"system",
"temperature",
"thinking",
"tool_choice",
"tools",
"top_k",
"inference_geo",
"top_p",
"mcp_servers",
"context_management",
"compaction",
"container",
"output_format",
"speed",
"output_config",
"cache_control",
"reasoning_effort",
"safeguards",
];
fn merge_headers(
forwarded: Option<Map<String, Value>>,
extra_headers: Option<Map<String, Value>>,
@ -143,7 +118,7 @@ impl MessagesPythonHost {
let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?;
let messages =
argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?;
let fields = project_optional_fields(BODY_FIELDS, argument)?;
let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?;
let body = [
("model".to_string(), Value::String(model.clone())),
("messages".to_string(), from_py(&messages)?),

18
litellm-rust/crates/python-bridge/src/routes/mod.rs
View file

@ -13,6 +13,7 @@ use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCal
use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
use litellm_host_python::{
CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
from_py,
};
use litellm_host_python::{Error as HostError, present};
use pyo3::{
@ -20,6 +21,7 @@ use pyo3::{
prelude::*,
types::{PyDict, PyMapping, PyTuple},
};
use serde_json::{Map, Value};
/// The public call as Python bound it: `base` holds the positional arguments by name plus
/// the signature defaults, `kwargs` the caller's keyword dict.
@ -92,6 +94,22 @@ impl RequestView {
present(arguments, self.request(py)?, name)
}
/// The body params `module.body_params` keeps from the keywords the call would send, so
/// each route's Python host decides which params reach the provider.
pub(crate) fn body_params(
&self,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
module: &str,
) -> PyResult<Map<String, Value>> {
let request = effective_py_args(self.request(py)?, arguments)?;
from_py(
&py.import(module)?
.getattr("body_params")?
.call1((request,))?,
)
}
pub(crate) fn close(&mut self) {
self.0 = None;
}

4
litellm-rust/crates/python-bridge/src/routes/ocr/host.rs
View file

@ -77,7 +77,7 @@ impl OcrPythonHost {
_ => "",
};
let mapped = py
.import("litellm.rust_bridge.ocr.route_host")
.import("litellm.rust_bridge.host.routes.ocr")
.and_then(|module| module.getattr("map_failure"))
.and_then(|map| map.call1((error.value(py), self.view.request(py)?, provider)))
.and_then(|mapped| mapped.extract::<Py<PyBaseException>>().map_err(PyErr::from));
@ -115,7 +115,7 @@ impl PythonBinding for OcrPythonHost {
py: Python<'_>,
response: LiteLLMOcrResponse,
) -> PyResult<Py<PyAny>> {
public_response(py, "litellm.rust_bridge.ocr.route_host", &response)
public_response(py, "litellm.rust_bridge.host.routes.ocr", &response)
}
fn encode_stream_head(

2
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
View file

@ -21,7 +21,7 @@ use super::{
};
use crate::errors::RustBridgeDeclined;
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.responses.route_host";
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
let resolved = call.resolved()?;

3
litellm/constants.py
View file

@ -380,6 +380,7 @@ REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = get_env_int_or_none("REALTIME
REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = get_env_float(
"REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 20.0
)
BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS: Final = get_env_float("BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS", 5.0)
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
@ -968,7 +969,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES: Final = {
"context_management": None,
}
openai_compatible_endpoints: Final[list] = [
openai_compatible_endpoints: Final[list[str]] = [
"api.perplexity.ai",
"api.endpoints.anyscale.com/v1",
"api.deepinfra.com/v1/openai",

4
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -61,6 +61,10 @@ def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool:
return url_path == endpoint_path or url_path.startswith(endpoint_path + "/")
def is_openai_compatible_endpoint(api_base: str) -> bool:
return any(_endpoint_matches_api_base(endpoint, api_base) for endpoint in litellm.openai_compatible_endpoints)
def is_non_openai_azure_model(model: str) -> bool:
try:
model_name: Final = model.split("/", 1)[1]

42
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4552,6 +4552,32 @@ def get_masked_values(
_get_masked_values = get_masked_values
@lru_cache(maxsize=1)
def start_sentry(dsn: str) -> None:
try:
import sentry_sdk
except ImportError:
print_verbose("Package 'sentry_sdk' is missing. Installing it...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "sentry_sdk"])
import sentry_sdk
from litellm.litellm_core_utils.sentry_scrubbing import build_sentry_init_options
sentry_sdk.init(**build_sentry_init_options({**os.environ, "SENTRY_DSN": dsn}))
def initialize_sentry() -> None:
global sentry_sdk_instance, capture_exception, add_breadcrumb
dsn: Final = os.environ.get("SENTRY_DSN")
if not dsn:
return
start_sentry(dsn)
import sentry_sdk
sentry_sdk_instance = sentry_sdk # rebind-ok: process-wide Sentry SDK handle read by every logging path
capture_exception = sentry_sdk.capture_exception # rebind-ok: process-wide Sentry SDK handle
add_breadcrumb = sentry_sdk.add_breadcrumb # rebind-ok: process-wide Sentry SDK handle
def set_callbacks(
callback_list: Iterable[str | Callable[..., object] | CustomLogger],
function_id: str | None = None,
@ -4560,9 +4586,6 @@ def set_callbacks(
Globally sets the callback client
"""
global \
sentry_sdk_instance, \
capture_exception, \
add_breadcrumb, \
slack_app, \
alerts_channel, \
traceloopLogger, \
@ -4586,18 +4609,7 @@ def set_callbacks(
try:
for callback in callback_list:
if callback == "sentry":
try:
import sentry_sdk
except ImportError:
print_verbose("Package 'sentry_sdk' is missing. Installing it...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "sentry_sdk"])
import sentry_sdk
from litellm.litellm_core_utils.sentry_scrubbing import build_sentry_init_options
sentry_sdk_instance = sentry_sdk
sentry_sdk_instance.init(**build_sentry_init_options(os.environ))
capture_exception = sentry_sdk_instance.capture_exception
add_breadcrumb = sentry_sdk_instance.add_breadcrumb
initialize_sentry()
elif callback == "slack":
try:
from slack_bolt import App

5
litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
View file

@ -16,6 +16,7 @@ from litellm.litellm_core_utils.hidden_params import get_hidden_params
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_reasoning_content,
)
from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer
from litellm.types.llms.databricks import DatabricksTool
from litellm.types.llms.openai import (
ChatCompletionReasoningItem,
@ -833,7 +834,9 @@ def convert_to_model_response_object(
if k not in _MODEL_RESPONSE_FIELDS:
setattr(model_response_object, k, v)
return model_response_object
ollama_ids: Final = OllamaResponseIdNormalizer()
ollama_ids.observe(model_response_object)
return ollama_ids.apply(model_response_object)
elif response_type == "embedding" and (
model_response_object is None or isinstance(model_response_object, EmbeddingResponse)
):

6
litellm/litellm_core_utils/sentry_scrubbing.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
import re
from collections.abc import Callable, Mapping, Sequence
from functools import reduce
from typing import TYPE_CHECKING, Final, TypeAlias, cast
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast
from pydantic import JsonValue
from sentry_sdk.scrubber import DEFAULT_DENYLIST, DEFAULT_PII_DENYLIST, EventScrubber
@ -65,6 +65,8 @@ class SentryInitOptions(TypedDict):
traces_sample_rate: ReadOnly[float]
sample_rate: ReadOnly[float]
send_default_pii: ReadOnly[bool]
max_request_body_size: ReadOnly[Literal["never"]]
include_local_variables: ReadOnly[bool]
event_scrubber: ReadOnly[EventScrubber]
before_send: ReadOnly[EventScrubFn]
before_send_transaction: ReadOnly[EventScrubFn]
@ -140,6 +142,8 @@ def build_sentry_init_options(env: Mapping[str, str]) -> SentryInitOptions:
traces_sample_rate=float(env.get("SENTRY_API_TRACE_RATE") or "1.0"),
sample_rate=float(env.get("SENTRY_API_SAMPLE_RATE") or "1.0"),
send_default_pii=send_default_pii,
max_request_body_size="never",
include_local_variables=send_default_pii,
event_scrubber=EventScrubber(
denylist=list(SECRET_FIELD_NAMES),
pii_denylist=list(PII_FIELD_NAMES),

7
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -1,4 +1,5 @@
import base64
import builtins
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass, replace
@ -118,6 +119,7 @@ class _BaseChunk(TypedDict, total=False):
created: ReadOnly[int]
model: ReadOnly[str]
system_fingerprint: ReadOnly[str | None]
provider_specific_fields: ReadOnly[Mapping[str, builtins.object] | None]
service_tier: ReadOnly[str | None]
choices: ReadOnly[Required[Sequence[StreamingChoices]]]
_hidden_params: ReadOnly[_ChunkHiddenParams]
@ -430,6 +432,11 @@ class ChunkProcessor:
"created": created,
"model": model,
"system_fingerprint": system_fingerprint,
**(
{"provider_specific_fields": chunks[0].get("provider_specific_fields")}
if system_fingerprint == "fp_ollama"
else {}
),
**(
MappingProxyType({"service_tier": service_tier})
if service_tier is not None

11
litellm/litellm_core_utils/streaming_handler.py
View file

@ -14,6 +14,7 @@ from typing import Any, Final, NoReturn, Protocol, TypeVar, cast
import anyio
import httpx
from openai.types.chat import ChatCompletionChunk
from pydantic import BaseModel, ValidationError
from typing_extensions import NotRequired, TypedDict
@ -27,6 +28,7 @@ from litellm.litellm_core_utils.model_response_utils import (
)
from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import OpenAIChatCompletionChunk
from litellm.types.router import GenericLiteLLMParams
@ -285,6 +287,7 @@ class CustomStreamWrapper:
self._response_headers = _response_headers
self.response_id: str | None = None
self._ollama_response_ids = OllamaResponseIdNormalizer()
self.received_model_id: str | None = None
self.logging_loop = None
self.rules = Rules()
@ -824,7 +827,7 @@ class CustomStreamWrapper:
pass
else:
model_response.choices = [StreamingChoices(finish_reason=None)]
return model_response
return self._ollama_response_ids.apply(model_response)
def is_delta_empty(self, delta: Delta) -> bool:
is_empty = True
@ -1449,8 +1452,14 @@ class CustomStreamWrapper:
return _ProviderChunkParsed(response_obj)
def chunk_creator(self, chunk: Any):
if self.custom_llm_provider != "cached_response" and isinstance(
chunk, (ChatCompletionChunk, ModelResponseStream)
):
self._ollama_response_ids.observe(chunk)
if hasattr(chunk, "id"):
self.response_id = chunk.id
if self._ollama_response_ids.response_id is not None:
self.response_id = self._ollama_response_ids.response_id
provider_response_model: Final = _provider_response_model(chunk)
if provider_response_model is not None:
self._provider_response_model = provider_response_model

67
litellm/llms/bedrock/realtime/handler.py
View file

@ -17,7 +17,9 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import redact_string, verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS,
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
BEDROCK_REALTIME_SDK_DISTRIBUTION,
@ -35,7 +37,7 @@ from litellm.types.realtime import RealtimeResponseTransformInput
from ..base_aws_llm import BaseAWSLLM, run_aws_signing
from ..common_utils import BedrockError
from .transformation import BedrockRealtimeConfig
from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection
_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None)
_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
@ -232,12 +234,43 @@ async def _ack_session_update(
return True
async def _reject_session_update(
client_ws: RealtimeClientWebSocket,
parsed_client_message: Mapping[str, JsonValue],
rejection: BedrockSessionUpdateRejection,
) -> bool:
"""Send an error event for a session.update Bedrock cannot apply; False means the client is gone."""
client_event_id: Final = parsed_client_message.get("event_id")
error_event: Final = {
"type": "error",
"event_id": f"event_{uuid.uuid4()}",
"error": {
"type": "invalid_request_error",
"code": rejection.code,
"message": rejection.message,
"param": rejection.param,
"event_id": client_event_id if isinstance(client_event_id, str) else None,
},
}
try:
await client_ws.send_text(json.dumps(error_event))
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
return False
return True
class BedrockRealtime(BaseAWSLLM):
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
def __init__(self, sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version):
def __init__(
self,
sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version,
close_timeout: float = BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS,
):
super().__init__()
self._sdk_version_lookup: Final = sdk_version_lookup
self._close_timeout: Final = close_timeout
async def async_realtime(
self,
@ -259,6 +292,7 @@ class BedrockRealtime(BaseAWSLLM):
aws_bedrock_runtime_endpoint: str | None = None,
aws_external_id: str | None = None,
aws_session_tags: object = None,
drop_params: bool = False,
**kwargs: object,
):
"""
@ -341,7 +375,9 @@ class BedrockRealtime(BaseAWSLLM):
)
try:
await self._run_session(websocket, open_bidirectional_stream, model, logging_obj, pending_session_update)
await self._run_session(
websocket, open_bidirectional_stream, model, logging_obj, pending_session_update, drop_params
)
finally:
await _close_bedrock_client(bedrock_client)
@ -352,8 +388,9 @@ class BedrockRealtime(BaseAWSLLM):
model: str,
logging_obj: LiteLLMLogging,
pending_session_update: str | None,
drop_params: bool = False,
) -> None:
transformation_config: Final = BedrockRealtimeConfig()
transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params)
bedrock_stream: Final = await open_bidirectional_stream()
@ -489,6 +526,11 @@ class BedrockRealtime(BaseAWSLLM):
parsed_client_message = _parse_client_message(message)
is_session_update = _json_str(parsed_client_message.get("type")) == "session.update"
if is_session_update:
rejection = transformation_config.session_update_rejection(parsed_client_message)
if rejection is not None:
if not await _reject_session_update(client_ws, parsed_client_message, rejection):
break
continue
client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = message
transformed_messages = transformation_config.transform_realtime_request(
@ -504,11 +546,20 @@ class BedrockRealtime(BaseAWSLLM):
):
break
finally:
for close_message in transformation_config.session_close_messages():
async def close_bedrock_input() -> None:
for close_message in transformation_config.session_close_messages():
with contextlib.suppress(Exception):
await send_to_bedrock(close_message)
with contextlib.suppress(Exception):
await send_to_bedrock(close_message)
with contextlib.suppress(Exception):
await bedrock_stream.input_stream.close()
await bedrock_stream.input_stream.close()
try:
await asyncio.wait_for(close_bedrock_input(), self._close_timeout)
except (asyncio.TimeoutError, TimeoutError):
verbose_proxy_logger.debug(
"Bedrock Realtime: timed out closing Bedrock input stream after %.1fs", self._close_timeout
)
async def _forward_bedrock_to_client(
self,

617
litellm/llms/bedrock/realtime/transformation.py
View file

@ -5,13 +5,16 @@ Transforms between OpenAI Realtime API format and Bedrock Nova Sonic format.
"""
import base64
import dataclasses
import json
import math
import uuid as uuid_lib
from collections.abc import Mapping
from typing import Final, cast
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final, Literal
import httpx
from pydantic import TypeAdapter
from pydantic import ConfigDict, JsonValue, TypeAdapter
from litellm._logging import verbose_logger
from litellm._uuid import uuid
@ -24,6 +27,8 @@ from litellm.types.llms.openai import (
OpenAIRealtimeContentPartDone,
OpenAIRealtimeDoneEvent,
OpenAIRealtimeEvents,
OpenAIRealtimeFunctionCallArgumentsDelta,
OpenAIRealtimeFunctionCallArgumentsDone,
OpenAIRealtimeInputAudioBufferSpeechEvent,
OpenAIRealtimeInputAudioTranscriptionCompleted,
OpenAIRealtimeInputAudioTranscriptionDelta,
@ -35,7 +40,9 @@ from litellm.types.llms.openai import (
OpenAIRealtimeResponseTextDone,
OpenAIRealtimeResponseUsage,
OpenAIRealtimeStreamResponseBaseObject,
OpenAIRealtimeStreamResponseOutputItem,
OpenAIRealtimeStreamResponseOutputItemAdded,
OpenAIRealtimeStreamResponseOutputItemContent,
OpenAIRealtimeStreamSession,
OpenAIRealtimeStreamSessionEvents,
OpenAIRealtimeUsageTokenDetails,
@ -52,6 +59,26 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
class BedrockContentEnd(LiteLLMBaseModel):
stopReason: str | None = None
type: str | None = None
def _tool_function(tool: JsonValue) -> Mapping[str, JsonValue] | None:
if not isinstance(tool, dict) or tool.get("type") != "function":
return None
function: Final = tool.get("function")
return function if isinstance(function, dict) and function else tool
class BedrockToolUse(LiteLLMBaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
toolUseId: str = ""
toolName: str = ""
content: JsonValue = "{}"
@property
def arguments(self) -> str:
return self.content if isinstance(self.content, str) else json.dumps(self.content)
class BedrockUsageTokenDetails(LiteLLMBaseModel):
@ -80,12 +107,63 @@ TRIGGER_AUDIO_BYTES_PER_SECOND: Final = TRIGGER_AUDIO_SAMPLE_RATE_HERTZ * 2
TRIGGER_LEADING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND // 2)
TRIGGER_TRAILING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND * 3)
TRIGGER_AUDIO_CHUNK_SIZE: Final = 1024
PROMPT_STARTING_CLIENT_EVENT_TYPES: Final = frozenset(
("input_audio_buffer.append", "input_audio_buffer.commit", "conversation.item.create", "response.create")
)
def _number_setting(value: JsonValue) -> float | None:
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
return None
return float(value)
def _integer_setting(value: JsonValue) -> int | None:
number: Final = _number_setting(value)
return None if number is None else int(number)
def _malformed_number_settings(session_config: Mapping[str, JsonValue]) -> tuple[str, ...]:
temperature: Final = session_config.get("temperature")
max_output_tokens: Final = session_config.get("max_response_output_tokens")
output_rate_hertz: Final = session_config.get("output_sample_rate_hertz")
malformed: Final = (
("temperature", temperature is not None and _number_setting(temperature) is None),
(
"max_response_output_tokens",
max_output_tokens is not None
and max_output_tokens != "inf"
and _integer_setting(max_output_tokens) is None,
),
("output_sample_rate_hertz", output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None),
)
return tuple(name for name, is_malformed in malformed if is_malformed)
@dataclasses.dataclass(frozen=True, slots=True)
class BedrockSessionUpdateRejection:
code: Literal["invalid_type", "session_update_not_applied"]
message: str
param: str
@dataclasses.dataclass(frozen=True, slots=True)
class BedrockLockedSessionSettings:
"""The session settings Bedrock Nova Sonic locks in once the prompt starts."""
tools_json: str
system_prompt: str
voice_id: str
output_sample_rate_hertz: int
temperature: float
max_tokens: int | None
class BedrockRealtimeConfig(BaseRealtimeConfig):
"""Configuration for Bedrock Nova Sonic realtime transformations."""
def __init__(self):
def __init__(self, drop_params: bool = False):
self.drop_params = drop_params
# Track session state
self.prompt_name = str(uuid_lib.uuid4())
self.content_name = str(uuid_lib.uuid4())
@ -95,7 +173,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
# Default configuration values
# Inference configuration
self.max_tokens = 1024
self.max_tokens: int | None = 1024
self.top_p = 0.9
self.temperature = 0.7
@ -128,6 +206,19 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
self._user_transcript_buffer = ""
self._cumulative_usage = BedrockUsageEvent()
self._reported_usage = BedrockUsageEvent()
self._response_done_items: Mapping[int, OpenAIRealtimeStreamResponseOutputItem] = MappingProxyType({})
self._response_item_count = 0
self._message_output_index = 0
self._assistant_block_open = False
self._response_done_pending = False
self._started_settings = BedrockLockedSessionSettings(
tools_json="[]",
system_prompt="",
voice_id=self.voice_id,
output_sample_rate_hertz=self.output_sample_rate_hertz,
temperature=self.temperature,
max_tokens=self.max_tokens,
)
def get_error_class(
self,
@ -149,7 +240,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
"""Bedrock requires session configuration."""
return True
def session_configuration_request(self, model: str, tools: list[dict] | None = None) -> str:
def session_configuration_request(self, model: str, tools: Sequence[JsonValue] | None = None) -> str:
"""
Create initial session configuration for Bedrock Nova Sonic.
@ -159,19 +250,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
Returns JSON string with session start and prompt start events.
"""
session_start: Final = {
"event": {
"sessionStart": {
"inferenceConfiguration": {
"maxTokens": self.max_tokens,
"topP": self.top_p,
"temperature": self.temperature,
}
}
}
}
session_start: Final = {"event": {"sessionStart": {"inferenceConfiguration": self._inference_configuration()}}}
prompt_start_config: Final = {
prompt_start_config: Final[Mapping[str, object]] = {
"promptName": self.prompt_name,
"textOutputConfiguration": {"mediaType": self.text_media_type},
"audioOutputConfiguration": {
@ -183,41 +264,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
"encoding": self.output_encoding,
"audioType": self.output_audio_type,
},
**(
{
"toolUseOutputConfiguration": {"mediaType": "application/json"},
"toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)},
}
if tools
else {}
),
}
# Add tool configuration if tools are provided
if tools:
prompt_start_config["toolUseOutputConfiguration"] = {"mediaType": "application/json"}
prompt_start_config["toolConfiguration"] = {"tools": self._transform_tools_to_bedrock_format(tools)}
prompt_start: Final = {"event": {"promptStart": prompt_start_config}}
# Return as a marker that we've sent the configuration
return json.dumps({"session_start": session_start, "prompt_start": prompt_start})
def _transform_tools_to_bedrock_format(self, tools: list[dict]) -> list[dict]:
"""
Transform OpenAI tool format to Bedrock tool format.
Args:
tools: List of OpenAI format tools
Returns:
List of Bedrock format tools
"""
bedrock_tools: Final = []
for tool in tools:
if tool.get("type") == "function":
function = tool.get("function", {})
bedrock_tool = {
"toolSpec": {
"name": function.get("name", ""),
"description": function.get("description", ""),
"inputSchema": {"json": json.dumps(function.get("parameters", {}))},
}
def _transform_tools_to_bedrock_format(self, tools: Sequence[JsonValue]) -> tuple[Mapping[str, JsonValue], ...]:
functions: Final = (_tool_function(tool) for tool in tools)
return tuple(
{
"toolSpec": {
"name": function.get("name", ""),
"description": function.get("description", ""),
"inputSchema": {"json": json.dumps(function.get("parameters", {}))},
}
bedrock_tools.append(bedrock_tool)
return bedrock_tools
}
for function in functions
if function is not None
)
def _map_audio_format_to_sample_rate(self, audio_format: str, is_output: bool = True) -> int:
"""
@ -252,47 +326,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
messages: Final[list[str]] = []
session_config: Final = json_message.get("session", {})
locked_settings: Final = self._apply_locked_settings(session_config, self._started_settings)
# Update inference configuration from session if provided
if "max_response_output_tokens" in session_config:
self.max_tokens = session_config["max_response_output_tokens"]
if "temperature" in session_config:
self.temperature = session_config["temperature"]
# Update audio output configuration from session if provided
if "voice" in session_config:
self.voice_id = session_config["voice"]
if "output_audio_format" in session_config:
output_format: Final = session_config["output_audio_format"]
self.output_sample_rate_hertz = self._map_audio_format_to_sample_rate(output_format, is_output=True)
# Settings Bedrock locks at prompt start may only change before it
if not self.prompt_started:
self.max_tokens = locked_settings.max_tokens
self.temperature = locked_settings.temperature
self.voice_id = locked_settings.voice_id
self.output_sample_rate_hertz = locked_settings.output_sample_rate_hertz
# Update audio input configuration from session if provided
if "input_audio_format" in session_config:
input_format: Final = session_config["input_audio_format"]
self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False)
# Allow direct override of sample rates if provided (custom extension)
if "output_sample_rate_hertz" in session_config:
self.output_sample_rate_hertz = session_config["output_sample_rate_hertz"]
# Allow direct override of the input sample rate if provided (custom extension)
if "input_sample_rate_hertz" in session_config:
self.input_sample_rate_hertz = session_config["input_sample_rate_hertz"]
if self.prompt_started:
return messages
# Send session start
session_start: Final = {
"event": {
"sessionStart": {
"inferenceConfiguration": {
"maxTokens": self.max_tokens,
"topP": self.top_p,
"temperature": self.temperature,
}
}
}
}
session_start: Final = {"event": {"sessionStart": {"inferenceConfiguration": self._inference_configuration()}}}
messages.append(json.dumps(session_start))
# Send prompt start
prompt_start_config: Final = {
tools: Final[Sequence[JsonValue] | None] = session_config.get("tools")
prompt_start_config: Final[Mapping[str, object]] = {
"promptName": self.prompt_name,
"textOutputConfiguration": {"mediaType": self.text_media_type},
"audioOutputConfiguration": {
@ -304,63 +365,138 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
"encoding": self.output_encoding,
"audioType": self.output_audio_type,
},
**(
{
"toolUseOutputConfiguration": {"mediaType": "application/json"},
"toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)},
}
if tools
else {}
),
}
# Add tool configuration if tools are provided
tools: Final = session_config.get("tools")
if tools:
prompt_start_config["toolUseOutputConfiguration"] = {"mediaType": "application/json"}
prompt_start_config["toolConfiguration"] = {"tools": self._transform_tools_to_bedrock_format(tools)}
prompt_start: Final = {"event": {"promptStart": prompt_start_config}}
messages.append(json.dumps(prompt_start))
self.prompt_started = True
self._started_settings = locked_settings
# Send system prompt if provided
instructions: Final = session_config.get("instructions")
if instructions:
text_content_name: Final = str(uuid_lib.uuid4())
# Bedrock requires the first content block in a prompt to be SYSTEM role
instructions: Final = locked_settings.system_prompt
text_content_name: Final = str(uuid_lib.uuid4())
# Content start
text_content_start: Final = {
"event": {
"contentStart": {
"promptName": self.prompt_name,
"contentName": text_content_name,
"type": "TEXT",
"interactive": False,
"role": "SYSTEM",
"textInputConfiguration": {"mediaType": self.text_media_type},
}
# Content start
text_content_start: Final = {
"event": {
"contentStart": {
"promptName": self.prompt_name,
"contentName": text_content_name,
"type": "TEXT",
"interactive": False,
"role": "SYSTEM",
"textInputConfiguration": {"mediaType": self.text_media_type},
}
}
messages.append(json.dumps(text_content_start))
}
messages.append(json.dumps(text_content_start))
# Text input
text_input: Final = {
"event": {
"textInput": {
"promptName": self.prompt_name,
"contentName": text_content_name,
"content": instructions,
}
# Text input
text_input: Final = {
"event": {
"textInput": {
"promptName": self.prompt_name,
"contentName": text_content_name,
"content": instructions,
}
}
messages.append(json.dumps(text_input))
}
messages.append(json.dumps(text_input))
# Content end
text_content_end: Final = {
"event": {
"contentEnd": {
"promptName": self.prompt_name,
"contentName": text_content_name,
}
# Content end
text_content_end: Final = {
"event": {
"contentEnd": {
"promptName": self.prompt_name,
"contentName": text_content_name,
}
}
messages.append(json.dumps(text_content_end))
}
messages.append(json.dumps(text_content_end))
return messages
def _inference_configuration(self) -> Mapping[str, int | float]:
return {
"topP": self.top_p,
"temperature": self.temperature,
**({"maxTokens": self.max_tokens} if self.max_tokens is not None else {}),
}
def _apply_locked_settings(
self, session_config: Mapping[str, JsonValue], base: BedrockLockedSessionSettings
) -> BedrockLockedSessionSettings:
output_audio_format: Final = session_config.get("output_audio_format")
temperature: Final = session_config.get("temperature")
max_output_tokens: Final = session_config.get("max_response_output_tokens")
requested_tools: Final = session_config.get("tools")
instructions: Final = session_config.get("instructions")
voice: Final = session_config.get("voice")
tools_json: Final = (
json.dumps(self._transform_tools_to_bedrock_format(requested_tools))
if isinstance(requested_tools, list)
else base.tools_json
)
system_prompt: Final = base.system_prompt if instructions is None else str(instructions)
voice_id: Final = base.voice_id if voice is None else str(voice)
format_sample_rate: Final = (
base.output_sample_rate_hertz
if output_audio_format is None
else self._map_audio_format_to_sample_rate(str(output_audio_format), is_output=True)
)
requested_output_rate: Final = _integer_setting(session_config.get("output_sample_rate_hertz"))
output_sample_rate_hertz: Final = (
requested_output_rate if requested_output_rate is not None else format_sample_rate
)
requested_temperature: Final = _number_setting(temperature)
requested_max_tokens: Final = _integer_setting(max_output_tokens)
locked_temperature: Final = requested_temperature if requested_temperature is not None else base.temperature
explicit_max_tokens: Final = requested_max_tokens if requested_max_tokens is not None else base.max_tokens
max_tokens: Final = None if max_output_tokens == "inf" else explicit_max_tokens
return dataclasses.replace(
base,
tools_json=tools_json,
system_prompt=system_prompt,
voice_id=voice_id,
output_sample_rate_hertz=output_sample_rate_hertz,
temperature=locked_temperature,
max_tokens=max_tokens,
)
def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> BedrockSessionUpdateRejection | None:
"""The error for a session.update this session cannot apply, or None when it applies."""
session_config: Final = json_message.get("session")
if not isinstance(session_config, Mapping):
return None
malformed: Final = () if self.drop_params else _malformed_number_settings(session_config)
if malformed:
return BedrockSessionUpdateRejection(
code="invalid_type",
message=(
f"Invalid type for 'session.{malformed[0]}': Bedrock Nova Sonic needs a number for temperature "
"and output_sample_rate_hertz, and a number or 'inf' for max_response_output_tokens. "
"Set drop_params to drop it instead"
),
param=f"session.{malformed[0]}",
)
if not self.prompt_started:
return None
if self._apply_locked_settings(session_config, self._started_settings) != self._started_settings:
return BedrockSessionUpdateRejection(
code="session_update_not_applied",
message="Bedrock Nova Sonic cannot change tools, instructions, voice, output audio format, temperature or max output tokens after the session has started, so this session.update was not applied",
param="session",
)
return None
def transform_input_audio_buffer_append_event(self, json_message: dict) -> list[str]:
"""
Transform input_audio_buffer.append event to Bedrock audio input.
@ -655,18 +791,21 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
return []
message_type: Final = json_message.get("type")
# Route to appropriate transformation method
prompt_start_messages: Final = (
self.transform_session_update_event({"session": {}})
if message_type in PROMPT_STARTING_CLIENT_EVENT_TYPES and not self.prompt_started
else []
)
if message_type == "session.update":
return self.transform_session_update_event(json_message)
elif message_type == "input_audio_buffer.append":
return self.transform_input_audio_buffer_append_event(json_message)
return prompt_start_messages + self.transform_input_audio_buffer_append_event(json_message)
elif message_type == "input_audio_buffer.commit":
return self.transform_input_audio_buffer_commit_event(json_message)
return prompt_start_messages + self.transform_input_audio_buffer_commit_event(json_message)
elif message_type == "conversation.item.create":
return self.transform_conversation_item_create_event(json_message)
return prompt_start_messages + self.transform_conversation_item_create_event(json_message)
elif message_type == "response.create":
return self.transform_response_create_event(json_message)
return prompt_start_messages + self.transform_response_create_event(json_message)
elif message_type == "response.cancel":
return self.transform_response_cancel_event(json_message)
else:
@ -718,6 +857,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id: str | None,
current_output_item_id: str | None,
current_conversation_id: str | None,
current_delta_type: ALL_DELTA_TYPES | None,
) -> tuple[
list[OpenAIRealtimeEvents],
str | None,
@ -751,7 +891,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id,
current_output_item_id,
current_conversation_id,
None,
current_delta_type,
)
generation_stage: Final = self._parse_generation_stage(content_start.get("additionalModelFields"))
@ -776,12 +916,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id = f"resp_{uuid.uuid4()}"
if not current_output_item_id:
current_output_item_id = f"item_{uuid.uuid4()}"
self._message_output_index = self._response_item_count
self._response_item_count += 1
if not current_conversation_id:
current_conversation_id = f"conv_{uuid.uuid4()}"
# Determine content type
content_type: Final = content_start.get("type", "TEXT")
current_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio"
next_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio"
returned_messages: Final[list[OpenAIRealtimeEvents]] = []
@ -805,7 +947,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded(
type="response.output_item.added",
response_id=current_response_id,
output_index=0,
output_index=self._message_output_index,
item={
"id": current_output_item_id,
"object": "realtime.item",
@ -815,18 +957,17 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
"content": [],
},
)
self._assistant_block_open = True
returned_messages.append(output_item_added)
# Send response.content_part.added
content_part_added: Final = OpenAIRealtimeResponseContentPartAdded(
type="response.content_part.added",
content_index=0,
output_index=0,
output_index=self._message_output_index,
event_id=f"event_{uuid.uuid4()}",
item_id=current_output_item_id,
part=(
{"type": "text", "text": ""} if current_delta_type == "text" else {"type": "audio", "transcript": ""}
),
part=({"type": "text", "text": ""} if next_delta_type == "text" else {"type": "audio", "transcript": ""}),
response_id=current_response_id,
)
returned_messages.append(content_part_added)
@ -836,7 +977,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id,
current_output_item_id,
current_conversation_id,
current_delta_type,
next_delta_type,
)
@staticmethod
@ -973,7 +1114,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
content_index=0,
event_id=f"event_{uuid.uuid4()}",
item_id=current_output_item_id,
output_index=0,
output_index=self._message_output_index,
response_id=current_response_id,
delta=text_content,
)
@ -1013,7 +1154,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
content_index=0,
event_id=f"event_{uuid.uuid4()}",
item_id=current_output_item_id,
output_index=0,
output_index=self._message_output_index,
response_id=current_response_id,
delta=audio_content,
)
@ -1049,22 +1190,21 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
return [], current_delta_chunks
if not current_output_item_id or not current_response_id:
return [], current_delta_chunks
if not self._assistant_block_open:
return [], current_delta_chunks
self._assistant_block_open = False
returned_messages: Final[list[OpenAIRealtimeEvents]] = []
accumulated_text: Final = "".join(chunk.get("delta", "") for chunk in current_delta_chunks or [])
# Send appropriate done event based on type
if current_delta_type == "text":
# Accumulate text
accumulated_text = ""
if current_delta_chunks:
accumulated_text = "".join([chunk.get("delta", "") for chunk in current_delta_chunks])
text_done: Final = OpenAIRealtimeResponseTextDone(
type="response.text.done",
content_index=0,
event_id=f"event_{uuid.uuid4()}",
item_id=current_output_item_id,
output_index=0,
output_index=self._message_output_index,
response_id=current_response_id,
text=accumulated_text,
)
@ -1076,7 +1216,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
content_index=0,
event_id=f"event_{uuid.uuid4()}",
item_id=current_output_item_id,
output_index=0,
output_index=self._message_output_index,
part={"type": "text", "text": accumulated_text},
response_id=current_response_id,
)
@ -1088,7 +1228,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
content_index=0,
event_id=f"event_{uuid.uuid4()}",
item_id=current_output_item_id,
output_index=0,
output_index=self._message_output_index,
response_id=current_response_id,
)
returned_messages.append(audio_done)
@ -1099,28 +1239,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
content_index=0,
event_id=f"event_{uuid.uuid4()}",
item_id=current_output_item_id,
output_index=0,
output_index=self._message_output_index,
part={"type": "audio", "transcript": ""},
response_id=current_response_id,
)
returned_messages.append(content_part_done)
# Send output_item.done
completed_parts: Final = self._completed_parts(current_delta_type, accumulated_text)
existing_item: Final = self._response_done_items.get(self._message_output_index)
existing_content: Final = existing_item.get("content", []) if existing_item is not None else []
completed_message_item: Final[OpenAIRealtimeStreamResponseOutputItem] = {
"id": current_output_item_id,
"object": "realtime.item",
"type": "message",
"status": "completed",
"role": "assistant",
"content": [*existing_content, *completed_parts],
}
output_item_done: Final = OpenAIRealtimeOutputItemDone(
type="response.output_item.done",
event_id=f"event_{uuid.uuid4()}",
output_index=0,
output_index=self._message_output_index,
response_id=current_response_id,
item={
"id": current_output_item_id,
"object": "realtime.item",
"type": "message",
"status": "completed",
"role": "assistant",
"content": [],
},
item=completed_message_item,
)
returned_messages.append(output_item_done)
self._response_done_items = MappingProxyType(
{**self._response_done_items, self._message_output_index: completed_message_item}
)
# Reset delta chunks
return returned_messages, None
@ -1165,6 +1311,12 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
return [], None, None, None
usage: Final = self._take_usage_delta()
output_items: Final = tuple(item for _, item in sorted(self._response_done_items.items()))
self._response_done_items = MappingProxyType({})
self._response_item_count = 0
self._message_output_index = 0
self._assistant_block_open = False
self._response_done_pending = False
response_done: Final = OpenAIRealtimeDoneEvent(
type="response.done",
event_id=f"event_{uuid.uuid4()}",
@ -1172,7 +1324,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
object="realtime.response",
id=current_response_id,
status="completed",
output=[],
output=list(output_items),
conversation_id=current_conversation_id,
usage=dict(usage),
),
@ -1184,56 +1336,148 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
def transform_tool_use_event(
self,
event: dict,
current_output_item_id: str | None,
current_response_id: str | None,
current_conversation_id: str | None,
) -> tuple[list[OpenAIRealtimeEvents], str, str]:
"""
Transform Bedrock toolUse event to OpenAI format.
Args:
event: Bedrock toolUse event
current_output_item_id: Current output item ID
current_response_id: Current response ID
current_conversation_id: Current conversation ID
Returns:
Tuple of (events, tool_call_id, tool_name) for tracking
Tuple of (events, response_id, conversation_id)
"""
verbose_logger.debug("Handling toolUse")
tool_use: Final = event["toolUse"]
tool_use: Final = BedrockToolUse.model_validate(event["toolUse"])
verbose_logger.debug("Tool use event: %s (ID: %s)", tool_use.toolName, tool_use.toolUseId)
if not current_output_item_id or not current_response_id:
return [], "", ""
response_id: Final = current_response_id or f"resp_{uuid.uuid4()}"
conversation_id: Final = current_conversation_id or f"conv_{uuid.uuid4()}"
item_id: Final = f"item_{uuid.uuid4()}"
output_index: Final = self._response_item_count
self._response_item_count += 1
# Parse the tool input
tool_input = {}
if "input" in tool_use:
try:
tool_input = json.loads(tool_use["input"]) if isinstance(tool_use["input"], str) else tool_use["input"]
except json.JSONDecodeError:
tool_input = {}
response_created_events: Final[tuple[OpenAIRealtimeEvents, ...]] = (
(
OpenAIRealtimeStreamResponseBaseObject(
type="response.created",
event_id=f"event_{uuid.uuid4()}",
response={
"object": "realtime.response",
"id": response_id,
"status": "in_progress",
"output": [],
"conversation_id": conversation_id,
},
),
)
if current_response_id is None
else ()
)
tool_call_id: Final = tool_use.get("toolUseId", "")
tool_name: Final = tool_use.get("toolName", "")
output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded(
type="response.output_item.added",
response_id=response_id,
output_index=output_index,
item={
"id": item_id,
"object": "realtime.item",
"type": "function_call",
"status": "in_progress",
"call_id": tool_use.toolUseId,
"name": tool_use.toolName,
"arguments": "",
},
)
# Create a function call arguments done event
# This is a custom event format that matches what clients expect
function_call_event: Final[dict[str, object]] = {
"type": "response.function_call_arguments.done",
"event_id": f"event_{uuid.uuid4()}",
"response_id": current_response_id,
"item_id": current_output_item_id,
"output_index": 0,
"call_id": tool_call_id,
"name": tool_name,
"arguments": json.dumps(tool_input),
arguments_delta: Final = OpenAIRealtimeFunctionCallArgumentsDelta(
type="response.function_call_arguments.delta",
event_id=f"event_{uuid.uuid4()}",
response_id=response_id,
item_id=item_id,
output_index=output_index,
call_id=tool_use.toolUseId,
delta=tool_use.arguments,
)
function_call_event: Final = OpenAIRealtimeFunctionCallArgumentsDone(
type="response.function_call_arguments.done",
event_id=f"event_{uuid.uuid4()}",
response_id=response_id,
item_id=item_id,
output_index=output_index,
call_id=tool_use.toolUseId,
name=tool_use.toolName,
arguments=tool_use.arguments,
)
completed_item: Final[OpenAIRealtimeStreamResponseOutputItem] = {
"id": item_id,
"object": "realtime.item",
"type": "function_call",
"status": "completed",
"call_id": tool_use.toolUseId,
"name": tool_use.toolName,
"arguments": tool_use.arguments,
}
output_item_done: Final = OpenAIRealtimeOutputItemDone(
type="response.output_item.done",
event_id=f"event_{uuid.uuid4()}",
output_index=output_index,
response_id=response_id,
item=completed_item,
)
self._response_done_items = MappingProxyType({**self._response_done_items, output_index: completed_item})
return (
[cast(OpenAIRealtimeEvents, function_call_event)],
tool_call_id,
tool_name,
[
*response_created_events,
output_item_added,
arguments_delta,
function_call_event,
output_item_done,
],
response_id,
conversation_id,
)
def _tool_content_end_events(
self,
current_output_item_id: str | None,
current_response_id: str | None,
current_conversation_id: str | None,
current_delta_type: ALL_DELTA_TYPES | None,
) -> tuple[
tuple[OpenAIRealtimeEvents, ...],
str | None,
str | None,
ALL_DELTA_TYPES | None,
]:
if self._assistant_block_open:
self._response_done_pending = True
return (), current_output_item_id, current_response_id, current_delta_type
done_events, output_item_id, response_id, delta_type = self._response_done_events(
current_response_id, current_conversation_id
)
return tuple(done_events), output_item_id, response_id, delta_type
@staticmethod
def _completed_parts(
delta_type: str | None, accumulated_text: str
) -> tuple[OpenAIRealtimeStreamResponseOutputItemContent, ...]:
if delta_type == "text":
return ({"type": "text", "text": accumulated_text},)
elif delta_type == "audio":
return ({"type": "audio", "transcript": ""},)
return ()
def _should_close_response(self, stop_reason: str | None) -> bool:
return stop_reason == "END_TURN" or self._response_done_pending
def transform_conversation_item_create_tool_result_event(self, json_message: dict) -> list[str]:
"""
Transform conversation.item.create with tool result to Bedrock format.
@ -1363,6 +1607,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id,
current_output_item_id,
current_conversation_id,
current_delta_type,
)
returned_messages.extend(events)
@ -1382,6 +1627,17 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
events = self.transform_audio_output_event(event, current_output_item_id, current_response_id)
returned_messages.extend(events)
elif "contentEnd" in event and BedrockContentEnd.model_validate(event["contentEnd"]).type == "TOOL":
(
done_events,
current_output_item_id, # rebind-ok: same state-threading pattern as the adjacent event branches
current_response_id, # rebind-ok: same state-threading pattern as the adjacent event branches
current_delta_type, # rebind-ok: same state-threading pattern as the adjacent event branches
) = self._tool_content_end_events(
current_output_item_id, current_response_id, current_conversation_id, current_delta_type
)
returned_messages.extend(done_events)
elif "contentEnd" in event and self._user_transcript_active:
self._user_transcript_active = False
self._user_transcript_generation_stage = None
@ -1396,7 +1652,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_delta_chunks,
)
returned_messages.extend(events)
if BedrockContentEnd.model_validate(event["contentEnd"]).stopReason == "END_TURN":
content_end: Final = BedrockContentEnd.model_validate(event["contentEnd"])
if self._should_close_response(content_end.stopReason):
(
done_events,
current_output_item_id,
@ -1412,12 +1669,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
self.transform_usage_event(BedrockUsageEvent.model_validate(event["usageEvent"]))
elif "toolUse" in event:
events, tool_call_id, tool_name = self.transform_tool_use_event(
event, current_output_item_id, current_response_id
events, current_response_id, current_conversation_id = self.transform_tool_use_event(
event, current_response_id, current_conversation_id
)
returned_messages.extend(events)
# Store tool call info for potential use
verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id)
elif "promptEnd" in event or "completionEnd" in event:
(

47
litellm/llms/ollama/common_utils.py
View file

@ -1,11 +1,56 @@
import base64
import io
from typing import Final
from typing import Final, overload
import httpx
from openai.types.chat import ChatCompletionChunk
from pydantic import BaseModel, TypeAdapter
from litellm import verbose_logger
from litellm._uuid import uuid
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import ModelResponse, ModelResponseStream
_PROVIDER_FIELDS: Final = TypeAdapter(dict[str, object])
class OllamaResponseIdNormalizer:
def __init__(self) -> None:
self.response_id: str | None = None
self.original_response_id: str | None = None
def observe(self, response: object) -> None:
if self.response_id is not None:
return
if not isinstance(response, (ModelResponse, ModelResponseStream, ChatCompletionChunk)):
return
if response.system_fingerprint != "fp_ollama":
return
fields: Final = _PROVIDER_FIELDS.validate_python(
BaseModel.model_dump(response).get("provider_specific_fields") or {}
)
original_id: Final = fields.get("original_response_id")
self.original_response_id = original_id if isinstance(original_id, str) else response.id
self.response_id = response.id if isinstance(original_id, str) else f"chatcmpl-{uuid.uuid4()}"
@overload
def apply(self, response: ModelResponse) -> ModelResponse: ...
@overload
def apply(self, response: ModelResponseStream) -> ModelResponseStream: ...
def apply(self, response: ModelResponse | ModelResponseStream) -> ModelResponse | ModelResponseStream:
if self.response_id is None:
return response
fields: Final = _PROVIDER_FIELDS.validate_python(
BaseModel.model_dump(response).get("provider_specific_fields") or {}
)
return response.model_copy(
update={
"id": self.response_id,
"provider_specific_fields": {**fields, "original_response_id": self.original_response_id},
}
)
class OllamaError(BaseLLMException):

2
litellm/llms/openai/chat/gpt_transformation.py
View file

@ -958,9 +958,11 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator):
)
choices = chunk.get("choices", [])
choices = self._map_reasoning_to_reasoning_content(choices)
fingerprint: Final = _OBJECT_MAPPING.validate_python(chunk).get("system_fingerprint")
kwargs: Final[dict[str, Any]] = {
"id": chunk.get("id"),
"system_fingerprint": fingerprint,
"object": "chat.completion.chunk",
"created": chunk.get("created"),
"model": chunk.get("model"),

61
litellm/model_prices_and_context_window_backup.json
View file

@ -32498,11 +32498,13 @@
"supports_vision": true
},
"github_copilot/claude-opus-5": {
"deprecation_date": "2026-11-16",
"litellm_provider": "github_copilot",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/messages"
@ -32828,11 +32830,13 @@
"supports_vision": true
},
"github_copilot/gpt-5.6-terra": {
"deprecation_date": "2026-11-16",
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november",
"supported_endpoints": [
"/v1/responses"
],
@ -75076,14 +75080,14 @@
"supports_vision": true
},
"vertex_ai/deepseek-ai/deepseek-v4.1-flash-maas": {
"cache_read_input_token_cost": 0.0,
"input_cost_per_token": 0.0,
"cache_read_input_token_cost": 3e-08,
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-deepseek_models",
"max_input_tokens": 1048576,
"max_output_tokens": 1048576,
"max_tokens": 1048576,
"mode": "chat",
"output_cost_per_token": 0.0,
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_regions": [
"global"
@ -75128,5 +75132,56 @@
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"thinking_always_on": true
},
"baseten/meta-models/muse-spark-open": {
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_token": 1.25e-06,
"litellm_provider": "baseten",
"max_input_tokens": 1048576,
"max_output_tokens": 1048576,
"max_tokens": 1048576,
"mode": "chat",
"output_cost_per_token": 4.25e-06,
"source": "https://inference.baseten.co/v1/models",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_tool_choice": true,
"supports_vision": true
},
"baseten/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": {
"cache_read_input_token_cost": 4e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "baseten",
"max_input_tokens": 1048576,
"max_output_tokens": 1048576,
"max_tokens": 1048576,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://inference.baseten.co/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "wandb",
"max_input_tokens": 1049000,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 7.5e-07,
"source": "https://openrouter.ai/api/v1/models/nvidia/nemotron-3.5-super-vl/endpoints",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
}
}

7
litellm/proxy/_types.py
View file

@ -3027,6 +3027,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"(see GitHub issue #12905)."
),
)
agent_access_default_deny: bool = Field(
False,
description=(
"When true, a key with no agent grants on itself or its team can reach no agents instead of all agents. "
"Proxy admins are exempt."
),
)
user_url_validation: bool | None = Field(
None,
description=(

51
litellm/proxy/agent_endpoints/auth/agent_permission_handler.py
View file

@ -9,9 +9,10 @@ import asyncio
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, TypeAlias
from typing import Final, Protocol, TypeAlias
from fastapi import HTTPException
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts
@ -47,6 +48,29 @@ class RestrictedAgentAccess:
AgentAccess: TypeAlias = UnrestrictedAgentAccess | RestrictedAgentAccess
class AgentAccessLookup(Protocol):
def __call__(self, user_api_key_auth: UserAPIKeyAuth | None, /, *, strict: bool) -> Awaitable[AgentAccess]: ...
_DEFAULT_DENY_FLAG: Final = TypeAdapter(bool)
def agent_access_default_deny() -> bool:
from litellm.proxy.proxy_server import general_settings_view
try:
return _DEFAULT_DENY_FLAG.validate_python(general_settings_view().get("agent_access_default_deny", False))
except ValidationError:
return True
def _is_proxy_admin(user_api_key_auth: UserAPIKeyAuth | None) -> bool:
return user_api_key_auth is not None and user_api_key_auth.user_role in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN.value,
)
def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]:
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
@ -80,7 +104,7 @@ class AgentRequestHandler:
- If team has restrictions and key has restrictions: use intersection
- If team has restrictions and key has none: inherit from team
- If team has no restrictions: use key restrictions
- If no restrictions: allow all agents
- If no restrictions: allow all agents unless agent_access_default_deny is enabled for a non-admin key
"""
@staticmethod
@ -89,15 +113,22 @@ class AgentRequestHandler:
resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
*,
strict: bool = False,
default_deny: Callable[[], bool] = agent_access_default_deny,
lookup_key_access: AgentAccessLookup | None = None,
lookup_team_access: AgentAccessLookup | None = None,
) -> AgentAccess:
"""Agents the key may reach: key and team grants, intersected with the agent's access group ceiling
and, for an agent key acting on behalf of an invoking user, with that user's team grants."""
if managed_agent_policy(user_api_key_auth) is not None:
return await _managed_actor_agent_access(user_api_key_auth)
key_team_access: Final = await AgentRequestHandler.resolve_key_team_agent_access(
user_api_key_auth, strict=strict
user_api_key_auth,
strict=strict,
lookup_key_access=lookup_key_access,
lookup_team_access=lookup_team_access,
)
if strict and isinstance(key_team_access, UnrestrictedAgentAccess):
deny_undefined: Final = strict or (not _is_proxy_admin(user_api_key_auth) and default_deny())
if deny_undefined and isinstance(key_team_access, UnrestrictedAgentAccess):
return RestrictedAgentAccess(frozenset())
caller_access: Final = await AgentRequestHandler.agent_caller_access(user_api_key_auth, strict=strict)
own_access: Final = _intersect_agent_access(key_team_access, caller_access)
@ -122,12 +153,18 @@ class AgentRequestHandler:
user_api_key_auth: UserAPIKeyAuth | None,
*,
strict: bool = False,
lookup_key_access: AgentAccessLookup | None = None,
lookup_team_access: AgentAccessLookup | None = None,
) -> AgentAccess:
try:
key_access: Final = await AgentRequestHandler.get_allowed_agents_for_key(user_api_key_auth, strict=strict)
team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(
user_api_key_auth, strict=strict
key_lookup: Final = (
AgentRequestHandler.get_allowed_agents_for_key if lookup_key_access is None else lookup_key_access
)
team_lookup: Final = (
AgentRequestHandler._get_allowed_agents_for_team if lookup_team_access is None else lookup_team_access
)
key_access: Final = await key_lookup(user_api_key_auth, strict=strict)
team_access: Final = await team_lookup(user_api_key_auth, strict=strict)
except Exception as e:
if strict:
raise HTTPException(503, "Agent invocation policy is unavailable") from e

5
litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py
View file

@ -38,6 +38,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardrail"):
import litellm
extras: Final = litellm_params.model_extra or {}
_noma_v2_callback: Final = NomaV2Guardrail(
guardrail_name=guardrail.get("guardrail_name", ""),
api_key=litellm_params.api_key,
@ -49,6 +50,10 @@ def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardra
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
timeout=litellm_params.timeout,
streaming_buffer_until_moderated=extras.get("streaming_buffer_until_moderated"),
streaming_end_of_stream_only=extras.get("streaming_end_of_stream_only"),
streaming_sampling_rate=extras.get("streaming_sampling_rate"),
streaming_buffer_release_on_scan=extras.get("streaming_buffer_release_on_scan"),
)
litellm.logging_callback_manager.add_litellm_callback(_noma_v2_callback)

23
litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py
View file

@ -26,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage
from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.noma import NomaV2StreamingParams
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
if TYPE_CHECKING:
@ -63,6 +64,10 @@ class NomaV2Guardrail(CustomGuardrail):
monitor_mode: bool | None = None,
block_failures: bool | None = None,
gateway_name: str | None = None,
streaming_buffer_until_moderated: bool | None = None,
streaming_end_of_stream_only: bool | None = None,
streaming_sampling_rate: int | None = None,
streaming_buffer_release_on_scan: bool | None = None,
**kwargs: Any,
) -> None:
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
@ -86,6 +91,24 @@ class NomaV2Guardrail(CustomGuardrail):
if self._requires_api_key(api_base=self.api_base) and not self.api_key:
raise ValueError("Noma v2 guardrail requires api_key when using Noma SaaS endpoint")
streaming: Final = NomaV2StreamingParams.model_validate(
{
"streaming_buffer_until_moderated": streaming_buffer_until_moderated,
"streaming_end_of_stream_only": streaming_end_of_stream_only,
"streaming_sampling_rate": streaming_sampling_rate,
"streaming_buffer_release_on_scan": streaming_buffer_release_on_scan,
},
context={"guardrail_name": kwargs.get("guardrail_name")},
)
if streaming.streaming_buffer_until_moderated is not None:
self.streaming_buffer_until_moderated: bool = streaming.streaming_buffer_until_moderated
if streaming.streaming_end_of_stream_only is not None:
self.streaming_end_of_stream_only: bool = streaming.streaming_end_of_stream_only
if streaming.streaming_sampling_rate is not None:
self.streaming_sampling_rate: int = streaming.streaming_sampling_rate
if streaming.streaming_buffer_release_on_scan is not None:
self.streaming_buffer_release_on_scan: bool = streaming.streaming_buffer_release_on_scan
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs

48
litellm/proxy/middleware/gzip_middleware.py
View file

@ -4,11 +4,50 @@ from typing import Final
import anyio.to_thread
from starlette.datastructures import Headers, MutableHeaders
from starlette.routing import get_route_path
from starlette.types import ASGIApp, Message, Receive, Scope, Send
MINIMUM_SIZE_BYTES: Final = 500
OFF_LOOP_SIZE_BYTES: Final = 16 * 1024
COMPRESS_LEVEL: Final = 6
DASHBOARD_GZIP_PATHS: Final = frozenset(
{
"/user/daily/activity/aggregated",
"/user/daily/activity/aggregated/keys",
"/user/daily/activity/aggregated/search",
"/user/daily/activity/aggregated/model_top_keys",
"/user/daily/activity/aggregated/users",
"/user/daily/activity/aggregated/cache_leakage_keys",
"/team/daily/activity/aggregated",
"/team/daily/activity/aggregated/keys",
"/team/daily/activity/aggregated/search",
"/team/daily/activity/aggregated/model_top_keys",
"/organization/daily/activity/aggregated",
"/organization/daily/activity/aggregated/keys",
"/organization/daily/activity/aggregated/search",
"/organization/daily/activity/aggregated/model_top_keys",
"/tag/daily/activity/aggregated",
"/tag/daily/activity/aggregated/keys",
"/tag/daily/activity/aggregated/search",
"/tag/daily/activity/aggregated/model_top_keys",
"/customer/daily/activity/aggregated",
"/customer/daily/activity/aggregated/keys",
"/customer/daily/activity/aggregated/search",
"/customer/daily/activity/aggregated/model_top_keys",
"/end_user/daily/activity/aggregated",
"/end_user/daily/activity/aggregated/keys",
"/end_user/daily/activity/aggregated/search",
"/end_user/daily/activity/aggregated/model_top_keys",
"/agent/daily/activity/aggregated",
"/agent/daily/activity/aggregated/keys",
"/agent/daily/activity/aggregated/search",
"/agent/daily/activity/aggregated/model_top_keys",
"/gateway/daily/activity",
"/gateway/errors/activity",
"/team/spend/by_user",
"/global/spend/provider",
"/auto_router/benchmarks",
}
)
def _coding_weight(part: str) -> tuple[str, float]:
@ -26,8 +65,6 @@ def accepts_gzip(accept_encoding: str) -> bool:
async def _compress(body: bytes) -> bytes:
if len(body) < OFF_LOOP_SIZE_BYTES:
return gzip.compress(body, compresslevel=COMPRESS_LEVEL)
return await anyio.to_thread.run_sync(gzip.compress, body, COMPRESS_LEVEL)
@ -83,11 +120,12 @@ class _BufferedBodyGzipResponder:
class GZipBufferedResponseMiddleware:
def __init__(self, app: ASGIApp) -> None:
def __init__(self, app: ASGIApp, paths: frozenset[str] = DASHBOARD_GZIP_PATHS) -> None:
self.app = app
self.paths: Final = paths
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
if scope["type"] != "http" or scope["method"] != "GET" or get_route_path(scope) not in self.paths:
await self.app(scope, receive, send)
return
gzip_accepted: Final = accepts_gzip(Headers(scope=scope).get("accept-encoding", ""))

15
litellm/proxy/proxy_server.py
View file

@ -90,6 +90,7 @@ from litellm.litellm_core_utils.bug_report import (
)
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
initialize_sentry,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
@ -4962,6 +4963,16 @@ def _is_remote_module_url(value: object) -> bool:
return isinstance(value, str) and (value.startswith("s3://") or value.startswith("gcs://"))
def sentry_is_configured(general_settings: Mapping[str, object]) -> bool:
alerting: Final = general_settings.get("alerting")
return (
"sentry" in litellm.input_callback
or "sentry" in litellm.success_callback
or "sentry" in litellm.failure_callback
or (isinstance(alerting, list) and "sentry" in alerting)
)
def _scrub_guardrail_inner(inner: dict[str, JsonValue]) -> None:
"""Strip remote-URL entries from a guardrail's ``callbacks`` list
and ``guardrail`` (v2 module-path) field. Mutates in place."""
@ -6911,6 +6922,9 @@ class ProxyConfig:
if general_settings is None:
general_settings = {}
if sentry_is_configured(general_settings):
initialize_sentry()
typed_general_settings: Final = _GENERAL_SETTINGS_VIEW.validate_python(general_settings)
ConfigGeneralSettings.model_validate(
MappingProxyType(
@ -18824,6 +18838,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
"cancel_on_disconnect": "Boolean",
"disable_auto_add_proxy_admin_to_teams": "Boolean",
"apply_user_budget_to_team_keys": "Boolean",
"agent_access_default_deny": "Boolean",
"user_api_key_cache_max_size": "Integer",
"transcribe_media_buckets": "List",
}

6
litellm/proxy/tracing_runtime.py
View file

@ -46,10 +46,8 @@ async def manage_tracing(
return
try:
connection: Final = LensConnection.from_env()
except ValueError:
verbose_proxy_logger.warning(
"Agent tracing unavailable: configure LITELLM_LENS_URL and LITELLM_LENS_SERVICE_TOKEN"
)
except ValueError as error:
verbose_proxy_logger.error("Agent tracing store lens failed to start, tracing is disabled: %s", error)
yield None
return
async with client_factory(connection) as client:

19
litellm/proxy/utils.py
View file

@ -3250,14 +3250,14 @@ class ProxyLogging:
alerting_metadata=alerting_metadata,
**extra_kwargs,
)
for client in self.alerting:
if client == "sentry":
if litellm.utils.sentry_sdk_instance is not None:
litellm.utils.sentry_sdk_instance.capture_message(formatted_message)
else:
raise Exception("Missing SENTRY_DSN from environment")
if "sentry" in self.alerting:
litellm.litellm_core_utils.litellm_logging.initialize_sentry()
sentry: Final = litellm.litellm_core_utils.litellm_logging.sentry_sdk_instance
if sentry is None:
raise Exception("Missing SENTRY_DSN from environment")
sentry.capture_message(formatted_message)
async def failure_handler(self, original_exception, duration: float, call_type: str, traceback_str=""):
async def failure_handler(self, original_exception: Exception, duration: float, call_type: str, traceback_str=""):
"""
Log failed db read/writes
@ -3298,8 +3298,9 @@ class ProxyLogging:
call_type=call_type,
)
if litellm.utils.capture_exception:
litellm.utils.capture_exception(error=original_exception)
capture: Final = litellm.litellm_core_utils.litellm_logging.capture_exception
if capture is not None:
capture(error=original_exception)
async def post_call_failure_hook(
self,

3
litellm/realtime_api/main.py
View file

@ -33,6 +33,7 @@ from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CallTypes, LlmProviders
from litellm.utils import ProviderConfigManager
from ..litellm_core_utils.core_helpers import normalize_drop_params
from ..litellm_core_utils.credential_accessor import CredentialAccessor
from ..litellm_core_utils.get_litellm_params import get_litellm_params
from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@ -584,6 +585,7 @@ async def _arealtime(
aws_bedrock_runtime_endpoint: Final = kwargs.get("aws_bedrock_runtime_endpoint")
aws_external_id: Final = kwargs.get("aws_external_id")
aws_session_tags: Final = kwargs.get("aws_session_tags")
drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True
await bedrock_realtime.async_realtime(
model=model,
@ -604,6 +606,7 @@ async def _arealtime(
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
aws_external_id=aws_external_id,
aws_session_tags=aws_session_tags,
drop_params=drop_params,
)
elif _custom_llm_provider == "xai":
api_base = (

26
litellm/responses/main.py
View file

@ -59,8 +59,8 @@ if TYPE_CHECKING:
else:
ResponseText = str # Fallback for ResponseText import
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.get_llm_provider_logic import is_openai_compatible_endpoint
from litellm.llms.openai.data_residency import infer_openai_data_residency
from litellm.secret_managers.main import get_secret_str
from litellm.types.responses.main import *
from litellm.types.responses.streaming_websocket import (
ResponsesWebSocketRequestDefaults,
@ -2464,6 +2464,13 @@ def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) ->
)
def _websocket_custom_llm_provider(custom_llm_provider: str | None, api_base: str | None) -> str | None:
is_openai_on_known_provider_host: Final = (
custom_llm_provider == "openai" and api_base is not None and is_openai_compatible_endpoint(api_base)
)
return None if is_openai_on_known_provider_host else custom_llm_provider
@client
async def _aresponses_websocket(
model: str,
@ -2495,6 +2502,7 @@ async def _aresponses_websocket(
dynamic_api_base,
) = litellm.get_llm_provider(
model=model,
custom_llm_provider=_websocket_custom_llm_provider(litellm_params.custom_llm_provider, api_base),
api_base=api_base,
api_key=api_key,
litellm_params=litellm_params,
@ -2516,22 +2524,18 @@ async def _aresponses_websocket(
)
resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None
responses_api_provider_config: BaseResponsesAPIConfig | None = None
if _custom_llm_provider is not None:
responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config(
responses_api_provider_config: Final = (
None
if _custom_llm_provider is None
else ProviderConfigManager.get_provider_responses_api_config(
model=resolved_model,
provider=litellm.LlmProviders(_custom_llm_provider),
api_base=resolved_api_base,
)
resolved_api_key: Final = (
dynamic_api_key
or litellm_params.api_key
or litellm.api_key
or litellm.openai_key
or get_secret_str("OPENAI_API_KEY")
)
resolved_api_key: Final = api_key or dynamic_api_key
# Extract params that we're passing explicitly to avoid duplicates in **kwargs
_explicit_keys: Final = {
"user_api_key_dict",

2
litellm/rust_bridge/host/AGENTS.md
View file

@ -8,4 +8,4 @@ Python state Rust needs but does not own, behind one module per question. Rust i
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green
- Inputs and outputs use values and frozen snapshots. Retain caller-owned clients and caches as opaque references when identity is required; Rust never discovers them through module globals
- Moving a question into Rust means replacing the function body with a native call, not adding a second resolver. The Python signature stays
- A `*/route_host.py` module is not this: it projects one route's kwargs and maps that route's failures. Shared questions belong here
- `routes/<route>.py` holds the questions one route asks, under the same function names on every route (`response`, `map_failure`, `connection_defaults`), so Rust names only the module. A question two routes ask goes in a top-level module here instead

0
tests/unit/rust_bridge/chat_completions/__init__.py → litellm/rust_bridge/host/routes/__init__.py
View file

6
litellm/rust_bridge/chat_completions/route_host.py → litellm/rust_bridge/host/routes/chat_completions.py
View file

@ -6,7 +6,7 @@ from typing import Final
import litellm
from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS
from litellm.rust_bridge import failures
from litellm.rust_bridge.public_call import optional_str
from litellm.rust_bridge.public_call import known_params, optional_str
from litellm.types.utils import ModelResponse
_TRANSPORT_PARAMETERS: Final = frozenset(
@ -27,6 +27,10 @@ _TRANSPORT_PARAMETERS: Final = frozenset(
PARAMETERS: Final = tuple(name for name in OPENAI_CHAT_COMPLETION_PARAMS if name not in _TRANSPORT_PARAMETERS)
def body_params(request: Mapping[str, object]) -> Mapping[str, object]:
return known_params(PARAMETERS, request)
def connection_defaults(provider: str) -> tuple[str | None, str | None]:
if provider == "anthropic":
return litellm.anthropic_key or litellm.api_key, litellm.api_base

35
litellm/rust_bridge/messages/route_host.py → litellm/rust_bridge/host/routes/messages.py
View file

@ -9,12 +9,39 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout
from litellm.llms.anthropic.pass_through.utils import is_reasoning_auto_summary_enabled
from litellm.rust_bridge import failures
from litellm.rust_bridge.public_call import optional_str
from litellm.rust_bridge.public_call import known_params, optional_str
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
_DROP_PATHS: Final = TypeAdapter(list[object])
_BODY_FIELDS: Final = frozenset(
{
"max_tokens",
"metadata",
"stop_sequences",
"stream",
"system",
"temperature",
"thinking",
"tool_choice",
"tools",
"top_k",
"inference_geo",
"top_p",
"mcp_servers",
"context_management",
"compaction",
"container",
"output_format",
"speed",
"output_config",
"cache_control",
"reasoning_effort",
"safeguards",
}
)
@dataclass(frozen=True, slots=True)
@ -22,6 +49,7 @@ class MessagesSettings:
drop_params: bool
reasoning_auto_summary: bool
additional_drop_params: Sequence[str]
request_timeout: float | None
def response(value: Mapping[str, object]) -> AnthropicMessagesResponse:
@ -31,6 +59,10 @@ def response(value: Mapping[str, object]) -> AnthropicMessagesResponse:
)
def body_params(request: Mapping[str, object]) -> Mapping[str, object]:
return known_params(_BODY_FIELDS, request)
def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]:
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
anthropic_messages_stream_hidden_params,
@ -69,5 +101,6 @@ def settings(kwargs: Mapping[str, object]) -> dict[str, object]:
drop_params=_drop_params(kwargs),
reasoning_auto_summary=is_reasoning_auto_summary_enabled(),
additional_drop_params=_additional_drop_params(kwargs),
request_timeout=get_configured_request_timeout(),
)
)

0
litellm/rust_bridge/ocr/route_host.py → litellm/rust_bridge/host/routes/ocr.py
View file

6
litellm/rust_bridge/responses/route_host.py → litellm/rust_bridge/host/routes/responses.py
View file

@ -6,12 +6,16 @@ from typing import Final
import litellm
from litellm import get_llm_provider
from litellm.rust_bridge import failures
from litellm.rust_bridge.public_call import inference_decline_reason, optional_str
from litellm.rust_bridge.public_call import inference_decline_reason, known_params, optional_str
from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams, ResponsesAPIResponse
PARAMETERS: Final = tuple(ResponsesAPIOptionalRequestParams.__annotations__)
def body_params(request: Mapping[str, object]) -> Mapping[str, object]:
return known_params(PARAMETERS, request)
def connection_defaults(provider: str) -> tuple[str | None, str | None]:
provider_key: Final = litellm.openai_key if provider == "openai" else None
return litellm.api_key or provider_key, litellm.api_base

6
litellm/rust_bridge/public_call.py
View file

@ -3,7 +3,7 @@
from __future__ import annotations
import inspect
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Callable, Collection, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, TypeVar, cast # noqa: TID251 # narrows caller-owned containers without copying them
@ -71,6 +71,10 @@ def optional_sequence(value: object) -> Sequence[object] | None:
return cast("Sequence[object]", value) # cast-ok: the same caller-owned object is handed on unchanged
def known_params(names: Collection[str], request: Mapping[str, object]) -> Mapping[str, object]:
return {name: value for name, value in request.items() if name in names and value is not None}
def inference_decline_reason(parameters: tuple[str, ...], kwargs: Mapping[str, object]) -> str | None:
if litellm.drop_params or litellm.modify_params:
return "native inference does not implement the configured parameter rewrites"

93
litellm/tracing/exporter.py
View file

@ -3,9 +3,11 @@ import json
from collections import deque
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from contextlib import suppress
from dataclasses import dataclass
from enum import Enum
from functools import cache
from io import BytesIO
from typing import Final
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import TypeAdapter, ValidationError
@ -16,6 +18,9 @@ from litellm.integrations.clickhouse.clickhouse_spend_logger import spend_log_ro
from litellm.integrations.custom_logger import CustomLogger
from litellm.tracing.types import SpendLogPayload
if TYPE_CHECKING:
from prometheus_client import Counter as PrometheusCounter
MAX_EVENT_BYTES: Final = 1024 * 1024
MAX_BUFFER_BYTES: Final = 32 * 1024 * 1024
MAX_BUFFER_EVENTS: Final = 1000
@ -89,11 +94,49 @@ def encode_record(value: Mapping[str, object]) -> bytes | ExportFailure:
return ExportFailure.INVALID
@dataclass(frozen=True, slots=True)
class ExportCounters:
written: "PrometheusCounter | None"
dropped: "PrometheusCounter | None"
def add_written(self, count: int) -> None:
if self.written is not None:
self.written.inc(count)
def add_dropped(self, count: int) -> None:
if self.dropped is not None:
self.dropped.inc(count)
@cache
def lens_export_counters() -> ExportCounters:
try:
from prometheus_client import Counter
except ImportError:
return ExportCounters(written=None, dropped=None)
return ExportCounters(
written=Counter(
"litellm_lens_export_rows_written_total",
"Request records the Lens exporter delivered to the Lens service",
),
dropped=Counter(
"litellm_lens_export_rows_dropped_total",
"Request records the Lens exporter dropped before the Lens service accepted them",
),
)
class LensExporter(CustomLogger):
def __init__(self, client: httpx.AsyncClient, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep) -> None:
def __init__(
self,
client: httpx.AsyncClient,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
counters: ExportCounters | None = None,
) -> None:
super().__init__()
self.client: Final = client
self.sleep: Final = sleep
self.counters: Final = counters if counters is not None else lens_export_counters()
self.queue: Final[deque[bytes]] = deque() # mutable-ok: bounded producer-consumer queue
self.wake: Final = asyncio.Event()
self.closed = False
@ -108,6 +151,14 @@ class LensExporter(CustomLogger):
if self.task is None:
self.task = asyncio.create_task(self._run())
def _wrote(self, count: int) -> None:
self.rows_written += count
self.counters.add_written(count)
def _drop(self, count: int) -> None:
self.rows_dropped += count
self.counters.add_dropped(count)
def enqueue(self, record: bytes) -> bool:
if (
self.closed
@ -115,7 +166,7 @@ class LensExporter(CustomLogger):
or self.buffered_events >= MAX_BUFFER_EVENTS
or self.buffered_bytes + len(record) > MAX_BUFFER_BYTES
):
self.rows_dropped += 1
self._drop(1)
return False
self.queue.append(record)
self.buffered_events += 1
@ -138,13 +189,13 @@ class LensExporter(CustomLogger):
if raw is None or self.closed:
return
if self.buffered_events >= MAX_BUFFER_EVENTS or self.buffered_bytes >= MAX_BUFFER_BYTES:
self.rows_dropped += 1
self._drop(1)
return
try:
checked: Final = _check_size(raw, MAX_EVENT_BYTES)
if isinstance(checked, ExportFailure):
self.rows_dropped += 1
self._warn(checked.value)
self._drop(1)
self._log_failure(checked.value)
return
payload: Final = _PAYLOAD.validate_python(raw)
if str(payload.get("call_type", "")).startswith(("/v1/traces", "/v1/logs")):
@ -152,24 +203,26 @@ class LensExporter(CustomLogger):
row: Final = spend_log_row_from_payload(payload, kwargs)
record: Final = encode_record(row)
if isinstance(record, ExportFailure):
self.rows_dropped += 1
self._warn(record.value)
self._drop(1)
self._log_failure(record.value)
return
self.enqueue(record)
except ValidationError as error:
self.rows_dropped += 1
self._drop(1)
fields: Final = tuple(
str(issue["loc"][0]) if issue["loc"] else "$"
for issue in error.errors(include_input=False, include_context=False, include_url=False)[:5]
)
self._warn("invalid request record fields: " + ", ".join(fields))
self._log_failure("invalid request record fields: " + ", ".join(fields))
except (ValueError, TypeError, OverflowError, RecursionError) as error:
self.rows_dropped += 1
self._warn(type(error).__name__)
self._drop(1)
self._log_failure(type(error).__name__)
def _warn(self, reason: str) -> None:
def _log_failure(self, reason: str) -> None:
if reason != self.last_error:
verbose_proxy_logger.warning("Lens request export failed (%s); model requests continue", reason)
verbose_proxy_logger.error(
"Agent tracing store lens dropped request records (%s); model requests continue", reason
)
self.last_error = reason
def _batch(self) -> tuple[bytes, ...]:
@ -196,13 +249,13 @@ class LensExporter(CustomLogger):
self.last_error = ""
return True
if response.status_code not in (429, 502, 503, 504):
self._warn(f"HTTP {response.status_code}")
self._log_failure(f"HTTP {response.status_code}")
return False
except httpx.HTTPError:
pass
if attempt < 2:
await self.sleep(float(1 << attempt))
self._warn("retry limit reached")
self._log_failure("retry limit reached")
return False
async def _run(self) -> None:
@ -217,11 +270,11 @@ class LensExporter(CustomLogger):
batch: Final = self._batch()
try:
if await self._send(batch):
self.rows_written += len(batch)
self._wrote(len(batch))
else:
self.rows_dropped += len(batch)
self._drop(len(batch))
except asyncio.CancelledError:
self.rows_dropped += len(batch)
self._drop(len(batch))
raise
finally:
self.buffered_events -= len(batch)
@ -237,7 +290,7 @@ class LensExporter(CustomLogger):
self.task.cancel()
with suppress(asyncio.CancelledError):
await self.task
self.rows_dropped += len(self.queue)
self._drop(len(self.queue))
self.queue.clear()
self.buffered_bytes = 0
self.buffered_events = 0

11
litellm/types/llms/openai.py
View file

@ -2271,6 +2271,16 @@ class OpenAIRealtimeContentPartDone(TypedDict):
type: Literal["response.content_part.done"]
class OpenAIRealtimeFunctionCallArgumentsDelta(TypedDict):
type: ReadOnly[Literal["response.function_call_arguments.delta"]]
event_id: ReadOnly[str]
response_id: ReadOnly[str]
item_id: ReadOnly[str]
output_index: ReadOnly[int]
call_id: ReadOnly[str]
delta: ReadOnly[str]
class OpenAIRealtimeFunctionCallArgumentsDone(TypedDict):
type: Literal["response.function_call_arguments.done"]
event_id: str
@ -2438,6 +2448,7 @@ OpenAIRealtimeEvents = (
| OpenAIRealtimeResponseAudioDone
| OpenAIRealtimeContentPartDone
| OpenAIRealtimeOutputItemDone
| OpenAIRealtimeFunctionCallArgumentsDelta
| OpenAIRealtimeFunctionCallArgumentsDone
| OpenAIRealtimeDoneEvent
| OpenAIRealtimeInputAudioBufferSpeechEvent

46
litellm/types/proxy/guardrails/guardrail_hooks/noma.py
View file

@ -1,4 +1,16 @@
from pydantic import Field
from typing import Final
from pydantic import (
BaseModel,
ConfigDict,
Field,
ValidationError,
ValidationInfo,
ValidatorFunctionWrapHandler,
field_validator,
)
from litellm._logging import verbose_proxy_logger
from .base import GuardrailConfigModel
@ -55,3 +67,35 @@ class NomaV2GuardrailConfigModel(GuardrailConfigModel):
@staticmethod
def ui_friendly_name() -> str:
return "Noma Security v2"
class NomaV2StreamingParams(BaseModel):
"""Streaming settings a Noma v2 guardrail reads from the top level of its litellm_params.
A key left unset, or set to a value that does not validate, stays None, so the unified
streaming hook keeps its own default for it and the guardrail still loads.
"""
model_config = ConfigDict(frozen=True)
streaming_buffer_until_moderated: bool | None = None
streaming_end_of_stream_only: bool | None = None
streaming_sampling_rate: int | None = Field(default=None, ge=1)
streaming_buffer_release_on_scan: bool | None = None
@field_validator("*", mode="wrap")
@classmethod
def _ignore_invalid_value(
cls, value: object, handler: ValidatorFunctionWrapHandler, info: ValidationInfo
) -> object:
try:
validated: Final = handler(value)
except ValidationError:
verbose_proxy_logger.warning(
"Noma v2 guardrail %s: ignoring %s=%.100r, which is not a valid value; the default applies",
info.context.get("guardrail_name") if info.context else None,
info.field_name,
value,
)
return None
return validated

4
litellm/types/router.py
View file

@ -1084,6 +1084,8 @@ class RouterRateLimitErrorBasic(ValueError):
Raise a basic error inside helper functions.
"""
status_code: Final = 429
def __init__(
self,
model: str,
@ -1099,6 +1101,8 @@ class RouterErrorTypes(str, enum.Enum):
class RouterRateLimitError(ValueError):
status_code: Final = 429
def __init__(
self,
model: str,

61
model_prices_and_context_window.json
View file

@ -32498,11 +32498,13 @@
"supports_vision": true
},
"github_copilot/claude-opus-5": {
"deprecation_date": "2026-11-16",
"litellm_provider": "github_copilot",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/messages"
@ -32828,11 +32830,13 @@
"supports_vision": true
},
"github_copilot/gpt-5.6-terra": {
"deprecation_date": "2026-11-16",
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november",
"supported_endpoints": [
"/v1/responses"
],
@ -75076,14 +75080,14 @@
"supports_vision": true
},
"vertex_ai/deepseek-ai/deepseek-v4.1-flash-maas": {
"cache_read_input_token_cost": 0.0,
"input_cost_per_token": 0.0,
"cache_read_input_token_cost": 3e-08,
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-deepseek_models",
"max_input_tokens": 1048576,
"max_output_tokens": 1048576,
"max_tokens": 1048576,
"mode": "chat",
"output_cost_per_token": 0.0,
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_regions": [
"global"
@ -75128,5 +75132,56 @@
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"thinking_always_on": true
},
"baseten/meta-models/muse-spark-open": {
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_token": 1.25e-06,
"litellm_provider": "baseten",
"max_input_tokens": 1048576,
"max_output_tokens": 1048576,
"max_tokens": 1048576,
"mode": "chat",
"output_cost_per_token": 4.25e-06,
"source": "https://inference.baseten.co/v1/models",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_tool_choice": true,
"supports_vision": true
},
"baseten/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": {
"cache_read_input_token_cost": 4e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "baseten",
"max_input_tokens": 1048576,
"max_output_tokens": 1048576,
"max_tokens": 1048576,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://inference.baseten.co/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "wandb",
"max_input_tokens": 1049000,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 7.5e-07,
"source": "https://openrouter.ai/api/v1/models/nvidia/nemotron-3.5-super-vl/endpoints",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
}
}

56
scripts/check_test_quality.py
View file

@ -88,6 +88,13 @@ TQ011 An integration-test assertion that compares an elapsed time against a nu
is a literal when it is built only from numbers, or names a module-level number.
Only `assert` statements are examined, so a polling deadline in a `while` loop is
left alone. Same path scope as TQ010.
TQ012 A port read back from a port-0 bind: `sock.bind((host, 0))` in a function
that also calls `getsockname`. The number is free only until the socket
closes, so a server started on it later can lose it to another process and
fail with "Address already in use". Take the port from `allocate_port()` in
tests/integration/_support/ports.py, which hands each pytest worker its own
block of ports. Every file under tests/integration is examined, including
tests/integration/_support.
Every rule is suppressible with `# test-quality-ok: <reason>` on the reported
line, following the repo's `*-ok: <reason>` convention. A suppression without a
@ -172,7 +179,7 @@ SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_cal
INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P"))
RULE_CODES: Final = frozenset((
"TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", "TQ010", "TQ011",
"TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", "TQ010", "TQ011", "TQ012",
))
INTEGRATION_ROOT: Final = ("tests", "integration")
@ -181,6 +188,8 @@ SIGNAL_SENDERS: Final = frozenset(("os.kill", "os.killpg"))
PROCESS_CONTROL_SIGNALS: Final = frozenset(("SIGKILL", "SIGSTOP", "SIGCONT"))
ZERO_ARGUMENT_PROCESS_CONTROLS: Final = frozenset(("kill", "terminate", "suspend", "resume"))
SIGNAL_METHOD: Final = "send_signal"
BIND_METHOD: Final = "bind"
PORT_READBACK: Final = "getsockname"
CLOCK_READERS: Final = frozenset((
"time.monotonic", "time.perf_counter", "time.time",
"time.monotonic_ns", "time.perf_counter_ns", "time.time_ns",
@ -877,6 +886,50 @@ def iter_wall_clock_bound_violations(path: Path, tree: ast.Module) -> Iterator[V
)
def is_integration_tree_path(path: Path) -> bool:
parts: Final = path.parts
return any(parts[index:index + 2] == INTEGRATION_ROOT for index in range(len(parts) - 1))
def _calls_in_scope(node: ast.AST, scope: ast.AST) -> Iterator[tuple[ast.Call, ast.AST]]:
for child in ast.iter_child_nodes(node):
if isinstance(child, ast.Call):
yield child, scope
inner: Final = child if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) else scope
yield from _calls_in_scope(child, inner)
def _binds_port_zero(call: ast.Call) -> bool:
if not isinstance(call.func, ast.Attribute) or call.func.attr != BIND_METHOD or not call.args:
return False
address: Final = call.args[0]
return (
isinstance(address, ast.Tuple)
and len(address.elts) == 2
and isinstance(address.elts[1], ast.Constant)
and address.elts[1].value == 0
)
def _reads_back_port(scope: ast.AST) -> bool:
return any(isinstance(node, ast.Attribute) and node.attr == PORT_READBACK for node in ast.walk(scope))
def iter_port_readback_violations(path: Path, tree: ast.Module) -> Iterator[Violation]:
if not is_integration_tree_path(path):
return
for call, scope in _calls_in_scope(tree, tree):
if _binds_port_zero(call) and _reads_back_port(scope):
yield Violation(
path,
call.lineno,
"TQ012",
"a port read back from a port-0 bind is free only until the socket closes, so a server started on "
"it later can lose it to another process; take the port from allocate_port() in "
f"tests/integration/_support/ports.py (suppress: `# {SUPPRESSION_TOKEN}: <reason>`)",
)
def iter_conftest_inventory_violations(path: Path, tree: ast.Module) -> Iterator[Violation]:
if path.name != CONFTEST_NAME:
return
@ -916,6 +969,7 @@ def check_file(path: Path) -> tuple[Violation, ...]:
*iter_child_interpreter_violations(path, tree),
*iter_process_signal_violations(path, tree),
*iter_wall_clock_bound_violations(path, tree),
*iter_port_readback_violations(path, tree),
)
if violation.line not in skip
)

2
terraform/provider/CHANGELOG.md
View file

@ -63,7 +63,9 @@ longer signal it.
- **key**: Updates no longer send an empty `budget_duration`, which the proxy rejects with a 400; any update to a key without a configured `budget_duration` previously failed outright
- **key**: A config-supplied `key` value (write-only) is now forwarded to `/key/generate`; previously it was silently dropped and the proxy generated a random key instead
- **security**: The `litellm_key` data source and `litellm_key_block` resource normalize raw `sk-` keys to their SHA-256 token hash before building request URLs and resource IDs, so plaintext keys no longer land in reverse-proxy access logs, Terraform plan output, or state IDs
- **mcp_server, model, key, organization_member**: Requests now accept any 2xx response instead of requiring exactly HTTP 200; the proxy legitimately returns 201 from `POST /v1/mcp/server` (and other create endpoints), so `litellm_mcp_server` create previously succeeded on the proxy but failed in the provider, leaving the resource out of state and risking a duplicate server on retry
- **unified_access_group**: create now accepts any 2xx response instead of requiring exactly HTTP 200; `POST /v1/unified_access_group` legitimately returns 201, so creation previously succeeded on the proxy but failed in the provider, leaving the group out of state and forcing a `terraform import` to recover on the next apply's 409. Same fix and shape as the one already applied to `mcp_server`/`model`/`key`/`organization_member`; `handleResponse` (shared by several other resources) was the one status-check helper that fix didn't reach. The legacy `litellm_access_group` resource calls the unrelated `/access_group/new` endpoint, which already returns 200, so it was never affected
- **mcp_server**: A plan right after a clean apply of a minimal litellm_mcp_server no longer shows changes. alias is now computed so the proxy's default (the server name) is kept, spec_version keeps the configured value when the proxy returns none, and an mcp_info block the proxy fills with fields the provider does not model is no longer written to state
### Changed

2
terraform/provider/litellm/client.go
View file

@ -455,7 +455,7 @@ func (c *Client) sendRequest(method, path string, body interface{}) (map[string]
log.Printf("Response status: %d", resp.StatusCode)
log.Printf("Response body: %s", c.redactSensitiveData(string(bodyBytes)))
if resp.StatusCode != http.StatusOK {
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, &apiError{StatusCode: resp.StatusCode, Body: string(bodyBytes)}
}

57
terraform/provider/litellm/client_test.go
View file

@ -1,10 +1,67 @@
package litellm
import (
"errors"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestSendRequestAcceptsFullSuccessRange(t *testing.T) {
tests := []struct {
name string
statusCode int
wantErr bool
wantValue string
}{
{name: "200 OK", statusCode: http.StatusOK, wantErr: false, wantValue: "ok"},
{name: "201 Created", statusCode: http.StatusCreated, wantErr: false, wantValue: "created"},
{name: "202 Accepted", statusCode: http.StatusAccepted, wantErr: false, wantValue: "accepted"},
{name: "400 Bad Request", statusCode: http.StatusBadRequest, wantErr: true},
{name: "404 Not Found", statusCode: http.StatusNotFound, wantErr: true},
{name: "500 Internal Server Error", statusCode: http.StatusInternalServerError, wantErr: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(tt.statusCode)
w.Write([]byte(`{"value":"` + tt.wantValue + `"}`))
}))
defer srv.Close()
client := NewClient(srv.URL, "test-key", true)
result, err := client.sendRequest("POST", "/whatever", map[string]string{"foo": "bar"})
if tt.wantErr {
if err == nil {
t.Fatalf("sendRequest returned no error for status %d", tt.statusCode)
}
var apiErr *apiError
if !errors.As(err, &apiErr) {
t.Fatalf("error is not *apiError: %v", err)
}
if apiErr.StatusCode != tt.statusCode {
t.Errorf("apiErr.StatusCode = %d, want %d", apiErr.StatusCode, tt.statusCode)
}
if tt.statusCode == http.StatusNotFound && !isNotFound(err) {
t.Errorf("isNotFound(err) = false for 404, want true")
}
return
}
if err != nil {
t.Fatalf("sendRequest returned unexpected error: %v", err)
}
if result["value"] != tt.wantValue {
t.Errorf("result[value] = %v, want %q", result["value"], tt.wantValue)
}
})
}
}
func TestRedactSensitiveDataNestedCredentialValues(t *testing.T) {
c := NewClient("http://localhost:4000", "sk-test", false)

1
terraform/provider/litellm/resource_mcp_server.go
View file

@ -24,6 +24,7 @@ func resourceLiteLLMMCPServer() *schema.Resource {
"alias": {
Type: schema.TypeString,
Optional: true,
Computed: true,
Description: "Alias for the MCP server",
},
"description": {

19
terraform/provider/litellm/resource_mcp_server_crud.go
View file

@ -130,7 +130,9 @@ func updateSchemaFromResponse(d *schema.ResourceData, resp *MCPServerResponse) e
d.Set("description", resp.Description)
d.Set("url", resp.URL)
d.Set("transport", resp.Transport)
d.Set("spec_version", resp.SpecVersion)
if resp.SpecVersion != "" {
d.Set("spec_version", resp.SpecVersion)
}
d.Set("auth_type", resp.AuthType)
d.Set("created_at", resp.CreatedAt)
d.Set("created_by", resp.CreatedBy)
@ -151,8 +153,9 @@ func updateSchemaFromResponse(d *schema.ResourceData, resp *MCPServerResponse) e
d.Set("args", resp.Args)
}
// Set mcp_info
if resp.MCPInfo != nil {
if !mcpInfoHasModeledFields(resp.MCPInfo) {
d.Set("mcp_info", []interface{}{})
} else {
mcpInfoList := make([]map[string]interface{}, 1)
mcpInfoMap := make(map[string]interface{})
@ -180,6 +183,14 @@ func updateSchemaFromResponse(d *schema.ResourceData, resp *MCPServerResponse) e
return nil
}
func mcpInfoHasModeledFields(info *MCPInfo) bool {
return info != nil &&
(info.ServerName != "" ||
info.Description != "" ||
info.LogoURL != "" ||
info.MCPServerCostInfo != nil)
}
func resourceLiteLLMMCPServerCreate(d *schema.ResourceData, m interface{}) error {
client, ok := m.(*Client)
if !ok {
@ -287,7 +298,7 @@ func resourceLiteLLMMCPServerDelete(d *schema.ResourceData, m interface{}) error
defer resp.Body.Close()
// For delete operations, we expect a simple string response
if resp.StatusCode != 200 {
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return fmt.Errorf("failed to delete MCP server: unexpected status code %d", resp.StatusCode)
}

211
terraform/provider/litellm/resource_mcp_server_crud_test.go
View file

@ -1,6 +1,7 @@
package litellm
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
@ -12,6 +13,52 @@ import (
"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
)
func TestMCPServerDeleteAcceptsFullSuccessRange(t *testing.T) {
tests := []struct {
name string
statusCode int
wantErr bool
}{
{name: "200 OK", statusCode: http.StatusOK, wantErr: false},
{name: "202 Accepted", statusCode: http.StatusAccepted, wantErr: false},
{name: "404 Not Found", statusCode: http.StatusNotFound, wantErr: true},
{name: "500 Internal Server Error", statusCode: http.StatusInternalServerError, wantErr: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(tt.statusCode)
}))
defer srv.Close()
d := schema.TestResourceDataRaw(t, resourceLiteLLMMCPServer().Schema, map[string]interface{}{
"server_name": "gh",
"transport": "stdio",
"command": "npx",
})
d.SetId("srv-1")
client := NewClient(srv.URL, "test-key", true)
err := resourceLiteLLMMCPServerDelete(d, client)
if tt.wantErr {
if err == nil {
t.Fatalf("resourceLiteLLMMCPServerDelete returned no error for status %d", tt.statusCode)
}
return
}
if err != nil {
t.Fatalf("resourceLiteLLMMCPServerDelete returned unexpected error: %v", err)
}
if d.Id() != "" {
t.Errorf("resource ID not cleared after successful delete, got %q", d.Id())
}
})
}
}
func TestMCPServerReadDoesNotPersistServerEnv(t *testing.T) {
d := schema.TestResourceDataRaw(t, resourceLiteLLMMCPServer().Schema, map[string]interface{}{
"server_name": "gh",
@ -184,3 +231,167 @@ func TestMCPServerSchemaWriteOnlyCredentialRules(t *testing.T) {
})
}
}
func TestMCPServerReadProducesStablePlan(t *testing.T) {
const responseWithoutModeledMCPInfo = `{"alias":"lit9695_proof","approval_status":"active","args":[],"auth_type":"none","command":null,"created_at":"2026-10-15T22:48:21.530000Z","created_by":"default_user_id","credentials":null,"description":null,"env":{},"mcp_access_groups":[],"mcp_info":{"is_public":false,"is_public_explicit":false},"server_id":"dc577aa3-22ac-4fdf-a56b-f6c40e35c389","server_name":"lit9695_proof","status":"unknown","transport":"http","updated_at":"2026-10-15T22:48:21.530000Z","updated_by":"default_user_id","url":"https://mcp.example.com/mcp"}`
responseWithLogoURL := strings.Replace(
responseWithoutModeledMCPInfo,
`"mcp_info":{"is_public":false,"is_public_explicit":false}`,
`"mcp_info":{"is_public":false,"is_public_explicit":false,"logo_url":"https://x/logo.png"}`,
1,
)
if responseWithLogoURL == responseWithoutModeledMCPInfo {
t.Fatal("failed to add logo_url to MCP info response")
}
cases := []struct {
name string
config map[string]interface{}
response string
wantEmptyDiff bool
wantChangedKey string
wantAlias string
wantSpecVersion string
wantMCPInfoLogoURL string
}{
{
name: "minimal config",
config: map[string]interface{}{
"server_name": "lit9695_proof",
"url": "https://mcp.example.com/mcp",
"transport": "http",
},
response: responseWithoutModeledMCPInfo,
wantEmptyDiff: true,
wantAlias: "lit9695_proof",
wantSpecVersion: "2024-11-05",
},
{
name: "configured spec version",
config: map[string]interface{}{
"server_name": "lit9695_proof",
"url": "https://mcp.example.com/mcp",
"transport": "http",
"spec_version": "2025-03-26",
},
response: responseWithoutModeledMCPInfo,
wantEmptyDiff: true,
wantSpecVersion: "2025-03-26",
},
{
name: "configured alias drift remains visible",
config: map[string]interface{}{
"server_name": "lit9695_proof",
"alias": "custom",
"url": "https://mcp.example.com/mcp",
"transport": "http",
},
response: responseWithoutModeledMCPInfo,
wantChangedKey: "alias",
},
{
name: "configured MCP info drift remains visible",
config: map[string]interface{}{
"server_name": "lit9695_proof",
"url": "https://mcp.example.com/mcp",
"transport": "http",
"mcp_info": []interface{}{
map[string]interface{}{"logo_url": "https://x/logo.png"},
},
},
response: responseWithoutModeledMCPInfo,
wantChangedKey: "mcp_info",
},
{
name: "modeled MCP info is retained",
config: map[string]interface{}{
"server_name": "lit9695_proof",
"url": "https://mcp.example.com/mcp",
"transport": "http",
"mcp_info": []interface{}{
map[string]interface{}{"logo_url": "https://x/logo.png"},
},
},
response: responseWithLogoURL,
wantEmptyDiff: true,
wantMCPInfoLogoURL: "https://x/logo.png",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost && r.Method != http.MethodGet {
t.Errorf("unexpected request method: %s", r.Method)
w.WriteHeader(http.StatusNotFound)
return
}
if r.URL.Path != "/v1/mcp/server" && r.URL.Path != "/v1/mcp/server/dc577aa3-22ac-4fdf-a56b-f6c40e35c389" {
t.Errorf("unexpected request path: %s", r.URL.Path)
w.WriteHeader(http.StatusNotFound)
return
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(tc.response))
}))
defer server.Close()
resource := resourceLiteLLMMCPServer()
d := schema.TestResourceDataRaw(t, resource.Schema, tc.config)
client := NewClient(server.URL, "test-key", true)
if err := resourceLiteLLMMCPServerCreate(d, client); err != nil {
t.Fatalf("create failed: %v", err)
}
if err := resourceLiteLLMMCPServerRead(d, client); err != nil {
t.Fatalf("read failed: %v", err)
}
diff, err := resource.Diff(context.Background(), d.State(), terraform.NewResourceConfigRaw(tc.config), nil)
if err != nil {
t.Fatalf("plan diff failed: %v", err)
}
if tc.wantEmptyDiff {
if diff != nil && !diff.Empty() {
t.Fatalf("expected empty plan diff, got %v", diff.Attributes)
}
if diff != nil {
for key := range diff.Attributes {
if key == "alias" || key == "spec_version" || strings.HasPrefix(key, "mcp_info.") {
t.Fatalf("unexpected diff on %s", key)
}
}
}
} else {
if diff == nil || diff.Empty() {
t.Fatal("expected a plan diff")
}
hasExpectedDiff := false
for key := range diff.Attributes {
hasExpectedDiff = hasExpectedDiff || key == tc.wantChangedKey ||
(tc.wantChangedKey == "mcp_info" && strings.HasPrefix(key, "mcp_info."))
}
if !hasExpectedDiff {
t.Fatalf("expected a plan diff on %s, got %v", tc.wantChangedKey, diff.Attributes)
}
}
if tc.wantAlias != "" && d.Get("alias") != tc.wantAlias {
t.Fatalf("alias = %v, want %s", d.Get("alias"), tc.wantAlias)
}
if tc.wantSpecVersion != "" && d.Get("spec_version") != tc.wantSpecVersion {
t.Fatalf("spec_version = %v, want %s", d.Get("spec_version"), tc.wantSpecVersion)
}
if tc.wantMCPInfoLogoURL != "" {
mcpInfo := d.Get("mcp_info").([]interface{})
if len(mcpInfo) != 1 {
t.Fatalf("mcp_info has %d blocks, want 1", len(mcpInfo))
}
gotLogoURL := mcpInfo[0].(map[string]interface{})["logo_url"]
if gotLogoURL != tc.wantMCPInfoLogoURL {
t.Fatalf("mcp_info logo_url = %v, want %s", gotLogoURL, tc.wantMCPInfoLogoURL)
}
}
})
}
}

3
terraform/provider/litellm/resource_model_crud.go
View file

@ -337,6 +337,9 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error {
}
return fmt.Errorf("failed to read model: %w", err)
}
if modelResp == nil {
return fmt.Errorf("failed to read model: empty response body")
}
// Update the state with values from the response or fall back to the data passed in during creation
d.Set("model_name", GetStringValue(modelResp.ModelName, d.Get("model_name").(string)))

36
terraform/provider/litellm/resource_model_test.go
View file

@ -153,6 +153,42 @@ func TestResourceLiteLLMModelReadDisplayName(t *testing.T) {
}
}
func TestResourceLiteLLMModelReadEmptyBodyKeepsState(t *testing.T) {
cases := []struct {
name string
serverBody string
}{
{name: "empty body"},
{name: "null body", serverBody: "null"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet || r.URL.Path != "/model/info" {
t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path)
}
_, _ = w.Write([]byte(tc.serverBody))
}))
defer srv.Close()
d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{
"model_name": "sonnet-4-5-anthropic",
"custom_llm_provider": "anthropic",
"base_model": "claude-sonnet-4-5",
"display_name": "Saved Name",
})
d.SetId("model-123")
if err := resourceLiteLLMModelRead(d, NewClient(srv.URL, "test-key", true)); err == nil {
t.Fatal("expected read error for empty response body")
}
if got := d.Get("display_name").(string); got != "Saved Name" {
t.Errorf("expected display_name to remain %q, got %q", "Saved Name", got)
}
})
}
}
func updateResourceData(t *testing.T, oldDisplayName, newDisplayName string) *schema.ResourceData {
t.Helper()
res := resourceLiteLLMModel()

8
terraform/provider/litellm/utils.go
View file

@ -46,7 +46,7 @@ func handleAPIResponse(resp *http.Response, reqBody interface{}, client *Client)
return nil, fmt.Errorf("failed to read response body: %v", err)
}
if resp.StatusCode != http.StatusOK {
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
var errResp ErrorResponse
if err := json.Unmarshal(bodyBytes, &errResp); err == nil {
if isModelNotFoundError(errResp) {
@ -65,6 +65,10 @@ func handleAPIResponse(resp *http.Response, reqBody interface{}, client *Client)
bodyBytes = envelope.Data[0]
}
if len(bodyBytes) == 0 || string(bodyBytes) == "null" {
return nil, nil
}
var modelResp ModelResponse
if err := json.Unmarshal(bodyBytes, &modelResp); err != nil {
return nil, fmt.Errorf("failed to parse response: %v", err)
@ -143,7 +147,7 @@ func handleMCPAPIResponse(resp *http.Response, result interface{}, client *Clien
return fmt.Errorf("failed to read response body: %v", err)
}
if resp.StatusCode != http.StatusOK {
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
var errResp ErrorResponse
if err := json.Unmarshal(bodyBytes, &errResp); err == nil {
if isMCPServerNotFoundError(errResp) {

128
terraform/provider/litellm/utils_test.go Normal file
View file

@ -0,0 +1,128 @@
package litellm
import (
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
)
func TestHandleAPIResponseAcceptsFullSuccessRange(t *testing.T) {
tests := []struct {
name string
statusCode int
wantErr bool
}{
{name: "200 OK", statusCode: http.StatusOK, wantErr: false},
{name: "201 Created", statusCode: http.StatusCreated, wantErr: false},
{name: "202 Accepted", statusCode: http.StatusAccepted, wantErr: false},
{name: "400 Bad Request", statusCode: http.StatusBadRequest, wantErr: true},
{name: "404 Not Found", statusCode: http.StatusNotFound, wantErr: true},
{name: "500 Internal Server Error", statusCode: http.StatusInternalServerError, wantErr: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
rec := httptest.NewRecorder()
rec.WriteHeader(tt.statusCode)
rec.WriteString(`{"model_name":"gpt-4o"}`)
resp := rec.Result()
client := NewClient("http://localhost:4000", "test-key", true)
got, err := handleAPIResponse(resp, map[string]interface{}{"model_name": "gpt-4o"}, client)
if tt.wantErr {
if err == nil {
t.Fatalf("handleAPIResponse returned no error for status %d", tt.statusCode)
}
if !strings.Contains(err.Error(), strconv.Itoa(tt.statusCode)) {
t.Errorf("error %q does not mention status %d", err.Error(), tt.statusCode)
}
return
}
if err != nil {
t.Fatalf("handleAPIResponse returned unexpected error: %v", err)
}
if got.ModelName != "gpt-4o" {
t.Errorf("got.ModelName = %q, want gpt-4o", got.ModelName)
}
})
}
}
func TestHandleAPIResponseAcceptsEmptyBodyOn2xx(t *testing.T) {
rec := httptest.NewRecorder()
rec.WriteHeader(http.StatusNoContent)
httpResp := rec.Result()
client := NewClient("http://localhost:4000", "test-key", true)
modelResp, err := handleAPIResponse(httpResp, map[string]interface{}{"model_name": "gpt-4o"}, client)
if err != nil {
t.Fatalf("handleAPIResponse returned unexpected error for empty-body 204: %v", err)
}
if modelResp != nil {
t.Fatalf("handleAPIResponse returned non-nil ModelResponse for empty-body 204: %#v", modelResp)
}
}
func TestHandleMCPAPIResponseAcceptsFullSuccessRange(t *testing.T) {
tests := []struct {
name string
statusCode int
wantErr bool
}{
{name: "200 OK", statusCode: http.StatusOK, wantErr: false},
{name: "201 Created", statusCode: http.StatusCreated, wantErr: false},
{name: "202 Accepted", statusCode: http.StatusAccepted, wantErr: false},
{name: "400 Bad Request", statusCode: http.StatusBadRequest, wantErr: true},
{name: "404 Not Found", statusCode: http.StatusNotFound, wantErr: true},
{name: "500 Internal Server Error", statusCode: http.StatusInternalServerError, wantErr: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
rec := httptest.NewRecorder()
rec.WriteHeader(tt.statusCode)
rec.WriteString(`{"server_id":"srv-1","server_name":"gh"}`)
resp := rec.Result()
client := NewClient("http://localhost:4000", "test-key", true)
var mcpResp MCPServerResponse
err := handleMCPAPIResponse(resp, &mcpResp, client)
if tt.wantErr {
if err == nil {
t.Fatalf("handleMCPAPIResponse returned no error for status %d", tt.statusCode)
}
if !strings.Contains(err.Error(), strconv.Itoa(tt.statusCode)) {
t.Errorf("error %q does not mention status %d", err.Error(), tt.statusCode)
}
return
}
if err != nil {
t.Fatalf("handleMCPAPIResponse returned unexpected error: %v", err)
}
if mcpResp.ServerID != "srv-1" {
t.Errorf("mcpResp.ServerID = %q, want srv-1", mcpResp.ServerID)
}
})
}
}
func TestHandleMCPAPIResponseRejectsEmptyBodyOn2xx(t *testing.T) {
rec := httptest.NewRecorder()
rec.WriteHeader(http.StatusNoContent)
resp := rec.Result()
client := NewClient("http://localhost:4000", "test-key", true)
var mcpResp MCPServerResponse
err := handleMCPAPIResponse(resp, &mcpResp, client)
if err == nil {
t.Fatal("handleMCPAPIResponse returned no error for empty-body 204; every MCP caller (create/read/update) writes the parsed result straight into Terraform state with no fallback, so a silently-accepted empty body would blank out or empty-ID the resource")
}
}

1
tests/e2e/coverage_registry/llm_nonconversational.yaml
View file

@ -62,6 +62,7 @@
- {id: llm.files.vertex.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Vertex unified file streams the GCS object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"}
- {id: llm.files.bedrock.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Bedrock unified file streams the S3 object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"}
- {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"}
- {id: llm.realtime.bedrock_converse.tool_use.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: tool_use, streaming: stream, assertions: [works], source: "test_realtime_e2e.py", rationale: "Nova Sonic realtime forwards each tool call with its arguments and completes the tool-result round trip (LIT-9465)"}
- {id: llm.google_native.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "LIT-4076 / proxy/google_endpoints/endpoints.py", fail_before_fix: proven, rationale: "google-native generateContent must stamp x-litellm-response-cost so SDK traffic reconciles against spend"}
- {id: llm.google_native.gemini.basic.stream.works, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: stream, assertions: [works], source: "PR #28213 / proxy/proxy_server.py async_data_generator", fail_before_fix: proven, rationale: "streamGenerateContent must relay single-prefixed SSE frames with no [DONE] sentinel; doubled data: prefixes and the OpenAI terminator both break the Vertex Java SDK"}
- {id: llm.realtime.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: realtime, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "vendor strategy §9.19 / LIT-4778", rationale: "HTTP /v1/realtime/client_secrets returns an ephemeral credential"}

11
tests/e2e/llm_translation/realtime/realtime_client.py
View file

@ -85,7 +85,16 @@ PROVIDERS = (
vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
),
),
# RealtimeProvider("bedrock", "bedrock-realtime", ...) # TODO: Enable when Bedrock is passing
RealtimeProvider(
"bedrock",
"bedrock-realtime",
LiteLLMParamsBody(
model="bedrock/amazon.nova-2-sonic-v1:0",
aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
aws_region_name="os.environ/AWS_REGION",
),
),
# RealtimeProvider(
# "xai",
# "xai-realtime",

84
tests/e2e/llm_translation/realtime/test_realtime_e2e.py
View file

@ -29,13 +29,14 @@ from realtime_client import (
JsonSchemaProperty,
RealtimeClient,
RealtimeProvider,
RealtimeSession,
ReceivedEvent,
ResponseCreate,
ResponseDone,
ServerEnvelope,
SessionConfig,
SessionUpdate,
as_text,
function_call_item,
parse_last,
realtime_model,
transcript,
@ -64,19 +65,26 @@ TEXT_PARAMS: Final = tuple(
for p in PROVIDERS
)
BEDROCK_TOOL_COVERS: Final = pytest.mark.covers(
"llm.realtime.bedrock_converse.tool_use.stream.works", exercised_on=["realtime"]
)
TOOL_PARAMS: Final = tuple(
pytest.param(
p,
id=p.id,
marks=meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.REALTIME,
providers=(Provider(p.id),),
models=(p.litellm_params.model,),
capabilities=(Capability.FUNCTION_CALLING,),
mode=Mode.WEBSOCKET,
)
marks=(
meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.REALTIME,
providers=(Provider(p.id),),
models=(p.litellm_params.model,),
capabilities=(Capability.FUNCTION_CALLING,),
mode=Mode.WEBSOCKET,
)
),
*((BEDROCK_TOOL_COVERS,) if p.id == "bedrock" else ()),
),
)
for p in PROVIDERS
@ -164,31 +172,51 @@ def test_tool_call_round_trip(
session.send(user_message("What's the weather in Paris right now?"))
session.send(ResponseCreate())
first = session.collect_until("response.done", timeout=60)
final_answer = _tool_call_loop(session, max_rounds=4)
args_event = parse_last(
first, "response.function_call_arguments.done", FunctionCallArgumentsDone
assert "72" in transcript(final_answer), (
"the answer after the tool results did not use them"
)
assert args_event is not None, "model did not emit a function call"
args = WeatherArgs.model_validate_json(args_event.arguments)
item = function_call_item(first)
assert item is not None, "no completed function_call output item"
assert item.name == "get_weather"
assert item.call_id == args_event.call_id
tool_result = WeatherResult(city=args.city, temperature_f=72)
session.send(
ConversationItemCreate(
item=FunctionCallOutputItem(
call_id=args_event.call_id, output=tool_result.model_dump_json()
def _tool_call_loop(session: RealtimeSession, *, max_rounds: int) -> tuple[ReceivedEvent, ...]:
"""Answer every function_call in each response.done until one arrives without any.
Nova 2 Sonic emits two tool calls per question, each inside its own response.done,
so a one-call round trip cannot wait for "the" follow-up: the second response.done
is another call, not the answer. The loop sends one function_call_output per call
then a single response.create, and returns the events of the first call-free round."""
for round_index in range(max_rounds):
events = session.collect_until("response.done", timeout=60)
done = parse_last(events, "response.done", ResponseDone)
assert done is not None
calls = tuple(
item for item in done.response.output or () if item.type == "function_call"
)
if not calls:
assert round_index > 0, "model did not emit a function call"
return events
done_args = tuple(
FunctionCallArgumentsDone.model_validate_json(event.payload)
for event in events
if event.type == "response.function_call_arguments.done"
)
for item in calls:
assert item.name == "get_weather", item
args_event = next(
event for event in done_args if event.call_id == item.call_id
)
args = WeatherArgs.model_validate_json(args_event.arguments)
tool_result = WeatherResult(city=args.city, temperature_f=72)
session.send(
ConversationItemCreate(
item=FunctionCallOutputItem(
call_id=args_event.call_id, output=tool_result.model_dump_json()
)
)
)
)
session.send(ResponseCreate())
second = session.collect_until("response.done", timeout=60)
assert "72" in transcript(second), "follow-up did not use the tool result"
raise AssertionError(f"still calling tools after {max_rounds} rounds")
_REFUSED_UPSTREAMS = (

3
tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py
View file

@ -91,6 +91,7 @@ TOOL_PARAMS: Final = tuple(
),
)
for p in PROVIDERS
if p.id != "bedrock"
)
AUDIO_OUTPUT_PARAMS: Final = tuple(
@ -109,6 +110,7 @@ AUDIO_OUTPUT_PARAMS: Final = tuple(
),
)
for p in PROVIDERS
if p.id != "bedrock"
)
AUDIO_INPUT_PARAMS: Final = tuple(
@ -127,6 +129,7 @@ AUDIO_INPUT_PARAMS: Final = tuple(
),
)
for p in PROVIDERS
if p.id != "bedrock"
)
# PCM16 24 kHz mono WAV of "What is the weather in Paris?" (generated via macOS

5
tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py
View file

@ -64,6 +64,9 @@ from pipecat_service import LiteLLMRealtimeLLMService # noqa: E402
# Vertex native-audio live is flaky through pipecat tool calling (upstream
# pipecat-ai/pipecat#2544); raw-ws tool_call_round_trip[vertex_ai] is the
# source of truth for that provider. Keep openai/azure/gemini here.
# Bedrock event shapes are not yet pipecat-compatible (missing GA fields and
# event names, and no tool-result round trip reaches Nova); raw-ws
# tool_call_round_trip[bedrock] is the source of truth for that provider.
PROVIDER_PARAMS = [
pytest.param(
p,
@ -79,7 +82,7 @@ PROVIDER_PARAMS = [
)
),
)
for p in PROVIDERS if p.id != "vertex_ai"
for p in PROVIDERS if p.id not in ("vertex_ai", "bedrock")
]
WEATHER_TOOL = ToolsSchema(

39
tests/e2e/perf/cases/small_bodies_5rps.py
View file

@ -1,20 +1,23 @@
"""Small bodies at a low rate on the small gateway: the shape of github issue #46276 (short chat
requests, master key and virtual keys, no JWT, no guardrails) at 5 rps on one 4-worker pod, far
under what the pod sustains, so nothing queues and gateway_added_seconds is the gateway's own
per-request overhead. A slower middleware chain or an extra Redis round trip per request shows up
directly instead of hiding behind queue wait."""
"""Small bodies at a low rate on the small gateway: the request sizes of github issue #46276 (short
chat requests, half streamed) from the same callers as the mixed case (two teams, JWT users in and
out of a team, virtual keys, the pre and post guardrail) at 5 rps on one 4-worker pod, far under
what the pod sustains, so nothing queues and gateway_added_seconds is the gateway's own per-request
overhead. A slower middleware chain or an extra Redis round trip per request shows up directly
instead of hiding behind queue wait."""
from __future__ import annotations
from typing import Final
from setup.performance_test_case import (
Guardrail,
Limit,
MockLatency,
PerformanceTestCase,
RequestShape,
Team,
Traffic,
User,
VirtualKey,
)
@ -22,16 +25,32 @@ TINY: Final = RequestShape("tiny", input_tokens=(2_500, 4_000))
CASE: Final = PerformanceTestCase(
name="small_bodies_5rps",
teams=(Team("team-a", models=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock")),),
users=(),
teams=(
Team("team-a", models=("gpt-6-luna", "claude-sonnet-5"), max_budget=5000.0, rpm_limit=6000),
Team("team-b", models=("claude-sonnet-5-bedrock", "claude-sonnet-5"), max_budget=5000.0),
),
users=(
User("alice", team="team-a"),
User("bob", team="team-a"),
User("carol", team="team-a"),
User("dave", team="team-b"),
User("erin", team="team-b"),
User("frank", team="team-b"),
User("grace"),
User("heidi"),
User("ivan"),
User("judy", max_budget=1000.0),
),
virtual_keys=(
VirtualKey("key-a1", team="team-a"),
VirtualKey("key-a2", team="team-a"),
VirtualKey("key-a1", team="team-a", models=("gpt-6-luna",), rpm_limit=1200),
VirtualKey("key-a2", team="team-a", user="alice"),
VirtualKey("key-b1", team="team-b"),
VirtualKey("key-b2", team="team-b", user="dave"),
VirtualKey("key-plain"),
),
model_groups=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock"),
mock_latency=MockLatency(seconds=1.0, jitter_seconds=0.25),
guardrails=(),
guardrails=(Guardrail("perf-pre", mode="pre_call"), Guardrail("perf-post", mode="post_call")),
traffic=Traffic(
seed=11,
rps=5.0,

58
tests/e2e/perf/setup/locustfile.py
View file

@ -10,13 +10,14 @@ line), paces to `PERF_RPS` over `PERF_USERS` users, and quits once the last requ
from __future__ import annotations
import os
import random
import time
from collections.abc import Iterator
from pathlib import Path
from typing import Final, cast
import gevent
from locust import FastHttpUser, constant_throughput, task
from locust import FastHttpUser, task
from locust.contrib.fasthttp import FastResponse, ResponseContextManager
from setup.wire import Endpoint, WireRecord, WireRequest, request_body
from pydantic import TypeAdapter
@ -28,6 +29,9 @@ _RECORDS_PATH: Final = Path(os.environ["PERF_RECORDS_FILE"])
_RPS: Final = float(os.environ["PERF_RPS"])
_USERS: Final = int(os.environ["PERF_USERS"])
_NEXT: Final[Iterator[WireRequest]] = iter(_PLAN)
_INTERVAL_S: Final = _USERS / _RPS
_JITTER: Final = 0.25
_RNG: Final = random.Random(_PLAN[0].marker if _PLAN else "empty-plan")
_RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8")
_written = 0
BLOCK_REASON: Final = "perf-test-block"
@ -84,6 +88,32 @@ def delta_of_event(endpoint: Endpoint, payload: dict[str, object]) -> str:
return ""
_RESPONSES_TOKEN_EVENTS: Final = frozenset(
{
"response.output_text.delta",
"response.refusal.delta",
"response.reasoning_text.delta",
"response.reasoning_summary_text.delta",
"response.function_call_arguments.delta",
}
)
def token_in_event(endpoint: Endpoint, payload: dict[str, object]) -> bool:
"""Whether one SSE event carries a generated token of any kind: visible text, reasoning or tool-call arguments.
The first such event is the stream's first token. `delta_of_event` stays visible text only because the body
check compares it with the text the mock generated, and a tool-call reply has none."""
if delta_of_event(endpoint, payload):
return True
if endpoint == "chat_completions":
choices: Final = _as_list(payload.get("choices"))
delta: Final = _as_dict(_as_dict(choices[0]).get("delta")) if choices else {}
return any(delta.get(key) for key in ("reasoning_content", "reasoning", "tool_calls", "function_call"))
if endpoint == "responses":
return payload.get("type") in _RESPONSES_TOKEN_EVENTS
return payload.get("type") == "content_block_delta"
def _events(buffer: bytes) -> tuple[list[dict[str, object]], bytes]:
"""Complete SSE events in `buffer` as their JSON data, plus the unfinished tail."""
events: list[dict[str, object]] = []
@ -169,12 +199,26 @@ def _authorization(user: FastHttpUser, item: WireRequest) -> str:
class PlanUser(FastHttpUser):
"""One simulated client; all users share the plan and together hold the planned rate."""
"""One simulated client; all users share the plan and together hold the planned rate. Each user starts at
a random point of its first interval and jitters every later interval by up to `_JITTER` either way, so the
users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate."""
wait_time = constant_throughput(_RPS / _USERS)
network_timeout = 120.0
connection_timeout = 30.0
def on_start(self) -> None:
self._last_run = 0.0
self._last_wait = 0.0
gevent.sleep(_RNG.uniform(0.0, _INTERVAL_S))
def wait_time(self) -> float:
target: Final = _INTERVAL_S * _RNG.uniform(1.0 - _JITTER, 1.0 + _JITTER)
now: Final = time.monotonic()
run_s: Final = now - self._last_run - self._last_wait if self._last_run else 0.0
self._last_wait = max(0.0, target - run_s)
self._last_run = now
return self._last_wait
@task
def send(self) -> None:
item: Final = next(_NEXT, None)
@ -221,11 +265,9 @@ class PlanUser(FastHttpUser):
buffer += line
events, buffer = _events(buffer)
for event in events:
delta = delta_of_event(item.endpoint, event)
if delta:
if first_token_s is None:
first_token_s = now
text += delta
if first_token_s is None and token_in_event(item.endpoint, event):
first_token_s = now
text += delta_of_event(item.endpoint, event)
else:
body_text = context.text or ""
first_byte_s = time.perf_counter() - started

2
tests/e2e/ui/integration.config.ts
View file

@ -23,6 +23,8 @@ export default defineConfig({
use: {
...devices["Desktop Chrome"],
baseURL: UI_BASE_URL,
// The proxy closes idle keep-alive sockets after 5s and a reused one fails with "socket hang up"
extraHTTPHeaders: { Connection: "close" },
actionTimeout: 15_000,
navigationTimeout: 30_000,
trace: "retain-on-failure",

2
tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts
View file

@ -12,7 +12,7 @@ import { captureRequestBody } from "../../helpers/roundTrip";
import { masterKey } from "../../helpers/traffic";
const master = masterKey();
const headers = { Authorization: `Bearer ${master}`, Connection: "close" };
const headers = { Authorization: `Bearer ${master}` };
const TOKEN = "USER_TOKEN";
type EnvVarStatus = {

2
tests/integration/README.md
View file

@ -28,6 +28,8 @@ Owned proxies disable `.env` loading by default. A contract testing dotenv behav
Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline by default, so uvicorn kills a worker that stays paused longer than that. Tests never signal proxy processes themselves. `OwnedProxy.kill_worker` SIGKILLs a worker and returns once its replacement answers readiness with its own pid, `replacing_worker` does the same for a kill that has to land mid-burst, and `suspended` pauses a worker or the supervisor, which only works on a proxy launched with `healthcheck=WorkerHealthcheck.LENIENT`. Helper processes that are not the proxy stop through `stop_child`, `stop_processes` and `stop_process_group` and pause through `paused`. `scripts/check_test_quality.py` rule TQ010 rejects a raw `os.kill`, SIGKILL / SIGSTOP / SIGCONT, `.kill()`, `.terminate()`, `.suspend()` or `.send_signal()` anywhere under `tests/integration` outside `_support`, and TQ011 rejects an assertion that caps an elapsed time with a fixed number. Bound timings with `_support/timing.py` instead, relative to a measured control or to the injected delay. `run.py` measures how long this box takes to import the proxy application before it starts pytest, records it in `boot-calibration.json` beside the results and exports it as `INTEGRATION_IMPORT_SECONDS`. Readiness deadlines, graceful stops and the default per-test timeout are multiples of that measurement (`_support/budget.py`), so a slower box gets longer budgets and a two-worker proxy gets twice the readiness deadline of a one-worker one. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing, for example by sending each half while the other worker is suspended
Any port the harness picks before the server that listens on it has started comes from `allocate_port()` in `_support/ports.py`. That covers owned proxies, scripted upstreams, fake providers, OTLP sinks, database relays and Redis. A server that binds port 0 itself and keeps that socket, such as `wire_server` without a port, needs no allocation. Each pytest process claims a block of 200 ports between 10000 and 32767 by holding an exclusive lock file per block in the system temp directory, so two runs on one machine, such as parallel worktrees, never share a block, and the OS drops the lock when the process exits. The range sits above the fixed service ports and below the Linux and macOS ephemeral ranges, so a port-0 bind or an outbound connection never lands in a block. Ports inside a block go out in order and are never handed out twice in a run, a port that something outside the harness already holds is skipped, and a process that uses up its block claims the next free one. A test that restarts a server on the same port allocates that port up front instead of reading it back from the first server. TQ012 in `scripts/check_test_quality.py` rejects reading a port back from a port-0 bind anywhere under `tests/integration`
The CircleCI workflow starts its own database and Redis, restricts test-phase egress to its owned services and writes JUnit plus an executed-node manifest. Missing setup, failed cleanup or a selected test with neither a passed call nor a skip fail qualification. Skipped nodes are listed under `skipped` in `execution.json`, so the skip reasons double as the open bug list. `GITHUB_FILES` in `run.py` lists the files excluded from the `run_integration.sh` selection
There is no hand-maintained list of test cases. A positional argument is a file of the group or a pytest node id inside one (`path::test[param]`), so one cell of a parametrized file can run alone. The runner fails only when pytest fails, when collection errors, or when a selected file collects zero tests. Older tests still carry `@pytest.mark.covers(...)` decorators; the marker stays registered so they collect, but the IDs are not checked against anything and new tests should not use it. The GitHub Actions coverage census reads the `GROUPS` literal in `run.py` and treats every `tests/integration/<directory>/test_*.py` file in a scheduled group as owned by CircleCI

8
tests/integration/_support/async_server.py
View file

@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import ssl
from collections.abc import Awaitable, Callable
from contextlib import suppress
from threading import Thread
@ -16,10 +17,11 @@ async def close_writer(writer: asyncio.StreamWriter) -> None:
class LoopbackServer:
def __init__(self, port: int, handler: Handler) -> None:
def __init__(self, port: int, handler: Handler, ssl_context: ssl.SSLContext | None = None) -> None:
self.loop: Final = asyncio.new_event_loop()
self._port: Final = port
self._handler: Final = handler
self._ssl_context: Final = ssl_context
self._server: asyncio.Server | None = None
self._thread: Final = Thread(target=self._run, name=f"integration-relay-{port}", daemon=True)
self._connections: tuple[asyncio.Task[None], ...] = ()
@ -42,7 +44,7 @@ class LoopbackServer:
finally:
self.loop.call_soon_threadsafe(self.loop.stop)
self._thread.join(10)
assert not self._thread.is_alive(), f"Database relay thread on port {self._port} did not stop"
assert not self._thread.is_alive(), f"Loopback server thread on port {self._port} did not stop"
def _run(self) -> None:
asyncio.set_event_loop(self.loop)
@ -53,7 +55,7 @@ class LoopbackServer:
self.loop.close()
async def _bind(self) -> None:
self._server = await asyncio.start_server(self._serve, "127.0.0.1", self._port)
self._server = await asyncio.start_server(self._serve, "127.0.0.1", self._port, ssl=self._ssl_context)
async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
task: Final = asyncio.current_task()

429
tests/integration/_support/bedrock_realtime_peer.py Normal file
View file

@ -0,0 +1,429 @@
"""A scripted Bedrock Nova Sonic realtime peer.
The handler dials ``api_base`` with ``invoke_model_with_bidirectional_stream`` through
``AWSCRTHTTPClient``, which only speaks h2 over TLS ALPN, so this peer is a loopback TLS
server offering ``h2`` on ``async_server.LoopbackServer``'s thread-owned asyncio loop.
It answers ``POST /model/{id}/invoke-with-bidirectional-stream``, decodes the client's
signed event-stream chunks into the ``{"event": {...}}`` documents Bedrock would read,
and answers each consumed ``Turn`` with event-stream ``chunk`` frames carrying
``{"bytes": <base64 event json>}`` payloads, the reverse of what the SDK produces.
Turn triggers are inbound events. ``user_input`` fires on the ``contentEnd`` closing the
first inbound USER-role TEXT content block, the end of a forwarded conversation.item.create
text message (the injected trigger-audio block is USER/AUDIO, so it never trips it).
``tool_result`` fires on the Nth inbound ``toolResult`` event, which the proxy sends inside
a TOOL contentStart/toolResult/contentEnd triple whose ``toolResultInputConfiguration``
``toolUseId`` names the call being answered.
"""
from __future__ import annotations
import asyncio
import base64
import json
import struct
import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager, suppress
from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from queue import SimpleQueue
from typing import Final, Literal
import h2.config
import h2.connection
import h2.events
import h2.exceptions
from integration._support.async_server import Handler, LoopbackServer
from integration._support.process import _free_port
from integration._support.tls import server_context
from integration._support.upstream import aws_event_stream_frame
from pydantic import JsonValue, TypeAdapter
BedrockEvent = Mapping[str, JsonValue]
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _ids() -> tuple[str, str, str, str]:
return (uuid.uuid4().hex, uuid.uuid4().hex, uuid.uuid4().hex, uuid.uuid4().hex)
def _block(prompt: str, content: str, session: str, completion: str) -> dict[str, JsonValue]:
return {"promptName": prompt, "contentName": content, "sessionId": session, "completionId": completion}
def user_transcript(text: str) -> tuple[BedrockEvent, ...]:
"""The USER ASR transcript block Bedrock emits once it has consumed the user turn."""
prompt, content, session, completion = _ids()
return (
{
"event": {
"contentStart": {
**_block(prompt, content, session, completion),
"role": "USER",
"type": "TEXT",
"interactive": False,
"textOutputConfiguration": {"mediaType": "text/plain"},
"additionalModelFields": json.dumps({"generationStage": "FINAL"}),
}
}
},
{"event": {"textOutput": {**_block(prompt, content, session, completion), "role": "USER", "content": text}}},
{
"event": {
"contentEnd": {
**_block(prompt, content, session, completion),
"stopReason": "PARTIAL_TURN",
"type": "TEXT",
}
}
},
)
def tool_use(tool_use_id: str, name: str, content: str | None) -> tuple[BedrockEvent, ...]:
"""The role-TOOL contentStart/toolUse/contentEnd triple Bedrock emits per tool call."""
prompt, content_name, session, completion = _ids()
tool_use_body: Final[dict[str, JsonValue]] = {
**_block(prompt, content_name, session, completion),
"role": "TOOL",
"toolName": name,
"toolUseId": tool_use_id,
}
if content is not None:
tool_use_body["content"] = content
return (
{
"event": {
"contentStart": {
**_block(prompt, content_name, session, completion),
"role": "TOOL",
"type": "TOOL",
"toolUseOutputConfiguration": {"mediaType": "application/json"},
}
}
},
{"event": {"toolUse": tool_use_body}},
{
"event": {
"contentEnd": {
**_block(prompt, content_name, session, completion),
"stopReason": "TOOL_USE",
"type": "TOOL",
}
}
},
)
def assistant_text(text: str, stop_reason: str = "END_TURN") -> tuple[BedrockEvent, ...]:
"""An ASSISTANT text block (SPECULATIVE stage) closed with `stop_reason`."""
prompt, content, session, completion = _ids()
return (
{
"event": {
"contentStart": {
**_block(prompt, content, session, completion),
"role": "ASSISTANT",
"type": "TEXT",
"interactive": False,
"textOutputConfiguration": {"mediaType": "text/plain"},
"additionalModelFields": json.dumps({"generationStage": "SPECULATIVE"}),
}
}
},
{
"event": {
"textOutput": {**_block(prompt, content, session, completion), "role": "ASSISTANT", "content": text}
}
},
{
"event": {
"contentEnd": {
**_block(prompt, content, session, completion),
"stopReason": stop_reason,
"type": "TEXT",
}
}
},
)
def usage(input_tokens: int, output_tokens: int) -> BedrockEvent:
total: Final = {
"input": {"speechTokens": 0, "textTokens": input_tokens},
"output": {"speechTokens": 0, "textTokens": output_tokens},
}
return {
"event": {
"usageEvent": {
"promptName": uuid.uuid4().hex,
"sessionId": uuid.uuid4().hex,
"completionId": uuid.uuid4().hex,
"details": {"delta": total["output"], "total": total},
"totalInputTokens": input_tokens,
"totalOutputTokens": output_tokens,
"totalTokens": input_tokens + output_tokens,
}
}
}
@dataclass(frozen=True, slots=True)
class Turn:
after: Literal["user_input", "tool_result"]
events: tuple[BedrockEvent, ...]
def _event_body(event: BedrockEvent) -> tuple[str, Mapping[str, JsonValue]]:
body: Final = event["event"]
assert isinstance(body, Mapping) and len(body) == 1, event
name: Final = next(iter(body))
fields: Final = body[name]
assert isinstance(fields, Mapping), event
return name, fields
def _outbound_chunk(event: BedrockEvent) -> bytes:
wrapped: Final = json.dumps({"bytes": base64.b64encode(json.dumps(event).encode()).decode()}).encode()
return aws_event_stream_frame(
{":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, wrapped
)
def _parse_frames(buffer: bytes) -> tuple[tuple[bytes, ...], bytes]:
"""Slice complete event-stream frames out of `buffer`; return (payloads, remainder)."""
payloads: Final[list[bytes]] = []
view = memoryview(buffer)
while len(view) >= 12:
total, header_length = struct.unpack_from("!II", view, 0)
if total > len(view):
break
payloads.append(view[12 + header_length : total - 4].tobytes())
view = view[total:]
return tuple(payloads), view.tobytes()
def _frame_payload(raw: bytes) -> bytes:
"""The payload slice of one complete framed message (prelude + headers + payload + CRC)."""
if len(raw) < 16:
return b""
_total, header_length = struct.unpack_from("!II", raw, 0)
return raw[12 + header_length : -4]
def _decode_inbound(chunk_payload: bytes) -> BedrockEvent | None:
"""One signed chunk frame's payload -> the inner chunk's {"bytes": ...} event."""
inner_payload: Final = _frame_payload(chunk_payload)
try:
wrapped: Final = _JSON_OBJECT.validate_json(inner_payload)
except ValueError:
return None
encoded: Final = wrapped.get("bytes")
if not isinstance(encoded, str):
return None
try:
return _JSON_OBJECT.validate_json(base64.b64decode(encoded))
except ValueError:
return None
class _Inbound:
"""Inbound event bookkeeping for turn triggers and accessors."""
def __init__(self) -> None:
self.events: Final = SimpleQueue[BedrockEvent]()
self.open_blocks: Final[dict[str, tuple[str, str]]] = {}
self.tool_use_ids: Final[dict[str, str]] = {}
self.tool_results_seen = 0
self.user_text_closed = False
def record(self, event: BedrockEvent) -> None:
self.events.put(event)
name, fields = _event_body(event)
content_name = fields.get("contentName")
if name == "contentStart" and isinstance(content_name, str):
self.open_blocks[content_name] = (str(fields.get("role")), str(fields.get("type")))
configuration = fields.get("toolResultInputConfiguration")
if isinstance(configuration, Mapping) and isinstance(configuration.get("toolUseId"), str):
self.tool_use_ids[content_name] = configuration["toolUseId"]
elif name == "contentEnd" and isinstance(content_name, str):
role_type: Final = self.open_blocks.pop(content_name, ("", ""))
if role_type == ("USER", "TEXT"):
self.user_text_closed = True
elif name == "toolResult":
self.tool_results_seen += 1
class _Script:
"""Consumes turns in order as their inbound triggers arrive."""
def __init__(self, turns: tuple[Turn, ...]) -> None:
self._turns: Final = turns
self._next = 0
@property
def exhausted(self) -> bool:
return self._next >= len(self._turns)
def ready_events(self, state: _Inbound, event: BedrockEvent) -> tuple[BedrockEvent, ...]:
if self._next >= len(self._turns):
return ()
turn: Final = self._turns[self._next]
if turn.after == "user_input":
if not state.user_text_closed:
return ()
else:
name, _fields = _event_body(event)
if name != "toolResult":
return ()
self._next += 1
return turn.events
class _PeerState:
def __init__(self, turns: tuple[Turn, ...]) -> None:
self.inbound: Final = _Inbound()
self.script: Final = _Script(turns)
@dataclass(frozen=True, slots=True)
class BedrockRealtimePeer:
api_base: str
_state: _PeerState
_drained: list[BedrockEvent]
def received(self) -> tuple[BedrockEvent, ...]:
while not self._state.inbound.events.empty():
self._drained.append(self._state.inbound.events.get()) # drained log, not a built value
return tuple(self._drained)
def tool_specs(self) -> tuple[Mapping[str, JsonValue], ...]:
return tuple(chain.from_iterable(_tool_specs(event) for event in self.received()))
def tool_results(self) -> tuple[tuple[str, str], ...]:
"""(toolUseId, content) per inbound toolResult in arrival order."""
return tuple(result for event in self.received() if (result := self._tool_result(event)) is not None)
def _tool_result(self, event: BedrockEvent) -> tuple[str, str] | None:
name, fields = _event_body(event)
if name != "toolResult":
return None
content_name: Final = fields.get("contentName")
content: Final = fields.get("content")
if not isinstance(content_name, str) or not isinstance(content, str):
return None
return (self._state.inbound.tool_use_ids.get(content_name, ""), content)
def _tool_specs(event: BedrockEvent) -> tuple[Mapping[str, JsonValue], ...]:
name, fields = _event_body(event)
if name != "promptStart":
return ()
configuration: Final = fields.get("toolConfiguration")
if not isinstance(configuration, Mapping):
return ()
tools: Final = configuration.get("tools")
if not isinstance(tools, list):
return ()
return tuple(
entry["toolSpec"]
for entry in tools
if isinstance(entry, Mapping) and isinstance(entry.get("toolSpec"), Mapping)
)
class _H2Session:
"""One bidirectional-stream connection: decode inbound, answer scripted turns."""
def __init__(self, state: _PeerState, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
self._state: Final = state
self._reader: Final = reader
self._writer: Final = writer
self._h2: Final = h2.connection.H2Connection(config=h2.config.H2Configuration(client_side=False))
self._buffers: Final[dict[int, bytes]] = {}
self._ended: set[int] = set()
self._closed = False
async def run(self) -> None:
self._h2.initiate_connection()
self._writer.write(self._h2.data_to_send())
await self._writer.drain()
while True:
incoming: Final = await self._reader.read(65536)
if not incoming:
return
for event in self._h2.receive_data(incoming):
self._apply(event)
pending: Final = self._h2.data_to_send()
if pending:
self._writer.write(pending)
await self._writer.drain()
if self._closed:
self._writer.close()
return
def _apply(self, event: h2.events.Event) -> None:
if isinstance(event, h2.events.RequestReceived):
self._h2.send_headers(
event.stream_id,
[(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")],
)
elif isinstance(event, h2.events.DataReceived):
self._h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id)
self._buffers[event.stream_id] = self._buffers.get(event.stream_id, b"") + event.data
payloads, remainder = _parse_frames(self._buffers[event.stream_id])
self._buffers[event.stream_id] = remainder
for payload in payloads:
decoded: Final = _decode_inbound(payload)
if decoded is None:
continue
self._state.inbound.record(decoded)
if event.stream_id not in self._ended:
for outbound in self._state.script.ready_events(self._state.inbound, decoded):
self._h2.send_data(event.stream_id, _outbound_chunk(outbound))
elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)):
self._close_stream(event.stream_id)
elif isinstance(event, h2.events.ConnectionTerminated):
self._writer.close()
def _close_stream(self, stream_id: int) -> None:
self._end_stream(stream_id)
self._closed = True
self._h2.close_connection()
def _end_stream(self, stream_id: int) -> None:
if stream_id in self._ended:
return
self._ended.add(stream_id)
with suppress(h2.exceptions.ProtocolError):
self._h2.end_stream(stream_id)
def _session_handler(state: _PeerState) -> Handler:
async def serve(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
with suppress(ConnectionError, OSError, asyncio.IncompleteReadError, h2.exceptions.StreamClosedError):
await _H2Session(state, reader, writer).run()
return serve
@contextmanager
def bedrock_realtime_peer(turns: tuple[Turn, ...], cert: tuple[Path, Path]) -> Iterator[BedrockRealtimePeer]:
"""Serve scripted `turns` on a loopback TLS h2 endpoint; yield the peer.
`cert` is the self-signed pair the owned proxy already trusts via SSL_CERT_FILE.
SigV4 headers are ignored; the stream closes cleanly when the client ends it.
"""
port: Final = _free_port()
context: Final = server_context(cert[0], cert[1])
context.set_alpn_protocols(["h2"])
state: Final = _PeerState(turns)
server: Final = LoopbackServer(port, _session_handler(state), ssl_context=context)
server.start()
try:
yield BedrockRealtimePeer(api_base=f"https://localhost:{port}", _state=state, _drained=[])
finally:
server.stop()

18
tests/integration/_support/database_relay.py
View file

@ -1,5 +1,4 @@
import asyncio
import socket
import threading
import time
from collections.abc import Generator
@ -7,23 +6,16 @@ from contextlib import contextmanager
from typing import Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from pydantic import TypeAdapter
from tests.integration._support.async_server import LoopbackServer, close_writer
from tests.integration._support.ports import allocate_port
PORT: Final = TypeAdapter(int)
OUTAGE_SECONDS: Final = 10.0
def _free_port() -> int:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
return PORT.validate_python(reserve.getsockname()[1])
class DatabaseRelay:
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
self.port: Final = _free_port()
self.port: Final = allocate_port()
self._upstream_host: Final = upstream_host
self._upstream_port: Final = upstream_port
self._trigger: Final = trigger
@ -83,7 +75,7 @@ class DatabaseRelay:
class HeldStatementRelay:
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
self.port: Final = _free_port()
self.port: Final = allocate_port()
self._upstream_host: Final = upstream_host
self._upstream_port: Final = upstream_port
self._trigger: Final = trigger
@ -144,7 +136,7 @@ class TriggerScanner:
class DroppedConnectionRelay:
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
self.port: Final = _free_port()
self.port: Final = allocate_port()
self._upstream_host: Final = upstream_host
self._upstream_port: Final = upstream_port
self._trigger: Final = trigger
@ -231,7 +223,7 @@ class FrontendStatementScanner:
class StatementCountingRelay:
def __init__(self, upstream_host: str, upstream_port: int) -> None:
self.port: Final = _free_port()
self.port: Final = allocate_port()
self._upstream_host: Final = upstream_host
self._upstream_port: Final = upstream_port
self._loop: Final = asyncio.new_event_loop()

19
tests/integration/_support/otlp_sink.py
View file

@ -13,7 +13,6 @@ import datetime
import json
import os
import signal
import socket
import ssl
import subprocess
import sys
@ -35,6 +34,8 @@ import yaml
from pydantic import JsonValue, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from integration._support.ports import allocate_port
INTERNAL_MARKERS: Final = ("gen_ai.operation.name", "mcp.method.name", "litellm.guardrail_name")
@ -421,12 +422,6 @@ class ConnectSink:
ca_pem: str
def _free_port() -> int:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
return int(reserve.getsockname()[1])
def _pid_reachable(url: str) -> bool:
try:
return httpx.get(f"{url}/__pid", trust_env=False, timeout=2).status_code == 200
@ -439,7 +434,7 @@ def owned_sinks(directory: Path) -> Iterator[SpanSinks]:
from integration._support.process import group_members, signal_group, stop_root_process
directory.mkdir(parents=True, exist_ok=True)
ports: Final = tuple(_free_port() for _ in range(3))
ports: Final = tuple(allocate_port() for _ in range(3))
root: Final = Path(__file__).resolve().parents[3]
with ExitStack() as stack:
processes: Final = tuple(
@ -509,8 +504,8 @@ def _await_sink(url: str) -> None:
@contextmanager
def owned_grpc_sink(directory: Path) -> Iterator[GrpcSink]:
http_port: Final = _free_port()
grpc_port: Final = _free_port()
http_port: Final = allocate_port()
grpc_port: Final = allocate_port()
with _spawn_sink(directory, "otlp-grpc-sink.log", ["--port", str(http_port), "--grpc-port", str(grpc_port)]):
control_url: Final = f"http://127.0.0.1:{http_port}"
_await_sink(control_url)
@ -519,8 +514,8 @@ def owned_grpc_sink(directory: Path) -> Iterator[GrpcSink]:
@contextmanager
def owned_connect_sink(directory: Path) -> Iterator[ConnectSink]:
http_port: Final = _free_port()
tunnel_port: Final = _free_port()
http_port: Final = allocate_port()
tunnel_port: Final = allocate_port()
ca_dir: Final = directory / "mitm"
with _spawn_sink(
directory,

61
tests/integration/_support/ports.py Normal file
View file

@ -0,0 +1,61 @@
import fcntl
import os
import socket
import tempfile
import threading
from collections.abc import Iterator
from pathlib import Path
from typing import Final
FIRST_PORT: Final = 10_000
LAST_PORT: Final = 32_767
BLOCK_SIZE: Final = 200
LOCK_DIRECTORY: Final = Path(tempfile.gettempdir()) / "litellm-integration-ports"
def _blocks() -> Iterator[range]:
return (range(start, start + BLOCK_SIZE) for start in range(FIRST_PORT, LAST_PORT + 2 - BLOCK_SIZE, BLOCK_SIZE))
def _claim_until_exit(block: range, directory: Path) -> bool:
descriptor: Final = os.open(directory / f"{block.start}.lock", os.O_RDWR | os.O_CREAT, 0o600)
try:
fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError:
os.close(descriptor)
return False
return True
def _bindable(port: int) -> bool:
with socket.socket() as probe:
try:
probe.bind(("127.0.0.1", port))
except OSError:
return False
return True
def _free_ports(directory: Path) -> Iterator[int]:
directory.mkdir(parents=True, exist_ok=True)
for block in _blocks():
if _claim_until_exit(block, directory):
yield from filter(_bindable, block)
class PortAllocator:
def __init__(self, directory: Path) -> None:
self._directory: Final = directory
self._guard: Final = threading.Lock()
self._ports: Final = _free_ports(directory)
def allocate(self) -> int:
with self._guard:
port: Final = next(self._ports, None)
assert port is not None, (
f"Every integration port block in {FIRST_PORT}-{LAST_PORT} is claimed (locks in {self._directory})"
)
return port
allocate_port: Final = PortAllocator(LOCK_DIRECTORY).allocate

77
tests/integration/_support/process.py
View file

@ -1,10 +1,8 @@
import errno
import json
import threading
import os
import re
import signal
import socket
import subprocess
import sys
import time
@ -33,6 +31,7 @@ from pydantic import TypeAdapter
from tests.integration._support.budget import import_seconds, readiness_for, readiness_seconds
from tests.integration._support.client import GATEWAY_LIMITS, Gateway
from tests.integration._support.ports import allocate_port
DB_PUSH: Final = ("--use_prisma_db_push",)
MIGRATE_DEPLOY: Final = ()
@ -406,16 +405,6 @@ def _stop(process: subprocess.Popen[bytes]) -> None:
raise
_PORT_ATTEMPTS: Final = 3
_BIND_COLLISION: Final = os.strerror(errno.EADDRINUSE).lower()
def _free_port() -> int:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
return reserve.getsockname()[1]
@dataclass(frozen=True, slots=True)
class _Launch:
process: subprocess.Popen[bytes]
@ -458,7 +447,7 @@ def _stop_launch(launch: _Launch) -> None:
def _launch(command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path) -> _Launch:
port: Final = _free_port()
port: Final = allocate_port()
log_path: Final = output / f"owned-proxy-{uuid.uuid4().hex}.log"
started: Final = time.monotonic()
with log_path.open("w") as log:
@ -473,10 +462,6 @@ def _launch(command: tuple[str, ...], root: Path, environment: Mapping[str, str]
return _Launch(process, port, log_path, started)
def _lost_port_race(exit_code: int | None, log: Path) -> bool:
return exit_code is not None and _BIND_COLLISION in log.read_text().lower()
def _owns_listener(launch: _Launch) -> bool:
for member in group_members(launch.process.pid):
try:
@ -600,42 +585,36 @@ def _wait_until_ready(launch: _Launch, workers: int = 1) -> BootPhases | None:
)
def _launch_until_bound(
command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path, attempts: int, workers: int
def _launch_until_ready(
command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path, workers: int
) -> _Launch:
launch: Final = _launch(command, root, environment, output)
try:
phases: Final = _wait_until_ready(launch, workers)
exit_code: Final = launch.process.poll()
assert exit_code is None or (attempts > 1 and _lost_port_race(exit_code, launch.log)), (
"Owned proxy exited before readiness"
)
assert launch.process.poll() is None, "Owned proxy exited before readiness"
except BaseException:
_stop_launch(launch)
raise
if exit_code is None:
assert phases is not None
launch.log.with_suffix(".startup.json").write_text(
json.dumps(
{
"root_pid": launch.process.pid,
"readiness_seconds": time.monotonic() - launch.started,
"logical_cpus": os.cpu_count(),
"workers": workers,
"spawned_at": _wall_clock(launch.started),
"phases": {
"prisma_started": phases.prisma_started,
"prisma_finished": phases.prisma_finished,
"listening": phases.listening,
"first_ready": phases.first_ready,
},
}
)
+ "\n"
assert phases is not None
launch.log.with_suffix(".startup.json").write_text(
json.dumps(
{
"root_pid": launch.process.pid,
"readiness_seconds": time.monotonic() - launch.started,
"logical_cpus": os.cpu_count(),
"workers": workers,
"spawned_at": _wall_clock(launch.started),
"phases": {
"prisma_started": phases.prisma_started,
"prisma_finished": phases.prisma_finished,
"listening": phases.listening,
"first_ready": phases.first_ready,
},
}
)
return launch
_stop_launch(launch)
return _launch_until_bound(command, root, environment, output, attempts - 1, workers)
+ "\n"
)
return launch
def _proxy_root() -> Path:
@ -785,7 +764,7 @@ def owned_proxy_process(
*database_setup,
*extra_arguments,
)
launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS, workers)
launch: Final = _launch_until_ready(command, root, environment, output, workers)
process: Final = launch.process
try:
with httpx.Client(
@ -825,7 +804,7 @@ def configless_proxy_process(
str(healthcheck_seconds(WorkerHealthcheck.STRICT)),
*DB_PUSH,
)
launch: Final = _launch_until_bound(command, _proxy_root(), environment, output, _PORT_ATTEMPTS, workers)
launch: Final = _launch_until_ready(command, _proxy_root(), environment, output, workers)
try:
with httpx.Client(
base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS
@ -858,7 +837,7 @@ def owned_gateway_image(
"--timeout-worker-healthcheck",
str(int(graceful_stop_seconds())),
)
launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS, workers)
launch: Final = _launch_until_ready(command, root, environment, output, workers)
try:
with httpx.Client(
base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS
@ -1040,7 +1019,7 @@ def owned_upstream(
directory: Path, *, tls: bool = False, port: int | None = None, inherited: Mapping[str, str] = os.environ
) -> Generator[UpstreamSlot]:
certificate: Final = self_signed_certificate(directory) if tls else None
slot: Final = UpstreamSlot(directory, port or _free_port(), _harness_root(), certificate, inherited)
slot: Final = UpstreamSlot(directory, port or allocate_port(), _harness_root(), certificate, inherited)
try:
slot.start()
yield slot

106
tests/integration/_support/redis_process.py
View file

@ -1,7 +1,6 @@
import os
import shutil
import signal
import socket
import subprocess
import time
import uuid
@ -14,18 +13,7 @@ from typing import Final, TextIO
from redis import Redis
from redis.exceptions import ConnectionError as RedisConnectionError
PORT_ATTEMPTS: Final = 5
@dataclass(frozen=True, slots=True)
class PortUnavailable:
reason: str
def free_loopback_port() -> int:
with socket.socket() as reservation:
reservation.bind(("127.0.0.1", 0))
return reservation.getsockname()[1]
from tests.integration._support.ports import allocate_port
@dataclass
@ -60,13 +48,8 @@ class OwnedRedis:
return int(result.stdout)
def start(self) -> None:
unavailable: Final = self.try_start()
assert unavailable is None, unavailable.reason
def try_start(self) -> PortUnavailable | None:
assert self.process is None
if self.serves():
return PortUnavailable(f"Another Redis already serves {self.host}:{self.port} before the owned one starts")
assert not self.serves(), f"Another Redis already serves {self.host}:{self.port} before the owned one starts"
process: Final = subprocess.Popen(
self.command, stdout=self.log, stderr=subprocess.STDOUT, start_new_session=True
)
@ -77,7 +60,7 @@ class OwnedRedis:
code: Final = process.poll()
if code is not None:
self.process = None
return PortUnavailable(
raise AssertionError(
f"Owned Redis exited with code {code} before readiness on {self.host}:{self.port}"
)
try:
@ -86,13 +69,13 @@ class OwnedRedis:
expected: Final = self.owned_pid()
if actual == expected:
self.server_pid = actual
return None
return
try:
process.wait(timeout=max(deadline - time.monotonic(), 0))
except subprocess.TimeoutExpired as error:
raise AssertionError("Redis readiness reached a different process") from error
self.process = None
return PortUnavailable(f"A different Redis answered on {self.host}:{self.port} during startup")
raise AssertionError(f"A different Redis answered on {self.host}:{self.port} during startup")
except RedisConnectionError:
pass
assert time.monotonic() < deadline, "Owned Redis readiness deadline exceeded"
@ -148,33 +131,10 @@ class OwnedRedis:
self.signal(signal.SIGCONT)
@contextmanager
def _started(
build: Callable[[int], OwnedRedis],
choose_port: Callable[[], int],
attempts: int,
log_name: str,
) -> Iterator[OwnedRedis]:
server: Final = build(choose_port())
try:
unavailable: Final = server.try_start()
if unavailable is None:
yield server
return
assert attempts > 1, (
f"Owned Redis found no usable port after {PORT_ATTEMPTS} attempts: {unavailable.reason} (log {log_name})"
)
with _started(build, choose_port, attempts - 1, log_name) as started:
yield started
finally:
if server.process is not None:
server.stop()
@contextmanager
def owned_redis(
directory: Path,
choose_port: Callable[[], int] = free_loopback_port,
choose_port: Callable[[], int] = allocate_port,
inherited: Mapping[str, str] = os.environ,
) -> Iterator[OwnedRedis]:
binary: Final = shutil.which("redis-server")
@ -198,31 +158,31 @@ def owned_redis(
output.mkdir(parents=True, exist_ok=True)
directory.mkdir(parents=True, exist_ok=True)
with (output / f"owned-redis-{uuid.uuid4().hex}.log").open("w") as log:
def build(port: int) -> OwnedRedis:
pid_file: Final = (
str(directory / f"owned-redis-{port}.pid")
if binary
else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid"
)
command: Final = (
*prefix,
"--port",
str(port),
"--set-proc-title",
"no",
"--pidfile",
pid_file,
"--bind",
"0.0.0.0" if not binary else "127.0.0.1",
"--protected-mode",
"no",
"--save",
"",
"--appendonly",
"no",
)
return OwnedRedis(host, port, command, log, pid_file)
with _started(build, choose_port, PORT_ATTEMPTS, log.name) as server:
port: Final = choose_port()
pid_file: Final = (
str(directory / f"owned-redis-{port}.pid") if binary else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid"
)
command: Final = (
*prefix,
"--port",
str(port),
"--set-proc-title",
"no",
"--pidfile",
pid_file,
"--bind",
"0.0.0.0" if not binary else "127.0.0.1",
"--protected-mode",
"no",
"--save",
"",
"--appendonly",
"no",
)
server: Final = OwnedRedis(host, port, command, log, pid_file)
try:
server.start()
yield server
finally:
if server.process is not None:
server.stop()

23
tests/integration/_support/sdk_models.py Normal file
View file

@ -0,0 +1,23 @@
from __future__ import annotations
from collections.abc import Iterator
from typing import Final
import anthropic
import openai
from pydantic import BaseModel
_SDK_PACKAGES: Final = frozenset({openai.__name__, anthropic.__name__})
def _subclasses(model: type[BaseModel]) -> Iterator[type[BaseModel]]:
for subclass in model.__subclasses__():
yield subclass
yield from _subclasses(subclass)
def build_sdk_models() -> None:
"""Build every pydantic model the OpenAI and Anthropic SDKs define now rather than on its first use."""
for model in frozenset(_subclasses(BaseModel)):
if model.__module__.partition(".")[0] in _SDK_PACKAGES:
_ = model.model_rebuild(raise_errors=False)

17
tests/integration/_support/workers.py
View file

@ -1,7 +1,6 @@
from __future__ import annotations
import os
import socket
import time
import uuid
from collections.abc import Callable, Iterator, Mapping
@ -23,6 +22,7 @@ from redis.exceptions import TimeoutError as RedisTimeoutError
from tests.integration._support.client import Gateway
from tests.integration._support.database import scratch_database
from tests.integration._support.ports import allocate_port
from tests.integration._support.redis_process import owned_redis
from tests.integration._support.routing import WORKER_DATABASE_PREFIX
@ -78,15 +78,12 @@ def wait_for_services(upstream: str, proxies: tuple[Gateway, ...], cache: Redis)
def worker_endpoints() -> Mapping[str, str]:
with socket.socket() as upstream, socket.socket() as provider:
upstream.bind(("127.0.0.1", 0))
provider.bind(("127.0.0.1", 0))
return MappingProxyType(
{
"INTEGRATION_UPSTREAM_URL": f"http://127.0.0.1:{upstream.getsockname()[1]}",
"INTEGRATION_PROVIDER_PORT": str(provider.getsockname()[1]),
}
)
return MappingProxyType(
{
"INTEGRATION_UPSTREAM_URL": f"http://127.0.0.1:{allocate_port()}",
"INTEGRATION_PROVIDER_PORT": str(allocate_port()),
}
)
def _replica_urls(database: str) -> Mapping[str, str]:

338
tests/integration/authorization/test_agent_least_privilege.py Normal file
View file

@ -0,0 +1,338 @@
import json
import uuid
from contextlib import ExitStack
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
def _write_access_config(directory: Path, *, default_deny: bool, model_name: str, api_base: str) -> Path:
source: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config: Final = {
**source,
"model_list": [
{
"model_name": model_name,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "integration-provider-key",
"api_base": api_base,
},
}
],
"general_settings": {
**source["general_settings"],
"agent_access_default_deny": default_deny,
},
}
path: Final = directory / f"agent-access-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(config))
return path
def _peer_response(request: Request, wire: Wire, agent_name: str) -> Reply:
if request.method == "GET":
assert request.target in ("/.well-known/agent-card.json", "/.well-known/agent.json")
return Reply(
body=json.dumps(
{
"protocolVersion": "0.3",
"name": agent_name,
"description": "Synthetic permission peer",
"version": "1.0.0",
"url": wire.url + "/",
"capabilities": {"streaming": False},
"defaultInputModes": ["text"],
"defaultOutputModes": ["text"],
"skills": [],
}
).encode()
)
assert request.method == "POST" and request.target == "/"
body: Final = json.loads(request.body)
assert body["jsonrpc"] == "2.0" and body["method"] == "message/send"
message: Final = body["params"]["message"]
assert message["role"] == "user"
assert message["kind"] == "message"
assert message["parts"] in (
[{"kind": "text", "text": "synthetic ping"}],
[{"kind": "text", "text": "user: synthetic ping"}],
)
return Reply(
body=json.dumps(
{
"jsonrpc": "2.0",
"id": body["id"],
"result": {
"kind": "message",
"role": "agent",
"messageId": message["messageId"] + "-out",
"parts": [{"kind": "text", "text": "synthetic pong"}],
},
}
).encode()
)
def _create_agent(candidate: Gateway, wire: Wire, agent_name: str, cleanups: ExitStack) -> str:
created: Final = candidate.post(
"/v1/agents",
{
"agent_name": agent_name,
"agent_card_params": {
"protocolVersion": "0.3",
"name": agent_name,
"description": "Synthetic permission peer",
"version": "1.0.0",
"url": wire.url + "/",
"capabilities": {"streaming": False},
"defaultInputModes": ["text"],
"defaultOutputModes": ["text"],
"skills": [],
},
"litellm_params": {"make_public": False},
},
)
agent_id: Final = string_value(created["agent_id"])
cleanups.callback(_delete_agent, candidate, agent_id)
return agent_id
def _delete_agent(candidate: Gateway, agent_id: str) -> None:
deleted: Final = candidate.request("DELETE", f"/v1/agents/{agent_id}")
assert deleted.status_code == 200, deleted.text
assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (agent_id,)) == []
def _card(candidate: Gateway, agent_id: str, key: str) -> httpx.Response:
return candidate.request(
"GET",
f"/a2a/{agent_id}/.well-known/agent-card.json",
key=key,
headers={"Connection": "close"},
)
def _send(candidate: Gateway, agent_id: str, key: str, marker: str) -> httpx.Response:
return candidate.request(
"POST",
f"/a2a/{agent_id}",
{
"jsonrpc": "2.0",
"id": marker,
"method": "message/send",
"params": {
"message": {
"kind": "message",
"role": "user",
"messageId": marker + "-in",
"parts": [{"kind": "text", "text": "synthetic ping"}],
}
},
},
key=key,
headers={"Connection": "close"},
)
def _chat_completion(candidate: Gateway, agent_name: str, key: str) -> httpx.Response:
return candidate.request(
"POST",
"/v1/chat/completions",
{
"model": f"a2a/{agent_name}",
"messages": [{"role": "user", "content": "synthetic ping"}],
},
key=key,
headers={"Connection": "close"},
)
def _model_names(candidate: Gateway, path: str, key: str) -> frozenset[str]:
response: Final = candidate.request("GET", path, key=key, headers={"Connection": "close"})
assert response.status_code == 200, response.text
payload: Final = response.json()
raw_entries: Final = payload if isinstance(payload, list) else payload["data"]
entries: Final = tuple(object_value(entry) for entry in raw_entries)
if path == "/model_group/info":
return frozenset(string_value(entry["model_group"]) for entry in entries)
return frozenset(string_value(entry["model_name"]) for entry in entries)
def _listed_agent_ids(candidate: Gateway, key: str) -> frozenset[str]:
response: Final = candidate.request("GET", "/v1/agents", key=key, headers={"Connection": "close"})
assert response.status_code == 200, response.text
return frozenset(string_value(object_value(agent)["agent_id"]) for agent in response.json())
def _assert_successful_send(response: httpx.Response, marker: str) -> None:
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["jsonrpc"] == "2.0" and body["id"] == marker and "error" not in body, response.text
assert body["result"]["messageId"] == marker + "-in-out", response.text
assert body["result"]["parts"] == [{"kind": "text", "text": "synthetic pong"}], response.text
def _assert_denial(response: httpx.Response, agent_id: str) -> None:
assert response.status_code == 403, response.text
assert response.json() == {
"detail": f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access."
}, response.text
@pytest.mark.timeout(180)
def test_agent_access_default_deny_denies_ungranted_key_but_keeps_explicit_grants(
gateway: Gateway, tmp_path: Path
) -> None:
agent_name: Final = "least-privilege-" + uuid.uuid4().hex
v2_model: Final = "listing-" + agent_name
config: Final = _write_access_config(
tmp_path, default_deny=True, model_name=v2_model, api_base=gateway.upstream_url + "/v1"
)
def upstream(request: Request) -> Reply:
return _peer_response(request, wire, agent_name)
with (
wire_server(upstream) as wire,
owned_proxy(gateway, tmp_path, {}, config=config, workers=2) as candidate,
candidate.scenario() as scenario,
):
agent_id: Final = _create_agent(candidate, wire, agent_name, scenario.cleanups)
ungranted_team: Final = scenario.team()
key_a: Final = scenario.key(team_id=ungranted_team)
key_b: Final = scenario.key(object_permission={"agents": [agent_id]})
granted_team: Final = scenario.team(object_permission={"agents": [agent_id]})
key_c: Final = scenario.key(team_id=granted_team)
wire.drain()
denied_cards: Final = tuple(_card(candidate, agent_id, key_a) for _ in range(8))
_assert_denial(denied_cards[0], agent_id)
assert all(response.json() == denied_cards[0].json() for response in denied_cards), [
response.text for response in denied_cards
]
denied_sends: Final = tuple(
_send(candidate, agent_id, key_a, f"denied-{index}-{uuid.uuid4().hex}") for index in range(8)
)
_assert_denial(denied_sends[0], agent_id)
assert all(response.json() == denied_sends[0].json() for response in denied_sends), [
response.text for response in denied_sends
]
denied_completion: Final = _chat_completion(candidate, agent_name, key_a)
assert denied_completion.status_code == 403, denied_completion.text
assert denied_completion.json() == {
"error": {
"message": f"Agent '{agent_name}' is not allowed for your key/team. Contact proxy admin for access.",
"type": "permission_error",
"param": None,
"code": "403",
}
}, denied_completion.text
denied_agents: Final = tuple(
candidate.request("GET", f"/v1/agents/{agent_id}", key=key_a, headers={"Connection": "close"})
for _ in range(8)
)
_assert_denial(denied_agents[0], agent_id)
assert all(response.json() == denied_agents[0].json() for response in denied_agents), [
response.text for response in denied_agents
]
assert v2_model in _model_names(candidate, "/v2/model/info", candidate.key)
assert _listed_agent_ids(candidate, key_a) == frozenset()
assert f"a2a/{agent_name}" not in _model_names(candidate, "/v2/model/info", key_a)
assert f"a2a/{agent_name}" not in _model_names(candidate, "/model_group/info", key_a)
assert tuple(item for item in wire.drain() if item.method == "POST") == ()
for key in (key_b, key_c, candidate.key):
cards: Final = tuple(_card(candidate, agent_id, key) for _ in range(8))
assert all(response.status_code == 200 for response in cards), [response.text for response in cards]
completion: Final = _chat_completion(candidate, agent_name, key)
assert completion.status_code == 200, completion.text
assert completion.json()["choices"][0]["message"]["content"] == "synthetic pong", completion.text
wire.drain()
for index in range(8):
marker: Final = f"granted-{index}-{uuid.uuid4().hex}"
_assert_successful_send(_send(candidate, agent_id, key, marker), marker)
posts: Final = tuple(item for item in wire.drain() if item.method == "POST")
assert len(posts) == 1, posts
if key != candidate.key:
assert agent_id in _listed_agent_ids(candidate, key)
assert f"a2a/{agent_name}" in _model_names(candidate, "/v2/model/info", key)
@pytest.mark.timeout(180)
@pytest.mark.parametrize("workers", [2, 4])
def test_agent_access_default_deny_toggles_at_runtime_through_config_api(
gateway: Gateway, tmp_path: Path, workers: int
) -> None:
agent_name: Final = "runtime-agent-" + uuid.uuid4().hex
def upstream(request: Request) -> Reply:
return _peer_response(request, wire, agent_name)
with (
wire_server(upstream) as wire,
owned_proxy(
gateway,
tmp_path,
{},
config=Path("tests/integration/proxy_config.yaml"),
workers=workers,
) as candidate,
candidate.scenario() as scenario,
):
agent_id: Final = _create_agent(candidate, wire, agent_name, scenario.cleanups)
ungranted_key: Final = scenario.key()
wire.drain()
assert _card(candidate, agent_id, ungranted_key).status_code == 200
enabled: Final = candidate.request(
"POST",
"/config/field/update",
{
"config_type": "general_settings",
"field_name": "agent_access_default_deny",
"field_value": "true",
},
)
assert enabled.status_code == 200, enabled.text
config_fields: Final = candidate.request(
"GET",
"/config/list",
params={"config_type": "general_settings"},
)
assert config_fields.status_code == 200, config_fields.text
setting: Final = next(
object_value(field)
for field in config_fields.json()
if object_value(field)["field_name"] == "agent_access_default_deny"
)
assert setting["field_type"] == "Boolean" and setting["field_value"] == "true", setting
eventually(
lambda: tuple(_card(candidate, agent_id, ungranted_key) for _ in range(8)),
lambda responses: all(response.status_code == 403 for response in responses),
seconds=40,
)
disabled: Final = candidate.request(
"POST",
"/config/field/update",
{
"config_type": "general_settings",
"field_name": "agent_access_default_deny",
"field_value": False,
},
)
assert disabled.status_code == 200, disabled.text
eventually(
lambda: tuple(_card(candidate, agent_id, ungranted_key) for _ in range(8)),
lambda responses: all(response.status_code == 200 for response in responses),
seconds=40,
)

9
tests/integration/authorization/test_key_guardrail_opt_out_chaos.py
View file

@ -1,5 +1,4 @@
import json
import socket
import threading
import time
import uuid
@ -15,6 +14,7 @@ from integration._support.client import Gateway, eventually, string_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy, owned_proxy_process
from integration._support.timing import assert_cut_short, timed
from integration._support.ports import allocate_port
from integration.authorization._guardrail_opt_out import (
chat,
guardrail_config,
@ -34,14 +34,9 @@ class _GuardrailSink:
self._action: Final = action
self._server: ThreadingHTTPServer | None = None
self._thread: threading.Thread | None = None
self._port: Final = self._claim_port()
self._port: Final = allocate_port()
self.start()
def _claim_port(self) -> int:
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
return probe.getsockname()[1]
@property
def url(self) -> str:
return f"http://127.0.0.1:{self._port}"

6
tests/integration/authorization/test_responses_compact_id_ownership_chaos.py
View file

@ -15,6 +15,7 @@ import pytest
from integration._support.client import Gateway, eventually
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Wire, wire_server
from integration._support.ports import allocate_port
from integration.authorization._responses_compact import (
ALIASES,
OTHER_TEAM_DETAIL,
@ -158,9 +159,9 @@ def _settled(target: _Target) -> bool:
def test_upstream_outage_mid_burst_keeps_refusals_and_lands_each_owner_compact_once(gateway: Gateway) -> None:
upstream: Final = CompactUpstream()
created: Final = marker("mint")
port: Final = allocate_port()
with gateway.scenario() as scenario:
with wire_server(upstream.respond) as first:
port: Final = urlsplit(first.url).port
with wire_server(upstream.respond, port=port) as first:
actors: Final = tenants(scenario, first)
target: Final = _Target(gateway, actors, mint(gateway, actors.owner, actors.model, created), vendor_id())
upstreamed(first)
@ -168,7 +169,6 @@ def test_upstream_outage_mid_burst_keeps_refusals_and_lands_each_owner_compact_o
before_received: Final = upstreamed(first)
during: Final = _send_all(target, _calls("during", 16))
alive: Final = gateway.request("GET", "/health/liveliness")
assert port is not None
with wire_server(upstream.respond, port=port) as second:
eventually(lambda: _settled(target), lambda settled: settled, seconds=60)
after: Final = _send_all(target, _calls("after", 16))

10
tests/integration/compatibility/test_a2a_protocol_version_override.py
View file

@ -10,13 +10,13 @@ from __future__ import annotations
import json
import os
import socket
import time
from typing import Final
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.ports import allocate_port
from integration.compatibility._a2a_peers import (
HINT,
KIND_ERROR,
@ -284,14 +284,8 @@ def test_registered_agent_persists_protocol_version(gateway: Gateway) -> None:
_assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
def _closed_local_port() -> int:
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
return int(probe.getsockname()[1])
def test_pinned_agent_survives_localhost_retry(gateway: Gateway) -> None:
dead: Final = f"http://localhost:{_closed_local_port()}/a2a/unreachable"
dead: Final = f"http://localhost:{allocate_port()}/a2a/unreachable"
with peer("lgretry", "langgraph", card_url=dead) as agent, gateway.scenario() as scenario:
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED)
_send_ok(gateway, identity, agent.marker)

12
tests/integration/compatibility/test_granian_admission_saturation.py
View file

@ -1,6 +1,5 @@
import asyncio
import os
import socket
import subprocess
import sys
import time
@ -12,22 +11,17 @@ import pytest
from tests._master_key import MASTER_KEY
from tests.integration._support.process import stop_child
from tests.integration._support.ports import allocate_port
MAX_IN_FLIGHT: Final = 8
MAX_QUEUED: Final = 8
REQUESTS: Final = 200
def _free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
listener.bind(("127.0.0.1", 0))
return int(listener.getsockname()[1])
@pytest.mark.asyncio
async def test_granian_admission_control_saturation(tmp_path: Path) -> None:
fake_port: Final = _free_port()
proxy_port: Final = _free_port()
fake_port: Final = allocate_port()
proxy_port: Final = allocate_port()
fake_script: Final = Path(__file__).parents[2] / "_fake_openai_endpoint_server.py"
config_path: Final = tmp_path / "saturation_config.yaml"
proxy_env: Final = {**{k: v for k, v in os.environ.items() if k != "DATABASE_URL"}, "LITELLM_MASTER_KEY": MASTER_KEY}

6
tests/integration/compatibility/test_missing_body_param_status.py
View file

@ -3,7 +3,6 @@ from __future__ import annotations
import asyncio
import json
import os
import socket
import subprocess
import sys
import uuid
@ -24,6 +23,7 @@ from integration._support.database import read_rows
from integration._support.process import owned_proxy_process, paused, stop_child
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.ports import allocate_port
from openai import APIStatusError, AsyncOpenAI, BadRequestError, OpenAI
from openai.types import Video
from openai.types.eval_create_response import EvalCreateResponse as _Eval
@ -1504,9 +1504,7 @@ def _healthy(url: str) -> int:
@contextmanager
def _upstream(directory: Path) -> Iterator[tuple[subprocess.Popen[bytes], str]]:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
port: Final = int(reserve.getsockname()[1])
port: Final = allocate_port()
root: Final = Path(__file__).resolve().parents[2]
url: Final = f"http://127.0.0.1:{port}"
with (directory / "upstream.log").open("w") as log:

2
tests/integration/conftest.py
View file

@ -25,6 +25,7 @@ from tests.integration._support.manifest import OWNED_DIRECTORIES
from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups
from tests.integration._support.provider import SharedProvider, shared_provider
from tests.integration._support.routing import RoutingPlugin
from tests.integration._support.sdk_models import build_sdk_models
from tests.integration._support.workers import gateway_profile, start_worker_stacks, worker_stack
from tests.integration.run import GITHUB_FILES
@ -70,6 +71,7 @@ def pytest_configure(config: pytest.Config) -> None:
config.stash[REPORTS] = []
config.stash[SHARD_FILES] = _shard_files(SHARD_PATH.validate_python(config.getoption("integration_shard_files")))
config.pluginmanager.register(IntegrationReportPlugin(config))
build_sdk_models()
if os.environ.get("INTEGRATION_ROUTING"):
config.pluginmanager.register(RoutingPlugin(config))

8
tests/integration/database/test_harness_lifecycle.py
View file

@ -44,6 +44,7 @@ from tests.integration._support.database_relay import (
from tests.integration._support.runtime import configure_executor
from tests.integration._support.sharding import ShardExecution, shard_errors
from tests.integration._support.wire import Reply, wire_server
from tests.integration._support.ports import allocate_port
from tests.integration.run import _without_provider_api_keys
@ -143,10 +144,11 @@ async def _unused_peer(reader: asyncio.StreamReader, writer: asyncio.StreamWrite
def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> None:
port: Final = allocate_port()
with socket.socket() as occupied:
occupied.bind(("127.0.0.1", 0))
occupied.bind(("127.0.0.1", port))
occupied.listen()
relay: Final = LoopbackServer(occupied.getsockname()[1], _unused_peer)
relay: Final = LoopbackServer(port, _unused_peer)
with pytest.raises(OSError, match=os.strerror(errno.EADDRINUSE).lower()):
relay.start()
relay.stop()
@ -156,7 +158,7 @@ def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> Non
def test_upstream_readiness_failure_reaps_the_started_process(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
monkeypatch.setattr(process_support, "_UPSTREAM_READY_SECONDS", 0)
children_before: Final = frozenset(child.pid for child in psutil.Process().children(recursive=True))
slot: Final = process_support.UpstreamSlot(tmp_path, process_support._free_port(), process_support._harness_root())
slot: Final = process_support.UpstreamSlot(tmp_path, allocate_port(), process_support._harness_root())
try:
with pytest.raises(AssertionError, match="readiness deadline exceeded"):
slot.start()

6
tests/integration/database/test_v1_migration_error_logging.py
View file

@ -25,7 +25,7 @@ from psycopg import sql
from psycopg.types.json import Jsonb
from tests.integration._support.client import eventually
from tests.integration._support.process import _free_port as free_port
from tests.integration._support.ports import allocate_port
from tests.integration._support.process import stop_process_group
REPO_ROOT: Final = Path(__file__).resolve().parents[3]
@ -94,7 +94,7 @@ def _unreachable_database_url(database_url: str) -> str:
parsed.path.lstrip("/"),
encode_password=False,
)
return _replace_port(url, free_port())
return _replace_port(url, allocate_port())
@contextmanager
@ -388,7 +388,7 @@ def test_unreachable_database_recovers_after_postgres_forwarder_starts(
port: Final = admin.port
assert hostname is not None and port is not None
target_host: Final = "127.0.0.1" if hostname == "localhost" else hostname
forwarding_port: Final = free_port()
forwarding_port: Final = allocate_port()
forwarded_url: Final = _replace_port(database_url, forwarding_port, "127.0.0.1")
command, environment = _migration_invocation(forwarded_url, tmp_path, {})
output_path: Final = _migration_log_path(request.node.name, tmp_path)

6
tests/integration/management/test_vector_store_file_list_managed_ids.py
View file

@ -22,6 +22,7 @@ from integration._support.client import Gateway, Scenario, eventually, object_va
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 integration._support.ports import allocate_port
from openai import APIStatusError, AsyncOpenAI, OpenAI
from pydantic import JsonValue, TypeAdapter
@ -1202,7 +1203,8 @@ def test_provider_outage_mid_burst_fails_loudly_and_mapping_resumes_after_recove
provider_a: Final = _provider_file_id(bearer, "a.txt")
respond: Final = _provider(store, _constant_listing(store, provider_a))
with gateway.scenario() as scenario:
with wire_server(respond) as wire:
port: Final = allocate_port()
with wire_server(respond, port=port) as wire:
model: Final = scenario.model(api_base=wire.url + "/v1", api_key=bearer)
_wait_until_every_worker_serves(gateway, model)
wire.drain()
@ -1212,8 +1214,6 @@ def test_provider_outage_mid_burst_fails_loudly_and_mapping_resumes_after_recove
served: Final = _burst(gateway, store, member.key, model, 40)
assert [_ids(_listed(response)) for response in served] == [(managed_a,)] * 40
assert sum(1 for request in wire.drain() if request.method == "GET") == 40
port: Final = urlsplit(wire.url).port
assert port is not None
failed: Final = _burst(gateway, store, member.key, model, 20)
assert [response.status_code for response in failed] == [500] * 20, [r.text for r in failed[:3]]
for response in failed:

10
tests/integration/mcp/test_proxy_mcp_e2e.py
View file

@ -45,6 +45,7 @@ from tests.integration._support.process import stop_child
from tests.integration._support.timing import assert_enforced_within
from tests.integration._support.wire import Reply, Wire, wire_server
from tests.integration._support.wire import Request as WireRequest
from tests.integration._support.ports import allocate_port
CONFIG_TEMPLATE_PATH = Path("tests/integration/mcp/_config_mcp_e2e.yaml")
MCP_SERVER_SCRIPT = Path("tests/integration/mcp/_math_server.py")
@ -294,8 +295,9 @@ def _start_proxy_server(
sock: Final = socket_factory()
startup.callback(sock.close)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(("127.0.0.1", 0))
host, port = sock.getsockname()
host: Final = "127.0.0.1"
port: Final = allocate_port()
sock.bind((host, port))
config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning", lifespan="off")
server = uvicorn.Server(config)
@ -369,9 +371,7 @@ def test_proxy_startup_failure_closes_acquired_resources(tmp_path: Path) -> None
@contextmanager
def _math_http_server(offset: int) -> typing.Iterator[str]:
host = "127.0.0.1"
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind((host, 0))
_, port = sock.getsockname()
port: Final = allocate_port()
with tempfile.TemporaryFile() as server_log:
process = subprocess.Popen(

6
tests/integration/observability/test_azure_content_safety_dispatch_resilience.py
View file

@ -1,5 +1,4 @@
import concurrent.futures
import socket
import threading
import uuid
from contextlib import ExitStack
@ -13,6 +12,7 @@ from integration._support.client import Gateway, Scenario, eventually
from integration._support.database import read_rows
from integration._support.redis_process import owned_redis
from integration._support.wire import Wire, wire_server
from integration._support.ports import allocate_port
from integration.observability.azure_dispatch_support import (
ATTACK_MARKER,
AzureBehavior,
@ -251,9 +251,7 @@ def test_x1_azure_server_stop_restart_during_mixed_burst(
with ExitStack() as resources:
redis: Final = resources.enter_context(owned_redis(tmp_path))
provider: Final = resources.enter_context(wire_server(provider_handler))
with socket.socket() as reservation:
reservation.bind(("127.0.0.1", 0))
azure_port: Final = int(reservation.getsockname()[1])
azure_port: Final = allocate_port()
first_azure_lifetime: Final = ExitStack()
try:
azure: Final = first_azure_lifetime.enter_context(wire_server(azure_handler(), port=azure_port))

7
tests/integration/observability/test_guardrail_effects.py
View file

@ -1,7 +1,6 @@
import asyncio
import json
import re
import socket
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
@ -22,6 +21,7 @@ from integration._support.database import read_rows
from integration._support.mcp import mcp_peer, register_mcp, tool_names
from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.ports import allocate_port
from openai import AsyncOpenAI, OpenAI
from pydantic import BaseModel, JsonValue
@ -1285,10 +1285,7 @@ def _blocked_stream_events(text: str) -> tuple[dict[str, object], ...]:
def _dead_api_base() -> str:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
port: Final = reserve.getsockname()[1]
return f"http://127.0.0.1:{port}/v1"
return f"http://127.0.0.1:{allocate_port()}/v1"
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_openai_sdk_streams_typed_message")

12
tests/integration/observability/test_guardrail_stream_scope_chaos.py
View file

@ -3,7 +3,6 @@ from __future__ import annotations
import json
import os
import shutil
import socket
import subprocess
import uuid
from collections.abc import Iterator, Mapping, Sequence
@ -23,6 +22,7 @@ from integration._support.client import Gateway, JsonValue, eventually
from integration._support.database import read_rows
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.ports import allocate_port
from pydantic import TypeAdapter
ChaosEndpoint: TypeAlias = Literal["chat", "messages", "responses"]
@ -455,12 +455,6 @@ def _owned_proxy(
yield owned
def _free_port() -> int:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
return int(reserve.getsockname()[1])
def _assert_down_wave(
plans: Sequence[CallPlan],
responses: Sequence[httpx.Response],
@ -485,7 +479,7 @@ def _assert_down_wave(
def test_h1_sink_outage_keeps_scope_isolated_through_recovery(rig: ChaosRig, tmp_path: Path) -> None:
port_a: Final = _free_port()
port_a: Final = allocate_port()
started: Final = Event()
unavailable: Final = Event()
release: Final = Event()
@ -763,7 +757,7 @@ def _owned_postgres(directory: Path) -> Iterator[PostgresCluster]:
)
assert docker_info.returncode == 0, docker_info.stderr
directory.mkdir(parents=True, exist_ok=True)
port: Final = _free_port()
port: Final = allocate_port()
container_name: Final = f"litellm-stream-scope-h5-{uuid.uuid4().hex}"
password: Final = uuid.uuid4().hex
created: Final = subprocess.run(

4
tests/integration/observability/test_langtrace_delivery.py
View file

@ -21,6 +21,7 @@ from integration._support.client import Gateway, eventually
from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process
from integration._support.timing import assert_not_delayed_by, timed
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.ports import allocate_port
from openai import AsyncOpenAI, OpenAI
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
from opentelemetry.proto.trace.v1.trace_pb2 import Span, Status
@ -595,8 +596,7 @@ def test_langtrace_two_workers_deliver_every_burst_span_exactly_once(gateway: Ga
def test_langtrace_sink_outage_mid_burst_recovers_on_the_same_port(gateway: Gateway, tmp_path: Path) -> None:
key: Final = "synthetic-langtrace-key-" + uuid.uuid4().hex
with wire_server(_accepted) as probe:
port: Final = int(probe.url.rsplit(":", 1)[1])
port: Final = allocate_port()
host: Final = f"http://127.0.0.1:{port}"
with wire_server(_upstream) as provider:
overrides: Final = {"LANGTRACE_API_KEY": key, "LANGTRACE_API_HOST": host, "OTEL_BSP_SCHEDULE_DELAY": "300"}

6
tests/integration/observability/test_logging_only_scope_chaos.py
View file

@ -1,6 +1,5 @@
from __future__ import annotations
import socket
import threading
import time
import uuid
@ -46,6 +45,7 @@ from integration._support.client import Gateway, eventually
from integration._support.process import owned_proxy, owned_proxy_process
from integration._support.timing import Timed, assert_not_delayed_by, timed
from integration._support.wire import Reply, Request
from integration._support.ports import allocate_port
from openai import APIConnectionError as OpenAIAPIConnectionError
from pydantic import JsonValue
@ -60,9 +60,7 @@ def test_K1_policy_edge_restart_mid_burst_keeps_output_observation_fail_open(gat
def policy(_request: Request) -> Reply:
return Reply(body=b'{"action":"NONE"}')
with socket.socket() as reservation:
reservation.bind(("127.0.0.1", 0))
policy_port: Final = reservation.getsockname()[1]
policy_port: Final = allocate_port()
with gateway.scenario() as scenario:
deployments: Final = _chaos_models(scenario, marker)
models: Final = tuple(deployment.model_name for deployment in deployments)

5
tests/integration/observability/test_passthrough_upstream_error_chaos.py
View file

@ -15,6 +15,7 @@ from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from integration._support.ports import allocate_port
from pydantic import JsonValue
_GENERATE_CONTENT: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}
@ -120,8 +121,8 @@ async def _fire_burst(
async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gateway: Gateway, tmp_path: Path) -> None:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
path: Final = tmp_path / "chaos-outage.yaml"
with wire_server(_chaos_reply) as wire:
port: Final = int(wire.url.rsplit(":", 1)[1])
port: Final = allocate_port()
with wire_server(_chaos_reply, port=port) as wire:
config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"}
path.write_text(yaml.safe_dump(config))
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:

7
tests/integration/observability/test_straiker_v3_platform.py
View file

@ -10,7 +10,6 @@ from __future__ import annotations
import hashlib
import itertools
import json
import socket
import threading
import uuid
from collections.abc import Callable, Iterator
@ -30,6 +29,7 @@ from integration._support.client import Gateway, eventually, gateway_from_enviro
from integration._support.database import read_rows
from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy, owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from integration._support.ports import allocate_port
V3_KEY: Final = "sk_agt_synthetic_integration_key"
V1_KEY: Final = "synthetic-v1-collection-key"
@ -426,10 +426,7 @@ class Rig:
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
root: Final = tmp_path_factory.mktemp("straiker")
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
port: Final = reserve.getsockname()[1]
sink: Final = Sink(port)
sink: Final = Sink(allocate_port())
sink.start()
with gateway_from_environment() as gateway, wire_server(_provider) as provider:
config: Final = _rig_config(sink.url, root)

7
tests/integration/providers/_assistants_route_errors_support.py
View file

@ -3,7 +3,6 @@ from __future__ import annotations
import asyncio
import json
import re
import socket
import threading
import uuid
from collections.abc import AsyncIterator, Callable, Generator, Mapping
@ -84,12 +83,6 @@ def new_marker() -> str:
return f"m{uuid.uuid4().hex}"
def free_port() -> int:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
return reserve.getsockname()[1]
def proxy_path(route: Route, marker: str) -> str:
return route.path.format(marker=marker)

7
tests/integration/providers/_coralbricks.py
View file

@ -2,7 +2,6 @@ from __future__ import annotations
import json
import re
import socket
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
@ -380,12 +379,6 @@ def raw_deployment(scenario: Scenario, model_name: str, litellm_params: Mapping[
return model_name
def closed_port() -> int:
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
return int(probe.getsockname()[1])
def provider_calls(wire: Wire) -> tuple[Request, ...]:
return tuple(request for request in wire.drain() if not is_readiness_probe(request))

7
tests/integration/providers/_vertex_batch_rows.py
View file

@ -1,5 +1,4 @@
import json
import socket
import threading
from collections.abc import Callable
from email.message import Message
@ -132,12 +131,6 @@ def stored_custom_ids(request: Request) -> tuple[str, ...]:
)
def closed_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
probe.bind(("127.0.0.1", 0))
return int(probe.getsockname()[1])
def multipart_file(request: Request) -> bytes:
envelope: Final = f"content-type: {request.headers['content-type']}\r\n\r\n".encode() + request.body
parts: Final[tuple[Message, ...]] = tuple(BytesParser(policy=HTTP).parsebytes(envelope).iter_parts())

10
tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py
View file

@ -1,6 +1,5 @@
import asyncio
import json
import socket
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
@ -18,6 +17,7 @@ import pytest
from anthropic.types import MessageParam
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.ports import allocate_port
from integration.providers._count_tokens_system_lift import (
ANTHROPIC_VERSION,
ASSISTANT,
@ -131,12 +131,6 @@ def _listening(deployment: _Deployment, count: Callable[[Request], Reply] = _cou
return wire_server(_peer(deployment.provider, count), port=deployment.port)
def _reserved_port() -> int:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
return _CLIENT_ADDRESS.validate_python(reserve.getsockname())[1]
def _cmdline(process: psutil.Process) -> tuple[str, ...]:
try:
return _ARGUMENTS.validate_python(process.cmdline())
@ -211,7 +205,7 @@ def _settled_on_every_worker(gateway: Gateway, body: Mapping[str, JsonValue]) ->
def _deploy(gateway: Gateway, scenario: Scenario, provider: _Provider) -> _Deployment:
port: Final = _reserved_port()
port: Final = allocate_port()
model: Final = scenario.model(
model_info=None, model=f"{provider.prefix}/{_MODEL}", api_base=f"http://127.0.0.1:{port}", api_key=_API_KEY
)

12
tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py
View file

@ -33,6 +33,7 @@ from integration._support.client import Gateway, eventually, held
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.timing import assert_not_delayed_by
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.ports import allocate_port
pytestmark: Final = [
pytest.mark.timeout(2 * graceful_stop_seconds() + 120),
@ -180,12 +181,6 @@ def _open_upstream_connections(pid: int, upstream: str) -> int:
)
def _port(wire: Wire) -> int:
port: Final = urlsplit(wire.url).port
assert port is not None, wire.url
return port
def _accept(_request: Request) -> Reply:
return Reply()
@ -200,8 +195,9 @@ async def _timed_chat(gateway: Gateway) -> float:
async def test_sink_outage_mid_burst_delivers_every_refusal_event_once(gateway: Gateway, tmp_path: Path) -> None:
calls: Final = _calls(30)
markers: Final = frozenset(call.marker for call in calls)
sink_port: Final = allocate_port()
with wire_server(respond) as upstream, ExitStack() as first_sink:
sink: Final = first_sink.enter_context(wire_server(_accept))
sink: Final = first_sink.enter_context(wire_server(_accept, port=sink_port))
first: Final = Ledger(sink)
config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=6)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
@ -213,7 +209,7 @@ async def test_sink_outage_mid_burst_delivers_every_refusal_event_once(gateway:
first.events()
first_sink.close()
served_during: Final = await _burst(owned.gateway, calls[10:20])
with wire_server(_accept, port=_port(sink)) as restarted:
with wire_server(_accept, port=sink_port) as restarted:
second: Final = Ledger(restarted)
served_after: Final = await _burst(owned.gateway, calls[20:])
for item in (*served_first, *served_during, *served_after):

Some files were not shown because too many files have changed in this diff Show more