feat(otel v2): excluded_services opt-out for datastore spans on tenant destinations (#43278)

* feat(otel): excluded_services opt-out for datastore spans on tenant destinations

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): keep upstream support unchanged

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel): excluded_services resolves from the otel callback config only

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel): name the otel callback logger so excluded_services owner lookup matches

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): assert no aux datastore traces reach the tenant sink

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): read bogus-start proxy log from the results dir

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): read only this invocation's bogus-start proxy log

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): assert operator kept db spans over the whole recorded window

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): split operator db-span asserts by trace scope

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel): build the otel logger after preset callbacks and validate the exclusion env at boot

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel): tolerate a bogus exclusion env when callback config wins

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel): keep bogus exclusion env fatal when a preset parses it

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel): hoist the preset check out of the callback loop

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(otel): use a rule-scoped pyright suppression

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel): log and drop unknown excluded_services instead of failing boot

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): wait for the operator spend-writer span before checking the tenant for postgres

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel v2): leave callback init and boot untouched when excluded_services is unset

Read callback_settings.otel.excluded_services directly instead of making the otel callback build its own logger, and drop the new boot-time parse of callback_settings.otel, so a proxy without the setting behaves exactly as on main

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel v2): normalize callback_settings excluded_services without rereading OTel env vars

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(otel v2): log and ignore malformed excluded_services instead of failing startup

Lowercase and trim names, drop non-string items, and add an integration matrix over endpoints, clients, cache hits, destination outages and setting shapes

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(otel v2): cover failed upstream calls in the excluded_services matrix

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(otel v2): pin operator Langfuse credentials in preset-only excluded_services tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: mrinal <mrinal@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-09-30 12:20:27 -07:00 • committed by GitHub
parent e78845afcf
commit 50f5cc9bbb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 2018 additions and 21 deletions

View file

@ -213,6 +213,15 @@ nothing here imports outside it:
`config.yaml` — the latter reach the config through the logger's constructor
kwargs. `baggage_team_metadata_keys` is empty by default, so none of a team's
free-form metadata is promoted until each sub-key is explicitly allowlisted.
`excluded_services` withholds datastore spans from key/team `callback_vars`
destinations while the operator's own exporters keep them: set
`LITELLM_OTEL_EXCLUDED_SERVICES` (comma-separated) or `excluded_services`
(a YAML list) under `callback_settings.otel`, naming the datastore services
to withhold (`redis`, `postgres`, `batch_write_to_db`, `redis_*`, or their
`db.system.name` spellings `redis` / `postgresql`). Unknown names are logged
as an error and ignored. A span is withheld when its `db.system.name` /
`db.system` attribute is in the set, so request root, auth, guardrail and
model spans can never be excluded.
- [`baggage.py`](./model/baggage.py) — the single definition of which request-identity
values are promoted into Baggage (so child spans inherit them) and under which
attribute keys.

View file

@ -30,7 +30,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.otel.emitter import SpanEmitter, stamp_error
from litellm.integrations.otel.mappers import resolve_mappers
from litellm.integrations.otel.model.baggage import promoted_baggage
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
from litellm.integrations.otel.model.config import OpenTelemetryV2Config, excluded_db_systems_from
from litellm.integrations.otel.model.metadata import (
LLMCallEvent,
RequestIdentity,
@ -898,12 +898,29 @@ def publish_global_otel_v2_provider(
"""
global _published_v2_provider
logger: Final = select_global_otel_v2_logger(in_memory_loggers, registered=registered)
attach_tenant_fan_out(logger.tracer_provider, *_v2_configs(in_memory_loggers, logger))
attach_tenant_fan_out(
logger.tracer_provider,
*_v2_configs(in_memory_loggers, logger),
excluded_db_systems=_excluded_db_systems(logger),
)
set_global_provider(logger.tracer_provider)
_published_v2_provider = logger.tracer_provider # rebind-ok: startup records the one provider carrying the fan-out
return logger
def _excluded_db_systems(logger: "OpenTelemetryV2") -> frozenset[str]:
"""The datastore services withheld from tenant destinations.
``callback_settings.otel.excluded_services`` wins over the env var whichever
logger got published: with ``callbacks: [langfuse_otel, otel]`` the ``otel``
callback folds into the preset, whose config is env-only.
"""
configured: Final = litellm.callback_settings.get("otel", {}).get("excluded_services")
if configured is None:
return logger.config.excluded_services
return excluded_db_systems_from(configured)
def _v2_configs(in_memory_loggers: Sequence[object], logger: "OpenTelemetryV2") -> tuple[OpenTelemetryV2Config, ...]:
"""Every v2 logger's config, the published logger's first.
@ -963,7 +980,11 @@ def fan_out_provider() -> ApiTracerProvider:
return published
logger: Final = _registered_v2_logger()
if logger is not None:
attach_tenant_fan_out(logger.tracer_provider, logger.config)
attach_tenant_fan_out(
logger.tracer_provider,
logger.config,
excluded_db_systems=_excluded_db_systems(logger),
)
return logger.tracer_provider
return get_tracer_provider()

View file

@ -4,14 +4,16 @@ from enum import Enum
from functools import lru_cache
from typing import Annotated, Any, Final
from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
from pydantic import AliasChoices, BaseModel, Field, TypeAdapter, ValidationError, field_validator, model_validator
from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict
from litellm._logging import verbose_logger
from litellm.integrations.otel.model.baggage import (
BAGGAGE_PROMOTED_KEYS,
DEFAULT_BAGGAGE_METADATA_KEYS,
DEFAULT_BAGGAGE_TEAM_METADATA_KEYS,
)
from litellm.integrations.otel.model.spans import POSTGRESQL, db_system
from litellm.types.utils import OtelSpanScope
#: Master feature-flag env var. The logger is inert until this is truthy.
@ -174,6 +176,19 @@ class OpenTelemetryV2Config(BaseSettings):
"key/team destinations are not affected."
),
)
excluded_services: Annotated[frozenset[str], NoDecode] = Field(
default_factory=frozenset,
validation_alias=AliasChoices("excluded_services", "LITELLM_OTEL_EXCLUDED_SERVICES"),
description=(
"Datastore services whose spans are withheld from key/team ``callback_vars`` "
"OTel destinations (the operator's own exporters still receive them). Accepted "
"values are the datastore ``ServiceTypes`` names (``redis``, ``postgres``, "
"``batch_write_to_db``, ``redis_*``) or their ``db.system.name`` spellings "
"(``redis``, ``postgresql``); stored normalized to ``db.system.name`` values. "
"Configure via the ``LITELLM_OTEL_EXCLUDED_SERVICES`` env var (comma-separated) "
"or ``callback_settings.otel.excluded_services`` in config.yaml (a YAML list)."
),
)
# ----- explicit multi-destination / vocabulary configuration ------------ #
@ -284,6 +299,11 @@ class OpenTelemetryV2Config(BaseSettings):
return [item.strip() for item in value.split(",") if item.strip()]
return value
@field_validator("excluded_services", mode="before")
@classmethod
def _read_excluded_services(cls, value: object) -> frozenset[str]:
return excluded_service_names(value)
@model_validator(mode="after")
def _normalize(self) -> "OpenTelemetryV2Config":
# An endpoint with the default exporter kind implies OTLP/HTTP.
@ -316,6 +336,7 @@ class OpenTelemetryV2Config(BaseSettings):
if self.legacy_compat and "legacy" not in names:
names.append("legacy")
self.mapper_names = names
self.excluded_services = _normalize_excluded_services(self.excluded_services)
return self
@property
@ -334,3 +355,55 @@ class OpenTelemetryV2Config(BaseSettings):
@classmethod
def from_env(cls) -> "OpenTelemetryV2Config":
return cls()
_EXCLUDED_SERVICES_INPUT: Final[TypeAdapter[str | tuple[object, ...]]] = TypeAdapter(str | tuple[object, ...])
def excluded_db_systems_from(value: object) -> frozenset[str]:
"""Normalize a raw ``excluded_services`` value without building a settings model that rereads the env"""
return _normalize_excluded_services(excluded_service_names(value))
def excluded_service_names(value: object) -> frozenset[str]:
"""Read a YAML list or comma-separated string of service names, logging and dropping unusable input
so a malformed value cannot stop the OTel logger from being built"""
if value is None:
return frozenset()
try:
parsed: Final = _EXCLUDED_SERVICES_INPUT.validate_python(value)
except ValidationError:
verbose_logger.error("excluded_services must be a list or comma-separated string; %r ignored", value)
return frozenset()
items: Final = tuple(parsed.split(",")) if isinstance(parsed, str) else parsed
return frozenset(name for item in items if (name := _service_name(item)))
def _service_name(item: object) -> str:
if not isinstance(item, str):
verbose_logger.error("excluded_services must be a list of service names; %r ignored", item)
return ""
return item.strip().lower()
def _normalize_excluded_services(services: frozenset[str]) -> frozenset[str]:
"""Fold each accepted spelling to its ``db.system.name`` value.
``postgres`` and ``postgresql`` name the same system, as do every
``ServiceTypes`` member that ``db_system`` maps. Anything else means the
operator pointed the setting at a span family it cannot cover; those names
are logged and dropped so a typo cannot take the proxy down.
"""
resolved: Final = frozenset(
system for service in services if (system := _db_system_for_excluded_service(service)) is not None
)
return resolved
def _db_system_for_excluded_service(service: str) -> str | None:
resolved: Final = db_system(service) if service != POSTGRESQL else POSTGRESQL
if resolved is None:
verbose_logger.error(
"excluded_services: %r is not a datastore service; ignored. Allowed: postgres, redis", service
)
return resolved

View file

@ -418,6 +418,13 @@ def _is_database_span(attributes: Mapping[str, AttributeValue]) -> bool:
return any(key in attributes for key in _DB_SYSTEM_KEYS)
def _is_excluded_database_span(attributes: Mapping[str, AttributeValue], excluded: frozenset[str]) -> bool:
if not excluded:
return False
system: Final = attributes.get(DB.SYSTEM_NAME) or attributes.get(DB.SYSTEM_LEGACY)
return isinstance(system, str) and system in excluded
def _is_tenant_owned_span(attributes: Mapping[str, AttributeValue]) -> bool:
return any(key in attributes for key in _TENANT_OWNED_KEYS)
@ -549,10 +556,12 @@ class TenantFanOutSpanProcessor(SpanProcessor):
processor_factory: 'Callable[["OtelDestination"], SpanProcessor | None] | None' = None,
shutdown_drain_seconds: float = _SHUTDOWN_DRAIN_SECONDS,
operator_sinks: 'Mapping[_SinkKey, "OtelSpanScope"]' = MappingProxyType({}),
excluded_db_systems: frozenset[str] = frozenset(),
pending_drains: int = _MAX_PENDING_DRAINS,
drain_pool: _DrainPool | None = None,
) -> None:
self._operator_sinks: Final = operator_sinks
self._excluded_db_systems: Final = excluded_db_systems
self._drain_seconds: Final = shutdown_drain_seconds
self._lock: Final = threading.Condition()
self._closed = False # guarded by ``_lock``: an unlocked read races the teardown it gates
@ -567,9 +576,12 @@ class TenantFanOutSpanProcessor(SpanProcessor):
def on_end(self, span: ReadableSpan) -> None:
suppressed: Final = suppressed_backends()
attributes: Final = span.attributes or _NO_ATTRIBUTES
for destination in request_destinations():
if self._operator_already_writes(span, destination, suppressed) or not _in_scope(
span, destination.span_scope
if (
self._operator_already_writes(span, destination, suppressed)
or not _in_scope(span, destination.span_scope)
or _is_excluded_database_span(attributes, self._excluded_db_systems)
):
continue
processor = self._acquire(destination)
@ -1155,7 +1167,9 @@ def build_tracer_provider(
_FAN_OUT_ATTACH_LOCK: Final = threading.Lock()
def attach_tenant_fan_out(provider: TracerProvider, *configs: OpenTelemetryV2Config) -> None:
def attach_tenant_fan_out(
provider: TracerProvider, *configs: OpenTelemetryV2Config, excluded_db_systems: frozenset[str] = frozenset()
) -> None:
"""Give ``provider`` the fan-out that delivers spans to key/team destinations.
Called on the one provider published as the OTel global, and idempotent so a
@ -1164,12 +1178,18 @@ def attach_tenant_fan_out(provider: TracerProvider, *configs: OpenTelemetryV2Con
so exactly one fan-out lands. ``configs`` name the operator's own exporters, one
config per v2 logger since each keeps its own provider and still writes its
account, so an additive destination pointing at any of them is delivered once
rather than twice.
rather than twice. ``excluded_db_systems`` only filters what the fan-out
delivers, never the operator's own exporters.
"""
with _FAN_OUT_ATTACH_LOCK:
if any(isinstance(processor, TenantFanOutSpanProcessor) for processor in _attached_processors(provider)):
return
provider.add_span_processor(TenantFanOutSpanProcessor(operator_sinks=operator_sink_scopes(*configs)))
provider.add_span_processor(
TenantFanOutSpanProcessor(
operator_sinks=operator_sink_scopes(*configs),
excluded_db_systems=excluded_db_systems,
)
)
def deliverable_destinations(

View file

@ -0,0 +1,528 @@
"""OTLP/HTTP trace sink: records exported spans and exposes them over a control API.
Accepts ``application/x-protobuf`` ``ExportTraceServiceRequest`` bodies and OTLP
``http/json`` bodies on any path. Tests read spans through ``recorded_spans`` and
steer the sink through ``configure``; the process can also be frozen with
``SIGSTOP``/``SIGCONT`` after reading its pid from ``/__pid``.
"""
from __future__ import annotations
import argparse
import datetime
import json
import os
import signal
import socket
import ssl
import subprocess
import sys
import threading
import time
from collections.abc import Iterator, Mapping, Sequence
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass, field
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final
from urllib.parse import urlparse
import httpx
import psutil
from pydantic import JsonValue, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
INTERNAL_MARKERS: Final = ("gen_ai.operation.name", "mcp.method.name", "litellm.guardrail_name")
class Span(TypedDict):
trace_id: ReadOnly[str]
span_id: ReadOnly[str]
parent_span_id: ReadOnly[str]
kind: ReadOnly[int]
name: ReadOnly[str]
attributes: ReadOnly[Mapping[str, JsonValue]]
resource: ReadOnly[Mapping[str, JsonValue]]
class _SpanListing(TypedDict):
next: ReadOnly[int]
spans: ReadOnly[list[Span]]
_SPAN_LISTING: Final = TypeAdapter(_SpanListing)
def _proto_spans(body: bytes) -> list[Span]:
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
from opentelemetry.proto.common.v1.common_pb2 import AnyValue
def scalar(value: AnyValue) -> JsonValue:
match value.WhichOneof("value"):
case "string_value":
return value.string_value
case "bool_value":
return value.bool_value
case "int_value":
return int(value.int_value)
case "double_value":
return value.double_value
case "bytes_value":
return value.bytes_value.decode("utf-8", errors="replace")
case "array_value":
return [scalar(item) for item in value.array_value.values]
case "kvlist_value":
return {pair.key: scalar(pair.value) for pair in value.kvlist_value.values}
case _:
return None
request: Final = ExportTraceServiceRequest()
request.ParseFromString(body)
return [
Span(
trace_id=span.trace_id.hex(),
span_id=span.span_id.hex(),
parent_span_id=span.parent_span_id.hex(),
kind=span.kind,
name=span.name,
attributes={attribute.key: scalar(attribute.value) for attribute in span.attributes},
resource={attribute.key: scalar(attribute.value) for attribute in resource.resource.attributes},
)
for resource in request.resource_spans
for scope in resource.scope_spans
for span in scope.spans
]
def _json_spans(body: bytes) -> list[Span]:
payload: Final = json.loads(body)
def scalar(value: object) -> JsonValue:
if not isinstance(value, dict):
return value if isinstance(value, (str, int, float, bool)) or value is None else str(value)
for key in ("stringValue", "intValue", "doubleValue", "boolValue", "bytesValue"):
if key in value:
return value[key]
if "arrayValue" in value:
return [scalar(item) for item in value["arrayValue"].get("values", [])]
if "kvlistValue" in value:
return {pair["key"]: scalar(pair["value"]) for pair in value["kvlistValue"].get("values", [])}
return None
return [
Span(
trace_id=str(span.get("traceId", "")),
span_id=str(span.get("spanId", "")),
parent_span_id=str(span.get("parentSpanId", "")),
kind=int(span.get("kind", 0)),
name=str(span.get("name", "")),
attributes={attribute["key"]: scalar(attribute.get("value")) for attribute in span.get("attributes", [])},
resource={
attribute["key"]: scalar(attribute.get("value"))
for attribute in resource.get("resource", {}).get("attributes", [])
},
)
for resource in payload.get("resourceSpans", [])
for scope in resource.get("scopeSpans", [])
for span in scope.get("spans", [])
]
def decode_spans(body: bytes, content_type: str) -> list[Span]:
if "protobuf" in content_type:
return _proto_spans(body)
return _json_spans(body)
def span_class(span: Span) -> str:
if span["kind"] == 2:
return "root"
if any(marker in span["attributes"] for marker in INTERNAL_MARKERS):
return "tenant"
return "internal"
def spans_for_trace(spans: tuple[Span, ...], trace_id: str) -> tuple[Span, ...]:
return tuple(span for span in spans if span["trace_id"] == trace_id)
@dataclass(slots=True)
class _State:
spans: list[Span] = field(default_factory=list)
requests: list[dict[str, JsonValue]] = field(default_factory=list)
status: int = 200
delay_seconds: float = 0.0
pause: threading.Event = field(default_factory=threading.Event)
def __post_init__(self) -> None:
self.pause.set()
class _Handler(BaseHTTPRequestHandler):
state: _State
protocol_version = "HTTP/1.1"
def _read_body(self) -> bytes:
return self.rfile.read(int(self.headers.get("content-length", "0")))
def _send_json(self, payload: object, status: int = 200) -> None:
body: Final = json.dumps(payload).encode()
self.send_response(status)
self.send_header("content-type", "application/json")
self.send_header("content-length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def _record(self) -> None:
body: Final = self._read_body()
self.state.pause.wait(timeout=120)
if self.state.delay_seconds > 0:
time.sleep(self.state.delay_seconds)
recorded: Final = decode_spans(body, self.headers.get("content-type", ""))
self.state.spans.extend(recorded)
self.state.requests.append(
{
"path": self.path,
"count": len(recorded),
"host": self.headers.get("host", ""),
"headers": dict(self.headers),
}
)
self._send_json({"recorded": len(recorded)}, status=self.state.status)
do_POST = _record
do_PUT = _record
def do_GET(self) -> None:
parsed: Final = urlparse(self.path)
if parsed.path == "/__spans":
since: Final = int(dict(part.split("=", 1) for part in parsed.query.split("&") if part).get("since", "0"))
self._send_json({"next": len(self.state.spans), "spans": self.state.spans[since:]})
return
if parsed.path == "/__pid":
self._send_json({"pid": os.getpid()})
return
if parsed.path == "/__requests":
self._send_json({"requests": self.state.requests})
return
self._send_json({"error": "unknown"}, status=404)
def do_DELETE(self) -> None:
if urlparse(self.path).path == "/__spans":
self.state.spans.clear()
self.state.requests.clear()
self._send_json({"cleared": True})
return
self._send_json({"error": "unknown"}, status=404)
def do_PATCH(self) -> None:
if urlparse(self.path).path != "/__control":
self._send_json({"error": "unknown"}, status=404)
return
fields: Final = json.loads(self._read_body() or b"{}")
if "status" in fields:
self.state.status = int(fields["status"])
if "delay_seconds" in fields:
self.state.delay_seconds = float(fields["delay_seconds"])
if fields.get("paused") is True:
self.state.pause.clear()
if fields.get("paused") is False:
self.state.pause.set()
self._send_json({"status": self.state.status, "delay_seconds": self.state.delay_seconds})
def log_message(self, format: str, *args: object) -> None:
pass
class _ConnectHandler(_Handler):
tunnel_context: ssl.SSLContext
def do_CONNECT(self) -> None:
self.state.requests.append({"connect": self.path})
self.connection.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
wrapped: Final = self.tunnel_context.wrap_socket(self.connection, server_side=True)
self.close_connection = True
type(self)(wrapped, self.client_address, self.server)
_MITM_HOSTS: Final = ("otlp.nr-data.net", "otlp.eu01.nr-data.net")
def _mitm_context(directory: Path) -> ssl.SSLContext:
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import NameOID
directory.mkdir(parents=True, exist_ok=True)
now: Final = datetime.datetime.now(datetime.timezone.utc)
window: Final = datetime.timedelta(days=2)
ca_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
ca_name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "otlp-sink test CA")])
ca_cert: Final = (
x509.CertificateBuilder()
.subject_name(ca_name)
.issuer_name(ca_name)
.public_key(ca_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(now - window)
.not_valid_after(now + window)
.add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
.sign(ca_key, hashes.SHA256())
)
leaf_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
leaf_cert: Final = (
x509.CertificateBuilder()
.subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, _MITM_HOSTS[0])]))
.issuer_name(ca_cert.subject)
.public_key(leaf_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(now - window)
.not_valid_after(now + window)
.add_extension(x509.SubjectAlternativeName([x509.DNSName(host) for host in _MITM_HOSTS]), critical=False)
.sign(ca_key, hashes.SHA256())
)
ca_pem: Final = directory / "ca.pem"
ca_pem.write_bytes(ca_cert.public_bytes(serialization.Encoding.PEM))
leaf_pem: Final = directory / "leaf.pem"
leaf_pem.write_bytes(leaf_cert.public_bytes(serialization.Encoding.PEM))
leaf_key_pem: Final = directory / "leaf-key.pem"
leaf_key_pem.write_bytes(
leaf_key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.TraditionalOpenSSL,
serialization.NoEncryption(),
)
)
context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
context.load_cert_chain(str(leaf_pem), str(leaf_key_pem))
return context
def _grpc_trace_server(state: _State, port: int) -> object:
from concurrent import futures
import grpc
from opentelemetry.proto.collector.trace.v1 import trace_service_pb2, trace_service_pb2_grpc
class _TraceService(trace_service_pb2_grpc.TraceServiceServicer):
def Export(self, request: object, context: grpc.ServicerContext) -> object:
state.pause.wait(timeout=120)
if state.delay_seconds > 0:
time.sleep(state.delay_seconds)
recorded: Final = _proto_spans(request.SerializeToString())
state.spans.extend(recorded)
state.requests.append(
{
"grpc": "Export",
"metadata": {key: value for key, value in context.invocation_metadata()},
"count": len(recorded),
}
)
return trace_service_pb2.ExportTraceServiceResponse()
server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
trace_service_pb2_grpc.add_TraceServiceServicer_to_server(_TraceService(), server)
server.add_insecure_port(f"127.0.0.1:{port}")
server.start()
return server
def recorded_spans(url: str, since: int = 0) -> tuple[int, tuple[Span, ...]]:
response: Final = httpx.get(f"{url}/__spans", params={"since": since}, trust_env=False, timeout=15)
response.raise_for_status()
listing: Final = _SPAN_LISTING.validate_python(response.json())
return listing["next"], tuple(listing["spans"])
def configure_sink(url: str, **fields: JsonValue) -> None:
httpx.request("PATCH", f"{url}/__control", json=dict(fields), trust_env=False, timeout=15).raise_for_status()
def reset_sink(url: str) -> None:
httpx.delete(f"{url}/__spans", trust_env=False, timeout=15).raise_for_status()
def sink_pid(url: str) -> int:
return int(httpx.get(f"{url}/__pid", trust_env=False, timeout=15).json()["pid"])
_REQUEST_LISTING: Final = TypeAdapter(list[dict[str, JsonValue]])
def recorded_requests(url: str) -> tuple[Mapping[str, JsonValue], ...]:
response: Final = httpx.get(f"{url}/__requests", trust_env=False, timeout=15)
response.raise_for_status()
return tuple(_REQUEST_LISTING.validate_python(response.json()["requests"]))
@dataclass(frozen=True, slots=True)
class SpanSinks:
operator: str
tenant: str
arize: str
@dataclass(frozen=True, slots=True)
class GrpcSink:
url: str
control_url: str
@dataclass(frozen=True, slots=True)
class ConnectSink:
proxy_url: str
control_url: str
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
except httpx.TransportError:
return False
@contextmanager
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))
root: Final = Path(__file__).resolve().parents[3]
with ExitStack() as stack:
processes: Final = tuple(
subprocess.Popen(
[sys.executable, "-P", "-m", "integration._support.otlp_sink", "--port", str(port)],
cwd=root,
stdout=stack.enter_context((directory / f"otlp-sink-{port}.log").open("w")),
stderr=subprocess.STDOUT,
start_new_session=True,
)
for port in ports
)
try:
urls: Final = tuple(f"http://127.0.0.1:{port}" for port in ports)
deadline: Final = time.monotonic() + 30
while True:
alive: Final = all(process.poll() is None for process in processes)
assert alive, "OTLP sink exited before readiness"
if all(_pid_reachable(url) for url in urls):
break
assert time.monotonic() < deadline, "OTLP sink readiness deadline exceeded"
time.sleep(0.05)
yield SpanSinks(operator=urls[0], tenant=urls[1], arize=urls[2])
finally:
for process in processes:
stopped: Final = stop_root_process(process)
residual: Final = group_members(process.pid)
if residual:
signal_group(process.pid, signal.SIGKILL)
psutil.wait_procs(residual, timeout=5)
survivors: Final = group_members(process.pid)
assert not survivors and stopped, "OTLP sink required forced cleanup"
@contextmanager
def _spawn_sink(directory: Path, log_name: str, argv: Sequence[str]) -> Iterator[None]:
from integration._support.process import group_members, signal_group, stop_root_process
directory.mkdir(parents=True, exist_ok=True)
root: Final = Path(__file__).resolve().parents[3]
with (directory / log_name).open("w") as log:
process: Final = subprocess.Popen(
[sys.executable, "-P", "-m", "integration._support.otlp_sink", *argv],
cwd=root,
stdout=log,
stderr=subprocess.STDOUT,
start_new_session=True,
)
try:
yield
finally:
stopped: Final = stop_root_process(process)
residual: Final = group_members(process.pid)
if residual:
signal_group(process.pid, signal.SIGKILL)
psutil.wait_procs(residual, timeout=5)
survivors: Final = group_members(process.pid)
assert not survivors and stopped, "OTLP sink required forced cleanup"
def _await_sink(url: str) -> None:
deadline: Final = time.monotonic() + 30
while not _pid_reachable(url):
assert time.monotonic() < deadline, "OTLP sink readiness deadline exceeded"
time.sleep(0.05)
@contextmanager
def owned_grpc_sink(directory: Path) -> Iterator[GrpcSink]:
http_port: Final = _free_port()
grpc_port: Final = _free_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)
yield GrpcSink(url=f"http://127.0.0.1:{grpc_port}", control_url=control_url)
@contextmanager
def owned_connect_sink(directory: Path) -> Iterator[ConnectSink]:
http_port: Final = _free_port()
tunnel_port: Final = _free_port()
ca_dir: Final = directory / "mitm"
with _spawn_sink(
directory,
"otlp-connect-sink.log",
["--port", str(http_port), "--connect-port", str(tunnel_port), "--ca-dir", str(ca_dir)],
):
control_url: Final = f"http://127.0.0.1:{http_port}"
_await_sink(control_url)
yield ConnectSink(
proxy_url=f"http://127.0.0.1:{tunnel_port}",
control_url=control_url,
ca_pem=str(ca_dir / "ca.pem"),
)
def main() -> None:
parser: Final = argparse.ArgumentParser()
parser.add_argument("--port", type=int, required=True)
parser.add_argument("--grpc-port", type=int, default=0)
parser.add_argument("--connect-port", type=int, default=0)
parser.add_argument("--ca-dir", type=Path, default=None)
arguments: Final = parser.parse_args()
bound_state: Final = _State()
class BoundHandler(_Handler):
state = bound_state
if arguments.grpc_port:
grpc_server: Final = _grpc_trace_server(bound_state, arguments.grpc_port)
assert grpc_server is not None
if arguments.connect_port:
assert arguments.ca_dir is not None, "--connect-port needs --ca-dir"
bound_context: Final = _mitm_context(arguments.ca_dir)
class BoundConnectHandler(_ConnectHandler):
state = bound_state
tunnel_context = bound_context # pyright: ignore[reportIncompatibleVariableOverride] # bound context, not a new field
tunnel: Final = ThreadingHTTPServer(("127.0.0.1", arguments.connect_port), BoundConnectHandler)
tunnel.daemon_threads = True
threading.Thread(target=tunnel.serve_forever, daemon=True).start()
server: Final = ThreadingHTTPServer(("127.0.0.1", arguments.port), BoundHandler)
server.daemon_threads = True
server.serve_forever()
if __name__ == "__main__":
main()

View file

@ -0,0 +1,53 @@
from __future__ import annotations
import uuid
from collections.abc import Callable, Iterator, Mapping
from pathlib import Path
from typing import Final
from urllib.parse import urlparse
import pytest
import yaml
from integration._support.otlp_sink import SpanSinks, owned_sinks
from pydantic import JsonValue
AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path]
@pytest.fixture(scope="module")
def audit_sinks(tmp_path_factory: pytest.TempPathFactory) -> Iterator[SpanSinks]:
directory: Final = tmp_path_factory.mktemp("otel-audit-sinks")
with owned_sinks(directory) as sinks:
yield sinks
@pytest.fixture(scope="module")
def otel_audit_config(audit_sinks: SpanSinks) -> AuditConfigWriter:
tenant_host: Final = urlparse(audit_sinks.tenant).netloc
def write(directory: Path, litellm_settings: Mapping[str, JsonValue] = {}) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"] = {
**config.get("litellm_settings", {}),
"callbacks": ["otel"],
"provider_url_destination_allowed_hosts": [tenant_host],
**dict(litellm_settings),
}
config["callback_settings"] = {
"otel": {"exporter": "http/json", "endpoint": audit_sinks.operator, "use_simple_processor": True}
}
config["general_settings"] = {**config.get("general_settings", {}), "user_api_key_cache_ttl": 2}
path: Final = directory / f"otel-audit-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(config))
return path
return write
@pytest.fixture(scope="module")
def langfuse_vars(audit_sinks: SpanSinks) -> dict[str, JsonValue]:
return {
"langfuse_public_key": "pk-lf-audit",
"langfuse_secret_key": "sk-lf-audit",
"langfuse_host": audit_sinks.tenant,
}

View file

@ -0,0 +1,392 @@
from __future__ import annotations
import time
import uuid
from collections.abc import Callable, Iterator, Mapping
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import (
Gateway,
eventually,
gateway_from_environment,
)
from integration._support.otlp_sink import (
Span,
SpanSinks,
recorded_spans,
spans_for_trace,
)
from integration._support.process import owned_proxy, owned_proxy_process
from pydantic import JsonValue
AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path]
DB_SYSTEM_KEYS: Final = frozenset({"db.system.name", "db.system"})
@pytest.fixture(scope="module")
def gateway(audit_sinks: SpanSinks) -> Iterator[Gateway]:
with gateway_from_environment() as base:
yield base
def _config_with(
directory: Path,
otel_audit_config: AuditConfigWriter,
*,
otel: Mapping[str, JsonValue] = MappingProxyType({}),
extra: Callable[[dict[str, JsonValue]], None] | None = None,
) -> Path:
config: Final = yaml.safe_load(otel_audit_config(directory, {}).read_text())
config["callback_settings"]["otel"].update(dict(otel))
if extra is not None:
extra(config)
path: Final = directory / f"otel-excl-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(config))
return path
def _operator_langfuse(audit_sinks: SpanSinks) -> dict[str, str]:
return {
"LANGFUSE_HOST": audit_sinks.operator,
"LANGFUSE_PUBLIC_KEY": "pk-lf-operator",
"LANGFUSE_SECRET_KEY": "sk-lf-operator",
"OTEL_EXPORTER": "http/json",
"OTEL_ENDPOINT": audit_sinks.operator,
}
def _add_callback(gateway: Gateway, team_id: str, callback_vars: Mapping[str, JsonValue]) -> httpx.Response:
return gateway.request(
"POST",
f"/team/{team_id}/callback",
{"callback_name": "langfuse_otel", "callback_vars": dict(callback_vars)},
)
def _drive(candidate: Gateway, langfuse_vars: Mapping[str, JsonValue]) -> httpx.Response:
with candidate.scenario() as scenario:
model: Final = scenario.model(model="openai/audit-chat", api_base=f"{candidate.upstream_url}/v1")
team_id: Final = scenario.team()
callback: Final = _add_callback(candidate, team_id, langfuse_vars)
assert callback.status_code == 200, callback.text
key: Final = scenario.key(team_id=team_id)
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"otel-excl-{uuid.uuid4().hex}"}]},
key=key,
)
assert response.status_code == 200, response.text
return response
def _trace_id(sink_url: str, response: httpx.Response, seconds: float = 40) -> str:
call_id: Final = response.headers.get("x-litellm-call-id")
response_id: Final = response.json().get("id")
def look() -> str | None:
_, spans = recorded_spans(sink_url)
return next(
(
str(span["trace_id"])
for span in spans
if (call_id is not None and span["attributes"].get("litellm.call_id") == call_id)
or (response_id is not None and span["attributes"].get("gen_ai.response.id") == response_id)
),
None,
)
found: Final = eventually(look, lambda value: value is not None, seconds=seconds)
assert found is not None
return found
def _trace_spans(sink_url: str, trace_id: str, seconds: float = 30) -> tuple[Span, ...]:
"""The trace's spans once the post-call tail has landed.
The spend-writer and other post-response spans flush after the request
answers, so absence assertions poll for the whole window instead of
settling at the first glimpse of the root span.
"""
deadline: Final = time.monotonic() + seconds
group: tuple[Span, ...] = () # rebind-ok: drains samples until the post-call tail lands
while time.monotonic() < deadline:
_, spans = recorded_spans(sink_url)
group = spans_for_trace(spans, trace_id)
time.sleep(0.5)
assert group, f"trace {trace_id} never reached {sink_url}"
return group
def _await_db_span(sink_url: str, trace_id: str | None, needle: str, seconds: float = 40, since: int = 0) -> None:
def seen() -> bool:
_, spans = recorded_spans(sink_url, since)
group: Final = spans if trace_id is None else spans_for_trace(spans, trace_id)
return any(
needle in str(span["name"]) or needle in {str(span["attributes"].get(k)) for k in DB_SYSTEM_KEYS}
for span in group
)
landed: Final = eventually(seen, bool, seconds=seconds)
assert landed, f"{needle} span never landed at {sink_url}"
def _db_systems(spans: tuple[Span, ...]) -> set[str]:
return {str(span["attributes"][key]) for span in spans for key in DB_SYSTEM_KEYS if key in span["attributes"]}
def _assert_core_spans_present(spans: tuple[Span, ...]) -> None:
attributes_by_span: Final = tuple(span["attributes"] for span in spans)
assert any(span["kind"] == 2 for span in spans), "request root span missing"
assert any("gen_ai.operation.name" in attrs for attrs in attributes_by_span), "model span missing"
assert any("litellm.guardrail.name" in attrs for attrs in attributes_by_span), "guardrail span missing"
names: Final = sorted(str(span["name"]) for span in spans)
assert any(name.startswith("auth") for name in names), f"auth span missing in {names}"
def _assert_tenant_keeps_redis_without_postgres(
candidate: Gateway, audit_sinks: SpanSinks, langfuse_vars: Mapping[str, JsonValue]
) -> None:
tenant_start, _ = recorded_spans(audit_sinks.tenant)
operator_start, _ = recorded_spans(audit_sinks.operator)
traffic: Final = _drive(candidate, langfuse_vars)
_await_db_span(audit_sinks.operator, None, "batch_write_to_db", seconds=60, since=operator_start)
tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic)
_await_db_span(audit_sinks.tenant, tenant_trace, "redis", seconds=60)
systems: Final = _db_systems(_trace_spans(audit_sinks.tenant, tenant_trace, seconds=15))
assert "redis" in systems, f"redis spans missing at tenant: {systems}"
_, all_tenant = recorded_spans(audit_sinks.tenant, tenant_start)
assert "postgresql" not in _db_systems(all_tenant), f"postgresql spans reached tenant: {_db_systems(all_tenant)}"
def _guardrail_block(config: dict) -> None:
config["guardrails"] = [
{
"guardrail_name": f"excl-filter-{uuid.uuid4().hex[:8]}",
"litellm_params": {
"guardrail": "litellm_content_filter",
"mode": "pre_call",
"default_on": True,
"patterns": [
{
"pattern_type": "regex",
"pattern_name": "excl_secret",
"pattern": "TOPSECRET\\d{9}",
"action": "BLOCK",
}
],
},
}
]
@pytest.mark.timeout(180)
def test_excluded_services_drops_db_spans_at_tenant_only(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
config: Final = _config_with(
tmp_path, otel_audit_config, otel={"excluded_services": ["redis", "postgres"]}, extra=_guardrail_block
)
with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate:
ten_start, _ = recorded_spans(audit_sinks.tenant)
op_start, _ = recorded_spans(audit_sinks.operator)
traffic: Final = _drive(candidate, langfuse_vars)
_await_db_span(audit_sinks.operator, None, "postgresql", seconds=60, since=op_start)
tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic)
tenant_spans: Final = _trace_spans(audit_sinks.tenant, tenant_trace)
_assert_core_spans_present(tenant_spans)
assert _db_systems(tenant_spans) == set(), (
f"db spans reached tenant: {sorted(str(s['name']) for s in tenant_spans)}"
)
operator_trace: Final = _trace_id(audit_sinks.operator, traffic)
assert operator_trace == tenant_trace
trace_systems: Final = _db_systems(_trace_spans(audit_sinks.operator, operator_trace))
assert "redis" in trace_systems, f"operator trace lost redis spans: {trace_systems}"
_, all_operator = recorded_spans(audit_sinks.operator, op_start)
operator_systems: Final = _db_systems(all_operator)
assert "postgresql" in operator_systems, f"operator lost aux db spans: {operator_systems}"
_, all_tenant = recorded_spans(audit_sinks.tenant, ten_start)
names: Final = sorted(str(span["name"]) for span in all_tenant)
assert _db_systems(all_tenant) == set(), f"aux db spans reached tenant: {names}"
assert not any("batch_write_to_db" in name for name in names), f"spend writer reached tenant: {names}"
@pytest.mark.timeout(180)
def test_without_excluded_services_the_tenant_still_gets_redis_and_postgres_spans(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
config: Final = _config_with(tmp_path, otel_audit_config, extra=_guardrail_block)
with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate:
tenant_start, _ = recorded_spans(audit_sinks.tenant)
traffic: Final = _drive(candidate, langfuse_vars)
_await_db_span(audit_sinks.tenant, None, "batch_write_to_db", seconds=60, since=tenant_start)
tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic)
_await_db_span(audit_sinks.tenant, tenant_trace, "redis", seconds=60)
_assert_core_spans_present(_trace_spans(audit_sinks.tenant, tenant_trace, seconds=15))
_, all_tenant = recorded_spans(audit_sinks.tenant, tenant_start)
systems: Final = _db_systems(all_tenant)
assert {"redis", "postgresql"} <= systems, f"datastore spans missing at tenant: {systems}"
def test_env_excluded_services_drops_only_redis(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
config: Final = _config_with(tmp_path, otel_audit_config)
with owned_proxy(
gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2
) as candidate:
start, _ = recorded_spans(audit_sinks.tenant)
_drive(candidate, langfuse_vars)
_await_db_span(audit_sinks.tenant, None, "postgresql", seconds=60, since=start)
_, tenant_spans = recorded_spans(audit_sinks.tenant, start)
systems: Final = _db_systems(tenant_spans)
assert "postgresql" in systems, f"postgresql spans missing at tenant: {systems}"
assert "redis" not in systems, f"redis spans reached tenant: {sorted(str(s['name']) for s in tenant_spans)}"
@pytest.mark.timeout(180)
def test_config_excluded_services_wins_over_env(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
def with_langfuse_otel(config: dict) -> None:
config["litellm_settings"]["callbacks"] = ["otel", "langfuse_otel"]
config: Final = _config_with(
tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=with_langfuse_otel
)
with owned_proxy(
gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2
) as candidate:
_assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars)
@pytest.mark.timeout(180)
def test_excluded_services_applies_with_preset_ordered_first(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
def preset_first(config: dict) -> None:
config["litellm_settings"]["callbacks"] = ["langfuse_otel", "otel"]
config: Final = _config_with(
tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=preset_first
)
overrides: Final = {"LITELLM_OTEL_V2": "1", **_operator_langfuse(audit_sinks)}
with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate:
_assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars)
@pytest.mark.timeout(180)
def test_bogus_excluded_service_logs_error_and_drops_at_proxy_start(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["auth", "postgres"]})
with owned_proxy_process(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as owned:
assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:]
_assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars)
@pytest.mark.timeout(180)
def test_valid_config_excluded_services_tolerates_bogus_env(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]})
with owned_proxy(
gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth"}, config=config, workers=2
) as candidate:
_assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars)
@pytest.mark.timeout(180)
def test_bogus_excluded_services_env_logs_and_drops_with_preset_alongside_otel(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
def with_langfuse_otel(config: dict) -> None:
config["litellm_settings"]["callbacks"] = ["otel", "langfuse_otel"]
config: Final = _config_with(tmp_path, otel_audit_config, extra=with_langfuse_otel)
overrides: Final = {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres"}
with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned:
assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:]
_assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars)
@pytest.mark.timeout(180)
def test_bogus_excluded_services_env_logs_and_drops_without_otel_callback(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
def presets_only(config: dict) -> None:
config["litellm_settings"]["callbacks"] = ["langfuse_otel"]
config: Final = _config_with(tmp_path, otel_audit_config, extra=presets_only)
overrides: Final = {
"LITELLM_OTEL_V2": "1",
"LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres",
**_operator_langfuse(audit_sinks),
}
with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned:
assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:]
_assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars)
def test_postgres_exclusion_covers_batch_write_to_db(
gateway: Gateway,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]})
with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate:
op_start, _ = recorded_spans(audit_sinks.operator)
ten_start, _ = recorded_spans(audit_sinks.tenant)
traffic: Final = _drive(candidate, langfuse_vars)
_await_db_span(audit_sinks.operator, None, "batch_write_to_db", seconds=60, since=op_start)
tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic)
_await_db_span(audit_sinks.tenant, tenant_trace, "redis", seconds=60)
tenant_spans: Final = _trace_spans(audit_sinks.tenant, tenant_trace, seconds=15)
_, all_tenant = recorded_spans(audit_sinks.tenant, ten_start)
names: Final = sorted(str(span["name"]) for span in all_tenant)
assert "redis" in _db_systems(tenant_spans), f"redis spans missing at tenant: {names}"
assert not any("batch_write_to_db" in name for name in names), f"spend writer reached tenant: {names}"

View file

@ -0,0 +1,713 @@
import asyncio
import json
import os
import re
import signal
import uuid
from collections.abc import Callable, Generator, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final, Literal
import anthropic
import httpx
import openai
import psutil
import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value
from integration._support.otlp_sink import Span, SpanSinks, configure_sink, recorded_spans, spans_for_trace
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
MARKER: Final = re.compile(rb"excl-[0-9a-f]{32}")
FAILING: Final = re.compile(rb"excl-fail-[0-9a-f]{32}")
JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
REPLY_TEXT: Final = "excluded ok"
SERVER: Final = 2
INVALID_NAME_LOG: Final = "is not a datastore service"
INVALID_VALUE_LOG: Final = "excluded_services must be"
Endpoint = Literal["chat", "responses", "messages"]
Client = Literal["raw", "sdk", "async_sdk"]
ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "responses", "messages")
CLIENTS: Final[tuple[Client, ...]] = ("raw", "sdk", "async_sdk")
AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path]
def _marker() -> str:
return "excl-" + uuid.uuid4().hex
def _chat_reply(identity: str, stream: bool) -> Reply:
if not stream:
return Reply(
body=json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": REPLY_TEXT}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
}
).encode()
)
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"}
first, _, rest = REPLY_TEXT.partition(" ")
deltas: Final[tuple[dict[str, JsonValue], ...]] = (
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": first}}]},
{**chunk, "choices": [{"index": 0, "delta": {"content": " " + rest}, "finish_reason": "stop"}]},
{**chunk, "choices": [], "usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}},
)
return Reply(
content_type="text/event-stream",
chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"),
)
def _responses_reply(identity: str, stream: bool) -> Reply:
response: Final[dict[str, JsonValue]] = {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"id": "msg_" + identity,
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": REPLY_TEXT, "annotations": []}],
}
],
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
}
if not stream:
return Reply(body=json.dumps(response).encode())
events: Final[tuple[dict[str, JsonValue], ...]] = (
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
},
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": "msg_" + identity,
"output_index": 0,
"content_index": 0,
"delta": REPLY_TEXT,
},
{"type": "response.completed", "sequence_number": 2, "response": response},
)
return Reply(
content_type="text/event-stream",
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
)
def _upstream(request: Request) -> Reply:
if FAILING.search(request.body) is not None:
return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}')
found: Final = MARKER.search(request.body)
if found is None:
return Reply(status=404, body=b'{"error":"no marker"}')
marker: Final = found.group(0).decode()
stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True
if request.target.endswith("/responses"):
return _responses_reply(f"resp_{marker}", stream)
return _chat_reply(f"chatcmpl-{marker}", stream)
def _at(payload: JsonValue, *path: str | int) -> JsonValue:
if not path:
return payload
step: Final = path[0]
if isinstance(step, int):
assert isinstance(payload, list), payload
return _at(payload[step], *path[1:])
return _at(object_value(payload)[step], *path[1:])
def _sse(body: str) -> tuple[JsonValue, ...]:
return tuple(
JSON.validate_json(line[6:])
for line in body.splitlines()
if line.startswith("data: ") and line != "data: [DONE]"
)
def _raw_text(endpoint: Endpoint, stream: bool, body: str) -> str:
if not stream:
path: Final[tuple[str | int, ...]] = {
"chat": ("choices", 0, "message", "content"),
"responses": ("output", 0, "content", 0, "text"),
"messages": ("content", 0, "text"),
}[endpoint]
return str(_at(JSON.validate_json(body), *path))
events: Final = _sse(body)
if endpoint == "chat":
return "".join(
str(object_value(_at(event, "choices", 0, "delta")).get("content") or "")
for event in events
if _at(event, "choices")
)
if endpoint == "responses":
return "".join(
str(_at(event, "delta")) for event in events if _at(event, "type") == "response.output_text.delta"
)
return "".join(
str(_at(event, "delta", "text"))
for event in events
if _at(event, "type") == "content_block_delta" and _at(event, "delta", "type") == "text_delta"
)
def _body(model: str, endpoint: Endpoint, marker: str, stream: bool) -> tuple[str, dict[str, JsonValue]]:
if endpoint == "chat":
return "/v1/chat/completions", {
"model": model,
"messages": [{"role": "user", "content": marker}],
"stream": stream,
}
if endpoint == "responses":
return "/v1/responses", {"model": model, "input": marker, "stream": stream}
return "/v1/messages", {
"model": model,
"max_tokens": 16,
"messages": [{"role": "user", "content": marker}],
"stream": stream,
}
@dataclass(frozen=True, slots=True)
class Sent:
call_id: str
text: str
@dataclass(frozen=True, slots=True)
class Cursors:
operator: int
tenant: int
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
owned: OwnedProxy
scenario: Scenario
model: str
key: str
upstream: Wire
sinks: SpanSinks
def cursors(self) -> Cursors:
self.upstream.drain()
return Cursors(recorded_spans(self.sinks.operator)[0], recorded_spans(self.sinks.tenant)[0])
def upstream_hits(self, marker: str) -> int:
return sum(1 for request in self.upstream.drain() if marker.encode() in request.body)
def base_url(self) -> str:
return str(self.proxy.client.base_url)
def raw(
self, endpoint: Endpoint, marker: str, stream: bool, key: str | None = None, trace_id: str | None = None
) -> Sent:
path, body = _body(self.model, endpoint, marker, stream)
auth: Final = {"Authorization": f"Bearer {key or self.key}"}
parent: Final = {} if trace_id is None else {"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"}
with self.proxy.client.stream("POST", path, json=body, headers={**auth, **parent}) as response:
text: Final = response.read().decode()
assert response.status_code == 200, text
return Sent(response.headers["x-litellm-call-id"], _raw_text(endpoint, stream, text))
def sdk(self, endpoint: Endpoint, marker: str, stream: bool) -> Sent:
if endpoint == "messages":
messages: Final = anthropic.Anthropic(base_url=self.base_url(), api_key=self.key, max_retries=0).messages
if not stream:
reply: Final = messages.with_raw_response.create(
model=self.model, max_tokens=16, messages=[{"role": "user", "content": marker}]
)
block: Final = reply.parse().content[0]
assert isinstance(block, anthropic.types.TextBlock), block
return Sent(reply.headers["x-litellm-call-id"], block.text)
with messages.with_streaming_response.create(
model=self.model, max_tokens=16, messages=[{"role": "user", "content": marker}], stream=True
) as streamed:
return Sent(
streamed.headers["x-litellm-call-id"],
"".join(
event.delta.text
for event in streamed.parse()
if event.type == "content_block_delta" and event.delta.type == "text_delta"
),
)
client: Final = openai.OpenAI(base_url=self.base_url() + "/v1", api_key=self.key, max_retries=0)
if endpoint == "chat":
if not stream:
completion: Final = client.chat.completions.with_raw_response.create(
model=self.model, messages=[{"role": "user", "content": marker}]
)
return Sent(
completion.headers["x-litellm-call-id"], completion.parse().choices[0].message.content or ""
)
with client.chat.completions.with_streaming_response.create(
model=self.model, messages=[{"role": "user", "content": marker}], stream=True
) as chunks:
return Sent(
chunks.headers["x-litellm-call-id"],
"".join(chunk.choices[0].delta.content or "" for chunk in chunks.parse() if chunk.choices),
)
if not stream:
created: Final = client.responses.with_raw_response.create(model=self.model, input=marker)
return Sent(created.headers["x-litellm-call-id"], created.parse().output_text)
with client.responses.with_streaming_response.create(model=self.model, input=marker, stream=True) as events:
return Sent(
events.headers["x-litellm-call-id"],
"".join(event.delta for event in events.parse() if event.type == "response.output_text.delta"),
)
async def async_sdk(self, endpoint: Endpoint, marker: str, stream: bool) -> Sent:
if endpoint == "messages":
messages: Final = anthropic.AsyncAnthropic(
base_url=self.base_url(), api_key=self.key, max_retries=0
).messages
if not stream:
reply: Final = await messages.with_raw_response.create(
model=self.model, max_tokens=16, messages=[{"role": "user", "content": marker}]
)
block: Final = reply.parse().content[0]
assert isinstance(block, anthropic.types.TextBlock), block
return Sent(reply.headers["x-litellm-call-id"], block.text)
async with messages.with_streaming_response.create(
model=self.model, max_tokens=16, messages=[{"role": "user", "content": marker}], stream=True
) as streamed:
pieces: Final = [
event.delta.text
async for event in await streamed.parse()
if event.type == "content_block_delta" and event.delta.type == "text_delta"
]
return Sent(streamed.headers["x-litellm-call-id"], "".join(pieces))
client: Final = openai.AsyncOpenAI(base_url=self.base_url() + "/v1", api_key=self.key, max_retries=0)
if endpoint == "chat":
if not stream:
completion: Final = await client.chat.completions.with_raw_response.create(
model=self.model, messages=[{"role": "user", "content": marker}]
)
return Sent(
completion.headers["x-litellm-call-id"], completion.parse().choices[0].message.content or ""
)
async with client.chat.completions.with_streaming_response.create(
model=self.model, messages=[{"role": "user", "content": marker}], stream=True
) as chunks:
deltas: Final = [
chunk.choices[0].delta.content or "" async for chunk in await chunks.parse() if chunk.choices
]
return Sent(chunks.headers["x-litellm-call-id"], "".join(deltas))
if not stream:
created: Final = await client.responses.with_raw_response.create(model=self.model, input=marker)
return Sent(created.headers["x-litellm-call-id"], created.parse().output_text)
async with client.responses.with_streaming_response.create(
model=self.model, input=marker, stream=True
) as events:
texts: Final = [
event.delta async for event in await events.parse() if event.type == "response.output_text.delta"
]
return Sent(events.headers["x-litellm-call-id"], "".join(texts))
def send(self, endpoint: Endpoint, client: Client, marker: str, stream: bool) -> Sent:
if client == "raw":
return self.raw(endpoint, marker, stream)
if client == "sdk":
return self.sdk(endpoint, marker, stream)
return asyncio.run(self.async_sdk(endpoint, marker, stream))
def _db_systems(spans: tuple[Span, ...]) -> set[str]:
return {
str(system)
for span in spans
if (system := span["attributes"].get("db.system.name") or span["attributes"].get("db.system")) is not None
}
def _names(spans: tuple[Span, ...]) -> list[str]:
return sorted(span["name"] for span in spans)
def _has_root(spans: tuple[Span, ...]) -> bool:
return any(span["kind"] == SERVER for span in spans)
def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]:
_, spans = recorded_spans(sink, since)
traces: Final = {span["trace_id"] for span in spans if span["attributes"].get("litellm.call_id") == call_id}
return tuple(span for span in spans if span["trace_id"] in traces)
def _operator_trace(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]:
trace: Final = eventually(
lambda: _trace_of_call(rig.sinks.operator, sent.call_id, cursors.operator),
lambda spans: _has_root(spans) and "redis" in _db_systems(spans),
seconds=40,
)
assert len({span["trace_id"] for span in trace}) == 1, _names(trace)
return trace
def _traced_raw(rig: Rig, endpoint: Endpoint, marker: str) -> tuple[str, Sent]:
trace_id: Final = uuid.uuid4().hex
return trace_id, rig.raw(endpoint, marker, stream=False, trace_id=trace_id)
def _operator_trace_by_id(rig: Rig, trace_id: str, cursors: Cursors) -> tuple[Span, ...]:
return eventually(
lambda: spans_for_trace(recorded_spans(rig.sinks.operator, cursors.operator)[1], trace_id),
_has_root,
seconds=40,
)
def _tenant_mirror(rig: Rig, operator: tuple[Span, ...], cursors: Cursors) -> tuple[Span, ...]:
kept: Final = frozenset(span["name"] for span in operator if not _db_systems((span,)))
return eventually(
lambda: spans_for_trace(recorded_spans(rig.sinks.tenant, cursors.tenant)[1], operator[0]["trace_id"]),
lambda spans: kept <= {span["name"] for span in spans},
seconds=40,
)
def _assert_tenant_mirrors(rig: Rig, operator: tuple[Span, ...], cursors: Cursors) -> tuple[Span, ...]:
tenant: Final = _tenant_mirror(rig, operator, cursors)
assert _db_systems(tenant) == set(), f"datastore spans reached the tenant: {_names(tenant)}"
assert sum(1 for span in tenant if span["kind"] == SERVER) == 1, _names(tenant)
return tenant
def _assert_withheld(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]:
tenant: Final = _assert_tenant_mirrors(rig, _operator_trace(rig, sent, cursors), cursors)
assert any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant)
return tenant
def _config(directory: Path, otel_audit_config: AuditConfigWriter, otel: Mapping[str, JsonValue], name: str) -> Path:
written: Final = otel_audit_config(directory, {})
loaded: Final = object_value(JSON.validate_python(yaml.safe_load(written.read_text())))
settings: Final = object_value(loaded["callback_settings"])
config: Final = {**loaded, "callback_settings": {**settings, "otel": {**object_value(settings["otel"]), **otel}}}
path: Final = directory / f"{name}.yaml"
path.write_text(yaml.safe_dump(config))
return path
@contextmanager
def _started(
provider: Wire,
sinks: SpanSinks,
config: Path,
directory: Path,
langfuse_vars: Mapping[str, JsonValue],
workers: int,
) -> Generator[Rig]:
with (
gateway_from_environment() as gateway,
owned_proxy_process(
gateway,
directory,
{"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"},
config=config,
remove_environment=("LITELLM_OTEL_EXCLUDED_SERVICES",),
workers=workers,
) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1")
team: Final = scenario.team()
attached: Final = owned.gateway.request(
"POST", f"/team/{team}/callback", {"callback_name": "langfuse_otel", "callback_vars": dict(langfuse_vars)}
)
assert attached.status_code == 200, attached.text
yield Rig(owned.gateway, owned, scenario, model, scenario.key(team_id=team), provider, sinks)
@pytest.fixture(scope="module")
def provider() -> Iterator[Wire]:
with wire_server(_upstream) as wire:
yield wire
@pytest.fixture(scope="module")
def rig(
provider: Wire,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path_factory: pytest.TempPathFactory,
) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("excluded-matrix")
config: Final = _config(directory, otel_audit_config, {"excluded_services": ["redis", "postgres"]}, "matrix")
with _started(provider, audit_sinks, config, directory, langfuse_vars, workers=2) as started:
yield started
@pytest.mark.timeout(120)
@pytest.mark.parametrize("stream", [False, True], ids=["unary", "stream"])
@pytest.mark.parametrize("client", CLIENTS)
@pytest.mark.parametrize("endpoint", ENDPOINTS)
def test_tenant_trace_keeps_request_spans_without_datastore_spans(
rig: Rig, endpoint: Endpoint, client: Client, stream: bool
) -> None:
cursors: Final = rig.cursors()
marker: Final = _marker()
sent: Final = rig.send(endpoint, client, marker, stream)
assert sent.text == REPLY_TEXT, sent
assert rig.upstream_hits(marker) == 1
_assert_withheld(rig, sent, cursors)
@pytest.mark.timeout(120)
@pytest.mark.parametrize("endpoint", ["chat", "messages"])
def test_cache_hit_twin_keeps_datastore_spans_off_the_tenant(rig: Rig, endpoint: Endpoint) -> None:
marker: Final = _marker()
first: Final = rig.raw(endpoint, marker, stream=False)
assert first.text == REPLY_TEXT, first
assert rig.upstream_hits(marker) == 1
cursors: Final = rig.cursors()
trace_id, hit = eventually(
lambda: _traced_raw(rig, endpoint, marker), lambda sent: rig.upstream_hits(marker) == 0, seconds=20
)
assert hit.text == REPLY_TEXT, hit
_assert_tenant_mirrors(rig, _operator_trace_by_id(rig, trace_id, cursors), cursors)
@pytest.mark.timeout(120)
@pytest.mark.parametrize("endpoint", ENDPOINTS)
def test_failed_upstream_call_keeps_datastore_spans_off_the_tenant(rig: Rig, endpoint: Endpoint) -> None:
cursors: Final = rig.cursors()
marker: Final = "excl-fail-" + uuid.uuid4().hex
trace_id: Final = uuid.uuid4().hex
path, body = _body(rig.model, endpoint, marker, stream=False)
failed: Final = rig.proxy.client.post(
path,
json=body,
headers={"Authorization": f"Bearer {rig.key}", "traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
)
assert failed.status_code == 500, failed.text
assert rig.upstream_hits(marker) >= 1
operator: Final = eventually(
lambda: spans_for_trace(recorded_spans(rig.sinks.operator, cursors.operator)[1], trace_id),
lambda spans: _has_root(spans) and "redis" in _db_systems(spans),
seconds=40,
)
_assert_tenant_mirrors(rig, operator, cursors)
@pytest.mark.timeout(120)
def test_key_level_callback_vars_destination_is_filtered_too(rig: Rig, langfuse_vars: dict[str, JsonValue]) -> None:
key: Final = rig.scenario.key(
metadata={
"logging": [
{"callback_name": "langfuse_otel", "callback_type": "success", "callback_vars": dict(langfuse_vars)}
]
}
)
cursors: Final = rig.cursors()
marker: Final = _marker()
sent: Final = rig.raw("chat", marker, stream=False, key=key)
assert sent.text == REPLY_TEXT, sent
assert rig.upstream_hits(marker) == 1
_assert_withheld(rig, sent, cursors)
@pytest.mark.timeout(120)
@pytest.mark.parametrize("status", [403, 404])
def test_rejecting_tenant_destination_leaves_serving_and_the_operator_trace_intact(rig: Rig, status: int) -> None:
configure_sink(rig.sinks.tenant, status=status)
try:
cursors: Final = rig.cursors()
marker: Final = _marker()
sent: Final = rig.raw("chat", marker, stream=True)
assert sent.text == REPLY_TEXT, sent
assert rig.upstream_hits(marker) == 1
_assert_withheld(rig, sent, cursors)
finally:
configure_sink(rig.sinks.tenant, status=200)
after: Final = rig.cursors()
_assert_withheld(rig, rig.raw("responses", _marker(), stream=False), after)
def _burst(rig: Rig, count: int) -> tuple[Sent | str, ...]:
def one(index: int) -> Sent | str:
try:
return rig.raw(ENDPOINTS[index % 3], _marker(), stream=index % 2 == 0)
except (httpx.HTTPError, AssertionError) as error:
return repr(error)
with ThreadPoolExecutor(max_workers=10) as pool:
return tuple(pool.map(one, range(count)))
def _served(results: tuple[Sent | str, ...]) -> tuple[Sent, ...]:
return tuple(result for result in results if isinstance(result, Sent))
def _assert_operator_exactly_once(rig: Rig, served: tuple[Sent, ...], cursors: Cursors) -> set[str]:
wanted: Final = {sent.call_id for sent in served}
def roots() -> dict[str, int]:
_, spans = recorded_spans(rig.sinks.operator, cursors.operator)
traced: Final = {
span["trace_id"]: str(span["attributes"]["litellm.call_id"])
for span in spans
if span["attributes"].get("litellm.call_id") in wanted
}
counts: Final = {call: 0 for call in wanted}
for span in spans:
if span["kind"] == SERVER and span["trace_id"] in traced:
counts[traced[span["trace_id"]]] += 1
return counts
landed: Final = eventually(roots, lambda counts: all(count >= 1 for count in counts.values()), seconds=90)
assert landed == {call: 1 for call in wanted}, landed
_, spans = recorded_spans(rig.sinks.operator, cursors.operator)
return {span["trace_id"] for span in spans if span["attributes"].get("litellm.call_id") in wanted}
def _assert_tenant_never_saw_datastore_spans(rig: Rig, cursors: Cursors, traces: set[str]) -> None:
tenant: Final = eventually(
lambda: recorded_spans(rig.sinks.tenant, cursors.tenant)[1],
lambda spans: traces <= {span["trace_id"] for span in spans if span["kind"] == SERVER},
seconds=90,
)
assert _db_systems(tenant) == set(), _names(tenant)
@pytest.mark.timeout(300)
def test_tenant_outage_during_a_mixed_burst_keeps_serving_and_never_leaks_datastore_spans(rig: Rig) -> None:
cursors: Final = rig.cursors()
configure_sink(rig.sinks.tenant, status=503)
try:
results: Final = _burst(rig, 30)
finally:
configure_sink(rig.sinks.tenant, status=200)
served: Final = _served(results)
assert len(served) == 30, [result for result in results if isinstance(result, str)]
assert all(sent.text == REPLY_TEXT for sent in served), served
traces: Final = _assert_operator_exactly_once(rig, served, cursors)
_assert_tenant_never_saw_datastore_spans(rig, cursors, traces)
after: Final = rig.cursors()
_assert_withheld(rig, rig.raw("messages", _marker(), stream=True), after)
@pytest.mark.timeout(300)
def test_stalled_tenant_destination_during_a_burst_does_not_block_responses(rig: Rig) -> None:
cursors: Final = rig.cursors()
configure_sink(rig.sinks.tenant, paused=True)
try:
results: Final = _burst(rig, 20)
finally:
configure_sink(rig.sinks.tenant, paused=False)
served: Final = _served(results)
assert len(served) == 20, [result for result in results if isinstance(result, str)]
traces: Final = _assert_operator_exactly_once(rig, served, cursors)
_assert_tenant_never_saw_datastore_spans(rig, cursors, traces)
@pytest.mark.timeout(300)
def test_killing_one_of_two_workers_mid_burst_keeps_the_filter_on_the_survivor(rig: Rig) -> None:
root: Final = psutil.Process(rig.owned.process.pid)
workers: Final = eventually(
lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())),
lambda found: len(found) == 2,
seconds=30,
)
cursors: Final = rig.cursors()
def one(index: int) -> Sent | str:
if index == 6:
os.kill(workers[0].pid, signal.SIGKILL)
try:
return rig.raw("chat", _marker(), stream=index % 2 == 0)
except (httpx.HTTPError, AssertionError) as error:
return repr(error)
with ThreadPoolExecutor(max_workers=6) as pool:
results: Final = tuple(pool.map(one, range(18)))
assert rig.owned.process.poll() is None, "Proxy root exited after a worker was killed"
failures: Final = tuple(result for result in results if isinstance(result, str))
assert all(failure.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for failure in failures), (
failures
)
assert len(failures) <= 6, failures
settled: Final = tuple(result for index, result in enumerate(results) if index > 12 and isinstance(result, Sent))
traces: Final = _assert_operator_exactly_once(rig, settled, cursors)
_assert_tenant_never_saw_datastore_spans(rig, cursors, traces)
after: Final = rig.cursors()
_assert_withheld(rig, rig.raw("chat", _marker(), stream=False), after)
@dataclass(frozen=True, slots=True)
class Setting:
otel: Mapping[str, JsonValue]
withholds_redis: bool
logs: str | None
SETTINGS: Final[dict[str, Setting]] = {
"missing": Setting({}, False, None),
"null": Setting({"excluded_services": None}, False, None),
"empty_list": Setting({"excluded_services": []}, False, None),
"empty_string": Setting({"excluded_services": ""}, False, None),
"yaml_string": Setting({"excluded_services": "redis"}, True, None),
"duplicates": Setting({"excluded_services": ["redis", "redis"]}, True, None),
"case_and_space": Setting({"excluded_services": ["REDIS", " Postgres "]}, True, None),
"integer": Setting({"excluded_services": 7}, False, INVALID_VALUE_LOG),
"mapping": Setting({"excluded_services": {"redis": True}}, False, INVALID_VALUE_LOG),
"non_string_item": Setting({"excluded_services": [7, "redis"]}, True, INVALID_VALUE_LOG),
"oversized_name": Setting({"excluded_services": "x" * 5000}, False, INVALID_NAME_LOG),
}
@pytest.mark.timeout(180)
@pytest.mark.parametrize("name", SETTINGS)
def test_excluded_services_setting_shapes_boot_and_resolve(
name: str,
provider: Wire,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
langfuse_vars: dict[str, JsonValue],
tmp_path: Path,
) -> None:
setting: Final = SETTINGS[name]
config: Final = _config(tmp_path, otel_audit_config, setting.otel, name)
with _started(provider, audit_sinks, config, tmp_path, langfuse_vars, workers=1) as started:
cursors: Final = started.cursors()
marker: Final = _marker()
sent: Final = started.raw("chat", marker, stream=False)
assert sent.text == REPLY_TEXT, sent
assert started.upstream_hits(marker) == 1
operator: Final = _operator_trace(started, sent, cursors)
tenant: Final = _tenant_mirror(started, operator, cursors)
if setting.withholds_redis:
assert "redis" not in _db_systems(tenant), _names(tenant)
else:
eventually(
lambda: _db_systems(
spans_for_trace(recorded_spans(started.sinks.tenant, cursors.tenant)[1], tenant[0]["trace_id"])
),
lambda systems: "redis" in systems,
seconds=30,
)
log: Final = started.owned.log.read_text()
if setting.logs is None:
assert INVALID_NAME_LOG not in log and INVALID_VALUE_LOG not in log, log[-2000:]
else:
assert setting.logs in log, log[-4000:]

View file

@ -11,6 +11,7 @@
"""
import asyncio
import logging
import pytest
@ -22,17 +23,18 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E4
)
from litellm.integrations.otel import LiteLLM, OpenTelemetryV2Config # noqa: E402
from litellm.integrations.otel.plumbing import providers # noqa: E402
from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
from litellm.integrations.otel.model.baggage import ( # noqa: E402
BAGGAGE_PROMOTED_KEYS,
DEFAULT_BAGGAGE_METADATA_KEYS,
)
from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
from litellm.integrations.otel.model.config import excluded_db_systems_from # noqa: E402
from litellm.integrations.otel.model.payloads import GuardrailSpanData # noqa: E402
from litellm.integrations.otel.model.spans import ( # noqa: E402
LITELLM_PROXY_REQUEST_SPAN_NAME,
SpanRole,
)
from litellm.integrations.otel.plumbing import providers # noqa: E402
# --------------------------------------------------------------------------- #
# Area 1 — baggage allowlists configurable
@ -74,13 +76,11 @@ def test_baggage_keys_from_config_yaml_kwargs():
def test_baggage_processor_allowlist_uses_config_keys():
cfg = OpenTelemetryV2Config(
exporter="in_memory", baggage_promoted_keys=[LiteLLM.TEAM_ID]
)
cfg = OpenTelemetryV2Config(exporter="in_memory", baggage_promoted_keys=[LiteLLM.TEAM_ID])
provider, exporter = providers.in_memory_provider(cfg)
from litellm.integrations.otel.plumbing import context as ctx_mod
from litellm.integrations.otel.emitter import SpanEmitter
from litellm.integrations.otel.model.payloads import ServiceSpanData
from litellm.integrations.otel.plumbing import context as ctx_mod
engine = SpanEmitter(providers.get_tracer(provider, "t"), cfg)
ctx = ctx_mod.set_request_baggage({LiteLLM.TEAM_ID: "t1", LiteLLM.TEAM_ALIAS: "ta"})
@ -90,6 +90,68 @@ def test_baggage_processor_allowlist_uses_config_keys():
assert LiteLLM.TEAM_ALIAS not in span.attributes # not in this allowlist
@pytest.mark.parametrize(
"given,expected",
[
(["redis"], frozenset({"redis"})),
(["postgres"], frozenset({"postgresql"})),
(["postgresql"], frozenset({"postgresql"})),
(["batch_write_to_db"], frozenset({"postgresql"})),
(["redis_spend_update_queue"], frozenset({"redis"})),
(["redis", "postgres"], frozenset({"redis", "postgresql"})),
],
)
def test_excluded_services_normalize_to_db_system_names(given, expected):
assert OpenTelemetryV2Config(excluded_services=given).excluded_services == expected
def test_excluded_services_from_env_csv(monkeypatch):
monkeypatch.setenv("LITELLM_OTEL_EXCLUDED_SERVICES", "redis, postgres")
assert OpenTelemetryV2Config().excluded_services == frozenset({"redis", "postgresql"})
def test_excluded_services_config_wins_over_env(monkeypatch):
monkeypatch.setenv("LITELLM_OTEL_EXCLUDED_SERVICES", "redis")
assert OpenTelemetryV2Config(excluded_services=["postgres"]).excluded_services == frozenset({"postgresql"})
def test_excluded_services_drops_a_non_datastore_service_and_logs(caplog):
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
config = OpenTelemetryV2Config(excluded_services=["auth", "redis"])
assert config.excluded_services == frozenset({"redis"})
assert any("'auth' is not a datastore service; ignored" in record.message for record in caplog.records)
def test_excluded_services_env_drops_a_bad_value_and_logs(monkeypatch, caplog):
monkeypatch.setenv("LITELLM_OTEL_EXCLUDED_SERVICES", "auth,postgres")
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
config = OpenTelemetryV2Config()
assert config.excluded_services == frozenset({"postgresql"})
assert any("'auth' is not a datastore service; ignored" in record.message for record in caplog.records)
@pytest.mark.parametrize(
"given,expected,logged",
[
(None, frozenset(), None),
("", frozenset(), None),
([], frozenset(), None),
(["REDIS", " Postgres "], frozenset({"redis", "postgresql"}), None),
(7, frozenset(), "excluded_services must be a list or comma-separated string; 7 ignored"),
({"redis": True}, frozenset(), "excluded_services must be a list or comma-separated string"),
([7, "redis"], frozenset({"redis"}), "excluded_services must be a list of service names; 7 ignored"),
],
)
def test_malformed_excluded_services_logs_and_still_builds_the_config(given, expected, logged, caplog):
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
config = OpenTelemetryV2Config(excluded_services=given)
resolved = excluded_db_systems_from(given)
assert config.excluded_services == expected
assert resolved == expected
messages = [record.message for record in caplog.records]
assert (logged is None and messages == []) or any(logged in message for message in messages), messages
# --------------------------------------------------------------------------- #
# Area 2 — pass-through LLM span parents to the ambient server span
# --------------------------------------------------------------------------- #
@ -124,9 +186,7 @@ def test_passthrough_llm_span_parents_to_ambient_server_span():
later (possibly detached) success callback only closes the already-parented
span, so it never becomes a separate root trace."""
logger, exporter = _logger()
server = logger._emitter.start_span(
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
)
server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
kwargs = {
"standard_logging_object": _payload(),
"litellm_params": {"metadata": {}},
@ -150,9 +210,7 @@ def test_llm_span_unaffected_by_phase_span_active_at_close():
successor to the old auth-failure-401 case where the LLM log nested under
``auth``: the span is now born after auth, parented to the request root."""
logger, exporter = _logger()
server = logger._emitter.start_span(
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
)
server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
kwargs = {
"standard_logging_object": _payload(),
"litellm_params": {"metadata": {}},

View file

@ -514,6 +514,53 @@ class TestFanOut:
for child in ("auth /v1/chat/completions", "chat gpt-4"):
assert by_name[child].parent.span_id == root.context.span_id
def test_excluded_services_drop_only_the_datastore_spans_at_the_tenant(self):
"""The exclusion is per ``db.system.*`` value: a span naming an excluded
datastore never reaches the tenant, while every span of the request's
own work (root, auth, guardrail, model) still does, and the operator's
own exporter keeps the full tree."""
dest_exporter, operator_exporter = InMemorySpanExporter(), InMemorySpanExporter()
provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(operator_exporter))
provider.add_span_processor(
TenantFanOutSpanProcessor(
processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter),
excluded_db_systems=frozenset({"redis", "postgresql"}),
)
)
tracer = get_tracer(provider, "litellm")
def run():
set_request_destinations((LANGFUSE_DEST,))
with tracer.start_as_current_span("POST /v1/chat/completions"):
with tracer.start_as_current_span("auth /v1/chat/completions"):
pass
with tracer.start_as_current_span("execute_guardrail pii"):
pass
with tracer.start_as_current_span("redis async_get_cache") as redis_span:
redis_span.set_attribute("db.system.name", "redis")
with tracer.start_as_current_span("batch_write_to_db _PROXY_track_cost_callback") as spend_span:
spend_span.set_attribute("db.system", "postgresql")
with tracer.start_as_current_span("chat gpt-4"):
pass
in_fresh_context(run)
assert {s.name for s in dest_exporter.get_finished_spans()} == {
"POST /v1/chat/completions",
"auth /v1/chat/completions",
"execute_guardrail pii",
"chat gpt-4",
}
assert {s.name for s in operator_exporter.get_finished_spans()} == {
"POST /v1/chat/completions",
"auth /v1/chat/completions",
"execute_guardrail pii",
"redis async_get_cache",
"batch_write_to_db _PROXY_track_cost_callback",
"chat gpt-4",
}
def test_a_team_naming_two_backends_gets_the_trace_at_both(self):
"""The fan-out rides one provider, so it cannot skip a destination on the
grounds that some other backend owns it: nothing else would deliver it."""
@ -1022,6 +1069,89 @@ class TestProviderWiring:
assert kinds(published).count("TenantFanOutSpanProcessor") == 1
assert "TenantFanOutSpanProcessor" not in kinds(other)
@staticmethod
def _fan_out_of(logger: OpenTelemetryV2) -> TenantFanOutSpanProcessor:
return next(
processor
for processor in logger._tracer_provider._active_span_processor._span_processors
if isinstance(processor, TenantFanOutSpanProcessor)
)
def test_callback_settings_excluded_services_win_over_the_published_preset_env_config(self, monkeypatch):
"""A preset builds its config env-only, so the fan-out must read
``callback_settings.otel.excluded_services`` itself rather than the
published logger's config, or the env value would win."""
monkeypatch.setattr(litellm, "callback_settings", {"otel": {"excluded_services": ["postgres"]}}, raising=False)
preset = OpenTelemetryV2(
config=OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory")], excluded_services=["redis"]),
callback_name="langfuse_otel",
)
publish_global_otel_v2_provider([], lambda _p: None, registered=preset)
assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"postgresql"})
def test_callback_settings_excluded_services_apply_even_when_other_otel_env_vars_are_malformed(self, monkeypatch):
"""Reading the setting must not rebuild the whole settings model, or an unrelated bad env
value the operator overrode in config would stop publication before the fan-out is attached"""
preset = OpenTelemetryV2(
config=OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory")]),
callback_name="langfuse_otel",
)
monkeypatch.setenv("LITELLM_OTEL_LEGACY_COMPAT", "not-a-bool")
monkeypatch.setattr(litellm, "callback_settings", {"otel": {"excluded_services": ["postgres"]}}, raising=False)
publish_global_otel_v2_provider([], lambda _p: None, registered=preset)
assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"postgresql"})
def test_excluded_services_fall_back_to_the_published_logger_config_without_callback_settings(self, monkeypatch):
monkeypatch.setattr(litellm, "callback_settings", {"otel": {"exporter": "in_memory"}}, raising=False)
preset = OpenTelemetryV2(
config=OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory")], excluded_services=["redis"]),
callback_name="langfuse_otel",
)
publish_global_otel_v2_provider([], lambda _p: None, registered=preset)
assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"redis"})
def test_otel_after_a_preset_reuses_it_and_still_takes_callback_settings_exclusions(self, monkeypatch):
"""``callbacks: [langfuse_otel, otel]`` keeps one v2 logger, exactly as
before ``excluded_services`` existed, and the exclusion still comes from
``callback_settings.otel`` rather than the preset's env-only config."""
from litellm.litellm_core_utils import litellm_logging as logging_module
logging_module._in_memory_loggers.clear()
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk")
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk")
monkeypatch.setenv("LITELLM_OTEL_EXCLUDED_SERVICES", "redis")
is_otel_v2_enabled.cache_clear()
monkeypatch.setattr(litellm, "callback_settings", {"otel": {"excluded_services": ["postgres"]}}, raising=False)
try:
def init(name: str) -> CustomLogger | None:
return logging_module._init_custom_logger_compatible_class(
logging_integration=name, # pyright: ignore[reportArgumentType] # test passes a literal callback name
internal_usage_cache=None,
llm_router=None,
custom_logger_init_args={},
)
preset = init("langfuse_otel")
otel_cb = init("otel")
assert isinstance(preset, OpenTelemetryV2)
assert otel_cb is preset
v2_loggers = [cb for cb in logging_module._in_memory_loggers if isinstance(cb, OpenTelemetryV2)]
assert v2_loggers == [preset], v2_loggers
publish_global_otel_v2_provider(logging_module._in_memory_loggers, lambda _p: None, registered=preset)
assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"postgresql"})
finally:
logging_module._in_memory_loggers.clear()
is_otel_v2_enabled.cache_clear()
@pytest.mark.parametrize("canonical", ["langfuse_otel", "arize"])
def test_publishing_tells_the_fan_out_about_every_v2_loggers_account(self, monkeypatch, canonical):
monkeypatch.setenv("LITELLM_OTEL_TENANT_DESTINATION_MODE", "additive")