diff --git a/enterprise/litellm_enterprise/proxy/hooks/azure_passthrough_permissions.py b/enterprise/litellm_enterprise/proxy/hooks/azure_passthrough_permissions.py new file mode 100644 index 00000000000..2de70667e50 --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/hooks/azure_passthrough_permissions.py @@ -0,0 +1,411 @@ +""" +Permission management for Azure OpenAI pass-through routes. + +Intercepts /batches, /files, and /responses operations on Azure pass-through +to enforce per-user ownership: + +- Create: After Azure returns, store ownership in ManagedObjectTable/ManagedFileTable +- Read/Delete/Cancel: Check ownership before forwarding, rewrite managed IDs to real IDs +- List: Filter results by user_id from the managed tables +""" + +import json +import re +from dataclasses import dataclass +from typing import Any, Dict, Literal, Optional, Tuple + +from fastapi import HTTPException, Response + +from litellm import verbose_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + get_original_file_id, + is_model_embedded_id, +) + +# Regex to match Azure OpenAI deployment-based resource endpoints +# Captures: resource_type (batches|files|responses), resource_id (optional), sub_action (cancel|content, optional) +_AZURE_RESOURCE_PATTERN = re.compile( + r"openai/deployments/[^/]+/(batches|files|responses)(?:/([^/?]+))?(?:/(cancel|content))?$" +) + + +ResourceType = Literal["batch", "file", "response"] +OperationType = Literal["create", "retrieve", "delete", "list", "cancel", "content"] + + +@dataclass +class AzurePassthroughOp: + """Classified Azure pass-through operation.""" + + resource_type: ResourceType + operation: OperationType + resource_id: Optional[str] # Raw ID from URL path, if present + + +def classify_azure_passthrough_request( + endpoint: str, method: str +) -> Optional[AzurePassthroughOp]: + """ + Classify an Azure pass-through request by resource type and operation. + + Args: + endpoint: The path portion after /azure/ (e.g. "openai/deployments/gpt-4/batches/batch_abc123") + method: HTTP method (GET, POST, DELETE, etc.) + + Returns: + AzurePassthroughOp if matched, None otherwise + """ + match = _AZURE_RESOURCE_PATTERN.search(endpoint) + if not match: + return None + + resource_type_raw = match.group(1) # batches, files, responses + resource_id = match.group(2) # e.g. batch_abc123, file-xyz, resp_123 + sub_action = match.group(3) # cancel, content + + # Map plural URL segment to singular resource type + resource_type_map: Dict[str, ResourceType] = { + "batches": "batch", + "files": "file", + "responses": "response", + } + resource_type = resource_type_map[resource_type_raw] + + method_upper = method.upper() + + # Determine operation + if sub_action == "cancel": + operation: OperationType = "cancel" + elif sub_action == "content": + operation = "content" + elif resource_id is not None: + # Has resource ID in path + if method_upper == "GET": + operation = "retrieve" + elif method_upper == "DELETE": + operation = "delete" + elif method_upper == "POST" and resource_type == "batch": + # POST to /batches/{id} is not standard; treat as retrieve for safety + operation = "retrieve" + else: + operation = "retrieve" + else: + # No resource ID + if method_upper == "POST": + operation = "create" + elif method_upper == "GET": + operation = "list" + else: + return None + + return AzurePassthroughOp( + resource_type=resource_type, + operation=operation, + resource_id=resource_id, + ) + + +async def passthrough_pre_request( + endpoint: str, + request_method: str, + user_api_key_dict: UserAPIKeyAuth, + managed_files_obj: Any, +) -> Tuple[str, Optional[AzurePassthroughOp]]: + """ + Pre-request hook for Azure pass-through permission management. + + For read/delete/cancel/content operations: + - Checks if the resource ID is a managed (encoded) ID + - Verifies ownership + - Rewrites the endpoint to use the real provider ID + + Args: + endpoint: URL path after /azure/ + request_method: HTTP method + user_api_key_dict: Authenticated user info + managed_files_obj: Enterprise managed files hook instance + + Returns: + Tuple of (potentially rewritten endpoint, classified operation) + """ + op = classify_azure_passthrough_request(endpoint, request_method) + if op is None: + return endpoint, None + + # Only intercept operations that target a specific resource + if op.operation not in ("retrieve", "delete", "cancel", "content"): + return endpoint, op + + if op.resource_id is None: + return endpoint, op + + if managed_files_obj is None: + return endpoint, op + + # Check if the resource ID is a managed (encoded) ID + resource_id = op.resource_id + + # Check if the resource ID is a model-embedded (managed) ID + if is_model_embedded_id(resource_id): + if op.resource_type == "file": + can_access = await managed_files_obj.can_user_call_unified_file_id( + unified_file_id=resource_id, + user_api_key_dict=user_api_key_dict, + ) + else: + can_access = await managed_files_obj.can_user_call_unified_object_id( + unified_object_id=resource_id, + user_api_key_dict=user_api_key_dict, + ) + if not can_access: + raise HTTPException( + status_code=403, + detail=f"You don't have permission to access this {op.resource_type}.", + ) + # Decode real provider ID and rewrite endpoint + real_id = get_original_file_id(resource_id) + endpoint = endpoint.replace(resource_id, real_id) + op.resource_id = real_id + elif op.resource_type in ("batch", "response"): + # Raw provider ID - check if it exists in the managed object table + try: + db_obj = await managed_files_obj.prisma_client.db.litellm_managedobjecttable.find_first( + where={"model_object_id": resource_id} + ) + if db_obj and db_obj.created_by != user_api_key_dict.user_id: + raise HTTPException( + status_code=403, + detail=f"You don't have permission to access this {op.resource_type}.", + ) + except HTTPException: + raise + except Exception: + # DB lookup failed - pass through (backward compat) + pass + elif op.resource_type == "file": + # Raw file ID - check if it exists in the managed file table + try: + db_obj = await managed_files_obj.prisma_client.db.litellm_managedfiletable.find_first( + where={"flat_model_file_ids": {"has": resource_id}} + ) + if db_obj and db_obj.created_by != user_api_key_dict.user_id: + raise HTTPException( + status_code=403, + detail="You don't have permission to access this file.", + ) + except HTTPException: + raise + except Exception: + # DB lookup failed - pass through (backward compat) + pass + + return endpoint, op + + +async def passthrough_post_response( + response_body: bytes, + op: AzurePassthroughOp, + user_api_key_dict: UserAPIKeyAuth, + managed_files_obj: Any, + deployment_name: Optional[str] = None, +) -> bytes: + """ + Post-response hook for Azure pass-through permission management. + + For create operations: + - Parses the response body + - Encodes provider IDs into managed IDs + - Stores ownership in the database + + Args: + response_body: Raw response body bytes from Azure + op: Classified operation + user_api_key_dict: Authenticated user info + managed_files_obj: Enterprise managed files hook instance + deployment_name: Azure deployment name (used as model for encoding) + + Returns: + Potentially modified response body bytes + """ + if op.operation != "create": + return response_body + + if managed_files_obj is None: + return response_body + + try: + response_json = json.loads(response_body) + except (json.JSONDecodeError, UnicodeDecodeError): + return response_body + + provider_id = response_json.get("id") + if not provider_id: + return response_body + + # Use deployment_name as the model for encoding, fallback to "azure-passthrough" + model = deployment_name or "azure-passthrough" + + if op.resource_type == "file": + # Encode file ID + managed_id = encode_file_id_with_model( + file_id=provider_id, model=model, id_type="file" + ) + response_json["id"] = managed_id + + # Store ownership in DB + try: + from litellm.types.llms.openai import OpenAIFileObject + + file_object = OpenAIFileObject(**response_json) + file_object.id = provider_id # Store with original ID + await managed_files_obj.store_unified_file_id( + file_id=managed_id, + file_object=file_object, + litellm_parent_otel_span=None, + model_mappings={model: provider_id}, + user_api_key_dict=user_api_key_dict, + ) + except Exception as e: + verbose_logger.error( + f"Failed to store managed file ID for passthrough: {e}" + ) + + elif op.resource_type == "batch": + # Encode batch ID + managed_id = encode_file_id_with_model( + file_id=provider_id, model=model, id_type="batch" + ) + response_json["id"] = managed_id + + # Also encode input_file_id, output_file_id, error_file_id if present + for file_field in ("input_file_id", "output_file_id", "error_file_id"): + raw_file_id = response_json.get(file_field) + if raw_file_id: + response_json[file_field] = encode_file_id_with_model( + file_id=raw_file_id, model=model, id_type="file" + ) + + # Store ownership in DB + try: + from litellm.types.utils import LiteLLMBatch + + batch_object = LiteLLMBatch(**response_json) + batch_object.id = provider_id # Store with original ID + await managed_files_obj.store_unified_object_id( + unified_object_id=managed_id, + file_object=batch_object, + litellm_parent_otel_span=None, + model_object_id=provider_id, + file_purpose="batch", + user_api_key_dict=user_api_key_dict, + ) + except Exception as e: + verbose_logger.error( + f"Failed to store managed batch ID for passthrough: {e}" + ) + + elif op.resource_type == "response": + # Encode response ID + managed_id = encode_file_id_with_model( + file_id=provider_id, model=model, id_type="batch" + ) + response_json["id"] = managed_id + + # Store ownership in DB + try: + from litellm.types.llms.openai import ResponsesAPIResponse + + response_obj = ResponsesAPIResponse(**response_json) + response_obj.id = provider_id # Store with original ID + await managed_files_obj.store_unified_object_id( + unified_object_id=managed_id, + file_object=response_obj, + litellm_parent_otel_span=None, + model_object_id=provider_id, + file_purpose="response", + user_api_key_dict=user_api_key_dict, + ) + except Exception as e: + verbose_logger.error( + f"Failed to store managed response ID for passthrough: {e}" + ) + + return json.dumps(response_json).encode() + + +async def passthrough_list_filter( + op: AzurePassthroughOp, + user_api_key_dict: UserAPIKeyAuth, + managed_files_obj: Any, +) -> Optional[Response]: + """ + Handle list operations by filtering from the managed tables. + + Returns a Response if handled, None to fall through to Azure. + """ + if op.operation != "list": + return None + + if managed_files_obj is None: + return None + + if op.resource_type == "batch": + try: + result = await managed_files_obj.list_user_batches( + user_api_key_dict=user_api_key_dict, + ) + return Response( + content=json.dumps(result), + media_type="application/json", + status_code=200, + ) + except Exception as e: + verbose_logger.warning( + f"Failed to list managed batches, falling through to Azure: {e}" + ) + return None + + elif op.resource_type == "file": + try: + where_clause: Dict[str, Any] = {} + if user_api_key_dict.user_id: + where_clause["created_by"] = user_api_key_dict.user_id + + files = await managed_files_obj.prisma_client.db.litellm_managedfiletable.find_many( + where=where_clause, + order={"created_at": "desc"}, + ) + + file_objects = [] + for f in files: + try: + file_data = ( + json.loads(f.file_object) + if isinstance(f.file_object, str) + else f.file_object + ) + if isinstance(file_data, dict): + file_data["id"] = f.unified_file_id + file_objects.append(file_data) + except Exception: + continue + + result = { + "object": "list", + "data": file_objects, + "has_more": False, + } + return Response( + content=json.dumps(result), + media_type="application/json", + status_code=200, + ) + except Exception as e: + verbose_logger.warning( + f"Failed to list managed files, falling through to Azure: {e}" + ) + return None + + return None diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 534022cc133..f55fff89664 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1284,7 +1284,47 @@ async def azure_proxy_route( Checks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route. """ - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + + # --- Permission management hooks (enterprise) --- + _passthrough_pre_request = None + _passthrough_post_response = None + _passthrough_list_filter = None + _managed_files_obj = None + try: + from litellm_enterprise.proxy.hooks.azure_passthrough_permissions import ( + passthrough_list_filter as _passthrough_list_filter, + passthrough_post_response as _passthrough_post_response, + passthrough_pre_request as _passthrough_pre_request, + ) + + if proxy_logging_obj is not None: + _managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") + except ImportError: + pass + + passthrough_op = None + if _passthrough_pre_request is not None and _managed_files_obj is not None: + endpoint, passthrough_op = await _passthrough_pre_request( + endpoint=endpoint, + request_method=request.method, + user_api_key_dict=user_api_key_dict, + managed_files_obj=_managed_files_obj, + ) + + # Handle list operations by filtering from managed tables + if ( + passthrough_op is not None + and passthrough_op.operation == "list" + and _passthrough_list_filter is not None + ): + list_response = await _passthrough_list_filter( + op=passthrough_op, + user_api_key_dict=user_api_key_dict, + managed_files_obj=_managed_files_obj, + ) + if list_response is not None: + return list_response parts = endpoint.split( "/" @@ -1353,9 +1393,24 @@ async def azure_proxy_route( ), ) - # Non-streaming response + # Non-streaming response - apply post-response permission hook result = cast(httpx.Response, result) content = await result.aread() + + if ( + passthrough_op is not None + and passthrough_op.operation == "create" + and _passthrough_post_response is not None + and _managed_files_obj is not None + ): + content = await _passthrough_post_response( + response_body=content, + op=passthrough_op, + user_api_key_dict=user_api_key_dict, + managed_files_obj=_managed_files_obj, + deployment_name=part, + ) + return Response( content=content, status_code=result.status_code, @@ -1439,16 +1494,49 @@ async def azure_proxy_route( "Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure." ) - return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( - endpoint=endpoint, - request=request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - base_target_url=base_target_url, - api_key=azure_api_key, - custom_llm_provider=litellm.LlmProviders.AZURE, + received_value = ( + await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( + endpoint=endpoint, + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + base_target_url=base_target_url, + api_key=azure_api_key, + custom_llm_provider=litellm.LlmProviders.AZURE, + ) ) + # Post-response permission hook for create operations (default path) + if ( + passthrough_op is not None + and passthrough_op.operation == "create" + and _passthrough_post_response is not None + and _managed_files_obj is not None + and isinstance(received_value, Response) + and not isinstance(received_value, StreamingResponse) + ): + # Extract deployment name from endpoint parts + _deployment_name = None + for _part in parts: + if _part not in ("openai", "deployments", "batches", "files", "responses"): + _deployment_name = _part + break + modified_body = await _passthrough_post_response( + response_body=received_value.body, + op=passthrough_op, + user_api_key_dict=user_api_key_dict, + managed_files_obj=_managed_files_obj, + deployment_name=_deployment_name, + ) + received_value = Response( + content=modified_body, + status_code=received_value.status_code, + headers=dict(received_value.headers), + media_type=received_value.media_type, + ) + + return received_value + from abc import ABC, abstractmethod diff --git a/tests/test_litellm/enterprise/proxy/test_azure_passthrough_permissions.py b/tests/test_litellm/enterprise/proxy/test_azure_passthrough_permissions.py new file mode 100644 index 00000000000..73a8871a4b3 --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_azure_passthrough_permissions.py @@ -0,0 +1,614 @@ +""" +Tests for Azure pass-through permission management. + +Tests cover: +- URL classification for all resource types and operations +- Pre-request ownership checks and ID rewriting +- Post-response ID encoding and ownership storage +- List filtering by user +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, +) +from litellm_enterprise.proxy.hooks.azure_passthrough_permissions import ( + AzurePassthroughOp, + classify_azure_passthrough_request, + passthrough_list_filter, + passthrough_post_response, + passthrough_pre_request, +) + + +# ============================================================================ +# URL CLASSIFICATION TESTS +# ============================================================================ + + +class TestClassifyAzurePassthroughRequest: + def test_create_batch(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/batches", "POST" + ) + assert op is not None + assert op.resource_type == "batch" + assert op.operation == "create" + assert op.resource_id is None + + def test_retrieve_batch(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/batches/batch_abc123", "GET" + ) + assert op is not None + assert op.resource_type == "batch" + assert op.operation == "retrieve" + assert op.resource_id == "batch_abc123" + + def test_cancel_batch(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/batches/batch_abc123/cancel", "POST" + ) + assert op is not None + assert op.resource_type == "batch" + assert op.operation == "cancel" + assert op.resource_id == "batch_abc123" + + def test_list_batches(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/batches", "GET" + ) + assert op is not None + assert op.resource_type == "batch" + assert op.operation == "list" + assert op.resource_id is None + + def test_create_file(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/files", "POST" + ) + assert op is not None + assert op.resource_type == "file" + assert op.operation == "create" + assert op.resource_id is None + + def test_retrieve_file(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/files/file-abc123", "GET" + ) + assert op is not None + assert op.resource_type == "file" + assert op.operation == "retrieve" + assert op.resource_id == "file-abc123" + + def test_delete_file(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/files/file-abc123", "DELETE" + ) + assert op is not None + assert op.resource_type == "file" + assert op.operation == "delete" + assert op.resource_id == "file-abc123" + + def test_get_file_content(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/files/file-abc123/content", "GET" + ) + assert op is not None + assert op.resource_type == "file" + assert op.operation == "content" + assert op.resource_id == "file-abc123" + + def test_list_files(self): + op = classify_azure_passthrough_request("openai/deployments/gpt-4/files", "GET") + assert op is not None + assert op.resource_type == "file" + assert op.operation == "list" + assert op.resource_id is None + + def test_create_response(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/responses", "POST" + ) + assert op is not None + assert op.resource_type == "response" + assert op.operation == "create" + assert op.resource_id is None + + def test_retrieve_response(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/responses/resp_abc123", "GET" + ) + assert op is not None + assert op.resource_type == "response" + assert op.operation == "retrieve" + assert op.resource_id == "resp_abc123" + + def test_delete_response(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/responses/resp_abc123", "DELETE" + ) + assert op is not None + assert op.resource_type == "response" + assert op.operation == "delete" + assert op.resource_id == "resp_abc123" + + def test_non_matching_endpoint(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/completions", "POST" + ) + assert op is None + + def test_non_matching_chat_completions(self): + op = classify_azure_passthrough_request( + "openai/deployments/gpt-4/chat/completions", "POST" + ) + assert op is None + + +# ============================================================================ +# PRE-REQUEST TESTS +# ============================================================================ + + +class TestPassthroughPreRequest: + @pytest.mark.asyncio + async def test_non_matching_endpoint_passes_through(self): + endpoint = "openai/deployments/gpt-4/completions" + new_endpoint, op = await passthrough_pre_request( + endpoint=endpoint, + request_method="POST", + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=None, + ) + assert new_endpoint == endpoint + assert op is None + + @pytest.mark.asyncio + async def test_create_operation_passes_through(self): + endpoint = "openai/deployments/gpt-4/batches" + new_endpoint, op = await passthrough_pre_request( + endpoint=endpoint, + request_method="POST", + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=MagicMock(), + ) + assert new_endpoint == endpoint + assert op is not None + assert op.operation == "create" + + @pytest.mark.asyncio + async def test_retrieve_with_managed_batch_id_owner_allowed(self): + """Owner can retrieve a managed batch.""" + # Encode a batch ID with model info + managed_id = encode_file_id_with_model( + file_id="batch_real123", model="gpt-4", id_type="batch" + ) + + managed_files_obj = AsyncMock() + managed_files_obj.can_user_call_unified_object_id = AsyncMock(return_value=True) + + endpoint = f"openai/deployments/gpt-4/batches/{managed_id}" + new_endpoint, op = await passthrough_pre_request( + endpoint=endpoint, + request_method="GET", + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + assert op is not None + assert op.operation == "retrieve" + # Endpoint should be rewritten with the real provider ID + assert "batch_real123" in new_endpoint + assert managed_id not in new_endpoint + + @pytest.mark.asyncio + async def test_retrieve_with_managed_batch_id_non_owner_blocked(self): + """Non-owner gets 403 when trying to retrieve a managed batch.""" + managed_id = encode_file_id_with_model( + file_id="batch_real123", model="gpt-4", id_type="batch" + ) + + managed_files_obj = AsyncMock() + managed_files_obj.can_user_call_unified_object_id = AsyncMock( + return_value=False + ) + + endpoint = f"openai/deployments/gpt-4/batches/{managed_id}" + with pytest.raises(HTTPException) as exc_info: + await passthrough_pre_request( + endpoint=endpoint, + request_method="GET", + user_api_key_dict=UserAPIKeyAuth( + user_id="user2", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_retrieve_with_managed_file_id_owner_allowed(self): + """Owner can retrieve a managed file.""" + managed_id = encode_file_id_with_model( + file_id="file-real456", model="gpt-4", id_type="file" + ) + + managed_files_obj = AsyncMock() + managed_files_obj.can_user_call_unified_file_id = AsyncMock(return_value=True) + + endpoint = f"openai/deployments/gpt-4/files/{managed_id}" + new_endpoint, op = await passthrough_pre_request( + endpoint=endpoint, + request_method="GET", + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + assert "file-real456" in new_endpoint + assert managed_id not in new_endpoint + + @pytest.mark.asyncio + async def test_retrieve_with_managed_file_id_non_owner_blocked(self): + """Non-owner gets 403 when trying to retrieve a managed file.""" + managed_id = encode_file_id_with_model( + file_id="file-real456", model="gpt-4", id_type="file" + ) + + managed_files_obj = AsyncMock() + managed_files_obj.can_user_call_unified_file_id = AsyncMock(return_value=False) + + endpoint = f"openai/deployments/gpt-4/files/{managed_id}" + with pytest.raises(HTTPException) as exc_info: + await passthrough_pre_request( + endpoint=endpoint, + request_method="DELETE", + user_api_key_dict=UserAPIKeyAuth( + user_id="user2", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_raw_provider_id_ownership_check_via_db(self): + """Raw (non-encoded) IDs are checked via model_object_id lookup.""" + managed_files_obj = AsyncMock() + db_obj = MagicMock() + db_obj.created_by = "user1" + managed_files_obj.prisma_client.db.litellm_managedobjecttable.find_first = ( + AsyncMock(return_value=db_obj) + ) + + endpoint = "openai/deployments/gpt-4/batches/batch_raw789" + new_endpoint, op = await passthrough_pre_request( + endpoint=endpoint, + request_method="GET", + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + # Should pass through since ownership matches + assert new_endpoint == endpoint + + @pytest.mark.asyncio + async def test_raw_provider_id_non_owner_blocked(self): + """Raw ID with different owner gets 403.""" + managed_files_obj = AsyncMock() + db_obj = MagicMock() + db_obj.created_by = "user1" + managed_files_obj.prisma_client.db.litellm_managedobjecttable.find_first = ( + AsyncMock(return_value=db_obj) + ) + + endpoint = "openai/deployments/gpt-4/batches/batch_raw789" + with pytest.raises(HTTPException) as exc_info: + await passthrough_pre_request( + endpoint=endpoint, + request_method="GET", + user_api_key_dict=UserAPIKeyAuth( + user_id="user2", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_untracked_raw_id_passes_through(self): + """Raw ID not in DB passes through (backward compat).""" + managed_files_obj = AsyncMock() + managed_files_obj.prisma_client.db.litellm_managedobjecttable.find_first = ( + AsyncMock(return_value=None) + ) + + endpoint = "openai/deployments/gpt-4/batches/batch_unknown" + new_endpoint, op = await passthrough_pre_request( + endpoint=endpoint, + request_method="GET", + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + assert new_endpoint == endpoint + + @pytest.mark.asyncio + async def test_no_managed_files_obj_passes_through(self): + """When enterprise is not available, everything passes through.""" + endpoint = "openai/deployments/gpt-4/batches/batch_abc" + new_endpoint, op = await passthrough_pre_request( + endpoint=endpoint, + request_method="GET", + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=None, + ) + assert new_endpoint == endpoint + assert op is not None + assert op.operation == "retrieve" + + +# ============================================================================ +# POST-RESPONSE TESTS +# ============================================================================ + + +class TestPassthroughPostResponse: + @pytest.mark.asyncio + async def test_create_batch_encodes_id_and_stores(self): + """Batch create response gets encoded ID and stored in DB.""" + response_json = { + "id": "batch_abc123", + "object": "batch", + "status": "validating", + "input_file_id": "file-input456", + "output_file_id": None, + "error_file_id": None, + "completion_window": "24h", + "created_at": 1234567890, + "endpoint": "/v1/chat/completions", + } + response_body = json.dumps(response_json).encode() + + managed_files_obj = AsyncMock() + managed_files_obj.store_unified_object_id = AsyncMock() + + op = AzurePassthroughOp( + resource_type="batch", operation="create", resource_id=None + ) + result = await passthrough_post_response( + response_body=response_body, + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + deployment_name="gpt-4", + ) + + result_json = json.loads(result) + # ID should be encoded (not the raw provider ID) + assert result_json["id"] != "batch_abc123" + assert result_json["id"].startswith("batch_") + # input_file_id should also be encoded + assert result_json["input_file_id"] != "file-input456" + assert result_json["input_file_id"].startswith("file-") + # DB storage should have been called + managed_files_obj.store_unified_object_id.assert_called_once() + + @pytest.mark.asyncio + async def test_create_file_encodes_id_and_stores(self): + """File create response gets encoded ID and stored in DB.""" + response_json = { + "id": "file-xyz789", + "object": "file", + "bytes": 1234, + "created_at": 1234567890, + "filename": "test.jsonl", + "purpose": "batch", + } + response_body = json.dumps(response_json).encode() + + managed_files_obj = AsyncMock() + managed_files_obj.store_unified_file_id = AsyncMock() + + op = AzurePassthroughOp( + resource_type="file", operation="create", resource_id=None + ) + result = await passthrough_post_response( + response_body=response_body, + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + deployment_name="gpt-4", + ) + + result_json = json.loads(result) + assert result_json["id"] != "file-xyz789" + assert result_json["id"].startswith("file-") + managed_files_obj.store_unified_file_id.assert_called_once() + + @pytest.mark.asyncio + async def test_non_create_operation_passes_through(self): + """Non-create operations don't modify the response.""" + response_body = b'{"id": "batch_abc", "status": "completed"}' + + op = AzurePassthroughOp( + resource_type="batch", operation="retrieve", resource_id="batch_abc" + ) + result = await passthrough_post_response( + response_body=response_body, + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=AsyncMock(), + ) + assert result == response_body + + @pytest.mark.asyncio + async def test_invalid_json_passes_through(self): + """Invalid JSON response is returned unchanged.""" + response_body = b"not valid json" + + op = AzurePassthroughOp( + resource_type="batch", operation="create", resource_id=None + ) + result = await passthrough_post_response( + response_body=response_body, + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=AsyncMock(), + ) + assert result == response_body + + @pytest.mark.asyncio + async def test_no_managed_files_passes_through(self): + """Without managed_files_obj, response passes through unchanged.""" + response_body = b'{"id": "batch_abc"}' + + op = AzurePassthroughOp( + resource_type="batch", operation="create", resource_id=None + ) + result = await passthrough_post_response( + response_body=response_body, + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=None, + ) + assert result == response_body + + +# ============================================================================ +# LIST FILTER TESTS +# ============================================================================ + + +class TestPassthroughListFilter: + @pytest.mark.asyncio + async def test_list_batches_filters_by_user(self): + """List batches returns user-filtered results from managed tables.""" + managed_files_obj = AsyncMock() + managed_files_obj.list_user_batches = AsyncMock( + return_value={ + "object": "list", + "data": [{"id": "batch_1", "status": "completed"}], + "has_more": False, + } + ) + + op = AzurePassthroughOp( + resource_type="batch", operation="list", resource_id=None + ) + result = await passthrough_list_filter( + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + assert result is not None + assert result.status_code == 200 + body = json.loads(result.body) + assert body["data"][0]["id"] == "batch_1" + + @pytest.mark.asyncio + async def test_list_files_filters_by_user(self): + """List files returns user-filtered results from managed tables.""" + mock_file = MagicMock() + mock_file.file_object = json.dumps( + { + "id": "file-raw", + "object": "file", + "bytes": 100, + "created_at": 1234567890, + "filename": "test.jsonl", + "purpose": "batch", + } + ) + mock_file.unified_file_id = "file-managed123" + + managed_files_obj = AsyncMock() + managed_files_obj.prisma_client.db.litellm_managedfiletable.find_many = ( + AsyncMock(return_value=[mock_file]) + ) + + op = AzurePassthroughOp( + resource_type="file", operation="list", resource_id=None + ) + result = await passthrough_list_filter( + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=managed_files_obj, + ) + assert result is not None + assert result.status_code == 200 + body = json.loads(result.body) + assert len(body["data"]) == 1 + assert body["data"][0]["id"] == "file-managed123" + + @pytest.mark.asyncio + async def test_non_list_operation_returns_none(self): + """Non-list operations return None (fall through to Azure).""" + op = AzurePassthroughOp( + resource_type="batch", operation="create", resource_id=None + ) + result = await passthrough_list_filter( + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=AsyncMock(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_no_managed_files_returns_none(self): + """Without managed_files_obj, returns None (fall through).""" + op = AzurePassthroughOp( + resource_type="batch", operation="list", resource_id=None + ) + result = await passthrough_list_filter( + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=None, + ) + assert result is None + + @pytest.mark.asyncio + async def test_response_list_returns_none(self): + """List responses not yet supported, returns None.""" + op = AzurePassthroughOp( + resource_type="response", operation="list", resource_id=None + ) + result = await passthrough_list_filter( + op=op, + user_api_key_dict=UserAPIKeyAuth( + user_id="user1", parent_otel_span=MagicMock() + ), + managed_files_obj=AsyncMock(), + ) + assert result is None