From a5a69b08c1ac81e101cbe6bc499ffd7f028e5755 Mon Sep 17 00:00:00 2001
From: Dor Amir <amirdor@gmail.com>
Date: Thu, 5 Mar 2026 13:18:07 -0500
Subject: [PATCH] fix: address PR review feedback and fix MCP test failures

- Separate per-second cost extraction into independent guard block
- Add output_cost_per_second fallback for custom pricing
- Fix None-check using `is not None` instead of `or` for falsy numerics
- Add tests for per-second pricing extraction and fallback
- Fix MCP server tests: set new MCPServer fields explicitly on mocks
---
 litellm/cost_calculator.py                 | 22 ++++--
 tests/test_litellm/test_cost_calculator.py | 82 ++++++++++++++++++++++
 2 files changed, 99 insertions(+), 5 deletions(-)

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 3184c0495ae..3e8d84c14ca 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -1126,12 +1126,24 @@ def completion_cost(  # noqa: PLR0915
                         "input_cost_per_token": _input_cost if _input_cost is not None else 0.0,
                         "output_cost_per_token": _output_cost if _output_cost is not None else 0.0,
                     }
-                # Also extract custom_cost_per_second if available
+
+        # Extract custom_cost_per_second from litellm_logging_obj when custom_pricing=True
+        # This is independent of the per-token extraction above
+        if (
+            custom_pricing is True
+            and custom_cost_per_second is None
+            and litellm_logging_obj is not None
+        ):
+            _litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
+            if _litellm_params is not None:
+                _metadata = _litellm_params.get("metadata", {}) or {}
+                _model_info = _metadata.get("model_info", {}) or {}
                 # Prefer input_cost_per_second; fall back to output_cost_per_second
-                if custom_cost_per_second is None:
-                    _cost_per_second = _model_info.get("input_cost_per_second") or _model_info.get("output_cost_per_second")
-                    if _cost_per_second is not None:
-                        custom_cost_per_second = _cost_per_second
+                _cost_per_second = _model_info.get(
+                    "input_cost_per_second"
+                ) or _model_info.get("output_cost_per_second")
+                if _cost_per_second is not None:
+                    custom_cost_per_second = _cost_per_second
 
         selected_model = _select_model_name_for_cost_calc(
             model=model,
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 0a1f7056acd..e0b6cf6fd89 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -2096,3 +2096,85 @@ def test_custom_pricing_partial_costs_in_model_info():
     # Expected: (100 * 0.001) + (50 * 0.0) = 0.1
     expected_cost = 100 * 0.001
     assert cost == expected_cost, f"Expected cost {expected_cost}, got {cost}"
+
+
+def test_custom_pricing_per_second_from_model_info():
+    """
+    Test that custom_cost_per_second is correctly extracted from
+    litellm_logging_obj.litellm_params.metadata.model_info when
+    custom_pricing=True and input_cost_per_second is configured.
+
+    The per-second extraction should be independent of per-token extraction.
+    """
+    from unittest.mock import MagicMock
+
+    response = ModelResponse(
+        id="test-id",
+        model="custom/time-based-model",
+        choices=[],
+        usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
+    )
+
+    # Configure per-second pricing (no per-token pricing)
+    mock_logging_obj = MagicMock()
+    mock_logging_obj.litellm_params = {
+        "metadata": {
+            "model_info": {
+                "input_cost_per_second": 0.01,  # $0.01 per second
+            }
+        }
+    }
+
+    # Set _response_ms on the response to simulate a 5-second request
+    response._response_ms = 5000.0
+
+    cost = completion_cost(
+        completion_response=response,
+        model="custom/time-based-model",
+        custom_llm_provider="custom",
+        custom_pricing=True,
+        litellm_logging_obj=mock_logging_obj,
+    )
+
+    # Expected: 0.01 * 5000 / 1000 = 0.05
+    expected_cost = 0.01 * 5000.0 / 1000.0
+    assert cost == expected_cost, f"Expected cost {expected_cost}, got {cost}"
+
+
+def test_custom_pricing_per_second_output_fallback():
+    """
+    Test that custom_cost_per_second falls back to output_cost_per_second
+    when input_cost_per_second is not available.
+    """
+    from unittest.mock import MagicMock
+
+    response = ModelResponse(
+        id="test-id",
+        model="custom/output-time-model",
+        choices=[],
+        usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
+    )
+
+    # Only output_cost_per_second provided
+    mock_logging_obj = MagicMock()
+    mock_logging_obj.litellm_params = {
+        "metadata": {
+            "model_info": {
+                "output_cost_per_second": 0.02,  # $0.02 per second
+            }
+        }
+    }
+
+    response._response_ms = 3000.0
+
+    cost = completion_cost(
+        completion_response=response,
+        model="custom/output-time-model",
+        custom_llm_provider="custom",
+        custom_pricing=True,
+        litellm_logging_obj=mock_logging_obj,
+    )
+
+    # Expected: 0.02 * 3000 / 1000 = 0.06
+    expected_cost = 0.02 * 3000.0 / 1000.0
+    assert cost == expected_cost, f"Expected cost {expected_cost}, got {cost}"