diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index d57c1a78f3d..fe02f5a46a0 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -73,6 +73,24 @@ else: PrismaClient = Any +def _parse_managed_file_object( + raw_file_object: object, unified_file_id: str +) -> Optional[OpenAIFileObject]: + if not raw_file_object: + return None + try: + return ( + OpenAIFileObject.model_validate_json(raw_file_object) + if isinstance(raw_file_object, str) + else OpenAIFileObject.model_validate(raw_file_object) + ) + except Exception as e: + verbose_logger.warning( + f"Failed to parse managed file object {unified_file_id}: {e}" + ) + return None + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -376,13 +394,43 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if owner_filter is None: return [] - file_ids = await self.prisma_client.db.litellm_managedfiletable.find_many( + rows = await self.prisma_client.db.litellm_managedfiletable.find_many( where={ **owner_filter, "flat_model_file_ids": {"hasSome": model_object_ids}, } ) - return [OpenAIFileObject.model_validate(file_object.file_object) for file_object in file_ids] + return [ + parsed + for row in rows + if ( + parsed := _parse_managed_file_object( + row.file_object, row.unified_file_id + ) + ) + is not None + ] + + async def _provider_file_ids_claimed_by_managed_rows( + self, model_object_ids: List[str] + ) -> set[str]: + """Provider file ids from ``model_object_ids`` that appear in any managed row. + + Used by list filtering so untracked (non-managed) provider uploads stay + visible, while managed rows owned by other callers can still be hidden. + """ + if not model_object_ids: + return set() + page_ids = set(model_object_ids) + rows = await self.prisma_client.db.litellm_managedfiletable.find_many( + where={"flat_model_file_ids": {"hasSome": model_object_ids}}, + ) + claimed: set[str] = set() + for row in rows: + for provider_id in row.flat_model_file_ids or []: + if provider_id in page_ids: + claimed.add(provider_id) + return claimed async def check_managed_file_id_access( self, data: Dict, user_api_key_dict: UserAPIKeyAuth @@ -1248,30 +1296,70 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) elif isinstance(response, AsyncCursorPage): """ - For listing files, filter for the ones created by the user + Filter a provider file list for multi-tenant isolation. + + Managed rows on a shared provider account must only be visible to + their owner. Untracked (non-managed) provider uploads, including + provider-scoped /{provider}/v1/files creates that never entered the + managed table, must stay on the page; dropping them is LIT-4820. """ - ## check if file object if hasattr(response, "data") and isinstance(response.data, list): if all( isinstance(file_object, FileObject) for file_object in response.data ): - ## Get all file id's - ## Check which file id's were created by the user - ## Filter the response to only include the files created by the user - ## Return the filtered response - file_ids = [ - file_object.id - for file_object in cast(List[FileObject], response.data) # type: ignore - ] - user_created_file_ids = await self.get_user_created_file_ids( - user_api_key_dict, file_ids + response.data = await self._filter_listed_provider_files( # type: ignore + provider_files=cast(List[FileObject], response.data), # type: ignore + user_api_key_dict=user_api_key_dict, ) - ## Filter the response to only include the files created by the user - response.data = user_created_file_ids # type: ignore return response return response return response + async def _filter_listed_provider_files( + self, + provider_files: List[FileObject], + user_api_key_dict: UserAPIKeyAuth, + ) -> List[Union[FileObject, OpenAIFileObject]]: + file_ids = [file_object.id for file_object in provider_files] + if not file_ids: + return [] + + claimed_by_anyone = await self._provider_file_ids_claimed_by_managed_rows( + file_ids + ) + owner_filter = build_owner_filter(user_api_key_dict) + user_owned_by_provider_id: Dict[str, OpenAIFileObject] = {} + if owner_filter is not None and claimed_by_anyone: + user_owned_rows = await self.prisma_client.db.litellm_managedfiletable.find_many( + where={ + **owner_filter, + "flat_model_file_ids": {"hasSome": list(claimed_by_anyone)}, + } + ) + page_id_set = set(file_ids) + for row in user_owned_rows: + parsed = _parse_managed_file_object(row.file_object, row.unified_file_id) + if parsed is None: + continue + for provider_id in row.flat_model_file_ids or []: + if provider_id in page_id_set: + user_owned_by_provider_id[provider_id] = parsed + + kept: List[Union[FileObject, OpenAIFileObject]] = [] + emitted_managed_ids: set[str] = set() + for provider_file in provider_files: + if provider_file.id not in claimed_by_anyone: + kept.append(provider_file) + continue + managed_obj = user_owned_by_provider_id.get(provider_file.id) + if managed_obj is None: + continue + if managed_obj.id in emitted_managed_ids: + continue + kept.append(managed_obj) + emitted_managed_ids.add(managed_obj.id) + return kept + async def afile_retrieve( self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router=None ) -> OpenAIFileObject: diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 51b9139cff1..960af09f433 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -56,6 +56,7 @@ from litellm.repositories.table_repositories import ManagedFileRepository from litellm.router import Router from litellm.types.llms.openai import ( CREATE_FILE_REQUESTS_PURPOSE, + AsyncCursorPage, FileExpiresAfter, OpenAIFileObject, OpenAIFilesPurpose, @@ -1397,7 +1398,9 @@ async def list_files( _response = await proxy_logging_obj.post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=response ) - if _response is not None and isinstance(_response, OpenAIFileObject): + if _response is not None and isinstance( + _response, (OpenAIFileObject, AsyncCursorPage) + ): response = _response ### ALERTING ### diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 180639b53e3..3fdfcc76351 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -33,8 +33,6 @@ Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datad - Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find; dual-read with `dd_logs` from conftest when delivery matters - Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down - If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream; the default is always Datadog -- The one standing exception is `test_mcp_chat_completion_oauth_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so it cannot exercise gateway-managed OAuth or per-user token seeding in any form. That test drives a real Linear MCP server instead; it is still a real remote upstream, so the no-mock, no-fixture rule above holds unchanged - ## Lay the pattern down in a class Keep the cases for one feature inside a class so the file reads as a spec for how that feature behaves in production. The class name says what is under test; each method is one behavior. Think of it as documenting the contract, with the rough intent being diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 5c25b7f2a93..f9cd2a3f15f 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -524,17 +524,6 @@ class TestOpenAIFiles: "llm.files.openai.list.nonstream.works", exercised_on=["files"], ) - @pytest.mark.skip( - reason=( - "LIT-4820 (https://linear.app/litellm-ai/issue/LIT-4820): GET /v1/files omits " - "newly uploaded files. The upload succeeds and " - "GET /v1/files/{id} returns the file, but it never appears in the listing; the " - "returned set is stable with its newest entry ~10h old, on both the managed " - "(/v1/files?model=) and provider-scoped (/openai/v1/files) routes. Skipped rather " - "than weakened because the assertion below is the correct contract. Remove this " - "marker when LIT-4820 is fixed; do not relax the assertion to make it pass." - ) - ) def test_uploaded_file_appears_in_list( self, client: BatchClient, resources: ResourceManager, batch_deployments: None ) -> None: diff --git a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py b/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py deleted file mode 100644 index 01e94f7b86f..00000000000 --- a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py +++ /dev/null @@ -1,197 +0,0 @@ -"""On-demand e2e: a chat completion drives a gateway-managed OAuth MCP server. - -The real end-user flow for MCP over an OAuth server: a user registers a Linear -authorization_code server, authorizes it once so the gateway stores their -upstream token, then sends a normal /chat/completions request with the Linear -MCP attached. The gateway resolves the user from the LiteLLM key, lists Linear's -tools with the stored per-user token, lets the model call one, executes it -upstream with that token, and returns the answer. This is proven against the -real Linear MCP server (mcp.linear.app) and a real Anthropic model, once per -documented ingress header (x-litellm-api-key and Authorization). - -The authorize dance is seeded through the mcp SDK's OAuthClientProvider; the one -step Linear cannot auto-approve is the human consent, so it is captured once out -of band (mcp/linear_session_capture.py) into a saved browser session and a -headless Chromium clicks Approve every run. The test therefore skips unless -E2E_LINEAR_STORAGE_STATE points at that session, so it never runs on the per-PR -CI path; it is a nightly/on-demand real-server smoke test. - -Fail-before-fix: without the stored per-user token the gateway lists no Linear -tools, so mcp_list_tools comes back empty, nothing is called, and the -assertions fail; a served, called, non-empty Linear tool proves the gateway -pulled and used the user's token. -""" - -from __future__ import annotations - -import os - -import pytest - -from e2e_config import CHEAP_ANTHROPIC_MODEL, LINEAR_MCP_URL, LINEAR_STORAGE_STATE, unique_marker -from e2e_http import AuthHeaders -from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, McpServerCreateBody, ObjectPermission -from proxy_client import ProxyClient - -pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`") -pytest.importorskip( - "playwright.async_api", - reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`", -) - -from oauth_chat_client import ChatMcpClient, build_chat_client # noqa: E402 # imports follow the importorskip guards - -pytestmark = [ - pytest.mark.e2e, - pytest.mark.skipif( - not LINEAR_STORAGE_STATE or not os.path.exists(LINEAR_STORAGE_STATE), - reason="set E2E_LINEAR_STORAGE_STATE to a Linear session captured via mcp/linear_session_capture.py", - ), -] - -# Pinned from a live dance during verification (never guessed); the gateway -# prefixes every upstream tool name with the server alias. list_teams is a -# read-only Linear tool that takes no arguments and returns the caller's teams. -LINEAR_READONLY_TOOL = "list_teams" -LINEAR_PROMPT = "Use the list_teams tool to list my Linear teams, then reply with the name of one of them." - - -@pytest.fixture(scope="session") -def chat_client(proxy: ProxyClient) -> ChatMcpClient: - return build_chat_client(proxy) - - -class TestMcpChatCompletionOauth: - """A scoped internal-user key on a real Linear authorization_code server, - used through /chat/completions once per ingress header: the gateway pulls - the user's stored upstream token, lists and executes Linear's tools during - the completion, and returns the answer.""" - - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - def test_chat_completion_uses_linear_with_x_litellm_api_key_header( - self, chat_client: ChatMcpClient, resources: ResourceManager - ) -> None: - marker = unique_marker() - alias = f"e2elinear{marker}" - created = chat_client.create_server( - McpServerCreateBody( - alias=alias, - url=LINEAR_MCP_URL, - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - ) - ) - resources.defer(lambda: chat_client.delete_server(created.server_id)) - - stored = chat_client.server_info(created.server_id) - assert stored.auth_type == "oauth2" - assert stored.oauth2_flow == "authorization_code" - assert stored.allow_all_keys is False - - key = chat_client.proxy.generate_key( - KeyGenerateBody( - user_id="e2e-test-user", - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - resources.defer(lambda: chat_client.proxy.delete_key(key)) - - seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) - assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( - f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" - ) - - response = chat_client.chat_with_mcp( - AuthHeaders.model_validate({"x-litellm-api-key": f"Bearer {key}"}), - ChatBody( - model=CHEAP_ANTHROPIC_MODEL, - messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], - tools=[ - McpChatTool( - server_url=f"litellm_proxy/mcp/{alias}", - server_label=alias, - require_approval="never", - ) - ], - ), - ) - - message = response.choices[0].message - assert message is not None and message.content, f"completion returned no answer: {response}" - meta = message.provider_specific_fields - assert meta is not None, f"no MCP metadata on the completion: {response}" - listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} - assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( - f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" - ) - results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] - assert results and results[0].result, ( - f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" - ) - - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - def test_chat_completion_uses_linear_with_authorization_bearer_header( - self, chat_client: ChatMcpClient, resources: ResourceManager - ) -> None: - marker = unique_marker() - alias = f"e2elinear{marker}" - created = chat_client.create_server( - McpServerCreateBody( - alias=alias, - url=LINEAR_MCP_URL, - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - ) - ) - resources.defer(lambda: chat_client.delete_server(created.server_id)) - - stored = chat_client.server_info(created.server_id) - assert stored.auth_type == "oauth2" - assert stored.oauth2_flow == "authorization_code" - assert stored.allow_all_keys is False - - key = chat_client.proxy.generate_key( - KeyGenerateBody( - user_id="e2e-test-user", - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - resources.defer(lambda: chat_client.proxy.delete_key(key)) - - seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) - assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( - f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" - ) - - response = chat_client.chat_with_mcp( - AuthHeaders.model_validate({"authorization": f"Bearer {key}"}), - ChatBody( - model=CHEAP_ANTHROPIC_MODEL, - messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], - tools=[ - McpChatTool( - server_url=f"litellm_proxy/mcp/{alias}", - server_label=alias, - require_approval="never", - ) - ], - ), - ) - - message = response.choices[0].message - assert message is not None and message.content, f"completion returned no answer: {response}" - meta = message.provider_specific_fields - assert meta is not None, f"no MCP metadata on the completion: {response}" - listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} - assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( - f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" - ) - results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] - assert results and results[0].result, ( - f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" - ) diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 2580197d6d2..ec6ae286351 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -443,3 +443,103 @@ async def test_store_unified_file_id_is_idempotent_via_upsert(): assert upsert_data["create"]["unified_file_id"] == file_id assert json.loads(upsert_data["create"]["model_mappings"]) == model_mappings assert json.loads(upsert_data["update"]["model_mappings"]) == model_mappings + + +def _make_managed_file_row( + *, + unified_file_id: str, + provider_ids: list[str], + file_object: object, + created_by: str = "test-user", + team_id: str | None = None, +): + row = MagicMock() + row.unified_file_id = unified_file_id + row.flat_model_file_ids = provider_ids + row.file_object = file_object + row.created_by = created_by + row.team_id = team_id + return row + + +@pytest.mark.asyncio +async def test_get_user_created_file_ids_skips_rows_without_file_object(): + """A managed row with file_object=None must not take down GET /v1/files.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + valid = _make_managed_file_row( + unified_file_id="unified-ok", + provider_ids=["file-ok"], + file_object=_make_file_object("file-ok").model_dump(), + ) + null_row = _make_managed_file_row( + unified_file_id="unified-null", + provider_ids=["file-null"], + file_object=None, + ) + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[valid, null_row] + ) + managed = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), prisma_client=mock_prisma + ) + + result = await managed.get_user_created_file_ids( + user_api_key_dict=_make_user_api_key_dict(), + model_object_ids=["file-ok", "file-null"], + ) + + assert [f.id for f in result] == ["file-ok"] + + +@pytest.mark.asyncio +async def test_filter_listed_provider_files_keeps_untracked_uploads(): + """LIT-4820: provider-scoped uploads that never entered the managed table + must still appear on GET /v1/files. The old filter replaced the page with + only managed rows owned by the caller, so raw uploads vanished.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + raw_upload = _make_file_object("file-raw-new") + other_user_managed = _make_file_object("file-other-user") + own_managed_provider = _make_file_object("file-own-managed") + own_managed_unified = _make_file_object("unified-own") + + other_row = _make_managed_file_row( + unified_file_id="unified-other", + provider_ids=["file-other-user"], + file_object=other_user_managed.model_dump(), + created_by="other-user", + ) + own_row = _make_managed_file_row( + unified_file_id="unified-own", + provider_ids=["file-own-managed"], + file_object=own_managed_unified.model_dump(), + created_by="test-user", + ) + + mock_prisma = MagicMock() + + async def _find_many(*, where): + if "created_by" in where or "OR" in where: + return [own_row] + return [other_row, own_row] + + mock_prisma.db.litellm_managedfiletable.find_many = AsyncMock(side_effect=_find_many) + managed = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), prisma_client=mock_prisma + ) + + kept = await managed._filter_listed_provider_files( + provider_files=[raw_upload, other_user_managed, own_managed_provider], + user_api_key_dict=_make_user_api_key_dict(), + ) + + kept_ids = [f.id for f in kept] + assert "file-raw-new" in kept_ids + assert "file-other-user" not in kept_ids + assert "unified-own" in kept_ids