diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index a4a40a7a3d7..1c433d04819 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -182,7 +182,7 @@ class ManagedObjectCreatorAttribution: team_id: Final = getattr(job, "team_id", None) rows: Final = _CreatorRows( key=await self._live_key_row(api_key, job.created_by), - team=await self._live_team_row(team_id), + team=await self._live_team_row(team_id, job.created_at), user=await self._user_row(object_id, _creator_user_id(job.created_by)), ) return ResponseCreator( @@ -269,10 +269,15 @@ class ManagedObjectCreatorAttribution: verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}") return None - async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: + async def _live_team_row( + self, team_id: str | None, created_before: datetime | None = None + ) -> prisma_models.LiteLLM_TeamTable | None: if not team_id: return None - return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + team: Final = await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + if team is not None and created_before is not None and team.created_at > created_before: + return None + return team def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index baef89912fa..41619f5d1b1 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_a from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError +from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS from litellm.types.llms.bedrock_invoke_agents import ( InvokeAgentChunkPayload, InvokeAgentEvent, @@ -159,7 +160,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): return { "inputText": query, "enableTrace": True, - **optional_params, + **{key: value for key, value in optional_params.items() if key not in AWS_AUTH_PARAM_KEYS}, } def _parse_aws_event_stream(self, raw_content: bytes) -> InvokeAgentEventList: diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index fb727367caf..934180ea964 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -108,7 +108,14 @@ def _bridge_request_data( if key in data and key in agent_kwargs } return _with_headers( - {**agent_kwargs, **data, **layered, **pinned, "model": bridge_model_name(litellm_params)}, + { + **agent_kwargs, + **data, + **layered, + **pinned, + "model": bridge_model_name(litellm_params), + "custom_llm_provider": agent_kwargs.get("custom_llm_provider"), + }, extra_headers, ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b40eb451979..4514882a45f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4969,11 +4969,9 @@ def _owning_team_names_for_team_deployment( if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"): return () deployment: Final = llm_router.registry().deployments(model, source="first_in_group") - public_name: Final = ( - deployment.model_info.team_public_model_name - if deployment is not None and deployment.model_info.team_id == team_id - else None - ) + if deployment is None or deployment.model_info.team_id != team_id: + return () + public_name: Final = deployment.model_info.team_public_model_name legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model) return (*((public_name,) if public_name else ()), *legacy_aliases) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 24602dec677..230dc7f6c58 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2124,7 +2124,7 @@ def _extract_models_from_managed_resource_id( decoded_video: Final = decode_video_id_with_provider(resource_id) _append_model_candidates( candidates=candidates, - value=_video_access_group(decoded_video, llm_router) or decoded_video.get("model_id"), + value=video_access_group(decoded_video, llm_router) or decoded_video.get("model_id"), ) else: model_id = decode_character_id_with_provider(resource_id).get("model_id") @@ -2138,7 +2138,9 @@ def _extract_models_from_managed_resource_id( return _dedupe_model_candidates(candidates) -def _video_access_group(decoded_video: DecodedVideoId, llm_router: Router | None) -> str | None: +def video_access_group( + decoded_video: DecodedVideoId, llm_router: Router | None, *, resolved_model: str | None = None +) -> str | None: """ The group a key must be allowed to call to use a video: the creating deployment's group, or, while router fallbacks skip model access checks, the group whose router fallback created it. @@ -2146,7 +2148,9 @@ def _video_access_group(decoded_video: DecodedVideoId, llm_router: Router | None from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.video_endpoints.utils import video_model_group - creating_group: Final = video_model_group(decoded_video, llm_router) + creating_group: Final = ( + resolved_model if resolved_model is not None else video_model_group(decoded_video, llm_router) + ) requested_model: Final = decoded_video.get("requested_model") if llm_router is None or creating_group is None or requested_model is None: return creating_group diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 0ab3d1ca7af..f99dd4c699b 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -16,6 +16,7 @@ from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.managed_resources.isolation import resolve_resource_owner_id from litellm.proxy._types import LitellmUserRoles from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, @@ -276,7 +277,12 @@ class ResponsesIDSecurity(CustomLogger): ): return True - if response_id_user_id and response_id_user_id != user_api_key_dict.user_id: + if not response_id_user_id and not response_id_team_id: + if general_settings.get("disable_responses_id_security", False): + return True + raise HTTPException(status_code=403, detail="Forbidden. This response id has no associated owner") + + if response_id_user_id and response_id_user_id != resolve_resource_owner_id(user_api_key_dict): if general_settings.get("disable_responses_id_security", False): verbose_proxy_logger.debug( "Responses ID Security is disabled. User %s is accessing response id %s which is not associated with them.", @@ -394,7 +400,7 @@ class ResponsesIDSecurity(CustomLogger): managed_id: Final = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( original_id, - user_api_key_dict.user_id or "", + resolve_resource_owner_id(user_api_key_dict) or "", user_api_key_dict.team_id or "", ) encrypted_id: Final = f"resp_{encrypt_value_helper(value=managed_id)}" diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index 834b032b04f..f13da0bbd1a 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -180,6 +180,9 @@ class TeamMemberPermissionChecks: caller_team_role: Final = get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict) + if not set(access_group_ids).issubset(team_table.access_group_ids or []): + raise HTTPException(status_code=403, detail="Access groups must be assigned to the key's team first") + # Team admins always bypass (consistent with other member-permission checks). if caller_team_role == "admin": return diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8c92a058f43..8367a63f8c5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16908,7 +16908,7 @@ async def _get_caller_byok_team_scope( must have `model_info.team_id` ∈ that set, otherwise they belong to a team the caller is not a member of and must be dropped. """ - if user_api_key_dict is None or prisma_client is None: + if user_api_key_dict is None: return None if user_api_key_dict.user_role in ( LitellmUserRoles.PROXY_ADMIN, @@ -16917,7 +16917,7 @@ async def _get_caller_byok_team_scope( return None key_team_scope: Final[set[str]] = {user_api_key_dict.team_id} if user_api_key_dict.team_id else set() user_id: Final = user_api_key_dict.user_id - if user_id is None: + if user_id is None or prisma_client is None: return key_team_scope try: user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) @@ -17599,7 +17599,7 @@ async def _find_model_by_id( llm_router: Router | None, prisma_client: PrismaClient | None, proxy_config: "ProxyConfig", -) -> tuple[list, int | None]: +) -> tuple[Sequence[dict[str, Any]], int | None]: """Find a model by its ID and optionally filter by search term.""" found_model = None @@ -17796,18 +17796,32 @@ async def model_info_v2( # Load existing config await proxy_config.get_config() + allowed_team_ids: Final = await _get_caller_byok_team_scope(user_api_key_dict, prisma_client) + # If modelId is provided, search for the specific model if modelId is not None: - all_models, search_total_count = await _find_model_by_id( + matched_models: Final = await _find_model_by_id( model_id=modelId, search=search, llm_router=llm_router, prisma_client=prisma_client, proxy_config=proxy_config, ) + all_models, search_total_count = ( + [ + m + for m in matched_models[0] + if not _byok_row_outside_caller_teams(m.get("model_info", {}), allowed_team_ids) + ], + matched_models[1], + ) else: # Normal flow when modelId is not provided - all_models = copy.deepcopy(llm_router.registry().all()) + all_models = [ + m + for m in copy.deepcopy(llm_router.registry().all()) + if not _byok_row_outside_caller_teams(m.get("model_info", {}), allowed_team_ids) + ] if user_model is not None: # if user does not use a config.yaml, https://github.com/BerriAI/litellm/issues/2061 @@ -18599,6 +18613,16 @@ async def model_info_v1( status_code=400, detail={"error": f"Model id = {litellm_model_id} not found on litellm proxy"}, ) + if _byok_row_outside_caller_teams( + deployment_info.model_info.model_dump(), + await _get_caller_byok_team_scope(user_api_key_dict, prisma_client), + ): + if include_team_models: + return _model_info_json_response([]) + raise HTTPException( + status_code=400, + detail={"error": f"Model id = {litellm_model_id} not found on litellm proxy"}, + ) _deployment_info_dict = _get_proxy_model_info(model=deployment_info.model_dump(exclude_none=True)) single_model_list: list[dict] = [_deployment_info_dict] if prisma_client is not None: diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 31d6af71594..4e14eb31762 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -22,13 +22,13 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.image_endpoints.endpoints import batch_to_bytesio from litellm.proxy.video_endpoints.utils import ( VideoRequestRoute, + authorized_video_model, encode_character_id_in_response, encode_video_id_in_response, extract_model_from_target_model_names, fallback_requested_model, get_custom_provider_from_data, routable_model_name, - video_model_group, video_reference_to_id, ) from litellm.types.videos.utils import ( @@ -283,7 +283,7 @@ async def video_status( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model @@ -384,7 +384,7 @@ async def video_content( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing @@ -490,7 +490,7 @@ async def video_remix( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model @@ -814,7 +814,7 @@ async def video_edit( ) data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model @@ -917,7 +917,7 @@ async def video_extension( ) data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 696ad4ab170..24fd443b68a 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -3,6 +3,9 @@ from typing import TYPE_CHECKING, Any, Final import orjson from pydantic import BaseModel, ConfigDict, TypeAdapter +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import video_access_group +from litellm.proxy.openai_files_endpoints.common_utils import authorize_model_for_key from litellm.types.videos.main import DecodedVideoId, VideoObject from litellm.types.videos.utils import ( decode_video_id_with_provider, @@ -122,6 +125,20 @@ def video_model_group(decoded: DecodedVideoId, llm_router: "Router | None") -> s return creating_group +async def authorized_video_model( + decoded: DecodedVideoId, llm_router: "Router | None", user_api_key_dict: UserAPIKeyAuth +) -> str | None: + model: Final = video_model_group(decoded, llm_router) + if model is None: + return None + await authorize_model_for_key( + model_id=video_access_group(decoded, llm_router, resolved_model=model) or model, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) + return model + + def fallback_requested_model(requested_model: str | None, response: object, llm_router: "Router | None") -> str | None: """The group a create asked for when one of that group's router fallbacks served it.""" from litellm.proxy.auth.auth_utils import router_fallback_model_names diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 9da89d96e73..900cfa3d597 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -647,12 +647,21 @@ class ResponsesSessionHandler: SELECT 1 FROM matching_session WHERE matching_session.session_id = logs.session_id - AND matching_session."user" IS NOT DISTINCT FROM logs."user" - AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id + AND ( + matching_session."user" = logs."user" + OR (matching_session."user" IS NULL AND logs."user" IS NULL) + ) + AND ( + matching_session.team_id = logs.team_id + OR (matching_session.team_id IS NULL AND logs.team_id IS NULL) + ) AND ( COALESCE(matching_session."user", '') <> '' OR COALESCE(matching_session.team_id, '') <> '' - OR matching_session.api_key = logs.api_key + OR ( + COALESCE(matching_session.api_key, '') <> '' + AND matching_session.api_key = logs.api_key + ) ) ) ORDER BY "endTime" ASC; diff --git a/litellm/utils.py b/litellm/utils.py index 9baf21befc4..66b5b274f97 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -91,6 +91,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ) from litellm.litellm_core_utils.attempt_settlement import mark_attempt_answered from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, bind_budget_reservation_to_callbacks, has_attributes, iter_items, @@ -5380,18 +5381,30 @@ def pre_process_non_default_params( return non_default_params -def remove_sensitive_keys_from_dict(d: dict) -> dict: +def remove_sensitive_keys_from_dict(d: Mapping[str, object]) -> dict[str, object]: """ Remove sensitive keys from a dictionary """ - sensitive_key_phrases: Final = ["key", "secret", "access", "credential"] - remove_keys: Final = [] - for key in d: - if any(phrase in key.lower() for phrase in sensitive_key_phrases): - remove_keys.append(key) - for key in remove_keys: - d.pop(key) - return d + sensitive_key_phrases: Final = ("key", "secret", "access", "credential") + result: Final = dict[str, object]() + pending: Final = [(d, result)] + while pending: + source, target = pending.pop() + for key, value in source.items(): + if ( + any(phrase in key.lower() for phrase in sensitive_key_phrases) + or key.lower() == "token" + or key.lower().endswith("_token") + ): + continue + mapping = as_str_mapping(value) if isinstance(value, dict) else None + if mapping is not None: + nested = dict[str, object]() + target[key] = nested + pending.append((mapping, nested)) + else: + target[key] = value + return result _TOOL_LIST: Final = TypeAdapter(list[object]) diff --git a/tests/unit/proxy/common_utils/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py index 21685f71923..4404a0d926c 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -67,6 +67,7 @@ class _StoredResponse: id: str unified_object_id: str status: str + created_at: datetime = datetime(2026, 1, 2) model_object_id: str = "" file_purpose: str = "response" file_object: Mapping[str, object] = dataclasses.field(default_factory=lambda: {"model": "gpt-5"}) @@ -2081,7 +2082,14 @@ def _key_row(metadata: Mapping[str, object]) -> object: def _team_row(metadata: Mapping[str, object]) -> object: from types import SimpleNamespace - return SimpleNamespace(team_alias=None, organization_id=None, spend=0.0, max_budget=None, metadata=dict(metadata)) + return SimpleNamespace( + team_alias=None, + organization_id=None, + spend=0.0, + max_budget=None, + metadata=dict(metadata), + created_at=datetime(2026, 1, 1), + ) def _langfuse_callback_settings(callback_vars: Mapping[str, object]) -> Mapping[str, object]: diff --git a/tests/unit/proxy/management_helpers/test_team_member_permission_checks.py b/tests/unit/proxy/management_helpers/test_team_member_permission_checks.py index 55d29724e38..544aa4ed21b 100644 --- a/tests/unit/proxy/management_helpers/test_team_member_permission_checks.py +++ b/tests/unit/proxy/management_helpers/test_team_member_permission_checks.py @@ -271,6 +271,7 @@ class TestEnforceMemberCanAssignAccessGroups: team = MagicMock() team.team_id = team_id team.team_member_permissions = team_member_permissions + team.access_group_ids = ["ag-1"] return team def test_no_access_group_ids_is_noop(self, monkeypatch): diff --git a/tests/unit/proxy/video_endpoints/test_utils.py b/tests/unit/proxy/video_endpoints/test_utils.py index 01a2b7c2840..fff88ab462f 100644 --- a/tests/unit/proxy/video_endpoints/test_utils.py +++ b/tests/unit/proxy/video_endpoints/test_utils.py @@ -20,8 +20,10 @@ import pytest from litellm.router import Router +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.videos.main import DecodedVideoId from litellm.proxy.video_endpoints.utils import ( + authorized_video_model, encode_character_id_in_response, encode_video_id_in_response, extract_model_from_target_model_names, @@ -548,6 +550,33 @@ def _wildcard_router() -> Router: ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("enforce_fallback_model_access", "allowed_model"), + [(False, "sora-pair"), (True, "sora-backup")], +) +async def test_authorized_video_model_keeps_fallback_routing_after_the_key_check( + monkeypatch: pytest.MonkeyPatch, enforce_fallback_model_access: bool, allowed_model: str +) -> None: + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"enforce_fallback_model_access": enforce_fallback_model_access}, + ) + decoded: Final = decode_video_id_with_provider( + encode_video_id_with_provider( + "video_raw", "azure", "sora-2", deployment_id="dep-sora-backup", requested_model="sora-pair" + ) + ) + token: Final = UserAPIKeyAuth(api_key="hashed", models=[allowed_model]) + + assert await authorized_video_model(decoded, _fallback_router(), token) == "sora-backup" + + +@pytest.mark.asyncio +async def test_authorized_video_model_without_a_router_has_no_routing_model() -> None: + assert await authorized_video_model(_decoded("sora-2", "dep-sora-backup"), None, UserAPIKeyAuth()) is None + + @pytest.mark.parametrize( ("requested_model", "answering_deployment", "expected"), [ diff --git a/tests/unit/test_video_generation.py b/tests/unit/test_video_generation.py index 9d9ca3bb734..93243607b6f 100644 --- a/tests/unit/test_video_generation.py +++ b/tests/unit/test_video_generation.py @@ -17,6 +17,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.gemini.videos.transformation import GeminiVideoConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.videos.main import VideoObject, VideoResponse from litellm.videos import main as videos_main from litellm.videos.main import ( @@ -2232,7 +2233,7 @@ def video_proxy_test_client(): app = FastAPI() app.include_router(video_router) - app.dependency_overrides[user_api_key_auth] = lambda: MagicMock() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(api_key="sk-9876") return TestClient(app)