From ce2ff17c926ce496d000e78a2539d368ac038d17 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 14 Oct 2026 03:18:50 +0000 Subject: [PATCH] feat(control_plane): ocsf audit events, sink protocol, queued delivery and hash chain Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/model/config.py | 12 +- .../integrations/otel/plumbing/providers.py | 7 +- litellm/litellm_core_utils/litellm_logging.py | 6 +- litellm/proxy/control_plane/app.py | 32 +- .../proxy/control_plane/audit/dependencies.py | 9 +- litellm/proxy/control_plane/audit/event.py | 47 +++ litellm/proxy/control_plane/audit/log.py | 210 ++++++++---- litellm/proxy/control_plane/audit/ocsf.py | 226 +++++++++++++ litellm/proxy/control_plane/audit/pipeline.py | 298 +++++++++++++++++ litellm/proxy/control_plane/audit/router.py | 29 +- litellm/proxy/control_plane/audit/sinks.py | 134 ++++++++ .../proxy/control_plane/auth/dependencies.py | 45 ++- litellm/proxy/control_plane/auth/router.py | 6 + litellm/proxy/control_plane/common/tracing.py | 21 +- litellm/proxy/control_plane/container.py | 43 ++- litellm/proxy/control_plane/proxy/router.py | 104 ++++-- litellm/proxy/control_plane/settings.py | 21 +- litellm/types/proxy/control_plane/audit.py | 9 + .../otel/test_otel_v2_components.py | 58 ++++ .../proxy/control_plane/audit/test_log.py | 82 +++++ .../proxy/control_plane/audit/test_ocsf.py | 100 ++++++ .../control_plane/audit/test_pipeline.py | 304 ++++++++++++++++++ .../proxy/control_plane/audit/test_router.py | 111 +++++++ .../proxy/control_plane/audit/test_sinks.py | 230 +++++++++++++ .../proxy/control_plane/proxy/test_router.py | 227 ++++++++++++- tests/unit/proxy/control_plane/test_app.py | 41 +++ .../unit/proxy/control_plane/test_settings.py | 35 ++ 27 files changed, 2334 insertions(+), 113 deletions(-) create mode 100644 litellm/proxy/control_plane/audit/event.py create mode 100644 litellm/proxy/control_plane/audit/ocsf.py create mode 100644 litellm/proxy/control_plane/audit/pipeline.py create mode 100644 litellm/proxy/control_plane/audit/sinks.py create mode 100644 tests/unit/proxy/control_plane/audit/test_ocsf.py create mode 100644 tests/unit/proxy/control_plane/audit/test_pipeline.py create mode 100644 tests/unit/proxy/control_plane/audit/test_router.py create mode 100644 tests/unit/proxy/control_plane/audit/test_sinks.py diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index 42dc22dfcbc..5976b86f7b1 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -84,6 +84,7 @@ class ExporterSpec(LiteLLMBaseModel): "``endpoint`` is a base URL the signal path is appended to." ), ) + logs_endpoint: str | None = None headers: str | None = None owner: ExporterOwner | None = Field( default=None, @@ -187,6 +188,10 @@ class OpenTelemetryV2Config(BaseSettings): "used verbatim instead of ``endpoint`` + ``/v1/traces``." ), ) + logs_endpoint: str | None = Field( + default=None, + validation_alias=AliasChoices("OTEL_LOGS_ENDPOINT", "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT"), + ) headers: str | None = Field( default=None, validation_alias=AliasChoices("OTEL_HEADERS", "OTEL_EXPORTER_OTLP_HEADERS"), @@ -345,7 +350,7 @@ class OpenTelemetryV2Config(BaseSettings): @model_validator(mode="after") def _normalize(self) -> "OpenTelemetryV2Config": # An endpoint with the default exporter kind implies OTLP/HTTP. - if (self.endpoint or self.traces_endpoint) and self.exporter == "console": + if (self.endpoint or self.traces_endpoint or self.logs_endpoint) and self.exporter == "console": self.exporter = "otlp_http" # When no explicit destinations are given, fold the single-destination # shorthand into one spec so the provider always has a destination. A spec @@ -357,9 +362,12 @@ class OpenTelemetryV2Config(BaseSettings): kind=self.exporter, endpoint=self.endpoint, traces_endpoint=self.traces_endpoint, + logs_endpoint=self.logs_endpoint, headers=self.headers, ) - if not self.model_fields_set.isdisjoint(("exporter", "endpoint", "traces_endpoint", "headers")) + if not self.model_fields_set.isdisjoint( + ("exporter", "endpoint", "traces_endpoint", "logs_endpoint", "headers") + ) else ExporterSpec() ] # Ensure ``genai`` is always present and first. diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index b5a5b85fc11..02a5a6aedbf 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -1379,7 +1379,7 @@ def build_log_exporter(config: OpenTelemetryV2Config) -> LogExporter: tls: Final = resolve_otlp_http_tls("LOGS") return HTTPLogExporter( - endpoint=_otlp_logs_endpoint(config.endpoint), + endpoint=config.logs_endpoint or _otlp_logs_endpoint(config.endpoint), headers=parse_headers(config.headers), certificate_file=tls.certificate_file, session=tls.session, @@ -1395,7 +1395,10 @@ def build_log_exporter(config: OpenTelemetryV2Config) -> LogExporter: "`opentelemetry-exporter-otlp` and `grpcio` (or `litellm[grpc]`)." ) from exc - return GRPCLogExporter(endpoint=config.endpoint, headers=parse_headers(config.headers)) + return GRPCLogExporter( + endpoint=config.logs_endpoint or config.endpoint, + headers=parse_headers(config.headers), + ) return ConsoleLogExporter() diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f8080531961..77b6c4e0c01 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6094,7 +6094,11 @@ 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, console_default=True + config=OpenTelemetryV2Config( + **otel_settings, # pyright: ignore[reportUnknownArgumentType] # Pydantic validates dynamic settings + ), + settings=otel_settings, + console_default=True, ) _in_memory_loggers.append(otel_logger_v2) _maybe_auto_initialize_arize_phoenix(_in_memory_loggers) diff --git a/litellm/proxy/control_plane/app.py b/litellm/proxy/control_plane/app.py index a4c5c833706..c278d15021e 100644 --- a/litellm/proxy/control_plane/app.py +++ b/litellm/proxy/control_plane/app.py @@ -1,8 +1,8 @@ """App factory for the LiteLLM global control plane. Enterprise only: refuses to build without a valid license.""" -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager from contextvars import ContextVar +from dataclasses import dataclass +from types import TracebackType from typing import Final, Protocol from fastapi import FastAPI @@ -81,23 +81,35 @@ class ControlPlaneRequestSizeLimitMiddleware: await self._middleware(scope, receive, send_problem) +@dataclass(frozen=True, slots=True) +class _Lifespan: + container: Container + + def __call__(self, app: FastAPI) -> "_Lifespan": + return self + + async def __aenter__(self) -> None: + self.container.audit_log.pipeline.start() + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None + ) -> None: + try: + await self.container.audit_log.pipeline.close() + finally: + self.container.telemetry.shutdown() + + def create_app(container: Container, license_check: License) -> FastAPI: if not license_check.is_premium(): raise LicenseRequired() - @asynccontextmanager - async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]: - try: - yield - finally: - container.telemetry.shutdown() - app: Final = FastAPI( title=TITLE, version="1", + lifespan=_Lifespan(container), docs_url=f"{container.settings.api_prefix}/docs", openapi_url=f"{container.settings.api_prefix}/openapi.json", - lifespan=lifespan, ) app.add_middleware(ControlPlaneRequestSizeLimitMiddleware) container.telemetry.instrument(app) diff --git a/litellm/proxy/control_plane/audit/dependencies.py b/litellm/proxy/control_plane/audit/dependencies.py index c260c4a65b3..f47cafed8f1 100644 --- a/litellm/proxy/control_plane/audit/dependencies.py +++ b/litellm/proxy/control_plane/audit/dependencies.py @@ -1,5 +1,10 @@ -from litellm.proxy.control_plane.audit.log import AuditLog +from typing import Final + +from litellm.proxy.control_plane.audit.log import AuditLog, build_audit_log + +_UNBOUND: Final = build_audit_log(capacity=1) def get_audit_log() -> AuditLog: - raise RuntimeError("get_audit_log is not bound, build the app with create_app") + """Overridden in create_app. A throwaway ring by default so a bare router still works in tests.""" + return _UNBOUND diff --git a/litellm/proxy/control_plane/audit/event.py b/litellm/proxy/control_plane/audit/event.py new file mode 100644 index 00000000000..67f1cf62690 --- /dev/null +++ b/litellm/proxy/control_plane/audit/event.py @@ -0,0 +1,47 @@ +from dataclasses import dataclass +from datetime import datetime +from functools import reduce +from typing import Final, Literal, TypeAlias +from urllib.parse import unquote + +Decision: TypeAlias = Literal["allowed", "denied", "unreachable"] +EventKind: TypeAlias = Literal["api_activity", "authentication"] +KEY_PREFIX: Final = "sk-" +REDACTED_KEY: Final = "sk-[redacted]" +REDACTION_DECODE_ROUNDS: Final = 8 +CONTROL_PLANE_ID: Final = "control-plane" +FINAL_KEY_ACTIONS: Final = frozenset({"reset_spend", "regenerate"}) + + +def _decode_segment(segment: str, remaining: int = REDACTION_DECODE_ROUNDS) -> str: + return reduce(lambda decoded, _: unquote(decoded), range(remaining), segment) + + +def redact_path(path: str) -> str: + """A virtual key in a path (``/key/sk-.../reset_spend``) must never reach the log or the audit API.""" + segments: Final = path.split("/") + key_index: Final = next( + (index for index, segment in enumerate(segments) if _decode_segment(segment).lower().startswith(KEY_PREFIX)), + None, + ) + if key_index is None: + return path + final_action: Final = segments[-1] if segments[-1] in FINAL_KEY_ACTIONS else None + return "/".join((*segments[:key_index], REDACTED_KEY, *((final_action,) if final_action is not None else ()))) + + +@dataclass(frozen=True, slots=True) +class AuditEvent: + at: datetime + actor: str + actor_email: str | None + dataplane_id: str + method: str + path: str + decision: Decision + status: int | None + credential: str | None + detail: str | None + kind: EventKind = "api_activity" + trace_id: str | None = None + span_id: str | None = None diff --git a/litellm/proxy/control_plane/audit/log.py b/litellm/proxy/control_plane/audit/log.py index 63605af048e..c0453fc9e86 100644 --- a/litellm/proxy/control_plane/audit/log.py +++ b/litellm/proxy/control_plane/audit/log.py @@ -1,63 +1,92 @@ -"""Who did what on which dataplane through the control plane. Bounded in memory plus a structured log line.""" +"""Who did what through the control plane. Every event becomes a chained OCSF record, kept in a bounded ring for the +audit API and handed to the configured sinks for delivery.""" import json +import logging +import sys from collections import deque from collections.abc import Callable -from dataclasses import asdict, dataclass from datetime import UTC, datetime -from functools import reduce -from typing import Final, Literal, TypeAlias -from urllib.parse import unquote +from typing import Final, TextIO, TypeAlias -from litellm._logging import verbose_proxy_logger +from opentelemetry.trace import format_span_id, format_trace_id, get_current_span + +from litellm._logging import redact_secrets +from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID, AuditEvent, Decision, EventKind, redact_path +from litellm.proxy.control_plane.audit.ocsf import AuditRecord, ChainState, Product, chain +from litellm.proxy.control_plane.audit.pipeline import NO_PIPELINE, Accepted, AuditPipeline, Offer, Reservation -Decision: TypeAlias = Literal["allowed", "denied", "unreachable"] DEFAULT_CAPACITY: Final = 10_000 -KEY_PREFIX: Final = "sk-" -REDACTED_KEY: Final = "sk-[redacted]" -REDACTION_DECODE_ROUNDS: Final = 8 -FINAL_KEY_ACTIONS: Final = frozenset({"reset_spend", "regenerate"}) +MAX_AUDIT_FIELD_CHARS: Final = 2048 +UNKNOWN_PRODUCT: Final = Product(version="unknown") +Trace: TypeAlias = tuple[str | None, str | None] -def _decode_segment(segment: str, remaining: int = REDACTION_DECODE_ROUNDS) -> str: - return reduce(lambda decoded, _: unquote(decoded), range(remaining), segment) +def build_audit_logger(stream: TextIO | None = None) -> logging.Logger: + logger: Final = logging.Logger("litellm.control_plane.audit", logging.INFO) + logger.propagate = False + handler: Final = logging.StreamHandler(sys.stderr if stream is None else stream) + handler.setLevel(logging.INFO) + logger.addHandler(handler) + return logger -def redact_path(path: str) -> str: - """A virtual key in a path (``/key/sk-.../reset_spend``) must never reach the log or the audit API.""" - segments: Final = path.split("/") - key_index: Final = next( - (index for index, segment in enumerate(segments) if _decode_segment(segment).lower().startswith(KEY_PREFIX)), - None, - ) - if key_index is None: - return path - final_action: Final = segments[-1] if segments[-1] in FINAL_KEY_ACTIONS else None - return "/".join((*segments[:key_index], REDACTED_KEY, *((final_action,) if final_action is not None else ()))) +AUDIT_LOGGER: Final = build_audit_logger() -@dataclass(frozen=True, slots=True) -class AuditEvent: - at: datetime - actor: str - actor_email: str | None - dataplane_id: str - method: str - path: str - decision: Decision - status: int | None - credential: str | None - detail: str | None +def _safe_text(value: str) -> str: + return value.encode("utf-8", "backslashreplace").decode("utf-8") + + +def _bounded_redacted(value: str) -> str: + return redact_secrets(value[:MAX_AUDIT_FIELD_CHARS])[:MAX_AUDIT_FIELD_CHARS] def utc_now() -> datetime: return datetime.now(UTC) -@dataclass(frozen=True, slots=True) +def current_trace() -> Trace: + context: Final = get_current_span().get_span_context() + if not context.is_valid: + return None, None + return format_trace_id(context.trace_id), format_span_id(context.span_id) + + class AuditLog: - events: deque[AuditEvent] # mutable-ok: a bounded ring buffer of events is the whole point - now: Callable[[], datetime] = utc_now + def __init__( + self, + records: deque[AuditRecord], # mutable-ok: a bounded ring buffer is the whole point + product: Product, + pipeline: AuditPipeline = NO_PIPELINE, + now: Callable[[], datetime] = utc_now, + trace: Callable[[], Trace] = current_trace, + unauthenticated_records: deque[AuditRecord] | None = None, # mutable-ok: bounded failure ring + logger: logging.Logger = AUDIT_LOGGER, + known_dataplane_ids: frozenset[str] = frozenset(), + ) -> None: + self._records = records + self._unauthenticated_records: deque[AuditRecord] = ( # mutable-ok: bounded failure ring + unauthenticated_records if unauthenticated_records is not None else deque[AuditRecord]() + ) + self._product = product + self._pipeline = pipeline + self._now = now + self._trace = trace + self._logger = logger + self._known_dataplane_ids = known_dataplane_ids + self._state = ChainState() + + @property + def pipeline(self) -> AuditPipeline: + return self._pipeline + + def accepting(self) -> bool: + return self._pipeline.accepting() + + def reserve(self) -> Reservation | None: + """Hold one queue slot per strict sink for a write that is about to go out, None when no slot is free.""" + return self._pipeline.reserve() def record( self, @@ -70,36 +99,99 @@ class AuditLog: status: int | None = None, credential: str | None = None, detail: str | None = None, + kind: EventKind = "api_activity", + reservation: Reservation | None = None, ) -> AuditEvent: + trace_id, span_id = self._trace() + safe_actor: Final = _bounded_redacted(_safe_text(actor)) + safe_email: Final = None if actor_email is None else _bounded_redacted(_safe_text(actor_email)) + safe_method: Final = _bounded_redacted(_safe_text(method)) + safe_path: Final = _bounded_redacted(redact_path(_safe_text(path)[:MAX_AUDIT_FIELD_CHARS])) + safe_detail: Final = None if detail is None else _bounded_redacted(_safe_text(detail)) + safe_dataplane_id: Final = _safe_text(dataplane_id) + recorded_dataplane_id: Final = ( + safe_dataplane_id[:MAX_AUDIT_FIELD_CHARS] + if dataplane_id in self._known_dataplane_ids + else _bounded_redacted(safe_dataplane_id) + ) event: Final = AuditEvent( - at=self.now(), - actor=actor, - actor_email=actor_email, - dataplane_id=dataplane_id, - method=method, - path=redact_path(path), + at=self._now(), + actor=safe_actor, + actor_email=safe_email, + dataplane_id=recorded_dataplane_id, + method=safe_method, + path=safe_path, decision=decision, status=status, credential=credential, - detail=detail, + detail=safe_detail, + kind=kind, + trace_id=trace_id, + span_id=span_id, ) - self.events.append(event) - verbose_proxy_logger.info("control_plane.audit %s", json.dumps(asdict(event), default=str)) + self._state, record = chain(event, self._product, self._state) + unauthenticated: Final = kind == "authentication" and decision == "denied" + (self._unauthenticated_records if unauthenticated else self._records).append(record) + if reservation is not None: + reservation.release() + if unauthenticated: + offers: Final = self._pipeline.publish_unauthenticated(record) + log_sink_accepted: Final = any( + worker.sink.name == "log" and isinstance(offer, Accepted) + for worker, offer in zip(self._pipeline.workers, offers, strict=True) + ) + if not log_sink_accepted: + self._logger.info("control_plane.audit %s", json.dumps(record.body, default=str, separators=(",", ":"))) + else: + self._pipeline.publish(record) return event + def record_login(self, username: str, decision: Decision, detail: str | None = None) -> AuditEvent: + return self.record( + username, None, CONTROL_PLANE_ID, "POST", "/auth/login", decision, detail=detail, kind="authentication" + ) + + def recent_records( + self, limit: int, dataplane_ids: frozenset[str] | None = None, actor: str | None = None + ) -> tuple[AuditRecord, ...]: + """``dataplane_ids`` of ``None`` means every dataplane, an empty set means nothing is visible.""" + ordered: Final = tuple( + sorted((*self._records, *self._unauthenticated_records), key=lambda record: record.sequence, reverse=True) + ) + matching: Final = ( + record + for record in ordered + if (dataplane_ids is None or record.event.dataplane_id in dataplane_ids) + and (actor is None or record.event.actor == actor) + ) + return tuple(record for _, record in zip(range(limit), matching, strict=False)) + def recent( self, limit: int, dataplane_ids: frozenset[str] | None = None, actor: str | None = None ) -> tuple[AuditEvent, ...]: - """``dataplane_ids`` of ``None`` means every dataplane, an empty set means nothing is visible.""" - newest_first: Final = reversed(self.events) - matching: Final = ( - event - for event in newest_first - if (dataplane_ids is None or event.dataplane_id in dataplane_ids) - and (actor is None or event.actor == actor) - ) - return tuple(event for _, event in zip(range(limit), matching, strict=False)) + return tuple(record.event for record in self.recent_records(limit, dataplane_ids, actor)) -def build_audit_log(capacity: int = DEFAULT_CAPACITY, now: Callable[[], datetime] = utc_now) -> AuditLog: - return AuditLog(events=deque(maxlen=capacity), now=now) +def build_audit_log( + capacity: int = DEFAULT_CAPACITY, + now: Callable[[], datetime] = utc_now, + pipeline: AuditPipeline = NO_PIPELINE, + product: Product = UNKNOWN_PRODUCT, + trace: Callable[[], Trace] = current_trace, + stream: TextIO | None = None, + known_dataplane_ids: frozenset[str] = frozenset(), +) -> AuditLog: + logger: Final = AUDIT_LOGGER if stream is None else build_audit_logger(stream) + return AuditLog( + deque[AuditRecord](maxlen=capacity), + product, + pipeline, + now, + trace, + deque[AuditRecord](maxlen=capacity), + logger, + known_dataplane_ids, + ) + + +__all__ = ["AuditEvent", "AuditLog", "Decision", "Offer", "Reservation", "build_audit_log", "redact_path", "utc_now"] diff --git a/litellm/proxy/control_plane/audit/ocsf.py b/litellm/proxy/control_plane/audit/ocsf.py new file mode 100644 index 00000000000..230dcf0b518 --- /dev/null +++ b/litellm/proxy/control_plane/audit/ocsf.py @@ -0,0 +1,226 @@ +"""Audit events as OCSF 1.5 objects: API Activity (6003) for calls through the control plane, Authentication (3002) +for logins. Each record carries a per-process sequence and a sha256 chain over the previous record so a dropped or +edited line is detectable in whatever sink holds it.""" + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass +from functools import reduce +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm.proxy.control_plane.audit.event import AuditEvent + +OCSF_VERSION: Final = "1.5.0" +API_ACTIVITY_CLASS: Final = 6003 +AUTHENTICATION_CLASS: Final = 3002 +APPLICATION_ACTIVITY_CATEGORY: Final = 6 +IAM_CATEGORY: Final = 3 +HASH_PREFIX: Final = "sha256:" +VENDOR: Final = "LiteLLM" +PRODUCT: Final = "LiteLLM Control Plane" +Json: TypeAlias = Mapping[str, object] +_JSON: Final = TypeAdapter(Mapping[str, object]) +_ACTIVITY_BY_METHOD: Final[Mapping[str, tuple[int, str]]] = { + "POST": (1, "Create"), + "GET": (2, "Read"), + "PUT": (3, "Update"), + "PATCH": (3, "Update"), + "DELETE": (4, "Delete"), +} +_LOGON: Final = (1, "Logon") +_UNKNOWN: Final = (0, "Unknown") + + +@dataclass(frozen=True, slots=True) +class ChainState: + sequence: int = 0 + prev_hash: str | None = None + + +@dataclass(frozen=True, slots=True) +class AuditRecord: + event: AuditEvent + sequence: int + hash: str + body: Json + + +@dataclass(frozen=True, slots=True) +class Product: + version: str + + +Outcome: TypeAlias = Literal["success", "failure"] + + +def _outcome(event: AuditEvent) -> Outcome: + if event.decision != "allowed": + return "failure" + return "failure" if event.status is not None and event.status >= 400 else "success" + + +def _status(event: AuditEvent) -> Json: + outcome: Final = _outcome(event) + return { + "status_id": 1 if outcome == "success" else 2, + "status": "Success" if outcome == "success" else "Failure", + **({} if event.status is None else {"status_code": str(event.status)}), + **({} if event.detail is None else {"status_detail": event.detail}), + } + + +def _severity(event: AuditEvent) -> Json: + if event.decision == "allowed": + return {"severity_id": 1, "severity": "Informational"} + return {"severity_id": 3, "severity": "Medium"} + + +def _actor(event: AuditEvent) -> Json: + user: Final = {"uid": event.actor, **({} if event.actor_email is None else {"email_addr": event.actor_email})} + return { + "user": user, + **({} if event.credential is None else {"session": {"credential_uid": event.credential}}), + "authorizations": [{"decision": event.decision}], + } + + +def _metadata(event: AuditEvent, product: Product, state: ChainState) -> Json: + return { + "version": OCSF_VERSION, + "product": {"name": PRODUCT, "vendor_name": VENDOR, "version": product.version}, + "sequence": state.sequence + 1, + **({} if event.trace_id is None else {"correlation_uid": event.trace_id}), + } + + +def _litellm(event: AuditEvent, state: ChainState) -> Json: + return { + "decision": event.decision, + "sequence": state.sequence + 1, + "prev_hash": state.prev_hash, + **({} if event.trace_id is None else {"trace_id": event.trace_id}), + **({} if event.span_id is None else {"span_id": event.span_id}), + } + + +def _api_activity(event: AuditEvent, product: Product, state: ChainState) -> Json: + activity_id, activity_name = _ACTIVITY_BY_METHOD.get(event.method, _UNKNOWN) + return { + "class_uid": API_ACTIVITY_CLASS, + "class_name": "API Activity", + "category_uid": APPLICATION_ACTIVITY_CATEGORY, + "category_name": "Application Activity", + "type_uid": API_ACTIVITY_CLASS * 100 + activity_id, + "activity_id": activity_id, + "activity_name": activity_name, + "time": int(event.at.timestamp() * 1000), + **_severity(event), + **_status(event), + "metadata": _metadata(event, product, state), + "actor": _actor(event), + "api": { + "operation": event.method, + "service": {"name": event.dataplane_id, "uid": "dataplane"}, + **({} if event.status is None else {"response": {"code": event.status}}), + }, + "http_request": {"http_method": event.method, "url": {"path": event.path}}, + **({} if event.status is None else {"http_response": {"code": event.status}}), + "unmapped": {"litellm": {"dataplane_id": event.dataplane_id, **_litellm(event, state)}}, + } + + +def _authentication(event: AuditEvent, product: Product, state: ChainState) -> Json: + activity_id, activity_name = _LOGON + return { + "class_uid": AUTHENTICATION_CLASS, + "class_name": "Authentication", + "category_uid": IAM_CATEGORY, + "category_name": "Identity & Access Management", + "type_uid": AUTHENTICATION_CLASS * 100 + activity_id, + "activity_id": activity_id, + "activity_name": activity_name, + "time": int(event.at.timestamp() * 1000), + **_severity(event), + **_status(event), + "metadata": _metadata(event, product, state), + "actor": _actor(event), + "user": {"uid": event.actor}, + "service": {"name": PRODUCT}, + "http_request": {"http_method": event.method, "url": {"path": event.path}}, + "unmapped": {"litellm": _litellm(event, state)}, + } + + +def to_ocsf(event: AuditEvent, product: Product, state: ChainState) -> Json: + match event.kind: + case "api_activity": + return _api_activity(event, product, state) + case "authentication": + return _authentication(event, product, state) + + +def canonical(body: Json) -> bytes: + return json.dumps(body, sort_keys=True, separators=(",", ":"), default=str).encode() + + +def chain(event: AuditEvent, product: Product, state: ChainState) -> tuple[ChainState, AuditRecord]: + """The next record and the state after it. The hash covers the whole body including the previous hash.""" + body: Final = to_ocsf(event, product, state) + digest: Final = HASH_PREFIX + hashlib.sha256(canonical(body)).hexdigest() + litellm_block: Final = dict(_unmapped_litellm(body)) + sealed: Final = {**body, "unmapped": {"litellm": {**litellm_block, "hash": digest}}} + record: Final = AuditRecord(event=event, sequence=state.sequence + 1, hash=digest, body=sealed) + return ChainState(sequence=state.sequence + 1, prev_hash=digest), record + + +def _child(body: Json, key: str) -> Json: + try: + return _JSON.validate_python(body.get(key)) + except ValidationError: + return {} + + +def section(body: Json, *keys: str) -> Json: + """The nested object at ``keys`` or an empty mapping when any step is missing or not an object.""" + return reduce(_child, keys, body) + + +def _unmapped_litellm(body: Json) -> Json: + return section(body, "unmapped", "litellm") + + +@dataclass(frozen=True, slots=True) +class ChainCheck: + """``valid`` is False only on tampering. ``gaps`` counts places where consecutive records are not consecutive + sequences, which is expected for a filtered or limited page and a dropped record for a complete export.""" + + valid: bool + gaps: int + + +def verify(records: tuple[AuditRecord, ...]) -> ChainCheck: + """Every record's hash must match its body, the first record of a process names no predecessor, and records + with adjacent sequences must link. Records separated by a gap are checked on their own only.""" + pairs: Final = tuple(zip(records, records[1:], strict=False)) + adjacent: Final = tuple(pair for pair in pairs if pair[1].sequence == pair[0].sequence + 1) + return ChainCheck( + valid=all(_sealed(record) for record in records) and all(_linked(prev, record) for prev, record in adjacent), + gaps=len(pairs) - len(adjacent), + ) + + +def _sealed(record: AuditRecord) -> bool: + litellm_block: Final = _unmapped_litellm(record.body) + unsealed: Final = { + **record.body, + "unmapped": {"litellm": {key: value for key, value in litellm_block.items() if key != "hash"}}, + } + hashed: Final = litellm_block.get("hash") == HASH_PREFIX + hashlib.sha256(canonical(unsealed)).hexdigest() + return hashed and (litellm_block.get("prev_hash") is None) == (record.sequence == 1) + + +def _linked(prev: AuditRecord, record: AuditRecord) -> bool: + return _unmapped_litellm(record.body).get("prev_hash") == prev.hash diff --git a/litellm/proxy/control_plane/audit/pipeline.py b/litellm/proxy/control_plane/audit/pipeline.py new file mode 100644 index 00000000000..9e2d1d91f83 --- /dev/null +++ b/litellm/proxy/control_plane/audit/pipeline.py @@ -0,0 +1,298 @@ +"""Asynchronous delivery of audit records to configured sinks. One bounded queue and one worker per sink, so a sink +that is down never blocks another one or the request path. Full queue: drop the oldest and count it, or in strict +mode refuse the record so the caller can fail the request before it does anything.""" + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Final, Literal, Protocol, TypeAlias + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.control_plane.audit.ocsf import AuditRecord + +Sleeper: TypeAlias = Callable[[float], Awaitable[None]] +Batch: TypeAlias = tuple[AuditRecord, ...] + + +@dataclass(frozen=True, slots=True) +class Delivered: + count: int + + +@dataclass(frozen=True, slots=True) +class Undelivered: + reason: str + settled: int = 0 + + +DeliveryResult: TypeAlias = Delivered | Undelivered + + +class AuditSink(Protocol): + @property + def name(self) -> str: ... + + async def deliver(self, batch: Batch) -> DeliveryResult: ... + + async def close(self) -> None: ... + + +@dataclass(frozen=True, slots=True) +class Accepted: + pass + + +@dataclass(frozen=True, slots=True) +class Dropped: + sink: str + + +@dataclass(frozen=True, slots=True) +class Refused: + sink: str + + +Offer: TypeAlias = Accepted | Dropped | Refused +Mode: TypeAlias = Literal["drop_oldest", "strict"] +EXPONENT_CAP: Final = 30 + + +@dataclass(frozen=True, slots=True) +class Backoff: + base_seconds: float = 0.5 + max_seconds: float = 30.0 + + def delay(self, failures: int) -> float: + return min(self.max_seconds, self.base_seconds * (2.0 ** min(failures, EXPONENT_CAP))) + + +@dataclass(slots=True) +class SinkWorker: + sink: AuditSink + queue: asyncio.Queue[AuditRecord] + batch_size: int + mode: Mode + backoff: Backoff = field(default_factory=Backoff) + sleep: Sleeper = asyncio.sleep + dropped: int = 0 + failures: int = 0 + task: asyncio.Task[None] | None = None + inflight: Batch | None = None + stopping: bool = False + reservations: frozenset[object] = frozenset() + + def offer(self, record: AuditRecord) -> Offer: + if self.mode == "strict" and self.full(): + return Refused(self.sink.name) + if not self.queue.full(): + self.queue.put_nowait(record) + return Accepted() + self.queue.get_nowait() + self.queue.put_nowait(record) + self.dropped += 1 + verbose_proxy_logger.warning( + "control_plane.audit sink %s queue full, dropped oldest (total %d)", self.sink.name, self.dropped + ) + return Dropped(self.sink.name) + + def offer_unauthenticated(self, record: AuditRecord) -> Offer: + if ( + self.queue.qsize() * 2 >= self.queue.maxsize + or self.queue.qsize() + len(self.reservations) + 1 >= self.queue.maxsize + ): + self.dropped += 1 + verbose_proxy_logger.warning( + "control_plane.audit sink %s dropped unauthenticated failure (total %d)", + self.sink.name, + self.dropped, + ) + return Dropped(self.sink.name) + self.queue.put_nowait(record) + return Accepted() + + def full(self) -> bool: + """Queued plus reserved slots. Reservations only matter in strict mode, where a refusal is the contract.""" + return self.queue.qsize() + len(self.reservations) >= self.queue.maxsize + + def reserve(self, token: object) -> None: + self.reservations = self.reservations | {token} + + def release(self, token: object) -> None: + self.reservations = self.reservations - {token} + + def _take_batch(self, first: AuditRecord) -> Batch: + rest: Final = tuple(self.queue.get_nowait() for _ in range(min(self.batch_size - 1, self.queue.qsize()))) + return (first, *rest) + + async def _deliver_until_done(self, batch: Batch) -> None: + """Keeps the batch as ``inflight`` until the sink takes it, so a cancelled delivery can be redone.""" + self.inflight = batch + current_batch: Batch = batch # rebind-ok: retries advance through immutable unsettled suffixes + while isinstance(result := await self.sink.deliver(current_batch), Undelivered): + self.failures += 1 + verbose_proxy_logger.warning( + "control_plane.audit sink %s failed (%s), retry %d", self.sink.name, result.reason, self.failures + ) + current_batch = current_batch[result.settled :] + self.inflight = current_batch + if not current_batch: + break + await self.sleep(self.backoff.delay(self.failures - 1)) + self.failures = 0 + self.inflight = None + + async def flush_once(self) -> int: + """Deliver what is queued right now, returns how many records went out. Tests drive this directly.""" + if self.queue.empty(): + return 0 + batch: Final = self._take_batch(self.queue.get_nowait()) + await self._deliver_until_done(batch) + return len(batch) + + async def _next_batch(self) -> None: + first: Final = await self.queue.get() + await self._deliver_until_done(self._take_batch(first)) + + async def run(self) -> None: + while not self.stopping: + await self._next_batch() + + def start(self) -> None: + if self.task is None or self.task.done(): + self.task = asyncio.create_task(self.run(), name=f"control_plane.audit.{self.sink.name}") + + async def close(self, drain_seconds: float) -> None: + """Stops the worker, then drains its queue within the deadline.""" + loop: Final = asyncio.get_running_loop() + deadline: Final = loop.time() + drain_seconds + task: Final = self.task + if task is not None: + self.stopping = True + delivery_finished: Final = ( + task is None + or self.inflight is None + or ( + task in (await asyncio.wait({task}, timeout=max(0.0, deadline - loop.time())))[0] + and not task.cancelled() + ) + ) + if task is not None: + if self.inflight is None or not delivery_finished: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + except Exception: # noqa: BLE001 # shutdown must continue after arbitrary sink worker failures + verbose_proxy_logger.exception("control_plane.audit sink %s worker crashed", self.sink.name) + if delivery_finished: + try: + await asyncio.wait_for(self._drain(), timeout=max(0.0, deadline - loop.time())) + except TimeoutError: + self._warn_undelivered() + except Exception: # noqa: BLE001 # continue closing sinks after arbitrary drain failures + verbose_proxy_logger.exception( + "control_plane.audit sink %s failed to drain at shutdown", self.sink.name + ) + else: + self._warn_undelivered() + try: + await self.sink.close() + except Exception: # noqa: BLE001 # close other sinks despite arbitrary sink close failures + verbose_proxy_logger.exception("control_plane.audit sink %s failed to close", self.sink.name) + + def _warn_undelivered(self) -> None: + verbose_proxy_logger.warning( + "control_plane.audit sink %s left %d records undelivered at shutdown", + self.sink.name, + self.queue.qsize() + len(self.inflight or ()), + ) + + async def _drain(self) -> None: + if self.inflight is not None: + await self._deliver_until_done(self.inflight) + while await self.flush_once(): + pass + + +@dataclass(frozen=True, slots=True) +class Reservation: + """One held slot on every strict worker. Releasing twice is harmless, so the request can release in a + ``finally`` after the audit record consumed it.""" + + token: object + workers: tuple[SinkWorker, ...] + + def release(self) -> None: + for worker in self.workers: + worker.release(self.token) + + +@dataclass(frozen=True, slots=True) +class AuditPipeline: + workers: tuple[SinkWorker, ...] + drain_seconds: float = 5.0 + + def publish(self, record: AuditRecord) -> tuple[Offer, ...]: + return tuple(worker.offer(record) for worker in self.workers) + + def publish_unauthenticated(self, record: AuditRecord) -> tuple[Offer, ...]: + return tuple(worker.offer_unauthenticated(record) for worker in self.workers) + + def accepting(self) -> bool: + """False only in strict mode with a full queue, the signal to answer 503 before a write goes out.""" + return not any(worker.mode == "strict" and worker.full() for worker in self.workers) + + def reserve(self) -> Reservation | None: + """Atomically (no await) takes one slot on every strict worker so a write that passes the gate always has + room for its record, however many writes are in flight. None when any strict queue is full.""" + if not self.accepting(): + return None + token: Final = object() + strict: Final = tuple(worker for worker in self.workers if worker.mode == "strict") + for worker in strict: + worker.reserve(token) + return Reservation(token, strict) + + def start(self) -> None: + for worker in self.workers: + worker.start() + + async def close(self) -> None: + results: Final = tuple( + await asyncio.gather( + *(worker.close(self.drain_seconds) for worker in self.workers), + return_exceptions=True, + ) + ) + for worker, result in zip(self.workers, results, strict=True): + if isinstance(result, BaseException): + verbose_proxy_logger.error( + "control_plane.audit sink %s failed during shutdown: %s", worker.sink.name, result + ) + + +def build_pipeline( + sinks: tuple[AuditSink, ...], + queue_size: int, + batch_size: int, + strict: bool, + sleep: Sleeper = asyncio.sleep, + drain_seconds: float = 5.0, +) -> AuditPipeline: + return AuditPipeline( + workers=tuple( + SinkWorker( + sink=sink, + queue=asyncio.Queue(maxsize=queue_size), + batch_size=batch_size, + mode="strict" if strict else "drop_oldest", + sleep=sleep, + ) + for sink in sinks + ), + drain_seconds=drain_seconds, + ) + + +NO_PIPELINE: Final = AuditPipeline(workers=()) diff --git a/litellm/proxy/control_plane/audit/router.py b/litellm/proxy/control_plane/audit/router.py index 06b2fb0bc8c..f7c85a929ce 100644 --- a/litellm/proxy/control_plane/audit/router.py +++ b/litellm/proxy/control_plane/audit/router.py @@ -3,22 +3,45 @@ from typing import Annotated, Final from fastapi import APIRouter, Depends, Query, Security from litellm.proxy.control_plane.audit.dependencies import get_audit_log +from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID from litellm.proxy.control_plane.audit.log import AuditLog +from litellm.proxy.control_plane.audit.ocsf import verify from litellm.proxy.control_plane.auth.dependencies import require_scopes from litellm.proxy.control_plane.auth.rbac import DataplaneBinding, Principal, Scope from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden -from litellm.types.proxy.control_plane.audit import AuditEventResponse, AuditListResponse +from litellm.types.proxy.control_plane.audit import AuditEventResponse, AuditListResponse, OcsfAuditListResponse router: Final = APIRouter(prefix="/audit", tags=["audit"]) def visible_dataplanes(binding: DataplaneBinding, requested: str | None) -> frozenset[str] | None: - """What this caller may read: everything when fleet-wide, else only bound ids, and a requested id must be bound.""" + """What this caller may read: everything when fleet-wide, else only bound ids, and a requested id must be bound. + Events against the control plane itself (logins, rejected tokens, denied scopes) are fleet-wide only.""" + if requested == CONTROL_PLANE_ID and not binding.all: + raise ControlPlaneException(forbidden("Control plane events are visible to fleet-wide auditors only.")) if requested is not None and not (binding.all or requested in binding.ids): raise ControlPlaneException(forbidden(f"Dataplane {requested!r} is not bound to this caller.")) if requested is not None: return frozenset({requested}) - return None if binding.all else binding.ids + return None if binding.all else binding.ids - {CONTROL_PLANE_ID} + + +@router.get("/ocsf", response_model=OcsfAuditListResponse) +async def list_audit_ocsf( + principal: Annotated[Principal, Security(require_scopes, scopes=[Scope.AUDIT_READ.value])], + audit: Annotated[AuditLog, Depends(get_audit_log)], + limit: Annotated[int, Query(ge=1, le=1000)] = 100, + dataplane: Annotated[str | None, Query(description="Only events against this dataplane id.")] = None, + actor: Annotated[str | None, Query(description="Only events by this control plane subject.")] = None, +) -> OcsfAuditListResponse: + """The same events as OCSF 1.5 objects, newest first. ``chain_valid`` is false only when a returned record was + altered, ``sequence_gaps`` counts pairs of returned records whose sequence numbers are not consecutive.""" + visible: Final = visible_dataplanes(principal.dataplanes, dataplane) + records: Final = audit.recent_records(limit, visible, actor) + check: Final = verify(tuple(reversed(records))) + return OcsfAuditListResponse( + items=tuple(record.body for record in records), chain_valid=check.valid, sequence_gaps=check.gaps + ) @router.get("", response_model=AuditListResponse) diff --git a/litellm/proxy/control_plane/audit/sinks.py b/litellm/proxy/control_plane/audit/sinks.py new file mode 100644 index 00000000000..5168b420984 --- /dev/null +++ b/litellm/proxy/control_plane/audit/sinks.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import time +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Final, TextIO + +from opentelemetry._logs import LogRecord, SeverityNumber +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState, set_span_in_context + +from litellm.proxy.control_plane.audit.log import AUDIT_LOGGER, build_audit_logger +from litellm.proxy.control_plane.audit.ocsf import AuditRecord +from litellm.proxy.control_plane.audit.pipeline import Batch, Delivered, DeliveryResult, Undelivered + +if TYPE_CHECKING: + from opentelemetry.sdk._logs import ReadableLogRecord + from opentelemetry.sdk._logs.export import LogRecordExporter + from opentelemetry.sdk.resources import Resource + +EVENT_NAME: Final = "litellm.control_plane.audit" +LOG_PREFIX: Final = "control_plane.audit" + + +def _line(record: AuditRecord) -> str: + return json.dumps(record.body, default=str, separators=(",", ":")) + + +@dataclass(frozen=True, slots=True) +class LogSink: + name: str = "log" + stream: TextIO | None = None + _logger: logging.Logger = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + object.__setattr__(self, "_logger", AUDIT_LOGGER if self.stream is None else build_audit_logger(self.stream)) + + async def deliver(self, batch: Batch) -> DeliveryResult: + for record in batch: + self._logger.info("%s %s", LOG_PREFIX, _line(record)) + return Delivered(len(batch)) + + async def close(self) -> None: + return None + + +def _severity(record: AuditRecord) -> tuple[str, SeverityNumber]: + return ("INFO", SeverityNumber.INFO) if record.event.decision == "allowed" else ("WARN", SeverityNumber.WARN) + + +def _ids(record: AuditRecord) -> tuple[int, int]: + trace_id: Final = int(record.event.trace_id, 16) if record.event.trace_id else 0 + span_id: Final = int(record.event.span_id, 16) if record.event.span_id else 0 + return trace_id, span_id + + +def _attributes(record: AuditRecord) -> Mapping[str, str | int]: + return { + "event.name": EVENT_NAME, + "enduser.id": record.event.actor, + "http.request.method": record.event.method, + "url.path": record.event.path, + "litellm.control_plane.decision": record.event.decision, + "litellm.control_plane.dataplane": record.event.dataplane_id, + "litellm.control_plane.audit.sequence": record.sequence, + "litellm.control_plane.audit.hash": record.hash, + } + + +def _log_data(record: AuditRecord, resource: Resource) -> ReadableLogRecord: + from opentelemetry.sdk._logs import ReadableLogRecord + from opentelemetry.sdk.util.instrumentation import InstrumentationScope + + trace_id, span_id = _ids(record) + severity_text, severity_number = _severity(record) + trace_flags: Final = TraceFlags(TraceFlags.SAMPLED if trace_id else TraceFlags.DEFAULT) + return ReadableLogRecord( + LogRecord( + timestamp=int(record.event.at.timestamp() * 1_000_000_000), + observed_timestamp=time.time_ns(), + context=set_span_in_context( + NonRecordingSpan( + SpanContext( + trace_id=trace_id, + span_id=span_id, + is_remote=False, + trace_flags=trace_flags, + trace_state=TraceState(), + ) + ) + ), + severity_text=severity_text, + severity_number=severity_number, + body=_line(record), + attributes=dict(_attributes(record)), + ), + resource, + InstrumentationScope(EVENT_NAME), + ) + + +@dataclass(frozen=True, slots=True) +class OtlpLogSink: + """One OTel log record per audit event, body is the OCSF object, trace and span ids point at the request span. + The batch goes straight to the exporter, off the event loop, so its verdict is the delivery result and a + collector that refuses the batch makes the worker retry instead of the SDK dropping it later.""" + + exporter: LogRecordExporter + resource: Resource + name: str = "otlp" + + async def deliver(self, batch: Batch) -> DeliveryResult: + from opentelemetry.sdk._logs.export import LogRecordExportResult + from pydantic import TypeAdapter + + try: + data: Final = tuple(_log_data(record, self.resource) for record in batch) + result: Final = await asyncio.to_thread(self.exporter.export, data) + except Exception as error: # noqa: BLE001 # an exporter that raises instead of returning FAILURE is still a failed delivery + return Undelivered(f"otlp export raised {type(error).__name__}") + try: + verdict: Final = TypeAdapter(LogRecordExportResult).validate_python(result) + except Exception as error: # noqa: BLE001 # invalid exporter verdicts must be retried, not kill the worker + return Undelivered(f"otlp export verdict raised {type(error).__name__}") + match verdict: + case LogRecordExportResult.SUCCESS: + return Delivered(len(batch)) + case LogRecordExportResult.FAILURE: + return Undelivered("otlp export failed") + + async def close(self) -> None: + await asyncio.to_thread(self.exporter.shutdown) diff --git a/litellm/proxy/control_plane/auth/dependencies.py b/litellm/proxy/control_plane/auth/dependencies.py index 499da3e86fe..4ebf9db040c 100644 --- a/litellm/proxy/control_plane/auth/dependencies.py +++ b/litellm/proxy/control_plane/auth/dependencies.py @@ -2,11 +2,14 @@ from typing import Annotated, Final -from fastapi import Depends +from fastapi import Depends, Request from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer, SecurityScopes from starlette.concurrency import run_in_threadpool from litellm._logging import verbose_proxy_logger +from litellm.proxy.control_plane.audit.dependencies import get_audit_log +from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID +from litellm.proxy.control_plane.audit.log import AuditLog from litellm.proxy.control_plane.auth.jwt import Rejected, TokenVerifier from litellm.proxy.control_plane.auth.rbac import Principal, Scope from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden, unauthorized @@ -24,27 +27,67 @@ def get_telemetry() -> Telemetry: return UNTRACED +ANONYMOUS: Final = "anonymous" + + async def get_principal( + request: Request, credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)], verifier: Annotated[TokenVerifier, Depends(get_token_verifier)], telemetry: Annotated[Telemetry, Depends(get_telemetry)], + audit: Annotated[AuditLog, Depends(get_audit_log)], ) -> Principal: if credentials is None: + audit.record( + ANONYMOUS, + None, + CONTROL_PLANE_ID, + request.method, + request.url.path, + "denied", + 401, + detail="missing bearer", + kind="authentication", + ) raise ControlPlaneException(unauthorized("Missing bearer credentials.")) verified: Final = await run_in_threadpool(verifier.verify, credentials.credentials) if isinstance(verified, Rejected): verbose_proxy_logger.debug("Control-plane bearer verification rejected: %s", verified.reason) + audit.record( + ANONYMOUS, + None, + CONTROL_PLANE_ID, + request.method, + request.url.path, + "denied", + 401, + detail=verified.reason, + kind="authentication", + ) raise ControlPlaneException(unauthorized("Invalid bearer credentials.")) annotate_principal(verified, telemetry.identity_attributes) return verified async def require_scopes( + request: Request, security_scopes: SecurityScopes, principal: Annotated[Principal, Depends(get_principal)], + audit: Annotated[AuditLog, Depends(get_audit_log)], ) -> Principal: required: Final = frozenset(Scope(scope) for scope in security_scopes.scopes) missing: Final = principal.missing_scopes(required) if missing: + detail: Final = f"missing scope: {', '.join(sorted(scope.value for scope in missing))}" + audit.record( + principal.subject, + principal.email, + CONTROL_PLANE_ID, + request.method, + request.url.path, + "denied", + 403, + detail=detail, + ) raise ControlPlaneException(forbidden(f"Missing scope: {', '.join(sorted(scope.value for scope in missing))}.")) return principal diff --git a/litellm/proxy/control_plane/auth/router.py b/litellm/proxy/control_plane/auth/router.py index 90f541dd641..4dcaeb0d7c2 100644 --- a/litellm/proxy/control_plane/auth/router.py +++ b/litellm/proxy/control_plane/auth/router.py @@ -3,6 +3,8 @@ from typing import Annotated, Final from fastapi import APIRouter, Depends, Request, Security from fastapi.responses import JSONResponse +from litellm.proxy.control_plane.audit.dependencies import get_audit_log +from litellm.proxy.control_plane.audit.log import AuditLog from litellm.proxy.control_plane.auth.dependencies import require_scopes from litellm.proxy.control_plane.auth.login import InvalidCredentials, LoginService, Throttled from litellm.proxy.control_plane.auth.rbac import Principal @@ -32,20 +34,24 @@ async def login( body: LoginRequest, service: Annotated[LoginService, Depends(get_login_service)], request: Request, + audit: Annotated[AuditLog, Depends(get_audit_log)], ) -> JSONResponse: client_source: Final = request.client.host if request.client is not None else "" match service.login(body.username, body.password, client_source): case SessionTokenResponse() as token: + audit.record_login(body.username, "allowed") return JSONResponse( content=token.model_dump(mode="json"), headers={"Cache-Control": "no-store", "Pragma": "no-cache"}, ) case Throttled(retry_after_seconds=seconds): + audit.record_login(body.username, "denied", detail="throttled") raise ControlPlaneException( too_many_requests(f"Too many failed logins, retry in {seconds} seconds."), headers={"Retry-After": str(seconds)}, ) case InvalidCredentials(): + audit.record_login(body.username, "denied", detail="invalid credentials") raise ControlPlaneException(unauthorized("Invalid username or password.")) diff --git a/litellm/proxy/control_plane/common/tracing.py b/litellm/proxy/control_plane/common/tracing.py index f47879fcb7d..077bcfb3e33 100644 --- a/litellm/proxy/control_plane/common/tracing.py +++ b/litellm/proxy/control_plane/common/tracing.py @@ -24,6 +24,8 @@ from litellm.proxy.control_plane.settings import ControlPlaneSettings if TYPE_CHECKING: from fastapi import FastAPI + from opentelemetry.sdk._logs.export import LogRecordExporter + from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SpanExporter @@ -275,14 +277,17 @@ def tracing_enabled(settings: ControlPlaneSettings, environ: Mapping[str, str] = ) +def _config(settings: ControlPlaneSettings) -> OpenTelemetryV2Config: + return OpenTelemetryV2Config.from_env().model_copy(update={"service_name": settings.otel_service_name}) + + def build_provider(settings: ControlPlaneSettings, exporter: SpanExporter | None = None) -> TracerProvider: from opentelemetry.sdk.trace import SpanProcessor from litellm.integrations.otel.plumbing.providers import build_tracer_provider - config: Final = OpenTelemetryV2Config.from_env().model_copy(update={"service_name": settings.otel_service_name}) return build_tracer_provider( - config, + _config(settings), exporter=exporter, baggage_processor=SpanProcessor(), use_simple_processor=None if exporter is None else True, @@ -290,6 +295,18 @@ def build_provider(settings: ControlPlaneSettings, exporter: SpanExporter | None ) +def build_log_exporter(settings: ControlPlaneSettings) -> LogRecordExporter: + from litellm.integrations.otel.plumbing.providers import build_log_exporter as build_provider_log_exporter + + return build_provider_log_exporter(_config(settings)) + + +def build_resource(settings: ControlPlaneSettings) -> Resource: + from litellm.integrations.otel.plumbing.providers import build_resource as build_provider_resource + + return build_provider_resource(_config(settings)) + + def build_telemetry( settings: ControlPlaneSettings, provider: TracerProvider | None = None, diff --git a/litellm/proxy/control_plane/container.py b/litellm/proxy/control_plane/container.py index 7bf246e4d4c..2fed38b5b21 100644 --- a/litellm/proxy/control_plane/container.py +++ b/litellm/proxy/control_plane/container.py @@ -7,6 +7,7 @@ from collections.abc import Mapping, Sequence from contextlib import AbstractAsyncContextManager from dataclasses import dataclass from datetime import timedelta +from importlib.metadata import PackageNotFoundError, version from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -17,6 +18,9 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # params is an untyped dict upstream ) from litellm.proxy.control_plane.audit.log import AuditLog, build_audit_log +from litellm.proxy.control_plane.audit.ocsf import Product +from litellm.proxy.control_plane.audit.pipeline import AuditPipeline, AuditSink, build_pipeline +from litellm.proxy.control_plane.audit.sinks import LogSink, OtlpLogSink from litellm.proxy.control_plane.auth.jwt import JwksKeyLookup, SessionTokenCodec, TokenVerifier, build_token_verifier from litellm.proxy.control_plane.auth.login import LoginService, LoginThrottle from litellm.proxy.control_plane.auth.rbac import ( @@ -29,7 +33,7 @@ from litellm.proxy.control_plane.auth.rbac import ( parse_role_definition, ) from litellm.proxy.control_plane.common.fanout import Fanout -from litellm.proxy.control_plane.common.tracing import Telemetry, build_telemetry +from litellm.proxy.control_plane.common.tracing import Telemetry, build_log_exporter, build_resource, build_telemetry from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, HttpTransport, RetryPolicy from litellm.proxy.control_plane.dataplanes.credentials import TokenExchanger, http_token_exchanger from litellm.proxy.control_plane.dataplanes.registry import DataplaneRegistry, build_registry @@ -115,12 +119,39 @@ def _shared_http_client(settings: ControlPlaneSettings) -> HttpTransport: return HandlerTransport(_shared_http_handler(settings)) +def _sink(kind: str, settings: ControlPlaneSettings) -> AuditSink: + match kind: + case "otlp": + return OtlpLogSink(build_log_exporter(settings), build_resource(settings)) + case _: + return LogSink() + + +def build_audit_pipeline(settings: ControlPlaneSettings, sinks: tuple[AuditSink, ...] | None = None) -> AuditPipeline: + chosen: Final = sinks if sinks is not None else tuple(_sink(kind, settings) for kind in settings.audit_sinks) + return build_pipeline( + chosen, + settings.audit_queue_size, + settings.audit_batch_size, + settings.audit_strict, + drain_seconds=settings.audit_drain_seconds, + ) + + +def litellm_product() -> Product: + try: + return Product(version=version("litellm")) + except PackageNotFoundError: + return Product(version="unknown") + + def build_container( settings: ControlPlaneSettings, jwks: JwksKeyLookup | None = None, transport: HttpTransport | None = None, exchanger: TokenExchanger | None = None, tracer_provider: TracerProvider | None = None, + audit_sinks: tuple[AuditSink, ...] | None = None, ) -> Container: master_key: Final = settings.master_key.get_secret_value() policy: Final = build_role_policy(settings) @@ -139,6 +170,7 @@ def build_container( ) password: Final = settings.ui_password telemetry: Final = build_telemetry(settings, tracer_provider) + registry: Final = build_registry(settings.dataplanes, exchange) return Container( settings=settings, policy=policy, @@ -151,7 +183,7 @@ def build_container( admin_password=None if password is None else password.get_secret_value(), throttle=LoginThrottle(time.monotonic), ), - registry=build_registry(settings.dataplanes, exchange), + registry=registry, dataplane_client=DataplaneClient( transport=http, timeout_seconds=settings.dataplane_timeout_seconds, @@ -159,6 +191,11 @@ def build_container( tracing=telemetry.recorder, ), fanout=Fanout(concurrency=settings.fanout_concurrency), - audit_log=build_audit_log(settings.audit_capacity), + audit_log=build_audit_log( + settings.audit_capacity, + pipeline=build_audit_pipeline(settings, audit_sinks), + product=litellm_product(), + known_dataplane_ids=frozenset(dataplane.id for dataplane in registry.all()), + ), telemetry=telemetry, ) diff --git a/litellm/proxy/control_plane/proxy/router.py b/litellm/proxy/control_plane/proxy/router.py index d4dcb40b252..6994fa9e536 100644 --- a/litellm/proxy/control_plane/proxy/router.py +++ b/litellm/proxy/control_plane/proxy/router.py @@ -9,8 +9,9 @@ from fastapi.security import SecurityScopes from pydantic import TypeAdapter, ValidationError from litellm.proxy.control_plane.audit.dependencies import get_audit_log -from litellm.proxy.control_plane.audit.log import AuditLog, redact_path -from litellm.proxy.control_plane.auth.dependencies import get_principal, require_scopes +from litellm.proxy.control_plane.audit.event import redact_path +from litellm.proxy.control_plane.audit.log import AuditLog, Reservation +from litellm.proxy.control_plane.auth.dependencies import get_principal from litellm.proxy.control_plane.auth.rbac import Principal, Scope from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, bad_gateway, error_body, forbidden from litellm.proxy.control_plane.dataplanes.client import ( @@ -28,8 +29,9 @@ from litellm.proxy.control_plane.dataplanes.dependencies import ( get_registry, ) from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry -from litellm.proxy.control_plane.proxy.allowlist import Denied, decide +from litellm.proxy.control_plane.proxy.allowlist import Allowed, Denied, decide +AUDIT_UNAVAILABLE: Final = "audit-unavailable" DATAPLANE_HEADER: Final = "X-LiteLLM-Dataplane" READ_METHODS: Final = ("GET", "HEAD") WRITE_METHODS: Final = ("POST", "PUT", "PATCH", "DELETE") @@ -40,12 +42,17 @@ _BODY: Final = TypeAdapter(Mapping[str, object]) _METHOD: Final[TypeAdapter[Method]] = TypeAdapter(Method) +def audit_unavailable() -> ControlPlaneException: + return ControlPlaneException( + error_body(AUDIT_UNAVAILABLE, "Audit Unavailable", 503, "No audit sink can record this write right now.") + ) + + def bad_request(detail: str) -> ControlPlaneException: return ControlPlaneException(error_body("bad-request", "Bad Request", 400, detail)) -async def _json_body(request: Request) -> Mapping[str, object] | None: - raw: Final = await request.body() +def _json_body(raw: bytes) -> Mapping[str, object] | None: if not raw: return None try: @@ -67,16 +74,47 @@ async def _forward( decision: Final = decide(path) if isinstance(decision, Denied): audit.record( - principal.subject, principal.email, dataplane.id, method, decision.path, "denied", detail=decision.reason + principal.subject, + principal.email, + dataplane.id, + method, + decision.path, + "denied", + detail=decision.reason, ) raise ControlPlaneException(forbidden(f"{redact_path(decision.path)}: {decision.reason}.")) + raw: Final = await request.body() + reservation: Final = audit.reserve() if method in WRITE_METHODS else None + if method in WRITE_METHODS and reservation is None: + raise audit_unavailable() + try: + return await _forward_reserved( + request, raw, decision, method, principal, dataplane, caller, client, audit, reservation + ) + finally: + if reservation is not None: + reservation.release() + + +async def _forward_reserved( + request: Request, + raw: bytes, + decision: Allowed, + method: Method, + principal: Principal, + dataplane: Dataplane, + caller: Caller, + client: DataplaneClient, + audit: AuditLog, + reservation: Reservation | None, +) -> Response: result: Final = await client.request( method, dataplane, decision.path, caller, params=tuple(request.query_params.multi_items()), - json=await _body_or_audit(request, principal, dataplane, method, decision.path, audit), + json=_body_or_audit(raw, principal, dataplane, method, decision.path, audit, reservation), ) match result: case DataplaneReply(status=status, credential=credential) if 300 <= status < 400: @@ -90,6 +128,7 @@ async def _forward( status, credential, detail="dataplane answered with a redirect, which the control plane never follows", + reservation=reservation, ) raise ControlPlaneException(bad_gateway("Dataplane answered with a redirect, which is not followed.")) case DataplaneReply( @@ -102,7 +141,15 @@ async def _forward( cache_control=cache_control, ): audit.record( - principal.subject, principal.email, dataplane.id, method, decision.path, "allowed", status, credential + principal.subject, + principal.email, + dataplane.id, + method, + decision.path, + "allowed", + status, + credential, + reservation=reservation, ) if raw_content is not None: headers: Final = { @@ -131,6 +178,7 @@ async def _forward( "unreachable", 502, detail="dataplane response exceeded the configured byte limit", + reservation=reservation, ) raise ControlPlaneException( error_body( @@ -142,26 +190,41 @@ async def _forward( ) case DataplaneUnreachable(detail=detail, attempts=attempts): audit.record( - principal.subject, principal.email, dataplane.id, method, decision.path, "unreachable", detail=detail + principal.subject, + principal.email, + dataplane.id, + method, + decision.path, + "unreachable", + detail=detail, + reservation=reservation, ) raise ControlPlaneException(bad_gateway(f"Dataplane unreachable after {attempts} attempt(s): {detail}")) case CredentialUnavailable(reason=reason): audit.record( - principal.subject, principal.email, dataplane.id, method, decision.path, "unreachable", detail=reason + principal.subject, + principal.email, + dataplane.id, + method, + decision.path, + "unreachable", + detail=reason, + reservation=reservation, ) raise ControlPlaneException(bad_gateway(f"No usable credential for this dataplane: {reason}")) -async def _body_or_audit( - request: Request, +def _body_or_audit( + raw: bytes, principal: Principal, dataplane: Dataplane, method: Method, path: str, audit: AuditLog, + reservation: Reservation | None, ) -> Mapping[str, object] | None: try: - return await _json_body(request) + return _json_body(raw) except ControlPlaneException as error: audit.record( principal.subject, @@ -172,6 +235,7 @@ async def _body_or_audit( "denied", error.body.status, detail=error.body.detail, + reservation=reservation, ) raise @@ -209,9 +273,10 @@ async def _passthrough_scope( principal: Annotated[Principal, Depends(get_principal)], audit: Annotated[AuditLog, Depends(get_audit_log)], ) -> Principal: - try: - return await require_scopes(security_scopes, principal) - except ControlPlaneException as error: + required: Final = frozenset(Scope(scope) for scope in security_scopes.scopes) + missing: Final = principal.missing_scopes(required) + if missing: + detail: Final = f"missing scope: {', '.join(sorted(scope.value for scope in missing))}" audit.record( principal.subject, principal.email, @@ -219,10 +284,11 @@ async def _passthrough_scope( request.method, "/" + path.lstrip("/"), "denied", - error.body.status, - detail=error.body.detail, + 403, + detail=detail, ) - raise + raise ControlPlaneException(forbidden(f"Missing scope: {', '.join(sorted(scope.value for scope in missing))}.")) + return principal @router.api_route("/dataplanes/{dataplane_id}/proxy/{path:path}", methods=list(READ_METHODS)) diff --git a/litellm/proxy/control_plane/settings.py b/litellm/proxy/control_plane/settings.py index 9d0b791dc88..b382135a946 100644 --- a/litellm/proxy/control_plane/settings.py +++ b/litellm/proxy/control_plane/settings.py @@ -3,17 +3,19 @@ import os from collections.abc import Mapping from dataclasses import dataclass -from typing import Final, TypeAlias +from typing import Final, Literal, TypeAlias import yaml from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, field_validator, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict from litellm.litellm_core_utils.oauth_token_exchange import TokenExchangeProfile +from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID from litellm.proxy.control_plane.auth.rbac import ADMIN_ROLE from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH DEFAULT_SESSION_TTL_SECONDS: Final = 12 * 60 * 60 +AuditSinkKind: TypeAlias = Literal["log", "otlp"] DEFAULT_DATAPLANE_TIMEOUT_SECONDS: Final = 10.0 @@ -117,6 +119,17 @@ class ControlPlaneSettings(BaseSettings): otel_enabled: bool | None = None otel_service_name: str = "litellm-control-plane" otel_identity_attributes: bool = True + audit_sinks: tuple[AuditSinkKind, ...] = ("log",) + audit_queue_size: int = Field(default=10_000, ge=2) + audit_batch_size: int = Field(default=100, ge=1) + audit_drain_seconds: float = Field(default=5.0, gt=0) + audit_strict: bool = False + + @model_validator(mode="after") + def _strict_audit_requires_a_sink(self) -> "ControlPlaneSettings": + if self.audit_strict and not self.audit_sinks: + raise ValueError("CONTROL_PLANE_AUDIT_STRICT requires at least one audit sink.") + return self @field_validator("master_key") @classmethod @@ -136,6 +149,12 @@ class ControlPlaneSettings(BaseSettings): return self raise ValueError("password login requires a custom admin role when default roles are disabled.") + @model_validator(mode="after") + def _dataplanes_exclude_control_plane_id(self) -> "ControlPlaneSettings": + if CONTROL_PLANE_ID not in self.dataplanes: + return self + raise ValueError(f"dataplane id {CONTROL_PLANE_ID!r} is reserved.") + @property def oidc(self) -> OidcSettings | None: if self.oidc_issuer is None or self.oidc_audience is None: diff --git a/litellm/types/proxy/control_plane/audit.py b/litellm/types/proxy/control_plane/audit.py index aab75106dd5..f95ebd57092 100644 --- a/litellm/types/proxy/control_plane/audit.py +++ b/litellm/types/proxy/control_plane/audit.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from typing import Literal, TypeAlias @@ -27,3 +28,11 @@ class AuditListResponse(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) items: tuple[AuditEventResponse, ...] + + +class OcsfAuditListResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + items: tuple[Mapping[str, object], ...] + chain_valid: bool + sequence_gaps: int diff --git a/tests/unit/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py index ce39c69bb16..1156fcf78d3 100644 --- a/tests/unit/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -939,6 +939,39 @@ def test_traces_endpoint_alone_creates_an_otlp_http_exporter_spec(): assert spec.kind == "otlp_http" and spec.traces_endpoint == endpoint +@pytest.mark.parametrize("env_name", ("OTEL_LOGS_ENDPOINT", "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT")) +def test_logs_endpoint_alias_alone_posts_to_the_configured_url( + monkeypatch: pytest.MonkeyPatch, + otlp_collector: tuple[str, list[str]], + env_name: str, +) -> None: + base_url, received_paths = otlp_collector + for var in ( + "OTEL_EXPORTER", + "OTEL_EXPORTER_OTLP_PROTOCOL", + "OTEL_ENDPOINT", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_TRACES_ENDPOINT", + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + "OTEL_LOGS_ENDPOINT", + "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", + ): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv(env_name, f"{base_url}/services/collector/logs") + + config: Final = OpenTelemetryV2Config.from_env() + assert config.exporter == "otlp_http" and config.logs_endpoint == f"{base_url}/services/collector/logs" + logger_provider: Final = providers.build_logger_provider( + config, + log_exporter=providers.build_log_exporter(config), + ) + logger_provider.get_logger("endpoint-test").emit(body="logs-endpoint") + assert logger_provider.force_flush() + logger_provider.shutdown() + + assert received_paths == ["/services/collector/logs"] + + def test_traces_endpoint_env_posts_to_the_configured_url_verbatim(monkeypatch, otlp_collector): base_url, received_paths = otlp_collector for var in ("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_EXPORTER_OTLP_ENDPOINT"): @@ -1045,6 +1078,31 @@ def test_build_log_exporter_variants(): assert "OTLPLogExporter" in type(http_exporter).__name__ +def test_grpc_log_exporter_uses_the_base_endpoint_without_an_http_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter as GRPCLogExporter + + for env_name in ( + "OTEL_EXPORTER", + "OTEL_ENDPOINT", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_LOGS_ENDPOINT", + "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", + ): + monkeypatch.delenv(env_name, raising=False) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "grpc") + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://host:4317") + + config: Final = OpenTelemetryV2Config.from_env() + exporter: Final = providers.build_log_exporter(config) + try: + assert isinstance(exporter, GRPCLogExporter) + assert exporter._endpoint == "host:4317" + finally: + exporter.shutdown() + + def test_build_logger_provider_picks_processor_by_exporter_kind(): """Console and in-memory exporters export synchronously (tests depend on it); every other destination gets the batch processor.""" diff --git a/tests/unit/proxy/control_plane/audit/test_log.py b/tests/unit/proxy/control_plane/audit/test_log.py index c1c1030d25a..1590feef0e9 100644 --- a/tests/unit/proxy/control_plane/audit/test_log.py +++ b/tests/unit/proxy/control_plane/audit/test_log.py @@ -1,3 +1,6 @@ +import io +import json +import logging from collections.abc import Callable from datetime import UTC, datetime, timedelta from itertools import count @@ -5,7 +8,9 @@ from typing import Final import pytest +from litellm.proxy.control_plane.audit.event import redact_path as redact_event_path from litellm.proxy.control_plane.audit.log import build_audit_log, redact_path +from litellm.proxy.control_plane.audit.ocsf import verify T0: Final = datetime(2026, 10, 13, 12, 0, tzinfo=UTC) @@ -37,6 +42,81 @@ def test_the_ring_buffer_drops_the_oldest_event_at_capacity() -> None: assert [e.path for e in audit.recent(10)] == ["/c", "/b"] +def test_unauthenticated_failures_do_not_evict_authenticated_history() -> None: + audit: Final = build_audit_log(capacity=2, now=_ticking()) + audit.record("alice", None, "eu", "GET", "/one", "allowed", 200) + audit.record("alice", None, "eu", "GET", "/two", "allowed", 200) + for index in range(2): + audit.record_login(f"bad-{index}", "denied") + assert [event.path for event in audit.recent(10)] == ["/auth/login", "/auth/login", "/two", "/one"] + + +def test_free_text_fields_are_safe_to_encode() -> None: + audit: Final = build_audit_log(capacity=2, now=_ticking()) + event: Final = audit.record("\ud800", "\ud800", "eu", "GET", "/\ud800", "denied", detail="\ud800") + assert (event.actor, event.actor_email, event.path, event.detail) == ( + "\\ud800", + "\\ud800", + "/\\ud800", + "\\ud800", + ) + + +def test_caller_controlled_audit_fields_are_bounded_before_redaction() -> None: + audit: Final = build_audit_log(capacity=2, now=_ticking()) + path: Final = "/" + "x" * 32_768 + + event: Final = audit.record("a" * 32_768, "e" * 32_768, "d" * 32_768, "GET", path, "denied", detail="z" * 32_768) + + assert all( + len(value) <= 2048 + for value in (event.actor, event.actor_email, event.dataplane_id, event.path, event.detail) + if value is not None + ) + malformed_path: Final = "/" + "-----BEGIN PRIVATE KEY-----" * 256 + bounded: Final = audit.record("alice", None, "eu", "GET", malformed_path, "denied") + assert len(bounded.path) <= 2048 + + +def test_unknown_secret_looking_dataplane_ids_are_redacted() -> None: + audit: Final = build_audit_log(known_dataplane_ids=frozenset({"registered"})) + unknown_id: Final = "sk-unknown-dataplane" + + event: Final = audit.record("alice", None, unknown_id, "GET", "/key/list", "allowed", 200) + + assert event.dataplane_id != unknown_id + assert "sk-" not in event.dataplane_id + + +def test_dropped_unauthenticated_records_use_an_injected_info_logger() -> None: + stream: Final = io.StringIO() + audit: Final = build_audit_log(capacity=2, now=_ticking(), stream=stream) + proxy_logger: Final = logging.getLogger("LiteLLM Proxy") + previous_level: Final = proxy_logger.level + proxy_logger.setLevel(logging.WARNING) + try: + audit.record_login("alice", "denied") + finally: + proxy_logger.setLevel(previous_level) + assert "control_plane.audit" in stream.getvalue() + + +def test_failed_login_redacts_before_sealing_and_logging() -> None: + stream: Final = io.StringIO() + audit: Final = build_audit_log(capacity=2, now=_ticking(), stream=stream) + username: Final = "sk-" + "A" * 24 + + audit.record_login(username, "denied") + + (record,) = audit.recent_records(1) + emitted_body: Final = stream.getvalue().split("control_plane.audit ", 1)[1].rstrip("\n") + stored_body: Final = json.dumps(record.body, default=str, separators=(",", ":")) + assert emitted_body == stored_body + assert username not in emitted_body + assert record.event.actor != username + assert verify((record,)).valid + + def test_virtual_keys_in_paths_are_redacted_before_they_are_stored_or_logged() -> None: audit: Final = build_audit_log(capacity=10, now=_ticking()) event: Final = audit.record( @@ -62,10 +142,12 @@ def test_recent_with_an_empty_dataplane_set_shows_nothing_and_none_shows_everyth "/key/SK-live-123/reset_spend", "/key/sk%2Dlive-123/reset_spend", "/key/sk%252Dlive-123/reset_spend", + "/key/%2573k-live-123/team/secret/reset_spend", ], ) def test_encoded_virtual_keys_are_redacted_after_repeated_decoding(path: str) -> None: assert redact_path(path) == "/key/sk-[redacted]/reset_spend" + assert redact_event_path(path) == "/key/sk-[redacted]/reset_spend" @pytest.mark.parametrize( diff --git a/tests/unit/proxy/control_plane/audit/test_ocsf.py b/tests/unit/proxy/control_plane/audit/test_ocsf.py new file mode 100644 index 00000000000..9d6ec9efd9e --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_ocsf.py @@ -0,0 +1,100 @@ +from datetime import UTC, datetime +from typing import Final + +from litellm.proxy.control_plane.audit.event import AuditEvent +from litellm.proxy.control_plane.audit.ocsf import ( + AuditRecord, + ChainCheck, + ChainState, + Product, + chain, + section, + to_ocsf, + verify, +) + +T0: Final = datetime(2026, 10, 13, 12, 0, tzinfo=UTC) +PRODUCT: Final = Product(version="1.2.3") + + +def _event(**overrides: object) -> AuditEvent: + base: Final = { + "at": T0, + "actor": "alice", + "actor_email": "alice@x.test", + "dataplane_id": "eu", + "method": "POST", + "path": "/key/generate", + "decision": "allowed", + "status": 200, + "credential": "service_token", + "detail": None, + "trace_id": "0af7651916cd43dd8448eb211c80319c", + "span_id": "b7ad6b7169203331", + } + return AuditEvent(**{**base, **overrides}) # pyright: ignore[reportArgumentType] # test literal + + +def test_an_allowed_write_is_an_api_activity_create_with_success_status() -> None: + body: Final = to_ocsf(_event(), PRODUCT, ChainState()) + assert (body["class_uid"], body["category_uid"], body["type_uid"], body["activity_id"]) == (6003, 6, 600301, 1) + assert (body["status_id"], body["severity_id"], body["time"]) == (1, 1, int(T0.timestamp() * 1000)) + assert body["actor"] == { + "user": {"uid": "alice", "email_addr": "alice@x.test"}, + "session": {"credential_uid": "service_token"}, + "authorizations": [{"decision": "allowed"}], + } + assert body["api"] == { + "operation": "POST", + "service": {"name": "eu", "uid": "dataplane"}, + "response": {"code": 200}, + } + assert body["http_request"] == {"http_method": "POST", "url": {"path": "/key/generate"}} + assert body["metadata"] == { + "version": "1.5.0", + "product": {"name": "LiteLLM Control Plane", "vendor_name": "LiteLLM", "version": "1.2.3"}, + "sequence": 1, + "correlation_uid": "0af7651916cd43dd8448eb211c80319c", + } + + +def test_denied_and_failed_calls_are_failures_with_medium_severity() -> None: + denied: Final = to_ocsf(_event(decision="denied", status=None, detail="no scope"), PRODUCT, ChainState()) + failed: Final = to_ocsf(_event(status=400), PRODUCT, ChainState()) + unreachable: Final = to_ocsf(_event(decision="unreachable", status=None, method="GET"), PRODUCT, ChainState()) + assert (denied["status_id"], denied["severity_id"], denied["status_detail"]) == (2, 3, "no scope") + assert (failed["status_id"], failed["status_code"]) == (2, "400") + assert (unreachable["status_id"], unreachable["activity_id"], unreachable["activity_name"]) == (2, 2, "Read") + + +def test_a_login_is_an_authentication_logon_event() -> None: + body: Final = to_ocsf( + _event(kind="authentication", method="POST", path="/auth/login", status=None), PRODUCT, ChainState() + ) + assert (body["class_uid"], body["category_uid"], body["type_uid"], body["activity_name"]) == ( + 3002, + 3, + 300201, + "Logon", + ) + assert body["user"] == {"uid": "alice"} + + +def test_records_chain_and_verify_detects_a_tampered_or_missing_link() -> None: + state0: Final = ChainState() + state1, first = chain(_event(), PRODUCT, state0) + state2, second = chain(_event(method="GET", path="/key/list"), PRODUCT, state1) + _, third = chain(_event(method="DELETE", path="/key/delete"), PRODUCT, state2) + assert (first.sequence, second.sequence, third.sequence) == (1, 2, 3) + assert section(first.body, "unmapped", "litellm")["prev_hash"] is None + assert section(second.body, "unmapped", "litellm")["prev_hash"] == first.hash + assert verify((first, second, third)) == ChainCheck(valid=True, gaps=0) + assert verify((first, third)) == ChainCheck(valid=True, gaps=1), "a missing record is a gap, not tampering" + assert verify((second, third)) == ChainCheck(valid=True, gaps=0), "a page may start after the first record" + tampered_body: Final = {**second.body, "actor": {"user": {"uid": "mallory"}}} + tampered: Final = AuditRecord(second.event, second.sequence, second.hash, tampered_body) + assert verify((first, tampered, third)).valid is False, "an edited record no longer matches its hash" + _, forged = chain(_event(method="GET", path="/key/list"), PRODUCT, ChainState(1, "sha256:" + "0" * 64)) + assert verify((first, forged)).valid is False, "a re-hashed record that points at the wrong predecessor" + assert verify((forged,)).valid is True, "alone it is consistent, only the link exposes it" + assert verify(()) == ChainCheck(valid=True, gaps=0) diff --git a/tests/unit/proxy/control_plane/audit/test_pipeline.py b/tests/unit/proxy/control_plane/audit/test_pipeline.py new file mode 100644 index 00000000000..94f4d90d0e5 --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_pipeline.py @@ -0,0 +1,304 @@ +import asyncio +from dataclasses import dataclass, field +from datetime import UTC, datetime +from itertools import accumulate +from typing import Final + +import pytest + +from litellm.proxy.control_plane.audit.event import AuditEvent +from litellm.proxy.control_plane.audit.log import build_audit_log +from litellm.proxy.control_plane.audit.ocsf import AuditRecord, ChainState, Product, chain +from litellm.proxy.control_plane.audit.pipeline import ( + Accepted, + Batch, + Delivered, + DeliveryResult, + Dropped, + Refused, + Undelivered, + build_pipeline, +) + +T0: Final = datetime(2026, 10, 13, 12, 0, tzinfo=UTC) + + +@dataclass +class FakeSink: # mutable-ok: records deliveries for assertions + name: str = "fake" + fail_first: int = 0 + fail_close: bool = False + batches: list[Batch] = field(default_factory=list) + closed: bool = False + + async def deliver(self, batch: Batch) -> DeliveryResult: + if self.fail_first > 0: + self.fail_first -= 1 + return Undelivered("sink down") + self.batches.append(batch) + return Delivered(len(batch)) + + async def close(self) -> None: + self.closed = True + if self.fail_close: + raise RuntimeError("sink refused to close") + + +@dataclass +class PrefixSettlingSink: # mutable-ok: records the batches received for assertions + batches: list[Batch] = field(default_factory=list) + name: str = "prefix" + + async def deliver(self, batch: Batch) -> DeliveryResult: + self.batches.append(batch) + return Undelivered("partial delivery", settled=2) if len(self.batches) == 1 else Delivered(len(batch)) + + async def close(self) -> None: + return None + + +@dataclass +class Clock: # mutable-ok: records sleeps for assertions + slept: list[float] = field(default_factory=list) + + async def sleep(self, seconds: float) -> None: + self.slept.append(seconds) + + +def _records(count: int) -> tuple[AuditRecord, ...]: + def append_record( + previous: tuple[ChainState, tuple[AuditRecord, ...]], event: AuditEvent + ) -> tuple[ChainState, tuple[AuditRecord, ...]]: + state, records = previous + next_state, record = chain(event, Product("t"), state) + return next_state, (*records, record) + + events: Final = tuple( + AuditEvent(T0, "alice", None, "eu", "GET", f"/key/{index}", "allowed", 200, None, None) + for index in range(count) + ) + results: Final = tuple(accumulate(events, append_record, initial=(ChainState(), ()))) + return results[-1][1] + + +@pytest.mark.asyncio +async def test_records_are_delivered_in_order_in_batches_of_the_configured_size() -> None: + sink, clock = FakeSink(), Clock() + pipeline: Final = build_pipeline((sink,), queue_size=10, batch_size=2, strict=False, sleep=clock.sleep) + for record in _records(5): + assert pipeline.publish(record) == (Accepted(),) + (worker,) = pipeline.workers + assert [await worker.flush_once() for _ in range(4)] == [2, 2, 1, 0] + assert [[r.sequence for r in batch] for batch in sink.batches] == [[1, 2], [3, 4], [5]] + + +@pytest.mark.asyncio +async def test_a_full_queue_drops_the_oldest_and_counts_it_unless_strict() -> None: + sink: Final = FakeSink() + lenient: Final = build_pipeline((sink,), queue_size=2, batch_size=10, strict=False) + one, two, three = _records(3) + assert [lenient.publish(r) for r in (one, two, three)] == [(Accepted(),), (Accepted(),), (Dropped("fake"),)] + (worker,) = lenient.workers + assert worker.dropped == 1 and lenient.accepting() is True + await worker.flush_once() + assert [r.sequence for r in sink.batches[0]] == [2, 3] + + strict: Final = build_pipeline((FakeSink(),), queue_size=2, batch_size=10, strict=True) + assert [strict.publish(r) for r in (one, two, three)] == [(Accepted(),), (Accepted(),), (Refused("fake"),)] + assert strict.accepting() is False and strict.workers[0].queue.qsize() == 2 + + +def test_a_strict_reservation_holds_the_slot_until_the_record_consumes_or_releases_it() -> None: + strict: Final = build_pipeline((FakeSink(),), queue_size=1, batch_size=10, strict=True) + one, two = _records(2) + held: Final = strict.reserve() + assert held is not None and strict.reserve() is None, "two in-flight writes cannot share one free slot" + assert strict.publish(one) == (Refused("fake"),), "an unreserved event does not steal a reserved slot" + held.release() + assert strict.publish(one) == (Accepted(),) + held.release() + assert strict.publish(two) == (Refused("fake"),), "a second release does not free a slot that is in use" + assert build_pipeline((FakeSink(),), queue_size=1, batch_size=10, strict=False).reserve() is not None + + +@pytest.mark.asyncio +async def test_a_failing_sink_is_retried_with_backoff_and_the_batch_is_never_lost() -> None: + sink, clock = FakeSink(fail_first=3), Clock() + pipeline: Final = build_pipeline((sink,), queue_size=10, batch_size=10, strict=False, sleep=clock.sleep) + for record in _records(2): + pipeline.publish(record) + (worker,) = pipeline.workers + assert await worker.flush_once() == 2 + assert clock.slept == [0.5, 1.0, 2.0] and [r.sequence for r in sink.batches[0]] == [1, 2] + assert worker.failures == 0 + + +@pytest.mark.asyncio +async def test_worker_retries_only_the_unsettled_suffix() -> None: + sink: Final = PrefixSettlingSink() + clock: Final = Clock() + pipeline: Final = build_pipeline((sink,), queue_size=10, batch_size=4, strict=False, sleep=clock.sleep) + records: Final = _records(4) + for record in records: + pipeline.publish(record) + + (worker,) = pipeline.workers + assert await worker.flush_once() == 4 + assert sink.batches == [records, records[2:]] + assert clock.slept == [0.5] + + +@pytest.mark.asyncio +async def test_a_sink_down_for_thousands_of_attempts_still_gets_the_batch_when_it_recovers() -> None: + sink, clock = FakeSink(fail_first=5_000), Clock() + pipeline: Final = build_pipeline((sink,), queue_size=10, batch_size=10, strict=False, sleep=clock.sleep) + pipeline.publish(_records(1)[0]) + assert await pipeline.workers[0].flush_once() == 1 + assert len(clock.slept) == 5_000 and clock.slept[-1] == 30.0 and sink.batches + + +@pytest.mark.asyncio +async def test_each_sink_has_its_own_queue_so_one_down_sink_does_not_block_another() -> None: + down, up, clock = FakeSink("down", fail_first=1), FakeSink("up"), Clock() + pipeline: Final = build_pipeline((down, up), queue_size=10, batch_size=10, strict=False, sleep=clock.sleep) + (record,) = _records(1) + assert pipeline.publish(record) == (Accepted(), Accepted()) + await pipeline.workers[1].flush_once() + assert up.batches and not down.batches + await pipeline.workers[0].flush_once() + assert down.batches + + +@pytest.mark.asyncio +async def test_start_restarts_a_worker_with_a_finished_task() -> None: + sink: Final = FakeSink() + pipeline: Final = build_pipeline((sink,), queue_size=2, batch_size=1, strict=False) + (worker,) = pipeline.workers + previous: Final = asyncio.create_task(asyncio.sleep(0)) + await previous + worker.task = previous + + worker.start() + + assert worker.task is not None and worker.task is not previous and not worker.task.done() + await worker.close(drain_seconds=0.1) + + +@pytest.mark.asyncio +async def test_start_runs_the_worker_and_close_drains_then_closes_the_sink() -> None: + sink: Final = FakeSink() + pipeline: Final = build_pipeline((sink,), queue_size=10, batch_size=10, strict=False) + pipeline.start() + for record in _records(3): + pipeline.publish(record) + await asyncio.sleep(0) + await pipeline.close() + assert sum(len(batch) for batch in sink.batches) == 3 and sink.closed is True + assert pipeline.workers[0].queue.empty() + + +@pytest.mark.asyncio +async def test_close_isolates_sink_failures_and_closes_every_sink(caplog: pytest.LogCaptureFixture) -> None: + first, second = FakeSink("first", fail_close=True), FakeSink("second") + pipeline: Final = build_pipeline((first, second), queue_size=2, batch_size=10, strict=False) + (record,) = _records(1) + pipeline.publish(record) + + await pipeline.close() + + assert first.closed and second.closed + assert first.batches == [(record,)] and second.batches == [(record,)] + assert "sink refused to close" in caplog.text + + +@dataclass +class StallingSink: # mutable-ok: the test controls a blocked delivery + stalled: asyncio.Event = field(default_factory=asyncio.Event) + released: asyncio.Event = field(default_factory=asyncio.Event) + calls: int = 0 + batches: list[Batch] = field(default_factory=list) + name: str = "stalling" + + async def deliver(self, batch: Batch) -> DeliveryResult: + self.calls += 1 + if self.calls == 1: + self.stalled.set() + await self.released.wait() + self.batches.append(batch) + return Delivered(len(batch)) + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_close_redelivers_the_batch_a_worker_was_holding_when_it_crashed() -> None: + sink: Final = CrashingSink() + pipeline: Final = build_pipeline((sink,), queue_size=10, batch_size=10, strict=False) + pipeline.start() + for record in _records(2): + pipeline.publish(record) + await sink.crashed.wait() + (worker,) = pipeline.workers + task: Final = worker.task + assert task is not None + with pytest.raises(RuntimeError, match="delivery crashed"): + await task + await pipeline.close() + assert [[r.sequence for r in batch] for batch in sink.batches] == [[1, 2]] + + +@dataclass +class CrashingSink: # mutable-ok: the test observes the failed delivery and its retry + crashed: asyncio.Event = field(default_factory=asyncio.Event) + calls: int = 0 + batches: list[Batch] = field(default_factory=list) + name: str = "crashing" + + async def deliver(self, batch: Batch) -> DeliveryResult: + self.calls += 1 + if self.calls == 1: + self.crashed.set() + raise RuntimeError("delivery crashed") + self.batches.append(batch) + return Delivered(len(batch)) + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_auth_failures_do_not_consume_strict_write_reservations() -> None: + sink: Final = StallingSink() + pipeline: Final = build_pipeline((sink,), queue_size=4, batch_size=1, strict=True) + audit: Final = build_audit_log(pipeline=pipeline) + pipeline.start() + audit.record("anonymous", None, "control-plane", "GET", "/auth/me", "denied", 401, kind="authentication") + await sink.stalled.wait() + for _ in range(6): + audit.record("anonymous", None, "control-plane", "GET", "/auth/me", "denied", 401, kind="authentication") + (worker,) = pipeline.workers + assert worker.queue.qsize() == 2 + assert worker.dropped == 4 + reservation: Final = audit.reserve() + assert reservation is not None + reservation.release() + sink.released.set() + await pipeline.close() + + +@pytest.mark.asyncio +async def test_unauthenticated_failures_keep_a_queue_slot_for_writes() -> None: + sink: Final = FakeSink() + pipeline: Final = build_pipeline((sink,), queue_size=2, batch_size=10, strict=False) + audit: Final = build_audit_log(pipeline=pipeline, now=lambda: T0) + + audit.record_login("anonymous", "denied") + audit.record_login("another-anonymous", "denied") + + (worker,) = pipeline.workers + assert worker.queue.qsize() == 1 + assert worker.dropped == 1 + assert await worker.flush_once() == 1 + assert len(sink.batches) == 1 and len(sink.batches[0]) == 1 + assert sink.batches[0][0].event.actor == "anonymous" diff --git a/tests/unit/proxy/control_plane/audit/test_router.py b/tests/unit/proxy/control_plane/audit/test_router.py new file mode 100644 index 00000000000..41ab9f50e8e --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_router.py @@ -0,0 +1,111 @@ +from typing import Final + +from litellm.proxy.control_plane.audit.dependencies import get_audit_log +from litellm.proxy.control_plane.audit.log import build_audit_log +from litellm.proxy.control_plane.auth.rbac import EDITOR_ROLE +from tests.unit.proxy.control_plane.support import response, rig + +EU: Final = "https://eu.test" + + +def test_ocsf_listing_returns_chained_records_newest_first_with_a_valid_chain() -> None: + client, _, minter = rig({f"{EU}/key/list": [response(200, {})], f"{EU}/key/generate": [response(200, {})]}) + client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")) + client.post("/api/v1/dataplanes/eu/proxy/key/generate", json={}, headers=minter.headers("eu", role=EDITOR_ROLE)) + page: Final = client.get("/api/v1/audit/ocsf?dataplane=eu", headers=minter.headers(role="admin")).json() + assert [(i["class_uid"], i["activity_name"], i["http_request"]["url"]["path"]) for i in page["items"]] == [ + (6003, "Create", "/key/generate"), + (6003, "Read", "/key/list"), + ] + assert (page["chain_valid"], page["sequence_gaps"]) == (True, 0) + assert page["items"][0]["unmapped"]["litellm"]["prev_hash"] == page["items"][1]["unmapped"]["litellm"]["hash"] + newest: Final = client.get("/api/v1/audit/ocsf?limit=1", headers=minter.headers(role="admin")).json() + assert (newest["chain_valid"], newest["sequence_gaps"], len(newest["items"])) == (True, 0, 1) + reads: Final = client.get( + "/api/v1/audit/ocsf?actor=alice&dataplane=eu", headers=minter.headers(role="admin") + ).json() + assert reads["chain_valid"] is True, "a filtered page is intact even though it skips records" + + +def test_registered_secret_looking_dataplane_ids_are_visible_only_to_their_auditor() -> None: + dataplane_id: Final = "sk-production-dataplane" + dataplane_url: Final = "https://production.test" + client, _, minter = rig( + {f"{dataplane_url}/key/list": [response(200, {"keys": []})]}, + dataplanes={ + dataplane_id: { + "url": dataplane_url, + "service_token": {"token": "sk-upstream"}, + } + }, + ) + client.get( + f"/api/v1/dataplanes/{dataplane_id}/proxy/key/list", + headers=minter.headers(dataplane_id), + ) + + registered: Final = client.get( + f"/api/v1/audit?dataplane={dataplane_id}", + headers=minter.headers(dataplane_id, role="admin"), + ).json()["items"] + redacted: Final = client.get( + "/api/v1/audit?dataplane=REDACTED", + headers=minter.headers("REDACTED", role="admin"), + ).json()["items"] + + assert len(registered) == 1 + assert registered[0]["dataplane_id"] == dataplane_id + assert redacted == [] + + +def test_logins_denied_scopes_and_rejected_tokens_are_audited_against_the_control_plane() -> None: + client, _, minter = rig({}, UI_USERNAME="admin", UI_PASSWORD="pw") + assert client.post("/api/v1/auth/login", json={"username": "admin", "password": "nope"}).status_code == 401 + assert client.post("/api/v1/auth/login", json={"username": "admin", "password": "pw"}).status_code == 200 + assert client.get("/api/v1/audit", headers={"Authorization": "Bearer garbage"}).status_code == 401 + assert client.get("/api/v1/audit", headers=minter.headers(role=EDITOR_ROLE)).status_code == 403 + trail: Final = client.get("/api/v1/audit?dataplane=control-plane", headers=minter.headers(role="admin")).json()[ + "items" + ] + assert [(e["actor"], e["path"], e["decision"], e["status"], e["detail"]) for e in trail] == [ + ("alice", "/api/v1/audit", "denied", 403, "missing scope: audit:read"), + ("anonymous", "/api/v1/audit", "denied", 401, trail[1]["detail"]), + ("admin", "/auth/login", "allowed", None, None), + ("admin", "/auth/login", "denied", None, "invalid credentials"), + ] + ocsf: Final = client.get( + "/api/v1/audit/ocsf?dataplane=control-plane&limit=1&actor=admin", headers=minter.headers(role="admin") + ) + assert (ocsf.json()["items"][0]["class_uid"], ocsf.json()["chain_valid"]) == (3002, True) + bound: Final = client.get( + "/api/v1/audit?dataplane=control-plane", headers=minter.headers("control-plane", role="admin") + ) + assert bound.status_code == 403, "control plane events are fleet-wide only" + implicit: Final = client.get("/api/v1/audit", headers=minter.headers("control-plane", "eu", role="admin")).json() + assert implicit["items"] == [], "a binding that names the control plane does not grant its events either" + + +def test_strict_mode_refuses_writes_but_not_reads_when_no_sink_can_take_the_record() -> None: + client, transport, minter = rig( + {f"{EU}/key/list": [response(200, {}), response(200, {}), response(200, {})]}, + audit_strict=True, + audit_queue_size=2, + ) + client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")) + client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")) + refused: Final = client.post( + "/api/v1/dataplanes/eu/proxy/key/generate", json={}, headers=minter.headers("eu", role=EDITOR_ROLE) + ) + assert (refused.status_code, refused.json()["type"]) == (503, "urn:litellm:error:audit-unavailable") + assert transport.writes == [], "a refused write never reaches the dataplane" + assert client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")).status_code != 503 + + +def test_audit_endpoints_serialize_free_text_with_lone_surrogates() -> None: + client, _, minter = rig({}) + audit: Final = build_audit_log() + audit.record("\ud800", "\ud800", "eu", "GET", "/\ud800", "denied", detail="\ud800") + client.app.dependency_overrides[get_audit_log] = lambda: audit + headers: Final = minter.headers(role="admin") + assert client.get("/api/v1/audit", headers=headers).status_code == 200 + assert client.get("/api/v1/audit/ocsf", headers=headers).status_code == 200 diff --git a/tests/unit/proxy/control_plane/audit/test_sinks.py b/tests/unit/proxy/control_plane/audit/test_sinks.py new file mode 100644 index 00000000000..4910f8b2ec3 --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_sinks.py @@ -0,0 +1,230 @@ +import asyncio +import hashlib +import io +import logging +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from datetime import UTC, datetime +from threading import Event, Lock +from typing import Final + +import pytest +from opentelemetry.sdk._logs import ReadableLogRecord +from opentelemetry.sdk._logs.export import InMemoryLogRecordExporter, LogRecordExporter, LogRecordExportResult +from pydantic import TypeAdapter + +from litellm.proxy.control_plane.audit.event import AuditEvent +from litellm.proxy.control_plane.audit.log import build_audit_log +from litellm.proxy.control_plane.audit.ocsf import HASH_PREFIX, ChainState, Product, canonical, chain, section +from litellm.proxy.control_plane.audit.pipeline import Accepted, Delivered, Undelivered, build_pipeline +from litellm.proxy.control_plane.audit.sinks import LogSink, OtlpLogSink +from litellm.proxy.control_plane.common.tracing import build_resource +from litellm.proxy.control_plane.settings import ControlPlaneSettings +from tests.unit.proxy.control_plane.support import MASTER_KEY + +T0: Final = datetime(2026, 10, 13, 12, 0, tzinfo=UTC) +TRACE: Final = "0af7651916cd43dd8448eb211c80319c" +SPAN: Final = "b7ad6b7169203331" + + +_JSON: Final = TypeAdapter(Mapping[str, object]) + + +def _record(): + event: Final = AuditEvent( + T0, "alice", None, "eu", "POST", "/key/generate", "denied", 403, None, "no", "api_activity", TRACE, SPAN + ) + return chain(event, Product("t"), ChainState())[1] + + +@pytest.mark.asyncio +async def test_the_otlp_sink_emits_one_log_record_per_event_tied_to_the_request_span() -> None: + exporter: Final = InMemoryLogRecordExporter() + settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY}) + sink: Final = OtlpLogSink(exporter, build_resource(settings)) + record: Final = _record() + assert await sink.deliver((record,)) == Delivered(1) + (emitted,) = exporter.get_finished_logs() + log: Final = emitted.log_record + assert (format(log.trace_id, "032x"), format(log.span_id, "016x")) == (TRACE, SPAN) + assert isinstance(log.body, str) and _JSON.validate_json(log.body) == record.body + assert log.severity_text == "WARN" + assert emitted.resource.attributes["service.name"] == settings.otel_service_name + assert log.attributes is not None and log.attributes["event.name"] == "litellm.control_plane.audit" + assert log.attributes["litellm.control_plane.audit.hash"] == record.hash + assert log.timestamp == int(T0.timestamp() * 1_000_000_000) + + +class RefusingExporter(LogRecordExporter): + def export(self, batch: Sequence[ReadableLogRecord]) -> LogRecordExportResult: + return LogRecordExportResult.FAILURE + + def shutdown(self) -> None: + return None + + def force_flush(self, timeout_millis: int = 30000) -> bool: + return True + + +@dataclass +class BlockingExporter(LogRecordExporter): # mutable-ok: coordinates exporter state with its worker thread + export_started: Event = field(default_factory=Event) + export_release: Event = field(default_factory=Event) + export_finished: Event = field(default_factory=Event) + lock: Lock = field(default_factory=Lock) + export_calls: int = 0 + shutdown_calls: int = 0 + exporting: bool = False + shutdown_overlapped_export: bool = False + + def export(self, batch: Sequence[ReadableLogRecord]) -> LogRecordExportResult: + with self.lock: + self.export_calls += 1 + self.exporting = True + self.export_started.set() + self.export_release.wait() + with self.lock: + self.exporting = False + self.export_finished.set() + return LogRecordExportResult.SUCCESS + + def shutdown(self) -> None: + with self.lock: + self.shutdown_calls += 1 + self.shutdown_overlapped_export = self.shutdown_overlapped_export or self.exporting + + def force_flush(self, timeout_millis: int = 30000) -> bool: + return True + + +@pytest.mark.asyncio +async def test_the_otlp_sink_reports_a_refused_export_so_the_worker_retries_it() -> None: + settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY}) + assert await OtlpLogSink(RefusingExporter(), build_resource(settings)).deliver((_record(),)) == Undelivered( + "otlp export failed" + ) + + +@pytest.mark.asyncio +async def test_worker_close_waits_for_an_inflight_otlp_export_before_shutdown() -> None: + settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY}) + exporter: Final = BlockingExporter() + sink: Final = OtlpLogSink(exporter, build_resource(settings)) + pipeline: Final = build_pipeline((sink,), queue_size=2, batch_size=1, strict=False) + pipeline.start() + assert pipeline.publish(_record()) == (Accepted(),) + (worker,) = pipeline.workers + + try: + assert await asyncio.to_thread(exporter.export_started.wait, 5) + close_task: Final = asyncio.create_task(worker.close(drain_seconds=5)) + await asyncio.sleep(0) + exporter.export_release.set() + await close_task + finally: + exporter.export_release.set() + if exporter.export_started.is_set(): + assert await asyncio.to_thread(exporter.export_finished.wait, 5) + + assert exporter.export_calls == 1 + assert exporter.shutdown_calls == 1 + assert not exporter.shutdown_overlapped_export + + +@pytest.mark.asyncio +async def test_worker_close_stops_waiting_for_an_inflight_export_at_the_drain_deadline() -> None: + settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY}) + exporter: Final = BlockingExporter() + sink: Final = OtlpLogSink(exporter, build_resource(settings)) + pipeline: Final = build_pipeline((sink,), queue_size=2, batch_size=1, strict=False) + pipeline.start() + assert pipeline.publish(_record()) == (Accepted(),) + (worker,) = pipeline.workers + assert await asyncio.to_thread(exporter.export_started.wait, 5) + + try: + await asyncio.wait_for(worker.close(drain_seconds=0.1), timeout=1) + assert not exporter.export_release.is_set() + assert exporter.shutdown_calls == 1 + finally: + exporter.export_release.set() + assert await asyncio.to_thread(exporter.export_finished.wait, 5) + + +@pytest.mark.asyncio +async def test_the_otlp_sink_retries_when_batch_construction_raises() -> None: + settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY}) + event: Final = AuditEvent( + T0, + "alice", + None, + "eu", + "POST", + "/key/generate", + "denied", + 403, + None, + "no", + "api_activity", + "not-hex", + SPAN, + ) + record: Final = chain(event, Product("t"), ChainState())[1] + result: Final = await OtlpLogSink(InMemoryLogRecordExporter(), build_resource(settings)).deliver((record,)) + assert result == Undelivered("otlp export raised ValueError") + + +@pytest.mark.asyncio +async def test_the_log_sink_writes_the_ocsf_object_as_one_line_even_when_proxy_logs_are_at_warning() -> None: + record: Final = _record() + proxy_logger: Final = logging.getLogger("LiteLLM Proxy") + stream: Final = io.StringIO() + previous_level: Final = proxy_logger.level + proxy_logger.setLevel(logging.WARNING) + try: + assert await LogSink(stream=stream).deliver((record,)) == Delivered(1) + finally: + proxy_logger.setLevel(previous_level) + line: Final = stream.getvalue().strip() + assert '"class_uid":6003' in line and record.hash in line and "\n" not in line + + +@pytest.mark.asyncio +async def test_the_log_sink_preserves_the_seal_and_registered_secret_looking_dataplane_id() -> None: + dataplane_id: Final = "sk-production-dataplane" + audit: Final = build_audit_log( + now=lambda: T0, + product=Product("t"), + trace=lambda: (TRACE, SPAN), + known_dataplane_ids=frozenset({dataplane_id}), + ) + audit.record("alice", None, dataplane_id, "POST", "/key/list", "allowed", 200, credential="service_token") + record: Final = audit.recent_records(1)[0] + stream: Final = io.StringIO() + assert await LogSink(stream=stream).deliver((record,)) == Delivered(1) + + line: Final = stream.getvalue().strip() + assert line.startswith("control_plane.audit ") + body: Final = _JSON.validate_json(line.removeprefix("control_plane.audit ")) + litellm_block: Final = section(body, "unmapped", "litellm") + printed_hash: Final = litellm_block["hash"] + hashless_body: Final = { + **body, + "unmapped": { + "litellm": {key: value for key, value in litellm_block.items() if key != "hash"}, + }, + } + + assert printed_hash == record.hash + assert printed_hash == HASH_PREFIX + hashlib.sha256(canonical(hashless_body)).hexdigest() + assert section(body, "api", "service")["name"] == dataplane_id + + +def test_audit_log_redacts_the_http_method_before_chaining() -> None: + method: Final = "sk-production-dataplane" + audit: Final = build_audit_log(now=lambda: T0, product=Product("t"), trace=lambda: (TRACE, SPAN)) + event: Final = audit.record("alice", None, "eu", method, "/key/list", "allowed", 200) + record: Final = audit.recent_records(1)[0] + + assert event.method != method + assert section(record.body, "http_request")["http_method"] == event.method diff --git a/tests/unit/proxy/control_plane/proxy/test_router.py b/tests/unit/proxy/control_plane/proxy/test_router.py index 0a53e0c382b..bcec95d23be 100644 --- a/tests/unit/proxy/control_plane/proxy/test_router.py +++ b/tests/unit/proxy/control_plane/proxy/test_router.py @@ -1,11 +1,29 @@ +import asyncio from typing import Final import httpx +import pytest +from starlette.requests import Request +from starlette.types import Message, Scope -from litellm.proxy.control_plane.auth.rbac import EDITOR_ROLE, USAGE_VIEWER_ROLE, VIEWER_ROLE +from litellm.proxy.control_plane.auth.rbac import ( + EDITOR_ROLE, + USAGE_VIEWER_ROLE, + VIEWER_ROLE, + build_principal, +) +from litellm.proxy.control_plane.common.exceptions import ControlPlaneException +from litellm.proxy.control_plane.container import build_container from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, RetryPolicy +from litellm.proxy.control_plane.dataplanes.credentials import Caller from litellm.proxy.control_plane.dataplanes.dependencies import get_dataplane_client +from litellm.proxy.control_plane.proxy.router import _forward +from litellm.proxy.control_plane.settings import ControlPlaneSettings from tests.unit.proxy.control_plane.support import ( + DATAPLANES, + MASTER_KEY, + RoutedTransport, + no_exchange, response, rig, ) @@ -60,7 +78,7 @@ def test_scope_denied_passthrough_reads_and_writes_are_audited() -> None: ("alice", "eu", "GET", "/key/sk-[redacted]", "denied"), ("alice", "eu", "POST", "/key/sk-[redacted]/regenerate", "denied"), } - assert all("Missing scope" in event["detail"] for event in trail) + assert {event["detail"] for event in trail} == {"missing scope: proxy:read", "missing scope: proxy:write"} def test_get_preserves_repeated_query_parameters_in_order() -> None: @@ -92,7 +110,7 @@ def test_writes_need_proxy_write_go_out_once_and_are_audited() -> None: assert timed_out.status_code == 502 and "1 attempt" in timed_out.json()["detail"] failed: Final = client.delete("/api/v1/dataplanes/eu/proxy/key/delete", headers=editor) assert (failed.status_code, failed.json()) == (400, {"error": {"message": "no such key"}}) - trail: Final = client.get("/api/v1/audit", headers=minter.headers(role="admin")).json()["items"] + trail: Final = client.get("/api/v1/audit?dataplane=eu", headers=minter.headers(role="admin")).json()["items"] assert [(e["method"], e["path"], e["decision"], e["status"]) for e in trail] == [ ("DELETE", "/key/delete", "allowed", 400), ("PATCH", "/team/update", "unreachable", None), @@ -119,7 +137,9 @@ def test_passthrough_write_accepts_a_200_kilobyte_body() -> None: def test_denied_paths_are_403_audited_and_never_reach_the_dataplane() -> None: client, transport, minter = rig({}) blocked: Final = client.post( - "/api/v1/dataplanes/eu/proxy/v1/chat/completions", json={"model": "x"}, headers=minter.headers(role=EDITOR_ROLE) + "/api/v1/dataplanes/eu/proxy/v1/chat/completions", + content=b"[", + headers={**minter.headers(role=EDITOR_ROLE), "Content-Type": "application/json"}, ) assert blocked.status_code == 403 and "inference" in blocked.json()["detail"] assert transport.calls == [] @@ -149,7 +169,10 @@ def test_audit_read_needs_its_scope_and_filters_by_actor() -> None: client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")) assert client.get("/api/v1/audit", headers=minter.headers(role=EDITOR_ROLE)).status_code == 403 assert client.get("/api/v1/audit?actor=bob", headers=minter.headers(role="admin")).json()["items"] == [] - assert len(client.get("/api/v1/audit?actor=alice", headers=minter.headers(role="admin")).json()["items"]) == 1 + assert ( + len(client.get("/api/v1/audit?actor=alice&dataplane=eu", headers=minter.headers(role="admin")).json()["items"]) + == 1 + ) def test_audit_reads_are_clamped_to_the_callers_dataplane_binding() -> None: @@ -169,9 +192,7 @@ def test_a_redirecting_dataplane_is_a_502_and_the_redirect_is_never_followed() - client, transport, minter = rig( {f"{EU}/v1/mcp/server/tools/call": [httpx.Response(302, headers={"Location": "http://127.0.0.1:9/"})]} ) - reply: Final = client.get( - "/api/v1/dataplanes/eu/proxy/v1/mcp/server/tools/call", headers=minter.headers("eu") - ) + reply: Final = client.get("/api/v1/dataplanes/eu/proxy/v1/mcp/server/tools/call", headers=minter.headers("eu")) assert reply.status_code == 502 and "redirect" in reply.json()["detail"] assert [url for url, _, _ in transport.calls] == [f"{EU}/v1/mcp/server/tools/call"] trail: Final = client.get("/api/v1/audit", headers=minter.headers(role="admin")).json()["items"] @@ -324,3 +345,193 @@ def test_guardrail_apply_is_denied_even_under_a_management_prefix() -> None: assert reply.status_code == 403 and "guardrail apply" in reply.json()["detail"] assert transport.calls == [] + + +@pytest.mark.asyncio +async def test_pending_request_body_does_not_hold_a_strict_audit_reservation() -> None: + settings: Final = ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "dataplanes": DATAPLANES, + "audit_queue_size": 2, + "audit_strict": True, + } + ) + transport: Final = RoutedTransport({f"{EU}/key/generate": iter([response(200, {"token_id": "t1"})])}) + container: Final = build_container(settings, transport=transport, exchanger=no_exchange) + principal: Final = build_principal( + container.policy, + "alice", + "alice@x.test", + frozenset({EDITOR_ROLE}), + None, + frozenset({"eu"}), + "session", + ) + caller: Final = Caller(principal, "session-token") + dataplane: Final = container.registry.get("eu") + assert dataplane is not None + body_started: Final = asyncio.Event() + release_body: Final = asyncio.Event() + + async def receive() -> Message: + body_started.set() + await release_body.wait() + return {"type": "http.request", "body": b'{"key_alias":"x"}', "more_body": False} + + scope: Final[Scope] = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/key/generate", + "raw_path": b"/key/generate", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + "client": ("testclient", 123), + "server": ("testserver", 80), + } + request: Final = Request(scope, receive) + write: Final = asyncio.create_task( + _forward( + request, + "key/generate", + principal, + dataplane, + caller, + container.dataplane_client, + container.audit_log, + ) + ) + await body_started.wait() + await asyncio.sleep(0) + try: + first: Final = container.audit_log.reserve() + second: Final = container.audit_log.reserve() + assert first is not None and second is not None + first.release() + second.release() + finally: + release_body.set() + reply: Final = await write + assert reply.status_code == 200 + assert [(event.path, event.decision) for event in container.audit_log.recent(1)] == [("/key/generate", "allowed")] + + +@pytest.mark.asyncio +async def test_malformed_write_body_uses_its_strict_audit_reservation() -> None: + settings: Final = ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "dataplanes": DATAPLANES, + "audit_queue_size": 2, + "audit_strict": True, + } + ) + transport: Final = RoutedTransport({}) + container: Final = build_container(settings, transport=transport, exchanger=no_exchange) + principal: Final = build_principal( + container.policy, + "alice", + "alice@x.test", + frozenset({EDITOR_ROLE}), + None, + frozenset({"eu"}), + "session", + ) + dataplane: Final = container.registry.get("eu") + assert dataplane is not None + + async def receive() -> Message: + return {"type": "http.request", "body": b"[", "more_body": False} + + scope: Final[Scope] = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/key/generate", + "raw_path": b"/key/generate", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + "client": ("testclient", 123), + "server": ("testserver", 80), + } + request: Final = Request(scope, receive) + + with pytest.raises(ControlPlaneException) as raised: + await _forward( + request, + "key/generate", + principal, + dataplane, + Caller(principal, "session-token"), + container.dataplane_client, + container.audit_log, + ) + + assert raised.value.body.status == 400 + assert container.audit_log.pipeline.workers[0].queue.qsize() == 1 + assert container.audit_log.pipeline.workers[0].queue.get_nowait().event.path == "/key/generate" + + +@pytest.mark.asyncio +async def test_denied_path_precedes_a_full_strict_audit_reservation() -> None: + settings: Final = ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "dataplanes": DATAPLANES, + "audit_queue_size": 2, + "audit_strict": True, + } + ) + container: Final = build_container(settings, transport=RoutedTransport({}), exchanger=no_exchange) + principal: Final = build_principal( + container.policy, + "alice", + "alice@x.test", + frozenset({EDITOR_ROLE}), + None, + frozenset({"eu"}), + "session", + ) + dataplane: Final = container.registry.get("eu") + assert dataplane is not None + container.audit_log.record("alice", None, "eu", "POST", "/occupied", "allowed", 200) + container.audit_log.record("alice", None, "eu", "POST", "/occupied", "allowed", 200) + assert container.audit_log.pipeline.workers[0].queue.full() + + async def receive() -> Message: + return {"type": "http.request", "body": b"", "more_body": False} + + scope: Final[Scope] = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/mcp/server/oauth/id/register", + "raw_path": b"/v1/mcp/server/oauth/id/register", + "query_string": b"", + "headers": [], + "client": ("testclient", 123), + "server": ("testserver", 80), + } + + with pytest.raises(ControlPlaneException) as raised: + await _forward( + Request(scope, receive), + "v1/mcp/server/oauth/id/register", + principal, + dataplane, + Caller(principal, "session-token"), + container.dataplane_client, + container.audit_log, + ) + + assert raised.value.body.status == 403 + assert [(event.path, event.decision) for event in container.audit_log.recent(1)] == [ + ("/v1/mcp/server/oauth/id/register", "denied") + ] diff --git a/tests/unit/proxy/control_plane/test_app.py b/tests/unit/proxy/control_plane/test_app.py index 89eef97d441..80227373ac2 100644 --- a/tests/unit/proxy/control_plane/test_app.py +++ b/tests/unit/proxy/control_plane/test_app.py @@ -5,9 +5,11 @@ from typing import Final import jwt import pytest from fastapi.testclient import TestClient +from opentelemetry.sdk.trace import SpanProcessor, TracerProvider from pydantic import SecretStr from litellm.proxy.control_plane.app import LicenseRequired, create_app +from litellm.proxy.control_plane.audit.pipeline import Batch, Delivered, DeliveryResult from litellm.proxy.control_plane.auth.dependencies import get_token_verifier from litellm.proxy.control_plane.auth.jwt import PyJwkClientLookup, SessionTokenCodec, build_token_verifier from litellm.proxy.control_plane.auth.login import LOCKOUT_FAILURES, LoginService, LoginThrottle @@ -29,6 +31,24 @@ class Unlicensed: return False +class BrokenCloseSink: + name: Final = "broken" + + async def deliver(self, batch: Batch) -> DeliveryResult: + return Delivered(len(batch)) + + async def close(self) -> None: + raise RuntimeError("sink refused to close") + + +class ShutdownCounter(SpanProcessor): # mutable-ok: counts provider shutdowns for assertions + def __init__(self) -> None: + self.shutdowns = 0 + + def shutdown(self) -> None: + self.shutdowns += 1 + + def _settings(**overrides: object) -> ControlPlaneSettings: return ControlPlaneSettings.model_validate( {"LITELLM_MASTER_KEY": MASTER_KEY, "UI_USERNAME": "admin", "UI_PASSWORD": "pw", **overrides} @@ -116,6 +136,17 @@ def test_garbage_bearer_is_401() -> None: assert response.json()["detail"] == "Invalid bearer credentials." +def test_missing_and_non_bearer_credentials_are_audited_as_authentication_failures() -> None: + client: Final = _client() + assert client.get("/api/v1/auth/me").status_code == 401 + assert client.get("/api/v1/auth/me", headers={"Authorization": "Basic dTpw"}).status_code == 401 + records: Final = client.get("/api/v1/audit/ocsf", headers=_bearer(MASTER_KEY)).json()["items"] + assert [(record["class_uid"], record["actor"]["user"]["uid"], record["status_detail"]) for record in records] == [ + (3002, "anonymous", "missing bearer"), + (3002, "anonymous", "missing bearer"), + ] + + def test_master_key_is_a_fleet_wide_cp_admin() -> None: body: Final = _client().get("/api/v1/auth/me", headers=_bearer(MASTER_KEY)).json() assert body["roles"] == ["admin"] @@ -257,3 +288,13 @@ def test_a_role_granting_an_unknown_scope_fails_startup() -> None: def test_a_role_inheriting_an_unknown_role_fails_startup() -> None: with pytest.raises(InvalidRolePolicy, match="inherits unknown role 'ghost'"): build_container(_settings(roles={"x": {"inherits": ["ghost"]}})) + + +def test_tracing_shuts_down_even_when_an_audit_sink_fails_to_close() -> None: + counter: Final = ShutdownCounter() + provider: Final = TracerProvider() + provider.add_span_processor(counter) + container: Final = build_container(_settings(), tracer_provider=provider, audit_sinks=(BrokenCloseSink(),)) + with TestClient(create_app(container, Licensed())): + pass + assert counter.shutdowns == 1 diff --git a/tests/unit/proxy/control_plane/test_settings.py b/tests/unit/proxy/control_plane/test_settings.py index b963d1409b2..e4318dce632 100644 --- a/tests/unit/proxy/control_plane/test_settings.py +++ b/tests/unit/proxy/control_plane/test_settings.py @@ -3,6 +3,7 @@ from typing import Final import pytest from pydantic import ValidationError +from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID from litellm.proxy.control_plane.settings import ControlPlaneSettings, load_settings from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH @@ -79,3 +80,37 @@ def test_usage_window_maximum_is_configurable() -> None: {"LITELLM_MASTER_KEY": MASTER_KEY, "usage_max_window_days": 12} ) assert settings.usage_max_window_days == 12 + + +def test_dataplane_id_cannot_be_reserved_for_the_control_plane() -> None: + with pytest.raises(ValidationError, match="dataplane id 'control-plane' is reserved"): + ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "dataplanes": {CONTROL_PLANE_ID: {"url": "https://control.test"}}, + } + ) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("audit_capacity", 0), + ("audit_queue_size", 1), + ("audit_batch_size", 0), + ("audit_drain_seconds", 0), + ], +) +def test_audit_capacity_and_delivery_limits_must_be_positive(field: str, value: int) -> None: + with pytest.raises(ValidationError): + ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY, field: value}) + + +def test_audit_queue_size_allows_two_slots() -> None: + settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY, "audit_queue_size": 2}) + assert settings.audit_queue_size == 2 + + +def test_strict_audit_requires_at_least_one_sink() -> None: + with pytest.raises(ValidationError, match="requires at least one audit sink"): + ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY, "audit_strict": True, "audit_sinks": ()})