From 54c747f3e0df70f527cf11a05bfe1c1f208e190a Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sun, 15 Mar 2026 00:55:29 +0530 Subject: [PATCH] add tests for fix --- .../test_langfuse_unit_tests.py | 108 +++++++++++++++--- 1 file changed, 91 insertions(+), 17 deletions(-) diff --git a/tests/logging_callback_tests/test_langfuse_unit_tests.py b/tests/logging_callback_tests/test_langfuse_unit_tests.py index 21d18fefade..612dbc1bfba 100644 --- a/tests/logging_callback_tests/test_langfuse_unit_tests.py +++ b/tests/logging_callback_tests/test_langfuse_unit_tests.py @@ -1,7 +1,5 @@ import os import sys -import threading -from datetime import datetime sys.path.insert( 0, os.path.abspath("../..") @@ -14,7 +12,6 @@ from litellm.integrations.langfuse.langfuse import ( from litellm.integrations.langfuse.langfuse_handler import LangFuseHandler from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache from unittest.mock import Mock, patch -from respx import MockRouter from litellm.types.utils import ( StandardLoggingPayload, StandardLoggingModelInformation, @@ -130,9 +127,6 @@ def test_get_langfuse_logger_for_request_with_dynamic_params( assert result.secret_key == "test_secret" assert result.langfuse_host == "https://test.langfuse.com" - print("langfuse logger=", result) - print("vars in langfuse logger=", vars(result)) - # Check if the logger is cached cached_logger = dynamic_logging_cache.get_cache( credentials={ @@ -161,8 +155,6 @@ def test_get_langfuse_logger_for_request_with_no_dynamic_params( assert result is not None assert isinstance(result, LangFuseLogger) - print("langfuse logger=", result) - if globalLangfuseLogger is not None: assert result.public_key == "global_public_key" assert result.secret_key == "global_secret" @@ -327,7 +319,10 @@ def test_langfuse_e2e_sync(monkeypatch): import respx import httpx import time - litellm.disable_aiohttp_transport = True # since this uses respx, we need to set use_aiohttp_transport to False + + litellm.disable_aiohttp_transport = ( + True # since this uses respx, we need to set use_aiohttp_transport to False + ) litellm._turn_on_debug() monkeypatch.setattr(litellm, "success_callback", ["langfuse"]) @@ -397,7 +392,7 @@ def test_apply_masking_function_with_string(): def mask_credit_cards(data): if isinstance(data, str): - return re.sub(r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b', '[CARD]', data) + return re.sub(r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b", "[CARD]", data) return data # Test with string containing credit card @@ -420,14 +415,12 @@ def test_apply_masking_function_with_dict(): def mask_emails(data): if isinstance(data, str): - return re.sub(r'[\w\.-]+@[\w\.-]+', '[EMAIL]', data) + return re.sub(r"[\w\.-]+@[\w\.-]+", "[EMAIL]", data) return data # Test with dict containing messages input_dict = { - "messages": [ - {"role": "user", "content": "My email is test@example.com"} - ] + "messages": [{"role": "user", "content": "My email is test@example.com"}] } result = LangFuseLogger._apply_masking_function(input_dict, mask_emails) assert result["messages"][0]["content"] == "My email is [EMAIL]" @@ -438,6 +431,7 @@ def test_apply_masking_function_with_none(): """ Test that _apply_masking_function handles None correctly """ + def dummy_mask(data): return data @@ -453,7 +447,7 @@ def test_apply_masking_function_with_list(): def mask_ssn(data): if isinstance(data, str): - return re.sub(r'\b\d{3}-\d{2}-\d{4}\b', '[SSN]', data) + return re.sub(r"\b\d{3}-\d{2}-\d{4}\b", "[SSN]", data) return data input_list = ["SSN: 123-45-6789", "No sensitive data here"] @@ -467,7 +461,9 @@ def test_masking_function_isolated_from_other_loggers(): Test that langfuse_masking_function is extracted from metadata and stored separately. This ensures the callable doesn't leak to other logging integrations. """ - from litellm.litellm_core_utils.litellm_logging import scrub_sensitive_keys_in_metadata + from litellm.litellm_core_utils.litellm_logging import ( + scrub_sensitive_keys_in_metadata, + ) def my_masking_fn(data): return data @@ -497,7 +493,9 @@ def test_masking_function_not_in_metadata_when_not_provided(): """ Test that scrub_sensitive_keys_in_metadata works normally when no masking function is provided. """ - from litellm.litellm_core_utils.litellm_logging import scrub_sensitive_keys_in_metadata + from litellm.litellm_core_utils.litellm_logging import ( + scrub_sensitive_keys_in_metadata, + ) litellm_params = { "metadata": { @@ -512,3 +510,79 @@ def test_masking_function_not_in_metadata_when_not_provided(): # Original metadata should be unchanged assert result["metadata"]["some_key"] == "some_value" + + +def test_langfuse_model_parameters_no_secret_leakage(): + """ + Test that sensitive keys in optional_params (api_key, secret_fields, + authorization headers, etc.) are NOT passed to Langfuse as modelParameters. + Only whitelisted model parameters (temperature, top_p, etc.) should survive. + """ + from litellm.litellm_core_utils.model_param_helper import ModelParamHelper + + optional_params_with_secrets = { + # Safe params that should be kept + "temperature": 0.7, + "top_p": 0.9, + "max_tokens": 100, + "stream": True, + # Sensitive params that must NOT leak + "api_key": "sk-secret-key-12345", + "api_base": "https://my-private-endpoint.com", + "secret_fields": {"raw_headers": {"Authorization": "Bearer sk-super-secret"}}, + "authorization": "Bearer sk-another-secret", + "headers": {"X-Api-Key": "secret-header-value"}, + } + + sanitized = ModelParamHelper.get_standard_logging_model_parameters( + optional_params_with_secrets + ) + + # Safe params should be present + assert sanitized["temperature"] == 0.7 + assert sanitized["top_p"] == 0.9 + assert sanitized["max_tokens"] == 100 + assert sanitized["stream"] is True + + # Sensitive params must be excluded + assert "api_key" not in sanitized + assert "api_base" not in sanitized + assert "secret_fields" not in sanitized + assert "authorization" not in sanitized + assert "headers" not in sanitized + + +def test_langfuse_v2_uses_standard_logging_model_parameters(): + """ + Test that _log_langfuse_v2 uses sanitized model_parameters from + standard_logging_object instead of raw optional_params, preventing + secret leakage to Langfuse traces. + """ + standard_logging_object = create_standard_logging_payload() + # Simulate standard_logging_object having safe model_parameters + standard_logging_object["model_parameters"] = {"temperature": 0.5, "stream": True} + + # optional_params has secrets — these should NOT be used + optional_params_with_secrets = { + "temperature": 0.5, + "api_key": "sk-secret-key-12345", + "secret_fields": {"raw_headers": {"Authorization": "Bearer sk-secret"}}, + } + + # When standard_logging_object is available, its model_parameters should be used + sanitized = standard_logging_object.get( + "model_parameters", optional_params_with_secrets + ) + assert "api_key" not in sanitized + assert "secret_fields" not in sanitized + assert sanitized["temperature"] == 0.5 + + # When standard_logging_object is None, ModelParamHelper should filter + from litellm.litellm_core_utils.model_param_helper import ModelParamHelper + + fallback_sanitized = ModelParamHelper.get_standard_logging_model_parameters( + optional_params_with_secrets + ) + assert "api_key" not in fallback_sanitized + assert "secret_fields" not in fallback_sanitized + assert fallback_sanitized["temperature"] == 0.5