chore: merge main with upstream OpenTelemetry test fix

This commit is contained in:
moe-berri 2026-10-29 15:11:27 -07:00
commit 3ce7c8529b
47 changed files with 2713 additions and 153 deletions

View file

@ -1,14 +1,14 @@
#!/usr/bin/env bash
set -euo pipefail
if [ "${GITHUB_ACTIONS:-}" = true ]; then
echo "Integration contracts are owned by CircleCI" >&2
exit 1
fi
suite="${1:?integration suite required}"
mode="${2:-standard}"
side="${3:-}"
requested_workers="${INTEGRATION_WORKERS:-}"
license_env=()
if [ -n "${LITELLM_LICENSE:-}" ]; then
license_env=("LITELLM_LICENSE=$LITELLM_LICENSE")
fi
if [ "$mode" = replica ]; then
results="test-results/integration-${suite}-replica"
elif [ "$mode" = parity ]; then
@ -84,7 +84,7 @@ if [ "$suite" = browser ]; then
export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
test -f "$LITELLM_UI_PATH/index.html"
fi
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256(((os.environ.get("CIRCLE_SHA1") or os.environ.get("GITHUB_SHA") or "local") + (os.environ.get("CIRCLE_WORKFLOW_ID") or os.environ.get("GITHUB_RUN_ID") or "local")).encode()).hexdigest()[:8],16))')"
export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
@ -107,7 +107,7 @@ sudo iptables -N integration_only
guard_created=true
sudo iptables -A integration_only -o lo -j ACCEPT
sudo iptables -A integration_only -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
for service in postgres-db redis-cache; do
for service in "${INTEGRATION_POSTGRES_CONTAINER:-postgres-db}" "${INTEGRATION_REDIS_CONTAINER:-redis-cache}"; do
address="$(docker inspect --format '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$service")"
sudo iptables -A integration_only -d "$address" -j ACCEPT
done
@ -145,6 +145,9 @@ fi
if [ "$suite" = mcp ]; then
export INTEGRATION_WORKERS=2 INTEGRATION_COVERAGE=1
fi
if [ -n "$requested_workers" ]; then
export INTEGRATION_WORKERS="$requested_workers"
fi
isolated_workers=false
if [ "${INTEGRATION_WORKERS:-1}" -gt 1 ]; then
isolated_workers=true
@ -185,7 +188,7 @@ start_proxy() {
"${database_env[@]}" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
"${license_env[@]}" \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True LITELLM_ENABLE_MCP_STDIO=true "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
@ -257,7 +260,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
"${shared_services[@]}" "${worker_coverage[@]}" \
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
LITELLM_MODE=PRODUCTION \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
"${license_env[@]}" \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
INTEGRATION_CASE_SHARDS="${CIRCLE_NODE_TOTAL:-1}" INTEGRATION_CASE_SHARD_INDEX="${CIRCLE_NODE_INDEX:-0}" \

119
.github/scripts/assert_secret_absent.py vendored Normal file
View file

@ -0,0 +1,119 @@
import base64
import io
import os
import sys
import urllib.parse
import zipfile
from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass
from functools import reduce
from itertools import chain
from pathlib import Path
from types import MappingProxyType
from typing import Final, TextIO
SECRET_VARIABLE: Final = "SCANNED_SECRET"
MIN_ENCODED_LENGTH: Final = 12
MAX_ARCHIVE_DEPTH: Final = 3
MASK: Final = "***"
Forms = tuple[tuple[str, str], ...]
@dataclass(frozen=True, slots=True)
class Hit:
location: str
form: str
def _aligned_encodings(raw: bytes, name: str, encode: Callable[[bytes], bytes]) -> Iterator[tuple[str, str]]:
for offset in range(3):
chunk: Final = raw[offset:]
yield f"{name} offset {offset}", encode(chunk[: len(chunk) - len(chunk) % 3]).decode()
def secret_forms(secret: str) -> Forms:
encoded: Final = (
("url-encoded", urllib.parse.quote(secret, safe="")),
*_aligned_encodings(secret.encode(), "base64", base64.b64encode),
*_aligned_encodings(secret.encode(), "base64url", base64.urlsafe_b64encode),
)
by_value: Final = MappingProxyType(
{value: name for name, value in encoded if value != secret and len(value) >= MIN_ENCODED_LENGTH}
)
return (("raw", secret), *((name, value) for value, name in by_value.items()))
def mask(text: str, forms: Forms) -> str:
return reduce(lambda masked, form: masked.replace(form[1], MASK), forms, text)
def _name_hits(location: str, name: str, forms: Forms) -> Iterator[Hit]:
for form, value in forms:
if value in name:
yield Hit(mask(location, forms), f"{form} in name")
def _content_hits(location: str, data: bytes, forms: Forms, depth: int) -> Iterator[Hit]:
for form, value in forms:
if value.encode() in data:
yield Hit(mask(location, forms), form)
if depth >= MAX_ARCHIVE_DEPTH or not zipfile.is_zipfile(io.BytesIO(data)):
return
with zipfile.ZipFile(io.BytesIO(data)) as archive:
for member in archive.infolist():
if member.is_dir():
continue
yield from _content_hits(f"{location}!{member.filename}", archive.read(member), forms, depth + 1)
def _walk(root: Path) -> Iterator[tuple[Path, bool]]:
for directory, subdirectories, files in os.walk(root):
yield from ((Path(directory, name), True) for name in subdirectories if Path(directory, name).is_symlink())
yield from ((Path(directory, name), False) for name in files)
def _path_hits(root: Path, path: Path, linked_directory: bool, forms: Forms) -> Iterator[Hit]:
location: Final = str(path.relative_to(root))
if linked_directory:
yield Hit(mask(location, forms), "directory symlink, not scanned")
return
yield from _name_hits(location, location, forms)
yield from _content_hits(location, path.read_bytes(), forms, 0)
def scan(root: Path, secret: str) -> tuple[int, tuple[Hit, ...]]:
forms: Final = secret_forms(secret)
entries: Final = tuple(_walk(root))
hits: Final = tuple(chain.from_iterable(_path_hits(root, path, linked, forms) for path, linked in entries))
return len(entries), hits
def main(argv: tuple[str, ...], environ: Mapping[str, str], out: TextIO) -> int:
if len(argv) != 1:
print(f"usage: assert_secret_absent.py <directory>, with the secret in ${SECRET_VARIABLE}", file=out)
return 2
secret: Final = environ.get(SECRET_VARIABLE, "")
root: Final = Path(argv[0])
if not secret:
print(f"{SECRET_VARIABLE} is empty, nothing to scan for", file=out)
return 0
if not root.is_dir():
print(f"{root} is not a directory, nothing to scan", file=out)
return 0
try:
scanned, hits = scan(root, secret)
except Exception as error:
print(f"scan failed with {type(error).__name__}, treating the results as unsafe", file=out)
return 1
for hit in hits:
print(f"secret found: {hit.location} ({hit.form})", file=out)
if hits:
print(f"{len(hits)} copies of the secret under {root}, refusing to upload", file=out)
return 1
print(f"no copy of the secret in {scanned} files under {root}", file=out)
return 0
if __name__ == "__main__":
sys.exit(main(tuple(sys.argv[1:]), os.environ, sys.stdout))

199
.github/workflows/test-integration.yml vendored Normal file
View file

@ -0,0 +1,199 @@
name: integration
on:
pull_request:
types: [labeled]
schedule:
- cron: "5 8-20/3 * * 1-6"
timezone: America/Los_Angeles
workflow_dispatch:
inputs:
suite:
description: Integration group to run
type: choice
options:
- all
- providers
- extensions
- accounting
- management
- management-replica
- database
- database-replica
- mcp
- security
- cost
- sdk
- browser
default: all
workers:
description: Workers per group, blank keeps each group's default
type: string
default: ""
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name }}-${{ inputs.suite }}-${{ inputs.workers }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
integration:
name: integration ${{ matrix.suite }}
if: github.event_name != 'pull_request' || github.event.label.name == 'run-ci'
environment: integration-license
runs-on: ${{ fromJSON('{"providers":"ubuntu-latest-64-cores","extensions":"ubuntu-latest-64-cores","accounting":"ubuntu-latest-32-cores","management":"ubuntu-latest-32-cores","management-replica":"ubuntu-latest-32-cores"}')[matrix.suite] || 'ubuntu-latest-16-cores' }}
timeout-minutes: 60
strategy:
fail-fast: false
matrix:
suite: ${{ fromJSON((inputs.suite || 'all') == 'all' && '["providers","extensions","accounting","management","management-replica","database","database-replica","mcp","security","cost","sdk","browser"]' || format('["{0}"]', inputs.suite)) }}
env:
UV_PYTHON: "3.12"
INTEGRATION_WORKERS: ${{ inputs.workers || fromJSON('{"providers":"32","extensions":"32","accounting":"16","management":"16","management-replica":"16","cost":"16","mcp":"8","security":"8","database":"4","database-replica":"4"}')[matrix.suite] || '' }}
services:
postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
command: -c max_connections=1000 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000
env:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: circle_test
ports:
- 5432:5432
options: >-
--shm-size 1g
--health-cmd "pg_isready -U postgres"
--health-interval 5s
--health-timeout 5s
--health-retries 20
redis:
image: redis:7-alpine@sha256:7aec734b2bb298a1d769fd8729f13b8514a41bf90fcdd1f38ec52267fbaa8ee6
ports:
- 6379:6379
options: >-
--health-cmd "redis-cli ping"
--health-interval 5s
--health-timeout 5s
--health-retries 20
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: ${{ env.UV_PYTHON }}
- name: Set up uv
timeout-minutes: 3
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
timeout-minutes: 5
uses: ./.github/actions/cache-cargo-build
- name: Cache Prisma binaries
timeout-minutes: 5
uses: ./.github/actions/cache-prisma-binaries
- name: Install dependencies
timeout-minutes: 20
run: .github/scripts/uv_sync_with_retries.sh --frozen --all-groups --all-extras
- name: Install redis-server 7.2.16
timeout-minutes: 3
run: |
mkdir -p "$RUNNER_TEMP/redis-bin"
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
docker cp "${cid}:/usr/local/bin/redis-server" "$RUNNER_TEMP/redis-bin/redis-server"
docker rm "$cid"
"$RUNNER_TEMP/redis-bin/redis-server" --version | grep -q 'v=7.2.16'
echo "$RUNNER_TEMP/redis-bin" >> "$GITHUB_PATH"
- name: Install MCP SDK1 peer
if: matrix.suite == 'mcp'
timeout-minutes: 5
run: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV"
- name: Set up Node.js
if: matrix.suite == 'browser'
timeout-minutes: 5
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
cache: "npm"
cache-dependency-path: |
ui/litellm-dashboard/package-lock.json
tests/e2e/ui/package-lock.json
- name: Install browser dependencies and build the dashboard
if: matrix.suite == 'browser'
timeout-minutes: 20
run: |
(cd ui/litellm-dashboard && npm ci)
(cd tests/e2e/ui && npm ci)
sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive \
node tests/e2e/ui/node_modules/@playwright/test/cli.js install-deps chromium
node tests/e2e/ui/node_modules/@playwright/test/cli.js install chromium
(cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build)
- name: Run integration contracts
timeout-minutes: 45
env:
SUITE: ${{ matrix.suite }}
INTEGRATION_POSTGRES_CONTAINER: ${{ job.services.postgres.id }}
INTEGRATION_REDIS_CONTAINER: ${{ job.services.redis.id }}
LITELLM_LICENSE: ${{ secrets.LITELLM_LICENSE }}
run: |
case "$SUITE" in
*-replica) bash .circleci/scripts/run_integration.sh "${SUITE%-replica}" replica ;;
*) bash .circleci/scripts/run_integration.sh "$SUITE" ;;
esac
- name: Verify every case ran exactly once
if: ${{ !cancelled() && matrix.suite != 'browser' }}
timeout-minutes: 2
run: |
record="$(find test-results -mindepth 2 -maxdepth 2 -name execution.json -print -quit)"
test -n "$record"
mkdir -p "$RUNNER_TEMP/shards/0"
cp "$record" "$RUNNER_TEMP/shards/0/execution.json"
.venv/bin/python -m tests.integration._support.sharding "$RUNNER_TEMP/shards" --shards 1
- name: Report MCP coverage
if: ${{ !cancelled() && matrix.suite == 'mcp' }}
timeout-minutes: 5
run: |
coverage_dir=test-results/integration-mcp/coverage
test -f "$coverage_dir/coverage.txt" && exit 0
export COVERAGE_FILE="$PWD/$coverage_dir/data"
.venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml
.venv/bin/python -m coverage report --rcfile=tests/integration/mcp_coverage.toml > "$coverage_dir/coverage.txt"
.venv/bin/python -m coverage html --rcfile=tests/integration/mcp_coverage.toml -d "$coverage_dir/html"
tail -n 1 "$coverage_dir/coverage.txt"
- name: Check results for the license
id: secret-scan
if: ${{ !cancelled() }}
timeout-minutes: 5
env:
SCANNED_SECRET: ${{ secrets.LITELLM_LICENSE }}
run: python .github/scripts/assert_secret_absent.py test-results
- name: Upload results
if: ${{ !cancelled() && steps.secret-scan.outcome == 'success' }}
timeout-minutes: 10
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: integration-${{ matrix.suite }}-${{ github.run_id }}-${{ github.run_attempt }}
path: test-results/
retention-days: 14

View file

@ -29,6 +29,9 @@ it does not belong to. Queued clients are therefore bucketed by what it takes to
close them, and each bucket is ordered by deadline, so a reap walks the entries
that are due rather than the whole queue.
A bucket whose loop has been collected is dropped, and its clients are left to the
collector, as they are when the queue is full.
The queue holds its clients weakly, so waiting out a grace window never keeps
alive anything the collector would have reclaimed first.
"""
@ -36,6 +39,7 @@ alive anything the collector would have reclaimed first.
import asyncio
import contextlib
import inspect
import itertools
import threading
import time
import weakref
@ -55,6 +59,9 @@ _CLOSABLE_ANYWHERE: Final = "closable-anywhere"
_CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
_BucketKey = str | int
_LOOP_TOKENS: Final[weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, int]] = weakref.WeakKeyDictionary()
_LOOP_TOKEN_COUNTER: Final = itertools.count()
_LOOP_TOKEN_LOCK: Final = threading.Lock()
@dataclass(frozen=True, slots=True)
@ -65,12 +72,12 @@ class _PendingClose:
collector would otherwise have reclaimed first.
``needs_loop`` is set for a client whose close is a coroutine; those can only
be closed from the event loop they were evicted on, recorded in ``loop_id``.
be closed from the event loop they were evicted on, recorded in ``loop_token``.
A client that closes synchronously carries neither constraint.
"""
client_ref: "weakref.ref[object]"
loop_id: int | None
loop_token: int | None
needs_loop: bool
close_after: float
@ -79,16 +86,24 @@ def _bucket_key(pending: _PendingClose) -> _BucketKey:
"""Which reaps can close this entry: any at all, any running a loop, or one loop's."""
if not pending.needs_loop:
return _CLOSABLE_ANYWHERE
if pending.loop_id is None:
if pending.loop_token is None:
return _CLOSABLE_ON_ANY_LOOP
return pending.loop_id
return pending.loop_token
def _running_loop_id() -> int | None:
def running_loop_token() -> int | None:
try:
return id(asyncio.get_running_loop())
loop: Final = asyncio.get_running_loop()
except RuntimeError:
return None
try:
token: Final[int | None] = _LOOP_TOKENS.get(loop)
if token is not None:
return token
with _LOOP_TOKEN_LOCK:
return _LOOP_TOKENS.setdefault(loop, next(_LOOP_TOKEN_COUNTER))
except TypeError:
return id(loop)
def _close_function(client: object) -> Callable[[], object] | None:
@ -186,6 +201,8 @@ class EvictedClientCloser:
self._clock = clock
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
self._dead_loop_tokens: deque[int] = deque() # mutable-ok: loop finalizers append without taking the queue lock
self._watched_loop_tokens: set[int] = set() # mutable-ok: tracks tokens until their dead buckets are discarded
self._pending_count = 0
self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop
self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes
@ -216,13 +233,19 @@ class EvictedClientCloser:
close_fn: Final = _close_function(client)
if close_fn is None:
return
loop_token: Final = running_loop_token()
needs_loop: Final = inspect.iscoroutinefunction(close_fn)
if self._pending_count >= self._max_pending:
return
self._discard_dead_loops()
if self._pending_count >= self._max_pending:
return
if needs_loop and loop_token is not None:
self._watch_loop(loop_token)
self._enqueue(
_PendingClose(
client_ref=weakref.ref(client),
loop_id=_running_loop_id(),
needs_loop=inspect.iscoroutinefunction(close_fn),
loop_token=loop_token,
needs_loop=needs_loop,
close_after=self._clock() + self._grace_seconds,
)
)
@ -236,7 +259,7 @@ class EvictedClientCloser:
if not self._pending_count:
return
now: Final = self._clock()
for pending in self._take_due(_running_loop_id(), now):
for pending in self._take_due(running_loop_token(), now):
client = pending.client_ref()
if client is None:
continue
@ -249,6 +272,29 @@ class EvictedClientCloser:
def pending_count(self) -> int:
return self._pending_count
def _watch_loop(self, loop_token: int) -> None:
if loop_token in self._watched_loop_tokens:
return
try:
weakref.finalize(asyncio.get_running_loop(), self._dead_loop_tokens.append, loop_token)
except TypeError:
return
self._watched_loop_tokens.add(loop_token)
def _discard_dead_loops(self) -> None:
with self._queue_lock:
self._discard_dead_loops_locked()
def _discard_dead_loops_locked(self) -> None:
while self._dead_loop_tokens:
self._discard_dead_loop_token_locked(self._dead_loop_tokens.popleft())
def _discard_dead_loop_token_locked(self, loop_token: int) -> None:
self._watched_loop_tokens.discard(loop_token)
bucket: Final = self._buckets.pop(loop_token, None)
if bucket is not None:
self._pending_count -= len(bucket)
def _enqueue(self, pending: _PendingClose) -> None:
"""Append to the entry's bucket, dropping any dead entries it queues behind.
@ -257,6 +303,7 @@ class EvictedClientCloser:
the front rather than having to be searched for.
"""
with self._queue_lock:
self._discard_dead_loops_locked()
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque())
while bucket and bucket[0].client_ref() is None:
bucket.popleft()
@ -264,9 +311,12 @@ class EvictedClientCloser:
bucket.append(pending)
self._pending_count += 1
def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
def _take_due(self, loop_token: int | None, now: float) -> tuple[_PendingClose, ...]:
buckets = (
(_CLOSABLE_ANYWHERE,) if loop_token is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_token)
)
with self._queue_lock:
self._discard_dead_loops_locked()
return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:

View file

@ -2,10 +2,9 @@
Add the event loop to the cache key, to prevent event loop closed errors.
"""
import asyncio
from typing import Final
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer, running_loop_token
from .in_memory_cache import InMemoryCache
@ -43,24 +42,22 @@ class LLMClientCache(InMemoryCache):
self.evicted_client_closer.schedule(evicted)
self.evicted_client_closer.reap()
def update_cache_key_with_event_loop(self, key):
def update_cache_key_with_event_loop(self, key: str) -> str:
"""
Add the event loop to the cache key, to prevent event loop closed errors.
If none, use the key as is.
"""
try:
event_loop: Final = asyncio.get_running_loop()
stringified_event_loop: Final = str(id(event_loop))
return f"{key}-{stringified_event_loop}"
except RuntimeError: # handle no current running event loop
return key
token: Final = running_loop_token()
return f"{key}-{token}" if token is not None else key
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
def _scoped_key(self, key: str, loop_bound: bool) -> str:
return self.update_cache_key_with_event_loop(key) if loop_bound else key
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, loop_bound: bool = True, **kwargs):
"""``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted."""
if litellm_owned_client:
self.evicted_client_closer.mark_owned(value)
key = self.update_cache_key_with_event_loop(key)
return super().set_cache(key, value, **kwargs)
return super().set_cache(self._scoped_key(key, loop_bound), value, **kwargs)
async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
if litellm_owned_client:
@ -68,13 +65,13 @@ class LLMClientCache(InMemoryCache):
key = self.update_cache_key_with_event_loop(key)
return await super().async_set_cache(key, value, **kwargs)
def get_cache(self, key, **kwargs):
key = self.update_cache_key_with_event_loop(key)
def get_cache(self, key: str, loop_bound: bool = True, **kwargs):
scoped_key: Final = self._scoped_key(key, loop_bound)
self.evicted_client_closer.reap()
return super().get_cache(key, **kwargs)
return super().get_cache(scoped_key, **kwargs)
async def async_get_cache(self, key, **kwargs):
async def async_get_cache(self, key: str, **kwargs):
key = self.update_cache_key_with_event_loop(key)
return await super().async_get_cache(key, **kwargs)

View file

@ -102,18 +102,14 @@ class AgentOps(OpenTelemetry):
"Connection": "keep-alive",
}
client: Final = get_httpx_client()
try:
response: Final = client.post(
url=auth_endpoint,
headers=headers,
json={"api_key": api_key},
timeout=10,
)
response: Final = get_httpx_client().post(
url=auth_endpoint,
headers=headers,
json={"api_key": api_key},
timeout=10,
)
if response.status_code != 200:
raise Exception(f"Failed to fetch auth token: {response.text}")
if response.status_code != 200:
raise Exception(f"Failed to fetch auth token: {response.text}")
return _DECODED_JSON.validate_python(response.json())
finally:
client.close()
return _DECODED_JSON.validate_python(response.json())

View file

@ -35,6 +35,8 @@ from litellm.types.utils import StandardCallbackDynamicParams
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy._types import UserAPIKeyAuth
Span = _Span
else:
Span = Any
@ -126,6 +128,26 @@ def _starts_its_own_trace(span: object, kwargs: Mapping[str, object]) -> bool:
)
def _a_destination_records_the_call(kwargs: Mapping[str, object]) -> bool:
"""Whether a key's or team's ``langfuse_otel`` destination gets this call from the OTel v2 fan-out, which records
every call but one whose logging payload shows it never reached a provider."""
from litellm.integrations.otel.model.metadata import LLMCallEvent
if not _a_destination_carries_the_trace():
return False
call: Final = LLMCallEvent.from_dict(kwargs)
return call.payload is None or call.reached_a_provider
def _a_destination_carries_the_trace() -> bool:
"""Whether this request resolved a key's or team's ``langfuse_otel`` destination, whose fan-out sends it the whole
trace."""
from litellm.integrations.otel.model.config import ExporterOwner
from litellm.integrations.otel.plumbing.context import destination_backends
return ExporterOwner.LANGFUSE_OTEL in destination_backends()
class LangfuseOtelLogger(OpenTelemetry):
def __init__(self, config=None, *args, **kwargs):
# Prevent LangfuseOtelLogger from modifying global environment variables by constructing config manually
@ -541,6 +563,27 @@ class LangfuseOtelLogger(OpenTelemetry):
"""
return None
def _emit_once(
self,
kwargs: dict[str, object], # mutable-ok: the v1 dedupe marker is written into the request's kwargs
*scope: object,
) -> bool:
"""Whether to emit this span: once per request, and never for a call a Langfuse destination records."""
return not _a_destination_records_the_call(kwargs) and super()._emit_once(kwargs, *scope) # pyright: ignore[reportUnknownMemberType] # the v1 hook takes a bare dict
async def async_post_call_failure_hook(
self,
request_data: dict[str, object], # mutable-ok: the v1 hook contract
original_exception: Exception,
user_api_key_dict: "UserAPIKeyAuth",
traceback_str: str | None = None,
) -> None:
"""Record the proxy failure of a request no Langfuse destination carries."""
if not _a_destination_carries_the_trace():
await super().async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # the v1 hook takes a bare dict
request_data, original_exception, user_api_key_dict, traceback_str
)
async def async_service_success_hook(self, *args, **kwargs):
"""
Langfuse should not receive service success logs.

View file

@ -13,9 +13,9 @@ from opentelemetry.trace.status import Status, StatusCode
from litellm.integrations.otel.mappers import resolve_mappers
from litellm.integrations.otel.mappers.base import AttributeMapper, AttrValue, SpanData
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES, LangfuseMapper
from litellm.integrations.otel.mappers.openinference import fit_indexed_messages
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
from litellm.integrations.otel.model.config import ExporterOwner, OpenTelemetryV2Config
from litellm.integrations.otel.model.payloads import (
GuardrailSpanData,
LLMCallSpanData,
@ -34,6 +34,7 @@ from litellm.integrations.otel.model.spans import (
mcp_tool_call_span_name,
service_span_name,
)
from litellm.integrations.otel.plumbing.context import destination_backends
from litellm.integrations.otel.plumbing.events import GenAIEventRecorder
from litellm.integrations.otel.plumbing.providers import to_otel_span_kind
@ -168,6 +169,9 @@ class SpanEmitter:
self._mappers: list[AttributeMapper] = (
list(mappers) if mappers is not None else resolve_mappers(config.mapper_names)
)
self._langfuse: Final = (
None if any(isinstance(mapper, LangfuseMapper) for mapper in self._mappers) else LangfuseMapper()
)
# Bounded LRU (ordered by insertion / most-recent touch). Storing keys
# only — the value is unused — so it behaves like a capped set.
self._emitted: OrderedDict[tuple[str, SpanRole], None] = OrderedDict()
@ -295,17 +299,27 @@ class SpanEmitter:
)
else None
)
mapped: Final = MappingProxyType(
{key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
configured: Final = {key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
for_langfuse: Final = (
dict(
sorted(
self._langfuse.map(data).items(),
key=lambda item: item[0] not in LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
)
)
if self._langfuse is not None and ExporterOwner.LANGFUSE_OTEL in destination_backends()
else _NO_ATTRIBUTES
)
mapped: Final = MappingProxyType({**configured, **for_langfuse})
stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES
reserved: Final = len(stamped_later.keys() - mapped.keys())
overwrites: Final = carried_keys(span) & mapped.keys()
# The Langfuse request metadata sheds first, so a span at its attribute limit loses nothing else to it.
# The Langfuse request metadata, and the Langfuse vocabulary only a destination asked for, shed first,
# so a span at its attribute limit loses nothing else to them.
fitted: Final = fit_indexed_messages(
MappingProxyType({key: value for key, value in mapped.items() if key not in overwrites}),
attribute_budget(span, reserved),
LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
LANGFUSE_REQUEST_METADATA_ATTRIBUTES.union(for_langfuse),
)
for key, value in mapped.items():
if key in overwrites or key in fitted:

View file

@ -198,6 +198,7 @@ class OpenTelemetryV2(CustomLogger):
tracer_provider: TracerProvider | None = None,
logger_provider: LoggerProvider | None = None,
meter_provider: "MeterProvider | None" = None,
console_default: bool = False,
**kwargs: Unpack[_CustomLoggerOptions],
) -> None:
super().__init__(**kwargs)
@ -206,7 +207,7 @@ class OpenTelemetryV2(CustomLogger):
self._tracer_provider: TracerProvider = (
tracer_provider
if tracer_provider is not None
else build_tracer_provider(self.config, tenant_overrides=True)
else build_tracer_provider(self.config, tenant_overrides=True, console_default=console_default)
)
self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME)
self._metrics_recorder = self._init_metrics(meter_provider)
@ -217,7 +218,9 @@ class OpenTelemetryV2(CustomLogger):
mappers=resolve_mappers(self.config.mapper_names),
event_recorder=self._init_events(logger_provider),
)
self._tenant_tracers = TenantTracerCache(self.config, callback_name, LITELLM_TRACER_NAME)
self._tenant_tracers = TenantTracerCache(
self.config, callback_name, LITELLM_TRACER_NAME, console_default=console_default
)
self._open_llm_calls: OrderedDict[str, _LLMCallSpan] = OrderedDict()
self._init_otel_logger_on_litellm_proxy()
@ -541,9 +544,7 @@ class OpenTelemetryV2(CustomLogger):
# gate rejection carries ``is_no_upstream_call`` and gets no span. A cache
# hit never hands off to a provider, so ``cache_hit`` substitutes for the
# upstream stamp, since the payload proves the call completed.
if carrier is None and (
call.is_no_upstream_call or call.payload is None or (not call.upstream_started and not call.cache_hit)
):
if carrier is None and not call.reached_a_provider:
return None
try:
return self._finish_carrier(carrier, call, start_time, end_time)
@ -1091,6 +1092,7 @@ def build_otel_v2_logger(
logger_provider: LoggerProvider | None = None,
meter_provider: "MeterProvider | None" = None,
settings: Mapping[str, object] = MappingProxyType({}),
console_default: bool = False,
) -> OpenTelemetryV2:
return _logger_class(config)(
config=config,
@ -1098,6 +1100,7 @@ def build_otel_v2_logger(
tracer_provider=tracer_provider,
logger_provider=logger_provider,
meter_provider=meter_provider,
console_default=console_default,
**settings,
)

View file

@ -103,6 +103,18 @@ class LangfuseMapper:
),
}
#: Every attribute this mapper can write on a model call.
LLM_CALL_ATTRIBUTES: Final[frozenset[str]] = frozenset(
{
*_LLM_CALL_ATTRS,
*_BLOB_ATTRS,
LANGFUSE_TRACE_NAME,
LANGFUSE_TRACE_USER_ID,
LANGFUSE_TRACE_SESSION_ID,
LANGFUSE_TRACE_TAGS,
}
)
def map(self, data: SpanData) -> AttributeMap:
match data:
case LLMCallSpanData():

View file

@ -121,6 +121,17 @@ class ExporterSpec(LiteLLMBaseModel):
)
def is_unconfigured_placeholder(spec: ExporterSpec) -> bool:
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
No field set is what says the operator asked for nothing: an exporter they did
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
and so does the gated spec ``credential_gated_exporters`` appends, which would
otherwise eat itself when one preset layers onto another.
"""
return not spec.model_fields_set
class _EnvWithoutBareExcludedServices(PydanticBaseSettingsSource):
def __init__(self, settings_cls: type[BaseSettings], env_settings: PydanticBaseSettingsSource) -> None:
super().__init__(settings_cls)

View file

@ -272,6 +272,11 @@ class LLMCallEvent:
trace: TraceControls
session_id: str | None
@property
def reached_a_provider(self) -> bool:
"""Whether the call was handed to a provider or answered from the response cache."""
return not self.is_no_upstream_call and self.payload is not None and (self.upstream_started or self.cache_hit)
@classmethod
def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent:
raw_payload: Final = kwargs.get("standard_logging_object")

View file

@ -47,6 +47,7 @@ from litellm.integrations.otel.mappers.langfuse import (
LANGFUSE_OBSERVATION_OUTPUT,
LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
LANGFUSE_TRACE_NAME,
LangfuseMapper,
)
from litellm.integrations.otel.model.config import (
CaptureMessageContent,
@ -55,6 +56,7 @@ from litellm.integrations.otel.model.config import (
OpenTelemetryV2Config,
capture_message_content_from_env,
captures_span_content,
is_unconfigured_placeholder,
)
from litellm.integrations.otel.model.semconv import (
DB,
@ -490,19 +492,25 @@ def _is_litellm_span(span: ReadableSpan) -> bool:
return scope is not None and scope.name == LITELLM_TRACER_NAME
def without_langfuse_request_metadata(span: ReadableSpan) -> ReadableSpan:
"""``span`` without the request metadata the Langfuse mapper writes for a Langfuse destination.
def without_destination_langfuse_attributes(span: ReadableSpan, langfuse_mapped: bool) -> ReadableSpan:
"""``span`` without the Langfuse attributes only a Langfuse destination asked for.
Only LiteLLM's own spans lose it: any other span keeps every key its own code set.
That is the request metadata the Langfuse mapper writes and, on a model call written by
a logger whose own mappers leave the Langfuse vocabulary out (``langfuse_mapped`` off),
every Langfuse attribute, since that logger wrote them for a key's or team's Langfuse
destination alone. Only LiteLLM's own spans lose them: any other span keeps every key
its own code set.
"""
attributes: Final = span.attributes or _NO_ATTRIBUTES
if not (_is_litellm_span(span) and any(key in LANGFUSE_REQUEST_METADATA_ATTRIBUTES for key in attributes)):
dropped: Final = (
LANGFUSE_REQUEST_METADATA_ATTRIBUTES
if langfuse_mapped or not is_llm_call_span(span)
else LangfuseMapper.LLM_CALL_ATTRIBUTES
)
if not (_is_litellm_span(span) and any(key in dropped for key in attributes)):
return span
return _CountedSpanView(
span,
MappingProxyType(
{key: value for key, value in attributes.items() if key not in LANGFUSE_REQUEST_METADATA_ATTRIBUTES}
),
span, MappingProxyType({key: value for key, value in attributes.items() if key not in dropped})
)
@ -658,6 +666,7 @@ def _for_destination(
span: ReadableSpan,
destination: "OtelDestination",
default_capture: str,
langfuse_mapped: bool,
) -> ReadableSpan:
"""The view of ``span`` a tenant destination receives.
@ -672,13 +681,16 @@ def _for_destination(
and the query string of the request URL, which can hold the same key. A
destination whose effective setting disables span content loses the prompt and
response content too, and a destination for any backend but Langfuse loses the
request metadata the Langfuse mapper writes on LiteLLM's own spans. The span itself
stays, so the tenant still gets the whole trace tree.
Langfuse attributes only a Langfuse destination asked for (see
:func:`without_destination_langfuse_attributes`). The span itself stays, so the
tenant still gets the whole trace tree.
"""
extra: Final = destination.resource_attributes
defaults: Final = destination.resource_defaults
source: Final = (
span if destination.callback_name == ExporterOwner.LANGFUSE_OTEL else without_langfuse_request_metadata(span)
span
if destination.callback_name == ExporterOwner.LANGFUSE_OTEL
else without_destination_langfuse_attributes(span, langfuse_mapped)
)
attributes: Final = source.attributes or _NO_ATTRIBUTES
database: Final = _is_database_span(attributes)
@ -728,8 +740,10 @@ class TenantFanOutSpanProcessor(SpanProcessor):
default_capture_by_backend: Mapping[str, str] = MappingProxyType({}),
pending_drains: int = _MAX_PENDING_DRAINS,
drain_pool: _DrainPool | None = None,
langfuse_mapped: bool = True,
) -> None:
self._operator_sinks: Final = operator_sinks
self._langfuse_mapped: Final = langfuse_mapped
self._excluded_db_systems: Final = excluded_db_systems
self._default_capture: Final = default_capture
self._default_capture_by_backend: Final = default_capture_by_backend
@ -822,7 +836,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
try:
processor.on_end(
_scoped(
_for_destination(span, destination, self._default_capture_for(destination)), destination.span_scope
_for_destination(span, destination, self._default_capture_for(destination), self._langfuse_mapped),
destination.span_scope,
)
)
except Exception as exc: # noqa: BLE001 # one destination's failure must not cost the others their span
@ -1198,17 +1213,19 @@ class _OverriddenBackendFilter(SpanProcessor):
return self._inner.force_flush(timeout_millis)
class _LangfuseRequestMetadataFilter(SpanProcessor):
"""Strip the Langfuse request metadata from LiteLLM's own spans before ``inner`` exports them."""
class _DestinationLangfuseAttributeFilter(SpanProcessor):
"""Strip the Langfuse attributes only a Langfuse destination asked for from LiteLLM's own spans before ``inner``
exports them."""
def __init__(self, inner: SpanProcessor) -> None:
def __init__(self, inner: SpanProcessor, langfuse_mapped: bool) -> None:
self._inner: Final = inner
self._langfuse_mapped: Final = langfuse_mapped
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
self._inner.on_start(span, parent_context)
def on_end(self, span: ReadableSpan) -> None:
self._inner.on_end(without_langfuse_request_metadata(span))
self._inner.on_end(without_destination_langfuse_attributes(span, self._langfuse_mapped))
def shutdown(self) -> None:
self._inner.shutdown()
@ -1491,13 +1508,18 @@ def _spec_processor(
spec: ExporterSpec,
use_simple_processor: bool | None,
strip_message_content: bool,
strip_langfuse_request_metadata: bool,
strip_destination_langfuse_attributes: bool,
langfuse_mapped: bool,
) -> SpanProcessor:
exporting: Final = _processor_for(
_exporter_from_spec(spec),
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
)
scoped: Final = _LangfuseRequestMetadataFilter(exporting) if strip_langfuse_request_metadata else exporting
scoped: Final = (
_DestinationLangfuseAttributeFilter(exporting, langfuse_mapped)
if strip_destination_langfuse_attributes
else exporting
)
return _MessageContentFilter(scoped) if strip_message_content else scoped
@ -1530,6 +1552,7 @@ def build_tracer_provider(
use_simple_processor: bool | None = None,
tenant_overrides: bool = False,
strip_message_content: bool = False,
console_default: bool = False,
) -> TracerProvider:
"""Build the shared :class:`TracerProvider`.
@ -1548,11 +1571,16 @@ def build_tracer_provider(
``config.langfuse_span_scope`` narrows the exporter owned by ``langfuse_otel``
alone; a collector or any other backend in the same config keeps the full tree.
In a config that writes the Langfuse vocabulary, every exporter but ``langfuse_otel``'s
goes without the request metadata the Langfuse mapper writes on LiteLLM's own spans.
Every exporter but ``langfuse_otel``'s goes without the Langfuse attributes only a
Langfuse destination asked for (see :func:`without_destination_langfuse_attributes`),
which any logger writes for a request with a Langfuse destination.
``strip_message_content`` removes prompt and response content before every
exporter, for a routed provider whose spans can carry content another
destination asked for.
The console spec ``config._normalize`` folds in when nothing was configured
prints only under ``console_default``, which the ``otel`` callback passes: a
vendor preset or the proxy's background global logger was never asked to print.
"""
provider: Final = TracerProvider(resource=build_resource(config))
provider.sampler = _RequestExclusionSampler(provider.sampler)
@ -1564,19 +1592,20 @@ def build_tracer_provider(
provider.add_span_processor(_processor_for(exporter, use_simple_processor))
return provider
langfuse_mapped: Final = _maps_langfuse(config)
# ``config._normalize`` guarantees at least one spec (it folds the top-level
# ``exporter``/``endpoint``/``headers`` fields in when ``exporters`` is empty).
writes_langfuse: Final = "langfuse" in config.mapper_names or any(
spec.owner is ExporterOwner.LANGFUSE_OTEL for spec in config.exporters
)
for spec in config.exporters:
if spec.requires_headers and not spec.headers:
continue
if is_unconfigured_placeholder(spec) and not console_default:
continue
processor = _spec_processor(
spec,
use_simple_processor,
strip_message_content,
strip_langfuse_request_metadata=writes_langfuse and spec.owner is not ExporterOwner.LANGFUSE_OTEL,
strip_destination_langfuse_attributes=spec.owner is not ExporterOwner.LANGFUSE_OTEL,
langfuse_mapped=langfuse_mapped,
)
owner = spec.owner.value if tenant_overrides and spec.owner is not None else None
scope = _operator_scope(config, spec)
@ -1613,8 +1642,9 @@ def attach_tenant_fan_out(
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. ``excluded_db_systems`` only filters what the fan-out
delivers, never the operator's own exporters.
rather than twice. The first is the config of the logger that owns ``provider``,
whose model calls the fan-out carries. ``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)):
@ -1625,10 +1655,16 @@ def attach_tenant_fan_out(
excluded_db_systems=excluded_db_systems,
default_capture=_operator_capture_mode(configs),
default_capture_by_backend=MappingProxyType(dict(_owned_capture_modes(configs))),
langfuse_mapped=any(_maps_langfuse(config) for config in configs[:1]),
)
)
def _maps_langfuse(config: OpenTelemetryV2Config) -> bool:
"""Whether a logger built from ``config`` writes the Langfuse vocabulary on every model call."""
return "langfuse" in config.mapper_names
def _operator_capture_mode(configs: Sequence[OpenTelemetryV2Config]) -> str:
"""The mode of the operator's generic ``otel`` callback, so ``callback_settings.otel`` counts as the global default"""
return next(

View file

@ -163,10 +163,12 @@ class TenantTracerCache:
config: OpenTelemetryV2Config,
callback_name: str | None,
tracer_name: str,
console_default: bool = False,
) -> None:
self._config = config
self._callback_name = callback_name
self._tracer_name = tracer_name
self._console_default = console_default
# Guards the three mutable structures below: ``pre_call`` can run on
# thread-pool workers concurrently with the event loop, so cache
# updates, span counts, and retirement must be atomic.
@ -284,6 +286,7 @@ class TenantTracerCache:
built: Final = build_tracer_provider(
self._routed_config(credential_headers, project_headers, endpoint, service_name),
strip_message_content=not self._config.capture_span_content,
console_default=self._console_default,
)
self._providers[cache_key] = built
return built

View file

@ -111,7 +111,7 @@ _REQUIRED_HEADERS_BY_CALLBACK: Final[Mapping[str, frozenset[str]]] = MappingProx
_NO_ATTRS: Final[Mapping[str, str]] = MappingProxyType({})
def _arize_project() -> Mapping[str, str]:
def _arize_project(_params: StandardCallbackDynamicParams) -> Mapping[str, str]:
from litellm.integrations.arize.arize import ArizeLogger
project: Final = ArizeLogger.get_arize_config().project_name
@ -124,10 +124,27 @@ def _arize_default_project() -> Mapping[str, str]:
return MappingProxyType({"model_id": ARIZE_DEFAULT_PROJECT})
#: Resource attributes the operator chose for a backend, applied to every export the
#: team receives over the span's own resource: ``ARIZE_PROJECT_NAME`` for Arize.
_RESOURCE_ATTRIBUTES_BY_CALLBACK: Final[Mapping[str, Callable[[], Mapping[str, str]]]] = MappingProxyType(
{"arize": _arize_project}
def _stands_in_for_the_legacy_langfuse_logger() -> bool:
"""Whether a Langfuse destination replaces the older Langfuse logger's copy, which a proxy without Langfuse keys
of its own sends for a key's or team's ``langfuse_otel``."""
return not (os.environ.get("LANGFUSE_PUBLIC_KEY") and os.environ.get("LANGFUSE_SECRET_KEY"))
def _langfuse_environment(params: StandardCallbackDynamicParams) -> Mapping[str, str]:
"""The key's or team's ``langfuse_environment``, else ``LANGFUSE_TRACING_ENVIRONMENT``, where the older Langfuse
logger wrote it."""
from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes
if not _stands_in_for_the_legacy_langfuse_logger():
return _NO_ATTRS
environment: Final = params.get("langfuse_environment") or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
return MappingProxyType({LangfuseSpanAttributes.LANGFUSE_ENVIRONMENT.value: environment} if environment else {})
#: Resource attributes a backend gets on every export the team receives, over the
#: span's own resource: ``ARIZE_PROJECT_NAME`` for Arize, the Langfuse environment.
_RESOURCE_ATTRIBUTES_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], Mapping[str, str]]]] = (
MappingProxyType({"arize": _arize_project, "langfuse_otel": _langfuse_environment})
)
#: Resource attributes a backend needs on every export to accept it, filled only where
@ -137,12 +154,14 @@ _RESOURCE_DEFAULTS_BY_CALLBACK: Final[Mapping[str, Callable[[], Mapping[str, str
)
def _resource_attributes(callback_name: str, service_name: str | None) -> Mapping[str, str]:
def _resource_attributes(
callback_name: str, service_name: str | None, params: StandardCallbackDynamicParams
) -> Mapping[str, str]:
backend: Final = _RESOURCE_ATTRIBUTES_BY_CALLBACK.get(callback_name)
return MappingProxyType(
{
**({"service.name": service_name} if service_name else {}),
**(backend() if backend is not None else {}),
**(backend(params) if backend is not None else {}),
}
)
@ -214,7 +233,7 @@ def destination_for(
return OtelDestination(
endpoint=endpoint,
headers=MappingProxyType(dict(headers)),
resource_attributes=_resource_attributes(callback_name, service_name),
resource_attributes=_resource_attributes(callback_name, service_name, params),
resource_defaults=_resource_defaults(callback_name),
callback_name=callback_name,
protocol=protocol,

View file

@ -3,7 +3,7 @@
from collections.abc import Iterable
from typing import Final
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec, is_unconfigured_placeholder
def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]:
@ -35,14 +35,3 @@ def credential_gated_exporters(
*(spec for spec in exporters if not is_unconfigured_placeholder(spec)),
ExporterSpec(owner=owner, requires_headers=True),
)
def is_unconfigured_placeholder(spec: "ExporterSpec") -> bool:
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
No field set is what says the operator asked for nothing: an exporter they did
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
and so does the gated spec this module appends, which would otherwise eat itself
when one preset layers onto another.
"""
return not spec.model_fields_set

View file

@ -5640,7 +5640,7 @@ def build_custom_logger_compatible_class(
return callback
otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
otel_logger_v2: Final = build_otel_v2_logger(
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings, console_default=True
)
_in_memory_loggers.append(otel_logger_v2)
_maybe_auto_initialize_arize_phoenix(_in_memory_loggers)

View file

@ -22,6 +22,7 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items
@ -196,7 +197,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
params: Final = cast( # cast-ok: the base signature leaves the params dict untyped
"ResponsesAPIOptionalRequestParams", response_api_optional_request_params
)
hoisted: Final = hoist_additional_tools(input, params.get("tools"))
hoisted: Final = hoist_additional_tools(
input, params.get("tools"), drop_params=drops_unsupported_params(litellm_params.drop_params)
)
normalized_input, rewritten_types = normalize_codex_input_items(hoisted.input)
if rewritten_types:
verbose_logger.warning(

View file

@ -1205,6 +1205,17 @@ class AsyncHTTPHandler:
# Mirrors LiteLLMAiohttpTransport._background_close_tasks.
_finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes
@classmethod
async def _close_client_at_finalization(cls, client: httpx.AsyncClient) -> None:
try:
await client.aclose()
except (OSError, RuntimeError) as exc:
verbose_logger.debug("Error closing client at finalization: %s", exc)
finally:
task: Final = asyncio.current_task()
if task is not None:
cls._finalizer_close_tasks.discard(task)
@classmethod
def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None:
cls._finalizer_close_tasks.discard(task)
@ -1272,8 +1283,8 @@ class AsyncHTTPHandler:
# here is the cross-loop path the transport refuses.
self._dispose_wrapped_aiohttp_session()
return
task: Final = loop.create_task(self._client.aclose())
cls: Final = type(self)
task: Final = loop.create_task(cls._close_client_at_finalization(self._client))
cls._finalizer_close_tasks.add(task)
task.add_done_callback(cls._on_finalizer_close_done)
except Exception:
@ -1973,7 +1984,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
cache = LLMClientCache()
setattr(litellm, "in_memory_llm_clients_cache", cache)
_cached_client: Final = cache.get_cache(_cache_key_name)
_cached_client: Final = cache.get_cache(_cache_key_name, loop_bound=False)
if _cached_client:
return _cached_client
@ -1989,6 +2000,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
value=_new_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=True,
loop_bound=False,
)
return _new_client

View file

@ -275,6 +275,10 @@ def build_output_token_limit_response(
return e.response.headers, (_aiter_once(chunk) if is_async else _iter_once(chunk))
def _is_loop_bound(client_initialization_params: Mapping[str, object]) -> bool:
return client_initialization_params.get("is_async") is not False
class BaseOpenAILLM:
"""
Base class for OpenAI LLMs for getting their httpx clients and SSL verification settings
@ -289,7 +293,9 @@ class BaseOpenAILLM:
client_initialization_params=client_initialization_params,
client_type=client_type,
)
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(_cache_key)
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(
_cache_key, loop_bound=_is_loop_bound(client_initialization_params)
)
return _cached_client
@staticmethod
@ -330,6 +336,7 @@ class BaseOpenAILLM:
value=openai_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=litellm_owned_client,
loop_bound=_is_loop_bound(client_initialization_params),
)
@staticmethod

View file

@ -287,6 +287,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
streaming_sampling_rate: int | None = None,
streaming_end_of_stream_only: bool | None = None,
streaming_buffer_release_on_scan: bool | None = None,
streaming_scan_in_background: bool | None = None,
async_handler: AsyncHTTPHandler | None = None,
**kwargs,
):
@ -301,6 +302,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"streaming_sampling_rate": streaming_sampling_rate,
"streaming_end_of_stream_only": streaming_end_of_stream_only,
"streaming_buffer_release_on_scan": streaming_buffer_release_on_scan,
"streaming_scan_in_background": streaming_scan_in_background,
}
)
)
@ -366,6 +368,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.streaming_sampling_rate = streaming_params.streaming_sampling_rate
self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only
self.streaming_buffer_release_on_scan = streaming_params.streaming_buffer_release_on_scan
self.streaming_scan_in_background = streaming_params.streaming_scan_in_background
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
super().update_in_memory_litellm_params(litellm_params)

View file

@ -137,6 +137,13 @@ def _chunk_choices(item: object) -> Sequence[object]:
return choices
def _ends_stream(endpoint_translation: object, item: object) -> bool:
has_ended: Final[Callable[[Sequence[object]], bool] | None] = getattr(
endpoint_translation, "_check_streaming_has_ended", None
)
return has_ended is not None and has_ended((item,))
def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]:
return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0)
@ -1179,6 +1186,13 @@ class UnifiedLLMGuardrails(CustomLogger):
if buffer_until_moderated and not release_on_scan:
end_of_stream_only = True
# Run each sampled scan as a background task so the stream keeps flowing; a failed scan ends
# the stream before the next chunk, and the terminal frames wait for the final verdict.
# Buffered and end-only streams already scan before releasing.
scan_in_background: Final = bool(_streaming_setting("streaming_scan_in_background", False)) and not (
buffer_until_moderated or end_of_stream_only
)
if guardrail_to_apply is None:
async for item in response:
yield item
@ -1242,6 +1256,10 @@ class UnifiedLLMGuardrails(CustomLogger):
last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round
tool_calls_in_flight = False # rebind-ok: tracks the latest scan key's unscanned tool calls
verdict_settled = False # rebind-ok: set once the end-of-stream scan or a block owns the verdict
pending_scan: asyncio.Task[object] | None = None # rebind-ok: the one background scan in flight
# Background scans read copies, because the translation rewrites the chunks it scans in place.
# Each chunk is copied once, when a scan first covers it.
scanned_responses: Final[list[object]] = [] # mutable-ok: grows with the scanned stream
try:
async for item in response:
@ -1291,10 +1309,21 @@ class UnifiedLLMGuardrails(CustomLogger):
withheld_items.append(item)
continue
if scan_in_background and (
pending_end_of_stream_items or _ends_stream(mappings[CallTypes(call_type)](), item)
):
pending_end_of_stream_items.append(item)
continue
if pending_scan is not None and pending_scan.done():
if pending_scan.exception() is not None:
break # the end-of-stream handling below ends the stream with this scan's verdict
pending_scan = None
# Process chunk based on sampling rate
if buffer_until_moderated:
withheld_items.append(item)
if chunk_counter % sampling_rate == 0:
if chunk_counter % sampling_rate == 0 and pending_scan is None:
endpoint_translation = mappings[CallTypes(call_type)]()
scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
if scan_key is not None:
@ -1327,18 +1356,30 @@ class UnifiedLLMGuardrails(CustomLogger):
guardrail_to_apply.guardrail_name,
)
if scan_in_background:
scanned_responses.extend(copy.deepcopy(responses_so_far[len(scanned_responses) :]))
scan = endpoint_translation.process_output_streaming_response(
responses_so_far=list(scanned_responses) if scan_in_background else responses_so_far,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=request_data.get("litellm_logging_obj"),
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
if scan_in_background:
pending_scan = asyncio.create_task(scan)
if scan_key is not None:
last_scan_key = scan_key
chunks_yielded = True
responses_yielded.append(item)
yield item
continue
original_items = (
tuple(copy.deepcopy(withheld_items)) if buffer_until_moderated else (copy.deepcopy(item),)
)
try:
await endpoint_translation.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=request_data.get("litellm_logging_obj"),
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
await scan
except ModifyResponseException as e:
verdict_settled = True
if e.original_response is None:
@ -1405,6 +1446,10 @@ class UnifiedLLMGuardrails(CustomLogger):
endpoint_translation = mappings[CallTypes(call_type)]()
if pending_scan is not None:
await asyncio.wait((pending_scan,))
background_scan_failed: Final = pending_scan is not None and pending_scan.exception() is not None
buffered_items: Final = (
tuple(copy.deepcopy(withheld_items))
if buffer_until_moderated and release_on_scan and not end_of_stream_only
@ -1414,8 +1459,10 @@ class UnifiedLLMGuardrails(CustomLogger):
)
end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
verdict_settled = True
if _is_redundant_scan(end_scan_key, last_scan_key) and not _owes_ended_scan(
endpoint_translation, guardrail_to_apply, responses_so_far, last_scan_key
if (
_is_redundant_scan(end_scan_key, last_scan_key)
and not _owes_ended_scan(endpoint_translation, guardrail_to_apply, responses_so_far, last_scan_key)
and not background_scan_failed
):
verbose_proxy_logger.debug(
"Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
@ -1430,6 +1477,8 @@ class UnifiedLLMGuardrails(CustomLogger):
try:
with anyio.CancelScope(shield=chunks_yielded):
if pending_scan is not None:
pending_scan.result() # a failed background scan raises into the handlers below
await endpoint_translation.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail_to_apply,
@ -1471,6 +1520,11 @@ class UnifiedLLMGuardrails(CustomLogger):
):
yield error_item
except (GeneratorExit, asyncio.CancelledError):
if pending_scan is not None:
with anyio.CancelScope(shield=True):
await asyncio.wait((pending_scan,))
if isinstance(pending_scan.exception(), (ModifyResponseException, HTTPException)):
verdict_settled = True # rebind-ok: the background scan already reached a verdict
translation_class: Final = None if call_type is None else mappings.get(CallTypes(call_type))
if (
chunks_yielded
@ -1487,3 +1541,8 @@ class UnifiedLLMGuardrails(CustomLogger):
request_data=typed_request_data,
)
raise
except Exception:
if pending_scan is not None:
pending_scan.cancel()
await asyncio.wait((pending_scan,))
raise

View file

@ -46,6 +46,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
streaming_sampling_rate=streaming_params.streaming_sampling_rate,
streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan,
streaming_scan_in_background=streaming_params.streaming_scan_in_background,
timeout=litellm_params.timeout,
)
litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback)

View file

@ -89,6 +89,8 @@ def _unique_tools(tools: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, o
def hoist_additional_tools(
input: str | ResponseInputParam,
tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None,
*,
drop_params: bool,
) -> HoistedAdditionalTools:
existing: Final = tuple(tools or ())
if isinstance(input, str):
@ -103,6 +105,8 @@ def hoist_additional_tools(
len(items),
)
remaining_input: Final = [item for item in input if not _is_additional_tools_item(item)]
if not drop_params:
return HoistedAdditionalTools(input=remaining_input, tools=(*existing, *hoisted), hoisted=hoisted)
unique: Final = tuple(
cast("ALL_RESPONSES_API_TOOL_PARAMS", tool) # cast-ok: a merged namespace is rebuilt from the same tool dicts
for tool in _unique_tools((*existing, *hoisted))

View file

@ -114,7 +114,7 @@ def with_provider_file_ids(
}
def _with_hoisted_additional_tools(responses_request: Mapping[str, object]) -> Mapping[str, object]:
def _with_hoisted_additional_tools(responses_request: Mapping[str, object], drop_params: bool) -> Mapping[str, object]:
responses_input: Final = responses_request.get("input")
if not isinstance(responses_input, list):
return responses_request
@ -122,7 +122,7 @@ def _with_hoisted_additional_tools(responses_request: Mapping[str, object]) -> M
top_level_tools: Final = cast( # cast-ok: tools are forwarded as sent, as inference does
"tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]", _ITEMS.validate_python(responses_request.get("tools") or ())
)
hoisted: Final = hoist_additional_tools(input_items, top_level_tools)
hoisted: Final = hoist_additional_tools(input_items, top_level_tools, drop_params=drop_params)
if not hoisted.hoisted:
return responses_request
return {**responses_request, "input": hoisted.input, "tools": list(hoisted.tools)}
@ -196,7 +196,7 @@ async def chat_count_request(
chat_body: Final = await _with_previous_turns(
responses_batch_body_to_chat_body(
_with_compaction_as_summary_text(
_with_emulated_file_search(_with_hoisted_additional_tools(responses_request))
_with_emulated_file_search(_with_hoisted_additional_tools(responses_request, policy.drop_params))
),
custom_llm_provider,
policy.drop_params,

View file

@ -9,6 +9,7 @@ from typing import Final
from pydantic import TypeAdapter
import litellm
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools
from litellm.responses.litellm_completion_transformation.client_tool_search import bridge_client_tool_search
from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary
@ -73,7 +74,11 @@ class LiteLLMCompletionTransformationHandler:
| BaseResponsesAPIStreamingIterator
| Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]
):
hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools"))
hoisted: Final = hoist_additional_tools(
input,
responses_api_request.get("tools"),
drop_params=drops_unsupported_params(_REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("drop_params")),
)
previous_response_id: Final = responses_api_request.get("previous_response_id")
if _is_async and previous_response_id and has_client_tool_search(hoisted.tools):
return self._async_response_api_handler_after_earlier_searches(

View file

@ -709,6 +709,12 @@ class BedrockGuardrailStreamingParams(LiteLLMBaseModel):
"and release the withheld chunks once the scan passes, instead of holding everything to end of stream. "
"Flagged content is never released. Ignored when streaming_end_of_stream_only is true.",
)
streaming_scan_in_background: bool = Field(
default=False,
description="When not buffering and not end-of-stream-only, run each sampled ApplyGuardrail scan in the "
"background so the sampled chunk and the ones after it stream while the scan runs. One scan runs at a time. "
"A flagged scan ends the stream before the next chunk, so content streamed while it ran reaches the client.",
)
@classmethod
def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams":

View file

@ -27,8 +27,11 @@ tests. A product bug the test exposes is `pytest.skip("BUG: <symptom>")` at the
fix in the test and not a deletion. Needs no network, proxy, DB or Redis: `tests/unit`
CircleCI's integration groups run every file here except the OCI and sandbox tests, which no CI runs, and
the files in `GITHUB_FILES` in `run.py`. CircleCI runs on main, not on pull requests, so a new file here
first runs after its PR merges. Wanting a file to run before merge is no reason to change CI, and a new
the files in `GITHUB_FILES` in `run.py`. `test-integration.yml` runs the same groups through the same
`run_integration.sh` on GitHub Actions, on main every three hours from 8am to 8pm Pacific Monday to Saturday,
and on a pull request when someone adds the `run-ci` label. Neither runs every pull request, so a new file here
first runs on main after its PR merges, or earlier if the PR gets `run-ci`. Wanting a file to run before merge
is no reason to change CI, and a new
dedicated workflow for it is the same mistake as a shared shard. Write the file, leave `GITHUB_FILES` and
every workflow alone, and tell the user it first runs on main. `GITHUB_FILES` is only for a file that
needs what CircleCI's groups cannot give it, like the images `image-scan.yml` builds or the plain pytest

View file

@ -1,6 +1,6 @@
# Integration contracts
These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite, and `.github/workflows/test-integration.yml` runs the same groups through `run_integration.sh` on GitHub Actions for `run-ci` pull requests and scheduled main runs. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
The ROI database contracts in `database/test_roi_observed.py` run as plain pytest outside `run_integration.sh` in CircleCI's `roi-database` Postgres job. The job uploads coverage with the `roi-postgres` flag. These contracts own temporary databases and script only the external provider transport

View file

@ -210,6 +210,8 @@ def langfuse_rig(
sink_reply: Callable[[Request], Reply] | None = None,
overrides: Mapping[str, str] | None = None,
model_list: Callable[[str], Sequence[Mapping[str, object]]] | None = None,
remove_environment: tuple[str, ...] = (),
operator_keys: bool = True,
) -> Iterator[LangfuseRig]:
def upstream(request: Request) -> Reply:
if request.method == "GET":
@ -253,8 +255,11 @@ def langfuse_rig(
gateway,
tmp_path,
{
"LANGFUSE_PUBLIC_KEY": "pk-lf-integration",
"LANGFUSE_SECRET_KEY": "sk-lf-integration",
**(
{"LANGFUSE_PUBLIC_KEY": "pk-lf-integration", "LANGFUSE_SECRET_KEY": "sk-lf-integration"}
if operator_keys
else {}
),
"LANGFUSE_HOST": collector.url,
"LITELLM_OTEL_V2": "1",
**{
@ -264,6 +269,11 @@ def langfuse_rig(
},
config=path,
workers=2,
remove_environment=(
remove_environment
if operator_keys
else (*remove_environment, "LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY")
),
) as owned:
yield LangfuseRig(owned.gateway, provider, collector, owned)

View file

@ -0,0 +1,98 @@
import json
import os
import uuid
from collections.abc import Mapping
from pathlib import Path
from typing import Final
import pytest
from _langfuse_otel_metadata_support import exported_llm_span, langfuse_rig, scenario_identity, watched_spans
from integration._support.client import Gateway, eventually
CONSOLE_SPAN: Final = '"span_id": "0x'
INHERITED_OTEL_ENV: Final = tuple(name for name in os.environ if name.startswith("OTEL_"))
@pytest.mark.parametrize(
("overrides", "printed"),
[({}, False), ({"OTEL_EXPORTER": "console"}, True)],
ids=["exporter-unset", "operator-asked-for-console"],
)
def test_langfuse_otel_prints_spans_to_stdout_only_when_the_operator_asks(
gateway: Gateway, tmp_path: Path, overrides: Mapping[str, str], printed: bool
) -> None:
marker: Final = "lf-console-" + uuid.uuid4().hex
with (
langfuse_rig(
gateway,
tmp_path,
{},
marker,
overrides={"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": "span_only", **overrides},
remove_environment=INHERITED_OTEL_ENV,
) as rig,
rig.candidate.scenario() as scenario,
):
key, identity = scenario_identity(scenario, rig.provider, team=False)
response: Final = rig.candidate.request(
"POST",
"/v1/chat/completions",
{"model": identity["model"], "messages": [{"role": "user", "content": marker}]},
key=key,
)
assert response.status_code == 200, response.text
exported_llm_span(rig.collector, response.json()["id"])
on_stdout: Final = CONSOLE_SPAN in rig.proxy.log.read_text()
assert on_stdout is printed, f"span JSON on proxy stdout: {on_stdout}, see {rig.proxy.log}"
@pytest.mark.parametrize("callbacks", [[], ["langfuse_otel"]], ids=["key-only", "callback-listed"])
def test_key_langfuse_without_operator_keys_prints_no_spans(
gateway: Gateway, tmp_path: Path, callbacks: list[str]
) -> None:
marker: Final = "lf-console-key-" + uuid.uuid4().hex
with (
langfuse_rig(
gateway,
tmp_path,
{"callbacks": callbacks},
marker,
overrides={"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": "span_only"},
remove_environment=INHERITED_OTEL_ENV,
operator_keys=False,
) as rig,
rig.candidate.scenario() as scenario,
):
key: Final = scenario.key(
key_alias="lf-console-key-" + uuid.uuid4().hex,
metadata={
"logging": [
{
"callback_name": "langfuse_otel",
"callback_type": "success",
"callback_vars": {
"langfuse_public_key": "pk-lf-key",
"langfuse_secret_key": "sk-lf-key",
"langfuse_host": rig.collector.url,
},
}
]
},
)
model: Final = scenario.model(api_base=rig.provider.url + "/v1")
response: Final = rig.candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": marker}]},
key=key,
)
assert response.status_code == 200, response.text
eventually(
watched_spans(rig.collector), lambda spans: any(marker in json.dumps(span) for span in spans), seconds=25
)
on_stdout: Final = CONSOLE_SPAN in rig.proxy.log.read_text()
assert not on_stdout, f"span JSON on proxy stdout, see {rig.proxy.log}"

View file

@ -0,0 +1,224 @@
"""Without operator Langfuse keys on OTel v2, Langfuse logging the OTel v2 fan-out does not carry still reaches the
tenant's own Langfuse project: a key's failure-only entry on a failed call, a config ``default_team_settings`` entry,
and a key destination's call the proxy rejected before any provider. A ``success_and_failure`` key entry sends that
project exactly one copy of a failed call."""
from __future__ import annotations
import json
import uuid
from collections.abc import Callable, Iterator
from typing import Final, Literal
import pytest
from integration._support.client import Gateway, gateway_from_environment, held
from integration._support.otlp_sink import (
AuditConfigWriter,
Span,
SpanSinks,
langfuse_callback_variables,
recorded_spans,
spans_for_trace,
)
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
OPERATION: Final = "gen_ai.operation.name"
TAGS: Final = ["agent:orders", "run:r1"]
CONFIG_TEAM: Final = f"lf-config-team-{uuid.uuid4().hex[:8]}"
FAIL: Final = "please-fail"
def _chat(request: Request) -> Reply:
if FAIL.encode() in request.body:
return Reply(
status=500, body=json.dumps({"error": {"message": "upstream boom", "type": "server_error"}}).encode()
)
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
}
).encode()
)
@pytest.fixture(scope="module")
def provider() -> Iterator[Wire]:
with wire_server(_chat) as wire:
yield wire
@pytest.fixture(scope="module")
def proxy(
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
tmp_path_factory: pytest.TempPathFactory,
) -> Iterator[Gateway]:
directory: Final = tmp_path_factory.mktemp("langfuse-legacy-destinations")
settings: Final[dict[str, JsonValue]] = {
"callbacks": [],
"default_team_settings": [
{
"team_id": CONFIG_TEAM,
"success_callback": ["langfuse_otel"],
"callback_vars": dict(langfuse_callback_variables(audit_sinks)),
}
],
}
with (
gateway_from_environment() as gateway,
owned_proxy_process(
gateway,
directory,
{
"LITELLM_OTEL_V2": "1",
"OTEL_BSP_SCHEDULE_DELAY": "300",
"OTEL_EXPORTER": "otlp_http",
"OTEL_ENDPOINT": audit_sinks.operator,
},
config=otel_audit_config(directory, settings),
remove_environment=("LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_HOST"),
workers=2,
) as owned,
):
yield owned.gateway
def _tenant_spans(url: str, since: int, trace_id: str) -> tuple[Span, ...]:
return spans_for_trace(recorded_spans(url, since)[1], trace_id)
def _tags(value: object) -> object:
"""The tags a span carries: the legacy logger writes them as a JSON string, OTel v2 as a list."""
return json.loads(value) if isinstance(value, str) else value
def _settled(spans: tuple[Span, ...]) -> bool:
"""The model span and the request's root span have both arrived."""
return any(OPERATION in span["attributes"] for span in spans) and any(
str(span["name"]).startswith("POST ") for span in spans
)
def _logging_entry(audit_sinks: SpanSinks, callback_type: str) -> list[JsonValue]:
return [
{
"callback_name": "langfuse_otel",
"callback_type": callback_type,
"callback_vars": langfuse_callback_variables(audit_sinks),
}
]
def _chat_status(proxy: Gateway, model: str, key: str, trace_id: str, content: str) -> int:
body: Final[dict[str, JsonValue]] = {
"model": model,
"messages": [{"role": "user", "content": content}],
"metadata": {"tags": [*TAGS], "session_id": f"sess-{trace_id}", "trace_user_id": "user-42"},
}
return proxy.request(
"POST",
"/v1/chat/completions",
body,
key=key,
headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
).status_code
def _trace_fields(spans: tuple[Span, ...]) -> list[tuple[object, object]]:
return [
(span["attributes"]["user.id"], _tags(span["attributes"].get("langfuse.trace.tags")))
for span in spans
if "user.id" in span["attributes"]
]
def _carries_the_trace_fields_once(spans: tuple[Span, ...]) -> bool:
return _trace_fields(spans) == [("user-42", TAGS)]
def _one_failed_call(spans: tuple[Span, ...]) -> bool:
return (
_settled(spans)
and len([span for span in spans if OPERATION in span["attributes"]]) == 1
and "litellm_request" not in {span["name"] for span in spans}
)
def _send(
proxy: Gateway,
provider: Wire,
audit_sinks: SpanSinks,
setup: Literal["failure", "config_team", "both"],
expected: Callable[[tuple[Span, ...]], bool],
) -> tuple[int, tuple[Span, ...]]:
logging_entry: Final = _logging_entry(audit_sinks, "failure" if setup == "failure" else "success_and_failure")
with proxy.scenario() as scenario:
model: Final = scenario.model(api_base=provider.url + "/v1", num_retries=0)
alias: Final = f"lf-legacy-{setup}-{uuid.uuid4().hex}"
key: Final = (
scenario.key(key_alias=alias, team_id=scenario.team(team_id=CONFIG_TEAM))
if setup == "config_team"
else scenario.key(key_alias=alias, metadata={"logging": logging_entry})
)
since: Final = recorded_spans(audit_sinks.tenant)[0]
trace_id: Final = uuid.uuid4().hex
content: Final = f"lf-legacy-{setup}-{trace_id}" + ("" if setup == "config_team" else f" {FAIL}")
status: Final = _chat_status(proxy, model, key, trace_id, content)
tenant: Final = held(
lambda: _tenant_spans(audit_sinks.tenant, since, trace_id), expected, holding=3, seconds=40
)
return status, tenant
@pytest.mark.parametrize("setup", ["failure", "config_team"])
def test_a_tenant_langfuse_project_without_a_destination_still_gets_the_trace(
setup: Literal["failure", "config_team"], proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
) -> None:
status, tenant = _send(proxy, provider, audit_sinks, setup, _carries_the_trace_fields_once)
assert status == (500 if setup == "failure" else 200)
assert _trace_fields(tenant) == [("user-42", TAGS)], [span["name"] for span in tenant]
def test_a_key_destinations_call_rejected_before_any_provider_still_gets_the_trace_fields(
proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
) -> None:
with proxy.scenario() as scenario:
model: Final = scenario.model(api_base=provider.url + "/v1", num_retries=0)
key: Final = scenario.key(
key_alias=f"lf-legacy-rejected-{uuid.uuid4().hex}",
metadata={"logging": _logging_entry(audit_sinks, "success_and_failure")},
rpm_limit=1,
)
since: Final = recorded_spans(audit_sinks.tenant)[0]
statuses: Final = {
trace_id: _chat_status(proxy, model, key, trace_id, f"lf-legacy-rejected-{trace_id}")
for trace_id in [uuid.uuid4().hex for _ in range(4)]
}
rejected: Final = [trace_id for trace_id, status in statuses.items() if status == 429]
assert rejected, statuses
tenant: Final = held(
lambda: _tenant_spans(audit_sinks.tenant, since, rejected[0]),
_carries_the_trace_fields_once,
holding=3,
seconds=40,
)
assert _trace_fields(tenant) == [("user-42", TAGS)], [span["name"] for span in tenant]
def test_a_failed_call_reaches_a_success_and_failure_destination_once(
proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
) -> None:
status, tenant = _send(proxy, provider, audit_sinks, "both", _one_failed_call)
assert status == 500
assert _one_failed_call(tenant), [span["name"] for span in tenant]

View file

@ -0,0 +1,186 @@
"""A key's or team's own Langfuse project gets one copy of the call with the Langfuse trace fields when the operator
has no Langfuse keys, whether or not ``langfuse_otel`` is listed in ``litellm_settings.callbacks``."""
from __future__ import annotations
import json
import uuid
from collections.abc import Iterator
from typing import Final, Literal, cast
import pytest
from integration._support.client import Gateway, gateway_from_environment, held
from integration._support.otlp_sink import (
AuditConfigWriter,
Span,
SpanSinks,
langfuse_callback_variables,
recorded_spans,
spans_for_trace,
)
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
OPERATION: Final = "gen_ai.operation.name"
TAGS: Final = ["agent:orders", "run:r1"]
LEGACY_SPANS: Final = frozenset({"litellm_request", "raw_gen_ai_request"})
def _chat(_request: Request) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
}
).encode()
)
@pytest.fixture(scope="module")
def provider() -> Iterator[Wire]:
with wire_server(_chat) as wire:
yield wire
@pytest.fixture(scope="module", params=[["langfuse_otel"], []], ids=["listed", "not-listed"])
def proxy(
request: pytest.FixtureRequest,
audit_sinks: SpanSinks,
otel_audit_config: AuditConfigWriter,
tmp_path_factory: pytest.TempPathFactory,
) -> Iterator[Gateway]:
directory: Final = tmp_path_factory.mktemp("langfuse-tenant-only")
with (
gateway_from_environment() as gateway,
owned_proxy_process(
gateway,
directory,
{
"LITELLM_OTEL_V2": "1",
"OTEL_BSP_SCHEDULE_DELAY": "300",
"OTEL_EXPORTER": "otlp_http",
"OTEL_ENDPOINT": audit_sinks.operator,
},
config=otel_audit_config(directory, {"callbacks": cast("list[JsonValue]", request.param)}),
remove_environment=("LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_HOST"),
workers=2,
) as owned,
):
yield owned.gateway
def _trace_spans(url: str, since: int, trace_id: str) -> tuple[Span, ...]:
return spans_for_trace(recorded_spans(url, since)[1], trace_id)
def _model_spans(spans: tuple[Span, ...]) -> tuple[Span, ...]:
return tuple(span for span in spans if OPERATION in span["attributes"])
def _one_model_span(spans: tuple[Span, ...]) -> bool:
return len(_model_spans(spans)) == 1 and not LEGACY_SPANS & {str(span["name"]) for span in spans}
@pytest.mark.parametrize("capture", [None, "span_only"])
@pytest.mark.parametrize("level", ["key", "team"])
def test_a_tenant_langfuse_project_gets_user_session_and_tags_without_operator_keys(
level: Literal["key", "team"], capture: str | None, proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
) -> None:
logging: Final[list[JsonValue]] = [
{
"callback_name": "langfuse_otel",
"callback_type": "success",
"callback_vars": {
**langfuse_callback_variables(audit_sinks),
**({"capture_message_content": capture} if capture else {}),
},
}
]
with proxy.scenario() as scenario:
model: Final = scenario.model(api_base=provider.url + "/v1")
alias: Final = f"lf-tenant-only-{uuid.uuid4().hex}"
key: Final = (
scenario.key(key_alias=alias, metadata={"logging": logging})
if level == "key"
else scenario.key(key_alias=alias, team_id=scenario.team(metadata={"logging": logging}))
)
since: Final = (recorded_spans(audit_sinks.tenant)[0], recorded_spans(audit_sinks.operator)[0])
trace_id: Final = uuid.uuid4().hex
session: Final = f"sess-{trace_id}"
body: Final[dict[str, JsonValue]] = {
"model": model,
"messages": [{"role": "user", "content": f"lf-tenant-only-{trace_id}"}],
"metadata": {"tags": [*TAGS], "session_id": session, "trace_user_id": "user-42"},
}
response: Final = proxy.request(
"POST",
"/v1/chat/completions",
body,
key=key,
headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
)
assert response.status_code == 200, response.text
tenant: Final = _model_spans(
held(lambda: _trace_spans(audit_sinks.tenant, since[0], trace_id), _one_model_span, holding=3, seconds=40)
)
operator: Final = _model_spans(
held(lambda: _trace_spans(audit_sinks.operator, since[1], trace_id), _one_model_span, holding=3, seconds=40)
)
assert [
{name: span["attributes"].get(name) for name in ("user.id", "session.id", "langfuse.trace.tags")}
for span in tenant
] == [{"user.id": "user-42", "session.id": session, "langfuse.trace.tags": TAGS}], [s["name"] for s in tenant]
assert "langfuse.observation.metadata" in tenant[0]["attributes"], sorted(tenant[0]["attributes"])
assert (f"lf-tenant-only-{trace_id}" in json.dumps(tenant[0]["attributes"])) == (capture == "span_only"), sorted(
tenant[0]["attributes"]
)
assert len(operator) == 1, [span["name"] for span in operator]
assert f"lf-tenant-only-{trace_id}" not in json.dumps(operator[0]["attributes"]), sorted(operator[0]["attributes"])
assert [
name for name in operator[0]["attributes"] if name in ("user.id", "session.id") or name.startswith("langfuse.")
] == [], sorted(operator[0]["attributes"])
def test_a_tenant_langfuse_project_gets_the_keys_langfuse_environment(
proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
) -> None:
logging: Final[list[JsonValue]] = [
{
"callback_name": "langfuse_otel",
"callback_type": "success",
"callback_vars": {**langfuse_callback_variables(audit_sinks), "langfuse_environment": "prod"},
}
]
with proxy.scenario() as scenario:
model: Final = scenario.model(api_base=provider.url + "/v1")
key: Final = scenario.key(key_alias=f"lf-tenant-env-{uuid.uuid4().hex}", metadata={"logging": logging})
since: Final = (recorded_spans(audit_sinks.tenant)[0], recorded_spans(audit_sinks.operator)[0])
trace_id: Final = uuid.uuid4().hex
response: Final = proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"lf-tenant-env-{trace_id}"}]},
key=key,
headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
)
assert response.status_code == 200, response.text
tenant: Final = held(
lambda: _trace_spans(audit_sinks.tenant, since[0], trace_id), _one_model_span, holding=3, seconds=40
)
operator: Final = held(
lambda: _trace_spans(audit_sinks.operator, since[1], trace_id), _one_model_span, holding=3, seconds=40
)
assert [span["resource"].get("langfuse.environment") for span in tenant] == ["prod"] * len(tenant), [
span["name"] for span in tenant
]
assert not any("langfuse.environment" in span["resource"] for span in operator), [span["name"] for span in operator]

View file

@ -70,7 +70,9 @@ def bedrock_converse_peer(request: Request) -> Reply:
)
def send_to_bedrock_converse(gateway: Gateway, history: list[JsonValue]) -> tuple[httpx.Response, tuple[Request, ...]]:
def send_to_bedrock_converse(
gateway: Gateway, history: list[JsonValue], drop_params: bool = True
) -> tuple[httpx.Response, tuple[Request, ...]]:
with wire_server(bedrock_converse_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/converse/{BEDROCK_GPT}",
@ -79,7 +81,15 @@ def send_to_bedrock_converse(gateway: Gateway, history: list[JsonValue]) -> tupl
aws_bedrock_runtime_endpoint=wire.url,
)
response: Final = gateway.request(
"POST", "/v1/responses", {"model": model, "input": history, "store": False, "cache": {"no-cache": True}}
"POST",
"/v1/responses",
{
"model": model,
"input": history,
"store": False,
"cache": {"no-cache": True},
"drop_params": drop_params,
},
)
return response, wire.drain()
@ -113,6 +123,20 @@ def test_codex_tool_list_repeated_on_every_turn_reaches_bedrock_converse_once(ga
assert sorted(names) == ["collaboration__followup_task", "exec_command"], names
def test_codex_tool_list_repeated_on_every_turn_reaches_bedrock_converse_as_sent_without_drop_params(
gateway: Gateway,
) -> None:
marker: Final = uuid.uuid4().hex
tools: Final = tool_list_item(function_tool("exec_command"))
history: Final[list[JsonValue]] = [tools, user_turn(f"first {marker}"), assistant_turn("ok"), tools]
response, forwarded = send_to_bedrock_converse(gateway, history, drop_params=False)
assert response.status_code == 400, response.text
assert "already defined" in response.text, response.text
assert [str(spec["name"]) for spec in converse_tool_specs(forwarded[0])] == ["exec_command", "exec_command"]
def test_newest_copy_of_a_repeated_codex_tool_is_what_bedrock_converse_receives(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
history: Final[list[JsonValue]] = [
@ -195,6 +219,7 @@ def test_codex_tool_list_repeated_on_every_turn_reaches_bedrock_mantle_once(gate
user_turn(f"second {marker}"),
],
"store": False,
"drop_params": True,
},
)
forwarded: Final = wire.drain()

View file

@ -0,0 +1,351 @@
import json
import threading
import time
import uuid
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from itertools import product
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, list_value, object_value
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_MODEL: Final = "background-scan-chat"
_CLAUDE: Final = "background-scan-claude"
_RATE: Final = 5
_TRIGGER: Final = "FORBIDDEN"
_DENIAL: Final = "Output withheld by policy."
_CLEAN: Final = tuple(f"w{index} " for index in range(40))
_LEAKY: Final = ("w0 ", f"{_TRIGGER} ", *_CLEAN[2:])
_OUTAGE: Final = ("w0 ", "OUTAGE ", *_CLEAN[2:])
_HOLD_SECONDS: Final = 8.0
_ALLOWED: Final = json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()
_INTERVENED: Final = json.dumps(
{
"action": "GUARDRAIL_INTERVENED",
"outputs": [{"text": _DENIAL}],
"assessments": [{"topicPolicy": {"topics": [{"name": "forbidden", "type": "DENY", "action": "BLOCKED"}]}}],
}
).encode()
@dataclass(slots=True)
class _Scans:
"""The scripted ApplyGuardrail peer holds each output scan until `release` is set or the hold runs out."""
release: threading.Event = field(default_factory=threading.Event)
started_at: list[float] = field(default_factory=list)
replied_at: list[float] = field(default_factory=list)
def answer(self, request: Request) -> Reply:
self.started_at.append(time.monotonic())
body: Final = request.body.decode()
if '"OUTPUT"' in body:
self.release.wait(_HOLD_SECONDS)
self.replied_at.append(time.monotonic())
if "OUTAGE" in body:
return Reply(status=500, body=b'{"message": "synthetic ApplyGuardrail outage"}')
return Reply(body=_INTERVENED if _TRIGGER in body else _ALLOWED)
@dataclass(frozen=True, slots=True)
class _Rig:
gateway: Gateway
scans: list[_Scans]
@dataclass(frozen=True, slots=True)
class _Streamed:
pieces: tuple[str, ...]
arrived_at: tuple[float, ...]
finished_after: tuple[int, ...]
errors: tuple[str, ...]
def _chat_sse(pieces: tuple[str, ...]) -> tuple[bytes, ...]:
def frame(delta: dict[str, JsonValue], finish: str | None) -> bytes:
chunk: Final = {
"id": "chatcmpl-background",
"object": "chat.completion.chunk",
"created": 1,
"model": _MODEL,
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
}
return b"data: " + json.dumps(chunk).encode() + b"\n\n"
return (
*(frame({"role": "assistant", "content": piece}, None) for piece in pieces),
frame({}, "stop"),
b"data: [DONE]\n\n",
)
def _anthropic_sse(pieces: tuple[str, ...]) -> tuple[bytes, ...]:
def event(payload: dict[str, JsonValue]) -> bytes:
return f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n".encode()
message: Final[dict[str, JsonValue]] = {
"id": "msg_background",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 5, "output_tokens": 1},
}
return (
event({"type": "message_start", "message": message}),
event({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
*(
event({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": piece}})
for piece in pieces
),
event({"type": "content_block_stop", "index": 0}),
event({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": len(pieces)}}),
event({"type": "message_stop"}),
)
def _model_reply(request: Request) -> Reply:
prompt: Final = str(object_value(list_value(object_value(json.loads(request.body))["messages"])[-1])["content"])
pieces: Final = _LEAKY if "leak" in prompt else _OUTAGE if "outage" in prompt else _CLEAN
return Reply(
chunks=_anthropic_sse(pieces) if request.target.endswith("/v1/messages") else _chat_sse(pieces),
content_type="text/event-stream",
pause_between_chunks=0.05,
)
def _bedrock_rail(name: str, endpoint: str, background: bool) -> dict[str, JsonValue]:
return {
"guardrail_name": name,
"litellm_params": {
"guardrail": "bedrock",
"mode": "post_call",
"default_on": False,
"disable_exception_on_block": True,
"streaming_buffer_until_moderated": False,
"streaming_end_of_stream_only": False,
"streaming_sampling_rate": _RATE,
**({"streaming_scan_in_background": True} if background else {}),
"guardrailIdentifier": "synthetic" + uuid.uuid4().hex[:8],
"guardrailVersion": "DRAFT",
"aws_region_name": "us-east-1",
"aws_access_key_id": "AKIASYNTHETICGUARDRAIL",
"aws_secret_access_key": "synthetic-secret",
"aws_bedrock_runtime_endpoint": endpoint,
},
}
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
directory: Final = tmp_path_factory.mktemp("background-stream-scan")
scans: Final = [_Scans()]
stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = directory / "background-stream-scan.yaml"
with wire_server(lambda request: scans[0].answer(request)) as policy, wire_server(_model_reply) as model:
config.write_text(
yaml.safe_dump(
{
**stock,
"model_list": [
*list_value(stock["model_list"]),
{
"model_name": _MODEL,
"litellm_params": {
"model": f"openai/{_MODEL}",
"api_base": f"{model.url}/v1",
"api_key": "synthetic",
},
},
{
"model_name": _CLAUDE,
"litellm_params": {
"model": "anthropic/claude-sonnet-4-6",
"api_base": model.url,
"api_key": "synthetic",
},
},
],
"guardrails": [
_bedrock_rail("background", policy.url, background=True),
_bedrock_rail("inline", policy.url, background=False),
],
}
)
)
with gateway_from_environment() as shared, owned_proxy(shared, directory, {}, config=config) as candidate:
yield _Rig(candidate, scans)
def _chat_text(event: dict[str, JsonValue]) -> str | None:
texts: Final = [
object_value(object_value(choice).get("delta") or {}).get("content")
for choice in list_value(event.get("choices") or [])
]
return "".join(text for text in texts if isinstance(text, str)) or None
def _messages_text(event: dict[str, JsonValue]) -> str | None:
delta: Final = object_value(event.get("delta") or {}) if event.get("type") == "content_block_delta" else {}
text: Final = delta.get("text")
return text if isinstance(text, str) and text else None
def _responses_text(event: dict[str, JsonValue]) -> str | None:
delta: Final = event.get("delta") if event.get("type") == "response.output_text.delta" else None
return delta if isinstance(delta, str) and delta else None
def _chat_finished(event: dict[str, JsonValue]) -> bool:
return any(object_value(choice).get("finish_reason") for choice in list_value(event.get("choices") or []))
def _messages_finished(event: dict[str, JsonValue]) -> bool:
return event.get("type") == "message_delta"
def _responses_finished(event: dict[str, JsonValue]) -> bool:
return event.get("type") in ("response.completed", "response.incomplete")
def _error_message(event: dict[str, JsonValue]) -> str | None:
error: Final = event.get("error")
if isinstance(error, dict):
return str(error.get("message"))
return str(event.get("message")) if event.get("type") == "error" else None
def _chat_body(prompt: str) -> dict[str, JsonValue]:
return {"model": _MODEL, "messages": [{"role": "user", "content": prompt}]}
def _messages_body(prompt: str) -> dict[str, JsonValue]:
return {**_chat_body(prompt), "model": _CLAUDE, "max_tokens": 256}
def _responses_body(prompt: str) -> dict[str, JsonValue]:
return {"model": _CLAUDE, "max_output_tokens": 256, "input": prompt}
_ROUTES: Final = {
"/v1/chat/completions": (_chat_body, _chat_text, _chat_finished),
"/v1/messages": (_messages_body, _messages_text, _messages_finished),
"/v1/responses": (_responses_body, _responses_text, _responses_finished),
}
def _stream(
rig: _Rig, guardrail: str, prompt: str, release_after: int, route: str = "/v1/chat/completions"
) -> tuple[_Streamed, _Scans]:
"""Stream `route`, releasing held scans once `release_after` text pieces have arrived."""
scans: Final = _Scans()
rig.scans[0] = scans
return _read(rig, guardrail, prompt, route, scans, release_after), scans
def _read(rig: _Rig, guardrail: str, prompt: str, route: str, scans: _Scans, release_after: int) -> _Streamed:
body_of, text_of, finished = _ROUTES[route]
pieces: Final[list[str]] = []
arrived_at: Final[list[float]] = []
finished_after: Final[list[int]] = []
errors: Final[list[str]] = []
with rig.gateway.client.stream(
"POST",
route,
json={**body_of(f"{prompt} {uuid.uuid4().hex}"), "stream": True, "guardrails": [guardrail]},
headers={"Authorization": f"Bearer {rig.gateway.key}"},
) as response:
assert response.status_code == 200, response.read()
for line in response.iter_lines():
if not line.startswith("data: ") or line == "data: [DONE]":
continue
event = object_value(json.loads(line[len("data: ") :]))
text = text_of(event)
if text is not None:
pieces.append(text)
arrived_at.append(time.monotonic())
if finished(event):
finished_after.append(len(pieces))
error = _error_message(event)
if error is not None:
errors.append(error)
if len(pieces) >= release_after:
scans.release.set()
return _Streamed(tuple(pieces), tuple(arrived_at), tuple(finished_after), tuple(errors))
def test_chunks_after_the_nth_keep_streaming_while_its_scan_is_pending(rig: _Rig) -> None:
streamed, scans = _stream(rig, "background", "clean", release_after=_RATE + 1)
assert streamed.pieces == _CLEAN
assert streamed.arrived_at[_RATE] < scans.replied_at[0], "the chunk after the Nth waited for the Nth chunk's scan"
@pytest.mark.parametrize("route", tuple(_ROUTES))
def test_block_verdict_stops_the_stream_at_the_next_chunk(rig: _Rig, route: str) -> None:
streamed, scans = _stream(rig, "background", "leak", release_after=_RATE + 1, route=route)
*before, last = streamed.pieces
assert last == _DENIAL
assert _DENIAL not in before
assert _RATE < len(before) < len(_LEAKY), before
assert tuple(before) == _LEAKY[: len(before)]
assert streamed.arrived_at[_RATE] < scans.replied_at[0]
@pytest.mark.parametrize("route", tuple(_ROUTES))
def test_block_verdict_that_lands_after_the_upstream_finished_replaces_the_finish(rig: _Rig, route: str) -> None:
streamed, scans = _stream(rig, "background", "leak", release_after=len(_LEAKY) + 1, route=route)
assert streamed.pieces == (*_LEAKY, _DENIAL)
assert streamed.arrived_at[len(_LEAKY) - 1] < scans.replied_at[0], "the verdict landed before the upstream finished"
assert streamed.finished_after == (len(streamed.pieces),), "the stream finished before the block verdict"
overlaps: Final = zip(scans.started_at[1:], scans.replied_at[:-1], strict=True)
assert all(started >= replied for started, replied in overlaps), "a scan started while another was in flight"
def test_without_the_flag_the_nth_chunk_still_waits_for_its_scan(rig: _Rig) -> None:
streamed, scans = _stream(rig, "inline", "clean", release_after=_RATE - 1)
assert streamed.pieces == _CLEAN
assert scans.replied_at[0] < streamed.arrived_at[_RATE - 1]
@pytest.mark.parametrize("route", tuple(_ROUTES))
def test_a_failed_background_scan_ends_the_stream_with_the_inline_scan_error(rig: _Rig, route: str) -> None:
background, _ = _stream(rig, "background", "outage", release_after=_RATE + 1, route=route)
inline, _ = _stream(rig, "inline", "outage", release_after=0, route=route)
assert len(inline.errors) == 1, inline
assert background.errors == inline.errors
assert _RATE < len(background.pieces) < len(_OUTAGE), background.pieces
assert background.pieces == _OUTAGE[: len(background.pieces)]
assert background.finished_after == ()
def test_a_mixed_burst_ends_every_stream_on_its_own_verdict_while_the_scan_peer_fails_some(rig: _Rig) -> None:
scans: Final = _Scans()
scans.release.set()
rig.scans[0] = scans
calls: Final = tuple(product(("clean", "leak", "outage"), (*_ROUTES, *_ROUTES)))
with ThreadPoolExecutor(len(calls)) as pool:
served: Final = tuple(pool.map(lambda call: _read(rig, "background", call[0], call[1], scans, 0), calls))
for (kind, route), streamed in zip(calls, served, strict=True):
match kind:
case "clean":
assert (streamed.pieces, streamed.errors) == (_CLEAN, ()), (route, streamed)
assert streamed.finished_after == (len(_CLEAN),), (route, streamed)
case "leak":
assert streamed.pieces[-1] == _DENIAL and streamed.errors == (), (route, streamed)
assert streamed.finished_after == (len(streamed.pieces),), (route, streamed)
case _:
assert len(streamed.errors) == 1 and streamed.finished_after == (), (route, streamed)
assert streamed.pieces == _OUTAGE[: len(streamed.pieces)], (route, streamed)
after: Final = _read(rig, "background", "clean", "/v1/chat/completions", scans, 0)
assert after.pieces == _CLEAN

View file

@ -0,0 +1,20 @@
from collections.abc import Callable
from typing import Final, TypeVar
import pytest
LoopResult = TypeVar("LoopResult")
def runs_until_a_loop_address_repeats(
run_on_new_loop: Callable[[], tuple[int, LoopResult]],
earlier: tuple[tuple[int, LoopResult], ...] = (),
limit: int = 50,
) -> tuple[tuple[int, LoopResult], ...]:
if len(earlier) == limit:
pytest.skip(f"none of {limit} collected loops had its address handed to a later loop")
latest: Final = run_on_new_loop()
runs: Final = (*earlier, latest)
if latest[0] in {loop_id for loop_id, _ in earlier}:
return runs
return runs_until_a_loop_address_repeats(run_on_new_loop, runs, limit)

View file

@ -10,6 +10,7 @@ never closed, because litellm does not own its lifecycle.
import asyncio
import gc
import weakref
from typing import Final
from unittest.mock import AsyncMock
import httpx
@ -18,6 +19,7 @@ from redis.asyncio import ConnectionPool, Redis
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
class FakeClock:
@ -41,6 +43,14 @@ class AsyncClient:
self.closed = True
class RecordingAsyncClient:
def __init__(self) -> None:
self.close_calls = 0
async def aclose(self) -> None:
self.close_calls += 1
class SyncClient:
def __init__(self) -> None:
self.closed = False
@ -266,6 +276,100 @@ def test_sync_client_evicted_outside_an_event_loop_is_still_closed():
assert closer.pending_count == 0
def test_async_client_from_collected_loop_is_not_closed_by_fresh_loops() -> None:
clock: Final = FakeClock()
closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock)
async def reap_then_schedule(client: RecordingAsyncClient) -> None:
closer.reap()
await asyncio.sleep(0)
closer.schedule(client)
def run_on_new_loop() -> tuple[int, RecordingAsyncClient]:
client: Final = RecordingAsyncClient()
closer.mark_owned(client)
loop: Final = asyncio.new_event_loop()
loop_id: Final = id(loop)
try:
loop.run_until_complete(reap_then_schedule(client))
finally:
loop.close()
del loop
gc.collect()
return loop_id, client
runs: Final = runs_until_a_loop_address_repeats(run_on_new_loop)
assert all(client.close_calls == 0 for _, client in runs)
def test_async_client_queued_on_a_live_loop_is_closed_only_by_that_loop() -> None:
clock: Final = FakeClock()
closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock)
client: Final = RecordingAsyncClient()
closer.mark_owned(client)
owner_loop: Final = asyncio.new_event_loop()
other_loop: Final = asyncio.new_event_loop()
async def schedule() -> None:
closer.schedule(client)
async def reap() -> None:
closer.reap()
await asyncio.sleep(0)
try:
owner_loop.run_until_complete(schedule())
other_loop.run_until_complete(reap())
closes_after_other_loop_reaped: Final = client.close_calls
pending_after_other_loop_reaped: Final = closer.pending_count
owner_loop.run_until_complete(reap())
finally:
other_loop.close()
owner_loop.close()
assert closes_after_other_loop_reaped == 0
assert pending_after_other_loop_reaped == 1
assert client.close_calls == 1
assert closer.pending_count == 0
def test_closes_queued_on_a_collected_loop_free_capacity() -> None:
clock: Final = FakeClock()
closer: Final = EvictedClientCloser(grace_seconds=0, max_pending=1, clock=clock)
first: Final = RecordingAsyncClient()
closer.mark_owned(first)
loop_a: Final = asyncio.new_event_loop()
async def schedule_first() -> None:
closer.schedule(first)
try:
loop_a.run_until_complete(schedule_first())
finally:
loop_a.close()
del loop_a
gc.collect()
second: Final = RecordingAsyncClient()
closer.mark_owned(second)
loop_b: Final = asyncio.new_event_loop()
async def schedule_and_reap_second() -> None:
closer.schedule(second)
closer.reap()
await asyncio.sleep(0)
try:
loop_b.run_until_complete(schedule_and_reap_second())
finally:
loop_b.close()
assert first.close_calls == 0
assert second.close_calls == 1
assert closer.pending_count == 0
@pytest.mark.asyncio
async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
clock = FakeClock()
@ -289,8 +393,7 @@ async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
@pytest.mark.asyncio
async def test_a_client_evicted_on_another_event_loop_is_left_alone():
"""Closing a client bound to a different loop would schedule work on that loop."""
async def test_a_client_evicted_on_a_collected_loop_is_dropped():
clock = FakeClock()
closer = make_closer(clock)
client = AsyncClient()
@ -310,7 +413,7 @@ async def test_a_client_evicted_on_another_event_loop_is_left_alone():
await asyncio.sleep(0.05)
assert client.closed is False
assert closer.pending_count == 1
assert closer.pending_count == 0
@pytest.mark.asyncio

View file

@ -10,6 +10,7 @@ See: https://github.com/BerriAI/litellm/pull/22247
import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref
import warnings
from typing import Final
import aiohttp
import pytest
@ -29,6 +30,41 @@ from litellm.llms.custom_httpx.http_handler import(
from litellm.types.utils import LlmProviders
from litellm.utils import _invalidate_model_cost_lowercase_map
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
async def _get_and_set_cache(cache: LLMClientCache) -> object | None:
cached: Final = await cache.async_get_cache("client")
await cache.async_set_cache("client", object())
return cached
def _cache_result_on_new_loop(cache: LLMClientCache) -> tuple[int, object | None]:
loop: Final = asyncio.new_event_loop()
loop_id: Final = id(loop)
try:
result: Final = loop.run_until_complete(_get_and_set_cache(cache))
finally:
loop.close()
del loop
gc.collect()
return loop_id, result
async def _get_cache_key(cache: LLMClientCache) -> str:
return cache.update_cache_key_with_event_loop("client")
def _cache_key_on_new_loop(cache: LLMClientCache) -> tuple[int, str]:
loop: Final = asyncio.new_event_loop()
loop_id: Final = id(loop)
try:
key: Final = loop.run_until_complete(_get_cache_key(cache))
finally:
loop.close()
del loop
gc.collect()
return loop_id, key
class MockAsyncClient:
@ -51,6 +87,44 @@ class MockSyncClient:
self.closed = True
def test_async_client_cache_does_not_reuse_clients_between_collected_loops() -> None:
cache: Final = LLMClientCache()
runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_result_on_new_loop(cache))
assert all(result is None for _, result in runs)
def test_async_client_cache_keys_differ_for_every_collected_loop() -> None:
cache: Final = LLMClientCache()
runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_key_on_new_loop(cache))
keys: Final = tuple(key for _, key in runs)
assert len(set(keys)) == len(keys)
class _LoopWithoutWeakReferences:
__slots__ = ()
def _cache_key_while_running(cache: LLMClientCache, loop: _LoopWithoutWeakReferences) -> str:
asyncio._set_running_loop(loop)
try:
return cache.update_cache_key_with_event_loop("client")
finally:
asyncio._set_running_loop(None)
def test_cache_keys_stay_per_loop_on_loops_without_weak_references() -> None:
cache: Final = LLMClientCache()
loop_a: Final = _LoopWithoutWeakReferences()
loop_b: Final = _LoopWithoutWeakReferences()
first_key_on_a: Final = _cache_key_while_running(cache, loop_a)
assert _cache_key_while_running(cache, loop_a) == first_key_on_a
assert _cache_key_while_running(cache, loop_b) != first_key_on_a
@pytest.mark.asyncio
async def test_remove_key_does_not_close_async_client():
"""

View file

@ -17,7 +17,7 @@ import pytest
from opentelemetry import trace as trace_api
from opentelemetry.context import Context
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
@ -27,6 +27,7 @@ import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.otel import GenAIOperation
from litellm.integrations.otel import logger as otel_logger
from litellm.integrations.otel.emitter import SpanEmitter
from litellm.integrations.otel.logger import (
OpenTelemetryV2,
build_otel_v2_logger,
@ -34,6 +35,7 @@ from litellm.integrations.otel.logger import (
publish_global_otel_v2_provider,
)
from litellm.integrations.otel.mappers import resolve_mappers
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES, LangfuseMapper
from litellm.integrations.otel.model.config import (
CaptureMessageContent,
ExporterOwner,
@ -52,6 +54,8 @@ from litellm.integrations.otel.model.payloads import (
ServerInfo,
ToolDefinition,
)
from litellm.integrations.otel.model.spans import SpanRole
from litellm.integrations.otel.model.trace_controls import TraceControls
from litellm.integrations.otel.plumbing import providers as otel_providers
from litellm.integrations.otel.plumbing.context import (
destination_backends,
@ -3348,8 +3352,7 @@ class TestArizeTenantOnly:
assert capture.received() == {
"space_id=space-operator,api_key=key-operator": ("chat gpt-4",),
None: ("chat gpt-4",),
}, "the operator's space, plus the stdout placeholder every credentialed preset keeps today"
}, "the operator's space, with no stdout copy the operator never asked for"
def test_the_standard_otlp_headers_still_reach_the_operators_space(self, monkeypatch):
credential_less_arize(monkeypatch)
@ -3362,8 +3365,7 @@ class TestArizeTenantOnly:
assert capture.received() == {
"space_id=space-operator,api_key=key-operator": ("chat gpt-4",),
None: ("chat gpt-4",),
}, "the operator's space, plus the stdout placeholder every credentialed preset keeps today"
}, "the operator's space, with no stdout copy the operator never asked for"
def test_a_credential_less_arize_still_delivers_a_team_destination(self, monkeypatch):
from litellm.integrations.otel.plumbing.providers import attach_tenant_fan_out
@ -3435,6 +3437,271 @@ class TestArizeTenantOnly:
assert loggers == [at_startup]
class TestLangfuseDestinationFields:
"""A key's or team's Langfuse destination gets the Langfuse fields from whichever logger the proxy published."""
def test_a_key_destination_gets_the_langfuse_fields_from_the_generic_startup_logger(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""With ``langfuse_otel`` only in a key's or team's logging, the startup logger is the
generic one; the key's project still gets the user, session and tags, and the
operator's stdout placeholder still gets no Langfuse request metadata."""
credential_less_proxy(monkeypatch)
capture: Final = _ExporterCapture()
for kind in ("otlp_grpc", "otlp_http", "console"):
monkeypatch.setitem(otel_providers._EXPORTER_FACTORIES, kind, capture.build) # pyright: ignore[reportPrivateUsage] # the factory table is the seam for capturing each exporter
config: Final = OpenTelemetryV2Config()
provider: Final = build_tracer_provider(config, tenant_overrides=True, console_default=True)
attach_tenant_fan_out(provider, config)
emitter: Final = SpanEmitter(get_tracer(provider, "litellm"), config)
call: Final = replace(
_MODEL_CALL_WITHOUT_CONTENT,
trace=TraceControls(user_id="user-42", session_id="sess-1", tags=("agent:orders",)),
request_metadata=MappingProxyType({"user_api_key_alias": "prod-key"}),
)
def with_destination() -> None:
set_request_destinations(deliverable_destinations((LANGFUSE_DEST,), provider))
emitter.emit(SpanRole.LLM_CALL, call)
in_fresh_context(with_destination)
provider.force_flush()
received: Final = {
headers: exporter.get_finished_spans()[0].attributes or {} for headers, exporter in capture.built
}
to_langfuse: Final = received[LANGFUSE_DEST.header_string()]
assert {key: to_langfuse.get(key) for key in ("user.id", "session.id", "langfuse.trace.tags")} == {
"user.id": "user-42",
"session.id": "sess-1",
"langfuse.trace.tags": ("agent:orders",),
}
assert "langfuse.observation.metadata" in to_langfuse
assert "langfuse.observation.metadata" not in received[None]
@staticmethod
def _received_by_account(
monkeypatch: pytest.MonkeyPatch,
config: OpenTelemetryV2Config,
destinations: tuple[OtelDestination, ...],
with_content: bool = False,
) -> Mapping[str | None, Mapping[str, object]]:
"""The attributes of the model call each account received, keyed by the headers it was sent with."""
credential_less_proxy(monkeypatch)
capture: Final = _ExporterCapture()
for kind in ("otlp_grpc", "otlp_http", "console"):
monkeypatch.setitem(otel_providers._EXPORTER_FACTORIES, kind, capture.build) # pyright: ignore[reportPrivateUsage] # the factory table is the seam for capturing each exporter
provider: Final = build_tracer_provider(config, tenant_overrides=True, console_default=True)
attach_tenant_fan_out(provider, config)
emitter: Final = SpanEmitter(get_tracer(provider, "litellm"), config)
call: Final = replace(
_MODEL_CALL_WITH_CONTENT if with_content else _MODEL_CALL_WITHOUT_CONTENT,
trace=TraceControls(user_id="user-42", session_id="sess-1", tags=("agent:orders",)),
request_metadata=MappingProxyType({"user_api_key_alias": "prod-key"}),
)
def with_destinations() -> None:
set_request_destinations(deliverable_destinations(destinations, provider))
emitter.emit(SpanRole.LLM_CALL, call)
in_fresh_context(with_destinations)
provider.force_flush()
return MappingProxyType(
{headers: dict(exporter.get_finished_spans()[0].attributes or {}) for headers, exporter in capture.built}
)
def test_only_the_langfuse_destination_gets_the_langfuse_fields_written_for_it(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The operator's own exporter and the key's Arize destination get the model call as they would with no
Langfuse destination on the request."""
received: Final = self._received_by_account(monkeypatch, OpenTelemetryV2Config(), (LANGFUSE_DEST, ARIZE_DEST))
assert received[LANGFUSE_DEST.header_string()]["user.id"] == "user-42"
assert {
headers: sorted(LangfuseMapper.LLM_CALL_ATTRIBUTES & attributes.keys())
for headers, attributes in received.items()
if headers != LANGFUSE_DEST.header_string()
} == {None: [], ARIZE_DEST.header_string(): []}
def test_an_exporter_whose_config_maps_langfuse_keeps_the_langfuse_fields_but_the_request_metadata(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
received: Final = self._received_by_account(
monkeypatch, OpenTelemetryV2Config(mapper_names=["genai", "langfuse"]), (LANGFUSE_DEST,)
)
assert {key: received[None].get(key) for key in ("user.id", "session.id", "langfuse.trace.tags")} == {
"user.id": "user-42",
"session.id": "sess-1",
"langfuse.trace.tags": ("agent:orders",),
}
assert not LANGFUSE_REQUEST_METADATA_ATTRIBUTES & received[None].keys()
@pytest.mark.parametrize(
("key_environment", "tracing_environment", "expected"),
[
("prod", None, {"langfuse.environment": "prod"}),
("prod", "staging", {"langfuse.environment": "prod"}),
(None, "staging", {"langfuse.environment": "staging"}),
(None, None, {}),
],
ids=["key", "key-over-env", "env", "neither"],
)
def test_a_langfuse_destination_names_the_environment_the_legacy_logger_wrote(
self,
monkeypatch: pytest.MonkeyPatch,
key_environment: str | None,
tracing_environment: str | None,
expected: Mapping[str, str],
) -> None:
credential_less_proxy(monkeypatch)
if tracing_environment is None:
monkeypatch.delenv("LANGFUSE_TRACING_ENVIRONMENT", raising=False)
else:
monkeypatch.setenv("LANGFUSE_TRACING_ENVIRONMENT", tracing_environment)
destination: Final = destination_for(
"langfuse_otel",
StandardCallbackDynamicParams(
langfuse_public_key="pk", langfuse_secret_key="sk", langfuse_environment=key_environment
),
)
assert destination is not None
assert dict(destination.resource_attributes) == expected
def test_a_langfuse_destination_on_a_proxy_with_langfuse_keys_names_no_environment(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""With operator keys the older Langfuse logger never runs, so there is no environment of its to keep."""
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-operator")
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-operator")
destination: Final = destination_for(
"langfuse_otel",
StandardCallbackDynamicParams(
langfuse_public_key="pk", langfuse_secret_key="sk", langfuse_environment="prod"
),
)
assert destination is not None
assert dict(destination.resource_attributes) == {}
@pytest.mark.parametrize("operator_keys", [False, True], ids=["no-operator-keys", "operator-keys"])
@pytest.mark.parametrize("configured", [None, "no_content", "span_only"])
def test_a_langfuse_destination_keeps_its_own_capture_setting_whatever_the_operator_keys(
self,
monkeypatch: pytest.MonkeyPatch,
configured: str | None,
operator_keys: bool,
) -> None:
"""An unset setting stays unset, so the destination follows the global capture mode either way."""
credential_less_proxy(monkeypatch)
if operator_keys:
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-operator")
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-operator")
destination: Final = destination_for(
"langfuse_otel",
StandardCallbackDynamicParams(langfuse_public_key="pk", langfuse_secret_key="sk"),
capture_message_content=configured,
)
assert destination is not None
assert destination.capture_message_content == configured
@pytest.mark.parametrize(("configured", "captured"), [(None, False), ("span_only", True)])
def test_a_langfuse_destination_gets_the_prompt_only_when_its_setting_asks_without_operator_keys(
self, monkeypatch: pytest.MonkeyPatch, configured: str | None, captured: bool
) -> None:
credential_less_proxy(monkeypatch)
destination: Final = destination_for(
"langfuse_otel",
StandardCallbackDynamicParams(langfuse_public_key="pk", langfuse_secret_key="sk"),
capture_message_content=configured,
)
arize: Final = destination_for("arize", StandardCallbackDynamicParams(arize_space_id="s", arize_api_key="k"))
assert destination is not None
assert arize is not None
received: Final = self._received_by_account(
monkeypatch, OpenTelemetryV2Config(), (destination, arize), with_content=True
)
assert {headers: "langfuse.observation.input" in attributes for headers, attributes in received.items()} == {
destination.header_string(): captured,
arize.header_string(): False,
None: False,
}
assert not [key for key in received[None] if key.startswith("gen_ai.input")]
def test_only_the_langfuse_destination_gets_the_keys_environment(self, monkeypatch: pytest.MonkeyPatch) -> None:
credential_less_proxy(monkeypatch)
capture: Final = _ExporterCapture()
for kind in ("otlp_grpc", "otlp_http", "console"):
monkeypatch.setitem(otel_providers._EXPORTER_FACTORIES, kind, capture.build) # pyright: ignore[reportPrivateUsage] # the factory table is the seam for capturing each exporter
config: Final = OpenTelemetryV2Config()
provider: Final = build_tracer_provider(config, tenant_overrides=True, console_default=True)
attach_tenant_fan_out(provider, config)
emitter: Final = SpanEmitter(get_tracer(provider, "litellm"), config)
destination: Final = destination_for(
"langfuse_otel", {"langfuse_public_key": "pk", "langfuse_secret_key": "sk", "langfuse_environment": "prod"}
)
assert destination is not None
def with_destination() -> None:
set_request_destinations(deliverable_destinations((destination,), provider))
emitter.emit(SpanRole.LLM_CALL, _MODEL_CALL_WITHOUT_CONTENT)
in_fresh_context(with_destination)
provider.force_flush()
environments: Final = {
headers: exporter.get_finished_spans()[0].resource.attributes.get("langfuse.environment")
for headers, exporter in capture.built
}
assert environments == {destination.header_string(): "prod", None: None}
@staticmethod
def _attributes(
destinations: tuple[OtelDestination, ...], limit: int | None, call: LLMCallSpanData
) -> dict[str, object]:
exporter: Final = InMemorySpanExporter()
provider: Final = TracerProvider(span_limits=SpanLimits(max_span_attributes=limit))
provider.add_span_processor(SimpleSpanProcessor(exporter))
emitter: Final = SpanEmitter(get_tracer(provider, "litellm"), OpenTelemetryV2Config())
def run() -> None:
set_request_destinations(destinations)
emitter.emit(SpanRole.LLM_CALL, call)
in_fresh_context(run)
return dict(exporter.get_finished_spans()[0].attributes or {})
def test_a_span_at_its_attribute_limit_sheds_the_destinations_langfuse_fields_first(self) -> None:
generic: Final = self._attributes((), None, _MODEL_CALL_WITH_CONTENT)
assert "user.id" not in generic
assert self._attributes((LANGFUSE_DEST,), len(generic), _MODEL_CALL_WITH_CONTENT) == generic, (
"the generic attributes all survive"
)
def test_a_span_at_its_attribute_limit_sheds_the_destinations_request_metadata_before_its_other_fields(
self,
) -> None:
call: Final = replace(
_MODEL_CALL_WITH_CONTENT, request_metadata=MappingProxyType({"user_api_key_alias": "prod-key"})
)
whole: Final = self._attributes((LANGFUSE_DEST,), None, call)
squeezed: Final = self._attributes((LANGFUSE_DEST,), len(whole) - 1, call)
assert len(squeezed) == len(whole) - 1
assert whole.keys() - squeezed.keys() <= LANGFUSE_REQUEST_METADATA_ATTRIBUTES
assert squeezed["langfuse.observation.type"] == "generation"
class TestContextIsolation:
def test_destinations_do_not_leak_between_requests(self):
def first():
@ -4556,6 +4823,115 @@ class TestCredentialGatedExporters:
assert kept[0] is operator_console
class TestConsolePlaceholder:
"""The console spec ``_normalize`` folds in for an empty exporter list prints spans
only for the ``otel`` callback; no other logger was asked to print."""
@pytest.fixture
def capture(self, monkeypatch: pytest.MonkeyPatch) -> "_ExporterCapture":
return TestArizeTenantOnly._capture_exporters(monkeypatch)
@staticmethod
def _received(capture: "_ExporterCapture", provider: TracerProvider) -> Mapping[str | None, tuple[str, ...]]:
emit(provider)
provider.force_flush()
return capture.received()
@staticmethod
def _no_otel_env(monkeypatch: pytest.MonkeyPatch) -> None:
for name in _OTEL_SHORTHAND_ENV:
monkeypatch.delenv(name, raising=False)
@staticmethod
def _operator_langfuse(monkeypatch: pytest.MonkeyPatch) -> None:
TestConsolePlaceholder._no_otel_env(monkeypatch)
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-operator")
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-operator")
@staticmethod
def _isolated_proxy_registration(monkeypatch: pytest.MonkeyPatch) -> None:
for name in ("service_callback", "input_callback", "_async_success_callback", "_async_failure_callback"):
monkeypatch.setattr(litellm, name, [])
monkeypatch.setattr("litellm.proxy.proxy_server.open_telemetry_logger", None)
def test_a_credentialed_langfuse_exports_to_langfuse_alone(
self, monkeypatch: pytest.MonkeyPatch, capture: "_ExporterCapture"
) -> None:
self._operator_langfuse(monkeypatch)
config: Final = langfuse_preset()
langfuse_headers: Final = next(
spec.headers for spec in config.exporters if spec.owner is ExporterOwner.LANGFUSE_OTEL
)
assert self._received(capture, build_tracer_provider(config)) == {langfuse_headers: ("chat gpt-4",)}
def test_a_console_exporter_the_operator_named_still_prints(
self, monkeypatch: pytest.MonkeyPatch, capture: "_ExporterCapture"
) -> None:
self._operator_langfuse(monkeypatch)
monkeypatch.setenv("OTEL_EXPORTER", "console")
assert self._received(capture, build_tracer_provider(langfuse_preset()))[None] == ("chat gpt-4",)
@pytest.fixture
def otel_callback(self, monkeypatch: pytest.MonkeyPatch, capture: "_ExporterCapture") -> OpenTelemetryV2:
"""The plain ``otel`` callback with nothing else configured, built the way the proxy builds it."""
from litellm.litellm_core_utils import litellm_logging as logging_module
self._no_otel_env(monkeypatch)
self._isolated_proxy_registration(monkeypatch)
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
monkeypatch.setattr(logging_module, "_in_memory_loggers", [])
is_otel_v2_enabled.cache_clear()
built: Final = logging_module.build_custom_logger_compatible_class(
"otel", internal_usage_cache=None, llm_router=None
)
assert isinstance(built, OpenTelemetryV2)
return built
def test_the_otel_callback_prints_when_nothing_else_is_configured(
self, otel_callback: OpenTelemetryV2, capture: "_ExporterCapture"
) -> None:
assert self._received(capture, otel_callback.tracer_provider) == {None: ("chat gpt-4",)}
def test_the_otel_callback_prints_a_call_for_a_key_with_its_own_service_name(
self, otel_callback: OpenTelemetryV2, capture: "_ExporterCapture"
) -> None:
"""A key or team ``otel_service_name`` moves the call onto a second provider built for that name."""
kwargs: Final = {
**_closed_chat_call_kwargs(),
"litellm_params": {"metadata": {"user_api_key_auth_metadata": {"otel_service_name": "payments-gateway"}}},
}
otel_callback.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
asyncio.run(otel_callback.async_log_success_event(kwargs, None, None, None))
assert capture.received() == {None: ("chat gpt-4o",)}
@pytest.mark.parametrize(
("env", "expected"),
[({}, {}), ({"OTEL_EXPORTER": "console"}, {None: ("chat gpt-4",)})],
ids=["nothing-configured", "operator-asked-for-console"],
)
def test_the_background_global_logger_prints_only_when_asked(
self,
monkeypatch: pytest.MonkeyPatch,
capture: "_ExporterCapture",
env: Mapping[str, str],
expected: Mapping[str | None, tuple[str, ...]],
) -> None:
"""The proxy publishes a generic logger even when no callback built one, which is
what key-only and team-only destinations run under."""
from litellm.integrations.otel.logger import select_global_otel_v2_logger
self._no_otel_env(monkeypatch)
self._isolated_proxy_registration(monkeypatch)
for name, value in env.items():
monkeypatch.setenv(name, value)
assert self._received(capture, select_global_otel_v2_logger([]).tracer_provider) == expected
class TestTenantHostSsrfGuard:
"""Anyone who can mint a key can write ``langfuse_host``, so the host it names has
to be one the operator approved."""

View file

@ -1,10 +1,14 @@
import os
from typing import Final
from unittest.mock import MagicMock, patch
import pytest
import respx
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.integrations.agentops.agentops import AgentOps, AgentOpsConfig
from litellm.llms.custom_httpx.http_handler import get_httpx_client
@pytest.fixture
@ -111,3 +115,17 @@ def test_agentops_initialization_no_auth():
assert agentops.resource_attributes["deployment.environment"] == "test_env"
assert agentops.resource_attributes["telemetry.sdk.name"] == "agentops"
assert "project.id" not in agentops.resource_attributes
def test_auth_fetch_leaves_the_shared_sync_client_open_for_other_callers(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
auth_endpoint: Final = "https://agentops.example/v3/auth/token"
respx_mock.post(auth_endpoint).respond(json={"token": "jwt", "project_id": "project-1"})
client_another_caller_is_using: Final = get_httpx_client().client
agentops: Final = AgentOps(config=AgentOpsConfig(api_key="key", auth_endpoint=auth_endpoint))
assert agentops.resource_attributes["project.id"] == "project-1"
assert not client_another_caller_is_using.is_closed

View file

@ -1,9 +1,12 @@
import asyncio
import base64
import contextvars
import json
import os
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Final
from unittest.mock import MagicMock, patch
@ -12,12 +15,17 @@ from openai.types.responses import ResponseFunctionToolCall
from opentelemetry import trace
from opentelemetry.context import Context
from opentelemetry.sdk.trace import SpanLimits, TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import Span
import litellm
from litellm.constants import REDACTED_BY_LITELLM
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, REDACTED_BY_LITELLM
from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME, OpenTelemetry, OpenTelemetryConfig
from litellm.integrations.otel.model.destination import OtelDestination
from litellm.integrations.otel.plumbing.context import set_request_destinations
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import Choices, ModelResponse, StandardCallbackDynamicParams
@ -1507,3 +1515,75 @@ def test_langfuse_otel_does_not_start_proxy_request_span() -> None:
finally:
langfuse_provider.shutdown()
generic_provider.shutdown()
_HOOK_CALL: Final[Mapping[str, object]] = MappingProxyType({})
def _spans_the_legacy_logger_exports(
destinations: tuple[OtelDestination, ...], call: Mapping[str, object] = _HOOK_CALL
) -> tuple[str, ...]:
exporter: Final = InMemorySpanExporter()
provider: Final = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
with patch.dict(os.environ, LANGFUSE_ENV_ONLY, clear=True):
logger: Final = LangfuseOtelLogger(tracer_provider=provider)
at: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
kwargs: Final[dict[str, object]] = {**CHAT_KWARGS, **call, "litellm_params": {"metadata": {}}}
server_span: Final = provider.get_tracer("proxy").start_span("POST /v1/chat/completions")
async def log() -> None:
set_request_destinations(destinations)
logger.log_success_event(dict(kwargs), _chat_response("ok"), at, at) # pyright: ignore[reportUnknownMemberType] # the v1 logger hooks are untyped
await logger.async_log_failure_event({**kwargs, "exception": ValueError("boom")}, None, at, at) # pyright: ignore[reportUnknownMemberType] # the v1 logger hooks are untyped
await logger.async_post_call_failure_hook({}, ValueError("boom"), UserAPIKeyAuth(parent_otel_span=server_span))
try:
contextvars.copy_context().run(asyncio.run, log())
return tuple(span.name for span in exporter.get_finished_spans())
finally:
provider.shutdown()
_LANGFUSE_DESTINATION: Final = OtelDestination(
endpoint="http://tenant.local/api/public/otel",
headers={"Authorization": "Basic dGVuYW50"},
callback_name="langfuse_otel",
)
_PROVIDER_CALL: Final[Mapping[str, object]] = MappingProxyType(
{"standard_logging_object": {"id": "c1"}, "api_call_start_time": datetime(2026, 1, 1, tzinfo=timezone.utc)}
)
@pytest.mark.parametrize("call", [_HOOK_CALL, _PROVIDER_CALL], ids=["hook", "provider-call"])
def test_a_request_with_a_langfuse_destination_exports_nothing_through_the_legacy_logger(
call: Mapping[str, object],
) -> None:
assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION,), call) == ()
def test_a_call_rejected_before_any_provider_still_exports_through_the_legacy_logger() -> None:
rejected: Final = {"standard_logging_object": {"id": "c1"}, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True}
assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION,), rejected) == (
"litellm_request",
"raw_gen_ai_request",
"litellm_request",
)
def test_a_request_without_a_langfuse_destination_still_exports_through_the_legacy_logger() -> None:
logged: Final = (
"litellm_request",
"raw_gen_ai_request",
"litellm_request",
"Failed Proxy Server Request",
"POST /v1/chat/completions",
)
assert _spans_the_legacy_logger_exports(()) == logged
assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION.model_copy(update={"callback_name": "arize"}),)) == (
logged
)

View file

@ -1388,6 +1388,55 @@ async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref():
assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled)
async def _drop_unused_handler() -> None:
AsyncHTTPHandler(timeout=61.0)
gc.collect()
def test_finalizer_close_does_not_keep_a_loop_closed_right_after_it_alive():
loop: Final = asyncio.new_event_loop()
loop.run_until_complete(_drop_unused_handler())
loop.close()
closed_loop: Final = weakref.ref(loop)
del loop
gc.collect()
assert closed_loop() is None
class _TransportFailingToClose(httpx.AsyncBaseTransport):
async def aclose(self) -> None:
raise RuntimeError("close failed")
async def _drop_handler_whose_close_fails() -> None:
AsyncHTTPHandler(timeout=61.0, transport=_TransportFailingToClose())
gc.collect()
def test_failed_finalizer_close_in_a_loops_last_step_reports_no_unretrieved_exception():
loop: Final = asyncio.new_event_loop()
exception_handler: Final = MagicMock()
loop.set_exception_handler(exception_handler)
loop.run_until_complete(_drop_handler_whose_close_fails())
loop.close()
del loop
gc.collect()
exception_handler.assert_not_called()
async def _sync_client_with_timeout_62() -> HTTPHandler:
return get_httpx_client(params={"timeout": 62.0})
def test_sync_client_is_shared_by_short_lived_loops_and_callers_outside_a_loop():
built_in_first_loop: Final = asyncio.run(_sync_client_with_timeout_62())
assert asyncio.run(_sync_client_with_timeout_62()) is built_in_first_loop
assert get_httpx_client(params={"timeout": 62.0}) is built_in_first_loop
@pytest.mark.asyncio
async def test_sync_close_helper_respects_session_ownership():
"""The loop-less fallback closes only sessions the transport owns; a

View file

@ -1,3 +1,4 @@
import asyncio
from typing import Final
from unittest.mock import MagicMock, patch
@ -154,6 +155,28 @@ def test_get_openai_client_initialization_param_fields(client_type):
assert "self" not in result
def _loop_test_client_params(is_async: bool) -> dict[str, object]:
return {"api_key": "sk-loop-scope-test", "is_async": is_async}
async def _cache_openai_client(client: openai.OpenAI | openai.AsyncOpenAI, is_async: bool) -> None:
BaseOpenAILLM.set_cached_openai_client(client, "openai", _loop_test_client_params(is_async))
async def _cached_openai_client(is_async: bool) -> object:
return BaseOpenAILLM.get_cached_openai_client(_loop_test_client_params(is_async), "openai")
def test_sync_openai_client_cached_in_one_loop_is_reused_by_the_next_while_async_is_not():
sync_client: Final = openai.OpenAI(api_key="sk-loop-scope-test")
async_client: Final = openai.AsyncOpenAI(api_key="sk-loop-scope-test")
asyncio.run(_cache_openai_client(sync_client, is_async=False))
asyncio.run(_cache_openai_client(async_client, is_async=True))
assert asyncio.run(_cached_openai_client(is_async=False)) is sync_client
assert asyncio.run(_cached_openai_client(is_async=True)) is None
@pytest.mark.parametrize("client_type", ["openai", "azure"])
def test_get_openai_client_cache_key(client_type):
"""Verify get_openai_client_cache_key doesn't raise on tuple + tuple concatenation."""

View file

@ -2,6 +2,7 @@
import asyncio
import contextlib
import gc
import io
import json
import logging
@ -2966,6 +2967,188 @@ class TestStreamingScanDedup:
assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]
class _BackgroundScanGuardrail(_MarkerBlockingScanGuardrail):
"""Background-scanned guardrail whose scans hold until released, then block BLOCKME and rewrite other text"""
def __init__(self) -> None:
super().__init__(sampling_rate=2)
self.streaming_scan_in_background = True
self.released = anyio.Event()
self.scanned = anyio.Event()
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
**kwargs: object,
) -> GenericGuardrailAPIInputs:
try:
await self.released.wait()
recorded = await super().apply_guardrail(inputs, request_data, input_type, **kwargs)
finally:
self.scanned.set()
return {**recorded, "texts": ["REWRITTEN" for _ in recorded.get("texts") or []]}
class TestStreamingBackgroundScan:
"""streaming_scan_in_background: the stream keeps flowing while a sampled scan runs, and a
failed scan ends it before the next chunk. LIT-9338."""
@pytest.fixture(autouse=True)
def _use_real_mappings(self, monkeypatch: pytest.MonkeyPatch) -> None:
_patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings())
@staticmethod
def _hook(guardrail: CustomGuardrail, upstream: AsyncIterator[ModelResponseStream]) -> AsyncGenerator[object, None]:
return UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"),
response=upstream,
request_data={"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}},
)
@classmethod
def _guarded_stream(cls, guardrail: CustomGuardrail, *texts: str) -> AsyncGenerator[object, None]:
async def upstream() -> AsyncIterator[ModelResponseStream]:
for text in texts:
yield _stream_chunk(text)
yield _stream_chunk("", finish_reason="stop")
return cls._hook(guardrail, upstream())
@staticmethod
async def _receive(stream: AsyncGenerator[object, None], count: int) -> list[object]:
with anyio.fail_after(5):
return [await stream.__anext__() for _ in range(count)]
@staticmethod
def _block_frames(items: list[object]) -> bytes:
assert not any(isinstance(item, ModelResponseStream) for item in items), items
return b"".join(item for item in items if isinstance(item, bytes))
@pytest.mark.asyncio
async def test_chunks_stream_while_the_scan_runs_and_its_rewrite_never_reaches_them(self) -> None:
guardrail: Final = _BackgroundScanGuardrail()
stream: Final = self._guarded_stream(guardrail, "a", "b", "c", "d")
sent: Final = await self._receive(stream, 4)
assert guardrail.scans == ()
guardrail.released.set()
rest: Final = [item async for item in stream]
assert [_delta_text(item) for item in (*sent, *rest)] == ["a", "b", "c", "d", ""]
assert [scan["texts"] for scan in guardrail.scans] == [["ab"], ["abcd"]], guardrail.scans
@pytest.mark.asyncio
async def test_a_block_ends_the_stream_before_the_next_chunk(self) -> None:
guardrail: Final = _BackgroundScanGuardrail()
stream: Final = self._guarded_stream(guardrail, "a", "BLOCKME", "c", "d", "e")
sent: Final = await self._receive(stream, 3)
guardrail.released.set()
await guardrail.scanned.wait()
rest: Final = [item async for item in stream]
assert [_delta_text(item) for item in sent] == ["a", "BLOCKME", "c"]
assert b"blocked" in self._block_frames(rest)
assert [scan["texts"] for scan in guardrail.scans] == [["aBLOCKME"]], guardrail.scans
@pytest.mark.asyncio
async def test_a_block_that_lands_after_the_last_text_replaces_the_finish(self) -> None:
guardrail: Final = _BackgroundScanGuardrail()
stream: Final = self._guarded_stream(guardrail, "a", "BLOCKME", "c", "d", "e", "f")
sent: Final = await self._receive(stream, 6)
guardrail.released.set()
rest: Final = [item async for item in stream]
assert [_delta_text(item) for item in sent] == ["a", "BLOCKME", "c", "d", "e", "f"]
assert b"blocked" in self._block_frames(rest)
assert [scan["texts"] for scan in guardrail.scans] == [["aBLOCKME"]], guardrail.scans
@pytest.mark.asyncio
async def test_a_block_from_a_scan_that_covered_all_the_text_still_ends_the_stream(self) -> None:
guardrail: Final = _BackgroundScanGuardrail()
guardrail.released.set()
out: Final = [item async for item in self._guarded_stream(guardrail, "a", "BLOCKME")]
assert [_delta_text(item) for item in out if isinstance(item, ModelResponseStream)] == ["a", "BLOCKME"]
assert b"blocked" in b"".join(item for item in out if isinstance(item, bytes))
assert [scan["texts"] for scan in guardrail.scans] == [["aBLOCKME"]], guardrail.scans
@pytest.mark.asyncio
async def test_an_upstream_error_cancels_the_running_scan(self) -> None:
guardrail: Final = _BackgroundScanGuardrail()
async def upstream() -> AsyncIterator[ModelResponseStream]:
yield _stream_chunk("a")
yield _stream_chunk("b")
await anyio.wait_all_tasks_blocked()
raise RuntimeError("upstream dropped")
with anyio.fail_after(5), pytest.raises(RuntimeError, match="upstream dropped"):
_ = [item async for item in self._hook(guardrail, upstream())]
assert guardrail.scanned.is_set(), "the scan was still running when the upstream error surfaced"
guardrail.released.set()
await anyio.wait_all_tasks_blocked()
assert guardrail.scans == ()
@pytest.mark.asyncio
async def test_an_upstream_error_after_a_failed_scan_reports_only_the_upstream_error(
self, caplog: pytest.LogCaptureFixture
) -> None:
guardrail: Final = _BackgroundScanGuardrail()
guardrail.released.set()
async def upstream() -> AsyncIterator[ModelResponseStream]:
yield _stream_chunk("a")
yield _stream_chunk("BLOCKME")
await guardrail.scanned.wait()
await anyio.wait_all_tasks_blocked()
raise RuntimeError("upstream dropped")
with pytest.raises(RuntimeError, match="upstream dropped"):
_ = [item async for item in self._hook(guardrail, upstream())]
gc.collect()
assert "never retrieved" not in caplog.text
@pytest.mark.asyncio
async def test_closing_mid_scan_waits_for_it_then_scans_what_was_sent(self) -> None:
guardrail: Final = _BackgroundScanGuardrail()
stream: Final = self._guarded_stream(guardrail, "a", "b", "c", "d")
await self._receive(stream, 3)
guardrail.released.set()
await stream.aclose()
assert [scan["texts"] for scan in guardrail.scans] == [["ab"], ["abc"]], guardrail.scans
@pytest.mark.asyncio
async def test_closing_after_a_background_block_does_not_scan_again(self) -> None:
guardrail: Final = _BackgroundScanGuardrail()
stream: Final = self._guarded_stream(guardrail, "a", "BLOCKME", "c", "d")
await self._receive(stream, 3)
guardrail.released.set()
await stream.aclose()
assert [scan["texts"] for scan in guardrail.scans] == [["aBLOCKME"]], guardrail.scans
@pytest.mark.asyncio
async def test_buffered_release_on_scan_ignores_the_flag(self) -> None:
guardrail: Final = _BackgroundScanGuardrail()
guardrail.streaming_buffer_until_moderated = True
guardrail.streaming_buffer_release_on_scan = True
guardrail.released.set()
out: Final = [item async for item in self._guarded_stream(guardrail, "BLOCKME", "b", "c", "d")]
assert b"blocked" in self._block_frames(out)
class TestTranslationMappingsAreReadLive:
"""The hooks must read the handler map on every call, never memoize it on the module.

View file

@ -95,8 +95,8 @@ async def test_claude_count_request_counts_additional_tools_input_items_as_tools
@pytest.mark.asyncio
async def test_claude_count_request_counts_a_tool_list_repeated_on_every_turn_once() -> None:
policy: Final = DropPolicy(drop_params=False, additional_drop_params=frozenset())
async def test_claude_count_request_counts_a_tool_list_repeated_on_every_turn_once_under_drop_params() -> None:
policy: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset())
tool_list: Final[dict[str, JsonValue]] = {"type": "additional_tools", "tools": [_WEATHER_TOOL]}
repeated: Final = await _claude_count(

View file

@ -15,7 +15,7 @@ _USER_MESSAGE = {"type": "message", "role": "user", "content": "Run ls"}
def test_string_input_passes_through_with_existing_tools():
hoisted = hoist_additional_tools("hello", [_TOP_LEVEL_TOOL])
hoisted = hoist_additional_tools("hello", [_TOP_LEVEL_TOOL], drop_params=False)
assert hoisted.input == "hello"
assert hoisted.tools == (_TOP_LEVEL_TOOL,)
@ -25,7 +25,7 @@ def test_string_input_passes_through_with_existing_tools():
def test_input_without_additional_tools_items_is_returned_untouched():
request_input = [_USER_MESSAGE]
hoisted = hoist_additional_tools(request_input, None)
hoisted = hoist_additional_tools(request_input, None, drop_params=False)
assert hoisted.input is request_input
assert hoisted.tools == ()
@ -39,7 +39,7 @@ def test_additional_tools_items_are_stripped_and_appended_after_top_level_tools_
{"type": "additional_tools", "id": "at_2", "role": "developer", "tools": [_WAIT_TOOL]},
]
hoisted = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL])
hoisted = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL], drop_params=False)
assert hoisted.input == [_USER_MESSAGE]
assert hoisted.tools == (_TOP_LEVEL_TOOL, _EXEC_TOOL, _WAIT_TOOL)
@ -49,7 +49,7 @@ def test_additional_tools_items_are_stripped_and_appended_after_top_level_tools_
def test_additional_tools_item_without_a_tools_list_is_stripped_and_contributes_nothing():
request_input = [{"type": "additional_tools", "id": "at_1", "role": "developer", "tools": "exec"}, _USER_MESSAGE]
hoisted = hoist_additional_tools(request_input, None)
hoisted = hoist_additional_tools(request_input, None, drop_params=False)
assert hoisted.input == [_USER_MESSAGE]
assert hoisted.tools == ()
@ -77,7 +77,7 @@ def test_tool_list_repeated_on_every_turn_keeps_each_tool_once_with_the_newest_d
_USER_MESSAGE,
]
hoisted: Final = hoist_additional_tools(request_input, None)
hoisted: Final = hoist_additional_tools(request_input, None, drop_params=True)
assert hoisted.tools == (collaboration, _function("exec", "new"))
@ -89,7 +89,7 @@ def test_namespace_sent_in_several_turns_keeps_every_distinct_member_once() -> N
_tool_list(_namespace("mcp__docs__", _function("fetch", "new"), _function("list"))),
]
hoisted: Final = hoist_additional_tools(request_input, None)
hoisted: Final = hoist_additional_tools(request_input, None, drop_params=True)
assert hoisted.tools == (
_namespace("mcp__docs__", _function("search"), _function("fetch", "new"), _function("list")),
@ -109,7 +109,7 @@ def test_unnamed_tools_collapse_only_when_identical() -> None:
}
request_input: Final = [_tool_list(docs_server, wiki_server), _USER_MESSAGE, _tool_list(docs_server)]
hoisted: Final = hoist_additional_tools(request_input, None)
hoisted: Final = hoist_additional_tools(request_input, None, drop_params=True)
assert hoisted.tools == (docs_server, wiki_server)
@ -117,6 +117,20 @@ def test_unnamed_tools_collapse_only_when_identical() -> None:
def test_top_level_tool_repeated_in_a_tool_list_item_is_kept_once() -> None:
request_input: Final = [_tool_list(_TOP_LEVEL_TOOL, _WAIT_TOOL), _USER_MESSAGE]
hoisted: Final = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL])
hoisted: Final = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL], drop_params=True)
assert hoisted.tools == (_TOP_LEVEL_TOOL, _WAIT_TOOL)
def test_tool_list_repeated_on_every_turn_is_sent_as_is_without_drop_params() -> None:
request_input: Final = [
_tool_list(_function("exec", "old")),
_USER_MESSAGE,
_tool_list(_function("exec", "new")),
_USER_MESSAGE,
]
hoisted: Final = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL], drop_params=False)
assert hoisted.tools == (_TOP_LEVEL_TOOL, _function("exec", "old"), _function("exec", "new"))
assert hoisted.input == [_USER_MESSAGE, _USER_MESSAGE]

View file

@ -0,0 +1,124 @@
import base64
import importlib.util
import io
import sys
import urllib.parse
import zipfile
from collections.abc import Callable, Mapping
from itertools import chain, product
from pathlib import Path
from types import MappingProxyType
from typing import Final
import pytest
_REPO_ROOT: Final = Path(__file__).resolve().parents[2]
_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "assert_secret_absent.py"
_spec: Final = importlib.util.spec_from_file_location("assert_secret_absent", _MODULE_PATH)
scanner: Final = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = scanner
_spec.loader.exec_module(scanner)
SECRET: Final = "lic/ense+KEY:9f8e~~??>>~?7d6c5b4a3210-zz"
PREFIXES: Final = ("", "x", "xy")
ENCODERS: Final = (base64.b64encode, base64.urlsafe_b64encode)
def _zip(members: Mapping[str, bytes], compression: int = zipfile.ZIP_DEFLATED) -> bytes:
buffer: Final = io.BytesIO()
with zipfile.ZipFile(buffer, "w", compression) as archive:
for name, data in members.items():
archive.writestr(name, data)
return buffer.getvalue()
def _blob(prefix: str) -> bytes:
return f"{prefix}{SECRET}tail".encode()
def _planted_forms() -> tuple[str, ...]:
encoded: Final = tuple(encode(_blob(prefix)).decode() for encode, prefix in product(ENCODERS, PREFIXES))
return (SECRET, urllib.parse.quote(SECRET, safe=""), *encoded)
def _assert_nothing_leaked(output: str) -> None:
windows: Final = frozenset(
chain.from_iterable((form[i : i + 8] for i in range(len(form) - 7)) for form in _planted_forms())
)
assert not [window for window in windows if window in output]
def _run(root: Path, secret: str = SECRET) -> tuple[int, str]:
out: Final = io.StringIO()
code: Final = scanner.main((str(root),), {scanner.SECRET_VARIABLE: secret}, out)
return code, out.getvalue()
def _encoded_blob(prefix: str, encode: Callable[[bytes], bytes]) -> bytes:
return b"trace " + encode(_blob(prefix)) + b" end"
QUOTED: Final = urllib.parse.quote(SECRET, safe="")
PLACEMENTS: Final = MappingProxyType(
{
"raw in a log": {"proxy.log": f"boot with key {SECRET} ok".encode()},
"url-encoded in a url": {"upstream.log": f"GET /verify_license/{QUOTED}".encode()},
"base64 aligned": {"a.json": _encoded_blob("", base64.b64encode)},
"base64 one byte in": {"a.json": _encoded_blob("x", base64.b64encode)},
"base64 two bytes in": {"a.json": _encoded_blob("xy", base64.b64encode)},
"base64url two bytes in": {"a.json": _encoded_blob("xy", base64.urlsafe_b64encode)},
"inside a zip member": {"trace.zip": _zip({"network.json": f'{{"auth": "{SECRET}"}}'.encode()})},
"inside a nested zip": {"outer.zip": _zip({"inner.zip": _zip({"env.txt": SECRET.encode()})})},
"named after it in a nested zip": {
"outer.zip": _zip({"inner.zip": _zip({f"dump-{QUOTED}.txt": SECRET.encode()})})
},
"in a file name": {f"run-{QUOTED}.log": b"nothing inside"},
}
)
@pytest.mark.parametrize("placement", PLACEMENTS)
def test_every_planted_copy_blocks_the_upload_without_printing_the_secret(tmp_path: Path, placement: str) -> None:
worker: Final = tmp_path / "gw0"
worker.mkdir()
for name, data in PLACEMENTS[placement].items():
(worker / name).write_bytes(data)
code, output = _run(tmp_path)
assert code == 1
assert "secret found:" in output
_assert_nothing_leaked(output)
def test_clean_results_pass_and_count_every_file(tmp_path: Path) -> None:
(tmp_path / "junit.xml").write_text("<testsuite/>")
(tmp_path / "gw0").mkdir()
(tmp_path / "gw0" / "proxy.log").write_text("LITELLM_LICENSE is not set")
code, output = _run(tmp_path)
assert code == 0
assert "no copy of the secret in 2 files" in output
def test_an_unset_secret_has_nothing_to_find(tmp_path: Path) -> None:
(tmp_path / "proxy.log").write_text(SECRET)
assert _run(tmp_path, secret="")[0] == 0
def test_a_symlinked_directory_fails_closed(tmp_path: Path) -> None:
outside: Final = tmp_path / "outside"
outside.mkdir()
results: Final = tmp_path / "results"
results.mkdir()
(results / "linked").symlink_to(outside, target_is_directory=True)
code, output = _run(results)
assert code == 1
assert "directory symlink" in output
def test_an_unreadable_archive_fails_without_echoing_its_contents(tmp_path: Path) -> None:
archive: Final = _zip({"env.txt": f"LITELLM_LICENSE={SECRET}".encode()}, zipfile.ZIP_STORED)
at: Final = archive.index(SECRET.encode())
(tmp_path / "trace.zip").write_bytes(archive[:at] + bytes([archive[at] ^ 0xFF]) + archive[at + 1 :])
code, output = _run(tmp_path)
assert code == 1
assert "scan failed with BadZipFile" in output
_assert_nothing_leaked(output)