diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..a9fb5c29bad 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -9,6 +9,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args import httpx +from pydantic import TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -218,6 +219,15 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = ) +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuardrailParams]]: + if not isinstance(metadata, dict): + return [] + return metadata.get("guardrails") or [] + + class CustomGuardrail(CustomLogger): # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. use_native_during_call_hook: ClassVar[bool] = False @@ -791,28 +801,22 @@ class CustomGuardrail(CustomLogger): return True raise - def get_guardrail_from_metadata(self, data: dict) -> list[str] | list[dict[str, DynamicGuardrailParams]]: + def get_guardrail_from_metadata(self, data: dict) -> list[str | dict[str, DynamicGuardrailParams]]: """ - Returns the guardrail(s) to be run from the metadata or root + Returns the guardrail(s) to be run from the root, else from both metadata fields together """ if "guardrails" in data: return data["guardrails"] - # Check both metadata locations. For regular endpoints move_guardrails_to_metadata - # writes to "metadata"; for thread/assistant endpoints it writes to - # "litellm_metadata". We check the one that actually contains the "guardrails" - # key so that a non-empty litellm_metadata without guardrails does not shadow - # the merged list stored in metadata (which would cause team guardrails to be - # silently skipped while default_on=True policy guardrails still fire). - for meta_key in ("metadata", "litellm_metadata"): - meta = data.get(meta_key) or {} - if isinstance(meta, dict) and "guardrails" in meta: - return meta.get("guardrails") or [] - return [] + request_fields: Final = _REQUEST_FIELDS.validate_python(data) + return [ + *_guardrails_listed_in(request_fields.get("metadata")), + *_guardrails_listed_in(request_fields.get("litellm_metadata")), + ] def _guardrail_is_in_requested_guardrails( self, - requested_guardrails: list[str] | list[dict[str, DynamicGuardrailParams]], + requested_guardrails: Sequence[str | dict[str, DynamicGuardrailParams]], ) -> bool: for _guardrail in requested_guardrails: if isinstance(_guardrail, dict): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1ef5b1072ed..b8d9530d2eb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -446,7 +446,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request, from litellm.proxy.litellm_pre_call_utils import ( _trace_id_from_traceparent, # pyright: ignore[reportPrivateUsage] # reuse W3C traceparent parser get_chain_id_from_headers, - metadata_variable_name_for_route, + metadata_variable_name_for_request, ) if get_chain_id_from_headers(dict(request.headers)): @@ -460,7 +460,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request, traceparent: Final[str | None] = request.headers.get("traceparent") if isinstance(traceparent, str) and _trace_id_from_traceparent(traceparent): return True - metadata_variable_name: Final = metadata_variable_name_for_route(get_request_route(request)) + metadata_variable_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body) selected_metadata: Final[object] = request_body.get(metadata_variable_name) if isinstance(selected_metadata, dict): selected_metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(selected_metadata) @@ -1367,11 +1367,15 @@ async def common_checks( if not skip_all_budget_checks: # Key metadata.tags are injected into request_body here so the tag budget # check can read them; this mutation must run before the gathered checks. - from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request if valid_token is not None: for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(request, request_body): LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=route) + LiteLLMProxyRequestSetup.merge_caller_body_tags( + request_data=tag_target, + metadata_variable_name=metadata_variable_name_for_request(route, tag_target), + ) LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(request_data=tag_target, user_api_key_dict=valid_token) LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_body) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7cd752c2297..a79d8fe29db 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2401,9 +2401,10 @@ class ProxyBaseLLMRequestProcessing: version=version, proxy_config=proxy_config, ) + request_fields: Final = _REQUEST_FIELDS.validate_python(self.data) + metadata_variable_name: Final = get_metadata_variable_name(request, request_fields) if logging_auth is not None: - metadata_variable_name: Final = get_metadata_variable_name(request) - request_metadata: Final = self.data.get(metadata_variable_name) + request_metadata: Final = request_fields.get(metadata_variable_name) if isinstance(request_metadata, dict): self.data[metadata_variable_name] = { **TypeAdapter(dict[str, object]).validate_python(request_metadata), @@ -2431,12 +2432,11 @@ class ProxyBaseLLMRequestProcessing: # Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved if queue_time_seconds is not None: - _metadata_variable_name: Final = get_metadata_variable_name(request) - if _metadata_variable_name not in self.data: - self.data[_metadata_variable_name] = {} - if not isinstance(self.data[_metadata_variable_name], dict): - self.data[_metadata_variable_name] = {} - self.data[_metadata_variable_name]["queue_time_seconds"] = queue_time_seconds + if metadata_variable_name not in self.data: + self.data[metadata_variable_name] = {} + if not isinstance(self.data[metadata_variable_name], dict): + self.data[metadata_variable_name] = {} + self.data[metadata_variable_name]["queue_time_seconds"] = queue_time_seconds if isinstance(model, str): reject_url_valued_destination("model", model) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 14b08a09116..15282679622 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -3312,7 +3312,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if key_agent_id: return key_agent_id metadata: Final = data.get("metadata") or {} - return metadata.get("agent_id") + litellm_metadata: Final = data.get("litellm_metadata") or {} + return metadata.get("agent_id") or litellm_metadata.get("agent_id") def _get_session_id_from_data(self, data: dict) -> str | None: """Extract session_id from request metadata or litellm_session_id.""" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 73eaa299530..40250b47c90 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -40,7 +40,7 @@ from litellm.proxy.db.db_spend_update_writer import ( get_llm_router, ) from litellm.proxy.hooks.search_free_tier import apply_search_free_tier -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope from litellm.proxy.spend_tracking.spend_event import ( @@ -140,19 +140,34 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( _OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) -def _proxy_stamped_used_client_oauth_token( - request_data: Mapping[str, object], request_route: str | None -) -> bool | None: +def _proxy_stamped_metadata(request_data: Mapping[str, object], request_route: str | None) -> Mapping[str, object]: proxy_bucket: Final = ( get_metadata_variable_name_from_kwargs(request_data) if request_route is None - else metadata_variable_name_for_route(request_route) + else metadata_variable_name_for_request(request_route, request_data) ) - proxy_metadata: Final = request_data.get(proxy_bucket) - stamped: Final = proxy_metadata.get("used_client_oauth_token") if isinstance(proxy_metadata, dict) else None + return _str_keyed_fields(request_data.get(proxy_bucket)) + + +def _proxy_stamped_used_client_oauth_token( + request_data: Mapping[str, object], request_route: str | None +) -> bool | None: + stamped: Final = _proxy_stamped_metadata(request_data, request_route).get("used_client_oauth_token") return stamped if isinstance(stamped, bool) else None +_PROXY_STAMPED_REQUEST_FIELDS: Final = ("requester_ip_address", "user_agent") + + +def _proxy_stamped_request_fields( + request_data: Mapping[str, object], request_route: str | None +) -> Mapping[str, object]: + if request_route is None: + return {} + proxy_metadata: Final = _proxy_stamped_metadata(request_data, request_route) + return {field: proxy_metadata[field] for field in _PROXY_STAMPED_REQUEST_FIELDS if field in proxy_metadata} + + def _str_keyed_fields(value: object) -> Mapping[str, object]: return value if is_str_keyed_mapping(value) else {} @@ -461,7 +476,13 @@ class _ProxyDBLogger(CustomLogger): _metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_fields, request_route) existing_metadata: Final[dict] = request_data.get("metadata", None) or {} - existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata}) + existing_metadata.update( + { + **_routed_metadata_fallback(request_fields), + **_proxy_stamped_request_fields(request_fields, request_route), + **_metadata, + } + ) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") litellm_metadata: Final[ObjectMapping] = ( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 33fb9766821..269c4ee4450 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -379,6 +379,8 @@ LITELLM_METADATA_ROUTES: Final = ( "interactions", ) +CURSOR_CHAT_COMPLETIONS_ROUTE: Final = "/cursor/chat/completions" + LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( { "mask_input", @@ -636,7 +638,7 @@ def _has_tag_budget(tag: "LiteLLM_TagTable | None") -> bool: return tag.litellm_budget_table.max_budget is not None -def _normalized_metadata_slot( +def normalized_metadata_slot( request_data: MutableMapping[str, object], metadata_variable_name: str ) -> dict[str, object]: """Return the request's metadata slot as a dict, normalising it in place first. @@ -653,6 +655,12 @@ def _normalized_metadata_slot( return normalized +def _holds_metadata_object(value: object) -> bool: + if value is None or isinstance(value, dict): + return True + return isinstance(value, str) and isinstance(_DECODED_JSON.validate_python(safe_json_loads(value)), dict) + + def _metadata_object_or_error(field: str, value: object) -> Mapping[str, object] | ProxyException: if value is None: return {} @@ -863,7 +871,7 @@ def _strip_router_reserved_metadata( ) -def get_metadata_variable_name(request: Request) -> str: +def get_metadata_variable_name(request: Request, data: Mapping[str, object] | None = None) -> str: """ Helper to return what the "metadata" field should be called in the request data @@ -874,7 +882,10 @@ def get_metadata_variable_name(request: Request) -> str: # Inline imports — auth_utils/route_checks participate in a proxy import cycle. from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 - return metadata_variable_name_for_route(get_request_route(request)) + route: Final = get_request_route(request) + if data is None: + return metadata_variable_name_for_route(route) + return metadata_variable_name_for_request(route, data) _get_metadata_variable_name: Final = get_metadata_variable_name @@ -890,6 +901,25 @@ def metadata_variable_name_for_route(route: str) -> Literal["metadata", "litellm return "metadata" +def is_chat_completions_body(data: Mapping[str, object]) -> bool: + messages: Final = data.get("messages") + if isinstance(messages, list) and messages: + return True + return "messages" in data and "input" not in data + + +def is_cursor_responses_request(route: str, body: Mapping[str, object]) -> bool: + return CURSOR_CHAT_COMPLETIONS_ROUTE in route and not is_chat_completions_body(body) + + +def metadata_variable_name_for_request( + route: str, body: Mapping[str, object] +) -> Literal["metadata", "litellm_metadata"]: + if is_cursor_responses_request(route, body): + return "litellm_metadata" + return metadata_variable_name_for_route(route) + + def _promoted_trace_control_fields( requester_metadata: Mapping[str, object], litellm_metadata: Mapping[str, object], @@ -2309,7 +2339,9 @@ class LiteLLMProxyRequestSetup: ``_tag_max_budget_check`` resolve to ``litellm_metadata``, leaving header tags invisible to per-tag budget enforcement. """ - if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES): + if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES) or is_cursor_responses_request( + route, request_data + ): request_data.setdefault("litellm_metadata", {}) @staticmethod @@ -2380,7 +2412,7 @@ class LiteLLMProxyRequestSetup: return tags: Final = _STRING_TAGS_ADAPTER.validate_python(raw_tags) slot: Final = metadata_key if metadata_key is not None else get_metadata_variable_name_from_kwargs(request_data) - metadata: Final = _normalized_metadata_slot(request_data, slot) + metadata: Final = normalized_metadata_slot(request_data, slot) existing: Final = metadata.get("tags") if existing and not isinstance(existing, list): return @@ -2393,6 +2425,34 @@ class LiteLLMProxyRequestSetup: tags=tuple(tag for tag in tags if tag not in sent_tags) ) + @staticmethod + def merge_caller_body_tags( + request_data: MutableMapping[str, object], # mutable-ok: the request body the tag merges write into + metadata_variable_name: str | None = None, + ) -> None: + slot: Final = ( + metadata_variable_name + if metadata_variable_name is not None + else get_metadata_variable_name_from_kwargs(request_data) + ) + if slot == "metadata" or not _holds_metadata_object(request_data.get(slot)): + return + raw_caller_metadata: Final = request_data.get("metadata") + caller_metadata: Final = ( + _DECODED_JSON.validate_python(safe_json_loads(raw_caller_metadata)) + if isinstance(raw_caller_metadata, str) + else raw_caller_metadata + ) + caller_tags: Final = caller_metadata.get("tags") if isinstance(caller_metadata, dict) else None + if not isinstance(caller_tags, list) or not caller_tags: + return + metadata: Final = normalized_metadata_slot(request_data, slot) + existing_tags: Final = metadata.get("tags") + metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags( + request_tags=existing_tags if isinstance(existing_tags, list) else None, + tags_to_add=caller_tags, + ) + @staticmethod def apply_key_tags_pre_auth( request_data: dict, @@ -2408,7 +2468,7 @@ class LiteLLMProxyRequestSetup: return _metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data) - metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name) + metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name) existing_tags: Final = metadata.get("tags") metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags( @@ -2460,7 +2520,7 @@ class LiteLLMProxyRequestSetup: # from (litellm_metadata vs metadata) so the merged tags are visible # to _tag_max_budget_check. _metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data) - metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name) + metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name) existing_tags: Final = metadata.get("tags") metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags( @@ -2637,7 +2697,7 @@ async def add_litellm_data_to_request( } safe_add_api_version_from_query_params(data, request) - _metadata_variable_name: Final = get_metadata_variable_name(request) + _metadata_variable_name: Final = get_metadata_variable_name(request, _REQUEST_FIELDS.validate_python(data)) if data.get(_metadata_variable_name, None) is None: data[_metadata_variable_name] = {} @@ -2974,17 +3034,7 @@ async def add_litellm_data_to_request( tags_to_add=tags, ) - _caller_body_metadata: Final = data.get("metadata") if _metadata_variable_name != "metadata" else None - _caller_body_tags: Final = ( - _caller_body_metadata.get("tags") - if isinstance(_caller_body_metadata, dict) and isinstance(_caller_body_metadata.get("tags"), list) - else None - ) - if _caller_body_tags: - data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup.merge_tags( # rebind-ok: matches file idiom - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=_caller_body_tags, - ) + LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data, metadata_variable_name=_metadata_variable_name) # caller_tags: the tags the request's headers and body carry (x-litellm-tags, body "tags", and body # "metadata.tags" / "litellm_metadata.tags" read before the mapped tags are added below), diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 9d0d89f681e..ef59deb2e5e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -79,6 +79,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry from litellm.proxy.litellm_pre_call_utils import ( _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + is_chat_completions_body, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_request_data, @@ -213,13 +214,6 @@ def _normalize_tool_dialect( return {**data, **{key: value for key, value in replaceable if key in data}} -def _is_chat_completions_body(data: Mapping[str, object]) -> bool: - messages: Final = data.get("messages") - if isinstance(messages, list) and messages: - return True - return "messages" in data and "input" not in data - - _CURSOR_THINKING_SEPARATOR: Final = "-thinking-" _CURSOR_FAST_SUFFIX: Final = "-fast" _CURSOR_THINKING_LEVELS: Final[frozenset[str]] = frozenset(get_args(REASONING_EFFORT)) @@ -261,7 +255,7 @@ def _resolve_cursor_model_variant( resolved: Final = {**data, "model": variant.base_model} if variant.reasoning_effort is None: return resolved - if _is_chat_completions_body(data): + if is_chat_completions_body(data): if "reasoning_effort" in data: return resolved return {**resolved, "reasoning_effort": variant.reasoning_effort} @@ -646,7 +640,7 @@ async def cursor_chat_completions( raw_body: Final = await read_request_body(request=request) - if _is_chat_completions_body(raw_body): + if is_chat_completions_body(raw_body): # Genuine chat completions body (Cursor sends these for models whose BYOK it # already fixed); delegate so behavior matches /chat/completions exactly. # Keyed on messages CONTENT, not key presence: Cursor can send a null or diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 99f16f0f6ce..1fed513f0a7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -255,6 +255,7 @@ from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata, key_or_team_allows_client_message_redaction_opt_out, + normalized_metadata_slot, ) from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import ( @@ -3986,9 +3987,7 @@ class ProxyLogging: ) request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP}) - if "metadata" not in request_data: - request_data["metadata"] = {} - request_data["metadata"].update(user_api_key_logged_metadata) + normalized_metadata_slot(request_data, "metadata").update(user_api_key_logged_metadata) if litellm_logging_obj is not None: ## UPDATE LOGGING INPUT diff --git a/tests/integration/providers/test_cursor_responses_wire.py b/tests/integration/providers/test_cursor_responses_wire.py index d77d35f2ff8..1bd79b5b575 100644 --- a/tests/integration/providers/test_cursor_responses_wire.py +++ b/tests/integration/providers/test_cursor_responses_wire.py @@ -3,11 +3,14 @@ from __future__ import annotations import json import os import uuid +from collections.abc import Mapping +from hashlib import sha256 from importlib.resources import files from typing import Final +import httpx import pytest -from integration._support.client import Gateway, Scenario +from integration._support.client import Gateway, Scenario, eventually, list_value, object_value from integration._support.responses_vendor import same_response from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.worker_readiness import wait_for_worker_readiness @@ -38,6 +41,7 @@ _CHAT_SUFFIXES: Final[tuple[tuple[str, str], ...]] = (("", ""),) + tuple( (f"-thinking-{level}", level) for level in _CHAT_THINKING_LEVELS ) _PROMPT: Final = "Inspect the requested change." +_REQUEST_TAG: Final = "cursor-integration" _INPUT: Final = [{"role": "user", "content": [{"type": "input_text", "text": _PROMPT}]}] _CUSTOM_TOOL: Final = { "type": "custom", @@ -107,16 +111,30 @@ def _stream(identity: str) -> tuple[bytes, ...]: ) + (b"data: [DONE]\n\n",) -def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]: +def _chat_completion(stream: bool, identity: str) -> bytes | tuple[bytes, ...]: if stream: - return ( - b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{"content":"Patch inspected."},"finish_reason":null}]}\n\n', - b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n', - b"data: [DONE]\n\n", - ) + return tuple( + f"data: {json.dumps(chunk)}\n\n".encode() + for chunk in ( + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _MODEL, + "choices": [{"index": 0, "delta": {"content": "Patch inspected."}, "finish_reason": None}], + }, + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _MODEL, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + ) + ) + (b"data: [DONE]\n\n",) return json.dumps( { - "id": "chat_cursor_1", + "id": identity, "object": "chat.completion", "created": 1, "model": _MODEL, @@ -133,11 +151,12 @@ def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]: ).encode() -def _deployment(scenario: Scenario, wire: Wire) -> str: +def _deployment(scenario: Scenario, wire: Wire, **pricing: JsonValue) -> str: model: Final = scenario.model( model=f"openai/{_MODEL}", api_key=_API_KEY, api_base=f"{wire.url}/v1", + **pricing, ) wait_for_worker_readiness(scenario.gateway, model, 1, _PROXY_WORKERS) return model @@ -167,7 +186,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output( assert request.method == "POST" assert request.headers["authorization"] == f"Bearer {_API_KEY}", dict(request.headers) if request.target == "/v1/chat/completions": - completion: Final = _chat_completion(stream) + completion: Final = _chat_completion(stream, raw_id) return Reply( content_type="text/event-stream" if stream else "application/json", chunks=completion if isinstance(completion, tuple) else None, @@ -175,7 +194,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output( ) assert request.target == "/v1/responses", request.target body: Final = _JSON_OBJECT.validate_json(request.body) - assert {key: value for key, value in body.items() if key != "metadata"} == { + assert body == { "model": _MODEL, "input": _INPUT, **reasoning, @@ -259,41 +278,69 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output( assert [request.target for request in drain_contract_requests(wire)] == ["/v1/responses"] -def test_cursor_responses_body_carries_no_proxy_metadata(gateway: Gateway) -> None: - pytest.skip( - "BUG: /cursor/chat/completions forwards proxy-internal metadata (key hash, team, user, client ip, user agent) " - "to the provider in the Responses body" - ) +@pytest.mark.parametrize( + "client_metadata", ({}, {"cursor_conversation": "integration"}), ids=("no_client_metadata", "client_metadata") +) +@pytest.mark.parametrize("stream", (False, True)) +def test_cursor_responses_body_sends_only_client_metadata_and_bills_the_key( + gateway: Gateway, + stream: bool, + client_metadata: Mapping[str, JsonValue], +) -> None: raw_id: Final = f"resp_cursor_metadata_{uuid.uuid4().hex}" + alias: Final = f"integration-cursor-{uuid.uuid4().hex}" + sent_metadata: Final = {"metadata": client_metadata} if client_metadata else {} def respond(request: Request) -> Reply: if request.method == "GET" and request.target == "/v1/models": return model_discovery_reply(_MODEL) assert request.method == "POST" and request.target == "/v1/responses", request.target + if stream: + return Reply(content_type="text/event-stream", chunks=_stream(raw_id)) return Reply(body=json.dumps(_response(raw_id)).encode()) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = _deployment(scenario, wire) + key: Final = scenario.key(key_alias=alias) + hashed_key: Final = sha256(key.encode()).hexdigest() response: Final = gateway.request( "POST", "/cursor/chat/completions", { - "model": f"{model}-thinking-high", + "model": model, "input": _INPUT, "messages": [], "tools": [_CUSTOM_TOOL], + "stream": stream, + **sent_metadata, }, + key=key, ) assert response.status_code == 200, response.text requests: Final = drain_contract_requests(wire) assert len(requests) == 1, requests - body: Final = _JSON_OBJECT.validate_json(requests[0].body) - assert body == { + assert _JSON_OBJECT.validate_json(requests[0].body) == { "model": _MODEL, "input": _INPUT, - "reasoning": {"effort": "high"}, "tools": [_CUSTOM_TOOL], + "stream": stream, + **sent_metadata, }, requests[0].body + (row,) = eventually( + lambda: gateway.rows( + 'SELECT call_type, spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (hashed_key,) + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert row["call_type"] == "aresponses", row + assert object_value(row["metadata"])["user_api_key_alias"] == alias, row + assert isinstance(row["spend"], float) and row["spend"] > 0, row + eventually( + lambda: gateway.rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed_key,)), + lambda tokens: len(tokens) == 1 and tokens[0]["spend"] == pytest.approx(row["spend"]), + seconds=70, + ) def test_cursor_stream_reports_usage_when_the_client_asks_for_it(gateway: Gateway) -> None: @@ -341,10 +388,12 @@ def test_cursor_messages_body_uses_chat_completions_wire( suffix: str, reasoning_effort: str, ) -> None: + chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}" + def respond(request: Request) -> Reply: if request.method == "GET" and request.target == "/v1/models": return model_discovery_reply(_MODEL) - completion: Final = _chat_completion(stream) + completion: Final = _chat_completion(stream, chat_id) return Reply( content_type="text/event-stream" if stream else "application/json", chunks=completion if isinstance(completion, tuple) else None, @@ -367,6 +416,7 @@ def test_cursor_messages_body_uses_chat_completions_wire( "stream": stream, }, key=key, + headers={"x-litellm-tags": _REQUEST_TAG}, ) assert response.status_code == 200, response.text if stream: @@ -393,7 +443,7 @@ def test_cursor_messages_body_uses_chat_completions_wire( payload: Final = _JSON_OBJECT.validate_json(response.content) identity: Final = payload.get("id") created: Final = payload.get("created") - assert isinstance(identity, str) and same_response(identity, "chat_cursor_1"), response.text + assert isinstance(identity, str) and same_response(identity, chat_id), response.text assert isinstance(created, int) and created > 0, response.text assert payload == { "id": identity, @@ -456,3 +506,72 @@ def test_cursor_messages_body_uses_chat_completions_wire( **({"stream_options": {"include_usage": True}} if stream else {}), **({"reasoning_effort": reasoning_effort} if reasoning_effort else {}), }, requests[0].body + (row,) = eventually( + lambda: gateway.rows( + 'SELECT call_type, team_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert (row["call_type"], row["team_id"]) == ("acompletion", team), row + assert _REQUEST_TAG in list_value(row["request_tags"]), row + + +@pytest.mark.parametrize( + ("body_shape", "metadata_as_json_string"), (("messages", False), ("input", False), ("input", True)) +) +def test_cursor_body_keeps_client_and_key_tags_and_enforces_the_client_tag_budget( + gateway: Gateway, body_shape: str, metadata_as_json_string: bool +) -> None: + chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}" + response_id: Final = f"resp_cursor_tags_{uuid.uuid4().hex}" + budget_tag: Final = f"cursor-tag-budget-{uuid.uuid4().hex}" + key_tag: Final = f"cursor-key-tag-{uuid.uuid4().hex}" + client_metadata: Final = {"tags": [budget_tag]} + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.target == "/v1/responses": + return Reply(body=json.dumps(_response(response_id)).encode()) + completion: Final = _chat_completion(False, chat_id) + return Reply(body=completion if isinstance(completion, bytes) else b"") + + def delete_tag() -> None: + gateway.post("/tag/delete", {"name": budget_tag}) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, input_cost_per_token=0.01, output_cost_per_token=0.01) + key: Final = scenario.key(models=[model], metadata={"tags": [key_tag]}) + gateway.post("/tag/new", {"name": budget_tag, "max_budget": 0.0001}) + scenario.cleanups.callback(delete_tag) + + def tagged_chat(text: str) -> httpx.Response: + return gateway.request( + "POST", + "/cursor/chat/completions", + { + "model": model, + body_shape: [{"role": "user", "content": text}], + "metadata": json.dumps(client_metadata) if metadata_as_json_string else client_metadata, + }, + key=key, + ) + + first: Final = tagged_chat("first tagged chat") + assert first.status_code == 200, first.text + (row,) = eventually( + lambda: gateway.rows( + 'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert {budget_tag, key_tag} <= set(list_value(row["request_tags"])), row + blocked: Final = eventually( + lambda: tagged_chat("tag budget probe"), lambda response: response.status_code != 200, seconds=70 + ) + assert blocked.status_code == 422, blocked.text + assert "budget" in blocked.text.lower(), blocked.text diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 614ff97553a..0483888533e 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3699,6 +3699,20 @@ def test_get_guardrail_from_metadata() -> None: ) == [{"test-guardrail": {"extra_body": {"key": "value"}}}] + +@pytest.mark.parametrize( + "caller_metadata", + [{"guardrails": []}, {"guardrails": ["not-a-guardrail"]}, {"guardrails": None}], +) +def test_caller_metadata_guardrails_never_hide_the_proxy_guardrails_in_litellm_metadata(caller_metadata: dict) -> None: + guardrail: Final = CustomGuardrail( + guardrail_name="key-guardrail", event_hook=GuardrailEventHooks.pre_call, default_on=False + ) + request_data: Final = {"metadata": caller_metadata, "litellm_metadata": {"guardrails": ["key-guardrail"]}} + + assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is True + + def test_guardrail_is_in_requested_guardrails() -> None: guardrail: Final = CustomGuardrail(guardrail_name="test-guardrail") diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index b66e7f3d091..3aad2bda9bf 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -2591,7 +2591,8 @@ async def test_agent_session_rate_limit_skipped_without_session_id(): @pytest.mark.asyncio -async def test_agent_rate_limit_from_metadata_agent_id(): +@pytest.mark.parametrize("metadata_field", ("metadata", "litellm_metadata")) +async def test_agent_rate_limit_from_metadata_agent_id(metadata_field): """ Test that agent rate limits work when agent_id comes from request metadata (header) rather than from the API key. @@ -2639,7 +2640,7 @@ async def test_agent_rate_limit_from_metadata_agent_id(): cache=local_cache, data={ "model": "gpt-4", - "metadata": {"agent_id": _agent_id}, + metadata_field: {"agent_id": _agent_id}, }, call_type="", ) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index b578a3e96b9..85f553db6ea 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -623,6 +623,56 @@ async def test_async_post_call_failure_hook_reads_used_client_oauth_token_from_t assert json.loads(payload["metadata"])["used_client_oauth_token"] is expected + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_route, body", + [ + ( + "/cursor/chat/completions", + { + "input": "Hello", + "metadata": {"requester_ip_address": "203.0.113.9", "user_agent": "spoofed"}, + "litellm_metadata": {"requester_ip_address": "127.0.0.1", "user_agent": "Cursor/1.0"}, + }, + ), + ( + "/v1/chat/completions", + { + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"requester_ip_address": "127.0.0.1", "user_agent": "Cursor/1.0"}, + "litellm_metadata": {"requester_ip_address": "203.0.113.9", "user_agent": "spoofed"}, + }, + ), + ], +) +async def test_async_post_call_failure_hook_logs_the_proxy_stamped_ip_and_user_agent(request_route: str, body: dict): + request_data = { + "model": "gpt-4o", + "proxy_server_request": {"request_id": "test_request_id"}, + **body, + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await ProxyDBLogger().async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("invalid api key"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key", request_route=request_route), + ) + + payload = get_logging_payload( + kwargs=mock_update_database.call_args[1]["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert payload["requester_ip_address"] == "127.0.0.1" + assert json.loads(payload["metadata"])["user_agent"] == "Cursor/1.0" + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_bills_guardrail_cost_on_blocked_request(): """LIT-5651: a request blocked by a guardrail never reaches the LLM, but the diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index a783ac70881..1e4dded801e 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3938,7 +3938,7 @@ class TestCursorInputArmFlattening: class TestChatCompletionsBodyDetection: def test_routing_matrix(self): - from litellm.proxy.response_api_endpoints.endpoints import _is_chat_completions_body + from litellm.proxy.litellm_pre_call_utils import is_chat_completions_body as _is_chat_completions_body assert _is_chat_completions_body({"messages": [{"role": "user", "content": "hi"}]}) is True assert _is_chat_completions_body({"messages": [{"role": "user", "content": "hi"}], "input": []}) is True diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index c867db2fd02..2d08609c93b 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9711,6 +9711,40 @@ async def test_saved_weights_override_caller_input_and_preserve_key_precedence( assert logging.pre_call_hook.call_args.kwargs["data"].get("_router_weights") == expected +def _hook_data(**kwargs: object) -> object: + return kwargs["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prompt, proxy_field, client_field", [ + ({"input": "hi"}, "litellm_metadata", "metadata"), + ({"messages": [{"role": "user", "content": "hi"}]}, "metadata", "litellm_metadata"), +]) +async def test_cursor_route_writes_proxy_fields_only_to_the_field_its_body_shape_picks( + prompt: dict[str, str | list[dict[str, str]]], + proxy_field: str, + client_field: str, +) -> None: + from litellm.proxy import proxy_server + + logging: Final = MagicMock(spec=ProxyLogging) + logging.pre_call_hook = AsyncMock(side_effect=_hook_data) + processor: Final = ProxyBaseLLMRequestProcessing(data={"model": "gpt-x", "metadata": {"client": "cursor"}, **prompt}) + data, _ = await processor.common_processing_pre_call_logic( + request=Request({"type": "http", "method": "POST", "path": "/cursor/chat/completions", "headers": []}), + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash"), + proxy_logging_obj=logging, + proxy_config=proxy_server.ProxyConfig(), + route_type="acompletion", + logging_auth=ProxyUserAPIKeyAuth(api_key="logging-hash", key_alias="logging-alias"), + ) + assert data[proxy_field]["user_api_key_alias"] == "logging-alias" + assert "queue_time_seconds" in data[proxy_field] + assert "user_api_key_alias" not in data.get(client_field, {}) + assert "queue_time_seconds" not in data.get(client_field, {}) + assert data["metadata"]["client"] == "cursor" + class TestPerRequestModelGroupAlias: """``router_settings.model_group_alias`` on a key or team has to be resolved by the proxy: the Router resolves aliases from its own shared instance diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 1ce61a8e9b3..2edb032b048 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -121,6 +121,23 @@ class TestGetMetadataVariableName: request = self._make_request("/chat/completions") assert get_metadata_variable_name(request) == "metadata" + @pytest.mark.parametrize( + "body, expected", + [ + ({"messages": [{"role": "user", "content": "hi"}]}, "metadata"), + ({"messages": []}, "metadata"), + ({"input": [{"role": "user", "content": "hi"}]}, "litellm_metadata"), + ({"input": "hi", "messages": []}, "litellm_metadata"), + ], + ) + def test_cursor_route_picks_the_metadata_field_by_body_shape(self, body, expected): + request = self._make_request("/cursor/chat/completions") + assert get_metadata_variable_name(request, body) == expected + + def test_cursor_route_without_a_body_keeps_the_route_answer(self): + request = self._make_request("/cursor/chat/completions") + assert get_metadata_variable_name(request) == "metadata" + def test_returns_metadata_for_completions(self): request = self._make_request("/v1/completions") assert get_metadata_variable_name(request) == "metadata" @@ -681,6 +698,45 @@ async def test_add_litellm_data_to_request_rejects_non_object_metadata(field, va assert field not in data +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/responses", "/cursor/chat/completions"]) +@pytest.mark.parametrize("value, received_type", [("junk", "a string"), (5, "an integer"), ([1], "an array")]) +async def test_invalid_litellm_metadata_answers_400_when_metadata_carries_tags(route, value, received_type): + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.utils import ProxyLogging + + data = {"model": "gpt-4o", "input": "hi", "metadata": {"tags": ["client-tag"]}, "litellm_metadata": value} + request = _batches_request_mock() + await common_checks( + request_body=data, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(api_key="hashed-key"), + request=request, + ) + + with pytest.raises(ProxyException) as exc_info: + await add_litellm_data_to_request( + data=data, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert exc_info.value.code == "400" + assert exc_info.value.message == ( + f"Invalid type for 'litellm_metadata': expected an object, but got {received_type} instead." + ) + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_removes_every_invalid_metadata_field_before_raising(): """When both fields are invalid, the raise for the first must not leave the @@ -11098,3 +11154,15 @@ def test_mapped_baggage_tags_keep_non_string_client_tags_as_sent(): data: Final[dict[str, object]] = {"metadata": {"tags": [1, {"team": "x"}, "client-tag", "client-tag"]}} LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data) assert data["metadata"]["tags"] == [1, {"team": "x"}, "client-tag", "client-tag", "tenant:a"] + + +@pytest.mark.parametrize("caller_metadata", ({"tags": ["client-tag"]}, json.dumps({"tags": ["client-tag"]}))) +def test_caller_tags_reach_litellm_metadata_whether_metadata_is_a_dict_or_a_json_string(caller_metadata): + data: Final[dict[str, object]] = { + "input": "hi", + "metadata": caller_metadata, + "litellm_metadata": {"tags": ["key-tag"]}, + } + LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data) + assert set(data["litellm_metadata"]["tags"]) == {"key-tag", "client-tag"} + assert data["metadata"] == caller_metadata diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index f880d3548f7..4386d5aec94 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -4,6 +4,7 @@ and ``_handle_logging_proxy_only_error``.""" from __future__ import annotations import asyncio +import json from collections.abc import Callable from datetime import datetime from types import MappingProxyType @@ -1192,3 +1193,31 @@ async def test_handle_logging_proxy_only_path_flags_requests_that_died_before_re blocked_by_guardrail.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP, blocked_by_guardrail["litellm_call_id"], ) == (True, "client-supplied", False, "call-1") + + +@pytest.mark.asyncio +async def test_auth_failure_on_a_responses_body_with_a_json_string_metadata_is_logged_against_the_key( + proxy_logging, monkeypatch +): + from litellm.proxy import proxy_server + + recorded: list[dict] = [] + + class _RecordingLogger(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + recorded.append(kwargs) + + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [_RecordingLogger()]) + proxy_logging.alert_types = [] + + await proxy_logging.post_call_failure_hook( + request_data={"model": "internal-model", "input": "hi", "metadata": json.dumps({"tags": ["checkout"]})}, + original_exception=litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-hashed", key_alias="checkout-key", request_route="/v1/responses"), + error_type=ProxyErrorTypes.auth_error, + route="/v1/responses", + ) + + assert len(recorded) == 1 + assert recorded[0]["standard_logging_object"]["metadata"]["user_api_key_alias"] == "checkout-key"