diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 92230fc8892..7fd66e3325e 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -80,7 +80,7 @@ jobs: - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml - name: Generate Prisma client if: steps.changes.outputs.decision != 'skip' diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 54a8e53d7a3..a69e50b5753 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -5,10 +5,24 @@ on: branches: - main - litellm_internal_staging + paths: + - "litellm/**" + - "tests/benchmarks/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/codspeed.yml" + - ".github/actions/setup-uv-with-retries/**" pull_request: branches: - main - litellm_internal_staging + paths: + - "litellm/**" + - "tests/benchmarks/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/codspeed.yml" + - ".github/actions/setup-uv-with-retries/**" # Allow CodSpeed to trigger backtest performance analysis # in order to generate initial data workflow_dispatch: diff --git a/.github/workflows/mutation-test.yml b/.github/workflows/mutation-test.yml index 6684952b998..da4fe073a6a 100644 --- a/.github/workflows/mutation-test.yml +++ b/.github/workflows/mutation-test.yml @@ -55,7 +55,7 @@ jobs: - name: Install dependencies run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml - name: Generate Prisma client env: diff --git a/Dockerfile b/Dockerfile index 9977ebb82d7..a127cdabd59 100644 --- a/Dockerfile +++ b/Dockerfile @@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 # Copy full source tree @@ -84,6 +85,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 34c9c606991..9ee076ce825 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -62,6 +62,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 # Copy full source tree @@ -82,6 +83,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 8e05f312ba0..946b4de6f5e 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -68,6 +68,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 # Copy full source tree @@ -94,6 +95,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 \ --no-sources-package litellm-proxy-extras; \ else \ @@ -102,6 +104,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3; \ fi diff --git a/gateway/Dockerfile b/gateway/Dockerfile index da2f2c9c1e0..4b000912393 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra bedrock-realtime \ --python python3 # Stage 2 — copy source and install the project + workspace members. @@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra bedrock-realtime \ --python python3 RUN mkdir -p /home/nonroot && \ diff --git a/litellm/__init__.py b/litellm/__init__.py index 55821012df9..3f8c742c5a2 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -211,6 +211,9 @@ filter_invalid_headers: Optional[bool] = False add_user_information_to_llm_headers: Optional[bool] = ( None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers ) +overwrite_user_with_key_hash: bool = ( + False # force the outgoing `user` param to the hashed api key, so providers see a stable, tamper-proof id +) store_audit_logs = False # Enterprise feature, allow users to see audit logs skip_system_message_in_guardrail: bool = False skip_tool_message_in_guardrail: bool = False diff --git a/litellm/constants.py b/litellm/constants.py index 9f60c635249..b9b9c0ba604 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1145,6 +1145,7 @@ BEDROCK_CONVERSE_MODELS = [ "anthropic.claude-sonnet-4-5-20250929-v1:0", "anthropic.claude-fable-5", "anthropic.claude-sonnet-5", + "anthropic.claude-opus-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6-v1:0", diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a40a8e1389c..96aed20529f 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2191,6 +2191,13 @@ def batch_cost_calculator( return total_prompt_cost, total_completion_cost +def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> List[str]: + field_names = list(type(prompt_tokens_details).model_fields) + if getattr(prompt_tokens_details, "cache_write_tokens", None) is None: + return field_names + return [attr for attr in field_names if attr != "cache_creation_tokens"] + + class BaseTokenUsageProcessor: @staticmethod def combine_usage_objects(usage_objects: List[Usage]) -> Usage: @@ -2225,7 +2232,7 @@ class BaseTokenUsageProcessor: # Check what keys exist in the model's prompt_tokens_details # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings - for attr in type(usage.prompt_tokens_details).model_fields: + for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): if ( hasattr(usage.prompt_tokens_details, attr) and not attr.startswith("_") diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index cf9dafcb222..f639ad49d5e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,3 +1,4 @@ +import contextvars import hashlib import os import secrets @@ -64,6 +65,10 @@ from litellm.exceptions import ( # proxy's metadata sanitizer. _PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16) +_guardrail_self_recorded: contextvars.ContextVar[bool] = contextvars.ContextVar( + "litellm_guardrail_self_recorded", default=False +) + def _strict_guardrail_modes_enabled() -> bool: """Whether guardrail-mode validation raises (default) or logs a warning. @@ -117,6 +122,7 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Optional[str] = None, sensitive_data_route_to_model: Optional[str] = None, sticky_session_routing: bool = True, + run_in_parallel: bool = False, only_scan_new_messages: bool = False, **kwargs, ): @@ -136,6 +142,9 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Action when sensitive data is detected. 'block' (default) or 'route' sensitive_data_route_to_model: Model to route to when on_sensitive_data='route' sticky_session_routing: When True, all subsequent requests in the session use the same model + run_in_parallel: When True, this pre_call or post_call guardrail runs concurrently with + other opted-in guardrails of the same hook. Only safe for block-only guardrails that + do not mutate the request or response. """ self.guardrail_name = guardrail_name self.supported_event_hooks = supported_event_hooks @@ -150,6 +159,7 @@ class CustomGuardrail(CustomLogger): self.on_sensitive_data: Optional[str] = on_sensitive_data self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model self.sticky_session_routing: bool = sticky_session_routing + self.run_in_parallel: bool = run_in_parallel self.only_scan_new_messages: bool = only_scan_new_messages if supported_event_hooks: @@ -956,6 +966,8 @@ class CustomGuardrail(CustomLogger): request_data["metadata"] = {} _append_guardrail_info(request_data["metadata"]) + _guardrail_self_recorded.set(True) + # Emit the otel guardrail span here, where every guardrail execution lands, # rather than relying on a post-call hook that does not fire on every path # (e.g. a pass-through request that passes its guardrails). @@ -1238,8 +1250,12 @@ def log_guardrail_information(func): (structured detections, tracing detail) than this decorator's "allow"/"mask"/raw-response default. To avoid double-recording in that case (which would emit two spans, two Datadog records, two spend-log - entries, etc.), snapshot the entry count before invocation: if the - wrapped function already appended its own entry, skip the auto-record. + entries, etc.), a context-local flag records whether the wrapped function + appended its own entry; if so, the auto-record is skipped. The flag is a + ``ContextVar`` rather than a count of entries in the shared ``request_data`` + so it stays correct when guardrails run concurrently (asyncio copies the + context into each gathered task): counting shared entries would let one + guardrail's append hide another guardrail's missing record. """ import functools import inspect @@ -1259,16 +1275,6 @@ def log_guardrail_information(func): return GuardrailEventHooks.post_call return None - def _count_recorded_guardrail_entries(request_data: dict) -> int: - total = 0 - for container_key in ("metadata", "litellm_metadata"): - container = request_data.get(container_key) - if isinstance(container, dict): - entries = container.get("standard_logging_guardrail_information") - if isinstance(entries, list): - total += len(entries) - return total - @functools.wraps(func) async def async_wrapper(*args, **kwargs): start_time = datetime.now() # Move start_time inside the wrapper @@ -1282,10 +1288,10 @@ def log_guardrail_information(func): original_inputs = kwargs.get("inputs") logging_obj = kwargs.get("logging_obj") - entries_before = _count_recorded_guardrail_entries(request_data) + self_recorded_token = _guardrail_self_recorded.set(False) try: response = await func(*args, **kwargs) - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): return response return self._process_response( response=response, @@ -1297,7 +1303,7 @@ def log_guardrail_information(func): original_inputs=original_inputs, ) except Exception as e: - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): raise return self._process_error( e=e, @@ -1308,6 +1314,7 @@ def log_guardrail_information(func): event_type=event_type, ) finally: + _guardrail_self_recorded.reset(self_recorded_token) _sync_guardrail_info_to_logging_obj(request_data, logging_obj) @functools.wraps(func) @@ -1323,10 +1330,10 @@ def log_guardrail_information(func): original_inputs = kwargs.get("inputs") logging_obj = kwargs.get("logging_obj") - entries_before = _count_recorded_guardrail_entries(request_data) + self_recorded_token = _guardrail_self_recorded.set(False) try: response = func(*args, **kwargs) - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): return response return self._process_response( response=response, @@ -1336,7 +1343,7 @@ def log_guardrail_information(func): original_inputs=original_inputs, ) except Exception as e: - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): raise return self._process_error( e=e, @@ -1345,6 +1352,7 @@ def log_guardrail_information(func): event_type=event_type, ) finally: + _guardrail_self_recorded.reset(self_recorded_token) _sync_guardrail_info_to_logging_obj(request_data, logging_obj) @functools.wraps(func) diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index c93f95ec97d..fa41070a8de 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -18,12 +18,14 @@ before the LLM call even starts), so a guardrail is a sibling of the LLM call, not a child of it. The emitter parents every span to the ambient OTel context (the active server span), which matches this. -MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) are intentionally NOT in this -tree. Per the OTel GenAI MCP semconv, MCP and the HTTP transport are independent -contexts, so an MCP span parents to the trace context the client propagated in -``params._meta`` (or starts its own root when none is propagated) and records the -``PROXY_REQUEST`` transport span as a span *link*, never a parent. The registry -encodes this as ``parent=None, links=PROXY_REQUEST``. +MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) have two shapes, chosen at emit +time by :func:`resolve_mcp_span_context`. When the client propagates trace context +in ``params._meta`` MCP and the HTTP transport are independent contexts per the +OTel GenAI MCP semconv, so the span parents to that propagated context and records +the ``PROXY_REQUEST`` transport span as a span *link*, never a parent — the shape +this registry's ``parent=None, links=PROXY_REQUEST`` entry encodes. When nothing is +propagated (the common case) the span nests under the transport span of the request +carrying that message, so the tool call stays in one trace. Not every service call becomes a span — :func:`span_role_for_service` decides: @@ -89,12 +91,13 @@ class SpanSpec: SPAN_REGISTRY: dict[SpanRole, SpanSpec] = { SpanRole.PROXY_REQUEST: SpanSpec(SpanRole.PROXY_REQUEST, LiteLLMSpanKind.SERVER, parent=None), SpanRole.LLM_CALL: SpanSpec(SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST), - # MCP and the HTTP transport are independent contexts (OTel GenAI MCP semconv), - # so an MCP span does not nest under the transport span. The proxy is an MCP - # client to the upstream server, so it's a CLIENT span; it parents to the trace - # context the client propagated in ``params._meta`` (or starts its own root when - # none is propagated) and records the PROXY_REQUEST transport span as a span - # *link*, never a parent — hence ``parent=None, links=PROXY_REQUEST``. + # The proxy is an MCP client to the upstream server, so MCP spans are CLIENT + # spans. With trace context propagated in ``params._meta``, MCP and the HTTP + # transport are independent contexts (OTel GenAI MCP semconv): the span parents + # to the propagated context and records the PROXY_REQUEST transport span as a + # span *link*, never a parent — the shape ``parent=None, links=PROXY_REQUEST`` + # encodes. With nothing propagated, ``resolve_mcp_span_context`` nests the span + # under that message's transport span instead, keeping the call in one trace. SpanRole.MCP_TOOL_CALL: SpanSpec( SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=None, links=SpanRole.PROXY_REQUEST ), diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 8acac112c3d..939559347b1 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -5,7 +5,14 @@ from typing import Mapping from opentelemetry import baggage from opentelemetry.context import Context, get_current -from opentelemetry.trace import Link, Span, get_current_span, set_span_in_context +from opentelemetry.trace import ( + Link, + NonRecordingSpan, + Span, + SpanContext, + get_current_span, + set_span_in_context, +) from opentelemetry.trace.propagation.tracecontext import ( TraceContextTextMapPropagator, ) @@ -72,6 +79,62 @@ def reset_mcp_message_trace_carrier(token: "Token[Mapping[str, str] | None]") -> _mcp_message_trace_carrier.reset(token) +# The transport span of the HTTP request carrying the CURRENT MCP message, as a +# plain ``SpanContext`` so it can cross a task boundary. +# +# ``_request_root_span`` above cannot be used for MCP: a *stateful* streamable-HTTP +# session runs every message on the single task spawned by that session's +# ``initialize`` POST, so the ContextVar the ASGI request task writes at auth time +# is frozen at ``initialize`` there and never sees the later ``tools/call`` POSTs. +# Reading it from the message handler would parent every tool call in the session +# to the first request's (already ended) server span. The gateway instead resolves +# the current message's transport span on the request task and hands it over the +# same way it hands over per-request auth, and the handler publishes it here for +# the span emitter to pick up. +_mcp_message_transport_span_context: "ContextVar[SpanContext | None]" = ContextVar( + "litellm_otel_mcp_message_transport_span_context", default=None +) + + +def set_mcp_message_transport_span_context( + span_context: "SpanContext | None", +) -> "Token[SpanContext | None]": + """Publish the transport span of the request carrying the current MCP message. + + Returns the reset token; the caller must reset it once the message is handled + so the transport never leaks to the next message on the same session task. + """ + return _mcp_message_transport_span_context.set(span_context) + + +def reset_mcp_message_transport_span_context(token: "Token[SpanContext | None]") -> None: + _mcp_message_transport_span_context.reset(token) + + +def request_root_span_context() -> "SpanContext | None": + """The anchored request root span's context, safe to hand to another task. + + A ``SpanContext`` is an immutable value, unlike the live ``Span``, so passing it + across the MCP session-task boundary cannot keep a finished span alive or invite + writes to it from the wrong request. + """ + span = request_root_span() + return span.get_span_context() if span is not None else None + + +def _mcp_transport_span_context() -> "SpanContext | None": + """The transport span an MCP message span should attach to. + + Prefers the transport the gateway published for this specific message; falls + back to the ambient request anchor for paths that emit an MCP span on the + request task itself (the REST MCP endpoints, the SDK). + """ + published = _mcp_message_transport_span_context.get() + if published is not None and published.is_valid: + return published + return request_root_span_context() + + def set_request_baggage(values: Mapping[str, str], context: Context | None = None) -> Context: """Return a context with ``values`` written into Baggage.""" ctx = context @@ -132,33 +195,44 @@ def resolve_request_span_context() -> Context: def resolve_mcp_span_context( carrier: "Mapping[str, str] | None" = None, ) -> "tuple[Context, tuple[Link, ...]]": - """Parent context + links for an MCP message span, per the OTel GenAI MCP semconv. + """Parent context + links for an MCP message span. - MCP and the underlying transport (HTTP) are independent lifecycles — one - streamable-HTTP session multiplexes many messages, so nesting the message span - under the HTTP/session span is wrong (it renders the message at the session's - start, skewed by however long the session has been open). Instead: + When the client propagates W3C trace context in the request's ``params._meta`` + (SEP-414), MCP and the underlying transport are independent lifecycles — one + streamable-HTTP session multiplexes many messages, and the client's own span is + the truthful parent. So, per the OTel GenAI MCP semconv: - * parent to the trace context the client propagated in the request's - ``params._meta`` (a *remote* parent), and - * record the transport/session span as a *link*, never the parent. + * parent to the trace context the client propagated (a *remote* parent), and + * record the transport span as a *link*, never the parent. + + Almost no client implements SEP-414 yet, so in practice nothing is propagated. + Rooting the span there splits a single tool call into two disconnected traces + joined only by a link, which is how it surfaces in APM: the ``POST`` transaction + and the ``tools/call`` span share no trace. With no remote parent to honor, + parent to the transport span of the request carrying this message instead, so + the call stays in one trace; no link is added since the transport is now the + real parent. The transport comes from :func:`_mcp_transport_span_context`, which + is the *current message's* POST rather than whatever request happened to open + the session, so a long-lived session does not glue every message under its + first request. With neither a remote parent nor a transport the returned context + carries no span and the span legitimately starts its own root trace. Only trace context (``traceparent``/``tracestate``) is extracted, never the client's W3C Baggage: ``params._meta`` is caller-controlled, and the otel baggage processor stamps allowlisted baggage keys (``litellm.team.id``, ``litellm.metadata.*``, ...) onto the span as attributes, so honoring remote - baggage would let a client spoof a span's identity attribution. - - With no propagated context the returned context carries no span, so the span - starts its own root trace (still linked to the transport). The base context is - explicitly empty so an absent ``traceparent`` can never fall through to the - ambient (stale session) span. + baggage would let a client spoof a span's identity attribution. The base context + for extraction is explicitly empty so an absent or malformed ``traceparent`` can + never fall through to the ambient (stale session) span. """ source = carrier if carrier is not None else _mcp_message_trace_carrier.get() parent = _PROPAGATOR.extract(dict(source or {}), context=Context()) - transport = request_root_span() - links = (Link(transport.get_span_context()),) if transport is not None else () - return parent, links + transport = _mcp_transport_span_context() + if is_recordable_span(get_current_span(parent)): + return parent, (Link(transport),) if transport is not None else () + if transport is not None: + return context_from_span(NonRecordingSpan(transport)), () + return parent, () def is_recordable_span(obj: object) -> bool: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3b3c6a6ce29..c9e70b7db73 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1612,6 +1612,35 @@ class Logging(LiteLLMLoggingBaseClass): **kwargs, ) + async def dispatch_failure_handlers( + self, + exception: Exception, + traceback_exception: str, + prefer_async_handlers: bool = False, + ) -> None: + """Route failure logging to async and/or sync handlers for this request. + + Mirrors ``dispatch_success_handlers``: the sync ``failure_handler`` never runs + concurrently with ``async_failure_handler`` on the shared logging object, so the + two paths cannot mutate it at the same time. ``prefer_async_handlers`` only + bypasses the sync-SDK-only shortcut (e.g. ``async for`` on a stream from + ``completion()``); legacy string callbacks still run via + ``executor.submit(failure_handler)`` when configured. + """ + litellm_params = self.model_call_details.get("litellm_params", {}) or {} + sync_sdk = self._is_sync_litellm_request(litellm_params) + passthrough = self.call_type == CallTypes.pass_through.value + if sync_sdk and not prefer_async_handlers and not passthrough: + self.failure_handler(exception, traceback_exception) + return + + await self.async_failure_handler(exception, traceback_exception) + + if not self._should_run_sync_failure_callbacks_for_async_calls(): + return + + executor.submit(self.failure_handler, exception, traceback_exception) + def should_run_logging( self, event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], @@ -3076,6 +3105,24 @@ class Logging(LiteLLMLoggingBaseClass): _filtered_success_callbacks = self._remove_internal_litellm_callbacks(_filtered_success_callbacks) return len(_filtered_success_callbacks) > 0 + def _should_run_sync_failure_callbacks_for_async_calls(self) -> bool: + """ + Returns: + - bool: True if sync failure callbacks should be run for async calls. eg. `langfuse`, `s3` + + Mirrors ``_should_run_sync_callbacks_for_async_calls`` but reads the failure + callback lists. Gating the legacy sync ``failure_handler`` on the success lists + would drop sync failure callbacks for any caller that configures only failure + callbacks, so streaming errors would be logged nowhere. + """ + _combined_sync_callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_failure_callbacks, + global_callbacks=litellm.failure_callback, + ) + _filtered_failure_callbacks = self._remove_internal_custom_logger_callbacks(_combined_sync_callbacks) + _filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks) + return len(_filtered_failure_callbacks) > 0 + def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: if dynamic_success_callbacks is None: return list(global_callbacks) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 33bf546c239..85ed0665ebf 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -457,7 +457,8 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: cache_creation_tokens = ( cast( Optional[int], - getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0), + getattr(usage.prompt_tokens_details, "cache_write_tokens", 0) + or getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0), ) or 0 ) @@ -906,10 +907,6 @@ def get_token_type_cost_breakdown( cache_read_tokens = prompt_tokens_details["cache_hit_tokens"] cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"] cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"] - # Some OpenAI-compatible providers (e.g. kimi-k2) report cache-write tokens - # under `cache_write_tokens`; mirror the total-cost normalization path. - if not cache_creation_tokens: - cache_creation_tokens = _coerce_token_count(getattr(usage.prompt_tokens_details, "cache_write_tokens", 0)) # Fall back to the private top-level counters the Usage constructor mirrors cache # tokens onto, so providers/callers that bypass prompt_tokens_details are covered. if not cache_read_tokens: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f518cbaadea..60dbf7c644a 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2008,12 +2008,9 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - threading.Thread( - target=self.logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming - asyncio.create_task(self.logging_obj.async_failure_handler(e, traceback_exception)) + asyncio.create_task( + self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) + ) self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: if self.received_finish_reason is None: @@ -2122,13 +2119,8 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - threading.Thread( - target=self.logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming asyncio.create_task( - self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore + self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) ) self._handle_stream_fallback_error(e) diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 6bbe8f75701..df903ba7ef0 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -123,7 +123,8 @@ class VertexAIBatchTransformation: Gets the output file id from the Vertex AI Batch response """ - output_file_id: str = response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "") + output_info = response.get("outputInfo") or OutputInfo() + output_file_id: str = output_info.get("gcsOutputDirectory", "") if output_file_id: output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" if output_file_id and output_file_id != "/predictions.jsonl": diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 624190a0b61..8dd0dc19b81 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3333,27 +3333,37 @@ class ModelResponseIterator: return self.chunk_parser(chunk=json_chunk) - def handle_accumulated_json_chunk(self, chunk: str) -> Optional["ModelResponseStream"]: - chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" - message = chunk.replace("\n\n", "") + def handle_accumulated_json_chunk(self, chunk: str, is_final: bool = False) -> Optional["ModelResponseStream"]: + message = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" + self.accumulated_json = (self.accumulated_json + message.replace("\n\n", "")).strip() - self.accumulated_json += message - - # json.loads on the whole buffer after every fragment is O(n^2) and - # holds the GIL, freezing the event loop for seconds on large responses - # (https://github.com/BerriAI/litellm/issues/26181). A complete Gemini - # chunk is a JSON object/array, so only attempt the parse once the - # buffer's last non-whitespace byte can close one. - stripped = self.accumulated_json.rstrip() - if not stripped or stripped[-1] not in "}]": + # Mid-stream, defer parsing until the buffer's last byte can close a value: + # attempting a parse after every fragment of one large object is O(n^2) and + # holds the GIL, freezing the event loop. At end of stream (is_final) no more + # data is coming, so drain whatever complete values remain regardless of the + # trailing byte, otherwise a complete leading value sitting behind a truncated + # trailing one would be silently dropped. + if not is_final and (not self.accumulated_json or self.accumulated_json[-1] not in "}]"): return None - try: - _data = json.loads(self.accumulated_json) - self.accumulated_json = "" # reset after successful parsing - return self.chunk_parser(chunk=_data) - except json.JSONDecodeError: - return None + # Peel one complete JSON value from the front of the buffer and keep the + # unconsumed tail. Running json.loads over the whole buffer would fail + # forever once it held more than one concatenated value ("Extra data") while + # never resetting the buffer, so the buffer grew without bound and pinned the + # core. raw_decode reports where the value ended, so concatenated values drain + # one call at a time. A leading non-dict value (never emitted by Gemini in + # practice) is consumed and skipped so it cannot block the dict values behind it. + decoder = json.JSONDecoder() + while self.accumulated_json: + try: + raw_value = decoder.raw_decode(self.accumulated_json) + except json.JSONDecodeError: + return None + decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode -> tuple[Any,int] + self.accumulated_json = self.accumulated_json[end_index:].strip() + if isinstance(decoded, dict): + return self.chunk_parser(chunk=decoded) + return None def _common_chunk_parsing_logic(self, chunk: str) -> Optional["ModelResponseStream"]: try: @@ -3378,7 +3388,9 @@ class ModelResponseIterator: chunk = self.response_iterator.__next__() except StopIteration: if self.chunk_type == "accumulated_json" and self.accumulated_json: - return self.handle_accumulated_json_chunk(chunk="") + result = self.handle_accumulated_json_chunk(chunk="", is_final=True) + if result is not None: + return result raise StopIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") @@ -3400,7 +3412,9 @@ class ModelResponseIterator: chunk = await self.async_response_iterator.__anext__() except StopAsyncIteration: if self.chunk_type == "accumulated_json" and self.accumulated_json: - return self.handle_accumulated_json_chunk(chunk="") + result = self.handle_accumulated_json_chunk(chunk="", is_final=True) + if result is not None: + return result raise StopAsyncIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d3917886060..d43eda39b1f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1502,6 +1502,222 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024 }, + "anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "global.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "us.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "eu.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "au.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "jp.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, @@ -2756,6 +2972,38 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true }, + "azure_ai/claude-opus-5": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "azure_ai/claude-opus-4-8": { "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, @@ -11846,6 +12094,44 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512 }, + "claude-opus-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1, + "fast": 2.0 + }, + "supports_output_config": true, + "supports_speed": true, + "prompt_cache_min_tokens": 512 + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -36896,6 +37182,70 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true }, + "vertex_ai/claude-opus-5": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, + "vertex_ai/claude-opus-5@default": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "vertex_ai/claude-opus-4-8": { "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, diff --git a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py index 7122c64ec64..f7bc14575c7 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -1,9 +1,12 @@ -from typing import Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Optional from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from litellm.proxy._types import UserAPIKeyAuth +if TYPE_CHECKING: + from opentelemetry.trace import SpanContext + class MCPAuthenticatedUser(AuthenticatedUser): """ @@ -16,6 +19,8 @@ class MCPAuthenticatedUser(AuthenticatedUser): 4. Server-specific authentication headers 5. OAuth2 headers 6. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin. + 7. Transport span context - the tracing span of the HTTP request carrying the current + message, which a stateful session's message handler cannot read from its own task. """ def __init__( @@ -28,6 +33,7 @@ class MCPAuthenticatedUser(AuthenticatedUser): mcp_protocol_version: Optional[str] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ): self.user_api_key_auth = user_api_key_auth self.mcp_auth_header = mcp_auth_header @@ -37,3 +43,4 @@ class MCPAuthenticatedUser(AuthenticatedUser): self.oauth2_headers = oauth2_headers self.raw_headers = raw_headers self.client_ip = client_ip + self.transport_span_context = transport_span_context diff --git a/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py b/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py deleted file mode 100644 index cd41dd648ee..00000000000 --- a/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers. - -Exchanges a user's incoming JWT (subject_token) for a scoped access token -at an IDP's token exchange endpoint. The exchanged token is then used to -authenticate requests to the upstream MCP server. - -See: https://datatracker.ietf.org/doc/html/rfc8693 -""" - -import asyncio -import hashlib -import weakref -from typing import TYPE_CHECKING, Dict, Tuple - -import httpx - -from litellm._logging import verbose_logger -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import ( - MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, - MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, - MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, -) -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( - build_token_endpoint_client_auth, -) -from litellm.types.llms.custom_http import httpxSpecialProvider -from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE - -if TYPE_CHECKING: - from litellm.types.mcp_server.mcp_server_manager import MCPServer - -# RFC 8693 grant type constant -TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" - - -class TokenExchangeHandler: - """Handles OAuth 2.0 Token Exchange (RFC 8693) for MCP servers. - - Caches exchanged tokens keyed by ``hash(subject_token + server_id)`` so - repeated calls with the same user token skip the IDP round-trip. - """ - - def __init__(self) -> None: - self._cache = InMemoryCache( - max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, - default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, - ) - # WeakValueDictionary so locks are GC'd once no coroutine holds a reference, - # preventing unbounded growth with many rotating user tokens. - self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() - - def _get_lock(self, cache_key: str) -> asyncio.Lock: - lock = self._locks.get(cache_key) - if lock is None: - lock = asyncio.Lock() - self._locks[cache_key] = lock - return lock - - @staticmethod - def _cache_key(subject_token: str, server_id: str) -> str: - raw = f"{subject_token}:{server_id}" - return hashlib.sha256(raw.encode()).hexdigest() - - async def exchange_token( - self, - subject_token: str, - server: "MCPServer", - ) -> str: - """Exchange *subject_token* for a scoped access token. - - Returns the exchanged ``access_token`` string (suitable for a - ``Bearer`` header). - - Raises ``ValueError`` on configuration or IDP errors. - """ - cache_key = self._cache_key(subject_token, server.server_id) - - # Fast path - cached = self._cache.get_cache(cache_key) - if cached is not None: - return cached - - # Slow path — one exchange at a time per (user, server) pair - async with self._get_lock(cache_key): - cached = self._cache.get_cache(cache_key) - if cached is not None: - return cached - - token, ttl = await self._do_exchange(subject_token, server) - self._cache.set_cache(cache_key, token, ttl=ttl) - return token - - async def _do_exchange( - self, - subject_token: str, - server: "MCPServer", - ) -> Tuple[str, int]: - """POST to the token exchange endpoint with RFC 8693 parameters. - - Returns ``(access_token, ttl_seconds)``. - """ - endpoint = server.token_exchange_endpoint or server.token_url - if not endpoint: - raise ValueError( - f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange " - f"but no token_exchange_endpoint or token_url configured" - ) - if not server.client_id or not server.client_secret: - raise ValueError( - f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange " - f"but missing client_id or client_secret" - ) - - client_auth = build_token_endpoint_client_auth( - auth_method=server.token_endpoint_auth_method, - client_id=server.client_id, - client_secret=server.client_secret, - ) - data: Dict[str, str] = { - "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, - "subject_token": subject_token, - "subject_token_type": server.subject_token_type or DEFAULT_SUBJECT_TOKEN_TYPE, - **client_auth.body, - } - if server.audience: - data["audience"] = server.audience - if server.scopes: - data["scope"] = " ".join(server.scopes) - - verbose_logger.debug( - "Exchanging token for MCP server %s at %s (audience=%s)", - server.server_id, - endpoint, - server.audience, - ) - - client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - post_kwargs = {"data": data, **({"headers": client_auth.headers} if client_auth.headers else {})} - try: - response = await client.post(endpoint, **post_kwargs) - response.raise_for_status() - except httpx.HTTPStatusError as exc: - verbose_logger.debug( - "Token exchange IDP error for MCP server %s (status %d)", - server.server_id, - exc.response.status_code, - ) - raise ValueError( - f"Token exchange for MCP server '{server.server_id}' failed with status {exc.response.status_code}" - ) from exc - - body = response.json() - if not isinstance(body, dict): - raise ValueError( - f"Token exchange response for MCP server '{server.server_id}' " - f"returned non-object JSON (got {type(body).__name__})" - ) - - access_token = body.get("access_token") - if not access_token: - raise ValueError(f"Token exchange response for MCP server '{server.server_id}' missing 'access_token'") - - raw_expires_in = body.get("expires_in") - try: - expires_in = int(raw_expires_in) if raw_expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL - except (TypeError, ValueError): - expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL - - ttl = max( - expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, - MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, - ) - - verbose_logger.info( - "Token exchange succeeded for MCP server %s (expires in %ds)", - server.server_id, - expires_in, - ) - return access_token, ttl - - def invalidate(self, subject_token: str, server_id: str) -> None: - """Remove a cached exchanged token (e.g. after a 401).""" - cache_key = self._cache_key(subject_token, server_id) - self._cache.delete_cache(cache_key) - - -# Module-level singleton -mcp_token_exchange_handler = TokenExchangeHandler() diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b442ea5de70..0ee74960293 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3086,9 +3086,7 @@ class MCPServerManager: ) ): spec = None - auth_value = ( - await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token) if spec is None else None - ) + auth_value = await resolve_mcp_auth(server, mcp_auth_header) if spec is None else None # Create sampling and elicitation callbacks for this client sampling_cb = _create_sampling_callback(user_api_key_auth=user_api_key_auth) if server.allow_sampling else None diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 43fe3999291..a6acaf8e1d6 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -26,7 +26,6 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) -from litellm.proxy._experimental.mcp_server.auth import token_exchange from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( build_token_endpoint_client_auth, ) @@ -58,17 +57,12 @@ class MCPOAuth2TokenCache(InMemoryCache): def _has_client_credentials_config(server: "MCPServer") -> bool: return bool(server.client_id and server.client_secret and server.token_url) - async def async_get_token( - self, - server: "MCPServer", - *, - require_client_credentials_flow: bool = True, - ) -> Optional[str]: + async def async_get_token(self, server: "MCPServer") -> Optional[str]: """Return a valid access token, fetching or refreshing as needed. Returns ``None`` when the server lacks client credentials config. """ - if require_client_credentials_flow and not server.has_client_credentials: + if not server.has_client_credentials: return None if not self._has_client_credentials_config(server): return None @@ -278,36 +272,16 @@ mcp_per_user_token_cache = MCPPerUserTokenCache() async def resolve_mcp_auth( server: "MCPServer", mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, - subject_token: Optional[str] = None, ) -> Optional[Union[str, Dict[str, str]]]: """Resolve the auth value for an MCP server. Priority: 1. ``mcp_auth_header`` — per-request/per-user override - 2. OAuth2 Token Exchange (OBO / RFC 8693) — exchange user token for scoped token - 3. OAuth2 client_credentials token — auto-fetched and cached - 4. ``server.authentication_token`` — static token from config/DB + 2. OAuth2 client_credentials token — auto-fetched and cached + 3. ``server.authentication_token`` — static token from config/DB """ if mcp_auth_header: return mcp_auth_header - if server.has_token_exchange_config: - if subject_token: - return await token_exchange.mcp_token_exchange_handler.exchange_token(subject_token, server) - # No subject_token — fall back to client_credentials using the same client - # credentials and token_url so M2M scenarios still work. - if server.client_id and server.client_secret and server.token_url: - return await mcp_oauth2_token_cache.async_get_token( - server, - require_client_credentials_flow=False, - ) - # OBO configured but no subject_token and missing client credentials — warn - # rather than silently proceeding unauthenticated. - verbose_logger.warning( - "MCP server '%s' is configured for token exchange (OBO) but no subject_token " - "was provided and client credentials (client_id/client_secret/token_url) are " - "incomplete. The request will proceed without authentication.", - server.server_id, - ) if server.has_client_credentials: return await mcp_oauth2_token_cache.async_get_token(server) return server.authentication_token diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 94271c54f4b..26e4176e09b 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -230,16 +230,33 @@ if MCP_AVAILABLE: return server_auth return mcp_auth_header - def _get_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]: - """Return the subset of *allowed_server_ids* whose servers use OAuth2 auth. + def _is_v1_resolved_oauth2_server(server: Optional[MCPServer]) -> bool: + """Whether this server's per-user OAuth2 token is still resolved by v1. - Used as a cheap pre-flight check to skip bulk credential fetching when no - OAuth2 servers are involved in the current request. + A server the v2 resolver owns reads its stored token from the resolver at connect + time and drops any Authorization built for it here, so the v1 lookup would be a DB + round-trip whose result is discarded. Mirrors the same guard on the protocol listing + path and in ``_resolve_oauth2_headers_for_tool_call``. + """ + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + to_server_spec, + ) + + if getattr(server, "auth_type", None) != MCPAuth.oauth2: + return False + return to_server_spec(server) is None + + def _v1_resolved_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]: + """Return the subset of *allowed_server_ids* whose per-user OAuth2 token is still + resolved by v1. + + Used as a cheap pre-flight check to skip bulk credential fetching when no such + server is involved in the current request. """ return { sid for sid in allowed_server_ids - if getattr(global_mcp_server_manager.get_mcp_server_by_id(sid), "auth_type", None) == MCPAuth.oauth2 + if _is_v1_resolved_oauth2_server(global_mcp_server_manager.get_mcp_server_by_id(sid)) } async def _get_user_oauth_extra_headers( @@ -253,11 +270,13 @@ if MCP_AVAILABLE: the MCP server the same way the admin "Add MCP / Authorize and Fetch" flow does. Returns None for non-OAuth2 servers or when no credential is stored. + A server the v2 resolver owns is skipped; see ``_is_v1_resolved_oauth2_server``. + Args: prefetched_creds: Optional dict keyed by server_id with credential payloads. When provided, avoids a per-server DB round-trip. """ - if getattr(server, "auth_type", None) != MCPAuth.oauth2: + if not _is_v1_resolved_oauth2_server(server): return None user_id = getattr(user_api_key_dict, "user_id", None) server_id = getattr(server, "server_id", None) @@ -320,38 +339,6 @@ if MCP_AVAILABLE: verbose_logger.warning(f"_prefetch_user_oauth_creds: failed to prefetch for user={user_id}: {e}") return {} - async def _get_bulk_user_oauth_headers( - user_api_key_dict: UserAPIKeyAuth, - ) -> Dict[str, Dict[str, str]]: - """ - Fetch ALL OAuth2 credentials for the current user in a single DB query and - return a mapping of server_id → {"Authorization": "Bearer "}. - - This is the batch alternative to calling _get_user_oauth_extra_headers - per-server inside a loop (N+1 DB queries). - """ - user_id = getattr(user_api_key_dict, "user_id", None) - if not user_id: - return {} - try: - from litellm.proxy._experimental.mcp_server.db import ( - list_user_oauth_credentials, - ) - from litellm.proxy.utils import get_prisma_client_or_throw - - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to use OAuth2 MCP tools." - ) - creds = await list_user_oauth_credentials(prisma_client, user_id) - return { - c["server_id"]: {"Authorization": f"Bearer {c['access_token']}"} - for c in creds - if c.get("access_token") and c.get("server_id") - } - except Exception: - verbose_logger.debug("Failed to bulk-fetch OAuth credentials", exc_info=True) - return {} - def _create_tool_response_objects(tools, server: MCPServer): """Helper function to create tool response objects. @@ -825,7 +812,7 @@ if MCP_AVAILABLE: # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. prefetched_oauth_creds = ( await _prefetch_user_oauth_creds(user_api_key_dict) - if _get_oauth2_server_ids(allowed_server_ids) + if _v1_resolved_oauth2_server_ids(allowed_server_ids) else {} ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 4fca4406a6f..483f57f9139 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -15,6 +15,7 @@ import types import uuid from datetime import datetime from typing import ( + TYPE_CHECKING, Any, AsyncIterator, Callable, @@ -107,6 +108,9 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER = 100 # arbitrarily large body just to make a routing decision. _MCP_ROUTING_PEEK_MAX_BYTES = 4096 +if TYPE_CHECKING: + from opentelemetry.trace import SpanContext + def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: """Remove a (user_id, server_id) entry from the BYOK credential cache. @@ -242,10 +246,12 @@ def _mcp_meta_trace_carrier(req_ctx: object) -> Optional[dict[str, str]]: """The W3C trace context (``traceparent``/``tracestate``) the MCP client propagated in the request's ``params._meta`` (SEP-414), or ``None``. - Per the OTel MCP semconv the MCP span parents to this propagated context rather - than to the HTTP/session transport (which is recorded as a link instead), so a - streamable-HTTP session that multiplexes many messages does not glue every - message under the session's first request. The client's W3C Baggage is + When present, per the OTel MCP semconv the MCP span parents to this propagated + context rather than to the HTTP transport (which is recorded as a link instead). + When absent, the span nests under the transport span of the request carrying + this specific message, so a streamable-HTTP session that multiplexes many + messages still does not glue every message under the session's first request; + see ``resolve_mcp_span_context``. The client's W3C Baggage is deliberately excluded: it is caller-controlled, and the otel baggage processor stamps allowlisted baggage keys (``litellm.team.id``, ``litellm.metadata.*``, ...) onto the span, so honoring remote baggage would let a client spoof a @@ -288,6 +294,56 @@ def _otel_reset_mcp_trace_carrier(token: object) -> None: return +def _otel_request_transport_span_context() -> Optional["SpanContext"]: + """The tracing span of the HTTP request being handled, as a portable value. + + Resolved on the ASGI request task, where the proxy's server span is anchored, + and carried to the MCP message handler on the authenticated-user object. A + stateful streamable-HTTP session handles every message on the task spawned by + its ``initialize`` POST, so the handler's own task cannot see later requests' + spans; this is the same reason per-request auth is carried across rather than + read from a ContextVar. Lazily imported so opentelemetry stays an optional + dependency; returns ``None`` when otel_v2 is unavailable or no request span is + anchored.""" + try: + from litellm.integrations.otel.plumbing.context import ( + request_root_span_context, + ) + + return request_root_span_context() + except ImportError: + return None + + +def _otel_set_mcp_transport_span_context(span_context: Optional["SpanContext"]) -> object: + """Publish the current message's transport span for the otel_v2 MCP span and + return a reset token, or ``None`` when otel_v2 is unavailable.""" + if span_context is None: + return None + try: + from litellm.integrations.otel.plumbing.context import ( + set_mcp_message_transport_span_context, + ) + + return set_mcp_message_transport_span_context(span_context) + except ImportError: + return None + + +def _otel_reset_mcp_transport_span_context(token: object) -> None: + """Paired with ``_otel_set_mcp_transport_span_context``.""" + if token is None: + return + try: + from litellm.integrations.otel.plumbing.context import ( + reset_mcp_message_transport_span_context, + ) + + reset_mcp_message_transport_span_context(token) + except ImportError: + return + + def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: """Map a ``ProxyException`` to an ``HTTPException`` that preserves its real status code and headers. @@ -654,6 +710,18 @@ if MCP_AVAILABLE: ############### MCP Server Routes ####################### ######################################################## + def _current_transport_span_context() -> Optional["SpanContext"]: + """The transport span of the HTTP request carrying the message being handled. + + Published by the ASGI request task onto the authenticated-user object, because + a stateful session's message handler runs on the task spawned by that session's + ``initialize`` POST and so cannot read later requests' spans from its own task. + """ + auth_user = auth_context_var.get() + if not isinstance(auth_user, MCPAuthenticatedUser): + auth_user = _recover_auth_from_session() + return auth_user.transport_span_context if auth_user is not None else None + @server.list_tools() async def handle_list_tools() -> "ListToolsResult | List[Tool]": """ @@ -670,9 +738,11 @@ if MCP_AVAILABLE: if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None + _transport_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) + _transport_token = _otel_set_mcp_transport_span_context(_current_transport_span_context()) # Get user authentication from context variable ( user_api_key_auth, @@ -728,6 +798,7 @@ if MCP_AVAILABLE: # This prevents the HTTP stream from failing and allows the client to get a response return [] finally: + _otel_reset_mcp_transport_span_context(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @@ -901,9 +972,11 @@ if MCP_AVAILABLE: if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None + _transport_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) + _transport_token = _otel_set_mcp_transport_span_context(_current_transport_span_context()) # Validate arguments ( user_api_key_auth, @@ -1042,6 +1115,7 @@ if MCP_AVAILABLE: return response finally: + _otel_reset_mcp_transport_span_context(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @@ -4197,6 +4271,7 @@ if MCP_AVAILABLE: session_id=session_id if use_stateful else None, touch_last_seen=(scope.get("method") or "").upper() != "DELETE", copy_existing_session_auth_context=is_initialize, + transport_span_context=_otel_request_transport_span_context(), ) local_send = send if use_stateful and is_initialize: @@ -4421,6 +4496,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ) -> None: auth_user.user_api_key_auth = user_api_key_auth auth_user.mcp_auth_header = mcp_auth_header @@ -4429,6 +4505,7 @@ if MCP_AVAILABLE: auth_user.oauth2_headers = oauth2_headers auth_user.raw_headers = raw_headers auth_user.client_ip = client_ip + auth_user.transport_span_context = transport_span_context def set_auth_context( user_api_key_auth: Optional[UserAPIKeyAuth], @@ -4438,6 +4515,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ) -> MCPAuthenticatedUser: """ Set the UserAPIKeyAuth in the auth context variable. @@ -4448,6 +4526,7 @@ if MCP_AVAILABLE: mcp_servers: Optional list of server names and access groups to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} client_ip: Client IP address for MCP access control + transport_span_context: Tracing span of the HTTP request carrying this message """ auth_user = MCPAuthenticatedUser( user_api_key_auth=user_api_key_auth, @@ -4457,6 +4536,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) auth_context_var.set(auth_user) return auth_user @@ -4472,6 +4552,7 @@ if MCP_AVAILABLE: session_id: Optional[str] = None, touch_last_seen: bool = True, copy_existing_session_auth_context: bool = False, + transport_span_context: Optional["SpanContext"] = None, ) -> MCPAuthenticatedUser: auth_user = _stateful_session_auth_contexts.get(session_id) if session_id else None if auth_user is not None and session_id is not None: @@ -4486,6 +4567,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) _update_auth_context( auth_user=auth_user, @@ -4496,6 +4578,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) auth_context_var.set(auth_user) return auth_user @@ -4507,6 +4590,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) def _wrap_send_with_stateful_session_auth_context( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 67d935e34e3..7e8d08e7cad 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -26762,6 +26762,113 @@ "title": "ToolPolicyUpdateResponse", "type": "object" }, + "ToolSpendDailyEntry": { + "description": "Spend attributed to one tool on one UTC day.", + "properties": { + "call_count": { + "default": 0, + "title": "Call Count", + "type": "integer" + }, + "date": { + "title": "Date", + "type": "string" + }, + "spend": { + "default": 0.0, + "title": "Spend", + "type": "number" + }, + "tool_name": { + "title": "Tool Name", + "type": "string" + } + }, + "required": [ + "date", + "tool_name" + ], + "title": "ToolSpendDailyEntry", + "type": "object" + }, + "ToolSpendEntry": { + "description": "Total spend attributed to one tool over the requested window.", + "properties": { + "call_count": { + "default": 0, + "title": "Call Count", + "type": "integer" + }, + "spend": { + "default": 0.0, + "description": "Attributed spend: a request that used several tools counts its full spend toward each of them", + "title": "Spend", + "type": "number" + }, + "tool_name": { + "title": "Tool Name", + "type": "string" + }, + "total_tokens": { + "default": 0, + "title": "Total Tokens", + "type": "integer" + } + }, + "required": [ + "tool_name" + ], + "title": "ToolSpendEntry", + "type": "object" + }, + "ToolSpendResponse": { + "properties": { + "by_tool": { + "items": { + "$ref": "#/components/schemas/ToolSpendEntry" + }, + "title": "By Tool", + "type": "array" + }, + "daily": { + "items": { + "$ref": "#/components/schemas/ToolSpendDailyEntry" + }, + "title": "Daily", + "type": "array" + }, + "end_date": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "End Date" + }, + "start_date": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Start Date" + }, + "total_spend": { + "default": 0.0, + "description": "Deduplicated spend of every request that called at least one tool in the window; less than the sum of per-tool attributed spend whenever multi-tool requests exist", + "title": "Total Spend", + "type": "number" + } + }, + "title": "ToolSpendResponse", + "type": "object" + }, "ToolUsageLogEntry": { "description": "One spend log row for a tool call (for UI \"recent logs\" table).", "properties": { @@ -26858,6 +26965,13 @@ }, "ValidationError": { "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, "loc": { "items": { "anyOf": [ @@ -27301,6 +27415,81 @@ ] } }, + "/v1/tool/spend": { + "get": { + "description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nJoins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to\n``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools\ncounts its full spend toward each of those tools, so per-tool numbers are\nattributions. ``total_spend`` is the deduplicated spend of every request that\ncalled at least one tool in the window, so it never double counts.", + "operationId": "get_tool_spend_v1_tool_spend_get", + "parameters": [ + { + "description": "YYYY-MM-DD (defaults to 30 days ago)", + "in": "query", + "name": "start_date", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "YYYY-MM-DD (defaults to 30 days ago)", + "title": "Start Date" + } + }, + { + "description": "YYYY-MM-DD (defaults to today)", + "in": "query", + "name": "end_date", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "YYYY-MM-DD (defaults to today)", + "title": "End Date" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ToolSpendResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Tool Spend", + "tags": [ + "tools" + ] + } + }, "/v1/tool/{tool_name}": { "get": { "description": "Get details for a single tool.", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 444e5ba0731..98efadc10a8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2616,6 +2616,17 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # key off. Server-only and stripped from validated input for the same reason as the marker # above: a forged entry would let a caller pick which team's rpm bucket it is charged against. mcp_source_team_rpm_limits: dict[str, dict[str, int]] | None = Field(default=None, exclude=True) + via_virtual_key: bool = Field( + default=False, + exclude=True, + description=( + "Server-only marker set exclusively by the DB virtual-key and master-key auth paths via " + "post-construction assignment. Stripped from validated input so custom auth handlers, JWT " + "claims, or key metadata cannot forge it. Gates overwrite_user_with_key_hash stamping: only " + "a credential the proxy itself validated as a key may be forwarded as the provider-facing " + "user id." + ), + ) budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True) budget_throttle_pct: Optional[float] = Field(default=None, exclude=True) user: Optional[Any] = None # Expanded user object when expand=user is used @@ -2641,6 +2652,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data. values.pop("mcp_admitted_user_subject", None) values.pop("mcp_source_team_rpm_limits", None) + values.pop("via_virtual_key", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 83a8a69511b..709df5e64df 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1497,6 +1497,13 @@ async def _user_api_key_auth_builder( check_cache_only=True, ).resolve(hashed_token=hash_token(api_key)) ) + # Key-cache entries are written only after the proxy validated a + # virtual key or the master key, but via_virtual_key is exclude=True + # so serialization drops it; restore it at this trusted boundary. + # The UI-login JWT fallback below constructs its token from a + # decrypted blob, not this cache, and stays unmarked. + if isinstance(valid_token, UserAPIKeyAuth): + valid_token.via_virtual_key = True except Exception: verbose_logger.debug("api key not found in cache.") valid_token = None @@ -1614,6 +1621,7 @@ async def _user_api_key_auth_builder( _user_api_key_obj = update_valid_token_with_end_user_params( valid_token=_user_api_key_obj, end_user_params=end_user_params ) + _user_api_key_obj.via_virtual_key = True return _user_api_key_obj @@ -2021,7 +2029,7 @@ async def _user_api_key_auth_builder( # No token was found when looking up in the DB raise Exception("Invalid proxy server token passed") if valid_token_dict is not None: - return await _return_user_api_key_auth_obj( + virtual_key_auth_obj = await _return_user_api_key_auth_obj( user_obj=user_obj, api_key=api_key, parent_otel_span=parent_otel_span, @@ -2029,6 +2037,8 @@ async def _user_api_key_auth_builder( route=route, start_time=start_time, ) + virtual_key_auth_obj.via_virtual_key = True + return virtual_key_auth_obj except Exception as e: return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( e=e, @@ -2442,6 +2452,7 @@ async def _reserve_budget_after_common_checks( end_user_id=end_user_id, end_user_object=end_user_object, skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True, + fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, ) diff --git a/litellm/proxy/config_resolvers/sso.py b/litellm/proxy/config_resolvers/sso.py index 3d83c06dd62..97c42106018 100644 --- a/litellm/proxy/config_resolvers/sso.py +++ b/litellm/proxy/config_resolvers/sso.py @@ -36,6 +36,10 @@ SSO_DESCRIPTORS: tuple[FieldDescriptor, ...] = ( FieldDescriptor("generic_token_endpoint", "generic_token_endpoint", "GENERIC_TOKEN_ENDPOINT"), FieldDescriptor("generic_userinfo_endpoint", "generic_userinfo_endpoint", "GENERIC_USERINFO_ENDPOINT"), FieldDescriptor("generic_scope", "generic_scope", "GENERIC_SCOPE", default="openid email profile"), + FieldDescriptor("saml_idp_metadata_url", "saml_idp_metadata_url", "SAML_IDP_METADATA_URL"), + FieldDescriptor("saml_idp_metadata_xml", "saml_idp_metadata_xml", "SAML_IDP_METADATA_XML"), + FieldDescriptor("saml_sp_entity_id", "saml_sp_entity_id", "SAML_SP_ENTITY_ID"), + FieldDescriptor("saml_allow_unsolicited", "saml_allow_unsolicited", "SAML_ALLOW_UNSOLICITED"), FieldDescriptor("proxy_base_url", "proxy_base_url", "PROXY_BASE_URL"), ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 31535a5b569..28b9dec100f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -432,7 +432,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): Override to store only the Model Armor API response, not the entire data dict. This prevents circular references in logging. """ - metadata = request_data.get("metadata", {}) if isinstance(request_data, dict) else {} + metadata = (request_data.get("metadata") or {}) if isinstance(request_data, dict) else {} guardrail_response = metadata.get("_model_armor_response", {}) # Determine status – default to "success" but prefer the explicit value if present. diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 1f2c9e0c182..bd00e9815a8 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -489,6 +489,9 @@ class InMemoryGuardrailHandler: "skip_tool_message_in_guardrail", getattr(litellm_params, "skip_tool_message_in_guardrail", None), ) + configured_run_in_parallel = getattr(litellm_params, "run_in_parallel", None) + if configured_run_in_parallel is not None: + custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel) parsed_guardrail = Guardrail( guardrail_id=guardrail.get("guardrail_id"), diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9d9ef28ec9b..a4cc4a62009 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -13,7 +13,7 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS, PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( iter_client_callback_metadata_dicts, @@ -48,6 +48,24 @@ _EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session- # Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores # (covers UUIDs and most common session-id formats). _SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") + +_SHA256_HEX_RE = re.compile(r"^[0-9a-f]{64}$") + + +def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: + """Only proxy-validated keys are stamped, proven by the unforgeable + via_virtual_key marker AND a known non-secret shape: the sha256 hex digest + UserAPIKeyAuth stores virtual keys in, or the master key's stable alias. + Custom-auth credentials arrive raw (never forward auth material) and hashed + JWTs rotate on re-issue (useless as a stable ban id), so both are skipped.""" + api_key = user_api_key_dict.api_key + if not user_api_key_dict.via_virtual_key or api_key is None: + return None + if api_key == LITELLM_PROXY_MASTER_KEY_ALIAS or _SHA256_HEX_RE.fullmatch(api_key): + return api_key + return None + + _ANTHROPIC_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]+$") @@ -1447,6 +1465,11 @@ async def add_litellm_data_to_request( if "user" not in data: data["user"] = user + if litellm.overwrite_user_with_key_hash is True: + stampable_hash = _stampable_key_hash(user_api_key_dict) + if stampable_hash is not None: + data["user"] = stampable_hash + data["secret_fields"] = SecretFields(raw_headers=_raw_headers) ## Dynamic api version (Azure OpenAI endpoints) ## diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py new file mode 100644 index 00000000000..37b641ca123 --- /dev/null +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -0,0 +1,493 @@ +""" +SAML 2.0 SSO for the LiteLLM proxy admin UI. + +Supports both SP-initiated and IdP-initiated login via the HTTP-POST binding, +using the OneLogin python3-saml toolkit for signature, audience and time +validation. The IdP is configured from its metadata (``SAML_IDP_METADATA_URL`` +or inline ``SAML_IDP_METADATA_XML``); a successful login is mapped to a +``CustomOpenID`` and handed to the shared post-login path used by every other +SSO provider. + +python3-saml pulls in the native ``xmlsec``/``libxml2`` libraries, so it is an +optional dependency. When it is not installed the SAML routes return a clear +error instead of breaking proxy startup. +""" + +# python3-saml ships no type stubs, so the type checker sees every onelogin call +# as Unknown and the guarded optional import as possibly-unbound. Values crossing +# that boundary are cast() to concrete types at each use site; these directives +# silence only the unavoidable noise from the untyped dependency in this module. +# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false +# pyright: reportUnknownArgumentType=false, reportUnknownParameterType=false +# pyright: reportMissingTypeStubs=false, reportPossiblyUnboundVariable=false +# pyright: reportConstantRedefinition=false + +import asyncio +import hashlib +import os +import secrets +import time +from typing import cast +from urllib.parse import parse_qsl + +from fastapi import HTTPException, Request, status +from fastapi.responses import RedirectResponse +from pydantic import ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role +from litellm.proxy.utils import get_custom_url + +try: + from onelogin.saml2.auth import OneLogin_Saml2_Auth + from onelogin.saml2.idp_metadata_parser import OneLogin_Saml2_IdPMetadataParser + from onelogin.saml2.settings import OneLogin_Saml2_Settings + from onelogin.saml2.xml_utils import OneLogin_Saml2_XML + + SAML_AVAILABLE = True +except ImportError: + SAML_AVAILABLE = False + +SAML_LOGIN_ROUTE = "sso/saml/login" +SAML_CALLBACK_ROUTE = "sso/saml/callback" +SAML_METADATA_ROUTE = "sso/saml/metadata" + +_SAML_AUTHN_STATE_COOKIE = "litellm_saml_authn" +_SAML_IDP_SETTINGS_CACHE_PREFIX = "saml_idp_settings" +_SAML_AUTHN_REQUEST_CACHE_PREFIX = "saml_authn_request" +_SAML_CONSUMED_ASSERTION_CACHE_PREFIX = "saml_consumed_assertion" +_SAML_AUTHN_REQUEST_TTL_SECONDS = 600 +_SAML_IDP_METADATA_TTL_SECONDS = 3600 +_SAML_METADATA_FETCH_TIMEOUT_SECONDS = 10 +_SAML_MAX_POST_BYTES = 5 * 1024 * 1024 +# The replay guard tracks each assertion's NotOnOrAfter so it spans the full +# validity window; the floor covers IdPs that issue hour-long assertions or omit +# the timestamp, and the cap bounds cache growth. +_SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS = 3600 +_SAML_REPLAY_GUARD_MAX_TTL_SECONDS = 86400 + +_EMAIL_ATTRIBUTE_CANDIDATES = ( + "urn:oid:0.9.2342.19200300.100.1.3", + "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/emailaddress", + "email", + "emailAddress", + "mail", + "Email", +) +_FIRST_NAME_ATTRIBUTE_CANDIDATES = ( + "urn:oid:2.5.4.42", + "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/givenname", + "givenName", + "first_name", + "firstName", +) +_LAST_NAME_ATTRIBUTE_CANDIDATES = ( + "urn:oid:2.5.4.4", + "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/surname", + "sn", + "surname", + "last_name", + "lastName", +) +_ROLE_ATTRIBUTE_CANDIDATES = ("role", "roles", "litellm_role") +_TEAM_IDS_ATTRIBUTE_CANDIDATES = ("teams", "team_ids", "groups") + + +def _saml_unavailable_error() -> HTTPException: + return HTTPException( + status_code=status.HTTP_501_NOT_IMPLEMENTED, + detail=( + "SAML SSO requires the optional 'python3-saml' dependency, which is " + "not installed. Re-install litellm with the saml extra: " + "'pip install litellm[saml]'. The saml extra bundles the native " + "xmlsec/libxml2 libraries, so no system packages are required." + ), + ) + + +class SAMLAuthHandler: + """SP- and IdP-initiated SAML 2.0 login for the admin UI.""" + + @staticmethod + def _env(name: str, default: str | None = None) -> str | None: + return os.getenv(name, default) + + @staticmethod + def is_saml_configured() -> bool: + return bool(SAMLAuthHandler._env("SAML_IDP_METADATA_URL") or SAMLAuthHandler._env("SAML_IDP_METADATA_XML")) + + @staticmethod + def _bool_env(name: str, default: bool) -> bool: + raw = SAMLAuthHandler._env(name) + if raw is None: + return default + return raw.strip().lower() in ("true", "1", "yes", "on") + + @staticmethod + def _base_url(request: Request) -> str: + base = get_custom_url(request_base_url=str(request.base_url)) + return base if base.endswith("/") else base + "/" + + @staticmethod + def _is_https(request: Request) -> bool: + return SAMLAuthHandler._base_url(request).startswith("https") + + @staticmethod + def _acs_url(request: Request) -> str: + return SAMLAuthHandler._base_url(request) + SAML_CALLBACK_ROUTE + + @staticmethod + def _metadata_url(request: Request) -> str: + return SAMLAuthHandler._base_url(request) + SAML_METADATA_ROUTE + + @staticmethod + def _sp_entity_id(request: Request) -> str: + return SAMLAuthHandler._env("SAML_SP_ENTITY_ID") or SAMLAuthHandler._metadata_url(request) + + @staticmethod + async def _load_idp_settings(cache: DualCache) -> dict[str, object]: + metadata_url = SAMLAuthHandler._env("SAML_IDP_METADATA_URL") + metadata_xml = SAMLAuthHandler._env("SAML_IDP_METADATA_XML") + source = metadata_url or metadata_xml + if source is None: + raise HTTPException( + status_code=status.HTTP_501_NOT_IMPLEMENTED, + detail="SAML SSO is not configured. Set SAML_IDP_METADATA_URL or SAML_IDP_METADATA_XML.", + ) + + cache_key = f"{_SAML_IDP_SETTINGS_CACHE_PREFIX}:{hashlib.sha256(source.encode()).hexdigest()}" + cached = cache.get_cache(key=cache_key) + if isinstance(cached, dict): + return cast(dict[str, object], cached) # cast-ok: untyped python3-saml + + if metadata_url is not None: + parsed = await asyncio.to_thread( + OneLogin_Saml2_IdPMetadataParser.parse_remote, + metadata_url, + validate_cert=SAMLAuthHandler._bool_env("SAML_IDP_METADATA_VALIDATE_CERT", True), + timeout=_SAML_METADATA_FETCH_TIMEOUT_SECONDS, + ) + else: + parsed = OneLogin_Saml2_IdPMetadataParser.parse(cast(str, metadata_xml)) # cast-ok: untyped python3-saml + + idp_settings = cast(dict[str, object], parsed) # cast-ok: untyped python3-saml + if not idp_settings.get("idp"): + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail="Could not parse an IdP entityID/SSO URL/certificate from the SAML metadata.", + ) + cache.set_cache(key=cache_key, value=idp_settings, ttl=_SAML_IDP_METADATA_TTL_SECONDS) + return idp_settings + + @staticmethod + def _build_settings(request: Request, idp_settings: dict[str, object]) -> dict[str, object]: + sp_settings: dict[str, object] = { + "strict": SAMLAuthHandler._bool_env("SAML_STRICT", True), + "debug": False, + "sp": { + "entityId": SAMLAuthHandler._sp_entity_id(request), + "assertionConsumerService": { + "url": SAMLAuthHandler._acs_url(request), + "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST", + }, + "NameIDFormat": SAMLAuthHandler._env( + "SAML_SP_NAME_ID_FORMAT", + "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress", + ), + }, + "security": { + "wantAssertionsSigned": SAMLAuthHandler._bool_env("SAML_WANT_ASSERTIONS_SIGNED", True), + "wantMessagesSigned": SAMLAuthHandler._bool_env("SAML_WANT_MESSAGES_SIGNED", False), + "authnRequestsSigned": SAMLAuthHandler._bool_env("SAML_AUTHN_REQUESTS_SIGNED", False), + "wantNameId": True, + "requestedAuthnContext": False, + "rejectUnsolicitedResponsesWithInResponseTo": False, + }, + } + return OneLogin_Saml2_IdPMetadataParser.merge_settings(sp_settings, idp_settings) + + @staticmethod + def _prepare_request_data(request: Request, post_data: dict[str, str] | None = None) -> dict[str, object]: + base = SAMLAuthHandler._base_url(request) + scheme, _, host_part = base.partition("://") + host = host_part.split("/", 1)[0] + return { + "https": "on" if scheme == "https" else "off", + "http_host": host, + "script_name": "/" + SAML_CALLBACK_ROUTE, + "get_data": dict(request.query_params), + "post_data": post_data or {}, + } + + @staticmethod + async def _build_auth( + request: Request, + cache: DualCache, + post_data: dict[str, str] | None = None, + ) -> "OneLogin_Saml2_Auth": + if not SAML_AVAILABLE: + raise _saml_unavailable_error() + idp_settings = await SAMLAuthHandler._load_idp_settings(cache) + settings = SAMLAuthHandler._build_settings(request, idp_settings) + request_data = SAMLAuthHandler._prepare_request_data(request, post_data) + try: + return OneLogin_Saml2_Auth(request_data, old_settings=settings) + except Exception as e: # noqa: BLE001 - toolkit exposes no common exception base; fail closed + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Invalid SAML configuration: {e}", + ) + + @staticmethod + async def build_login_redirect( + request: Request, cache: DualCache, relay_state: str | None = None + ) -> RedirectResponse: + auth = await SAMLAuthHandler._build_auth(request, cache) + redirect_url = cast(str, auth.login(return_to=relay_state)) # cast-ok: untyped python3-saml + response = RedirectResponse(url=redirect_url, status_code=303) + request_id = cast(str | None, auth.get_last_request_id()) # cast-ok: untyped python3-saml + if request_id is not None: + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", + value="1", + ttl=_SAML_AUTHN_REQUEST_TTL_SECONDS, + ) + secure = SAMLAuthHandler._is_https(request) + response.set_cookie( + key=_SAML_AUTHN_STATE_COOKIE, + value=request_id, + max_age=_SAML_AUTHN_REQUEST_TTL_SECONDS, + httponly=True, + secure=secure, + samesite="none" if secure else "lax", + ) + return response + + @staticmethod + async def build_sp_metadata(request: Request, cache: DualCache) -> str: + if not SAML_AVAILABLE: + raise _saml_unavailable_error() + idp_settings = await SAMLAuthHandler._load_idp_settings(cache) + settings = SAMLAuthHandler._build_settings(request, idp_settings) + saml_settings = OneLogin_Saml2_Settings(settings, sp_validation_only=True) + metadata = cast(str, saml_settings.get_sp_metadata()) # cast-ok: untyped python3-saml + errors = cast(list[str], saml_settings.validate_metadata(metadata)) # cast-ok: untyped python3-saml + if errors: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Invalid SP metadata: {', '.join(errors)}", + ) + return metadata + + @staticmethod + async def read_acs_post_data(request: Request) -> dict[str, str]: + """Read the ACS POST form under a hard size cap before any base64/XML decoding. + + Bounds both Content-Length-declared and chunked requests so an unauthenticated + caller cannot force unbounded buffering while decoding the SAMLResponse.""" + declared = request.headers.get("content-length") + if declared is not None and declared.isdigit() and int(declared) > _SAML_MAX_POST_BYTES: + raise HTTPException( + status_code=status.HTTP_413_CONTENT_TOO_LARGE, + detail="SAML response exceeds the maximum allowed size.", + ) + + body = bytearray() + async for chunk in request.stream(): + body += chunk + if len(body) > _SAML_MAX_POST_BYTES: + raise HTTPException( + status_code=status.HTTP_413_CONTENT_TOO_LARGE, + detail="SAML response exceeds the maximum allowed size.", + ) + + return dict(parse_qsl(body.decode("utf-8", "replace"))) + + @staticmethod + async def handle_acs(request: Request, cache: DualCache, post_data: dict[str, str]) -> CustomOpenID: + auth = await SAMLAuthHandler._build_auth(request, cache, post_data=post_data) + browser_request_id = request.cookies.get(_SAML_AUTHN_STATE_COOKIE) + try: + auth.process_response(request_id=browser_request_id) + except Exception as e: # noqa: BLE001 - toolkit exposes no common exception base; fail closed + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"Could not process SAML response: {e}", + ) + + errors = cast(list[str], auth.get_errors()) # cast-ok: untyped python3-saml + if errors or not auth.is_authenticated(): + reason = auth.get_last_error_reason() + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"SAML authentication failed: {reason or ', '.join(errors)}", + ) + + await SAMLAuthHandler._enforce_response_binding(auth, cache, browser_request_id) + return SAMLAuthHandler._result_from_auth(auth) + + @staticmethod + def _replay_guard_ttl(auth: "OneLogin_Saml2_Auth") -> int: + not_on_or_after = auth.get_last_assertion_not_on_or_after() + if not isinstance(not_on_or_after, int): + return _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS + remaining = not_on_or_after - int(time.time()) + return min( + max(remaining, _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS), + _SAML_REPLAY_GUARD_MAX_TTL_SECONDS, + ) + + @staticmethod + def _response_in_response_to(auth: "OneLogin_Saml2_Auth") -> str | None: + """The request id this response answers, read from the Response element or, when the + IdP only stamps it on the bearer SubjectConfirmationData, from there. A non-None value + marks the response as solicited (SP-initiated) and so requiring browser binding.""" + value = cast(str | None, auth.get_last_response_in_response_to()) # cast-ok: untyped python3-saml + if value: + return value + xml = cast(bytes | None, auth.get_last_response_xml()) # cast-ok: untyped python3-saml + if not xml: + return None + root = OneLogin_Saml2_XML.to_etree(xml) + for node in OneLogin_Saml2_XML.query(root, "//saml:SubjectConfirmationData[@InResponseTo]"): + irt = cast(str | None, node.get("InResponseTo")) # cast-ok: untyped python3-saml + if irt: + return irt + return None + + @staticmethod + async def _enforce_response_binding( + auth: "OneLogin_Saml2_Auth", + cache: DualCache, + browser_request_id: str | None, + ) -> None: + in_response_to = SAMLAuthHandler._response_in_response_to(auth) + + if in_response_to is not None: + authn_key = f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{in_response_to}" + if cache.get_cache(key=authn_key) is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML response references an unknown or already-used login request.", + ) + if browser_request_id is None or not secrets.compare_digest(browser_request_id, in_response_to): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML response is not bound to this browser's login request.", + ) + elif browser_request_id is not None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML response is not bound to this browser's login request.", + ) + elif not SAMLAuthHandler._bool_env("SAML_ALLOW_UNSOLICITED", False): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Unsolicited (IdP-initiated) SAML responses are disabled.", + ) + elif cache.redis_cache is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=( + "Unsolicited (IdP-initiated) SAML responses require a shared Redis cache " + "so the replay guard is enforced across every worker." + ), + ) + + assertion_id = cast(str | None, auth.get_last_assertion_id()) # cast-ok: untyped python3-saml + if assertion_id is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML assertion is missing the required ID attribute.", + ) + consumed_key = f"{_SAML_CONSUMED_ASSERTION_CACHE_PREFIX}:{assertion_id}" + consumed_count = await cache.async_increment_cache( + key=consumed_key, value=1, ttl=SAMLAuthHandler._replay_guard_ttl(auth) + ) + if consumed_count is not None and consumed_count > 1: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML assertion has already been used (replay detected).", + ) + + @staticmethod + def _result_from_auth(auth: "OneLogin_Saml2_Auth") -> CustomOpenID: + attributes = cast(dict[str, list[str]], auth.get_attributes()) # cast-ok: untyped python3-saml + name_id = cast(str | None, auth.get_nameid()) # cast-ok: untyped python3-saml + + email = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_EMAIL", _EMAIL_ATTRIBUTE_CANDIDATES) + if email is None and name_id is not None and "@" in name_id: + email = name_id + + if email is None and SAMLAuthHandler._env("ALLOWED_EMAIL_DOMAINS") is not None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=( + "SAML assertion did not contain an email address, but ALLOWED_EMAIL_DOMAINS " + "restricts sign-in by email domain." + ), + ) + + user_id = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_USER_ID", ()) or name_id or email + if user_id is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML assertion did not contain a usable subject (NameID) or email.", + ) + + first_name = SAMLAuthHandler._attribute_value( + attributes, "SAML_ATTRIBUTE_FIRST_NAME", _FIRST_NAME_ATTRIBUTE_CANDIDATES + ) + last_name = SAMLAuthHandler._attribute_value( + attributes, "SAML_ATTRIBUTE_LAST_NAME", _LAST_NAME_ATTRIBUTE_CANDIDATES + ) + role_value = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES) + team_ids = SAMLAuthHandler._attribute_values( + attributes, "SAML_ATTRIBUTE_TEAM_IDS", _TEAM_IDS_ATTRIBUTE_CANDIDATES + ) + + display_name = " ".join(part for part in (first_name, last_name) if part) or email + + verbose_proxy_logger.info(f"SAML login: subject={user_id}, email={email}, attributes={list(attributes.keys())}") + + try: + return CustomOpenID( + id=user_id, + email=email, + first_name=first_name, + last_name=last_name, + display_name=display_name, + picture=None, + provider="saml", + team_ids=team_ids, + user_role=get_litellm_user_role(role_value) if role_value else None, + ) + except ValidationError as e: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"SAML assertion contained an invalid subject or email: {e}", + ) + + @staticmethod + def _attribute_value( + attributes: dict[str, list[str]], + env_override: str, + candidates: tuple[str, ...], + ) -> str | None: + values = SAMLAuthHandler._attribute_values(attributes, env_override, candidates) + return values[0] if values else None + + @staticmethod + def _attribute_values( + attributes: dict[str, list[str]], + env_override: str, + candidates: tuple[str, ...], + ) -> list[str]: + override = SAMLAuthHandler._env(env_override) + keys = (override, *candidates) if override else candidates + for key in keys: + values = attributes.get(key) + if values: + return [v for v in values if v] + return [] diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 9d71761f115..ca606e07cee 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -10,16 +10,18 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a """ import uuid -from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, List, Optional +from datetime import datetime, timedelta, timezone +from itertools import groupby +from typing import TYPE_CHECKING, Annotated, Any, List, Optional from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import BaseModel, TypeAdapter if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import ( @@ -39,6 +41,9 @@ from litellm.types.tool_management import ( ToolPolicyOptionsResponse, ToolPolicyUpdateRequest, ToolPolicyUpdateResponse, + ToolSpendDailyEntry, + ToolSpendEntry, + ToolSpendResponse, ToolUsageLogEntry, ToolUsageLogsResponse, ) @@ -124,6 +129,147 @@ async def list_tools( raise HTTPException(status_code=500, detail=str(e)) +def _parse_day_start(value: str | None) -> datetime | None: + if not value: + return None + try: + return datetime.strptime(value.strip(), "%Y-%m-%d").replace(tzinfo=timezone.utc) + except ValueError: + raise HTTPException( + status_code=400, + detail=f"Invalid date format: {value}. Expected: 'YYYY-MM-DD'", + ) + + +class _ToolSpendRow(BaseModel): + date: str + tool_name: str + call_count: int + spend: float + total_tokens: int + + +class _RequestTotalRow(BaseModel): + total_spend: float + + +_TOOL_SPEND_ROWS = TypeAdapter(list[_ToolSpendRow]) +_REQUEST_TOTAL_ROWS = TypeAdapter(list[_RequestTotalRow]) + + +def _summarize_tool(name: str, grp: tuple[_ToolSpendRow, ...]) -> ToolSpendEntry: + return ToolSpendEntry( + tool_name=name, + spend=sum(r.spend for r in grp), + call_count=sum(r.call_count for r in grp), + total_tokens=sum(r.total_tokens for r in grp), + ) + + +def _build_tool_spend_response( + rows: list[_ToolSpendRow], + total_spend: float, + start_date: str, + end_date: str, +) -> ToolSpendResponse: + daily = [ + ToolSpendDailyEntry(date=r.date, tool_name=r.tool_name, spend=r.spend, call_count=r.call_count) for r in rows + ] + grouped = groupby(sorted(rows, key=lambda r: r.tool_name), key=lambda r: r.tool_name) + by_tool = sorted( + (_summarize_tool(name, tuple(grp)) for name, grp in grouped), + key=lambda e: e.spend, + reverse=True, + ) + return ToolSpendResponse( + by_tool=by_tool, + daily=daily, + total_spend=total_spend, + start_date=start_date, + end_date=end_date, + ) + + +@router.get( + "/v1/tool/spend", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolSpendResponse, +) +async def get_tool_spend( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + start_date: Annotated[str | None, Query(description="YYYY-MM-DD (defaults to 30 days ago)")] = None, + end_date: Annotated[str | None, Query(description="YYYY-MM-DD (defaults to today)")] = None, +): + """ + Spend attributed to each tool over a date range, for the Cost Optimization dashboard. + + Joins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to + ``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools + counts its full spend toward each of those tools, so per-tool numbers are + attributions. ``total_spend`` is the deduplicated spend of every request that + called at least one tool in the window, so it never double counts. + """ + from litellm.proxy.proxy_server import prisma_client + + if user_api_key_dict.user_role not in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ): + raise HTTPException( + status_code=403, + detail="Only proxy admin roles can view tool spend across the deployment", + ) + + if prisma_client is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + + now = datetime.now(timezone.utc) + end_day = _parse_day_start(end_date) + start_dt = _parse_day_start(start_date) or ((end_day or now) - timedelta(days=30)) + end_exclusive = (end_day + timedelta(days=1)) if end_day else now + + rows = await prisma_client.db.query_raw( + """ + SELECT to_char(ti.start_time, 'YYYY-MM-DD') AS date, + ti.tool_name AS tool_name, + COUNT(*)::int AS call_count, + COALESCE(SUM(sl.spend), 0)::double precision AS spend, + COALESCE(SUM(sl.total_tokens), 0)::bigint AS total_tokens + FROM "LiteLLM_SpendLogToolIndex" ti + JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id + WHERE ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') + AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') + GROUP BY date, ti.tool_name + ORDER BY date ASC, spend DESC + """, + start_dt.isoformat(), + end_exclusive.isoformat(), + ) + totals = await prisma_client.db.query_raw( + """ + SELECT COALESCE(SUM(sl.spend), 0)::double precision AS total_spend + FROM "LiteLLM_SpendLogs" sl + WHERE EXISTS ( + SELECT 1 + FROM "LiteLLM_SpendLogToolIndex" ti + WHERE ti.request_id = sl.request_id + AND ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') + AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') + ) + """, + start_dt.isoformat(), + end_exclusive.isoformat(), + ) + total_rows = _REQUEST_TOTAL_ROWS.validate_python(totals or []) + return _build_tool_spend_response( + rows=_TOOL_SPEND_ROWS.validate_python(rows or []), + total_spend=total_rows[0].total_spend if total_rows else 0.0, + start_date=start_dt.strftime("%Y-%m-%d"), + end_date=(end_day or now).strftime("%Y-%m-%d"), + ) + + @router.get( "/v1/tool/{tool_name:path}/detail", tags=["tool management"], diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 31b98bf20e4..8682b61f910 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -100,6 +100,7 @@ from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO +from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler from litellm.proxy.management_endpoints.sso_helper_utils import ( check_is_admin_only_access, has_admin_ui_access, @@ -857,6 +858,27 @@ def process_sso_jwt_access_token( return None +async def _raise_if_sso_exceeds_free_user_limit(premium_user: bool, prisma_client: PrismaClient | None) -> None: + """Free tier allows SSO for up to 5 billable users; beyond that requires an Enterprise license.""" + if premium_user is True: + return + if prisma_client is None: + raise ProxyException( + message=CommonProxyErrors.db_not_connected_error.value, + type=ProxyErrorTypes.auth_error, + param="premium_user", + code=status.HTTP_403_FORBIDDEN, + ) + billable_users = await UserRepository(prisma_client).count_billable_users() + if billable_users and billable_users > 5: + raise ProxyException( + message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You configured SSO (one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, `GENERIC_CLIENT_ID`, or SAML) in your env. Please unset it", + type=ProxyErrorTypes.auth_error, + param="premium_user", + code=status.HTTP_403_FORBIDDEN, + ) + + @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) async def google_login( request: Request, @@ -876,6 +898,7 @@ async def google_login( general_settings, premium_user, prisma_client, + user_api_key_cache, user_custom_ui_sso_sign_in_handler, ) @@ -891,25 +914,13 @@ async def google_login( return admin_ui_disabled() ####### Check if user is a Enterprise / Premium User ####### - if microsoft_client_id is not None or google_client_id is not None or generic_client_id is not None: - if premium_user is not True: - # Check if under 'free SSO user' limit - if prisma_client is not None: - billable_users = await UserRepository(prisma_client).count_billable_users() - if billable_users and billable_users > 5: - raise ProxyException( - message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) - else: - raise ProxyException( - message=CommonProxyErrors.db_not_connected_error.value, - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) + if ( + microsoft_client_id is not None + or google_client_id is not None + or generic_client_id is not None + or SAMLAuthHandler.is_saml_configured() + ): + await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) ####### Detect DB + MASTER KEY in .env ####### missing_env_vars = show_missing_vars_in_env() @@ -947,6 +958,19 @@ async def google_login( "Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise." ) + if ( + microsoft_client_id is None + and google_client_id is None + and generic_client_id is None + and SAMLAuthHandler.is_saml_configured() + ): + verbose_proxy_logger.info("Redirecting to SAML SSO login") + return await SAMLAuthHandler.build_login_redirect( + request=request, + cache=user_api_key_cache, + relay_state=return_to, + ) + # Check if we should use SSO handler if ( SSOAuthenticationHandler.should_use_sso_handler( @@ -1913,6 +1937,81 @@ async def auth_callback(request: Request, state: Optional[str] = None): ) +@router.get("/sso/saml/login", tags=["experimental"], include_in_schema=False) +async def saml_login(request: Request, return_to: str | None = None): + """SP-initiated SAML login. Redirects the user to the configured IdP.""" + from litellm.proxy.proxy_server import user_api_key_cache + + _disable_ui_flag = os.getenv("DISABLE_ADMIN_UI") + if _disable_ui_flag is not None and str_to_bool(value=_disable_ui_flag): + return admin_ui_disabled() + + return await SAMLAuthHandler.build_login_redirect(request=request, cache=user_api_key_cache, relay_state=return_to) + + +@router.get("/sso/saml/metadata", tags=["experimental"], include_in_schema=False) +async def saml_metadata(request: Request): + """Service Provider metadata XML, for registering this proxy at the IdP.""" + from litellm.proxy.proxy_server import user_api_key_cache + + metadata = await SAMLAuthHandler.build_sp_metadata(request=request, cache=user_api_key_cache) + return Response(content=metadata, media_type="application/xml") + + +@router.post("/sso/saml/callback", tags=["experimental"], include_in_schema=False) +async def saml_callback(request: Request): + """Assertion Consumer Service. Validates the IdP assertion and issues a UI session.""" + from litellm.proxy.proxy_server import ( + general_settings, + jwt_handler, + master_key, + premium_user, + prisma_client, + user_api_key_cache, + ) + + _disable_ui_flag = os.getenv("DISABLE_ADMIN_UI") + if _disable_ui_flag is not None and str_to_bool(value=_disable_ui_flag): + return admin_ui_disabled() + + if prisma_client is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + if master_key is None: + raise ProxyException( + message="Master Key not set for Proxy. Set `LITELLM_MASTER_KEY` in .env or general_settings:master_key in config.yaml.", + type=ProxyErrorTypes.auth_error, + param="master_key", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + post_data = await SAMLAuthHandler.read_acs_post_data(request) + if "SAMLResponse" not in post_data: + raise HTTPException(status_code=400, detail="Missing SAMLResponse in callback request.") + + result = await SAMLAuthHandler.handle_acs(request=request, cache=user_api_key_cache, post_data=post_data) + + await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) + + ui_access_mode = general_settings.get("ui_access_mode", None) + relay_state = post_data.get("RelayState") + cp_return_to: str | None = ( + relay_state + if isinstance(relay_state, str) and SSOAuthenticationHandler._validate_return_to(relay_state) + else None + ) + + return await SSOAuthenticationHandler.get_redirect_response_from_openid( + result=result, + request=request, + received_response=None, + generic_client_id=None, + ui_access_mode=ui_access_mode, + access_token_payload=None, + jwt_handler=jwt_handler, + return_to=cp_return_to, + ) + + async def _build_cli_sso_user_defined_values( result: Union[OpenID, dict], parsed_openid_result: ParsedOpenIDResult, diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 86beb063667..7c52b04c4eb 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -295,10 +295,16 @@ async def add_new_member( # same new user is provisioned concurrently), seeding teams on create. # The teams append lives in the filtered update below rather than the # upsert's update branch so an already-existing user does not get a - # duplicate team id. + # duplicate team id. The update branch still has to write something: + # Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it + # is non-empty, and falls back to a racy SELECT-then-INSERT when it is + # not, so this re-states user_id as a no-op rather than being empty. _returned_user = await UserRepository(prisma_client).table.upsert( where={"user_id": new_member.user_id}, - data={"create": {"teams": [team_id], **new_user_defaults}, "update": {}}, + data={ + "create": {"teams": [team_id], **new_user_defaults}, + "update": {"user_id": new_member.user_id}, + }, ) await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id) if _returned_user is not None: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 80fd8a1594e..013873179c6 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -4,7 +4,9 @@ import asyncio import json from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Mapping, Optional, Sequence, cast +from typing import Any, Dict, List, Mapping, NoReturn, Optional, Sequence, cast + +from fastapi import HTTPException, status import litellm from litellm._logging import verbose_proxy_logger @@ -59,6 +61,22 @@ class _CounterReservationUnavailable(Exception): super().__init__("Counter reservation unavailable") +def _raise_reservation_unavailable(counter_key: str) -> NoReturn: + verbose_proxy_logger.warning( + "fail_closed_budget_enforcement: rejecting request — budget reservation for %s could not be written", + counter_key, + ) + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + "Budget enforcement unavailable: the budget reservation could not " + "be written to the spend counter backend, and " + "fail_closed_budget_enforcement is enabled, so the request was " + "rejected to avoid exceeding the configured budget. Retry shortly." + ), + ) + + def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set: if not budget_reservation: return set() @@ -138,6 +156,7 @@ async def reserve_budget_for_request( end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, skip_user_budget_on_team_key: bool = False, + fail_closed_budget_enforcement: bool = False, ) -> Optional[dict]: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None @@ -193,6 +212,8 @@ async def reserve_budget_for_request( default_reserved_cost=reservation_cost, ) applied_entries.remove(entry) + if fail_closed_budget_enforcement: + _raise_reservation_unavailable(counter_key=counter.counter_key) continue if reserved_value is not None: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 55b50e7d9ff..0c525ee9466 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1695,13 +1695,10 @@ async def ui_view_spend_logs( code=status.HTTP_401_UNAUTHORIZED, ) - if start_date is None or end_date is None: - raise ProxyException( - message="Start date and end date are required", - type="bad_request", - param="None", - code=status.HTTP_400_BAD_REQUEST, - ) + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + + is_v2 = "/spend/logs/v2" in get_request_route(request) # Validate sort_by and sort_order valid_sort_fields = { @@ -1729,36 +1726,50 @@ async def ui_view_spend_logs( ) try: - # Inline import — auth_utils participates in a proxy import cycle. - from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + is_admin_view = _is_admin_view_safe(user_api_key_dict=user_api_key_dict) + is_request_id_lookup = request_id is not None and not is_v2 - is_v2 = "/spend/logs/v2" in get_request_route(request) - formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"] + if is_request_id_lookup: + # request_id is the @id primary key: it identifies a single row, so a + # time window is meaningless. The dashboard always sends a default 24h + # window, which hid ids copied from an older page (LIT-3981). Drop the + # window for the id lookup so it resolves across all time; every other + # query, including the public v2 route, still requires one (below). + start_date_obj: datetime | None = None + end_date_obj: datetime | None = None + else: + if start_date is None or end_date is None: + raise ProxyException( + message="Start date and end date are required", + type="bad_request", + param="None", + code=status.HTTP_400_BAD_REQUEST, + ) + formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"] - def parse_date(date_str: str) -> datetime: - date_str = date_str.strip() - for fmt in formats: - try: - return datetime.strptime(date_str, fmt).replace(tzinfo=timezone.utc) - except ValueError: - continue - expected = "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" if is_v2 else "'YYYY-MM-DD HH:MM:SS'" - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Invalid date format: {date_str}. Expected: {expected}", - ) + def parse_date(date_str: str) -> datetime: + date_str = date_str.strip() + for fmt in formats: + try: + return datetime.strptime(date_str, fmt).replace(tzinfo=timezone.utc) + except ValueError: + continue + expected = "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" if is_v2 else "'YYYY-MM-DD HH:MM:SS'" + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid date format: {date_str}. Expected: {expected}", + ) - start_date_obj = parse_date(start_date) - end_date_obj = parse_date(end_date) - - # Convert to ISO format strings for Prisma - start_date_iso = start_date_obj.isoformat() # Already in UTC, no need to add Z - end_date_iso = end_date_obj.isoformat() # Already in UTC, no need to add Z + start_date_obj = parse_date(start_date) + end_date_obj = parse_date(end_date) # Build where conditions - where_conditions: dict[str, Any] = { - "startTime": {"gte": start_date_iso, "lte": end_date_iso}, - } + where_conditions: dict[str, Any] = {} + if start_date_obj is not None and end_date_obj is not None: + where_conditions["startTime"] = { + "gte": start_date_obj.isoformat(), # Already in UTC, no need to add Z + "lte": end_date_obj.isoformat(), + } if team_id is not None: where_conditions["team_id"] = team_id @@ -1827,9 +1838,19 @@ async def ui_view_spend_logs( where_conditions["spend"]["gte"] = min_spend if max_spend is not None: where_conditions["spend"]["lte"] = max_spend - is_admin_view = _is_admin_view_safe(user_api_key_dict=user_api_key_dict) + # A request_id lookup drops the date window, so a non-admin could otherwise + # reach any single row by id; require they own it, mirroring the detail + # endpoint. That ownership check fully authorizes the one row, so the + # general scoping below is skipped for id lookups. Scoped to the UI route + # so the public v2 contract is unchanged. + if request_id is not None and not is_v2 and not is_admin_view: + await _assert_user_can_view_request_id( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + request_id=request_id, + ) permitted_team_ids: List[str] | None = None - if not is_admin_view: + if not is_request_id_lookup and not is_admin_view: if team_id is not None: can_view_team = await _can_team_member_view_log( prisma_client=prisma_client, @@ -1875,15 +1896,16 @@ async def ui_view_spend_logs( sql_params: List[Any] = [] p = 1 # parameter index counter - # Date range (always present). Wrap the param side with - # `AT TIME ZONE 'UTC'` so comparison against the plain `timestamp` - # column does not depend on the DB session timezone (see #22529). - sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')") - sql_params.append(start_date_obj) - p += 1 - sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')") - sql_params.append(end_date_obj) - p += 1 + # Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison + # against the plain `timestamp` column does not depend on the DB session + # timezone (see #22529). Absent for a request_id-only lookup (see above). + if start_date_obj is not None and end_date_obj is not None: + sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')") + sql_params.append(start_date_obj) + p += 1 + sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')") + sql_params.append(end_date_obj) + p += 1 # Equality filters - read effective values from where_conditions (post-authorization) for sql_col, wc_key in [ diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 50f6f791bc2..a6105b6dff9 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -374,12 +374,22 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if isinstance(v, BaseModel): v = v.model_dump() additional_usage_values.update({k: v}) - if "cache_read_input_tokens" not in additional_usage_values: - prompt_tokens_details = additional_usage_values.get("prompt_tokens_details") - if isinstance(prompt_tokens_details, dict): + prompt_tokens_details = additional_usage_values.get("prompt_tokens_details") + if not isinstance(prompt_tokens_details, dict): + usage_object = clean_metadata.get("usage_object") + if isinstance(usage_object, dict): + prompt_tokens_details = usage_object.get("prompt_tokens_details") + if isinstance(prompt_tokens_details, dict): + if "cache_read_input_tokens" not in additional_usage_values: cached_tokens = prompt_tokens_details.get("cached_tokens") if isinstance(cached_tokens, int) and cached_tokens > 0: additional_usage_values["cache_read_input_tokens"] = cached_tokens + if "cache_creation_input_tokens" not in additional_usage_values: + cache_write_tokens = prompt_tokens_details.get("cache_write_tokens") or prompt_tokens_details.get( + "cache_creation_tokens" + ) + if isinstance(cache_write_tokens, int) and cache_write_tokens > 0: + additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens clean_metadata["additional_usage_values"] = additional_usage_values if litellm.cache is not None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5b81d1f2da3..171e17ce650 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1400,6 +1400,14 @@ class ProxyLogging: self._process_guardrail_metadata(data) return data + parallel_guardrails: tuple[CustomGuardrail, ...] = tuple( + cb + for cb in caps.resolved_callbacks + if isinstance(cb, CustomGuardrail) + and getattr(cb, "run_in_parallel", False) + and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed) + ) + deferred_route_exc: Optional[SensitiveDataRouteException] = None for _callback in caps.resolved_callbacks: start_time = time.time() @@ -1409,6 +1417,9 @@ class ProxyLogging: if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed: continue + if getattr(_callback, "run_in_parallel", False): + continue + result = await self._process_guardrail_callback( callback=_callback, data=data, # type: ignore @@ -1465,6 +1476,14 @@ class ProxyLogging: if deferred_route_exc is not None and data is not None: data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict) + if parallel_guardrails and data is not None: + await self._run_parallel_pre_call_guardrails( + guardrails=parallel_guardrails, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + if data is not None: self._process_guardrail_metadata(data) @@ -1477,6 +1496,47 @@ class ProxyLogging: except Exception as e: raise e + async def _run_parallel_pre_call_guardrails( + self, + guardrails: tuple[CustomGuardrail, ...], + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: CallTypesLiteral, + ) -> None: + """ + Run opted-in pre_call guardrails concurrently against one shared payload + snapshot. These guardrails are declared block-only, so any modified data + they return is discarded; they run for their blocking side effect (raising + to reject the request before it reaches the LLM). Every guardrail is + awaited to completion (``return_exceptions=True``) so a raise by one never + leaves the others running as unobserved background tasks. A guardrail that + blocks (any exception other than a reroute or passthrough) takes precedence + over one that only changes the request flow, so a fast reroute can never + let a slower block be bypassed; the request is rejected before it reaches + the LLM, preserving the pre-call barrier that ``during_call`` guardrails + cannot provide. Per-guardrail latency is recorded by + ``_process_guardrail_callback``'s own metrics. + """ + results = await asyncio.gather( + *( + self._process_guardrail_callback( + callback=callback, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + event_type=GuardrailEventHooks.pre_call, + ) + for callback in guardrails + ), + return_exceptions=True, + ) + raised = tuple(result for result in results if isinstance(result, BaseException)) + blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) + if blocking is not None: + raise blocking + if raised: + raise raised[0] + async def _handle_sensitive_data_route_exception( self, exc: SensitiveDataRouteException, @@ -2277,9 +2337,16 @@ class ProxyLogging: # Merge model-level guardrails before checking which guardrails to run guardrail_data = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + parallel_guardrails: tuple[CustomGuardrail, ...] = tuple( + callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False) + ) + for callback in guardrail_callbacks: # Main - V2 Guardrails implementation + if getattr(callback, "run_in_parallel", False): + continue + if ( callback.should_run_guardrail( data=guardrail_data, @@ -2316,6 +2383,15 @@ class ProxyLogging: if guardrail_response is not None: response = guardrail_response + if parallel_guardrails: + await self._run_parallel_post_call_guardrails( + guardrails=parallel_guardrails, + data=data, + guardrail_data=guardrail_data, + response=response, + user_api_key_dict=user_api_key_dict, + ) + ############ Handle CustomLogger ############################### ################################################################# @@ -2329,6 +2405,65 @@ class ProxyLogging: raise e return response + async def _run_parallel_post_call_guardrails( + self, + guardrails: tuple[CustomGuardrail, ...], + data: dict, + guardrail_data: dict, + response: LLMResponseTypes, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """ + Run opted-in post_call guardrails concurrently against the response + produced by the sequential guardrails. These guardrails are declared + block-only, so any modified response they return is discarded; they run + for their blocking side effect (raising to reject the response before it + reaches the client). Every guardrail is awaited to completion + (``return_exceptions=True``) so a raise by one never leaves the others + running as unobserved background tasks. A guardrail that blocks (any + exception other than a passthrough) takes precedence over one that only + changes the response flow, so a fast passthrough can never let a slower + block be bypassed. Each per-guardrail coroutine sets ``guardrail_to_apply`` + immediately before awaiting, and the unified hook pops it before its first + suspension point, so concurrent guardrails never race on that key. + """ + + async def _run_one(callback: CustomGuardrail) -> None: + if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True: + return + if "apply_guardrail" in type(callback).__dict__: + data["guardrail_to_apply"] = callback + await self._run_guardrail_with_metrics( + callback, + unified_guardrail.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) + else: + await self._run_guardrail_with_metrics( + callback, + callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) + + results = await asyncio.gather( + *(_run_one(callback) for callback in guardrails), + return_exceptions=True, + ) + raised = tuple(result for result in results if isinstance(result, BaseException)) + blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) + if blocking is not None: + raise blocking + if raised: + raise raised[0] + async def post_call_response_headers_hook( self, data: dict, diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index ac92e5d6dcc..12c890ec91d 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1049,6 +1049,7 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), + cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), ) completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None output_tokens_details = getattr(response_api_usage, "output_tokens_details", None) diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index 10b4fb30f22..c6d0c1717a9 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -52,12 +52,13 @@ PROVIDERS: List[Dict] = [ { "id": "anthropic", "name": "Anthropic", - "description": "Claude Fable 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", + "description": "Claude Fable 5, Opus 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 5, Sonnet 4.6, Haiku 4.5", "env_key": "ANTHROPIC_API_KEY", "key_hint": "sk-ant-...", "test_model": "claude-haiku-4-5-20251001", "models": [ "claude-fable-5", + "claude-opus-5", "claude-sonnet-5", "claude-opus-4-8", "claude-opus-4-7", diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c86794b90f8..a324e71e289 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -910,6 +910,17 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up ), ) + run_in_parallel: Optional[bool] = Field( + default=None, + description=( + "When True, this pre_call or post_call guardrail runs concurrently with other opted-in " + "guardrails of the same hook, after the sequential guardrails have run. Use only for " + "block-only guardrails that inspect and reject; do not enable it for guardrails that " + "modify the request or response (e.g. PII masking or sensitive-data routing), since " + "parallel runs share one snapshot and their mutations would race." + ), + ) + @field_validator( "mode", "default_action", diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 8ae974b19a6..b0af22e7c3f 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -261,12 +261,3 @@ class MCPServer(BaseModel): if self.oauth_passthrough is not True: return False return any(h.lower() == "authorization" for h in self.extra_headers) - - @property - def has_token_exchange_config(self) -> bool: - """True if this server is configured for OAuth2 token exchange (OBO / RFC 8693).""" - return ( - self.auth_type == MCPAuth.oauth2_token_exchange - and bool(self.client_id and self.client_secret) - and bool(self.token_exchange_endpoint or self.token_url) - ) diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 742e0f7818f..d4b1d98f957 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -153,6 +153,24 @@ class SSOConfig(LiteLLMPydanticObjectBase): description="Space-separated OAuth scopes requested from the generic provider, e.g. 'openid email profile'", ) + # SAML SSO + saml_idp_metadata_url: Optional[str] = Field( + default=None, + description="URL of the SAML IdP metadata to fetch and parse for SSO authentication", + ) + saml_idp_metadata_xml: Optional[str] = Field( + default=None, + description="Inline SAML IdP metadata XML, used when a metadata URL is not available", + ) + saml_sp_entity_id: Optional[str] = Field( + default=None, + description="SAML Service Provider entityID; defaults to the proxy's /sso/saml/metadata URL", + ) + saml_allow_unsolicited: Optional[str] = Field( + default=None, + description="'true' to accept IdP-initiated (unsolicited) SAML responses, which cannot be browser-bound against login CSRF", + ) + # Common settings proxy_base_url: Optional[str] = Field( default=None, diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index 1c5e1df9e9a..71ec412e8ef 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -98,3 +98,38 @@ class ToolUsageLogsResponse(BaseModel): total: int page: int page_size: int + + +class ToolSpendEntry(BaseModel): + """Total spend attributed to one tool over the requested window.""" + + tool_name: str + spend: float = Field( + 0.0, + description="Attributed spend: a request that used several tools counts its full spend toward each of them", + ) + call_count: int = 0 + total_tokens: int = 0 + + +class ToolSpendDailyEntry(BaseModel): + """Spend attributed to one tool on one UTC day.""" + + date: str + tool_name: str + spend: float = 0.0 + call_count: int = 0 + + +class ToolSpendResponse(BaseModel): + by_tool: List[ToolSpendEntry] = Field(default_factory=list) + daily: List[ToolSpendDailyEntry] = Field(default_factory=list) + total_spend: float = Field( + 0.0, + description=( + "Deduplicated spend of every request that called at least one tool in the window; " + "less than the sum of per-tool attributed spend whenever multi-tool requests exist" + ), + ) + start_date: str | None = None + end_date: str | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 714ad372a5f..e4dfac48141 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1534,14 +1534,27 @@ class PromptTokensDetailsWrapper( audio_length_seconds: Optional[float] = None """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + cache_write_tokens: Optional[int] = None + """Number of cache write (creation) tokens sent to the model. OpenAI naming (prompt_tokens_details.cache_write_tokens); this is the canonical field.""" + cache_creation_tokens: Optional[int] = None - """Number of cache creation tokens sent to the model. Used for Anthropic prompt caching.""" + """Number of cache creation tokens sent to the model. Anthropic/Bedrock naming; kept in sync with cache_write_tokens (assigning either mirrors to the other).""" cache_creation_token_details: Optional[CacheCreationTokenDetails] = None """Details of cache creation tokens sent to the model. Used for tracking 5m/1h cache creation tokens for Anthropic prompt caching.""" + def __setattr__(self, name: str, value: object) -> None: + super().__setattr__(name, value) + if name == "cache_write_tokens": + super().__setattr__("cache_creation_tokens", value) + elif name == "cache_creation_tokens": + super().__setattr__("cache_write_tokens", value) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self.cache_write_tokens = ( + self.cache_write_tokens if self.cache_write_tokens is not None else self.cache_creation_tokens + ) if self.character_count is None: del self.character_count if self.image_count is None: @@ -1554,6 +1567,8 @@ class PromptTokensDetailsWrapper( del self.web_search_requests if self.tool_use_tokens is None: del self.tool_use_tokens + if self.cache_write_tokens is None: + del self.cache_write_tokens if self.cache_creation_tokens is None: del self.cache_creation_tokens if self.cache_creation_token_details is None: @@ -1662,10 +1677,10 @@ class Usage(SafeAttributeModel, CompletionUsage): if "cache_creation_input_tokens" in params and isinstance(params["cache_creation_input_tokens"], int): if _prompt_tokens_details is None: _prompt_tokens_details = PromptTokensDetailsWrapper( - cache_creation_tokens=params["cache_creation_input_tokens"] + cache_write_tokens=params["cache_creation_input_tokens"] ) else: - _prompt_tokens_details.cache_creation_tokens = params["cache_creation_input_tokens"] + _prompt_tokens_details.cache_write_tokens = params["cache_creation_input_tokens"] super().__init__( prompt_tokens=prompt_tokens or 0, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c9d871fc41d..749b2566c2a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1502,6 +1502,222 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024 }, + "anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "global.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "us.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "eu.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "au.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "jp.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, @@ -2756,6 +2972,38 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true }, + "azure_ai/claude-opus-5": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "azure_ai/claude-opus-4-8": { "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, @@ -11846,6 +12094,44 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512 }, + "claude-opus-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1, + "fast": 2.0 + }, + "supports_output_config": true, + "supports_speed": true, + "prompt_cache_min_tokens": 512 + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -36987,6 +37273,70 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true }, + "vertex_ai/claude-opus-5": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, + "vertex_ai/claude-opus-5@default": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "vertex_ai/claude-opus-4-8": { "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, diff --git a/pyproject.toml b/pyproject.toml index a448ab042b8..44c1967ad9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,6 +99,10 @@ utils = [ "numpydoc>=1.8.0,<2.0", ] caching = ["diskcache>=5.6.3,<6.0"] +# SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels +# bundle the native libxmlsec1/libxml2 libraries, so no system packages are +# required. Kept out of the base `proxy` extra so it stays optional. +saml = ["python3-saml>=1.16.0,<2.0"] semantic-router = [ "semantic-router>=0.1.15,<1.0; python_version < '3.14'", "aurelio-sdk>=0.0.19,<1.0; python_version < '3.14'", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index d3d70ff5ff4..39e9bf4773d 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -222,7 +222,7 @@ "limit": 38 }, "RET504": { - "limit": 721 + "limit": 719 }, "RUF010": { "limit": 874 @@ -324,7 +324,7 @@ "limit": 883 }, "UP006": { - "limit": 12792 + "limit": 12789 }, "UP007": { "limit": 2570 diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index 97ffa8c34a3..916ef623d3a 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -85,6 +85,7 @@ class A2ABridgeParams(BaseModel): custom_llm_provider: str model: str + api_key: str | None = None class AgentRegisterBody(BaseModel): @@ -190,11 +191,52 @@ class A2ATaskStatus(BaseModel): message: A2AResponseMessage | None = None +class A2AListingLocation(BaseModel): + """Only the location fields a test reads back off a returned listing.""" + + un_locode: str | None = None + + +class A2AListing(BaseModel): + """A single property card from the agent's `search_results` artifact; only the + identity/location fields a test asserts on are modelled.""" + + raia_id: str + property_type: str | None = None + service_type: str | None = None + location: A2AListingLocation = A2AListingLocation() + + +class A2ASearchResults(BaseModel): + """The DataPart payload the property agent's `search_properties` skill returns: + the run count plus the listing cards themselves. Proof the tool actually ran and + matched, not just that the task completed with some text.""" + + total: int + count: int + listings: list[A2AListing] = [] + + +class A2AArtifactPart(BaseModel): + kind: str | None = None + data: A2ASearchResults | None = None + + +class A2AArtifact(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + artifact_id: str | None = Field(default=None, alias="artifactId") + name: str | None = None + parts: list[A2AArtifactPart] = [] + + class A2AResult(BaseModel): """A message/send result. In 0.3 the message fields sit directly on the result (`kind`/`role`/`parts`); in 1.0 they are nested under `message`; a real agent that - runs a task replies with a `task` whose agent text lives on `status.message`. - `text` reads the agent's reply from whichever shape the served version produced.""" + runs a task replies with a `task` whose agent text lives on `status.message` and + whose tool output lives on `artifacts`. `text` reads the agent's reply from + whichever shape the served version produced; `search_results` reads the tool's + structured output when the agent ran a skill.""" model_config = ConfigDict(populate_by_name=True) @@ -204,6 +246,7 @@ class A2AResult(BaseModel): parts: list[A2AResponsePart] = [] message: A2AResponseMessage | None = None status: A2ATaskStatus | None = None + artifacts: list[A2AArtifact] = [] @property def text(self) -> str: @@ -221,6 +264,14 @@ class A2AResult(BaseModel): def is_nested_v1_shape(self) -> bool: return self.message is not None + @property + def search_results(self) -> A2ASearchResults | None: + for artifact in self.artifacts: + for part in artifact.parts: + if part.data is not None: + return part.data + return None + class A2AError(BaseModel): code: int diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py index aa60b57f99b..802b01455b1 100644 --- a/tests/e2e/a2a/test_a2a_agent_e2e.py +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -32,7 +32,11 @@ from e2e_config import unique_marker from e2e_http import Result, UnknownApiError, unwrap from lifecycle import ResourceManager -BRIDGE = A2ABridgeParams(custom_llm_provider="anthropic", model="claude-haiku-4-5") +BRIDGE = A2ABridgeParams( + custom_llm_provider="anthropic", + model="claude-haiku-4-5", + api_key="os.environ/ANTHROPIC_API_KEY", +) MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json" MOVEHOME_ORIGIN = "https://movehome.org" @@ -113,6 +117,7 @@ class TestA2AAgentLifecycle: agent = unwrap(client.register_agent(body)) resources.defer(lambda: client.delete_agent(agent.agent_id)) assert agent.agent_card_params.protocol_version == "0.3" + location = "GBLON" request = A2AJsonRpcRequest( id=f"e2e-{unique_marker()}", params=A2AMessageSendParams( @@ -121,7 +126,7 @@ class TestA2AAgentLifecycle: A2ADataPart( data=A2ASkillInvocation( skill="search_properties", - params=A2ASearchPropertiesParams(un_locode="USSFO", service_type="sale", asking_price_max=2_000_000, limit=3), + params=A2ASearchPropertiesParams(un_locode=location, service_type="long_term", limit=3), ) ) ], @@ -132,7 +137,12 @@ class TestA2AAgentLifecycle: response = unwrap(client.send_message(agent.agent_id, scoped_key, request)) assert response.error is None assert response.result is not None - assert response.result.text.strip() != "" + results = response.result.search_results + assert results is not None, "agent returned no search_results artifact; skill did not run" + assert results.total > 0 + assert results.listings, "search_properties matched nothing; agent returned no property cards" + assert all(listing.raia_id for listing in results.listings) + assert all(listing.location.un_locode == location for listing in results.listings) @pytest.mark.covers("other.a2a.discovery.proxy_fronted_card") def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index b0c53becb6b..f9cd2a3f15f 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -23,7 +23,7 @@ from typing import Callable import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from batch_client import ( UPLOAD_FILENAME, @@ -702,14 +702,7 @@ class TestBedrockBatchAssumeRole: def test_unified_batch_create_with_assume_role( self, client: BatchClient, resources: ResourceManager ) -> None: - (role_arn,) = require_env("AWS_ROLE_NAME") - require_env( - "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", - "AWS_REGION", - "AWS_BATCH_S3_BUCKET", - "AWS_BATCH_ROLE_ARN", - ) + role_arn = os.environ["AWS_ROLE_NAME"] session_name = f"e2e-batch-sts-{unique_marker()}"[:64] model_name = batch_model_name("bedrock-sts-batch") @@ -819,7 +812,7 @@ class TestHostedVllmBatch: def test_unified_file_and_batch_create( self, client: BatchClient, resources: ResourceManager ) -> None: - (api_base,) = require_env("HOSTED_VLLM_API_BASE") + api_base = os.environ["HOSTED_VLLM_API_BASE"] api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None model_id = ( os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct" diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 4ecc215a22d..feed680bd4a 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -102,22 +102,6 @@ ANOMALY_SPEND_SETTLE_SECONDS = float( ) -def require_env(*names: str) -> tuple[str, ...]: - """Return the non-empty values for each env name, or hard-fail naming which are missing. - - Live e2e never skips for missing credentials: a missing key is a red run so - ops knows the suite cannot prove the product path. - """ - missing = tuple(name for name in names if not (os.environ.get(name) or "").strip()) - if missing: - joined = ", ".join(missing) - raise AssertionError( - f"missing required env for e2e: {joined}. " - "Add them to tests/e2e/.env locally and to litellm ops for stage/CI." - ) - return tuple((os.environ.get(name) or "").strip() for name in names) - - def datadog_mcp_url(*, toolsets: str = "core") -> str: """Regional Datadog remote MCP endpoint for this process's DD_SITE. diff --git a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py index 9e41b8808e8..a2408f0021e 100644 --- a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py @@ -8,9 +8,11 @@ a 200 means the guardrail never ran. from __future__ import annotations +import os + import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import UnknownApiError from guardrails_client import GuardrailsClient from lifecycle import ResourceManager @@ -33,11 +35,8 @@ class TestBedrockGuardrail: def test_bedrock_pre_call_blocks_harmful_prompt( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - (identifier, version) = require_env( - "BEDROCK_GUARDRAIL_IDENTIFIER", - "BEDROCK_GUARDRAIL_VERSION", - ) - require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") + identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"] + version = os.environ["BEDROCK_GUARDRAIL_VERSION"] name = f"e2e-bedrock-guard-{unique_marker()}" guardrail_id = client.create_bedrock_guardrail( diff --git a/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py b/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py index e36fc7c3f9d..de087b190d0 100644 --- a/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py @@ -16,7 +16,7 @@ from __future__ import annotations import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import unwrap from guardrails_client import BlockCodeExecutionParamsBody, GuardrailsClient from lifecycle import ResourceManager @@ -46,7 +46,6 @@ class TestBlockCodeExecutionGuardrail: def test_blocks_execution_request_but_allows_explanation( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("GEMINI_API_KEY") model = client.create_backend_model(resources, prefix="e2e-blockcode-backend") name = f"e2e-block-code-{unique_marker()}" diff --git a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py index 4e2fcbf8fba..39950259fb5 100644 --- a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py @@ -14,7 +14,7 @@ from __future__ import annotations import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import UnknownApiError, unwrap from guardrails_client import GuardrailsClient, OpenAIModerationParamsBody from lifecycle import ResourceManager @@ -34,7 +34,6 @@ class TestOpenAIModerationGuardrail: def test_moderation_blocks_flagged_input( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("OPENAI_API_KEY", "GEMINI_API_KEY") model = client.create_backend_model(resources, prefix="e2e-moderation-backend") name = f"e2e-openai-moderation-{unique_marker()}" diff --git a/tests/e2e/guardrails/test_presidio_guardrail_e2e.py b/tests/e2e/guardrails/test_presidio_guardrail_e2e.py index a911f387382..d103714b1dd 100644 --- a/tests/e2e/guardrails/test_presidio_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_presidio_guardrail_e2e.py @@ -25,11 +25,12 @@ The chat backend is a gemini deployment created for the test. from __future__ import annotations +import os import time import pytest -from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, require_env, unique_marker +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker from e2e_http import NoBody, require_successful_call, unwrap from guardrails_client import GuardrailMode, GuardrailsClient, PresidioParamsBody from lifecycle import ResourceManager @@ -88,9 +89,8 @@ def _poll_logged_prompt(reader: OtelReader, *, call_id: str, genai_span: str) -> def _presidio_params( mode: GuardrailMode, *, apply_to_output: bool = False, logging_only: bool = False ) -> PresidioParamsBody: - analyzer, anonymizer = require_env( - "PRESIDIO_ANALYZER_API_BASE", "PRESIDIO_ANONYMIZER_API_BASE" - ) + analyzer = os.environ["PRESIDIO_ANALYZER_API_BASE"] + anonymizer = os.environ["PRESIDIO_ANONYMIZER_API_BASE"] return PresidioParamsBody( mode=mode, default_on=False, @@ -124,7 +124,6 @@ class TestPresidioGuardrail: def test_pre_call_masks_pii_before_the_model_sees_it( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("GEMINI_API_KEY") model = client.create_backend_model(resources, prefix="e2e-presidio-pre") name = f"e2e-presidio-pre-{unique_marker()}" guardrail_id = client.register(name, _presidio_params("pre_call")) @@ -149,7 +148,6 @@ class TestPresidioGuardrail: def test_post_call_masks_pii_in_model_output( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("GEMINI_API_KEY") model = client.create_backend_model(resources, prefix="e2e-presidio-post") name = f"e2e-presidio-post-{unique_marker()}" guardrail_id = client.register(name, _presidio_params("post_call", apply_to_output=True)) @@ -173,7 +171,6 @@ class TestPresidioGuardrail: def test_logging_only_masks_the_logged_prompt( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("GEMINI_API_KEY") _require_otel_v2_active(client) reader = build_otel_reader() diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index 8d3622e441a..af0e782e224 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -21,7 +21,7 @@ import os import pytest from pydantic import BaseModel -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import StreamingResponse, unwrap from lifecycle import ResourceManager from models import ( @@ -250,7 +250,7 @@ class TestCohereChat: def test_cohere_chat_returns_content( self, client: PassthroughClient, resources: ResourceManager ) -> None: - (cohere_key,) = require_env("COHERE_API_KEY") + cohere_key = os.environ["COHERE_API_KEY"] model = f"e2e-cohere-chat-{unique_marker()}" model_id = client.proxy.create_model( model, @@ -343,7 +343,7 @@ class TestHostedVllmChat: def test_hosted_vllm_chat_returns_content( self, client: PassthroughClient, resources: ResourceManager ) -> None: - (api_base,) = require_env("HOSTED_VLLM_API_BASE") + api_base = os.environ["HOSTED_VLLM_API_BASE"] api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None backend = ( os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct" @@ -395,7 +395,6 @@ class TestOpenAIChatCompletions: def test_openai_chat_streams_real_content( self, client: PassthroughClient, resources: ResourceManager ) -> None: - require_env("OPENAI_API_KEY") model = f"e2e-openai-chat-{unique_marker()}" model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") @@ -423,7 +422,6 @@ class TestOpenAIChatCompletions: def test_openai_chat_logs_cost( self, client: PassthroughClient, resources: ResourceManager ) -> None: - require_env("OPENAI_API_KEY") model = f"e2e-openai-cost-{unique_marker()}" model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") @@ -457,7 +455,6 @@ class TestOpenAIChatCompletions: def test_openai_chat_returns_tool_call( self, client: PassthroughClient, resources: ResourceManager ) -> None: - require_env("OPENAI_API_KEY") model = f"e2e-openai-tool-{unique_marker()}" model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") @@ -488,7 +485,6 @@ class TestOpenAIChatCompletions: def test_openai_chat_structured_output_conforms_to_schema( self, client: PassthroughClient, resources: ResourceManager ) -> None: - require_env("OPENAI_API_KEY") model = f"e2e-openai-schema-{unique_marker()}" model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") @@ -522,7 +518,6 @@ class TestOpenAIChatCompletions: def test_openai_chat_reasoning_reports_reasoning_tokens( self, client: PassthroughClient, resources: ResourceManager ) -> None: - require_env("OPENAI_API_KEY") model = f"e2e-openai-reasoning-{unique_marker()}" model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") @@ -561,7 +556,6 @@ class TestOpenAIChatCompletions: def test_openai_chat_vision_describes_image( self, client: PassthroughClient, resources: ResourceManager ) -> None: - require_env("OPENAI_API_KEY") model = f"e2e-openai-vision-{unique_marker()}" model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=OPENAI_VISION_BACKEND, api_key="os.environ/OPENAI_API_KEY") @@ -579,7 +573,6 @@ class TestOpenAIChatCompletions: def test_openai_chat_prompt_cache_hits_on_repeat( self, client: PassthroughClient, resources: ResourceManager ) -> None: - require_env("OPENAI_API_KEY") model = f"e2e-openai-cache-{unique_marker()}" model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") @@ -610,7 +603,6 @@ class TestOpenAIChatCompletions: def test_openai_chat_streams_tool_call( self, client: PassthroughClient, resources: ResourceManager ) -> None: - require_env("OPENAI_API_KEY") model = f"e2e-openai-tool-stream-{unique_marker()}" model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") @@ -645,7 +637,6 @@ class TestBedrockConverseChatCompletions: """ def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str: - require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") model = f"{prefix}-{unique_marker()}" model_id = client.proxy.create_model(model, _bedrock_params()) resources.defer(lambda: client.proxy.delete_model(model_id)) diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index 1ba78a7e083..45861d1e93a 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -9,7 +9,7 @@ from __future__ import annotations import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from endpoints_client import EndpointsClient, ImagesResult from lifecycle import ResourceManager @@ -50,7 +50,6 @@ class TestImageGeneration: def test_bedrock_image_generation_returns_image( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") model = f"e2e-bedrock-image-{unique_marker()}" model_id = endpoints_client.create_model( model, diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 44376218c6b..ef6ba5b95d3 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -10,7 +10,7 @@ from __future__ import annotations import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call, unwrap from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager @@ -73,7 +73,6 @@ class TestAnthropicMessages: def test_messages_logs_cost_matching_the_response_header( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - require_env("ANTHROPIC_API_KEY") model = f"e2e-messages-cost-{unique_marker()}" model_id = endpoints_client.create_model( model, diff --git a/tests/e2e/llm_translation/test_rerank_e2e.py b/tests/e2e/llm_translation/test_rerank_e2e.py index 0857ff65a52..c3614251e77 100644 --- a/tests/e2e/llm_translation/test_rerank_e2e.py +++ b/tests/e2e/llm_translation/test_rerank_e2e.py @@ -9,7 +9,7 @@ from __future__ import annotations import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from endpoints_client import EndpointsClient, RerankResult from lifecycle import ResourceManager @@ -56,7 +56,6 @@ class TestRerank: def test_bedrock_rerank_scores_top_n( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") model = f"e2e-bedrock-rerank-{unique_marker()}" model_id = endpoints_client.create_model( model, diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index d24d2b53b71..0b2ffce5b2a 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -13,7 +13,7 @@ from typing import cast import pytest from pydantic import BaseModel, ValidationError -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from endpoints_client import ( EndpointsClient, @@ -255,7 +255,6 @@ class TestResponses: def test_responses_bedrock_returns_completion( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") model = f"e2e-responses-{unique_marker()}" model_id = endpoints_client.create_model(model, _bedrock_params()) resources.defer(lambda: endpoints_client.delete_model(model_id)) @@ -270,7 +269,6 @@ class TestResponses: def test_responses_bedrock_returns_function_call( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") model = f"e2e-responses-{unique_marker()}" model_id = endpoints_client.create_model(model, _bedrock_params()) resources.defer(lambda: endpoints_client.delete_model(model_id)) diff --git a/tests/e2e/llm_translation/test_responses_metadata_e2e.py b/tests/e2e/llm_translation/test_responses_metadata_e2e.py index 6cf24348095..df854dcfa19 100644 --- a/tests/e2e/llm_translation/test_responses_metadata_e2e.py +++ b/tests/e2e/llm_translation/test_responses_metadata_e2e.py @@ -14,7 +14,7 @@ import time import pytest from pydantic import BaseModel, ConfigDict -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from endpoints_client import EndpointsClient, ResponsesResult from lifecycle import ResourceManager @@ -42,7 +42,7 @@ class RedisKeyInfo(BaseModel): def _redis_scan(marker: str) -> tuple[RedisKeyInfo, ...]: import redis - (host,) = require_env("REDIS_HOST") + host = os.environ["REDIS_HOST"] port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") try: with socket.create_connection((host, port), timeout=3): diff --git a/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py b/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py index ed6f0ce3b2c..a88f0ca546a 100644 --- a/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py +++ b/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py @@ -11,7 +11,7 @@ import socket import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from lifecycle import ResourceManager from models import KeyGenerateBody, LiteLLMParamsBody @@ -23,7 +23,7 @@ BACKEND = "anthropic/claude-haiku-4-5-20251001" def _require_redis_reachable() -> None: - (host,) = require_env("REDIS_HOST") + host = os.environ["REDIS_HOST"] port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") try: with socket.create_connection((host, port), timeout=3): diff --git a/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py b/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py index b509ae000f5..3e1bc662470 100644 --- a/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py +++ b/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py @@ -13,7 +13,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from lifecycle import ResourceManager from models import KeyGenerateBody, LiteLLMParamsBody @@ -28,7 +28,7 @@ RECOVERY_TIMEOUT = float( def _require_redis() -> None: - (host,) = require_env("REDIS_HOST") + host = os.environ["REDIS_HOST"] port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") try: with socket.create_connection((host, port), timeout=3): diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index 17ff1fc3f83..bb8806a9c01 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -1,5 +1,5 @@ import { test, expect } from "@playwright/test"; -import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_CRUD_ID } from "../../constants"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; import { Role, users } from "../../fixtures/users"; import { navigateToPage } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; @@ -56,17 +56,17 @@ test.describe("Add Model", () => { }, }); expect(createResponse.ok()).toBe(true); + const createdModelId = (await createResponse.json()).model_info?.id; + expect(createdModelId, "model id from /model/new").toBeTruthy(); // Navigate to Models + Endpoints await page.goto("/ui"); await page.getByText("Models + Endpoints").click(); - // Click the new model row to open its detail view. The table renders - // a clickable outer row plus a nested detail row for the same model, - // so we target the first match (outer row) explicitly. - const modelRow = page.locator("tr", { hasText: modelName }).first(); - await expect(modelRow).toBeVisible({ timeout: 10_000 }); - await modelRow.click(); + // The Model ID cell is the drill-in control; the row itself is not clickable. + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect(modelIdCell).toBeVisible({ timeout: 10_000 }); + await modelIdCell.click(); await expect(page.getByText("Back to Models").first()).toBeVisible({ timeout: 10_000 }); @@ -137,11 +137,11 @@ test.describe("Add Model", () => { await page.waitForTimeout(2000); // Search for the model we just added - await page.locator('input[placeholder="Search model names..."]').fill("claude-haiku-4-5"); + await page.getByPlaceholder("Search model names").fill("claude-haiku-4-5"); await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") - await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { timeout: 15_000, }); @@ -228,24 +228,24 @@ test.describe("Add Model", () => { // searching. await page.waitForTimeout(2000); - await page.locator('input[placeholder="Search model names..."]').fill("cohere"); + await page.getByPlaceholder("Search model names").fill("cohere"); await page.waitForTimeout(1000); // Confirm the search returned at least one result — gives a clear // failure message when the table is empty instead of timing out on a // row assertion. - await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { timeout: 15_000, }); - // Stronger than "alias appears somewhere in tbody" — pin the assertion + // Stronger than "the team appears somewhere in tbody" — pin the assertion // to a single row that has BOTH the cohere model_name AND the seeded - // team alias, so a stale cohere row from "Add wildcard route" (no team) - // can't satisfy the check. + // team, so a stale cohere row from "Add wildcard route" (no team) can't + // satisfy the check. The Team ID column renders the id, not the alias. const teamCohereRow = page .locator("table tbody tr") .filter({ hasText: "cohere/" }) - .filter({ hasText: E2E_TEAM_CRUD_ALIAS }); + .filter({ hasText: E2E_TEAM_CRUD_ID }); await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 }); } finally { await deleteTeamScopedCohereModels(); @@ -281,11 +281,11 @@ test.describe("Add Model", () => { await page.waitForTimeout(2000); // Search for the wildcard model - await page.locator('input[placeholder="Search model names..."]').fill("cohere"); + await page.getByPlaceholder("Search model names").fill("cohere"); await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") - await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { timeout: 15_000, }); diff --git a/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts index 877c7f8c555..e67dcb96f36 100644 --- a/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts @@ -67,9 +67,10 @@ test.describe("Clear custom pricing on a deployment", () => { await page.goto("/ui"); await page.getByText("Models + Endpoints").click(); - const modelRow = page.locator("tr", { hasText: modelName }).first(); - await expect(modelRow).toBeVisible({ timeout: 15_000 }); - await modelRow.click(); + // The Model ID cell is the drill-in control; the row itself is not clickable. + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect(modelIdCell).toBeVisible({ timeout: 15_000 }); + await modelIdCell.click(); await expect(page.getByText("Back to Models").first()).toBeVisible({ timeout: 10_000, }); diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index c47fddb1d8d..4fa77f38940 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -167,6 +167,13 @@ ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = ( "once the model is available." ), ), + ModelEntry( + alias="claude-opus-5", + model="anthropic/claude-opus-5", + mode="adaptive", + required_env=_ANTHROPIC_REQ, + caps=_CAPS_XHIGH_MAX, + ), ModelEntry( alias="claude-opus-4-8", model="anthropic/claude-opus-4-8", diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index 1fec7665daa..440fc36ed33 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -518,7 +518,7 @@ async def test_openai_codex_stream(sync_mode): from litellm.main import stream_chunk_builder kwargs = { - "model": "openai/gpt-5.2-codex", + "model": "openai/gpt-5.3-codex", "messages": [{"role": "user", "content": "Hey!"}], "stream": True, } @@ -550,7 +550,7 @@ async def test_openai_codex(sync_mode): { "model_name": "openai-codex-mini-latest", "litellm_params": { - "model": "openai/gpt-5.2-codex", + "model": "openai/gpt-5.3-codex", }, } ] @@ -838,7 +838,7 @@ def test_gpt_5_reasoning_streaming(): def test_openai_gpt_5_codex_reasoning(): litellm._turn_on_debug() completion_kwargs = { - "model": "gpt-5-codex", + "model": "gpt-5.3-codex", "messages": [ { "role": "system", diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index ee18c96c393..d2218b08386 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -7,10 +7,12 @@ from typing import Any, Dict, List, Optional, Union from unittest.mock import Mock import pytest -from fastapi import Request +from fastapi import HTTPException, Request from starlette.datastructures import State +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.utils import _get_docs_url, _get_openapi_url, _get_redoc_url +from litellm.types.guardrails import GuardrailEventHooks sys.path.insert( 0, os.path.abspath("../..") @@ -2638,6 +2640,463 @@ async def test_during_call_hook_parallel_execution_with_error(): litellm.callbacks = original_callbacks +class _PreCallGuardrail(CustomGuardrail): + """Test double for pre_call guardrails; records timing and observed payload.""" + + def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True): + super().__init__( + guardrail_name=name, + event_hook=GuardrailEventHooks.pre_call, + default_on=default_on, + run_in_parallel=run_in_parallel, + ) + self.name = name + self.sleep = sleep + self.execution_order = execution_order + self.observed_content = None + self.was_called = False + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.was_called = True + self.observed_content = data["messages"][0]["content"] + self.execution_order.append(f"{self.name}_start") + await asyncio.sleep(self.sleep) + self.execution_order.append(f"{self.name}_end") + return None + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_opted_in_guardrails_in_parallel(): + """run_in_parallel pre_call guardrails execute concurrently (all start before any ends).""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PreCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3) + ] + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item) + starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item) + assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}" + assert result["model"] == "gpt-4" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_default_guardrails_sequentially(): + """Guardrails without run_in_parallel keep the sequential, one-at-a-time behavior.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PreCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2) + ] + + start = asyncio.get_event_loop().time() + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + elapsed = asyncio.get_event_loop().time() - start + + assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"] + assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_sequential_mutations_precede_parallel_batch(): + """Sequential (mutating) guardrails run before the parallel batch, which sees their changes.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class MaskingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="masker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=False, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + data["messages"][0]["content"] = "MASKED" + return data + + parallel_observer = _PreCallGuardrail("observer", run_in_parallel=True, execution_order=execution_order) + + try: + litellm.callbacks = [parallel_observer, MaskingGuardrail()] + + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "secret"}]}, + call_type="completion", + ) + + assert parallel_observer.observed_content == "MASKED" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_guardrail_blocks_request(): + """A raising parallel guardrail blocks the request before it reaches the LLM.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class BlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise HTTPException(status_code=400, detail="blocked by guardrail") + + try: + litellm.callbacks = [BlockingGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "blocked by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_guardrail_skipped_when_should_not_run(): + """A parallel guardrail that should_run_guardrail rejects is never invoked.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + guardrail = _PreCallGuardrail( + "off_by_default", run_in_parallel=True, execution_order=execution_order, default_on=False + ) + litellm.callbacks = [guardrail] + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert guardrail.was_called is False + assert result["model"] == "gpt-4" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_block_wins_over_reroute(): + """A slower block must win over a faster reroute so crafted input cannot bypass a block.""" + from litellm.caching.caching import DualCache + from litellm.exceptions import SensitiveDataRouteException + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class FastRerouteGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="rerouter", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise SensitiveDataRouteException(route_to_model="on-prem", session_id="s1", guardrail_name="rerouter") + + class SlowBlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + await asyncio.sleep(0.1) + raise HTTPException(status_code=400, detail="blocked by guardrail") + + try: + litellm.callbacks = [FastRerouteGuardrail(), SlowBlockingGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "blocked by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_awaits_all_when_one_blocks(): + """A block must not orphan sibling guardrails; every parallel guardrail runs to completion.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + completed = [] + + class FastBlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="fast_blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise HTTPException(status_code=400, detail="blocked") + + class SlowGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="slow", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + await asyncio.sleep(0.1) + completed.append("slow") + return None + + try: + litellm.callbacks = [FastBlockingGuardrail(), SlowGuardrail()] + + with pytest.raises(HTTPException): + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert completed == ["slow"], "slow guardrail was orphaned instead of awaited to completion" + finally: + litellm.callbacks = original_callbacks + + +class _PostCallGuardrail(CustomGuardrail): + """Test double for post_call guardrails; records timing and invocation.""" + + def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True): + super().__init__( + guardrail_name=name, + event_hook=GuardrailEventHooks.post_call, + default_on=default_on, + run_in_parallel=run_in_parallel, + ) + self.name = name + self.sleep = sleep + self.execution_order = execution_order + self.was_called = False + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.was_called = True + self.execution_order.append(f"{self.name}_start") + await asyncio.sleep(self.sleep) + self.execution_order.append(f"{self.name}_end") + return None + + +@pytest.mark.asyncio +async def test_post_call_hook_runs_opted_in_guardrails_in_parallel(): + """run_in_parallel post_call guardrails execute concurrently (all start before any ends).""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PostCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3) + ] + + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item) + starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item) + assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_runs_default_guardrails_sequentially(): + """post_call guardrails without run_in_parallel keep the sequential behavior.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PostCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2) + ] + + start = asyncio.get_event_loop().time() + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + elapsed = asyncio.get_event_loop().time() - start + + assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"] + assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_parallel_guardrail_blocks_response(): + """A raising parallel post_call guardrail blocks the response before it reaches the client.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class BlockingPostCallGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="post_blocker", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + raise HTTPException(status_code=400, detail="blocked response by guardrail") + + try: + litellm.callbacks = [BlockingPostCallGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + assert exc_info.value.status_code == 400 + assert "blocked response by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_parallel_awaits_all_when_one_blocks(): + """A blocking post_call guardrail must not orphan its siblings; all run to completion.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + completed = [] + + class FastBlockingPostCall(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="fast_post_blocker", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + raise HTTPException(status_code=400, detail="blocked") + + class SlowPostCall(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="slow_post", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + await asyncio.sleep(0.1) + completed.append("slow") + return None + + try: + litellm.callbacks = [FastBlockingPostCall(), SlowPostCall()] + + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + assert completed == ["slow"], "slow post_call guardrail was orphaned instead of awaited to completion" + finally: + litellm.callbacks = original_callbacks + + @pytest.mark.asyncio async def test_handle_logging_proxy_only_error_preserves_pass_through_call_type(): """Ensure _handle_logging_proxy_only_error does not overwrite call_type diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 5f6002f4cdf..d3593f2c06b 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -29,7 +29,9 @@ from litellm.integrations.otel import ( # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.plumbing.context import ( # noqa: E402 reset_mcp_message_trace_carrier, + reset_mcp_message_transport_span_context, set_mcp_message_trace_carrier, + set_mcp_message_transport_span_context, set_request_root_span, ) from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402 @@ -56,9 +58,11 @@ def _reset_request_root_span(): _otel_context._request_root_span.set(None) _otel_context._mcp_message_trace_carrier.set(None) + _otel_context._mcp_message_transport_span_context.set(None) yield _otel_context._request_root_span.set(None) _otel_context._mcp_message_trace_carrier.set(None) + _otel_context._mcp_message_transport_span_context.set(None) def _payload(**overrides): @@ -528,15 +532,15 @@ _MCP_SPAN_CASES = [ @pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) -def test_mcp_span_roots_and_links_transport_without_propagated_context( +def test_mcp_span_nests_under_transport_without_propagated_context( make_payload, span_name ): - """MCP and the HTTP transport are independent lifecycles (one streamable-HTTP - session multiplexes many messages), so per the MCP semconv the message span - must NOT nest under the session/transport span — that is what made it render - skewed at the session's start. With no propagated ``params._meta`` context it - starts its own root trace and records the transport span as a *link*, never - the parent.""" + """Almost no MCP client implements SEP-414, so ``params._meta`` normally carries + no trace context. Rooting the span there split one tool call into two traces + joined only by a link, which is how it surfaced in APM: the ``POST`` transaction + and the ``tools/call`` span shared no ``trace_id``. With no remote parent to + honor the span nests under the transport span instead, and records no link since + the transport is now the real parent.""" logger, exporter = _logger() transport = logger._emitter.start_span( SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME @@ -549,11 +553,75 @@ def test_mcp_span_roots_and_links_transport_without_propagated_context( ) transport.end() span = next(s for s in exporter.get_finished_spans() if s.name == span_name) + assert span.parent is not None + assert span.parent.span_id == transport.get_span_context().span_id + assert span.context.trace_id == transport.get_span_context().trace_id + assert span.links == () + + +@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) +def test_mcp_span_nests_under_this_messages_transport_not_the_session_opener( + make_payload, span_name +): + """A *stateful* streamable-HTTP session runs every message on the single task + spawned by that session's ``initialize`` POST, so the ``_request_root_span`` + ContextVar the ASGI request task writes is frozen at ``initialize`` inside the + handler and never sees the later ``tools/call`` POST. Nesting on that anchor + would hang every tool call of the session off the first request's (already + ended) span, rendering skewed at the session's start. The gateway resolves the + current message's transport on the request task and publishes it, so the span + parents to the POST that actually carried this message.""" + logger, exporter = _logger() + session_opener = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + this_message = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + + async def session_task(): + token = set_mcp_message_transport_span_context( + this_message.get_span_context() + ) + try: + await logger.async_log_success_event( + {"standard_logging_object": make_payload()}, None, None, None + ) + finally: + reset_mcp_message_transport_span_context(token) + + async def initialize_request(): + # The anchor the session task inherits is the one ``initialize`` left behind; + # spawning here reproduces the SDK's session task, which outlives this request. + set_request_root_span(session_opener) + await asyncio.create_task(session_task()) + + asyncio.run(initialize_request()) + session_opener.end() + this_message.end() + span = next(s for s in exporter.get_finished_spans() if s.name == span_name) + assert span.parent is not None + assert span.parent.span_id == this_message.get_span_context().span_id + assert span.context.trace_id == this_message.get_span_context().trace_id + assert span.parent.span_id != session_opener.get_span_context().span_id + assert span.context.trace_id != session_opener.get_span_context().trace_id + + +@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) +def test_mcp_span_roots_without_transport_or_propagated_context( + make_payload, span_name +): + """With neither a remote parent nor a transport span there is nothing to nest + under, so the span legitimately starts its own root trace with no links.""" + logger, exporter = _logger() + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": make_payload()}, None, None, None + ) + ) + span = next(s for s in exporter.get_finished_spans() if s.name == span_name) assert span.parent is None - assert span.context.trace_id != transport.get_span_context().trace_id - assert [link.context.span_id for link in span.links] == [ - transport.get_span_context().span_id - ] + assert span.links == () @pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) @@ -641,10 +709,11 @@ def test_mcp_span_carries_authenticated_identity(make_payload, span_name): assert span.attributes[LiteLLM.TEAM_ID] == "t1" -def test_mcp_span_malformed_traceparent_starts_root(): +def test_mcp_span_malformed_traceparent_nests_under_transport(): """A malformed traceparent in ``params._meta`` must not crash or parent to a - bogus span: the propagator ignores it, so the span starts its own root trace and - still links the transport span.""" + bogus span: the propagator ignores it, leaving no remote parent, so the span + falls back to nesting under the transport span rather than starting a + disconnected root trace.""" logger, exporter = _logger() transport = logger._emitter.start_span( SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME @@ -661,9 +730,44 @@ def test_mcp_span_malformed_traceparent_starts_root(): reset_mcp_message_trace_carrier(token) transport.end() span = next(s for s in exporter.get_finished_spans() if s.name == "tools/list") - assert span.parent is None + assert span.parent is not None + assert span.parent.span_id == transport.get_span_context().span_id + assert span.links == () + + +def test_mcp_span_links_this_messages_transport_when_context_is_propagated(): + """On the semconv path the transport is recorded as a link, and that link must + point at the POST carrying this message too. Reading the stale session anchor + would attribute the tool call to whichever request opened the session.""" + logger, exporter = _logger() + session_opener = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + this_message = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + set_request_root_span(session_opener) + trace_token = set_mcp_message_trace_carrier( + {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"} + ) + transport_token = set_mcp_message_transport_span_context( + this_message.get_span_context() + ) + try: + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": _mcp_list_payload()}, None, None, None + ) + ) + finally: + reset_mcp_message_transport_span_context(transport_token) + reset_mcp_message_trace_carrier(trace_token) + session_opener.end() + this_message.end() + span = next(s for s in exporter.get_finished_spans() if s.name == "tools/list") + assert span.parent is not None and span.parent.span_id == 0x2222222222222222 assert [link.context.span_id for link in span.links] == [ - transport.get_span_context().span_id + this_message.get_span_context().span_id ] diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 64813c1eda7..4ea79f9e2a4 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock import pytest @@ -1394,6 +1395,38 @@ class TestEventTypeLogging: assert len(logged_info) == 1 assert logged_info[0]["guardrail_status"] == "guardrail_intervened" + @pytest.mark.asyncio + async def test_log_guardrail_information_records_every_concurrent_guardrail(self): + """Guardrails run concurrently (parallel pre_call/post_call, during_call) share one + request_data dict. Each must still record its own entry. The previous guard counted + entries in that shared dict, so a sibling's append made a guardrail think it had already + recorded and skip its own auto-record — silently dropping lifecycle logs the UI shows.""" + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class SleeperGuardrail(CustomGuardrail): + def __init__(self, name, sleep): + super().__init__(guardrail_name=name, event_hook=GuardrailEventHooks.pre_call) + self._sleep = sleep + + @log_guardrail_information + async def async_pre_call_hook(self, data: dict, **kwargs): + await asyncio.sleep(self._sleep) + return data + + request_data = {"metadata": {}} + # Different sleeps guarantee overlapping execution windows: the faster guardrail + # records while the slower one is still awaiting, which is exactly what tripped the + # old shared-count guard. + await asyncio.gather( + SleeperGuardrail("guardrail-a", 0.05).async_pre_call_hook(data=request_data), + SleeperGuardrail("guardrail-b", 0.15).async_pre_call_hook(data=request_data), + ) + + logged = request_data["metadata"]["standard_logging_guardrail_information"] + assert {entry["guardrail_name"] for entry in logged} == {"guardrail-a", "guardrail-b"} + assert len(logged) == 2 + def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none( self, ): diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 9ff67a82f40..d282e656ce8 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2111,6 +2111,37 @@ def test_token_type_cost_breakdown_reads_cache_write_tokens(): ) +def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(): + """ + Regression: OpenAI gpt-5.6 reports cache-write tokens under + prompt_tokens_details.cache_write_tokens (not the Anthropic cache_creation_tokens + name). Those tokens must be billed at the cache-write rate rather than the plain + input rate. Customer report: cache creation tokens were never counted for the + GPT-5.6 series, so cost was undercounted on cache-write requests. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gpt-5.6" + usage = Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_write_tokens=800), + ) + + assert usage.prompt_tokens_details.cache_write_tokens == 800 + assert usage.prompt_tokens_details.cache_creation_tokens == 800 + + prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + expected_prompt = (1000 - 800) * info["input_cost_per_token"] + 800 * info["cache_creation_input_token_cost"] + assert prompt_cost == pytest.approx(expected_prompt) + assert info["cache_creation_input_token_cost"] > info["input_cost_per_token"] + assert prompt_cost > 1000 * info["input_cost_per_token"] + + def test_token_type_cost_breakdown_reconciles_with_generic_total(): """ Both-ways check: the reasoning subset must sum with the remaining (text) output @@ -2166,6 +2197,65 @@ def test_token_type_cost_breakdown_zero_without_special_tokens(): ) +@pytest.mark.parametrize( + "raw_usage, expect_read, expect_write", + [ + ( + { + "input_tokens": 5000, + "output_tokens": 10, + "total_tokens": 5010, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 4012}, + }, + False, + True, + ), + ( + { + "input_tokens": 5000, + "output_tokens": 10, + "total_tokens": 5010, + "input_tokens_details": {"cached_tokens": 4012, "cache_write_tokens": 0}, + }, + True, + False, + ), + ], +) +def test_token_type_cost_breakdown_openai_responses_api_cache_write_read( + raw_usage, expect_read, expect_write +): + """Regression for #34309: OpenAI Responses API reports cache tokens under + input_tokens_details.{cached_tokens, cache_write_tokens}, not the Anthropic-style + top-level cache_creation_input_tokens. The itemized breakdown must still populate + cache_read_cost / cache_creation_cost from the transformed usage.""" + from litellm.responses.utils import ResponseAPILoggingUtils + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gpt-5.6" + usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_usage) + + breakdown = get_token_type_cost_breakdown( + model=model, custom_llm_provider="openai", usage=usage + ) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + if expect_write: + assert breakdown.cache_creation_cost == pytest.approx( + 4012 * info["cache_creation_input_token_cost"] + ) + assert breakdown.cache_creation_cost > 0 + assert breakdown.cache_read_cost == 0.0 + if expect_read: + assert breakdown.cache_read_cost == pytest.approx( + 4012 * info["cache_read_input_token_cost"] + ) + assert breakdown.cache_read_cost > 0 + assert breakdown.cache_creation_cost == 0.0 + + def test_token_type_cost_breakdown_handles_unknown_model_gracefully(): """A model with no pricing must yield zeros, never raise.""" breakdown = get_token_type_cost_breakdown( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 5bffda126fe..edc257f4c3f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1111,6 +1111,171 @@ async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through litellm._async_success_callback = original_async_callbacks +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_prefer_async_does_not_submit_sync_handler( + logging_obj, +): + """prefer_async_handlers must await async_failure_handler and never submit the sync failure_handler. + + Submitting the sync ``failure_handler`` while awaiting ``async_failure_handler`` + lets both mutate the shared logging_obj at once, which is the concurrent-mutation + crash this dispatch guard exists to prevent. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_failure_callbacks_for_async_calls", + return_value=False, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + mock_async.assert_awaited_once_with(exception, traceback_exception) + mock_sync.assert_not_called() + mock_submit.assert_not_called() + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_async_completes_before_sync_submit( + logging_obj, +): + """The async failure handler must fully finish before the legacy sync handler is scheduled. + + Ordering proves there is no window where both handlers touch the shared + logging_obj concurrently: the sync submit only happens after the await returns. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + events: list[str] = [] + + async def _async_failure(exc, tb, **kwargs): + events.append("async_start") + await asyncio.sleep(0) + events.append("async_end") + + def _submit(*args, **kwargs): + events.append("sync_submit") + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure), + patch.object(logging_obj, "failure_handler", new_callable=MagicMock), + patch.object( + logging_obj, + "_should_run_sync_failure_callbacks_for_async_calls", + return_value=True, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit", + side_effect=_submit, + ), + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + assert events == ["async_start", "async_end", "sync_submit"] + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_callbacks( + logging_obj, +): + """A sync failure callback must still run when only failure callbacks are configured. + + The legacy thread-based path always submitted the sync failure_handler, so gating it on + the success callback list would silently drop failure logging for any deployment that + registers only failure callbacks and no success callbacks. This drives the real predicate + (unmocked), so gating the sync failure handler on the success list fails this test. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + def _sync_failure_callback(*args, **kwargs): + return None + + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.dynamic_success_callbacks = None + logging_obj.dynamic_failure_callbacks = None + + with ( + patch.object(litellm, "success_callback", []), + patch.object(litellm, "failure_callback", [_sync_failure_callback]), + patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock), + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception) + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_sync_sdk_shortcut_runs_sync_handler_inline( + logging_obj, +): + """A sync-SDK request (prefer_async_handlers=False) runs failure_handler inline. + + ``async for`` over a stream from ``completion()`` passes prefer_async_handlers=True; a + plain sync request leaves it False, so the legacy sync handler runs directly and the + async handler is never awaited, matching dispatch_success_handlers. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=False, + ) + + mock_sync.assert_called_once_with(exception, traceback_exception) + mock_async.assert_not_awaited() + mock_submit.assert_not_called() + + def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj): """Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False.""" import datetime diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 71da1d39876..1b37ade6b30 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -226,6 +226,18 @@ def test_get_output_file_id_empty_output_info_falls_through_to_output_config(): assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" +def test_get_output_file_id_output_info_explicit_none_falls_through_to_output_config(): + resp = { + "outputInfo": None, + "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}}, + } + assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" + + +def test_get_output_file_id_output_info_explicit_none_and_no_output_config(): + assert T._get_output_file_id_from_vertex_ai_batch_response({"outputInfo": None}) == "" + + def test_get_output_file_id_no_output_info_and_no_output_config(): assert T._get_output_file_id_from_vertex_ai_batch_response({}) == "" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 95e8e6561f1..5871644ca5f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2087,6 +2087,9 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped(): async def async_failure_handler(self, *args, **kwargs): return None + async def dispatch_failure_handlers(self, *args, **kwargs): + return None + async def failing_make_call(client=None, **kwargs): raise VertexAIError(status_code=400, message="bad input", headers={}) @@ -5404,3 +5407,116 @@ def test_process_candidates_merges_thought_signatures_and_server_side_tools(): fields = model_response.choices[-1].message.provider_specific_fields assert fields["thought_signatures"] == ["sig-text"] assert fields["server_side_tool_invocations"][0]["id"] == "tool-1" + + +def _accumulating_gemini_iterator(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=[], sync_stream=True, logging_obj=MagicMock() + ) + iterator.chunk_type = "accumulated_json" + return iterator + + +def test_accumulated_json_chunk_multi_value_buffer_does_not_wedge(): + """Two complete Gemini objects buffered together must both surface. + + A whole-buffer json.loads raises "Extra data" on concatenated values and, since + the buffer was never reset on failure, returned None forever while growing without + bound. Peeling one value from the front keeps the remainder for the next call. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + + first = iterator.handle_accumulated_json_chunk(chunk=obj + obj) + assert first is not None + assert first.choices[0].delta.content == "a" + + second = iterator.handle_accumulated_json_chunk(chunk="") + assert second is not None + assert second.choices[0].delta.content == "a" + + assert iterator.accumulated_json.strip() == "" + + +def test_accumulated_json_end_of_stream_drains_all_buffered_values(): + """End of stream must drain every buffered value and then terminate. + + With concatenated values a whole-buffer parse never succeeds, so __next__ kept + returning None without shrinking the buffer - an unrecoverable per-request spin. + The bounded loop asserts the iterator both surfaces all values and terminates. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = obj + obj + obj + + out = [] + terminated = False + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + terminated = True + break + if chunk is not None: + out.append(chunk) + + assert terminated, "iterator did not terminate - accumulated buffer wedged" + assert len(out) == 3 + assert iterator.accumulated_json.strip() == "" + + +def test_accumulated_json_end_of_stream_surfaces_leading_value_before_truncated_tail(): + """A complete leading value must survive a truncated trailing value at end of stream. + + The mid-stream perf guard only inspects the buffer's last byte, so a complete leading + object followed by a truncated one (a server that cut the stream mid-object, last byte + not a closer) would keep the guard from ever parsing and drop the complete value. At end + of stream the drain ignores that guard, surfaces the complete value, and discards only + the truncated tail. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = obj + '{"candidates":' + + out = [] + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + break + if chunk is not None: + out.append(chunk) + + assert len(out) == 1 + assert out[0].choices[0].delta.content == "a" + + +def test_accumulated_json_skips_non_dict_leading_value(): + """A non-dict value at the front must not block the dict values behind it. + + raw_decode advances past a decoded value, so a leading non-dict (a JSON array or scalar, + which Gemini never emits but a malformed stream could) must be consumed and skipped. If + the drain stopped on it, the trailing objects would be lost at end of stream. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = "[1, 2]" + obj + + out = [] + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + break + if chunk is not None: + out.append(chunk) + + assert len(out) == 1 + assert out[0].choices[0].delta.content == "a" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_token_exchange.py deleted file mode 100644 index d2aa58e29ea..00000000000 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_token_exchange.py +++ /dev/null @@ -1,539 +0,0 @@ -""" -Tests for OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers. - -Covers: exchange flow, caching, error handling, resolve_mcp_auth integration, -bearer token extraction, and config loading. -""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest - -from litellm.proxy._experimental.mcp_server.auth.token_exchange import ( - TOKEN_EXCHANGE_GRANT_TYPE, - TokenExchangeHandler, -) -from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, -) -from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( - resolve_mcp_auth, -) -from litellm.proxy._types import LiteLLM_MCPServerTable, MCPTransport -from litellm.types.mcp import MCPAuth -from litellm.types.mcp_server.mcp_server_manager import MCPServer - - -def _obo_server(**overrides) -> MCPServer: - defaults = dict( - server_id="srv-obo-1", - name="test-obo", - url="https://mcp.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2_token_exchange, - client_id="litellm-client-id", - client_secret="litellm-client-secret", - token_exchange_endpoint="https://idp.example.com/oauth2/token", - audience="api://mcp-server", - scopes=["mcp.tools.read", "mcp.tools.execute"], - ) - defaults.update(overrides) - return MCPServer(**defaults) - - -def _exchange_response(token="exchanged-tok-abc", expires_in=3600): - resp = MagicMock() - resp.json.return_value = { - "access_token": token, - "token_type": "Bearer", - "expires_in": expires_in, - } - resp.raise_for_status = MagicMock() - resp.text = "" - return resp - - -# ── Exchange Flow ── - - -@pytest.mark.asyncio -async def test_exchange_token_success(): - """Token exchange sends correct RFC 8693 parameters and returns access_token.""" - handler = TokenExchangeHandler() - server = _obo_server() - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("scoped-token-1") - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - result = await handler.exchange_token("user-jwt-xyz", server) - - assert result == "scoped-token-1" - mock_client.post.assert_called_once() - - _, kwargs = mock_client.post.call_args - data = kwargs["data"] - assert data["grant_type"] == TOKEN_EXCHANGE_GRANT_TYPE - assert data["subject_token"] == "user-jwt-xyz" - assert data["subject_token_type"] == "urn:ietf:params:oauth:token-type:access_token" - assert data["audience"] == "api://mcp-server" - assert data["scope"] == "mcp.tools.read mcp.tools.execute" - assert data["client_id"] == "litellm-client-id" - assert data["client_secret"] == "litellm-client-secret" - - -@pytest.mark.asyncio -async def test_exchange_token_no_audience(): - """When audience is None, it is omitted from the request.""" - handler = TokenExchangeHandler() - server = _obo_server(audience=None) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response() - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - await handler.exchange_token("user-jwt", server) - - _, kwargs = mock_client.post.call_args - assert "audience" not in kwargs["data"] - - -@pytest.mark.asyncio -async def test_exchange_token_no_scopes(): - """When scopes is None, scope param is omitted from the request.""" - handler = TokenExchangeHandler() - server = _obo_server(scopes=None) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response() - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - await handler.exchange_token("user-jwt", server) - - _, kwargs = mock_client.post.call_args - assert "scope" not in kwargs["data"] - - -# ── Caching ── - - -@pytest.mark.asyncio -async def test_exchange_token_cached(): - """Second call with same user token uses cache — only 1 HTTP POST.""" - handler = TokenExchangeHandler() - server = _obo_server() - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("cached-exchange-tok") - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - t1 = await handler.exchange_token("same-jwt", server) - t2 = await handler.exchange_token("same-jwt", server) - - assert t1 == t2 == "cached-exchange-tok" - assert mock_client.post.call_count == 1 - - -@pytest.mark.asyncio -async def test_different_user_tokens_not_shared(): - """Different user JWTs get different exchanged tokens.""" - handler = TokenExchangeHandler() - server = _obo_server() - call_count = 0 - - async def mock_post(url, data=None): - nonlocal call_count - call_count += 1 - resp = MagicMock() - resp.json.return_value = { - "access_token": f"exchanged-{call_count}", - "expires_in": 3600, - } - resp.raise_for_status = MagicMock() - return resp - - mock_client = AsyncMock() - mock_client.post = mock_post - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - t1 = await handler.exchange_token("user-a-jwt", server) - t2 = await handler.exchange_token("user-b-jwt", server) - - assert t1 == "exchanged-1" - assert t2 == "exchanged-2" - assert call_count == 2 - - -# ── Error Handling ── - - -@pytest.mark.asyncio -async def test_exchange_token_http_error(): - """HTTP errors from the IDP are wrapped in a ValueError.""" - handler = TokenExchangeHandler() - server = _obo_server() - mock_response = MagicMock() - mock_response.status_code = 400 - mock_response.text = "invalid_grant" - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "Bad Request", - request=MagicMock(), - response=mock_response, - ) - mock_client = AsyncMock() - mock_client.post.return_value = mock_response - - with ( - patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ), - pytest.raises(ValueError, match="failed with status 400"), - ): - await handler.exchange_token("bad-jwt", server) - - -@pytest.mark.asyncio -async def test_exchange_token_http_error_does_not_log_response_body(): - """Raw IDP error bodies are not logged because they can contain credentials.""" - handler = TokenExchangeHandler() - server = _obo_server() - raw_response_body = "client_secret=do-not-log" - mock_response = MagicMock() - mock_response.status_code = 401 - mock_response.text = raw_response_body - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "Unauthorized", - request=MagicMock(), - response=mock_response, - ) - mock_client = AsyncMock() - mock_client.post.return_value = mock_response - - with ( - patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ), - patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.verbose_logger.debug" - ) as mock_debug, - pytest.raises(ValueError, match="failed with status 401"), - ): - await handler.exchange_token("bad-jwt", server) - - logged_values = " ".join( - str(value) - for call in mock_debug.call_args_list - for value in [*call.args, *call.kwargs.values()] - ) - assert raw_response_body not in logged_values - - -@pytest.mark.asyncio -async def test_exchange_token_missing_access_token(): - """Response without access_token raises ValueError.""" - handler = TokenExchangeHandler() - server = _obo_server() - resp = MagicMock() - resp.json.return_value = {"token_type": "Bearer"} - resp.raise_for_status = MagicMock() - mock_client = AsyncMock() - mock_client.post.return_value = resp - - with ( - patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ), - pytest.raises(ValueError, match="missing 'access_token'"), - ): - await handler.exchange_token("jwt", server) - - -@pytest.mark.asyncio -async def test_exchange_token_missing_endpoint(): - """Missing token_exchange_endpoint and token_url raises ValueError.""" - handler = TokenExchangeHandler() - server = _obo_server(token_exchange_endpoint=None, token_url=None) - - with pytest.raises(ValueError, match="no token_exchange_endpoint or token_url"): - await handler.exchange_token("jwt", server) - - -@pytest.mark.asyncio -async def test_exchange_token_missing_credentials(): - """Missing client_id or client_secret raises ValueError.""" - handler = TokenExchangeHandler() - server = _obo_server(client_id=None, client_secret=None) - # has_token_exchange_config will be False, so we call _do_exchange directly - with pytest.raises(ValueError, match="missing client_id or client_secret"): - await handler._do_exchange("jwt", server) - - -# ── resolve_mcp_auth Integration ── - - -@pytest.mark.asyncio -async def test_resolve_mcp_auth_with_token_exchange(): - """resolve_mcp_auth delegates to token exchange when server has OBO config and subject_token provided.""" - server = _obo_server() - mock_handler = AsyncMock() - mock_handler.exchange_token.return_value = "obo-scoped-token" - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.mcp_token_exchange_handler", - mock_handler, - ): - result = await resolve_mcp_auth(server, subject_token="user-jwt") - - assert result == "obo-scoped-token" - mock_handler.exchange_token.assert_called_once_with("user-jwt", server) - - -@pytest.mark.asyncio -async def test_resolve_mcp_auth_obo_without_subject_token_falls_through(): - """Without a subject_token, resolve_mcp_auth falls through to client_credentials.""" - server = _obo_server( - token_url="https://auth.example.com/token", - ) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("cc-token") - - with patch( - "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", - return_value=mock_client, - ): - result = await resolve_mcp_auth(server, subject_token=None) - - # Falls through to client_credentials since subject_token is None - # The server has client_id/client_secret/token_url so has_client_credentials is True - assert result == "cc-token" - - -@pytest.mark.asyncio -async def test_resolve_mcp_auth_obo_without_subject_token_uses_cached_client_credentials(): - """The M2M fallback for OBO servers reuses the client_credentials cache.""" - server = _obo_server( - server_id="srv-obo-m2m-cache", - token_url="https://auth.example.com/token", - ) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("cached-cc-token") - - with patch( - "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", - return_value=mock_client, - ): - first = await resolve_mcp_auth(server, subject_token=None) - second = await resolve_mcp_auth(server, subject_token=None) - - assert first == second == "cached-cc-token" - mock_client.post.assert_called_once() - - -@pytest.mark.asyncio -async def test_resolve_mcp_auth_header_beats_obo(): - """An explicit mcp_auth_header takes priority over OBO token exchange.""" - server = _obo_server() - result = await resolve_mcp_auth( - server, mcp_auth_header="Bearer override", subject_token="user-jwt" - ) - assert result == "Bearer override" - - -# ── Bearer Token Extraction ── - - -def test_extract_bearer_token_from_oauth2_headers(): - """Extracts token from oauth2_headers Authorization header.""" - result = MCPServerManager._extract_bearer_token( - oauth2_headers={"Authorization": "Bearer my-jwt-token"}, - raw_headers=None, - ) - assert result == "my-jwt-token" - - -def test_extract_bearer_token_from_raw_headers(): - """Falls back to raw_headers when oauth2_headers missing.""" - result = MCPServerManager._extract_bearer_token( - oauth2_headers=None, - raw_headers={"authorization": "Bearer raw-jwt"}, - ) - assert result == "raw-jwt" - - -def test_extract_bearer_token_no_bearer_prefix(): - """Returns token as-is when no Bearer prefix.""" - result = MCPServerManager._extract_bearer_token( - oauth2_headers={"Authorization": "some-opaque-token"}, - raw_headers=None, - ) - assert result == "some-opaque-token" - - -def test_extract_bearer_token_none(): - """Returns None when no auth headers present.""" - result = MCPServerManager._extract_bearer_token( - oauth2_headers=None, - raw_headers=None, - ) - assert result is None - - -# ── MCPServer Properties ── - - -def test_has_token_exchange_config_true(): - """has_token_exchange_config is True for a fully configured OBO server.""" - server = _obo_server() - assert server.has_token_exchange_config is True - - -def test_has_token_exchange_config_false_wrong_auth_type(): - """has_token_exchange_config is False when auth_type is not oauth2_token_exchange.""" - server = _obo_server(auth_type=MCPAuth.oauth2) - assert server.has_token_exchange_config is False - - -def test_has_token_exchange_config_false_missing_creds(): - """has_token_exchange_config is False when client_id/client_secret missing.""" - server = _obo_server(client_id=None) - assert server.has_token_exchange_config is False - - -def test_has_token_exchange_config_uses_token_url_fallback(): - """has_token_exchange_config is True when token_url is set instead of token_exchange_endpoint.""" - server = _obo_server( - token_exchange_endpoint=None, - token_url="https://idp.example.com/token", - ) - assert server.has_token_exchange_config is True - - -# ── Config Loading ── - - -@pytest.mark.asyncio -async def test_config_loading_token_exchange_fields(): - """load_servers_from_config correctly maps OBO config fields to MCPServer.""" - manager = MCPServerManager() - config = { - "my_obo_server": { - "url": "https://mcp.example.com/mcp", - "transport": "http", - "auth_type": "oauth2_token_exchange", - "client_id": "my-client", - "client_secret": "my-secret", - "token_exchange_endpoint": "https://idp.example.com/oauth2/token", - "audience": "api://my-mcp", - "scopes": ["read", "write"], - "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", - } - } - await manager.load_servers_from_config(config) - - servers = list(manager.config_mcp_servers.values()) - assert len(servers) == 1 - - server = servers[0] - assert server.auth_type == MCPAuth.oauth2_token_exchange - assert server.token_exchange_endpoint == "https://idp.example.com/oauth2/token" - assert server.audience == "api://my-mcp" - assert server.subject_token_type == "urn:ietf:params:oauth:token-type:jwt" - assert server.client_id == "my-client" - assert server.client_secret == "my-secret" - assert server.scopes == ["read", "write"] - assert server.has_token_exchange_config is True - - -@pytest.mark.asyncio -async def test_config_loading_default_subject_token_type(): - """subject_token_type defaults to access_token when not specified in config.""" - manager = MCPServerManager() - config = { - "obo_defaults": { - "url": "https://mcp.example.com/mcp", - "transport": "http", - "auth_type": "oauth2_token_exchange", - "client_id": "cid", - "client_secret": "csec", - "token_exchange_endpoint": "https://idp.example.com/token", - } - } - await manager.load_servers_from_config(config) - - server = list(manager.config_mcp_servers.values())[0] - assert server.subject_token_type == "urn:ietf:params:oauth:token-type:access_token" - - -@pytest.mark.asyncio -async def test_database_loading_token_exchange_scopes_from_credentials(): - """DB-loaded OBO server credentials retain configured scopes.""" - manager = MCPServerManager() - db_server = LiteLLM_MCPServerTable( - server_id="srv-obo-db", - server_name="obo_db_server", - url="https://mcp.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2_token_exchange, - credentials={ - "client_id": "db-client", - "client_secret": "db-secret", - "token_exchange_endpoint": "https://idp.example.com/oauth2/token", - "audience": "api://db-mcp", - "scopes": ["db.read", "db.write"], - }, - ) - - server = await manager.build_mcp_server_from_table( - db_server, - credentials_are_encrypted=False, - ) - - assert server.auth_type == MCPAuth.oauth2_token_exchange - assert server.client_id == "db-client" - assert server.client_secret == "db-secret" - assert server.token_exchange_endpoint == "https://idp.example.com/oauth2/token" - assert server.audience == "api://db-mcp" - assert server.scopes == ["db.read", "db.write"] - - -@pytest.mark.asyncio -async def test_exchange_token_uses_client_secret_basic_when_configured(): - """LIT-4091: token exchange with token_endpoint_auth_method=client_secret_basic sends the - client credentials as HTTP Basic and omits client_secret from the body.""" - import base64 - - handler = TokenExchangeHandler() - server = _obo_server( - server_id="srv-obo-basic", token_endpoint_auth_method="client_secret_basic" - ) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("scoped-basic") - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - result = await handler.exchange_token("user-jwt-basic", server) - - assert result == "scoped-basic" - _, kwargs = mock_client.post.call_args - expected = "Basic " + base64.b64encode(b"litellm-client-id:litellm-client-secret").decode() - assert kwargs["headers"]["Authorization"] == expected - assert "client_secret" not in kwargs["data"] - assert "client_id" not in kwargs["data"] - assert kwargs["data"]["grant_type"] == TOKEN_EXCHANGE_GRANT_TYPE diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 77f072b81d5..42b6cbee1c4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2333,6 +2333,59 @@ class TestMCPServerManager: assert emitted.headers["Authorization"] == "Bearer upstream-token" assert not kwargs["extra_headers"] or "authorization" not in {k.lower() for k in kwargs["extra_headers"]} + @pytest.mark.asyncio + async def test_create_mcp_client_token_exchange_never_falls_back_to_v1(self): + """A configured OBO server is owned end to end by the v2 token_exchange arm, even when the + caller supplies an x-mcp-* override. This is what makes the v1 OBO handler unreachable, so if + it ever defers to v1 again the deleted handler is silently needed back.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import ( + UpstreamCredentialProvider, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + class _StubExchanger: + def __init__(self): + self.subject_tokens = [] + + async def exchange(self, subject_token, server, config, *, tenant_id=""): + self.subject_tokens.append(subject_token) + return Ok(OAuthToken(access_token="exchanged-token")) + + async def invalidate(self, subject_token, server, config, *, tenant_id=""): + return None + + exchanger = _StubExchanger() + manager = MCPServerManager() + server = MCPServer( + server_id="obo-egress", + name="obo", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://idp.example.com/oauth2/token", + ) + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.resolve_mcp_auth", + new_callable=AsyncMock, + ) as mock_resolve, + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient") as mock_client_cls, + ): + await manager._create_mcp_client( + server=server, + mcp_auth_header="Bearer caller-override", + subject_token="eyJ-subject-token", + cred_provider=UpstreamCredentialProvider(token_exchanger=exchanger), + ) + mock_resolve.assert_not_awaited() + assert exchanger.subject_tokens == ["eyJ-subject-token"] + assert self._emitted_authorization(mock_client_cls) == "Bearer exchanged-token" + @staticmethod def _emitted_authorization(mock_client_cls) -> str: kwargs = mock_client_cls.call_args.kwargs diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 1da44029b5c..0e442102e53 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -237,7 +237,7 @@ class TestListToolRestApiWithToolSearch: return_value={}, ), patch( - "litellm.proxy._experimental.mcp_server.rest_endpoints._get_oauth2_server_ids", + "litellm.proxy._experimental.mcp_server.rest_endpoints._v1_resolved_oauth2_server_ids", return_value=[], ), patch( @@ -316,7 +316,7 @@ class TestListToolRestApiWithToolSearch: return_value={}, ), patch( - "litellm.proxy._experimental.mcp_server.rest_endpoints._get_oauth2_server_ids", + "litellm.proxy._experimental.mcp_server.rest_endpoints._v1_resolved_oauth2_server_ids", return_value=[], ), patch( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index d4ba66c4381..5c9612a055e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -2783,3 +2783,68 @@ class TestRestListToolsetFiltering: ) assert [tool.name for tool in result] == ["lookup_status"] + + +class TestV1ResolvedOauth2Gate: + """The REST surface must stop resolving per-user OAuth2 tokens for servers the v2 resolver owns. + + ``_resolve_v2_auth`` drops any Authorization built here for an ``authorization_code`` server and + injects the resolver's own token, so the v1 lookup was a DB round-trip whose result was discarded. + A server that still defers to v1 (upstream-delegated oauth2) must keep resolving, which is what + makes these assertions non-vacuous. + """ + + @staticmethod + def _oauth2_server(*, delegate_auth_to_upstream: bool) -> Any: + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + return MCPServer( + server_id="oauth2-srv", + name="oauth2-srv", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=delegate_auth_to_upstream, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "delegate_auth_to_upstream, expected_headers, expected_lookups", + [ + (False, None, 0), + (True, {"Authorization": "Bearer stored-token"}, 1), + ], + ) + async def test_user_oauth_headers_skip_v2_owned_servers( + self, delegate_auth_to_upstream, expected_headers, expected_lookups, monkeypatch + ): + from litellm.proxy._experimental.mcp_server import db as mcp_db + + server = self._oauth2_server(delegate_auth_to_upstream=delegate_auth_to_upstream) + resolve_token = AsyncMock(return_value={"access_token": "stored-token"}) + monkeypatch.setattr(mcp_db, "resolve_valid_user_oauth_token", resolve_token) + + headers = await rest_endpoints._get_user_oauth_extra_headers( + server, + UserAPIKeyAuth(user_id="alice", api_key="sk-1234"), + prefetched_creds={"oauth2-srv": {"access_token": "stored-token"}}, + ) + + assert headers == expected_headers + assert resolve_token.await_count == expected_lookups + + def test_prefetch_preflight_only_counts_v1_resolved_servers(self, monkeypatch): + v2_owned = self._oauth2_server(delegate_auth_to_upstream=False) + v1_resolved = self._oauth2_server(delegate_auth_to_upstream=True) + v1_resolved.server_id = "delegate-srv" + registry = {"oauth2-srv": v2_owned, "delegate-srv": v1_resolved} + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: registry.get(server_id), + ) + + assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv"]) == set() + assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv", "delegate-srv"]) == {"delegate-srv"} diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 2c1948adca1..0359d974d19 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -179,6 +179,45 @@ async def test_budget_reservation_runs_when_not_disabled(): assert user_api_key_auth_obj.budget_reservation == reservation +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings,expected_flag", + [ + ({"fail_closed_budget_enforcement": True}, True), + ({}, False), + ], +) +async def test_fail_closed_budget_enforcement_reaches_reservation( + general_settings, expected_flag +): + """#33923: the strict flag must be threaded into reserve_budget_for_request so a + failed reservation write can reject instead of failing open.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=None), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings=general_settings, + ) + + assert ( + mock_reserve.await_args.kwargs["fail_closed_budget_enforcement"] + is expected_flag + ) + + @pytest.mark.asyncio async def test_should_not_reuse_cached_key_object_for_request_state(): key_cache = DualCache() @@ -1290,6 +1329,250 @@ async def test_scim_deactivated_user_key_is_rejected(): setattr(_proxy_server_mod, attr, val) +@pytest.mark.asyncio +async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): + """Cached PROXY_ADMIN auth objects early-return before the marked DB and + master-key returns, and cache serialization drops the exclude=True marker; + the cache-hit boundary must restore it or cached admin traffic silently + bypasses overwrite_user_with_key_hash stamping.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.proxy_server import hash_token + + api_key = "sk-cached-admin-marker-test" + hashed_key = hash_token(api_key) + + cached_token = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + user_id="cached-admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + assert cached_token.via_virtual_key is False + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( + AsyncMock() + ) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs_to_set = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _original_values = { + attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set + } + try: + for attr, val in _attrs_to_set.items(): + setattr(_proxy_server_mod, attr, val) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + with patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=cached_token, + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert isinstance(result, UserAPIKeyAuth) + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + assert result.via_virtual_key is True + assert result.api_key == hashed_key + finally: + for attr, val in _original_values.items(): + setattr(_proxy_server_mod, attr, val) + + +@pytest.mark.asyncio +async def test_master_key_auth_sets_via_virtual_key_marker(): + """Master-key requests must also be stamped by overwrite_user_with_key_hash; + the auth path substitutes the stable alias for api_key and must mark the + result as proxy-validated.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + master_key = "sk-master-key" + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( + AsyncMock() + ) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs_to_set = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": master_key, + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _original_values = { + attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set + } + try: + for attr, val in _attrs_to_set.items(): + setattr(_proxy_server_mod, attr, val) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {master_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert isinstance(result, UserAPIKeyAuth) + assert result.via_virtual_key is True + assert result.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS + finally: + for attr, val in _original_values.items(): + setattr(_proxy_server_mod, attr, val) + + +@pytest.mark.asyncio +async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): + """via_virtual_key gates overwrite_user_with_key_hash stamping and is + forge-stripped from validated input, so the DB auth path setting it by + post-construction assignment is the only thing that turns stamping on.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.proxy_server import hash_token + + api_key = "sk-via-virtual-key-marker-test" + hashed_key = hash_token(api_key) + + valid_token = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + user_id="marker-test-user", + ) + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( + AsyncMock() + ) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + mock_prisma_client = MagicMock() + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs_to_set = { + "prisma_client": mock_prisma_client, + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _original_values = { + attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set + } + try: + for attr, val in _attrs_to_set.items(): + setattr(_proxy_server_mod, attr, val) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + with ( + patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=valid_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert isinstance(result, UserAPIKeyAuth) + assert result.via_virtual_key is True + assert result.api_key == hashed_key + finally: + for attr, val in _original_values.items(): + setattr(_proxy_server_mod, attr, val) + + @pytest.mark.asyncio async def test_return_user_api_key_auth_obj_user_spend_and_budget(): """ diff --git a/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py b/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py index 20bea98351f..9f91d9ee2c9 100644 --- a/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py +++ b/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py @@ -63,6 +63,27 @@ def test_sso_descriptor_mapping_is_single_sourced(): ) +def test_sso_descriptor_mapping_covers_saml_fields(): + # SAML config is stored and read through the same descriptor table as the + # OAuth providers; the login path reads these env vars, so the save path must + # map every SAML field to its uppercase env var. + assert SSO_FIELD_ENV_VARS["saml_idp_metadata_url"] == "SAML_IDP_METADATA_URL" + assert SSO_FIELD_ENV_VARS["saml_idp_metadata_xml"] == "SAML_IDP_METADATA_XML" + assert SSO_FIELD_ENV_VARS["saml_sp_entity_id"] == "SAML_SP_ENTITY_ID" + assert SSO_FIELD_ENV_VARS["saml_allow_unsolicited"] == "SAML_ALLOW_UNSOLICITED" + + +def test_resolve_sso_config_resolves_saml_fields(): + resolved = resolve_sso_config( + {"saml_idp_metadata_url": "https://idp.example.com/metadata"}, + {"SAML_ALLOW_UNSOLICITED": "true"}, + ) + assert resolved.config.saml_idp_metadata_url == "https://idp.example.com/metadata" + assert resolved.provenance["saml_idp_metadata_url"] == "db" + assert resolved.config.saml_allow_unsolicited == "true" + assert resolved.provenance["saml_allow_unsolicited"] == "env" + + def test_resolve_sso_config_returns_unmasked_secret_and_provenance(): # The resolver hands back plaintext; masking is the endpoint's job. If the # resolver masked, the login path would consume a masked secret and fail. diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 4021f922877..18b5bd92411 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3680,6 +3680,22 @@ async def test_pre_call_hook_skips_chat_traffic_when_configured_for_pre_mcp_call mock_post.assert_not_called() +def test_process_response_with_none_metadata_does_not_crash(): + guardrail = _make_guardrail() + response = {"id": "batch_123", "status": "validating"} + request_data = {"model": "gemini-2.5-flash", "metadata": None} + + result = guardrail._process_response( + response=response, + request_data=request_data, + event_type=GuardrailEventHooks.post_call, + ) + + assert result is response + assert isinstance(request_data["metadata"], dict) + assert "standard_logging_guardrail_information" in request_data["metadata"] + + @pytest.mark.asyncio async def test_moderation_hook_scans_mcp_tool_call_when_configured_for_during_mcp_call(): """A guardrail configured with mode `during_mcp_call` must scan MCP tool calls. diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 26feddadf79..14cab50f441 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -1,3 +1,5 @@ +import pytest + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_registry import ( get_guardrail_initializer_from_hooks, @@ -32,6 +34,44 @@ def test_noma_registry_resolution(): assert "noma_v2" in guardrail_initializer_registry +@pytest.mark.parametrize( + "configured, expected", + [(None, True), (False, False), (True, True)], +) +def test_initialize_guardrail_run_in_parallel_preserves_constructor_default(configured, expected): + """ + A guardrail whose constructor sets run_in_parallel=True must keep that default when + the config omits the key; only an explicit config value may override it. The + previous code wrote bool(None)==False on every instance, silently disabling the + opt-in for such guardrails. + """ + from litellm.proxy.guardrails import guardrail_registry as registry_module + + def _initializer(litellm_params, guardrail): + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + registry_module.guardrail_initializer_registry["parallel_default_test"] = _initializer + try: + params = {"guardrail": "parallel_default_test", "mode": "pre_call"} + if configured is not None: + params["run_in_parallel"] = configured + + handler = InMemoryGuardrailHandler() + result = handler.initialize_guardrail( + guardrail={"guardrail_name": "cf-parallel-default", "litellm_params": params}, + ) + + stored = handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]] + assert stored.run_in_parallel is expected + finally: + registry_module.guardrail_initializer_registry.pop("parallel_default_test", None) + + def test_update_in_memory_guardrail(): handler = InMemoryGuardrailHandler() handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail( diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 83593c20110..71e775842e3 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -62,3 +62,27 @@ def test_initialize_guardrail_preserves_guardrail_info(): assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"} stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]] assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"} + + +@pytest.mark.parametrize( + "config_value, expected", + [(True, True), (False, False), (None, False)], +) +def test_initialize_guardrail_sets_run_in_parallel(config_value, expected): + """run_in_parallel from litellm_params must reach the built guardrail instance.""" + litellm_params = { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": "pre_call", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + } + if config_value is not None: + litellm_params["run_in_parallel"] = config_value + + guardrail_handler = InMemoryGuardrailHandler() + result = guardrail_handler.initialize_guardrail( + guardrail={"guardrail_name": "test_parallel_flag", "litellm_params": litellm_params}, + ) + + custom_guardrail = guardrail_handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]] + assert custom_guardrail.run_in_parallel is expected diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py new file mode 100644 index 00000000000..57decc7d458 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -0,0 +1,644 @@ +""" +Regression tests for SAML 2.0 SSO (SP- and IdP-initiated) on the admin UI. + +These exercise the real OneLogin python3-saml validation by generating signed +SAML responses with a freshly minted IdP keypair, so a mutation that weakens +signature, signing-requirement, expiry, replay or attribute-mapping handling +makes a test fail. +""" + +import base64 +import datetime +import time + +import pytest +from fastapi import HTTPException, Request + +pytest.importorskip( + "onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests" +) + +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.x509.oid import NameOID +from onelogin.saml2.utils import OneLogin_Saml2_Utils +from starlette.datastructures import URL + +from typing import cast + +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy._types import LitellmUserRoles +from litellm.proxy.management_endpoints.sso.saml_sso import ( + _SAML_AUTHN_REQUEST_CACHE_PREFIX, + _SAML_AUTHN_STATE_COOKIE, + _SAML_MAX_POST_BYTES, + _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS, + _SAML_REPLAY_GUARD_MAX_TTL_SECONDS, + SAMLAuthHandler, +) + + +def _shared_cache(store=None): + """A DualCache whose replay guard is backed by a shared, atomic store. + + An InMemoryCache instance stands in for Redis; passing the same instance to + two DualCaches simulates two workers sharing one atomic backend.""" + return DualCache(redis_cache=cast(RedisCache, store or InMemoryCache())) + +IDP_ENTITY = "https://idp.example.com/metadata" +SP_ENTITY = "https://proxy.example.com/sso/saml/metadata" +ACS = "https://proxy.example.com/sso/saml/callback" +SSO_URL = "https://idp.example.com/sso" +PROXY_BASE_URL = "https://proxy.example.com" + + +def _make_idp_keypair(): + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "idp.example.com")]) + cert = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime.datetime.utcnow() - datetime.timedelta(days=1)) + .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365)) + .sign(key, hashes.SHA256()) + ) + key_pem = key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ).decode() + cert_pem = cert.public_bytes(serialization.Encoding.PEM).decode() + return key_pem, cert_pem + + +def _idp_metadata_xml(cert_pem): + cert_body = "".join( + line for line in cert_pem.splitlines() if "CERTIFICATE" not in line + ) + return ( + '' + f'' + '' + '' + f"{cert_body}" + "" + '' + "" + ) + + +def _saml_time(delta_seconds): + t = datetime.datetime.utcnow() + datetime.timedelta(seconds=delta_seconds) + return t.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _build_signed_response( + key_pem, + cert_pem, + *, + in_response_to=None, + response_level_in_response_to=True, + email="alice@example.com", + attributes=None, + not_before_delta=-60, + not_on_or_after_delta=300, + sign=True, +): + if attributes is None: + attributes = { + "email": [email], + "givenName": ["Alice"], + "sn": ["Smith"], + "role": ["internal_user"], + } + assertion_id = "_assertion_" + OneLogin_Saml2_Utils.generate_unique_id() + response_id = "_response_" + OneLogin_Saml2_Utils.generate_unique_id() + not_before = _saml_time(not_before_delta) + not_on_or_after = _saml_time(not_on_or_after_delta) + issue_instant = _saml_time(-1) + irt = f'InResponseTo="{in_response_to}"' if in_response_to else "" + response_irt = irt if response_level_in_response_to else "" + + attr_xml = "".join( + f'' + + "".join(f"{v}" for v in values) + + "" + for name, values in attributes.items() + ) + + assertion = ( + '' + f"{IDP_ENTITY}" + "" + '' + f"{email}" + '' + f'' + "" + f'' + f"{SP_ENTITY}" + "" + f'' + "" + "urn:oasis:names:tc:SAML:2.0:ac:classes:Password" + "" + f"{attr_xml}" + "" + ) + + if sign: + signed = OneLogin_Saml2_Utils.add_sign(assertion, key_pem, cert_pem) + assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace( + '', "" + ) + + return ( + '' + '' + f"{IDP_ENTITY}" + '' + "" + f"{assertion}" + ) + + +def _b64(xml): + return base64.b64encode(xml.encode()).decode() + + +def _fake_request(cookies=None): + return type( + "Req", + (), + { + "base_url": URL(PROXY_BASE_URL + "/"), + "query_params": {}, + "cookies": cookies or {}, + }, + )() + + +async def _acs(b64, cache, cookies=None): + return await SAMLAuthHandler.handle_acs( + _fake_request(cookies), cache, {"SAMLResponse": b64} + ) + + +@pytest.fixture +def saml_env(monkeypatch): + key_pem, cert_pem = _make_idp_keypair() + monkeypatch.setenv("SAML_IDP_METADATA_XML", _idp_metadata_xml(cert_pem)) + monkeypatch.setenv("SAML_SP_ENTITY_ID", SP_ENTITY) + monkeypatch.setenv("PROXY_BASE_URL", PROXY_BASE_URL) + for var in ( + "SAML_IDP_METADATA_URL", + "SAML_ATTRIBUTE_EMAIL", + "SAML_ATTRIBUTE_TEAM_IDS", + "SAML_ALLOW_UNSOLICITED", + "ALLOWED_EMAIL_DOMAINS", + ): + monkeypatch.delenv(var, raising=False) + return key_pem, cert_pem + + +@pytest.fixture +def saml_env_idp_initiated(saml_env, monkeypatch): + monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "true") + return saml_env + + +@pytest.mark.asyncio +async def test_valid_idp_initiated_login_maps_assertion_to_user(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem) + + result = await _acs(_b64(resp), _shared_cache()) + + assert result.email == "alice@example.com" + assert result.id == "alice@example.com" + assert result.first_name == "Alice" + assert result.last_name == "Smith" + assert result.user_role == LitellmUserRoles.INTERNAL_USER + assert result.provider == "saml" + + +@pytest.mark.asyncio +async def test_tampered_assertion_is_rejected(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response(key_pem, cert_pem) + tampered = resp.replace("alice@example.com", "attacker@example.com") + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(tampered), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_unsigned_assertion_is_rejected(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response(key_pem, cert_pem, sign=False) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_signature_from_untrusted_key_is_rejected(saml_env): + _, cert_pem = saml_env + attacker_key, attacker_cert = _make_idp_keypair() + resp = _build_signed_response(attacker_key, attacker_cert) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_expired_assertion_is_rejected(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response( + key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600 + ) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_sp_initiated_unknown_in_response_to_is_rejected(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response(key_pem, cert_pem, in_response_to="_never_issued") + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_sp_initiated_known_request_succeeds_once_then_replay_rejected(saml_env): + key_pem, cert_pem = saml_env + cache = DualCache() + request_id = "_authn_req_known" + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 + ) + resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id) + cookies = {_SAML_AUTHN_STATE_COOKIE: request_id} + + result = await _acs(_b64(resp), cache, cookies=cookies) + assert result.email == "alice@example.com" + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), cache, cookies=cookies) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env): + key_pem, cert_pem = saml_env + cache = DualCache() + request_id = "_authn_req_known" + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 + ) + resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), cache) + assert exc.value.status_code == 401 + + with pytest.raises(HTTPException) as exc: + await _acs( + _b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"} + ) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_subjectconfirmation_only_in_response_to_without_cookie_is_rejected( + saml_env_idp_initiated, +): + """An IdP that stamps InResponseTo only on the SubjectConfirmationData (not the + Response element) is still solicited and must be browser-bound: with unsolicited + explicitly allowed, a missing cookie must still 401 rather than slip through.""" + key_pem, cert_pem = saml_env_idp_initiated + cache = DualCache() + request_id = "_authn_req_known" + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 + ) + resp = _build_signed_response( + key_pem, + cert_pem, + in_response_to=request_id, + response_level_in_response_to=False, + ) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), cache) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml_env): + key_pem, cert_pem = saml_env + cache = DualCache() + request_id = "_authn_req_known" + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 + ) + resp = _build_signed_response( + key_pem, + cert_pem, + in_response_to=request_id, + response_level_in_response_to=False, + ) + + result = await _acs( + _b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id} + ) + assert result.email == "alice@example.com" + + +@pytest.mark.asyncio +async def test_unsolicited_response_rejected_by_default(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response(key_pem, cert_pem) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_idp_initiated_assertion_replay_is_rejected(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + cache = _shared_cache() + resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com") + + first = await _acs(_b64(resp), cache) + assert first.email == "bob@example.com" + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), cache) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_assertion_without_id_is_rejected(saml_env_idp_initiated): + """An assertion with no ID attribute has no stable replay key. On the unsolicited + path there is no browser binding, so the consumed-assertion guard is the only replay + defense; a missing ID must be rejected rather than silently skipping the guard.""" + + class _AuthNoAssertionId: + def get_last_response_in_response_to(self): + return None + + def get_last_response_xml(self): + return None + + def get_last_assertion_id(self): + return None + + with pytest.raises(HTTPException) as exc: + await SAMLAuthHandler._enforce_response_binding( + _AuthNoAssertionId(), _shared_cache(), None + ) + assert exc.value.status_code == 401 + assert "ID" in exc.value.detail + + +@pytest.mark.asyncio +async def test_unsolicited_response_rejected_when_disabled(saml_env, monkeypatch): + key_pem, cert_pem = saml_env + monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "false") + resp = _build_signed_response(key_pem, cert_pem) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_invalid_email_in_assertion_is_rejected_cleanly(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + email="not-an-email", + attributes={"email": ["not-an-email"], "givenName": ["X"]}, + ) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), _shared_cache()) + assert exc.value.status_code == 401 + assert "invalid subject or email" in exc.value.detail + + +@pytest.mark.asyncio +async def test_email_less_assertion_rejected_when_domain_restriction_configured( + saml_env_idp_initiated, monkeypatch +): + key_pem, cert_pem = saml_env_idp_initiated + monkeypatch.setenv("ALLOWED_EMAIL_DOMAINS", "example.com") + resp = _build_signed_response( + key_pem, + cert_pem, + email="opaque-persistent-id-123", + attributes={"givenName": ["Alice"]}, + ) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), _shared_cache()) + assert exc.value.status_code == 401 + assert "ALLOWED_EMAIL_DOMAINS" in exc.value.detail + + +@pytest.mark.asyncio +async def test_email_less_assertion_allowed_without_domain_restriction(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + email="opaque-persistent-id-123", + attributes={"givenName": ["Alice"]}, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.email is None + assert result.id == "opaque-persistent-id-123" + + +@pytest.mark.asyncio +async def test_custom_email_attribute_override(saml_env_idp_initiated, monkeypatch): + key_pem, cert_pem = saml_env_idp_initiated + monkeypatch.setenv("SAML_ATTRIBUTE_EMAIL", "corpMail") + resp = _build_signed_response( + key_pem, + cert_pem, + email="ignored@example.com", + attributes={ + "corpMail": ["real@corp.example.com"], + "givenName": ["Real"], + }, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.email == "real@corp.example.com" + + +@pytest.mark.asyncio +async def test_team_ids_extracted_from_groups_attribute(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + attributes={ + "email": ["carol@example.com"], + "groups": ["team-a", "team-b"], + }, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.team_ids == ["team-a", "team-b"] + + +@pytest.mark.asyncio +async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env): + cache = DualCache() + redirect = await SAMLAuthHandler.build_login_redirect(_fake_request(), cache) + + location = redirect.headers["location"] + assert location.startswith(SSO_URL) + assert "SAMLRequest=" in location + cached = [ + k + for k in cache.in_memory_cache.cache_dict + if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX) + ] + assert len(cached) == 1 + + request_id = cached[0].split(":", 1)[1] + set_cookie = redirect.headers["set-cookie"] + assert f"{_SAML_AUTHN_STATE_COOKIE}={request_id}" in set_cookie + assert "httponly" in set_cookie.lower() + + +@pytest.mark.asyncio +async def test_sp_metadata_contains_acs_and_entity_id(saml_env): + metadata = await SAMLAuthHandler.build_sp_metadata(_fake_request(), DualCache()) + assert ACS in metadata + assert SP_ENTITY in metadata + assert "AssertionConsumerService" in metadata + + +def test_replay_guard_ttl_tracks_assertion_validity(): + class _Auth: + def __init__(self, not_on_or_after): + self._not_on_or_after = not_on_or_after + + def get_last_assertion_not_on_or_after(self): + return self._not_on_or_after + + now = int(time.time()) + + long_lived = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 7200)) + assert long_lived >= 7200 + + short_lived = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 60)) + assert short_lived == _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS + + missing = SAMLAuthHandler._replay_guard_ttl(_Auth(None)) + assert missing == _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS + + capped = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 10 * 86400)) + assert capped == _SAML_REPLAY_GUARD_MAX_TTL_SECONDS + + +def test_is_saml_configured_reflects_env(monkeypatch): + monkeypatch.delenv("SAML_IDP_METADATA_URL", raising=False) + monkeypatch.delenv("SAML_IDP_METADATA_XML", raising=False) + assert SAMLAuthHandler.is_saml_configured() is False + + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml") + assert SAMLAuthHandler.is_saml_configured() is True + + +@pytest.mark.asyncio +async def test_idp_initiated_rejected_without_shared_cache(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + assert "shared Redis cache" in exc.value.detail + + +@pytest.mark.asyncio +async def test_idp_initiated_replay_rejected_across_workers(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + shared_store = InMemoryCache() + worker_one = _shared_cache(shared_store) + worker_two = _shared_cache(shared_store) + resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com") + + first = await _acs(_b64(resp), worker_one) + assert first.email == "bob@example.com" + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), worker_two) + assert exc.value.status_code == 401 + + +class _FakeChunkedRequest: + def __init__(self, chunks, content_length=None): + self._chunks = chunks + self.headers = {} if content_length is None else {"content-length": content_length} + + async def stream(self): + for chunk in self._chunks: + yield chunk + + +@pytest.mark.asyncio +async def test_read_acs_post_data_parses_form(): + body = b"SAMLResponse=abc123&RelayState=%2Fui%2F" + request = _FakeChunkedRequest([body], content_length=str(len(body))) + + post_data = await SAMLAuthHandler.read_acs_post_data(cast(Request, request)) + + assert post_data == {"SAMLResponse": "abc123", "RelayState": "/ui/"} + + +@pytest.mark.asyncio +async def test_read_acs_post_data_rejects_oversized_content_length(): + request = _FakeChunkedRequest([b""], content_length=str(_SAML_MAX_POST_BYTES + 1)) + + with pytest.raises(HTTPException) as exc: + await SAMLAuthHandler.read_acs_post_data(cast(Request, request)) + assert exc.value.status_code == 413 + + +@pytest.mark.asyncio +async def test_read_acs_post_data_rejects_oversized_stream_without_content_length(): + chunk = b"a" * (1024 * 1024) + chunk_count = _SAML_MAX_POST_BYTES // len(chunk) + 2 + request = _FakeChunkedRequest([chunk] * chunk_count) + + with pytest.raises(HTTPException) as exc: + await SAMLAuthHandler.read_acs_post_data(cast(Request, request)) + assert exc.value.status_code == 413 diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py index cf80ee5dee5..351f125052d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -13,12 +13,17 @@ from datetime import datetime, timezone from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient sys.path.insert(0, os.path.abspath("../../..")) -from litellm.proxy.management_endpoints.tool_management_endpoints import router +from litellm.proxy.management_endpoints.tool_management_endpoints import ( + _build_tool_spend_response, + _ToolSpendRow, + router, +) from litellm.types.tool_management import LiteLLM_ToolTableRow # --- helpers --- @@ -50,9 +55,9 @@ def _make_app() -> FastAPI: # Stub the auth dependency so we don't need a real proxy running. def _override_auth(): - from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - return UserAPIKeyAuth(api_key="sk-test", user_id="admin") + return UserAPIKeyAuth(api_key="sk-test", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) # A real (non-None) prisma stub for truthiness checks. @@ -147,3 +152,117 @@ class TestToolManagementEndpoints: json={"tool_name": "my_tool", "input_policy": "invalid_value"}, ) assert resp.status_code == 422 + + def test_tool_spend_route_not_shadowed_by_get_tool(self): + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend") + assert resp.status_code == 200 + assert resp.json()["by_tool"] == [] + + def test_tool_spend_aggregates_and_sorts(self): + rows = [ + {"date": "2026-07-01", "tool_name": "search", "call_count": 2, "spend": 1.0, "total_tokens": 100}, + {"date": "2026-07-02", "tool_name": "search", "call_count": 1, "spend": 4.0, "total_tokens": 50}, + {"date": "2026-07-01", "tool_name": "read_file", "call_count": 3, "spend": 2.0, "total_tokens": 300}, + ] + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(side_effect=[rows, [{"total_spend": 5.5}]]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + body = resp.json() + assert [t["tool_name"] for t in body["by_tool"]] == ["search", "read_file"] + search = body["by_tool"][0] + assert search["spend"] == 5.0 + assert search["call_count"] == 3 + assert search["total_tokens"] == 150 + assert len(body["daily"]) == 3 + assert body["start_date"] == "2026-07-01" + assert body["end_date"] == "2026-07-02" + assert body["total_spend"] == 5.5 + + @patch("litellm.proxy.proxy_server.prisma_client", None) + def test_tool_spend_no_db_returns_500(self): + resp = self.client.get("/v1/tool/spend") + assert resp.status_code == 500 + + def test_tool_spend_end_date_is_inclusive_via_exclusive_next_day_bound(self): + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + expected_binds = ( + datetime(2026, 7, 1, tzinfo=timezone.utc).isoformat(), + datetime(2026, 7, 3, tzinfo=timezone.utc).isoformat(), + ) + assert prisma.db.query_raw.await_count == 2 + for call in prisma.db.query_raw.await_args_list: + assert tuple(call.args[1:]) == expected_binds + assert resp.json()["end_date"] == "2026-07-02" + + @pytest.mark.parametrize( + "query", + [ + "start_date=not-a-date", + "start_date=2026-02-30", + "start_date=07/01/2026", + "end_date=2026-13-01", + "end_date=20260701", + ], + ) + def test_tool_spend_malformed_date_returns_400(self, query: str): + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get(f"/v1/tool/spend?{query}") + assert resp.status_code == 400 + assert "Invalid date format" in resp.json()["detail"] + prisma.db.query_raw.assert_not_awaited() + + def test_tool_spend_non_admin_returns_403(self): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = _make_app() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="sk-user", user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER + ) + client = TestClient(app, raise_server_exceptions=True) + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/tool/spend") + assert resp.status_code == 403 + prisma.db.query_raw.assert_not_awaited() + + +def _spend_row(date: str, tool_name: str, spend: float, call_count: int = 1, total_tokens: int = 10) -> _ToolSpendRow: + return _ToolSpendRow(date=date, tool_name=tool_name, call_count=call_count, spend=spend, total_tokens=total_tokens) + + +class TestBuildToolSpendResponse: + def test_multi_tool_attribution_double_counts_per_tool_but_not_total(self): + rows = [ + _spend_row("2026-07-01", "a", spend=3.0), + _spend_row("2026-07-01", "b", spend=3.0), + ] + resp = _build_tool_spend_response(rows, total_spend=3.0, start_date="2026-07-01", end_date="2026-07-01") + by_tool = {t.tool_name: t.spend for t in resp.by_tool} + assert by_tool == {"a": 3.0, "b": 3.0} + assert resp.total_spend == 3.0 + + def test_groups_across_days_and_sorts_by_spend(self): + rows = [ + _spend_row("2026-07-01", "b", spend=1.0, call_count=2, total_tokens=100), + _spend_row("2026-07-02", "b", spend=4.0, call_count=1, total_tokens=50), + _spend_row("2026-07-01", "a", spend=2.0, call_count=3, total_tokens=300), + ] + resp = _build_tool_spend_response(rows, total_spend=7.0, start_date="2026-07-01", end_date="2026-07-02") + assert [(t.tool_name, t.spend, t.call_count, t.total_tokens) for t in resp.by_tool] == [ + ("b", 5.0, 3, 150), + ("a", 2.0, 3, 300), + ] + assert len(resp.daily) == 3 diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 63a47428780..795b7cd5a9e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -7391,6 +7391,71 @@ async def test_legacy_login_page_hides_credentials_hint_via_general_settings(): assert "MASTER_KEY" not in body +@pytest.mark.asyncio +async def test_saml_callback_blocked_when_admin_ui_disabled(): + """An IdP-initiated assertion must not mint a UI session when the admin UI is + disabled; the ACS enforces DISABLE_ADMIN_UI like the SP-initiated login route.""" + from litellm.proxy.management_endpoints.ui_sso import saml_callback + + with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "true"}): + response = await saml_callback(SimpleNamespace(cookies={})) + + assert response.status_code == 200 + assert "Admin UI is Disabled" in response.body.decode() + + +@pytest.mark.asyncio +async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): + """An IdP-initiated assertion must not bypass the >5 free-SSO-user Enterprise gate + that /sso/key/generate enforces; the ACS re-checks it after validating the assertion, + so the entitlement DB query never runs on unvalidated input.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import saml_callback + from litellm.proxy.management_endpoints.types import CustomOpenID + + call_order: list[str] = [] + + async def _fake_handle_acs(**kwargs): + call_order.append("validate") + return CustomOpenID( + id="dana@litellm.ai", + email="dana@litellm.ai", + first_name=None, + last_name=None, + display_name="dana", + picture=None, + provider="saml", + team_ids=[], + user_role=None, + ) + + async def _fake_count_billable_users(): + call_order.append("count") + return 6 + + async def _stream(): + yield b"SAMLResponse=signed-response" + + request_double = SimpleNamespace(cookies={}, headers={}, stream=_stream) + + with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}), patch( + "litellm.proxy.proxy_server.premium_user", False + ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( + "litellm.proxy.proxy_server.master_key", "sk-1234" + ), patch( + "litellm.proxy.management_endpoints.sso.saml_sso.SAMLAuthHandler.handle_acs", + new=_fake_handle_acs, + ), patch( + "litellm.repositories.user_repository.UserRepository.count_billable_users", + new=AsyncMock(side_effect=_fake_count_billable_users), + ): + with pytest.raises(ProxyException) as exc: + await saml_callback(request_double) + + assert str(exc.value.code) == "403" + assert call_order == ["validate", "count"] + + @pytest.mark.asyncio async def test_cli_poll_key_tolerates_missing_user_row(): """The CLI poll must still mint the JWT when the user lookup raises, diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index 01e5414a469..504414ea635 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -1076,6 +1076,13 @@ async def test_add_new_member_creates_missing_user_atomically_via_upsert(): would race a check-then-create into a duplicate-key failure. The upsert seeds teams on create, and the filtered append is a no-op because the team is already present on the freshly created row. + + Calling upsert is not on its own enough to be atomic: Prisma only compiles it + down to a single INSERT ... ON CONFLICT when the update branch is non-empty, + and otherwise emits SELECT-then-INSERT, which loses the race. That is how + parallel /team/new calls naming the same new member started returning 500 + "Unique constraint failed on the fields: (user_id)", so the shape of both + branches is pinned here. """ from litellm.proxy._types import LitellmUserRoles @@ -1114,5 +1121,7 @@ async def test_add_new_member_creates_missing_user_atomically_via_upsert(): # non-atomic standalone create that could race under concurrent provisioning mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_usertable.create.assert_not_called() - create_data = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"]["create"] - assert create_data["teams"] == ["team-1"] + upsert_data = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"] + assert upsert_data["create"]["teams"] == ["team-1"] + assert upsert_data["update"], "empty update branch degrades the upsert to a racy SELECT-then-INSERT" + assert "teams" not in upsert_data["update"] diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index db72a7fb38c..67945436987 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1628,6 +1628,226 @@ async def test_ui_view_spend_logs_date_range_filter(client, monkeypatch): assert data["data"][0]["id"] == "log2" +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window( + client, monkeypatch +): + """ + LIT-3981: a request_id lookup on the UI route resolves across all time even + when the caller sends a date window that excludes the log (the dashboard + always sends a window). The window is dropped and request_id alone scopes + the query. Pre-fix the window was always applied, so an id from an older + page returned nothing. + """ + today = datetime.datetime.now(timezone.utc) + mock_spend_logs = [ + { + "id": "log_old", + "request_id": "req-old", + "api_key": "sk-test-key", + "user": "test_user_1", + "team_id": "team1", + "spend": 0.05, + "startTime": (today - datetime.timedelta(days=90)).isoformat(), + "model": "gpt-4", + }, + ] + + captured: dict = {} + + def filter_fn(where): + captured["where"] = where + rows = _filter_logs_by_date_range(mock_spend_logs, where) + if where.get("request_id"): + rows = [r for r in rows if r["request_id"] == where["request_id"]] + return rows + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn), + ) + + # A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends. + start_date = (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S") + end_date = today.strftime("%Y-%m-%d %H:%M:%S") + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs/ui", + params={ + "request_id": "req-old", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert data["data"][0]["request_id"] == "req-old" + # Query dropped the time window and scoped solely by the primary key. + assert "startTime" not in captured["where"] + assert captured["where"]["request_id"] == "req-old" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_requires_dates_without_request_id( + client, monkeypatch +): + """The date window stays mandatory on the UI route when no request_id is set.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([], lambda where: []), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs/ui", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 400 + assert "date" in response.text.lower() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_logs_v2_still_requires_dates_with_request_id(client, monkeypatch): + """The public /spend/logs/v2 contract is unchanged: dates remain required even + when request_id is supplied. Only the internal UI route relaxes the window.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([], lambda where: []), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs/v2", + params={"request_id": "req-old"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 400 + assert "date" in response.text.lower() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatch): + """A non-admin looking up a request_id they do not own is rejected (403), so + the relaxed date window cannot read another tenant's log by id.""" + + class _ForeignRow: + user = "other_user" + team_id = None + + class _SpendLogs: + async def find_unique(self, where, include=None): + return _ForeignRow() + + class _DB: + def __init__(self): + self.litellm_spendlogs = _SpendLogs() + + class _Prisma: + def __init__(self): + self.db = _DB() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _Prisma()) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" + ) + try: + response = client.get( + "/spend/logs/ui", + params={"request_id": "foreign-req"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( + client, monkeypatch +): + """A non-admin owner looking up their own request_id resolves across all time. + The ownership check authorizes the single row, so the query drops both the date + window and the general user/team scoping and filters by the primary key alone; + without that skip an internal user would have a `user`/`OR` clause added.""" + today = datetime.datetime.now(timezone.utc) + mock_spend_logs = [ + { + "id": "log_old", + "request_id": "req-old", + "api_key": "sk-test-key", + "user": "user_1", + "team_id": "team1", + "spend": 0.05, + "startTime": (today - datetime.timedelta(days=90)).isoformat(), + "model": "gpt-4", + }, + ] + + captured: dict = {} + + def filter_fn(where): + captured["where"] = where + rows = _filter_logs_by_date_range(mock_spend_logs, where) + if where.get("request_id"): + rows = [r for r in rows if r["request_id"] == where["request_id"]] + return rows + + mock_prisma = make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn) + + class _OwnedRow: + user = "user_1" + team_id = "team1" + + async def _find_unique(where, include=None): + return _OwnedRow() + + mock_prisma.db.find_unique = _find_unique + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends. + start_date = (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S") + end_date = today.strftime("%Y-%m-%d %H:%M:%S") + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" + ) + try: + response = client.get( + "/spend/logs/ui", + params={ + "request_id": "req-old", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert data["data"][0]["request_id"] == "req-old" + assert "startTime" not in captured["where"] + assert captured["where"]["request_id"] == "req-old" + assert "user" not in captured["where"] + assert "OR" not in captured["where"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_unauthorized(client): # Test without authorization header diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5d10bb33751..cc1e2943c8f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -109,6 +109,143 @@ def test_get_logging_payload_does_not_map_missing_or_zero_cached_tokens(prompt_t assert "cache_read_input_tokens" not in additional_usage_values +def test_get_logging_payload_maps_openai_cache_write_tokens_to_cache_creation_input_tokens(): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=1000, + completion_tokens=2, + total_tokens=1002, + prompt_tokens_details={"cached_tokens": 0, "cache_write_tokens": 800}, + ) + ) + + assert additional_usage_values["cache_creation_input_tokens"] == 800 + assert additional_usage_values["prompt_tokens_details"]["cache_write_tokens"] == 800 + + +def test_get_logging_payload_preserves_anthropic_cache_creation_input_tokens(): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=1000, + completion_tokens=2, + total_tokens=1002, + cache_creation_input_tokens=300, + ) + ) + + assert additional_usage_values["cache_creation_input_tokens"] == 300 + + +@pytest.mark.parametrize( + "prompt_tokens_details", + [None, {"cached_tokens": 100}, {"cached_tokens": 100, "cache_write_tokens": 0}], +) +def test_get_logging_payload_does_not_map_missing_or_zero_cache_write_tokens(prompt_tokens_details): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + prompt_tokens_details=prompt_tokens_details, + ) + ) + + assert "cache_creation_input_tokens" not in additional_usage_values + + +def _make_standard_logging_payload_with_usage_object(usage_object: dict) -> StandardLoggingPayload: + return StandardLoggingPayload( + id="test-id-responses", + call_type="responses", + stream=False, + response_cost=0.02, + status="success", + total_tokens=1010, + prompt_tokens=1000, + completion_tokens=10, + startTime=1234567890.0, + endTime=1234567891.0, + completionStartTime=None, + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-5.6", model_map_value=None), + model="gpt-5.6", + model_id="model-123", + model_group="openai", + custom_llm_provider="openai", + api_base="https://api.openai.com", + metadata=StandardLoggingMetadata( + user_api_key_hash="test_hash", + user_api_key_alias=None, + user_api_key_team_id=None, + user_api_key_org_id=None, + user_api_key_user_id=None, + user_api_key_team_alias=None, + spend_logs_metadata=None, + requester_ip_address=None, + requester_metadata=None, + user_api_key_end_user_id=None, + usage_object=usage_object, + ), + cache_hit=False, + cache_key=None, + saved_cache_cost=0.0, + request_tags=[], + end_user=None, + requester_ip_address=None, + messages=[], + response={}, + error_str=None, + model_parameters={}, + hidden_params=StandardLoggingHiddenParams( + model_id="model-123", + cache_key=None, + api_base="https://api.openai.com", + response_cost="0.02", + litellm_overhead_time_ms=None, + additional_headers=None, + batch_models=None, + litellm_model_name=None, + usage_object=None, + ), + ) + + +def test_get_logging_payload_maps_responses_api_cache_write_tokens_from_usage_object(): + """Responses API (/v1/responses) usage is not chat-Usage-shaped, so + additional_usage_values can't derive cache tokens from response_obj.usage. + The Admin UI Logs "Cache Creation Tokens" row reads + additional_usage_values.cache_creation_input_tokens, so it must be filled + from the normalized standard_logging usage_object (LIT-4633).""" + standard_logging_payload = _make_standard_logging_payload_with_usage_object( + usage_object={ + "prompt_tokens": 1000, + "completion_tokens": 10, + "total_tokens": 1010, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 800, "cache_creation_tokens": 800}, + } + ) + payload = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": standard_logging_payload, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 800}, + }, + }, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + additional_usage_values = json.loads(payload["metadata"])["additional_usage_values"] + assert additional_usage_values["cache_creation_input_tokens"] == 800 + + def test_sanitize_request_body_for_spend_logs_payload_basic(): request_body = { "messages": [{"role": "user", "content": "Hello, how are you?"}], diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1db76aed61d..e6ccbc579b5 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -3,6 +3,7 @@ from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache @@ -1562,6 +1563,100 @@ async def test_should_skip_reservation_when_counter_increment_fails( ) +@pytest.mark.asyncio +async def test_should_raise_503_when_counter_increment_fails_and_fail_closed( + spend_counter_state, + monkeypatch, +): + """#33923: with fail_closed_budget_enforcement on, a failed reservation write + must reject instead of silently degrading to read-time-only enforcement.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-reserve-fail-closed", + spend=0.0, + max_budget=1.0, + ) + + async def fail_increment_cache(*args, **kwargs): + raise RuntimeError("counter unavailable") + + monkeypatch.setattr(counter_cache, "async_increment_cache", fail_increment_cache) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(HTTPException) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.status_code == 503 + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-reserve-fail-closed" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_fail_closed_releases_earlier_counters_before_503( + spend_counter_state, +): + """#33923: when a later counter's reservation write fails in strict mode, the + counters that already reserved must be released before the 503 propagates.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-fail-closed-release", + spend=0.0, + max_budget=1.0, + budget_limits=[ + { + "budget_duration": "1h", + "max_budget": 1.0, + } + ], + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(HTTPException) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.status_code == 503 + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-fail-closed-release" + ) + == 0.0 + ) + + @pytest.mark.asyncio async def test_should_skip_reservation_when_counter_initialization_fails( spend_counter_state, diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 8bee7e9f33b..1437899f561 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -5226,3 +5226,198 @@ async def test_add_litellm_data_to_request_unions_metadata_tags_with_header_tags tags = updated["litellm_metadata"]["tags"] assert "header-tag" in tags assert "body-tag" in tags + + +def _make_chat_request_mock() -> MagicMock: + return _make_request_mock("/v1/chat/completions", {"Content-Type": "application/json"}) + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_clobbers_caller_supplied_user(monkeypatch): + """The flag exists so providers can ban by a tamper-proof id; a caller-chosen + `user` must never survive, and the raw sk- key must never be forwarded.""" + from litellm.proxy._types import hash_token + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + raw_key = "sk-overwrite-user-test-1234" + user_api_key_dict = UserAPIKeyAuth(api_key=raw_key) + user_api_key_dict.via_virtual_key = True + data = {"model": "gpt-4o", "user": "attacker-chosen-id"} + + updated_data = await add_litellm_data_to_request( + data=data, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == hash_token(raw_key) + assert updated_data["user"] != "attacker-chosen-id" + assert raw_key not in updated_data["user"] + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_sets_user_when_absent(monkeypatch): + from litellm.proxy._types import hash_token + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + raw_key = "sk-overwrite-user-test-5678" + user_api_key_dict = UserAPIKeyAuth(api_key=raw_key) + user_api_key_dict.via_virtual_key = True + data = {"model": "gpt-4o"} + + updated_data = await add_litellm_data_to_request( + data=data, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == hash_token(raw_key) + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_disabled_preserves_caller_user(): + assert litellm.overwrite_user_with_key_hash is False + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-overwrite-user-test-9999") + user_api_key_dict.via_virtual_key = True + data = {"model": "gpt-4o", "user": "caller-chosen-id"} + + updated_data = await add_litellm_data_to_request( + data=data, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_skips_custom_auth_credential(monkeypatch): + """Custom-auth credentials are not sk-prefixed or JWTs, so UserAPIKeyAuth stores + them raw; the stamp must skip them entirely so auth material never leaks.""" + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + raw_credential = "my-custom-auth-credential-abc123" + user_api_key_dict = UserAPIKeyAuth(api_key=raw_credential) + assert user_api_key_dict.api_key == raw_credential + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "caller-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_skips_jwt_auth(monkeypatch): + """A hashed JWT rotates on every token re-issue, so it is useless as a stable + ban id; JWT-authenticated requests are not stamped.""" + from litellm.proxy._types import hash_token + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + hashed_jwt = f"hashed-jwt-{hash_token('some-jwt-token')}" + user_api_key_dict = UserAPIKeyAuth(api_key=hashed_jwt) + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "caller-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_skips_hex_shaped_custom_credential(monkeypatch): + """A custom-auth credential that happens to be 64 hex chars is indistinguishable + from a key hash by shape alone; only the server-set via_virtual_key marker may + authorize stamping, so this raw credential must never be forwarded.""" + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + hex_shaped_credential = "a" * 64 + user_api_key_dict = UserAPIKeyAuth(api_key=hex_shaped_credential) + assert user_api_key_dict.api_key == hex_shaped_credential + assert user_api_key_dict.via_virtual_key is False + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "caller-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" + + +def test_via_virtual_key_cannot_be_forged_from_validated_input(): + from_kwargs = UserAPIKeyAuth(api_key="b" * 64, via_virtual_key=True) + assert from_kwargs.via_virtual_key is False + + from_dict = UserAPIKeyAuth.model_validate({"api_key": "b" * 64, "via_virtual_key": True}) + assert from_dict.via_virtual_key is False + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_stamps_master_key_alias(monkeypatch): + """Master-key requests carry the stable alias instead of a hash (so the master + key never propagates anywhere); the alias is the stampable id for them.""" + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + user_api_key_dict = UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS) + user_api_key_dict.via_virtual_key = True + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "attacker-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == LITELLM_PROXY_MASTER_KEY_ALIAS + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_rejects_alias_without_marker(monkeypatch): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + user_api_key_dict = UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS) + assert user_api_key_dict.via_virtual_key is False + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "caller-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 85dbf70b452..20451f5d0ac 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -617,6 +617,75 @@ class TestProxySettingEndpoints: create_sso_settings = json.loads(create_data["sso_settings"]) assert create_sso_settings["google_client_id"] == "new_google_client_id" + def test_update_sso_settings_maps_saml_fields_to_env_vars( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """SAML settings entered in the admin UI must be applied as the SAML_* env + vars the SAML handler reads, and the allow-unsolicited toggle must map to + the 'true'/'false' string the handler expects.""" + import json + import os + from unittest.mock import AsyncMock, MagicMock + + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + mock_prisma.db.litellm_config = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr( + proxy_config, + "_encrypt_env_variables", + lambda environment_variables: environment_variables, + ) + + for var in ( + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + "SAML_SP_ENTITY_ID", + "SAML_ALLOW_UNSOLICITED", + ): + monkeypatch.delenv(var, raising=False) + + new_sso_settings = { + "saml_idp_metadata_url": "https://idp.example.com/metadata", + "saml_sp_entity_id": "https://proxy.example.com/sso/saml/metadata", + "saml_allow_unsolicited": "true", + "proxy_base_url": "https://proxy.example.com", + "user_email": "admin@example.com", + } + + try: + response = client.patch("/update/sso_settings", json=new_sso_settings) + + assert response.status_code == 200 + + assert os.environ.get("SAML_IDP_METADATA_URL") == "https://idp.example.com/metadata" + assert os.environ.get("SAML_SP_ENTITY_ID") == "https://proxy.example.com/sso/saml/metadata" + assert os.environ.get("SAML_ALLOW_UNSOLICITED") == "true" + assert "SAML_IDP_METADATA_XML" not in os.environ + + stored = json.loads( + mock_prisma.db.litellm_ssoconfig.upsert.call_args.kwargs["data"]["create"]["sso_settings"] + ) + assert stored["saml_idp_metadata_url"] == "https://idp.example.com/metadata" + assert stored["saml_allow_unsolicited"] == "true" + finally: + for var in ( + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + "SAML_SP_ENTITY_ID", + "SAML_ALLOW_UNSOLICITED", + ): + os.environ.pop(var, None) + def test_update_sso_settings_audits_when_env_cleanup_fails( self, mock_proxy_config, mock_auth, monkeypatch ): diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 64c14abfd83..5c711fc6c34 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -626,6 +626,7 @@ def _moderation_guardrail() -> MagicMock: cb.should_run_guardrail = MagicMock(return_value=True) cb.async_moderation_hook = AsyncMock(return_value=None) cb.async_post_call_success_hook = AsyncMock(return_value=None) + cb.run_in_parallel = False return cb diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py index 6a339b37a80..715d66db181 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -28,6 +28,7 @@ def _make_guardrail(name="g", should_run=True, override=None): cb.event_hook = GuardrailEventHooks.post_call cb.should_run_guardrail = MagicMock(return_value=should_run) cb.async_post_call_success_hook = AsyncMock(return_value=override) + cb.run_in_parallel = False return cb diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 3a75a33fdc7..0141cf5d96a 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -369,6 +369,32 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details.image_tokens == 272 assert result.completion_tokens_details.text_tokens == 100 + def test_transform_response_api_usage_maps_cache_write_tokens(self): + """Responses API (/v1/responses) cache-write tokens must survive the usage transform. + + gpt-5.6 returns usage.input_tokens_details.cache_write_tokens (an extra field + not typed on InputTokensDetails). Before the fix the transform rebuilt the token + details and dropped it, leaving the cache-creation metric empty (LIT-4633). + """ + usage = { + "input_tokens": 10062, + "output_tokens": 16, + "total_tokens": 10078, + "input_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 10059, + }, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.cache_write_tokens == 10059 + assert result.prompt_tokens_details.cache_creation_tokens == 10059 + assert result.prompt_tokens_details.cached_tokens == 0 + def test_transform_response_api_usage_mixed_details(self): """Test transformation handles mixed token details (cached + image + audio).""" # Setup - hypothetical usage with mixed token types diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/test_litellm/test_claude_opus_5_config.py new file mode 100644 index 00000000000..84021a83a5a --- /dev/null +++ b/tests/test_litellm/test_claude_opus_5_config.py @@ -0,0 +1,277 @@ +""" +Validate Claude Opus 5 model configuration entries. + +Opus 5 carries Opus 4.8's pricing ($5 / $25 per MTok) and the gen-5 adaptive +thinking profile, but differs from 4.8 in two ways that are behavior-bearing in +LiteLLM: the cacheable-prefix minimum drops to 512 tokens, and Bedrock's Opus 5 +validator accepts the full effort ladder, so the entries must not carry the +``bedrock_output_config_effort_ceiling`` that silently clamps ``max`` to +``xhigh`` on 4.8. The cost-map entries are also what populate +``litellm.anthropic_models`` at import, which is what lets a bare +``claude-opus-5`` name resolve to the ``anthropic`` provider (and match an +``anthropic/*`` wildcard deployment). +""" + +import json +import os + +import pytest + +import litellm +from litellm.constants import BEDROCK_CONVERSE_MODELS +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + +REPO_ROOT = os.path.join(os.path.dirname(__file__), "../..") + +ALL_OPUS_5_VARIANTS = ( + "claude-opus-5", + "anthropic.claude-opus-5", + "global.anthropic.claude-opus-5", + "us.anthropic.claude-opus-5", + "eu.anthropic.claude-opus-5", + "au.anthropic.claude-opus-5", + "jp.anthropic.claude-opus-5", + "vertex_ai/claude-opus-5", + "vertex_ai/claude-opus-5@default", + "azure_ai/claude-opus-5", +) + +BEDROCK_OPUS_5_VARIANTS = ( + "anthropic.claude-opus-5", + "global.anthropic.claude-opus-5", + "us.anthropic.claude-opus-5", + "eu.anthropic.claude-opus-5", + "au.anthropic.claude-opus-5", + "jp.anthropic.claude-opus-5", +) + + +def _load_root_cost_map() -> dict: + json_path = os.path.join(REPO_ROOT, "model_prices_and_context_window.json") + with open(json_path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled backup cost map so assertions don't depend on the + network-fetched ``main`` copy (which lags this branch until merge).""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +def test_opus_5_pricing_and_capabilities(): + model_data = _load_root_cost_map() + + expected_providers = { + "claude-opus-5": "anthropic", + "anthropic.claude-opus-5": "bedrock_converse", + "vertex_ai/claude-opus-5": "vertex_ai-anthropic_models", + "azure_ai/claude-opus-5": "azure_ai", + } + + for model_name, provider in expected_providers.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + + assert info["litellm_provider"] == provider + assert info["mode"] == "chat" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + + # Opus 5 ships at Opus 4.8's rates: $5 / $25 per MTok, with the standard + # 1.25x cache-write, 2x 1-hour cache-write, and 0.1x cache-read multipliers. + assert info["input_cost_per_token"] == 5e-06 + assert info["output_cost_per_token"] == 2.5e-05 + assert info["cache_creation_input_token_cost"] == 6.25e-06 + assert info["cache_creation_input_token_cost_above_1hr"] == 1e-05 + assert info["cache_read_input_token_cost"] == 5e-07 + + # Flat rate across the full 1M window, no long-context premium. + assert "input_cost_per_token_above_200k_tokens" not in info + assert "output_cost_per_token_above_200k_tokens" not in info + + # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no + # assistant prefill. + assert info["supports_adaptive_thinking"] is True + assert info["supports_reasoning"] is True + assert info["supports_sampling_params"] is False + assert info["supports_assistant_prefill"] is False + assert info["supports_xhigh_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + + +def test_opus_5_bedrock_regional_pricing(): + """Global/base endpoints use base pricing; the us./eu./au./jp. regional + cross-region inference profiles carry a 10% premium.""" + model_data = _load_root_cost_map() + + base_pricing = { + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + } + regional_pricing = { + "input_cost_per_token": 5.5e-06, + "output_cost_per_token": 2.75e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + } + + expected = { + "anthropic.claude-opus-5": base_pricing, + "global.anthropic.claude-opus-5": base_pricing, + "us.anthropic.claude-opus-5": regional_pricing, + "eu.anthropic.claude-opus-5": regional_pricing, + "au.anthropic.claude-opus-5": regional_pricing, + "jp.anthropic.claude-opus-5": regional_pricing, + } + + for model_name, pricing in expected.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + assert info["litellm_provider"] == "bedrock_converse" + for key, value in pricing.items(): + assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}" + + +@pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) +def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name): + """Bedrock accepts every effort level for Opus 5, so no clamp belongs here. + + Opus 4.7/4.8 carry ``bedrock_output_config_effort_ceiling: "xhigh"``, which + is what ``normalize_bedrock_opus_output_config_effort`` reads to rewrite a + caller's effort down. Verified against Bedrock on 2026-07-24 that + ``output_config.effort="max"`` returns 200 for the Opus 5 profiles, so the + ceiling is deliberately absent; adding one back would silently downgrade + requests. + + This asserts the cost-map entry rather than calling the normalizer because + ``_BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER`` currently ranks ``max`` (3) below + ``xhigh`` (4), so an ``xhigh`` ceiling never clamps ``max`` and a behavioral + assertion would pass either way. Keeping the entry clean means Opus 5 stays + correct once that ordering is fixed.""" + info = _load_root_cost_map()[model_name] + assert "bedrock_output_config_effort_ceiling" not in info + + +@pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) +def test_opus_5_bedrock_rejects_strict_tools(model_name, local_model_cost_map): + """Bedrock Converse routes Opus through a validator that rejects + ``toolSpec.strict`` (``tools.0.custom.strict: Extra inputs are not + permitted``), same as Opus 4.7/4.8; verified against Bedrock on 2026-07-24. + Without the flag LiteLLM forwards ``strict`` and every tool call 400s.""" + from litellm.llms.bedrock.common_utils import bedrock_converse_supports_strict_tools + + assert bedrock_converse_supports_strict_tools(model_name) is False + + +def test_opus_5_prompt_cache_minimum_is_512(local_model_cost_map): + """Opus 5 halves the cacheable-prefix minimum (Opus 4.8 is 1024). + + The router's prompt-caching deployment check reads this value, so a stale + 1024 would route prompts of 512-1023 tokens away from a warm Opus 5 + deployment even though they cache fine.""" + from litellm.utils import get_prompt_cache_min_tokens + + assert get_prompt_cache_min_tokens(model="claude-opus-5") == 512 + assert get_prompt_cache_min_tokens(model="us.anthropic.claude-opus-5") == 512 + + +def test_opus_5_supports_fast_mode(local_model_cost_map): + """Fast mode is Opus 5 on the first-party API at $10 / $50 per MTok, i.e. 2x + base. ``supports_speed`` gates whether ``speed="fast"`` is forwarded at all, + and ``provider_specific_entry.fast`` is what prices the response.""" + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.anthropic.cost_calculation import ( + cost_per_token as anthropic_cost_per_token, + ) + from litellm.types.utils import Usage + + assert ( + AnthropicConfig._model_supports_speed_param("claude-opus-5", "anthropic") is True + ) + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + usage.speed = "fast" + prompt_cost, completion_cost = anthropic_cost_per_token( + model="claude-opus-5", usage=usage + ) + assert prompt_cost == pytest.approx(1000 * 5e-06 * 2.0) + assert completion_cost == pytest.approx(500 * 2.5e-05 * 2.0) + + +def test_opus_5_present_in_bundled_backup(): + """The bundled backup is the runtime fallback (and what tests load with + ``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the + root cost map, otherwise the model resolves on one path but not the other.""" + backup = GetModelCostMap.load_local_model_cost_map() + for model_name in ALL_OPUS_5_VARIANTS: + assert model_name in backup, f"Missing from backup cost map: {model_name}" + + +def test_opus_5_registered_for_bedrock_converse(): + assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS + + +def test_opus_5_provider_resolves_via_model_info(local_model_cost_map): + """Regression: ``claude-opus-5`` must resolve to provider ``anthropic``. + + Without the cost-map entry the model is unknown to LiteLLM, so it cannot be + tied to the ``anthropic`` provider and an ``anthropic/*`` wildcard deployment + would not match it.""" + info = litellm.get_model_info(model="claude-opus-5") + assert info["litellm_provider"] == "anthropic" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_opus_5_all_variants_carry_adaptive_thinking_flag(cost_map): + """Every Opus 5 entry must advertise ``supports_adaptive_thinking``. + + Adaptive-thinking detection is cost-map driven, so a single variant missing + the flag silently sends the legacy ``thinking.type='enabled'`` shape, which + Opus 5 rejects with a 400.""" + variants = [k for k in cost_map if "claude-opus-5" in k] + assert variants, "no claude-opus-5 entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_adaptive_thinking") is not True + ] + assert not missing, f"missing supports_adaptive_thinking: {missing}" + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_opus_5_all_variants_carry_512_token_cache_minimum(cost_map): + variants = [k for k in cost_map if "claude-opus-5" in k] + assert variants, "no claude-opus-5 entries found in cost map" + wrong = { + k: cost_map[k].get("prompt_cache_min_tokens") + for k in variants + if cost_map[k].get("prompt_cache_min_tokens") != 512 + } + assert not wrong, f"prompt_cache_min_tokens must be 512: {wrong}" diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 9636db4f4cd..276ee96ed65 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -12,6 +12,7 @@ from pydantic import BaseModel import litellm from litellm.cost_calculator import ( + BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, completion_cost, cost_per_token, @@ -3479,3 +3480,32 @@ def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate(): ) assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2) + + +def test_combine_usage_objects_sums_mirrored_cache_write_fields_once(): + """ + cache_write_tokens and cache_creation_tokens mirror each other on + PromptTokensDetailsWrapper, so field-iterating aggregation must sum the pair + once: a single 50-token usage stays 50 and two combine to 100, not double. + """ + single = Usage( + prompt_tokens=100, + completion_tokens=10, + total_tokens=110, + prompt_tokens_details=PromptTokensDetailsWrapper(cache_write_tokens=50), + ) + combined = BaseTokenUsageProcessor.combine_usage_objects([single]) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.cache_write_tokens == 50 + assert combined.prompt_tokens_details.cache_creation_tokens == 50 + + anthropic_style = Usage( + prompt_tokens=100, + completion_tokens=10, + total_tokens=110, + cache_creation_input_tokens=50, + ) + combined_pair = BaseTokenUsageProcessor.combine_usage_objects([anthropic_style, anthropic_style]) + assert combined_pair.prompt_tokens_details is not None + assert combined_pair.prompt_tokens_details.cache_write_tokens == 100 + assert combined_pair.prompt_tokens_details.cache_creation_tokens == 100 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index a1a9448cc58..edc0cfed63e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -17,7 +17,9 @@ from litellm.types.utils import ( Delta, LlmProviders, ModelResponseStream, + PromptTokensDetailsWrapper, StreamingChoices, + Usage, ) from litellm.utils import ( ProviderConfigManager, @@ -34,6 +36,57 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_usage_openai_cache_write_tokens_populates_both_names(): + """OpenAI reports cache-write tokens as prompt_tokens_details.cache_write_tokens. + The Usage constructor must expose it under both cache_write_tokens (canonical, + OpenAI naming) and cache_creation_tokens (legacy, Anthropic naming).""" + usage = Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details={"cached_tokens": 0, "cache_write_tokens": 800}, + ) + assert usage.prompt_tokens_details.cache_write_tokens == 800 + assert usage.prompt_tokens_details.cache_creation_tokens == 800 + + +def test_usage_anthropic_cache_creation_maps_to_cache_write_tokens(): + """Anthropic/Bedrock report the top-level cache_creation_input_tokens field. + It must be normalized onto the OpenAI cache_write_tokens name as well as the + legacy cache_creation_tokens name.""" + usage = Usage( + prompt_tokens=500, + completion_tokens=50, + total_tokens=550, + cache_creation_input_tokens=300, + cache_read_input_tokens=120, + ) + assert usage.prompt_tokens_details.cache_write_tokens == 300 + assert usage.prompt_tokens_details.cache_creation_tokens == 300 + assert usage.prompt_tokens_details.cached_tokens == 120 + + +def test_prompt_tokens_details_no_cache_write_tokens_when_absent(): + """A read-only cache hit (no cache write) must not surface cache-write fields.""" + details = PromptTokensDetailsWrapper(cached_tokens=800) + assert details.cached_tokens == 800 + assert not hasattr(details, "cache_write_tokens") + assert not hasattr(details, "cache_creation_tokens") + + +def test_prompt_tokens_details_cache_write_creation_stay_in_sync_on_assignment(): + """Assigning either name after construction must mirror to the other, so a + caller that sets only one field can't leave the pair silently out of sync.""" + details = PromptTokensDetailsWrapper(cache_write_tokens=100) + assert details.cache_write_tokens == details.cache_creation_tokens == 100 + + details.cache_write_tokens = 250 + assert details.cache_write_tokens == details.cache_creation_tokens == 250 + + details.cache_creation_tokens = 375 + assert details.cache_write_tokens == details.cache_creation_tokens == 375 + + @pytest.fixture def local_model_cost_map(monkeypatch): original_model_cost = litellm.model_cost diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 4fa5528aab8..ce743063309 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -4,11 +4,6 @@ "count": 1 } }, - "src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx": { "no-restricted-imports": { "count": 2 @@ -24,11 +19,6 @@ "count": 1 } }, - "src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx": { "no-restricted-imports": { "count": 2 @@ -37,16 +27,6 @@ "count": 1 } }, - "src/app/(dashboard)/agents/_components/AgentsPanel.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/agents/_components/AgentsTable.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/agents/_components/add_agent_form.tsx": { "local/filename-pascal-case": { "count": 1 @@ -71,9 +51,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/refs": { "count": 3 }, @@ -84,9 +61,6 @@ "src/app/(dashboard)/agents/_components/agent_cost_view.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 2 } }, "src/app/(dashboard)/agents/_components/agent_form_fields.tsx": { @@ -123,9 +97,6 @@ }, "no-nested-ternary": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/agents/_components/cost_config_fields.tsx": { @@ -163,9 +134,6 @@ "src/app/(dashboard)/budgets/_components/budget_panel.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx": { @@ -196,9 +164,6 @@ "src/app/(dashboard)/caching/_components/cache_health.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx": { @@ -206,11 +171,6 @@ "count": 1 } }, - "src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts": { "no-restricted-imports": { "count": 1 @@ -232,11 +192,6 @@ "count": 1 } }, - "src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/caching/_components/coordination_redis_settings/coordinationRedisFields.ts": { "no-restricted-imports": { "count": 1 @@ -857,15 +812,7 @@ "count": 1 } }, - "src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx": { - "no-restricted-imports": { - "count": 1 - }, "react-hooks/immutability": { "count": 2 } @@ -880,11 +827,6 @@ "count": 1 } }, - "src/app/(dashboard)/mcp-servers/_components/MCPServerCard.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, "src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -923,9 +865,6 @@ } }, "src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx": { - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -966,11 +905,6 @@ "count": 1 } }, - "src/app/(dashboard)/mcp-servers/_components/TruePassthroughWarning.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx": { "no-nested-ternary": { "count": 2 @@ -1018,9 +952,6 @@ }, "no-nested-ternary": { "count": 3 - }, - "no-restricted-imports": { - "count": 2 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx": { @@ -1030,9 +961,6 @@ "local/no-complex-jsx-arrow": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 2 } @@ -1040,17 +968,11 @@ "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 2 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx": { @@ -1076,9 +998,6 @@ "src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 2 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx": { @@ -1091,9 +1010,6 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 2 } @@ -1101,9 +1017,6 @@ "src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 2 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx": { @@ -1116,9 +1029,6 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 2 } @@ -1128,21 +1038,11 @@ "count": 1 } }, - "src/app/(dashboard)/memory/_components/MemoryDetailDrawer.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/memory/_components/MemoryEditModal.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/app/(dashboard)/memory/_components/MemoryView.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab.test.tsx": { "react/display-name": { "count": 1 @@ -1161,6 +1061,11 @@ "count": 1 } }, + "src/app/(dashboard)/models-and-endpoints/layout.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.ts": { "prefer-const": { "count": 6 @@ -1483,9 +1388,6 @@ "no-nested-ternary": { "count": 10 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/immutability": { "count": 1 } @@ -1497,9 +1399,6 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -1523,9 +1422,6 @@ "src/app/(dashboard)/policies/_components/impact_preview_alert.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/policies/_components/index.test.tsx": { @@ -1540,9 +1436,6 @@ "local/no-complex-jsx-arrow": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -1557,9 +1450,6 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 2 } @@ -1568,9 +1458,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -1578,9 +1465,6 @@ "src/app/(dashboard)/policies/_components/policy_templates.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/policies/_components/policy_test_panel.tsx": { @@ -1598,9 +1482,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/immutability": { "count": 1 }, @@ -1854,9 +1735,6 @@ } }, "src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx": { - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -1878,9 +1756,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -1902,9 +1777,6 @@ } }, "src/app/(dashboard)/ui-theme/UIThemeSettings.tsx": { - "no-restricted-imports": { - "count": 1 - }, "no-restricted-syntax": { "count": 3 }, @@ -2068,11 +1940,6 @@ "count": 1 } }, - "src/app/(dashboard)/vector-stores/_components/TestVectorStoreTab.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx": { "no-nested-ternary": { "count": 2 @@ -2084,18 +1951,10 @@ "count": 1 } }, - "src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/vector-stores/_components/index.tsx": { "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -2232,26 +2091,11 @@ "count": 1 } }, - "src/components/CloudZeroCostTracking/CloudZeroCostTracking.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CloudZeroCostTracking/CloudZeroEmptyPlaceholder.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx": { "no-restricted-imports": { "count": 1 @@ -2449,11 +2293,6 @@ "count": 1 } }, - "src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/Settings/AdminSettings/PluginSettings/PluginSettings.tsx": { "no-restricted-imports": { "count": 1 @@ -2619,20 +2458,7 @@ } }, "src/components/ToolDetail.tsx": { - "no-restricted-imports": { - "count": 1 - }, "unused-imports/no-unused-imports": { - "count": 2 - } - }, - "src/components/ToolPolicies/PolicySelect.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/ToolPolicies/ToolPoliciesTableColumns.tsx": { - "no-restricted-imports": { "count": 1 } }, @@ -3031,23 +2857,10 @@ } }, "src/components/common_components/Filters/FilterInput.tsx": { - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/common_components/Filters/FiltersButton.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/common_components/Filters/ResetFiltersButton.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/IconActionButton/BaseActionButton.tsx": { "no-restricted-imports": { "count": 1 @@ -3254,9 +3067,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/immutability": { "count": 1 } @@ -3265,9 +3075,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "prefer-const": { "count": 1 } @@ -3609,7 +3416,7 @@ "count": 1 }, "no-restricted-imports": { - "count": 3 + "count": 1 } }, "src/components/page_utils.test.ts": { @@ -4130,6 +3937,11 @@ "count": 1 } }, + "src/components/ui/radio-group.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/ui/scroll-area.tsx": { "local/filename-pascal-case": { "count": 1 @@ -4318,7 +4130,7 @@ }, "src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx": { "no-nested-ternary": { - "count": 3 + "count": 2 }, "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx index 72a89093bdb..9476a8d98af 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx @@ -1,68 +1,63 @@ import { useAccessGroupDetails } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails"; -import { - Button, - Card, - Col, - Descriptions, - Empty, - Flex, - Layout, - List, - Row, - Spin, - Tabs, - Tag, - theme, - Typography, -} from "antd"; import { ArrowLeftIcon, BotIcon, EditIcon, KeyIcon, LayersIcon, ServerIcon, UsersIcon } from "lucide-react"; import { useState } from "react"; import DefaultProxyAdminTag from "@/components/common_components/DefaultProxyAdminTag"; +import CopyButton from "@/components/shared/CopyButton"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card, CardAction, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { AccessGroupEditModal } from "./AccessGroupsModal/AccessGroupEditModal"; -const { Title, Text } = Typography; -const { Content } = Layout; - interface AccessGroupDetailProps { accessGroupId: string; onBack: () => void; } +const MAX_PREVIEW = 5; + +function ResourceList({ ids, emptyMessage }: { ids: string[]; emptyMessage: string }) { + if (ids.length === 0) { + return

{emptyMessage}

; + } + return ( +
+ {ids.map((id) => ( + + + {id} + + + ))} +
+ ); +} + export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailProps) { const { data: accessGroup, isLoading } = useAccessGroupDetails(accessGroupId); - const { token } = theme.useToken(); const [isEditModalVisible, setIsEditModalVisible] = useState(false); const [showAllKeys, setShowAllKeys] = useState(false); const [showAllTeams, setShowAllTeams] = useState(false); - const MAX_PREVIEW = 5; - if (isLoading) { return ( - - - - - +
+
+ +
+
); } if (!accessGroup) { return ( - - +

Access group not found

+ ); } @@ -75,224 +70,159 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr const displayedKeys = showAllKeys ? keyIds : keyIds.slice(0, MAX_PREVIEW); const displayedTeams = showAllTeams ? teamIds : teamIds.slice(0, MAX_PREVIEW); - const handleEdit = () => { - setIsEditModalVisible(true); - }; - - const tabItems = [ - { - key: "models", - label: ( - - - Models - {modelIds?.length} - - ), - children: - modelIds?.length > 0 ? ( - ( - - - {id} - - - )} - /> - ) : ( - - ), - }, - { - key: "mcp", - label: ( - - - MCP Servers - {mcpServerIds?.length} - - ), - children: - mcpServerIds?.length > 0 ? ( - ( - - - {id} - - - )} - /> - ) : ( - - ), - }, - { - key: "agents", - label: ( - - - Agents - {agentIds?.length} - - ), - children: - agentIds?.length > 0 ? ( - ( - - - {id} - - - )} - /> - ) : ( - - ), - }, - ]; - return ( - - {/* Header */} -
-
-
- - {accessGroup.access_group_name} - - - ID: {accessGroup.access_group_id} - +

{accessGroup.access_group_name}

+
+ ID: {accessGroup.access_group_id} + +
-
- {/* Group Details */} - - - - {accessGroup.description || "—"} - + + + Group Details + + +
+
Description
+
{accessGroup.description || "—"}
+
Created
+
{new Date(accessGroup.created_at).toLocaleString()} {accessGroup.created_by && ( - -  {"by"}  + <> + by - + )} - - +
+
Last Updated
+
{new Date(accessGroup.updated_at).toLocaleString()} {accessGroup.updated_by && ( - -  {"by"}  + <> + by - + )} - - - - - - {/* Attached Keys & Teams */} - - - - - Attached Keys - {keyIds?.length} - - } - extra={ - keyIds?.length > MAX_PREVIEW ? ( - - ) : null - } - > - {keyIds?.length > 0 ? ( - - {displayedKeys.map((id) => ( - - - {id.length > 20 ? `${id.slice(0, 10)}...${id.slice(-6)}` : id} - - - ))} - - ) : ( - - )} - - - - - - Attached Teams - {teamIds?.length} - - } - extra={ - teamIds?.length > MAX_PREVIEW ? ( - - ) : null - } - > - {teamIds?.length > 0 ? ( - - {displayedTeams.map((id) => ( - - - {id} - - - ))} - - ) : ( - - )} - - - - - {/* Resources Tabs */} - - +
+
+ + + +
+ + + + + Attached Keys + {keyIds.length} + + {keyIds.length > MAX_PREVIEW && ( + + + + )} + + + {keyIds.length > 0 ? ( +
+ {displayedKeys.map((id) => ( + + {id.length > 20 ? `${id.slice(0, 10)}...${id.slice(-6)}` : id} + + ))} +
+ ) : ( +

No keys attached

+ )} + + + + + + + + Attached Teams + {teamIds.length} + + {teamIds.length > MAX_PREVIEW && ( + + + + )} + + + {teamIds.length > 0 ? ( +
+ {displayedTeams.map((id) => ( + + {id} + + ))} +
+ ) : ( +

No teams attached

+ )} + + +
+ + + + + + + + Models + {modelIds.length} + + + + MCP Servers + {mcpServerIds.length} + + + + Agents + {agentIds.length} + + + + + + + + + + + + + - {/* Edit Modal */} setIsEditModalVisible(false)} /> - + ); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx index 0de6596f57c..f37acb3d85a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx @@ -1,10 +1,11 @@ import { AccessGroupResponse, useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; import { useDeleteAccessGroup } from "@/app/(dashboard)/hooks/accessGroups/useDeleteAccessGroup"; -import { PlusOutlined } from "@ant-design/icons"; -import { Button, Flex, Input, Layout, Space, theme, Typography } from "antd"; -import { SearchIcon } from "lucide-react"; +import { Plus, SearchIcon, X } from "lucide-react"; import { useMemo, useState } from "react"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import { PageHeader } from "@/components/shared/PageHeader"; +import { Button } from "@/components/ui/button"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { AccessGroupDetail } from "./AccessGroupsDetailsPage"; import { AccessGroupCreateModal } from "./AccessGroupsModal/AccessGroupCreateModal"; import { AccessGroupsTable } from "./AccessGroupsTable"; @@ -12,9 +13,6 @@ import { AccessGroup } from "./types"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { isProxyAdminRole } from "@/utils/roles"; -const { Title, Text } = Typography; -const { Content } = Layout; - function mapResponseToAccessGroup(r: AccessGroupResponse): AccessGroup { return { id: r.access_group_id, @@ -33,7 +31,6 @@ function mapResponseToAccessGroup(r: AccessGroupResponse): AccessGroup { } export function AccessGroupsPage() { - const { token } = theme.useToken(); const { userRole } = useAuthorized(); // Admin Viewer follows the read-parity rule: see access groups, no writes. const canModify = isProxyAdminRole(userRole ?? ""); @@ -62,31 +59,41 @@ export function AccessGroupsPage() { } return ( - - - - - Access Groups - - Manage resource permissions for your organization - - {canModify && ( - - )} - - - - } - placeholder="Search groups by name, ID, or description..." - style={{ maxWidth: 400 }} - value={searchText} - onChange={(e) => setSearchText(e.target.value)} - allowClear +
+
+ setIsCreateModalVisible(true)}> + + Create Access Group + + ) : undefined + } /> - +
+ +
+ + + + + setSearchText(e.target.value)} + /> + {searchText && ( + + setSearchText("")}> + + + + )} + +
- +
); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx index 441d300436a..d873687378b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx @@ -140,8 +140,9 @@ describe("AgentsPanel", () => { await user.click(await screen.findByTestId("agent-actions-agent-9")); await user.click(await screen.findByTestId("agent-action-delete")); - const modal = await screen.findByRole("dialog"); - await user.click(within(modal).getByRole("button", { name: /^delete$/i })); + const confirmPrompt = await screen.findByText(/are you sure you want to delete agent: Doomed Agent\?/i); + const confirmDialog = confirmPrompt.closest('[role="dialog"],[role="alertdialog"]') as HTMLElement; + await user.click(within(confirmDialog).getByRole("button", { name: /^delete$/i })); await waitFor(() => { expect(networking.deleteAgentCall).toHaveBeenCalledWith("test-token", "agent-9"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx index a4a71530c84..4459ee0c377 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx @@ -1,6 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Modal, Alert } from "antd"; -import { Plus } from "lucide-react"; +import { Info, Plus } from "lucide-react"; import { getAgentsList, deleteAgentCall } from "@/components/networking"; import AddAgentForm from "./add_agent_form"; import { isAdminRole } from "@/utils/roles"; @@ -9,6 +8,16 @@ import AgentsTable from "./AgentsTable"; import NotificationsManager from "@/components/molecules/notifications_manager"; import { Agent } from "@/components/agents/types"; import { Team } from "@/components/key_team_helpers/key_list"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { + AlertDialog, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; interface AgentsPanelProps { @@ -130,17 +139,18 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams

Agents

-

+

List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents public.

- + + + Why do agents need keys? + + Keys scope access to an agent and allow it to call MCP tools. Assign a key when creating an agent or from + the Virtual Keys page. + + {isAdmin && (
+ + + )}
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx index 824ae47f3e6..359cb49b910 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx @@ -1,13 +1,13 @@ "use client"; import { SortingState } from "@tanstack/react-table"; -import { Tooltip, Switch } from "antd"; -import { CheckCircleOutlined } from "@ant-design/icons"; -import { Bot } from "lucide-react"; +import { Bot, CircleCheck } from "lucide-react"; import React, { useMemo, useState } from "react"; import { Agent } from "@/components/agents/types"; import { DataTable } from "@/components/shared/DataTable"; +import { Switch } from "@/components/ui/switch"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { getAgentsTableColumns } from "./AgentsTableColumns"; @@ -67,18 +67,27 @@ const AgentsTable: React.FC = ({ size="compact" toolbar={() => (
- -
- - Health Check - + + + + Health Check + +
+ } /> -
- + When enabled, only agents with reachable URLs are shown + +
)} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx index 4ee6332c54e..7858bdb1cd4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx @@ -126,10 +126,7 @@ describe("AgentCardDiscovery", () => { expect(initialSelection.upstream_url).toBe("https://upstream.example.com"); expect(initialSelection.selected_card.skills).toHaveLength(2); - const summarizeLabel = screen.getByText("Summarize").closest("label"); - expect(summarizeLabel).toBeTruthy(); - const summarizeCheckbox = summarizeLabel!.querySelector("input[type='checkbox']") as HTMLInputElement; - await user.click(summarizeCheckbox); + await user.click(screen.getByRole("checkbox", { name: /Summarize/i })); await waitFor(() => { const latest = onApply.mock.calls.at(-1)?.[0]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx index e979b2dbe3f..017a9928f8b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx @@ -1,18 +1,20 @@ "use client"; import React, { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { Alert, Button, Checkbox, Collapse, Empty, Input, Space, Spin, Switch, Tag, Tooltip, Typography } from "antd"; -// Empty is used in the skills panel below. -import { - CheckCircleTwoTone, - InfoCircleOutlined, - LinkOutlined, - ReloadOutlined, - SearchOutlined, -} from "@ant-design/icons"; +import { ChevronDown, CircleAlert, CircleCheck, Info, Link as LinkIcon, RotateCw, Search, X } from "lucide-react"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { DiscoveredAgentCard, discoverAgentCardCall } from "@/components/networking"; +import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Input } from "@/components/ui/input"; +import { Switch } from "@/components/ui/switch"; +import { Textarea } from "@/components/ui/textarea"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { ALLOWED_CAPABILITY_KEYS, selectionsFromSavedAgentCard, @@ -20,9 +22,6 @@ import { skillId, } from "./agent_discovery_utils"; -const { Text, Paragraph } = Typography; -const { Panel } = Collapse; - const DISCOVERY_DEBOUNCE_WAIT_MS = 400; export interface DiscoveredAgentCardSelection { @@ -243,102 +242,115 @@ const AgentCardDiscovery: React.FC = ({ const skillCount = card?.skills?.length ?? 0; const selectedSkillCount = selectedSkillIds.size; + const renderDiscoverIcon = () => { + if (loading) return ; + if (card) return ; + return ; + }; + const discoverLabel = card ? "Re-discover" : "Discover"; + return ( -
-
- - Discover from agent URL - - - +
+
+ + Discover from agent URL + + + + + + } + /> + + LiteLLM will fetch /.well-known/agent-card.json from this URL and let you pick which skills and + capabilities to expose through the proxy. + + +
{isParentDriven ? ( <> - +

Using the connection details you entered above. We'll fetch: - -

+

+
{discoveryRequest!.display_url || effectiveUrl || ( - Fill in the fields above first + Fill in the fields above first )}
-
) : ( <> - +

Paste the upstream agent's base URL. We'll try /.well-known/agent-card.json,{" "} /.well-known/agent.json, and /agent.json in order. - +

- +
setManualUrl(e.target.value)} - onPressEnter={handleDiscover} - allowClear + onKeyDown={(e) => { + if (e.key === "Enter") handleDiscover(); + }} disabled={loading} /> - - +
)} {error && ( - setError(null)} - /> + + + Discovery failed + {error} + + + + )} {loading && !card && (
- +
)} {card && ( -
-
- - - Upstream card loaded - {card.version && v{card.version}} - {card.provider?.organization && {card.provider.organization}} - +
+
+ + Upstream card loaded + {card.version && v{card.version}} + {card.provider?.organization && {card.provider.organization}}
-
+
- + setEditedName(e.target.value)} placeholder="Agent name" />
- - Description +