From 4cd3d97eb18a79fe7ee946de52a386b0f0a09fc6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 22:30:55 +0000 Subject: [PATCH] test(otel v2): close audit gaps - C2 diagnostics cell, newrelic CONNECT delivery, arize gRPC cells, team metadata rejection cell Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/otlp_sink.py | 213 +++++++++++++++++- tests/integration/observability/conftest.py | 23 +- .../test_otel_tenant_internal_spans.py | 148 +++++++++++- .../test_otel_tenant_internal_spans_chaos.py | 35 +++ 4 files changed, 403 insertions(+), 16 deletions(-) diff --git a/tests/integration/_support/otlp_sink.py b/tests/integration/_support/otlp_sink.py index 4a0fbccd24b..eeabe9d886f 100644 --- a/tests/integration/_support/otlp_sink.py +++ b/tests/integration/_support/otlp_sink.py @@ -9,15 +9,17 @@ steer the sink through ``configure``; the process can also be frozen with from __future__ import annotations import argparse +import datetime import json import os import signal import socket +import ssl import subprocess import sys import threading import time -from collections.abc import Iterator, Mapping +from collections.abc import Iterator, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -178,7 +180,14 @@ class _Handler(BaseHTTPRequestHandler): time.sleep(self.state.delay_seconds) recorded: Final = decode_spans(body, self.headers.get("content-type", "")) self.state.spans.extend(recorded) - self.state.requests.append({"path": self.path, "count": len(recorded)}) + self.state.requests.append( + { + "path": self.path, + "count": len(recorded), + "host": self.headers.get("host", ""), + "headers": dict(self.headers), + } + ) self._send_json({"recorded": len(recorded)}, status=self.state.status) do_POST = _record @@ -225,6 +234,100 @@ class _Handler(BaseHTTPRequestHandler): pass +class _ConnectHandler(_Handler): + tunnel_context: ssl.SSLContext + + def do_CONNECT(self) -> None: + self.state.requests.append({"connect": self.path}) + self.connection.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + wrapped: Final = self.tunnel_context.wrap_socket(self.connection, server_side=True) + self.close_connection = True + type(self)(wrapped, self.client_address, self.server) + + +_MITM_HOSTS: Final = ("otlp.nr-data.net", "otlp.eu01.nr-data.net") + + +def _mitm_context(directory: Path) -> ssl.SSLContext: + from cryptography import x509 + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + + directory.mkdir(parents=True, exist_ok=True) + now: Final = datetime.datetime.now(datetime.timezone.utc) + window: Final = datetime.timedelta(days=2) + ca_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + ca_name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "otlp-sink test CA")]) + ca_cert: Final = ( + x509.CertificateBuilder() + .subject_name(ca_name) + .issuer_name(ca_name) + .public_key(ca_key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - window) + .not_valid_after(now + window) + .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True) + .sign(ca_key, hashes.SHA256()) + ) + leaf_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + leaf_cert: Final = ( + x509.CertificateBuilder() + .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, _MITM_HOSTS[0])])) + .issuer_name(ca_cert.subject) + .public_key(leaf_key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - window) + .not_valid_after(now + window) + .add_extension(x509.SubjectAlternativeName([x509.DNSName(host) for host in _MITM_HOSTS]), critical=False) + .sign(ca_key, hashes.SHA256()) + ) + ca_pem: Final = directory / "ca.pem" + ca_pem.write_bytes(ca_cert.public_bytes(serialization.Encoding.PEM)) + leaf_pem: Final = directory / "leaf.pem" + leaf_pem.write_bytes(leaf_cert.public_bytes(serialization.Encoding.PEM)) + leaf_key_pem: Final = directory / "leaf-key.pem" + leaf_key_pem.write_bytes( + leaf_key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ) + ) + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(str(leaf_pem), str(leaf_key_pem)) + return context + + +def _grpc_trace_server(state: _State, port: int) -> object: + from concurrent import futures + + import grpc + from opentelemetry.proto.collector.trace.v1 import trace_service_pb2, trace_service_pb2_grpc + + class _TraceService(trace_service_pb2_grpc.TraceServiceServicer): + def Export(self, request: object, context: grpc.ServicerContext) -> object: + state.pause.wait(timeout=120) + if state.delay_seconds > 0: + time.sleep(state.delay_seconds) + recorded: Final = _proto_spans(request.SerializeToString()) + state.spans.extend(recorded) + state.requests.append( + { + "grpc": "Export", + "metadata": {key: value for key, value in context.invocation_metadata()}, + "count": len(recorded), + } + ) + return trace_service_pb2.ExportTraceServiceResponse() + + server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + trace_service_pb2_grpc.add_TraceServiceServicer_to_server(_TraceService(), server) + server.add_insecure_port(f"127.0.0.1:{port}") + server.start() + return server + + def recorded_spans(url: str, since: int = 0) -> tuple[int, tuple[Span, ...]]: response: Final = httpx.get(f"{url}/__spans", params={"since": since}, trust_env=False, timeout=15) response.raise_for_status() @@ -244,6 +347,15 @@ def sink_pid(url: str) -> int: return int(httpx.get(f"{url}/__pid", trust_env=False, timeout=15).json()["pid"]) +_REQUEST_LISTING: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def recorded_requests(url: str) -> tuple[Mapping[str, JsonValue], ...]: + response: Final = httpx.get(f"{url}/__requests", trust_env=False, timeout=15) + response.raise_for_status() + return tuple(_REQUEST_LISTING.validate_python(response.json()["requests"])) + + @dataclass(frozen=True, slots=True) class SpanSinks: operator: str @@ -251,6 +363,19 @@ class SpanSinks: arize: str +@dataclass(frozen=True, slots=True) +class GrpcSink: + url: str + control_url: str + + +@dataclass(frozen=True, slots=True) +class ConnectSink: + proxy_url: str + control_url: str + ca_pem: str + + def _free_port() -> int: with socket.socket() as reserve: reserve.bind(("127.0.0.1", 0)) @@ -304,14 +429,96 @@ def owned_sinks(directory: Path) -> Iterator[SpanSinks]: assert not survivors and stopped, "OTLP sink required forced cleanup" +@contextmanager +def _spawn_sink(directory: Path, log_name: str, argv: Sequence[str]) -> Iterator[None]: + from integration._support.process import group_members, signal_group, stop_root_process + + directory.mkdir(parents=True, exist_ok=True) + root: Final = Path(__file__).resolve().parents[3] + with (directory / log_name).open("w") as log: + process: Final = subprocess.Popen( + [sys.executable, "-P", "-m", "integration._support.otlp_sink", *argv], + cwd=root, + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + try: + yield + finally: + stopped: Final = stop_root_process(process) + residual: Final = group_members(process.pid) + if residual: + signal_group(process.pid, signal.SIGKILL) + psutil.wait_procs(residual, timeout=5) + survivors: Final = group_members(process.pid) + assert not survivors and stopped, "OTLP sink required forced cleanup" + + +def _await_sink(url: str) -> None: + deadline: Final = time.monotonic() + 30 + while not _pid_reachable(url): + assert time.monotonic() < deadline, "OTLP sink readiness deadline exceeded" + time.sleep(0.05) + + +@contextmanager +def owned_grpc_sink(directory: Path) -> Iterator[GrpcSink]: + http_port: Final = _free_port() + grpc_port: Final = _free_port() + with _spawn_sink( + directory, "otlp-grpc-sink.log", ["--port", str(http_port), "--grpc-port", str(grpc_port)] + ): + control_url: Final = f"http://127.0.0.1:{http_port}" + _await_sink(control_url) + yield GrpcSink(url=f"http://127.0.0.1:{grpc_port}", control_url=control_url) + + +@contextmanager +def owned_connect_sink(directory: Path) -> Iterator[ConnectSink]: + http_port: Final = _free_port() + tunnel_port: Final = _free_port() + ca_dir: Final = directory / "mitm" + with _spawn_sink( + directory, + "otlp-connect-sink.log", + ["--port", str(http_port), "--connect-port", str(tunnel_port), "--ca-dir", str(ca_dir)], + ): + control_url: Final = f"http://127.0.0.1:{http_port}" + _await_sink(control_url) + yield ConnectSink( + proxy_url=f"http://127.0.0.1:{tunnel_port}", + control_url=control_url, + ca_pem=str(ca_dir / "ca.pem"), + ) + + def main() -> None: parser: Final = argparse.ArgumentParser() parser.add_argument("--port", type=int, required=True) + parser.add_argument("--grpc-port", type=int, default=0) + parser.add_argument("--connect-port", type=int, default=0) + parser.add_argument("--ca-dir", type=Path, default=None) arguments: Final = parser.parse_args() + bound_state: Final = _State() class BoundHandler(_Handler): - state = _State() + state = bound_state + if arguments.grpc_port: + grpc_server: Final = _grpc_trace_server(bound_state, arguments.grpc_port) + assert grpc_server is not None + if arguments.connect_port: + assert arguments.ca_dir is not None, "--connect-port needs --ca-dir" + bound_context: Final = _mitm_context(arguments.ca_dir) + + class BoundConnectHandler(_ConnectHandler): + state = bound_state + tunnel_context = bound_context # pyright: ignore[reportIncompatibleVariableOverride] # bound context, not a new field + + tunnel: Final = ThreadingHTTPServer(("127.0.0.1", arguments.connect_port), BoundConnectHandler) + tunnel.daemon_threads = True + threading.Thread(target=tunnel.serve_forever, daemon=True).start() server: Final = ThreadingHTTPServer(("127.0.0.1", arguments.port), BoundHandler) server.daemon_threads = True server.serve_forever() diff --git a/tests/integration/observability/conftest.py b/tests/integration/observability/conftest.py index f09a703f047..42ba22e455e 100644 --- a/tests/integration/observability/conftest.py +++ b/tests/integration/observability/conftest.py @@ -8,7 +8,14 @@ from urllib.parse import urlparse import pytest import yaml -from integration._support.otlp_sink import SpanSinks, owned_sinks +from integration._support.otlp_sink import ( + ConnectSink, + GrpcSink, + SpanSinks, + owned_connect_sink, + owned_grpc_sink, + owned_sinks, +) from pydantic import JsonValue AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path] @@ -21,6 +28,20 @@ def audit_sinks(tmp_path_factory: pytest.TempPathFactory) -> Iterator[SpanSinks] yield sinks +@pytest.fixture(scope="module") +def newrelic_sink(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ConnectSink]: + directory: Final = tmp_path_factory.mktemp("otel-audit-connect") + with owned_connect_sink(directory) as sink: + yield sink + + +@pytest.fixture(scope="module") +def arize_grpc_sink(tmp_path_factory: pytest.TempPathFactory) -> Iterator[GrpcSink]: + directory: Final = tmp_path_factory.mktemp("otel-audit-grpc") + with owned_grpc_sink(directory) as sink: + yield sink + + @pytest.fixture(scope="module") def otel_audit_config(audit_sinks: SpanSinks) -> AuditConfigWriter: tenant_host: Final = urlparse(audit_sinks.tenant).netloc diff --git a/tests/integration/observability/test_otel_tenant_internal_spans.py b/tests/integration/observability/test_otel_tenant_internal_spans.py index d1589588a14..d1a14d89222 100644 --- a/tests/integration/observability/test_otel_tenant_internal_spans.py +++ b/tests/integration/observability/test_otel_tenant_internal_spans.py @@ -21,9 +21,12 @@ from integration._support.client import ( ) from integration._support.database import read_rows from integration._support.otlp_sink import ( + ConnectSink, + GrpcSink, Span, SpanSinks, configure_sink, + recorded_requests, recorded_spans, span_class, spans_for_trace, @@ -850,6 +853,122 @@ def test_internal_spans_on_newrelic_accepted(gateway: Gateway) -> None: assert "newrelic" in data.get("success_callbacks", []) or "newrelic" in data.get("failure_callbacks", []), data +def test_newrelic_exclude_over_connect( + gateway: Gateway, audit_sinks: SpanSinks, newrelic_sink: ConnectSink, otel_audit_config: AuditConfigWriter, tmp_path: Path +) -> None: + _newrelic_flow(gateway, audit_sinks, newrelic_sink, otel_audit_config, tmp_path, internal_spans="exclude") + + +def test_newrelic_include_over_connect( + gateway: Gateway, audit_sinks: SpanSinks, newrelic_sink: ConnectSink, otel_audit_config: AuditConfigWriter, tmp_path: Path +) -> None: + _newrelic_flow(gateway, audit_sinks, newrelic_sink, otel_audit_config, tmp_path, internal_spans="include") + + +def _newrelic_flow( + gateway: Gateway, + audit_sinks: SpanSinks, + newrelic_sink: ConnectSink, + otel_audit_config: AuditConfigWriter, + tmp_path: Path, + *, + internal_spans: str, +) -> None: + overrides: Final = { + "LITELLM_OTEL_V2": "1", + "ARIZE_HTTP_ENDPOINT": audit_sinks.arize, + "HTTPS_PROXY": newrelic_sink.proxy_url, + "NO_PROXY": "127.0.0.1,localhost", + "OTEL_EXPORTER_OTLP_TRACES_CERTIFICATE": newrelic_sink.ca_pem, + } + with owned_proxy(gateway, tmp_path, overrides, config=otel_audit_config(tmp_path, {}), workers=2) as candidate: + with candidate.scenario() as scenario: + model: Final = _audit_model(scenario, candidate.upstream_url) + team_id: Final = scenario.team() + callback: Final = _add_callback( + candidate, + team_id, + {"newrelic_api_key": "nr-synthetic-audit", INTERNAL_SPANS_VAR: internal_spans}, + callback_name="newrelic", + ) + assert callback.status_code == 200, callback.text + key: Final = _key_on_team(scenario, team_id) + response: Final = _chat(candidate, key, model, _nonce()) + assert response.status_code == 200, response.text + call_id: Final = _call_id(response) + trace: Final = _trace_id(newrelic_sink.control_url, call_id=call_id) + group: Final = _trace_spans(newrelic_sink.control_url, trace) + if internal_spans == "exclude": + _assert_excluded(group) + else: + _assert_full(group) + requests: Final = recorded_requests(newrelic_sink.control_url) + assert any(entry.get("connect") == "otlp.nr-data.net:443" for entry in requests), requests + posted: Final = tuple(entry for entry in requests if entry.get("path") == "/v1/traces") + assert posted, f"tunnel saw no /v1/traces posts: {requests}" + for entry in posted: + assert entry.get("host") == "otlp.nr-data.net", entry + headers: Final = entry.get("headers") + assert isinstance(headers, Mapping) and headers.get("api-key") == "nr-synthetic-audit", entry + + +def test_arize_grpc_exclude( + gateway: Gateway, audit_sinks: SpanSinks, arize_grpc_sink: GrpcSink, otel_audit_config: AuditConfigWriter, tmp_path: Path +) -> None: + _arize_grpc_flow(gateway, audit_sinks, arize_grpc_sink, otel_audit_config, tmp_path, internal_spans="exclude") + + +def test_arize_grpc_include( + gateway: Gateway, audit_sinks: SpanSinks, arize_grpc_sink: GrpcSink, otel_audit_config: AuditConfigWriter, tmp_path: Path +) -> None: + _arize_grpc_flow(gateway, audit_sinks, arize_grpc_sink, otel_audit_config, tmp_path, internal_spans="include") + + +def _arize_grpc_flow( + gateway: Gateway, + audit_sinks: SpanSinks, + arize_grpc_sink: GrpcSink, + otel_audit_config: AuditConfigWriter, + tmp_path: Path, + *, + internal_spans: str, +) -> None: + overrides: Final = { + "LITELLM_OTEL_V2": "1", + "ARIZE_ENDPOINT": arize_grpc_sink.url, + "ARIZE_HTTP_ENDPOINT": audit_sinks.arize, + } + with owned_proxy(gateway, tmp_path, overrides, config=otel_audit_config(tmp_path, {}), workers=2) as candidate: + with candidate.scenario() as scenario: + model: Final = _audit_model(scenario, candidate.upstream_url) + team_id: Final = scenario.team() + callback: Final = _add_callback( + candidate, + team_id, + {**ARIZE_VARS, INTERNAL_SPANS_VAR: internal_spans}, + callback_name="arize", + ) + assert callback.status_code == 200, callback.text + key: Final = _key_on_team(scenario, team_id) + response: Final = _chat(candidate, key, model, _nonce()) + assert response.status_code == 200, response.text + call_id: Final = _call_id(response) + trace: Final = _trace_id(arize_grpc_sink.control_url, call_id=call_id) + group: Final = _trace_spans(arize_grpc_sink.control_url, trace) + if internal_spans == "exclude": + _assert_excluded(group) + else: + _assert_full(group) + requests: Final = recorded_requests(arize_grpc_sink.control_url) + exports: Final = tuple(entry for entry in requests if entry.get("grpc") == "Export") + assert exports, f"grpc sink saw no Export calls: {requests}" + for entry in exports: + metadata: Final = entry.get("metadata") + assert isinstance(metadata, Mapping), entry + assert metadata.get("arize-space-id") == ARIZE_VARS["arize_space_id"], entry + assert metadata.get("api_key") == ARIZE_VARS["arize_api_key"], entry + + def test_unauthenticated_callback_post_rejected(gateway: Gateway, langfuse_vars: dict[str, JsonValue]) -> None: with httpx.Client(base_url=str(gateway.client.base_url), timeout=15, trust_env=False) as client: response: Final = client.post( @@ -876,27 +995,32 @@ def test_key_update_bogus_internal_spans_rejected(gateway: Gateway, langfuse_var assert response.status_code == 400, f"expected 400, got {response.status_code}: {response.text}" -def test_team_update_bogus_internal_spans_drops_destination(gateway: Gateway, audit_sinks: SpanSinks, langfuse_vars: dict[str, JsonValue]) -> None: +def test_team_logging_metadata_rejects_bogus_internal_spans(gateway: Gateway, audit_sinks: SpanSinks, langfuse_vars: dict[str, JsonValue]) -> None: with gateway.scenario() as scenario: model: Final = _audit_model(scenario, gateway.upstream_url) team_id: Final = scenario.team() + registered: Final = _add_callback(gateway, team_id, langfuse_vars) + assert registered.status_code == 200, registered.text + bogus_metadata: Final[Mapping[str, JsonValue]] = { + "logging": _key_logging_entry({**langfuse_vars, INTERNAL_SPANS_VAR: "bogus"}) + } update: Final = gateway.request( - "POST", - "/team/update", - {"team_id": team_id, "metadata": {"logging": _key_logging_entry({**langfuse_vars, INTERNAL_SPANS_VAR: "bogus"})}}, + "POST", "/team/update", {"team_id": team_id, "metadata": bogus_metadata} ) - assert update.status_code == 200, update.text + assert update.status_code == 400, update.text + assert "otel_internal_spans" in update.text, update.text + created: Final = gateway.request("POST", "/team/new", {"metadata": bogus_metadata}) + assert created.status_code == 400, created.text + assert "otel_internal_spans" in created.text, created.text key: Final = _key_on_team(scenario, team_id) response: Final = _chat(gateway, key, model, _nonce()) assert response.status_code == 200, response.text call_id: Final = _call_id(response) - operator_trace: Final = _trace_id(audit_sinks.operator, call_id=call_id) - _assert_full(_trace_spans(audit_sinks.operator, operator_trace)) - _, observed = recorded_spans(audit_sinks.tenant, since=0) - leaked: Final = tuple( - span for span in observed if span["attributes"].get("litellm.call_id") == call_id - ) - assert not leaked, f"bogus metadata.logging still reached the tenant sink: {leaked}" + tenant_trace: Final = _trace_id(audit_sinks.tenant, call_id=call_id) + _assert_full(_trace_spans(audit_sinks.tenant, tenant_trace)) + unrelated_key: Final = scenario.key() + unrelated: Final = _chat(gateway, unrelated_key, model, _nonce()) + assert unrelated.status_code == 200, unrelated.text def _sink_status_flow( diff --git a/tests/integration/observability/test_otel_tenant_internal_spans_chaos.py b/tests/integration/observability/test_otel_tenant_internal_spans_chaos.py index 32e8b6ccf5e..3a406e0a72a 100644 --- a/tests/integration/observability/test_otel_tenant_internal_spans_chaos.py +++ b/tests/integration/observability/test_otel_tenant_internal_spans_chaos.py @@ -124,6 +124,41 @@ def test_frozen_tenant_sink_receives_every_span_after_resume(gateway: Gateway, a assert _classes(group)["internal"] == 0, f"internal spans leaked for {call_id}" +def test_sink_outage_keeps_diagnostics_green_and_recovers(gateway: Gateway, audit_sinks: SpanSinks, langfuse_vars: dict[str, JsonValue]) -> None: + pid: Final = sink_pid(audit_sinks.tenant) + with gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/audit-chat", api_base=f"{gateway.upstream_url}/v1") + team_id: Final = scenario.team() + callback: Final = gateway.request( + "POST", f"/team/{team_id}/callback", {"callback_name": "langfuse_otel", "callback_vars": {**langfuse_vars, INTERNAL_SPANS_VAR: "exclude"}} + ) + assert callback.status_code == 200, callback.text + key: Final = scenario.key(team_id=team_id) + os.kill(pid, signal.SIGSTOP) + try: + responses: Final = _send_burst(gateway, key, model, 20) + readiness: Final = gateway.request("GET", "/health/readiness") + assert readiness.status_code == 200, readiness.text + details: Final = gateway.request("GET", "/health/readiness/details") + assert details.status_code == 200, details.text + success_callbacks: Final = details.json().get("success_callbacks", []) + assert "OpenTelemetryV2" in success_callbacks, f"otel callback missing during outage: {details.json()}" + liveliness: Final = gateway.request("GET", "/health/liveliness") + assert liveliness.status_code == 200, liveliness.text + finally: + os.kill(pid, signal.SIGCONT) + assert all(response.status_code == 200 for response in responses), [r.status_code for r in responses] + call_ids: Final = [response.headers.get("x-litellm-call-id") for response in responses] + spans: Final = _wait_trace_count(audit_sinks.tenant, call_ids, seconds=120) + for call_id in call_ids: + group: Final = tuple( + span for span in spans if span["attributes"].get("litellm.call_id") == call_id + ) + assert group, f"call {call_id} never reached the tenant sink" + model_spans: Final = tuple(span for span in group if "gen_ai.operation.name" in span["attributes"]) + assert len(model_spans) == 1, f"call {call_id} exported {len(model_spans)} times" + + def test_slow_tenant_sink_exports_each_span_once(gateway: Gateway, audit_sinks: SpanSinks, langfuse_vars: dict[str, JsonValue]) -> None: configure_sink(audit_sinks.tenant, delay_seconds=2.0) try: