From 06cf0ecfd835f307e83d79d4cdd77bce3974117c Mon Sep 17 00:00:00 2001
From: Julio Quinteros Pro <jquinter@gmail.com>
Date: Tue, 10 Feb 2026 23:04:27 -0300
Subject: [PATCH] =?UTF-8?q?fix:=20address=20review=20feedback=20=E2=80=94?=
 =?UTF-8?q?=20check=20success=5Fcallback=20+=20move=20imports=20to=20top?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Also check `litellm_settings.success_callback` for prometheus (not just `callbacks`)
- Move atexit, shutil, tempfile imports to module-level per CLAUDE.md style guide
- Add test for prometheus detection via success_callback

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
---
 litellm/proxy/proxy_cli.py                 | 20 +++++---
 tests/test_litellm/proxy/test_proxy_cli.py | 55 +++++++++++++++++++++-
 2 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py
index ae4dfca6c76..7a9dd659da9 100644
--- a/litellm/proxy/proxy_cli.py
+++ b/litellm/proxy/proxy_cli.py
@@ -1,10 +1,13 @@
 # ruff: noqa: T201
+import atexit
 import importlib
 import json
 import os
 import random
+import shutil
 import subprocess
 import sys
+import tempfile
 import urllib.parse as urlparse
 from typing import TYPE_CHECKING, Any, Optional, Union
 
@@ -691,15 +694,20 @@ def run_server(  # noqa: PLR0915
                 if litellm_settings
                 else []
             )
+            _success_callbacks = (
+                litellm_settings.get("success_callback", [])
+                if litellm_settings
+                else []
+            )
+            _has_prometheus = (
+                "prometheus" in _callbacks
+                or "prometheus" in _success_callbacks
+            )
             if (
                 num_workers > 1
-                and "prometheus" in _callbacks
+                and _has_prometheus
                 and "PROMETHEUS_MULTIPROC_DIR" not in os.environ
             ):
-                import atexit
-                import shutil
-                import tempfile
-
                 _prom_dir = tempfile.mkdtemp(prefix="litellm_prometheus_")
                 os.environ["PROMETHEUS_MULTIPROC_DIR"] = _prom_dir
                 print(  # noqa
@@ -713,7 +721,7 @@ def run_server(  # noqa: PLR0915
                 atexit.register(_cleanup_prometheus_dir)
             elif (
                 num_workers > 1
-                and "prometheus" in _callbacks
+                and _has_prometheus
                 and "PROMETHEUS_MULTIPROC_DIR" in os.environ
             ):
                 print(  # noqa
diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py
index 5c4f02c3eed..4d5c5217234 100644
--- a/tests/test_litellm/proxy/test_proxy_cli.py
+++ b/tests/test_litellm/proxy/test_proxy_cli.py
@@ -485,7 +485,7 @@ class TestPrometheusMultiprocessSetup:
     """Test cases for auto-configuring PROMETHEUS_MULTIPROC_DIR with multiple workers"""
 
     @staticmethod
-    def _write_config(tmp_path, callbacks=None):
+    def _write_config(tmp_path, callbacks=None, success_callback=None):
         """Write a minimal config yaml and return the path."""
         import yaml
 
@@ -495,6 +495,8 @@ class TestPrometheusMultiprocessSetup:
         }
         if callbacks is not None:
             config["litellm_settings"]["callbacks"] = callbacks
+        if success_callback is not None:
+            config["litellm_settings"]["success_callback"] = success_callback
         config_path = os.path.join(str(tmp_path), "config.yaml")
         with open(config_path, "w") as f:
             yaml.dump(config, f)
@@ -614,6 +616,57 @@ class TestPrometheusMultiprocessSetup:
                 os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
 
 
+    @patch("uvicorn.run")
+    def test_prometheus_multiproc_dir_auto_created_via_success_callback(
+        self, mock_uvicorn_run, tmp_path
+    ):
+        """When num_workers > 1 and prometheus is in success_callback, PROMETHEUS_MULTIPROC_DIR should be auto-set"""
+        import shutil
+
+        from click.testing import CliRunner
+
+        from litellm.proxy.proxy_cli import run_server
+
+        runner = CliRunner()
+        config_path = self._write_config(
+            tmp_path, success_callback=["prometheus"]
+        )
+
+        orig_val = os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
+        try:
+            with patch(
+                "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
+            ) as mock_get_args:
+                mock_get_args.return_value = {
+                    "app": "litellm.proxy.proxy_server:app",
+                    "host": "localhost",
+                    "port": 8000,
+                }
+
+                result = runner.invoke(
+                    run_server,
+                    [
+                        "--config",
+                        config_path,
+                        "--num_workers",
+                        "4",
+                        "--skip_server_startup",
+                    ],
+                )
+
+                assert result.exit_code == 0, result.output
+                assert "PROMETHEUS_MULTIPROC_DIR" in os.environ
+                prom_dir = os.environ["PROMETHEUS_MULTIPROC_DIR"]
+                assert "litellm_prometheus_" in prom_dir
+                assert os.path.isdir(prom_dir)
+        finally:
+            prom_dir = os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
+            if prom_dir and os.path.exists(prom_dir):
+                shutil.rmtree(prom_dir, ignore_errors=True)
+            if orig_val is not None:
+                os.environ["PROMETHEUS_MULTIPROC_DIR"] = orig_val
+
+
 class TestHealthAppFactory:
     """Test cases for the health app factory module"""