fix(proxy): authorize batch files using upload target_model_names (LIT-3593) (#30009)

* fix(proxy): authorize batch files using upload target_model_names (LIT-3593) After replace_model_in_jsonl, body.model is a stripped provider id. Reverse-mapping it via resolve_model_name_from_model_id is first-match on model_list and caused false 403s when multiple deployments share the same stripped name. Use target_model_names from the unified file id instead. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593) Restores the reverse-lookup for the JSONL body.model fallback path so that legacy/pre-target_model_names managed files still map stripped provider IDs back to proxy aliases before auth. Also cleans up redundant `or None`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Revert "fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)" This reverts commit 30d2e96f77ef521ccaaf2193fe554980380eb669. --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> (cherry picked from commit 2cd7e87485)
2026-08-03 10:23:28 +00:00 · 2026-06-10 19:04:21 -07:00
parent 2d2fb8c131
commit 5cea68f3ae
2 changed files with 116 additions and 12 deletions
@@ -227,11 +227,17 @@ class _PROXY_BatchRateLimiter(CustomLogger):
            # Check if this is a managed file (base64 encoded unified file ID)
            from litellm.proxy.openai_files_endpoints.common_utils import (
                _is_base64_encoded_unified_file_id,
+                get_models_from_unified_file_id,
            )

            # Managed files require bypassing the HTTP endpoint (which runs access-check hooks)
            # and calling the managed files hook directly with the user's credentials.
            is_managed_file = _is_base64_encoded_unified_file_id(file_id)
+            target_model_names = (
+                get_models_from_unified_file_id(is_managed_file)
+                if is_managed_file
+                else []
+            )
            if is_managed_file and user_api_key_dict is not None:
                file_content = await self._fetch_managed_file_content(
                    file_id=file_id,
@@ -256,6 +262,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
                await self._enforce_batch_file_model_access(
                    user_api_key_dict=user_api_key_dict,
                    file_content_as_dict=file_content_as_dict,
+                    target_model_names=target_model_names or None,
                )

            input_file_usage = _get_batch_job_input_file_usage(
@@ -291,9 +298,13 @@ class _PROXY_BatchRateLimiter(CustomLogger):
        self,
        user_api_key_dict: UserAPIKeyAuth,
        file_content_as_dict: List[dict],
+        target_model_names: Optional[List[str]] = None,
    ) -> None:
-        """Reject the batch if the caller is not authorized for every
-        ``body.model`` named inside the JSONL.
+        """Reject the batch if the caller is not authorized for the upload target.
+
+        For managed files, ``target_model_names`` (from the unified file id) is
+        the proxy alias the file was uploaded for and is used directly for auth.
+        For legacy/non-managed files, falls back to ``body.model`` values in the JSONL.

        Reuses ``can_key_call_model`` so the same allowlist semantics
        (wildcards, access groups, ``all-proxy-models``, team aliases)
@@ -302,18 +313,16 @@ class _PROXY_BatchRateLimiter(CustomLogger):
        from litellm.proxy.auth.auth_checks import can_key_call_model
        from litellm.proxy.proxy_server import llm_router

-        models = _get_models_from_batch_input_file_content(file_content_as_dict)
-        if not models:
-            return
+        if target_model_names:
+            models = target_model_names
+        else:
+            models = _get_models_from_batch_input_file_content(file_content_as_dict)
+            if not models:
+                return

        llm_model_list = llm_router.model_list if llm_router is not None else None
        for model in models:
-            # body.model may be the provider id after replace_model_in_jsonl; map to proxy model_name for auth.
            model_to_check = model
-            if llm_router is not None:
-                proxy_model_name = llm_router.resolve_model_name_from_model_id(model)
-                if proxy_model_name is not None:
-                    model_to_check = proxy_model_name
            try:
                await can_key_call_model(
                    model=model_to_check,
@@ -262,7 +262,8 @@ async def test_pre_call_allows_authorized_model_in_batch_file():
@pytest.mark.asyncio
 async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias():
    """After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5).
-    Auth must check the proxy model_name the key was granted, not the stripped id."""
+    Auth must check target_model_names from the unified file id, not reverse-map
+    the stripped id."""
    from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter

    rate_limiter = _PROXY_BatchRateLimiter(
@@ -281,7 +282,6 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(
    )
    mock_router = MagicMock()
    mock_router.model_list = []
-    mock_router.resolve_model_name_from_model_id.return_value = proxy_alias
    can_key_call_model = AsyncMock(return_value=True)

    with (
@@ -294,10 +294,105 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(
        await rate_limiter._enforce_batch_file_model_access(
            user_api_key_dict=user,
            file_content_as_dict=file_dict,
+            target_model_names=[proxy_alias],
        )

    can_key_call_model.assert_awaited_once()
    assert can_key_call_model.await_args.kwargs["model"] == proxy_alias
+    mock_router.resolve_model_name_from_model_id.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_list_order",
+    [
+        [
+            "openai/openai/gpt-5.5",
+            "openai/openai/gpt-5.5-batch",
+            "us/azure/openai/gpt-5.5",
+        ],
+        [
+            "us/azure/openai/gpt-5.5",
+            "openai/openai/gpt-5.5",
+            "openai/openai/gpt-5.5-batch",
+        ],
+        [
+            "openai/openai/gpt-5.5-batch",
+            "us/azure/openai/gpt-5.5",
+            "openai/openai/gpt-5.5",
+        ],
+    ],
+)
+async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup(
+    model_list_order,
+):
+    """LIT-3593: three deployments strip to gpt-5.5; auth must use the upload
+    target alias from target_model_names, not first-match reverse lookup."""
+    from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
+
+    rate_limiter = _PROXY_BatchRateLimiter(
+        internal_usage_cache=MagicMock(),
+        parallel_request_limiter=MagicMock(),
+    )
+    batch_alias = "openai/openai/gpt-5.5-batch"
+    deployment_templates = {
+        "openai/openai/gpt-5.5": {
+            "model_name": "openai/openai/gpt-5.5",
+            "litellm_params": {"model": "openai/gpt-5.5"},
+            "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
+        },
+        "openai/openai/gpt-5.5-batch": {
+            "model_name": "openai/openai/gpt-5.5-batch",
+            "litellm_params": {"model": "openai/gpt-5.5"},
+            "model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"},
+        },
+        "us/azure/openai/gpt-5.5": {
+            "model_name": "us/azure/openai/gpt-5.5",
+            "litellm_params": {"model": "azure/gpt-5.5"},
+            "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
+        },
+    }
+    mock_router = MagicMock()
+    mock_router.model_list = [deployment_templates[name] for name in model_list_order]
+
+    def _resolve(model_id):
+        for deployment in mock_router.model_list:
+            actual_model = deployment.get("litellm_params", {}).get("model")
+            if actual_model == model_id or (
+                actual_model and actual_model.endswith(f"/{model_id}")
+            ):
+                return deployment.get("model_name")
+        return None
+
+    mock_router.resolve_model_name_from_model_id.side_effect = _resolve
+
+    file_dict = [
+        {"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}}
+    ]
+    user = UserAPIKeyAuth(
+        api_key="sk-ok",
+        user_id="alice",
+        models=[batch_alias],
+        user_role=LitellmUserRoles.INTERNAL_USER.value,
+    )
+    can_key_call_model = AsyncMock(return_value=True)
+
+    with (
+        patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new=can_key_call_model,
+        ),
+        patch("litellm.proxy.proxy_server.llm_router", mock_router),
+    ):
+        await rate_limiter._enforce_batch_file_model_access(
+            user_api_key_dict=user,
+            file_content_as_dict=file_dict,
+            target_model_names=[batch_alias],
+        )
+
+    can_key_call_model.assert_awaited_once()
+    assert can_key_call_model.await_args.kwargs["model"] == batch_alias
+    mock_router.resolve_model_name_from_model_id.assert_not_called()


@pytest.mark.asyncio