From 3596dee1447d58764aec8b8b45c2e69237b4b4e4 Mon Sep 17 00:00:00 2001 From: Rithvik Mysore Suresh Date: Fri, 31 Jul 2026 09:20:04 -0400 Subject: [PATCH 1/3] fix(managed-files): skip rows without file objects --- .../proxy/hooks/managed_files.py | 6 +++++- .../enterprise/proxy/test_managed_files_hook.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 8821736d0ff..b8c97dd2bb5 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -382,7 +382,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "flat_model_file_ids": {"hasSome": model_object_ids}, } ) - return [OpenAIFileObject(**file_object.file_object) for file_object in file_ids] + return [ + OpenAIFileObject(**file_object.file_object) + for file_object in file_ids + if file_object.file_object is not None + ] async def check_managed_file_id_access( self, data: Dict, user_api_key_dict: UserAPIKeyAuth 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..4a4aa7aa5ea 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -137,6 +137,23 @@ async def test_should_pass_credentials_to_afile_retrieve(): ) +@pytest.mark.asyncio +async def test_get_user_created_file_ids_skips_rows_without_file_object(): + managed_files = _make_managed_files_instance() + managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[ + MagicMock(file_object=_make_file_object().model_dump()), + MagicMock(file_object=None), + ] + ) + + files = await managed_files.get_user_created_file_ids( + _make_user_api_key_dict(), ["file-output-abc"] + ) + + assert [file.id for file in files] == ["file-output-abc"] + + @pytest.mark.asyncio async def test_should_fallback_when_no_router(): """ From 54fb717de19485dafd1647d865331ba05b8e3a23 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 5 Aug 2026 10:06:52 -0700 Subject: [PATCH 2/3] fix(router): redact fallback tracebacks at the call site and cover the sync deferred stream (#35843) Three follow-ups surfaced while merging current staging into this branch. `exc_info=True` at both fallback-failure log sites handed a live exception to the logging machinery. SecretRedactionFilter rewrites `record.exc_text`, but `record.exc_info` stays an exception object no filter can reach, so a handler that renders it itself (Datadog and OTel log bridges do) received the unredacted provider key. Both sites now pass `redact_string(traceback.format_exc())` as a `%s` arg, keeping staging's lazy-logging form. The existing test only asserted on `exc_text`, so it passed under the bug; it now renders `exc_info` the way a bridge handler would and covers every record the call emits. The eager deferred-stream fetch existed only on the async path. Vertex and Bedrock build the same `completion_stream=None` plus `make_call` wrapper on their sync branches, so `Router.completion(stream=True)` still surfaced the provider error on first iteration, outside `_completion`'s except block, and never reached the fallback chain. `_completion` now calls `fetch_sync_stream()` under the same guard `_acompletion` uses. The first of the three header-strip passes in the proxy error path was dead: only the custom-header update and the response-headers hook run before the second pass re-filters everything. Collapsed to one `safe_headers` binding. --- litellm/proxy/common_request_processing.py | 18 ++-- litellm/router.py | 16 +++- .../test_redact_string_in_error_paths.py | 58 +++++++++---- tests/test_litellm/test_router.py | 86 ++++++++++++++++++- 4 files changed, 146 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 21334507d02..0b9a2d5e4c0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2688,7 +2688,6 @@ class ProxyBaseLLMRequestProcessing: _response_headers: Final = getattr(_response, "headers", None) if _response_headers: headers = get_response_headers(dict(_response_headers)) - headers = {k: v for k, v in headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS} headers.update(custom_headers) # Call response headers hook for failure @@ -2704,16 +2703,15 @@ class ProxyBaseLLMRequestProcessing: except Exception: pass - headers = {k: v for k, v in headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS} + safe_headers: Final = {k: v for k, v in headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS} - self._apply_router_cooldown_retry_after(headers, e) + self._apply_router_cooldown_retry_after(safe_headers, e) if isinstance(e, ProxyException): - merged_headers = { - **e.headers, - **{k: v if isinstance(v, str) else str(v) for k, v in headers.items()}, + e.headers = { + **{k: v for k, v in e.headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS}, + **{k: v if isinstance(v, str) else str(v) for k, v in safe_headers.items()}, } - e.headers = {k: v for k, v in merged_headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS} raise e if isinstance(e, HTTPException): @@ -2730,7 +2728,7 @@ class ProxyBaseLLMRequestProcessing: param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), provider_specific_fields=merged_fields, - headers=headers, + headers=safe_headers, ) elif isinstance(e, httpx.HTTPStatusError): # Handle httpx.HTTPStatusError - extract actual error from response @@ -2756,7 +2754,7 @@ class ProxyBaseLLMRequestProcessing: type="invalid_request_error", param=None, code=status.HTTP_400_BAD_REQUEST, - headers=headers, + headers=safe_headers, ) # Extract status_code from the exception if it carries one. # Provider exceptions (NotFoundError, BadRequestError, GeminiError, @@ -2775,7 +2773,7 @@ class ProxyBaseLLMRequestProcessing: openai_code=getattr(e, "code", None), code=_code, provider_specific_fields=getattr(e, "provider_specific_fields", None), - headers=headers, + headers=safe_headers, ) ######################################################### diff --git a/litellm/router.py b/litellm/router.py index a0eb6e91c0c..39d5080a33d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1832,6 +1832,13 @@ class Router: llm_provider="", ) + if ( + isinstance(response, CustomStreamWrapper) + and response.completion_stream is None + and response.make_call is not None + ): + response.fetch_sync_stream() + # Wrap streaming responses so MidStreamFallbackError (raised # during iteration) triggers the Router's fallback chain. if isinstance(response, CustomStreamWrapper): @@ -6120,7 +6127,8 @@ class Router: """ Common utilities for async_function_with_fallbacks """ - verbose_router_logger.debug("Traceback", exc_info=True) + if verbose_router_logger.isEnabledFor(logging.DEBUG): + verbose_router_logger.debug("Traceback%s", redact_string(traceback.format_exc())) original_exception: Final = e fallback_model_group = None original_model_group: Final[str | None] = kwargs.get("model") @@ -6336,17 +6344,17 @@ class Router: except Exception as new_exception: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) fallback_failure_exception_str = redact_string(str(new_exception)) - cooldown_info = await _async_get_cooldown_deployments_with_debug_info( + cooldown_info: Final = await _async_get_cooldown_deployments_with_debug_info( litellm_router_instance=self, parent_otel_span=parent_otel_span, ) verbose_router_logger.error( "litellm.router.py::async_function_with_fallbacks() - " - "Error occurred while trying to do fallbacks - %s\n" + "Error occurred while trying to do fallbacks - %s\n%s\n" "Debug Information:\nCooldown Deployments=%s", fallback_failure_exception_str, + redact_string(traceback.format_exc()), cooldown_info, - exc_info=True, ) if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors: diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/test_litellm/test_redact_string_in_error_paths.py index 4a624017ea2..d01a9da6617 100644 --- a/tests/test_litellm/test_redact_string_in_error_paths.py +++ b/tests/test_litellm/test_redact_string_in_error_paths.py @@ -193,13 +193,22 @@ class TestProxyStreamingDataGeneratorRedaction: class TestRouterFallbackFailureTracebackRedaction: - """Test the fallback-failure error log in router.py's - async_function_with_fallbacks_common_utils. A prior version passed exc_info=True - alongside an already-redacted message, which bypasses redact_string() entirely - since the stdlib logging module renders exc_info separately from the message.""" + """Test the fallback-failure logs in router.py's + async_function_with_fallbacks_common_utils. Both call sites must redact the + traceback at the call site with redact_string() rather than hand a live + exception to exc_info=True. SecretRedactionFilter rewrites record.exc_text, + but record.exc_info stays an exception object no filter can rewrite, so any + handler that renders exc_info itself (Datadog and OTel log bridges do) would + receive the unredacted secret.""" @pytest.mark.asyncio async def test_fallback_failure_does_not_leak_secret_via_exc_info(self, caplog): + """The helper is driven from inside an `except` block because that is the only + way production reaches it, and the entry-point debug log takes its traceback + from the active exception. With no exception in flight sys.exc_info() is empty, + so an exc_info=True regression there would degrade to (None, None, None) and + this test would pass against it. + """ import litellm router = litellm.Router( @@ -221,24 +230,39 @@ class TestRouterFallbackFailureTracebackRedaction: "litellm.router.run_async_fallback", new=AsyncMock(side_effect=RuntimeError(f"boom api_key={secret}")), ): - with caplog.at_level(logging.ERROR, logger="LiteLLM Router"): - with pytest.raises(Exception): - await router.async_function_with_fallbacks_common_utils( - e=Exception("original failure"), - disable_fallbacks=False, - fallbacks=[{"gpt-3.5-turbo": ["claude-3-haiku"]}], - context_window_fallbacks=None, - content_policy_fallbacks=None, - model_group="gpt-3.5-turbo", - args=(), - kwargs={"model": "gpt-3.5-turbo"}, - ) + try: + raise ValueError(f"primary deployment failed api_key={secret}") + except ValueError as original_exception: + with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): + with pytest.raises(Exception): + await router.async_function_with_fallbacks_common_utils( + e=original_exception, + disable_fallbacks=False, + fallbacks=[{"gpt-3.5-turbo": ["claude-3-haiku"]}], + context_window_fallbacks=None, + content_policy_fallbacks=None, + model_group="gpt-3.5-turbo", + args=(), + kwargs={"model": "gpt-3.5-turbo"}, + ) + + debug_records = [r for r in caplog.records if r.levelno == logging.DEBUG] + assert debug_records, "expected the entry-point debug log, which carries the active traceback" error_records = [r for r in caplog.records if r.levelno == logging.ERROR] assert error_records, "expected an error log for the fallback failure" - for record in error_records: + assert any( + "Cooldown Deployments" in r.getMessage() for r in error_records + ), "expected the fallback-failure log, not an unrelated error" + + for record in caplog.records: assert secret not in record.getMessage() assert secret not in (record.exc_text or "") + rendered_exc_info = "".join(traceback.format_exception(*record.exc_info)) if record.exc_info else "" + assert secret not in rendered_exc_info, ( + f"{record.levelname} record passed a live exception to exc_info; " + "no logging filter can redact record.exc_info" + ) def _make_mock_ingest_options(): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 33aab1cf708..8f35597768f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -6343,8 +6343,9 @@ async def test_acompletion_deferred_stream_skipped_when_stream_already_set(): return yield + noop_stream = noop_aiter() already_set_wrapper = CustomStreamWrapper( - completion_stream=noop_aiter(), + completion_stream=noop_stream, model="openai/gpt-4o", logging_obj=logging_obj, custom_llm_provider="openai", @@ -6376,6 +6377,89 @@ async def test_acompletion_deferred_stream_skipped_when_stream_already_set(): ) assert result is not None, "should return a streaming wrapper without errors" + assert already_set_wrapper.completion_stream is noop_stream, "completion_stream must not be re-fetched" + await noop_stream.aclose() + + +def test_completion_deferred_stream_error_propagates_through_completion(): + """Regression: the sync router path needs the same eager fetch as the async one. + + A deferred-stream CustomStreamWrapper hands back a wrapper whose HTTP call has + not happened yet, so without fetch_sync_stream() the provider error surfaces on + first iteration, outside _completion's except block. The deployment is then never + marked failed and function_with_fallbacks never sees the error. + """ + import litellm as _litellm + + rate_limit_err = _litellm.RateLimitError( + message="Resource exhausted", + llm_provider="vertex_ai", + model="gemini-2.0-flash", + ) + make_call_invocations = [] + + def failing_make_call(**kwargs): + make_call_invocations.append(kwargs) + raise rate_limit_err + + router = _make_router_with_vertex_and_fallback() + deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) + + with patch("litellm.completion", return_value=deferred_wrapper): + with pytest.raises(_litellm.RateLimitError): + router._completion( + model="vertex_ai/gemini-2.0-flash", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + specific_deployment=router.model_list[0], + ) + + assert len(make_call_invocations) == 1, ( + "the deferred HTTP call must run inside _completion's try block; " + "without the eager fetch_sync_stream() fix it is deferred to first iteration" + ) + + +def test_completion_deferred_stream_skipped_when_stream_already_set(): + """A non-deferred sync provider already has completion_stream populated, so the + eager fetch must be skipped and make_call left untouched. + """ + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + def would_fail(**kwargs): + raise RuntimeError("should not be called") + + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}} + already_set_stream = iter([]) + + already_set_wrapper = CustomStreamWrapper( + completion_stream=already_set_stream, + model="openai/gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + make_call=would_fail, + ) + + router = litellm.Router( + model_list=[ + { + "model_name": "my-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + } + ], + ) + + with patch("litellm.completion", return_value=already_set_wrapper): + result = router._completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + specific_deployment=router.model_list[0], + ) + + assert result is not None, "should return a streaming wrapper without errors" + assert already_set_wrapper.completion_stream is already_set_stream, "completion_stream must not be re-fetched" class TestAdvisorSubCallCooldown: From 0659738b3e3e4c0064a94ffc3fe92df07a70b457 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 5 Aug 2026 10:14:46 -0700 Subject: [PATCH 3/3] fix(migrations): recover from an interrupted Prisma toolchain install (#35832) The Prisma CLI is a Node program that installs a private Node runtime on its first invocation. That one-time install shared the 60s budget that bounds each migration command, so on a slow or cold machine it was killed before it could finish. Prisma then decides whether to reinstall by testing the cache directory for existence alone, and a killed install leaves that directory behind, so every later attempt skipped the install and failed on a node binary that was never written. The existing four-attempt retry loop could not help: each attempt hit the same missing binary, which turned a slow start into a container that never migrated again. Migrations now prepare the toolchain as its own step under its own budget, and a cache directory that exists without a node binary is deleted first so an interrupted install reinstalls instead of persisting. Both budgets are overridable, LITELLM_PRISMA_BOOTSTRAP_TIMEOUT for the install and LITELLM_PRISMA_COMMAND_TIMEOUT for each Prisma command, and every previously hardcoded timeout now goes through one helper rather than thirteen literals. The per-command default stays at 60s. An override is only honoured when it parses as a finite positive number. Infinity and NaN parse as floats and survive a plain positivity check, and subprocess treats either as no deadline at all, so a value like `inf` or a fat-fingered `1e400` would have silently disabled the timeout it was meant to configure. --- .../litellm_proxy_extras/prisma_toolchain.py | 177 ++++++++++++++ .../litellm_proxy_extras/replica_identity.py | 3 +- .../litellm_proxy_extras/utils.py | 33 +-- .../test_prisma_toolchain.py | 221 ++++++++++++++++++ 4 files changed, 420 insertions(+), 14 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py create mode 100644 tests/proxy_migration_tests/test_prisma_toolchain.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py new file mode 100644 index 00000000000..76c22d0f8ce --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -0,0 +1,177 @@ +"""Prepare the Node toolchain the Prisma CLI needs, separately from migrations. + +The Prisma CLI is a Node program. The first invocation inside a fresh +container installs a private Node runtime and npm-installs the CLI itself, +which can take minutes on a cold or slow machine. Sharing one timeout between +that one-time bootstrap and the migration commands makes a slow bootstrap +indistinguishable from a slow migration, so the bootstrap gets killed long +before it can finish. + +A killed bootstrap does not correct itself. The installer leaves its cache +directory behind, and Prisma decides whether to install by testing that +directory for existence alone, so every later attempt skips the install and +then fails on a Node binary that was never written. Deleting a cache directory +that exists without a Node binary is what turns a killed bootstrap back into a +recoverable one. + +Both budgets are overridable so an operator can widen them without a release: +``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install and +``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every individual Prisma command. +""" + +import math +import os +import shutil +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +from litellm_proxy_extras._logging import logger + +try: + from prisma import config as prisma_config +except ImportError: + prisma_config = None + +PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT" +PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT" +NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR" + +DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0 +DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0 + +BOOTSTRAP_ARG = "--version" + + +@dataclass(frozen=True) +class ToolchainBootstrap: + """Outcome of preparing the Prisma toolchain.""" + + healed_incomplete_cache: bool + ready: bool + + +def _timeout_from_env(env_var: str, default: float) -> float: + raw = os.getenv(env_var) + if raw is None: + return default + try: + seconds = float(raw) + except ValueError: + logger.warning( + "%s=%r is not a number, falling back to %ss", env_var, raw, default + ) + return default + if not math.isfinite(seconds) or seconds <= 0: + logger.warning( + "%s=%r is not a finite positive number, falling back to %ss", + env_var, + raw, + default, + ) + return default + return seconds + + +def prisma_command_timeout() -> float: + """Seconds any single Prisma command may run for.""" + return _timeout_from_env( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_COMMAND_TIMEOUT + ) + + +def prisma_bootstrap_timeout() -> float: + """Seconds the one-time Node toolchain install may run for.""" + return _timeout_from_env( + PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT + ) + + +def nodeenv_cache_dir() -> Optional[Path]: + """Where Prisma installs its private Node runtime, or None if unknowable.""" + override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR) + if override: + return Path(override).absolute() + if prisma_config is not None: + try: + return Path(prisma_config.nodeenv_cache_dir).absolute() + except (OSError, ValueError) as e: + logger.warning("Could not read the Prisma nodeenv cache dir: %s", e) + try: + return Path.home() / ".cache" / "prisma-python" / "nodeenv" + except RuntimeError: + logger.warning( + "No resolvable home directory, cannot locate the Prisma nodeenv cache" + ) + return None + + +def node_binary_path(cache_dir: Path) -> Path: + """Path the Node binary occupies once the toolchain is fully installed.""" + if os.name == "nt": + return cache_dir / "Scripts" / "node.exe" + return cache_dir / "bin" / "node" + + +def heal_incomplete_nodeenv_cache() -> bool: + """Delete a nodeenv cache directory left without a Node binary. + + Returns True when a half-installed toolchain was removed, so the next + Prisma invocation reinstalls it instead of failing on a missing binary. + """ + cache_dir = nodeenv_cache_dir() + if cache_dir is None or not cache_dir.is_dir(): + return False + if node_binary_path(cache_dir).exists(): + return False + logger.warning( + "Node toolchain at %s has no %s, so a previous install was interrupted. " + "Removing it so it can be reinstalled.", + cache_dir, + node_binary_path(cache_dir).name, + ) + try: + shutil.rmtree(cache_dir) + except OSError as e: + logger.warning("Could not remove %s: %s", cache_dir, e) + return False + return True + + +def ensure_prisma_toolchain( + prisma_command: str, prisma_env: dict[str, str] +) -> ToolchainBootstrap: + """Install whatever the Prisma CLI needs to run, under its own timeout. + + Never raises. A toolchain that cannot be prepared is reported so the + caller can go on and let the real Prisma command produce the real error. + """ + healed = heal_incomplete_nodeenv_cache() + timeout = prisma_bootstrap_timeout() + logger.info("Preparing the Prisma CLI toolchain (timeout %ss)", timeout) + try: + subprocess.run( + [prisma_command, BOOTSTRAP_ARG], + timeout=timeout, + check=True, + capture_output=True, + text=True, + env=prisma_env, + ) + except subprocess.TimeoutExpired: + logger.warning( + "Preparing the Prisma CLI toolchain timed out after %ss. Raise %s " + "if this machine needs longer to install it.", + timeout, + PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, + ) + return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False) + except subprocess.CalledProcessError as e: + logger.warning("Preparing the Prisma CLI toolchain failed: %s", e.stderr) + return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False) + except OSError as e: + logger.warning("Could not run the Prisma CLI: %s", e) + return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False) + logger.info("Prisma CLI toolchain ready") + return ToolchainBootstrap(healed_incomplete_cache=healed, ready=True) diff --git a/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py b/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py index dc92e9dca6a..157d595404e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py +++ b/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py @@ -16,6 +16,7 @@ import tempfile from pathlib import Path from litellm_proxy_extras._logging import logger +from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL" @@ -75,7 +76,7 @@ def apply_replica_identity_full( "--schema", schema_path, ], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, text=True, diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index af822573322..5118865e43a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -14,6 +14,10 @@ from litellm_proxy_extras.replica_identity import ( REPLICA_IDENTITY_FULL_ENV_VAR, apply_replica_identity_full, ) +from litellm_proxy_extras.prisma_toolchain import ( + ensure_prisma_toolchain, + prisma_command_timeout, +) def str_to_bool(value: Optional[str]) -> bool: @@ -142,7 +146,7 @@ class ProxyExtrasDBManager: ], stdout=open(migration_file, "w"), check=True, - timeout=30, + timeout=prisma_command_timeout(), env=prisma_env, ) @@ -157,7 +161,7 @@ class ProxyExtrasDBManager: "0_init", ], check=True, - timeout=30, + timeout=prisma_command_timeout(), env=prisma_env, ) @@ -193,7 +197,7 @@ class ProxyExtrasDBManager: "--rolled-back", migration_name, ], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, env=prisma_env, @@ -205,7 +209,7 @@ class ProxyExtrasDBManager: prisma_env = _get_prisma_env() subprocess.run( [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, env=prisma_env, @@ -303,7 +307,7 @@ class ProxyExtrasDBManager: "--script", ], check=True, - timeout=60, + timeout=prisma_command_timeout(), stdout=f, env=_get_prisma_env(), ) @@ -335,7 +339,7 @@ class ProxyExtrasDBManager: "--schema", schema_path, ], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, text=True, @@ -364,7 +368,7 @@ class ProxyExtrasDBManager: "--schema", schema_path, ], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, text=True, @@ -393,7 +397,7 @@ class ProxyExtrasDBManager: "--applied", migration_name, ], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, text=True, @@ -530,7 +534,7 @@ class ProxyExtrasDBManager: try: subprocess.run( [_get_prisma_command(), "db", "push", "--accept-data-loss"], - timeout=60, + timeout=prisma_command_timeout(), check=True, env=_get_prisma_env(), ) @@ -555,7 +559,7 @@ class ProxyExtrasDBManager: try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, text=True, @@ -731,6 +735,9 @@ class ProxyExtrasDBManager: Returns: bool: True if setup was successful, False otherwise """ + ensure_prisma_toolchain( + prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env() + ) migrated = ProxyExtrasDBManager._run_migrations( use_migrate=use_migrate, use_v2_resolver=use_v2_resolver ) @@ -757,7 +764,7 @@ class ProxyExtrasDBManager: # Set migrations directory for Prisma result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, text=True, @@ -840,7 +847,7 @@ class ProxyExtrasDBManager: "--rolled-back", failed_migration, ], - timeout=60, + timeout=prisma_command_timeout(), check=True, capture_output=True, text=True, @@ -968,7 +975,7 @@ class ProxyExtrasDBManager: # Use prisma db push with increased timeout subprocess.run( [_get_prisma_command(), "db", "push", "--accept-data-loss"], - timeout=60, + timeout=prisma_command_timeout(), check=True, ) return True diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py new file mode 100644 index 00000000000..e68ec57028c --- /dev/null +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -0,0 +1,221 @@ +"""Migrations must survive a Node toolchain install that was killed mid-flight. + +The Prisma CLI installs a private Node runtime on its first invocation. If that +install is interrupted, the cache directory is left behind without a Node +binary and Prisma skips reinstalling it forever, so every later migration +attempt fails identically. These tests pin the two behaviours that keep a +container recoverable: an incomplete cache is deleted before Prisma is +invoked, and the install gets a budget of its own rather than sharing the one +that bounds each migration command. +""" + +import ast +import json +import os +import sys +import time +from pathlib import Path + +import pytest + +from litellm_proxy_extras.prisma_toolchain import ( + DEFAULT_PRISMA_COMMAND_TIMEOUT, + PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, + PRISMA_COMMAND_TIMEOUT_ENV_VAR, + ensure_prisma_toolchain, + heal_incomplete_nodeenv_cache, + node_binary_path, + prisma_bootstrap_timeout, + prisma_command_timeout, +) +from litellm_proxy_extras.utils import ProxyExtrasDBManager + +REPO_ROOT = Path(__file__).resolve().parents[2] +PROXY_EXTRAS = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" + +FAKE_PRISMA = """#!{python} +import json +import os +import pathlib +import sys +import time + +args = sys.argv[1:] +cache_dir = os.environ["PRISMA_NODEENV_CACHE_DIR"] +with pathlib.Path(os.environ["FAKE_PRISMA_LOG"]).open("a") as log: + log.write( + json.dumps({{"args": args, "cache_dir_present": os.path.isdir(cache_dir)}}) + + "\\n" + ) +time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0"))) +if args[:2] == ["migrate", "deploy"]: + print("No pending migrations to apply") +sys.exit(0) +""" + + +def _write_fake_prisma(tmp_path: Path) -> Path: + bin_dir = tmp_path / "fakebin" + bin_dir.mkdir() + script = bin_dir / "prisma" + script.write_text(FAKE_PRISMA.format(python=sys.executable)) + script.chmod(0o755) + return bin_dir + + +def _fake_prisma_calls(log_path: Path) -> list[dict[str, object]]: + if not log_path.exists(): + return [] + return [json.loads(line) for line in log_path.read_text().splitlines()] + + +@pytest.fixture +def toolchain_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, Path]: + """Point the toolchain at a scratch cache dir driven by a fake Prisma CLI.""" + cache_dir = tmp_path / "nodeenv" + log_path = tmp_path / "prisma-calls.jsonl" + bin_dir = _write_fake_prisma(tmp_path) + monkeypatch.setenv("PRISMA_NODEENV_CACHE_DIR", str(cache_dir)) + monkeypatch.setenv("FAKE_PRISMA_LOG", str(log_path)) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") + monkeypatch.delenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raising=False) + monkeypatch.delenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, raising=False) + return cache_dir, log_path + + +def _make_incomplete_cache(cache_dir: Path) -> None: + (cache_dir / "lib").mkdir(parents=True) + (cache_dir / "bin").mkdir() + + +def _make_complete_cache(cache_dir: Path) -> None: + node = node_binary_path(cache_dir) + node.parent.mkdir(parents=True) + node.write_text("") + + +def test_interrupted_toolchain_install_is_removed( + toolchain_env: tuple[Path, Path], +) -> None: + cache_dir, _ = toolchain_env + _make_incomplete_cache(cache_dir) + + assert heal_incomplete_nodeenv_cache() is True + assert not cache_dir.exists() + + +def test_installed_toolchain_is_left_alone(toolchain_env: tuple[Path, Path]) -> None: + cache_dir, _ = toolchain_env + _make_complete_cache(cache_dir) + + assert heal_incomplete_nodeenv_cache() is False + assert node_binary_path(cache_dir).exists() + + +def test_absent_toolchain_is_not_an_error(toolchain_env: tuple[Path, Path]) -> None: + cache_dir, _ = toolchain_env + + assert heal_incomplete_nodeenv_cache() is False + assert not cache_dir.exists() + + +def test_bootstrap_clears_the_cache_before_invoking_prisma( + toolchain_env: tuple[Path, Path], +) -> None: + cache_dir, log_path = toolchain_env + _make_incomplete_cache(cache_dir) + + result = ensure_prisma_toolchain( + prisma_command="prisma", prisma_env=dict(os.environ) + ) + + assert result.healed_incomplete_cache is True + assert result.ready is True + calls = _fake_prisma_calls(log_path) + assert len(calls) == 1 + assert calls[0]["cache_dir_present"] is False + + +def test_bootstrap_is_not_bounded_by_the_per_command_timeout( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + _, log_path = toolchain_env + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_SLEEP", "3") + + result = ensure_prisma_toolchain( + prisma_command="prisma", prisma_env=dict(os.environ) + ) + + assert result.ready is True + assert len(_fake_prisma_calls(log_path)) == 1 + + +def test_bootstrap_stops_at_its_own_timeout( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_SLEEP", "30") + + started = time.monotonic() + result = ensure_prisma_toolchain( + prisma_command="prisma", prisma_env=dict(os.environ) + ) + elapsed = time.monotonic() - started + + assert result.ready is False + assert elapsed < 15 + + +def test_setup_database_prepares_the_toolchain_before_migrating( + toolchain_env: tuple[Path, Path], +) -> None: + cache_dir, log_path = toolchain_env + _make_incomplete_cache(cache_dir) + + assert ProxyExtrasDBManager.setup_database(use_migrate=True) is True + + calls = _fake_prisma_calls(log_path) + assert [call["args"] for call in calls][:2] == [ + ["--version"], + ["migrate", "deploy"], + ] + assert calls[0]["cache_dir_present"] is False + + +@pytest.mark.parametrize( + "raw", + ["", "0", "-5", "not-a-number", "nan", "inf", "-inf", "1e400"], +) +def test_unusable_timeout_override_falls_back_to_the_default( + raw: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """A non-finite override would silently disable the timeout it configures.""" + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raw) + + assert prisma_command_timeout() == DEFAULT_PRISMA_COMMAND_TIMEOUT + + +def test_timeout_overrides_are_independent(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "12") + monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, "900") + + assert prisma_command_timeout() == 12 + assert prisma_bootstrap_timeout() == 900 + + +@pytest.mark.parametrize("module", ["utils.py", "replica_identity.py"]) +def test_every_prisma_command_timeout_is_overridable(module: str) -> None: + tree = ast.parse((PROXY_EXTRAS / module).read_text()) + literals = [ + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.keyword) + and node.arg == "timeout" + and isinstance(node.value, ast.Constant) + ] + + assert literals == [], ( + f"{module} still hardcodes a Prisma timeout at lines {literals}; " + "route it through prisma_command_timeout() so it can be raised without a release" + )