From 8f242c42a17f361a0b25b0cbf6ebdc164ac3514a Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Mon, 9 Feb 2026 20:58:06 -0600 Subject: [PATCH 01/15] fix(batch_completion): submit all model futures before waiting (#20705) * fix(batch_completion): submit all model futures before waiting * test: add batch_completion all responses concurrency regression * fix(batch_completion): continue collecting responses on per-model failures * fix(batch_completion): handle empty and string models in all responses * test(batch_completion): avoid blocking wait in concurrency regression --- litellm/batch_completion/main.py | 32 ++++- ...t_batch_completion_models_all_responses.py | 118 ++++++++++++++++++ 2 files changed, 144 insertions(+), 6 deletions(-) create mode 100644 tests/litellm/test_batch_completion_models_all_responses.py diff --git a/litellm/batch_completion/main.py b/litellm/batch_completion/main.py index 7100fb004f..446e3f2f99 100644 --- a/litellm/batch_completion/main.py +++ b/litellm/batch_completion/main.py @@ -237,17 +237,37 @@ def batch_completion_models_all_responses(*args, **kwargs): if "model" in kwargs: kwargs.pop("model") if "models" in kwargs: - models = kwargs["models"] - kwargs.pop("models") + models = kwargs.pop("models") else: raise Exception("'models' param not in kwargs") + if isinstance(models, str): + models = [models] + elif isinstance(models, (list, tuple)): + models = list(models) + else: + raise TypeError("'models' must be a string or list of strings") + + if len(models) == 0: + return [] + responses = [] with concurrent.futures.ThreadPoolExecutor(max_workers=len(models)) as executor: - for idx, model in enumerate(models): - future = executor.submit(litellm.completion, *args, model=model, **kwargs) - if future.result() is not None: - responses.append(future.result()) + futures = [ + executor.submit(litellm.completion, *args, model=model, **kwargs) + for model in models + ] + + for future in futures: + try: + result = future.result() + if result is not None: + responses.append(result) + except Exception as e: + print_verbose( + f"batch_completion_models_all_responses: model request failed: {str(e)}" + ) + continue return responses diff --git a/tests/litellm/test_batch_completion_models_all_responses.py b/tests/litellm/test_batch_completion_models_all_responses.py new file mode 100644 index 0000000000..2e96ada03f --- /dev/null +++ b/tests/litellm/test_batch_completion_models_all_responses.py @@ -0,0 +1,118 @@ +import concurrent.futures + +import litellm +from litellm.batch_completion.main import batch_completion_models_all_responses + + +def test_batch_completion_models_all_responses_submits_before_waiting(monkeypatch): + """ + Regression test for issue #20704. + Ensures all model calls are submitted to the thread pool before waiting on results. + """ + models = ["model-a", "model-b", "model-c"] + called_models = [] + + class _AssertingFuture: + def __init__(self, result, executor, expected_submissions): + self._result = result + self._executor = executor + self._expected_submissions = expected_submissions + + def result(self): + if self._executor.submit_count != self._expected_submissions: + raise AssertionError("Not all model calls were submitted before waiting") + return self._result + + class _RecordingThreadPoolExecutor: + def __init__(self, max_workers, *args, **kwargs): + self.max_workers = max_workers + self.submit_count = 0 + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def submit(self, fn, *args, **kwargs): + self.submit_count += 1 + result = fn(*args, **kwargs) + return _AssertingFuture( + result=result, + executor=self, + expected_submissions=len(models), + ) + + def _mock_completion(*args, model, **kwargs): + called_models.append(model) + return {"model": model} + + monkeypatch.setattr(litellm, "completion", _mock_completion) + monkeypatch.setattr( + concurrent.futures, "ThreadPoolExecutor", _RecordingThreadPoolExecutor + ) + + responses = batch_completion_models_all_responses( + models=models, + messages=[{"role": "user", "content": "hello"}], + ) + + assert sorted(called_models) == sorted(models) + assert len(responses) == len(models) + assert sorted(response["model"] for response in responses) == sorted(models) + + +def test_batch_completion_models_all_responses_continues_on_model_error(monkeypatch): + models = ["model-a", "model-error", "model-b"] + + def _mock_completion(*args, model, **kwargs): + if model == "model-error": + raise RuntimeError("simulated model failure") + return {"model": model} + + monkeypatch.setattr(litellm, "completion", _mock_completion) + + responses = batch_completion_models_all_responses( + models=models, + messages=[{"role": "user", "content": "hello"}], + ) + + assert len(responses) == 2 + assert sorted(response["model"] for response in responses) == ["model-a", "model-b"] + + +def test_batch_completion_models_all_responses_returns_empty_for_empty_models(monkeypatch): + called = False + + def _mock_completion(*args, model, **kwargs): + nonlocal called + called = True + return {"model": model} + + monkeypatch.setattr(litellm, "completion", _mock_completion) + + responses = batch_completion_models_all_responses( + models=[], + messages=[{"role": "user", "content": "hello"}], + ) + + assert responses == [] + assert called is False + + +def test_batch_completion_models_all_responses_accepts_single_model_string(monkeypatch): + called_models = [] + + def _mock_completion(*args, model, **kwargs): + called_models.append(model) + return {"model": model} + + monkeypatch.setattr(litellm, "completion", _mock_completion) + + responses = batch_completion_models_all_responses( + models="model-a", + messages=[{"role": "user", "content": "hello"}], + ) + + assert called_models == ["model-a"] + assert responses == [{"model": "model-a"}] From f77eeba186c31413d317ee534ff02d41cedb1269 Mon Sep 17 00:00:00 2001 From: Tsachi Shushan Date: Tue, 10 Feb 2026 06:05:53 +0200 Subject: [PATCH 02/15] fix: redaction headers ignored when sent via proxy (#20740) * fix: redaction headers ignored when sent via proxy When requests go through the proxy, `litellm_params["litellm_metadata"]` is always set (even when `None`), so `get_metadata_variable_name_from_kwargs` always returns "litellm_metadata". The redaction code then reads `None` instead of the actual metadata dict that contains the headers. Add a fallback to read from `metadata` when `litellm_metadata` is not a dict, so `x-litellm-enable-message-redaction` and related headers work correctly in the proxy flow. Fixes #20739 Co-Authored-By: Claude Opus 4.6 * fix: normalize non-dict metadata after fallback in redact_messages After falling back from litellm_metadata to metadata, ensure the value is always a dict so .get("headers") never raises on None/non-dict inputs. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- litellm/litellm_core_utils/redact_messages.py | 9 +- .../test_redact_messages.py | 145 ++++++++++++++++++ 2 files changed, 152 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_redact_messages.py diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index aa763dc989..5d6d1fbc1c 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -140,9 +140,14 @@ def should_redact_message_logging(model_call_details: dict) -> bool: metadata_field = get_metadata_variable_name_from_kwargs(litellm_params) metadata = litellm_params.get(metadata_field, {}) - + if not isinstance(metadata, dict): + # Fall back: litellm_metadata was None, try metadata + metadata = litellm_params.get("metadata", {}) + if not isinstance(metadata, dict): + metadata = {} + # Get headers from the metadata - request_headers = metadata.get("headers", {}) if isinstance(metadata, dict) else {} + request_headers = metadata.get("headers", {}) # Check for headers that explicitly control redaction if request_headers and bool( diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py new file mode 100644 index 0000000000..d7df7823ae --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -0,0 +1,145 @@ +""" +Tests for litellm.litellm_core_utils.redact_messages.should_redact_message_logging + +Covers the proxy flow where headers arrive in litellm_params["metadata"]["headers"] +but litellm_params["litellm_metadata"] is None. +""" + +import pytest + +import litellm +from litellm.litellm_core_utils.redact_messages import should_redact_message_logging + + +@pytest.fixture(autouse=True) +def _reset_global_redaction(): + """Ensure the global setting is off for every test.""" + original = litellm.turn_off_message_logging + litellm.turn_off_message_logging = False + yield + litellm.turn_off_message_logging = original + + +def _make_model_call_details( + metadata_headers=None, + litellm_metadata=None, + metadata=None, + standard_callback_dynamic_params=None, +): + """Build a model_call_details dict that mimics real proxy/SDK flows.""" + litellm_params = {} + if metadata is not None: + litellm_params["metadata"] = metadata + elif metadata_headers is not None: + litellm_params["metadata"] = {"headers": metadata_headers} + else: + litellm_params["metadata"] = {} + + # get_litellm_params always sets this key (even when value is None) + litellm_params["litellm_metadata"] = litellm_metadata + + details = {"litellm_params": litellm_params} + if standard_callback_dynamic_params is not None: + details["standard_callback_dynamic_params"] = standard_callback_dynamic_params + return details + + +class TestShouldRedactMessageLogging: + """Unit tests for should_redact_message_logging().""" + + # ---- proxy flow: headers in metadata, litellm_metadata is None ---- + + def test_enable_redaction_via_x_header_proxy_flow(self): + """x-litellm-enable-message-redaction header should enable redaction + even when litellm_metadata is None (proxy path).""" + details = _make_model_call_details( + metadata_headers={"x-litellm-enable-message-redaction": "true"}, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is True + + def test_enable_redaction_via_old_header_proxy_flow(self): + """litellm-enable-message-redaction header should enable redaction + even when litellm_metadata is None (proxy path).""" + details = _make_model_call_details( + metadata_headers={"litellm-enable-message-redaction": "true"}, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is True + + def test_disable_redaction_via_header_proxy_flow(self): + """litellm-disable-message-redaction should suppress redaction + even when global setting is on, and litellm_metadata is None.""" + litellm.turn_off_message_logging = True + details = _make_model_call_details( + metadata_headers={"litellm-disable-message-redaction": "true"}, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is False + + # ---- SDK direct-call flow: headers in litellm_metadata ---- + + def test_enable_redaction_via_header_in_litellm_metadata(self): + """Headers inside litellm_metadata (SDK direct call) should work.""" + details = _make_model_call_details( + litellm_metadata={"headers": {"x-litellm-enable-message-redaction": "true"}}, + ) + assert should_redact_message_logging(details) is True + + # ---- no headers at all ---- + + def test_no_headers_defaults_to_global_off(self): + """Without headers, falls back to global setting (False).""" + details = _make_model_call_details( + metadata_headers=None, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is False + + def test_no_headers_global_on(self): + """Without headers, respects global turn_off_message_logging=True.""" + litellm.turn_off_message_logging = True + details = _make_model_call_details( + metadata_headers=None, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is True + + # ---- dynamic params take precedence ---- + + def test_dynamic_param_enables_redaction(self): + """Dynamic turn_off_message_logging=True should enable redaction.""" + details = _make_model_call_details( + metadata_headers={}, + litellm_metadata=None, + standard_callback_dynamic_params={"turn_off_message_logging": True}, + ) + assert should_redact_message_logging(details) is True + + def test_dynamic_param_false_overrides_header(self): + """Dynamic turn_off_message_logging=False should take precedence over enable header.""" + details = _make_model_call_details( + metadata_headers={"x-litellm-enable-message-redaction": "true"}, + litellm_metadata=None, + standard_callback_dynamic_params={"turn_off_message_logging": False}, + ) + assert should_redact_message_logging(details) is False + + # ---- non-dict metadata safety ---- + + def test_both_metadata_fields_none(self): + """When both litellm_metadata and metadata are None, should not raise.""" + details = _make_model_call_details( + metadata=None, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is False + + def test_both_metadata_fields_none_global_on(self): + """When both metadata fields are None but global is on, should still return True.""" + litellm.turn_off_message_logging = True + details = _make_model_call_details( + metadata=None, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is True From 6bedc7decdf33260616c600f6ccd86b803c33f05 Mon Sep 17 00:00:00 2001 From: The Mavik <179817126+themavik@users.noreply.github.com> Date: Tue, 10 Feb 2026 22:37:45 +0530 Subject: [PATCH 03/15] fix(proxy): return early instead of raising ValueError when standard_logging_payload is missing (#20851) * fix: Preserved nullable object fields by carrying schema properties * Fix: _convert_schema_types * Fix all mypy issues * Add alert about email notifications * fixing tests * extending timeout for long running tests * Text changes * [Feat] MCP Oauth2 Fixes - Add support for MCP M2M Oauth2 support (#20788) * add has_client_credentials * MCPOAuth2TokenCache * init MCP Oauth2 constants * MCPOAuth2TokenCache * resolve_mcp_auth * test fixes * docs fix * address greptile review: min TTL, env-configurable constants, tests, docs - Fix zero-TTL edge case: floor at MCP_OAUTH2_TOKEN_CACHE_MIN_TTL (10s) - Make all MCP OAuth2 constants env-configurable via os.getenv() - Move test file to follow 1:1 mapping convention (test_oauth2_token_cache.py) - Add MCP OAuth doc page (mcp_oauth.md) with M2M and PKCE sections - Update FAQ in mcp.md to reflect M2M support - Add E2E test script and config Co-Authored-By: Claude Opus 4.6 * fix mypy lint * fix oauth2 * remove old files * docs fix * address greptile comments * fix: atomic lock creation + validate JSON response shape - Use dict.setdefault() for atomic per-server lock creation - Add isinstance(body, dict) check before accessing token response fields Co-Authored-By: Claude Opus 4.6 * fix: replace asserts with proper guards, wrap HTTP errors with context - Replace `assert` statements with `if/raise ValueError` (asserts can be disabled with python -O in production) - Wrap `httpx.HTTPStatusError` to provide a clear error message with server_id and status code - Add tests for HTTP error and non-dict JSON response error paths - Remove unused imports Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 * [UI] M2M OAuth2 UI Flow (#20794) * add has_client_credentials * MCPOAuth2TokenCache * init MCP Oauth2 constants * MCPOAuth2TokenCache * resolve_mcp_auth * test fixes * docs fix * address greptile review: min TTL, env-configurable constants, tests, docs - Fix zero-TTL edge case: floor at MCP_OAUTH2_TOKEN_CACHE_MIN_TTL (10s) - Make all MCP OAuth2 constants env-configurable via os.getenv() - Move test file to follow 1:1 mapping convention (test_oauth2_token_cache.py) - Add MCP OAuth doc page (mcp_oauth.md) with M2M and PKCE sections - Update FAQ in mcp.md to reflect M2M support - Add E2E test script and config Co-Authored-By: Claude Opus 4.6 * fix mypy lint * fix oauth2 * ui feat fixes * test M2M * test fix * ui feats * ui fixes * ui fix client ID * fix: backend endpoints * docs fix * fixes greptile --------- Co-authored-by: Claude Opus 4.6 * [Fix] prevent shared backend model key from being polluted by per-deployment custom pricing (#20679) * bug: custom price override for models * added associated test * fix(mcp): resolve OAuth2 root endpoints returning "MCP server not found" (#20784) When MCP SDK hits root-level /register, /authorize, /token without server name prefix, auto-resolve to the single configured OAuth2 server. Also fix WWW-Authenticate header to use correct public URL behind reverse proxy. * Add support for langchain_aws via litellm passthrough * fix(proxy): return early instead of raising ValueError when standard_logging_payload is missing The `_PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event` hook raises `ValueError` when `standard_logging_payload` is `None`. This breaks non-standard call types (e.g. vLLM `/classify`) that do not populate the payload, and the resulting exception disrupts downstream success callbacks like Langfuse. Return early with a debug log instead, matching the existing pattern used for missing `user_api_key_model_max_budget`. Fixes #18986 --------- Co-authored-by: Sameer Kankute Co-authored-by: yuneng-jiang Co-authored-by: Ishaan Jaff Co-authored-by: Claude Opus 4.6 Co-authored-by: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Co-authored-by: michelligabriele --- litellm/proxy/hooks/model_max_budget_limiter.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 69c7e92d82..b8c073dd06 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -153,7 +153,10 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): "standard_logging_object", None ) if standard_logging_payload is None: - raise ValueError("standard_logging_payload is required") + verbose_proxy_logger.debug( + "Skipping _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: standard_logging_payload is None" + ) + return _litellm_params: dict = kwargs.get("litellm_params", {}) or {} _metadata: dict = _litellm_params.get("metadata", {}) or {} From 95a8c2550c9ab1dfdf9cccc388cd41d13e91e208 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 11 Feb 2026 01:21:29 +0530 Subject: [PATCH 04/15] fix: handles edge case which are blocked by Lock (#20451) --- litellm/router.py | 34 +++++++++++++++++++--------------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d9de7e7fc5..37fa3926b4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -58,7 +58,6 @@ from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_metadata_variable_name_from_kwargs, ) -from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer @@ -619,11 +618,12 @@ class Router: self.retry_policy = RetryPolicy(**retry_policy) elif isinstance(retry_policy, RetryPolicy): self.retry_policy = retry_policy - verbose_router_logger.info( - "\033[32mRouter Custom Retry Policy Set:\n{}\033[0m".format( - self.retry_policy.model_dump(exclude_none=True) + if self.retry_policy is not None: + verbose_router_logger.info( + "\033[32mRouter Custom Retry Policy Set:\n{}\033[0m".format( + self.retry_policy.model_dump(exclude_none=True) + ) ) - ) self.model_group_retry_policy: Optional[ Dict[str, RetryPolicy] @@ -636,11 +636,12 @@ class Router: elif isinstance(allowed_fails_policy, AllowedFailsPolicy): self.allowed_fails_policy = allowed_fails_policy - verbose_router_logger.info( - "\033[32mRouter Custom Allowed Fails Policy Set:\n{}\033[0m".format( - self.allowed_fails_policy.model_dump(exclude_none=True) + if self.allowed_fails_policy is not None: + verbose_router_logger.info( + "\033[32mRouter Custom Allowed Fails Policy Set:\n{}\033[0m".format( + self.allowed_fails_policy.model_dump(exclude_none=True) + ) ) - ) self.alerting_config: Optional[AlertingConfig] = alerting_config @@ -1269,13 +1270,16 @@ class Router: if silent_model is not None: # Mirroring traffic to a secondary model - # Use shared thread pool for background calls - executor.submit( - self._silent_experiment_completion, - silent_model, - messages, - **kwargs, + # Use threading.Thread (not ThreadPoolExecutor) - executor.submit() + # requires pickling args, which fails when kwargs contain unpicklable + # objects (e.g. _thread.RLock from OTEL spans, loggers) in deployment. + thread = threading.Thread( + target=self._silent_experiment_completion, + args=(silent_model, messages), + kwargs=kwargs, + daemon=True, ) + thread.start() self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) kwargs.pop("silent_model", None) # Ensure it's not in kwargs either From 40aeda7a0ae2a00c32dbb27933e1eba55a6879fb Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 9 Feb 2026 22:28:27 -0500 Subject: [PATCH 05/15] fix(arize-phoenix): dynamic project naming from metadata + guardrails on image generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Allow per-request Phoenix project name via `metadata.phoenix_project_name`, falling back to PHOENIX_PROJECT_NAME env var - Add missing `post_call_success_hook` to `/images/generations` endpoint so guardrails and OTEL tracing apply to image generation 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- litellm/integrations/arize/arize_phoenix.py | 46 +++++++++++++++++---- litellm/proxy/image_endpoints/endpoints.py | 6 +++ 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index cd345a7f76..759dccdd8a 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -32,16 +32,48 @@ class ArizePhoenixLogger(OpenTelemetry): @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): + from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute + _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) - - # Set project name on the span for all traces to go to custom Phoenix projects - config = ArizePhoenixLogger.get_arize_phoenix_config() - if config.project_name: - from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute - safe_set_attribute(span, "openinference.project.name", config.project_name) - + + # Dynamic project name: check metadata first, then fall back to env var config + dynamic_project_name = ArizePhoenixLogger._get_dynamic_project_name(kwargs) + if dynamic_project_name: + safe_set_attribute(span, "openinference.project.name", dynamic_project_name) + else: + # Fall back to static config from env var + config = ArizePhoenixLogger.get_arize_phoenix_config() + if config.project_name: + safe_set_attribute(span, "openinference.project.name", config.project_name) + return + @staticmethod + def _get_dynamic_project_name(kwargs) -> Optional[str]: + """ + Retrieve dynamic Phoenix project name from request metadata. + + Users can set `metadata.phoenix_project_name` in their request to route + traces to different Phoenix projects dynamically. + """ + standard_logging_payload = kwargs.get("standard_logging_object") + if standard_logging_payload is not None: + metadata = standard_logging_payload.get("metadata") + if isinstance(metadata, dict): + project_name = metadata.get("phoenix_project_name") + if project_name: + return str(project_name) + + # Also check litellm_params.metadata for SDK usage + litellm_params = kwargs.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + if isinstance(metadata, dict): + project_name = metadata.get("phoenix_project_name") + if project_name: + return str(project_name) + + return None + @staticmethod def get_arize_phoenix_config() -> ArizePhoenixConfig: """ diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 4a2c05f859..4a8eb8e741 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -144,6 +144,12 @@ async def image_generation( litellm_call_id=data.get("litellm_call_id", ""), status="success" ) ) + + ### CALL HOOKS ### - modify outgoing data (guardrails, otel, etc.) + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + ### RESPONSE HEADERS ### hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = hidden_params.get("model_id", None) or "" From a422e8b9c9aa39ea1ee014323e3bb419903331f3 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 10 Feb 2026 00:29:50 -0500 Subject: [PATCH 06/15] fix(arize): allow OTEL and Arize Phoenix/Arize tracing to coexist in parallel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Arize Phoenix and Arize loggers now create dedicated TracerProviders instead of fighting over the global singleton, and the otel callback dedup check no longer incorrectly matches Arize subclasses. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- litellm/integrations/arize/arize.py | 35 ++++ litellm/integrations/arize/arize_phoenix.py | 49 +++++ litellm/litellm_core_utils/litellm_logging.py | 2 +- .../arize/test_arize_otel_coexistence.py | 171 ++++++++++++++++++ 4 files changed, 256 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 9c2f0d95d4..fe2f9f41f1 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -28,6 +28,41 @@ else: class ArizeLogger(OpenTelemetry): + """ + Arize logger that sends traces to an Arize endpoint. + + Creates its own dedicated TracerProvider so it can coexist with the + generic ``otel`` callback (or any other OTEL-based integration) without + fighting over the global ``opentelemetry.trace`` TracerProvider singleton. + """ + + def _init_tracing(self, tracer_provider): + """ + Override to always create a *private* TracerProvider for Arize. + + See ArizePhoenixLogger._init_tracing for full rationale. + """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import SpanKind + + if tracer_provider is not None: + self.tracer = tracer_provider.get_tracer("litellm") + self.span_kind = SpanKind + return + + provider = TracerProvider(resource=self._get_litellm_resource(self.config)) + provider.add_span_processor(self._get_span_processor()) + self.tracer = provider.get_tracer("litellm") + self.span_kind = SpanKind + + def _init_otel_logger_on_litellm_proxy(self): + """ + Override: Arize should NOT overwrite the proxy's + ``open_telemetry_logger``. That attribute is reserved for the + primary ``otel`` callback which handles proxy-level parent spans. + """ + pass + def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]): ArizeLogger.set_arize_attributes(span, kwargs, response_obj) return diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index 759dccdd8a..33f858cc42 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -26,6 +26,55 @@ ARIZE_HOSTED_PHOENIX_ENDPOINT = "https://otlp.arize.com/v1/traces" class ArizePhoenixLogger(OpenTelemetry): + """ + Arize Phoenix logger that sends traces to a Phoenix endpoint. + + Creates its own dedicated TracerProvider so it can coexist with the + generic ``otel`` callback (or any other OTEL-based integration) without + fighting over the global ``opentelemetry.trace`` TracerProvider singleton. + """ + + def _init_tracing(self, tracer_provider): + """ + Override to always create a *private* TracerProvider for Arize Phoenix. + + The base ``OpenTelemetry._init_tracing`` falls back to the global + TracerProvider when one already exists. That causes whichever + integration initialises second to silently reuse the first one's + exporter, so spans only reach one destination. + + By creating our own provider we guarantee Arize Phoenix always gets + its own exporter pipeline, regardless of initialisation order. + """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import SpanKind + + if tracer_provider is not None: + # Explicitly supplied (e.g. in tests) — honour it. + self.tracer = tracer_provider.get_tracer("litellm") + self.span_kind = SpanKind + return + + # Always create a dedicated provider — never touch the global one. + provider = TracerProvider(resource=self._get_litellm_resource(self.config)) + provider.add_span_processor(self._get_span_processor()) + self.tracer = provider.get_tracer("litellm") + self.span_kind = SpanKind + verbose_logger.debug( + "ArizePhoenixLogger: Created dedicated TracerProvider " + "(endpoint=%s, exporter=%s)", + self.config.endpoint, + self.config.exporter, + ) + + def _init_otel_logger_on_litellm_proxy(self): + """ + Override: Arize Phoenix should NOT overwrite the proxy's + ``open_telemetry_logger``. That attribute is reserved for the + primary ``otel`` callback which handles proxy-level parent spans. + """ + pass + def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]): ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj) return diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7ae62718af..82a7af64f9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3764,7 +3764,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 from litellm.integrations.opentelemetry import OpenTelemetry for callback in _in_memory_loggers: - if isinstance(callback, OpenTelemetry): + if type(callback) is OpenTelemetry: return callback # type: ignore otel_logger = OpenTelemetry( **_get_custom_logger_settings_from_proxy_server( diff --git a/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py b/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py new file mode 100644 index 0000000000..40b814e69e --- /dev/null +++ b/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py @@ -0,0 +1,171 @@ +""" +Tests that Arize Phoenix / Arize and the generic ``otel`` callback can +coexist, each sending spans to their own independent exporter. + +Covers the three root-cause fixes: +1. ArizePhoenixLogger / ArizeLogger create *dedicated* TracerProviders. +2. The ``otel`` dedup check does NOT match Arize subclasses. +3. Arize loggers do NOT overwrite ``proxy_server.open_telemetry_logger``. +""" + +import os +import unittest +from unittest.mock import patch + +import pytest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_otel_logger(exporter: InMemorySpanExporter) -> OpenTelemetry: + """Create a generic ``otel`` callback backed by an in-memory exporter. + + We build a dedicated TracerProvider explicitly so the test is isolated + from whatever global provider state may exist. + """ + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + config = OpenTelemetryConfig(exporter=exporter) + return OpenTelemetry(config=config, callback_name="otel", tracer_provider=provider) + + +def _make_arize_phoenix_logger(exporter: InMemorySpanExporter): + """Create an ``arize_phoenix`` callback backed by an in-memory exporter. + + ArizePhoenixLogger._init_tracing creates its own TracerProvider, so we + pass the exporter via config and let it build the provider internally. + """ + from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger + + config = OpenTelemetryConfig(exporter=exporter) + return ArizePhoenixLogger(config=config, callback_name="arize_phoenix") + + +def _make_arize_logger(exporter: InMemorySpanExporter): + """Create an ``arize`` callback backed by an in-memory exporter. + + ArizeLogger._init_tracing creates its own TracerProvider, so we pass + the exporter via config and let it build the provider internally. + """ + from litellm.integrations.arize.arize import ArizeLogger + + config = OpenTelemetryConfig(exporter=exporter) + return ArizeLogger(config=config, callback_name="arize") + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +class TestIndependentTracerProviders(unittest.TestCase): + """Each integration must get its own TracerProvider so spans go to the right exporter.""" + + def test_otel_and_arize_phoenix_have_different_tracer_providers(self): + otel_exporter = InMemorySpanExporter() + phoenix_exporter = InMemorySpanExporter() + + otel_logger = _make_otel_logger(otel_exporter) + phoenix_logger = _make_arize_phoenix_logger(phoenix_exporter) + + # The tracers must come from different providers + assert otel_logger.tracer is not phoenix_logger.tracer + + def test_otel_and_arize_have_different_tracer_providers(self): + otel_exporter = InMemorySpanExporter() + arize_exporter = InMemorySpanExporter() + + otel_logger = _make_otel_logger(otel_exporter) + arize_logger = _make_arize_logger(arize_exporter) + + assert otel_logger.tracer is not arize_logger.tracer + + def test_arize_phoenix_and_arize_have_different_tracer_providers(self): + phoenix_exporter = InMemorySpanExporter() + arize_exporter = InMemorySpanExporter() + + phoenix_logger = _make_arize_phoenix_logger(phoenix_exporter) + arize_logger = _make_arize_logger(arize_exporter) + + assert phoenix_logger.tracer is not arize_logger.tracer + + +class TestSpansRoutedToCorrectExporter(unittest.TestCase): + """Spans created by each logger must land in its own exporter, not the other's.""" + + def test_spans_go_to_respective_exporters(self): + otel_exporter = InMemorySpanExporter() + phoenix_exporter = InMemorySpanExporter() + + otel_logger = _make_otel_logger(otel_exporter) + phoenix_logger = _make_arize_phoenix_logger(phoenix_exporter) + + # Create a span on each — SimpleSpanProcessor exports synchronously on end() + otel_span = otel_logger.tracer.start_span("otel_test_span") + otel_span.end() + + phoenix_span = phoenix_logger.tracer.start_span("phoenix_test_span") + phoenix_span.end() + + # Read spans *before* shutdown (shutdown clears the in-memory store) + otel_span_names = [s.name for s in otel_exporter.get_finished_spans()] + phoenix_span_names = [s.name for s in phoenix_exporter.get_finished_spans()] + + assert "otel_test_span" in otel_span_names + assert "phoenix_test_span" not in otel_span_names + + assert "phoenix_test_span" in phoenix_span_names + assert "otel_test_span" not in phoenix_span_names + + +class TestOtelDedupCheck(unittest.TestCase): + """The ``otel`` callback dedup must use exact type check, not isinstance.""" + + def test_arize_phoenix_logger_is_not_matched_by_otel_dedup(self): + from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger + + phoenix_logger = _make_arize_phoenix_logger(InMemorySpanExporter()) + + # isinstance would match — but type() must not + assert isinstance(phoenix_logger, OpenTelemetry) + assert type(phoenix_logger) is not OpenTelemetry + + def test_arize_logger_is_not_matched_by_otel_dedup(self): + from litellm.integrations.arize.arize import ArizeLogger + + arize_logger = _make_arize_logger(InMemorySpanExporter()) + + assert isinstance(arize_logger, OpenTelemetry) + assert type(arize_logger) is not OpenTelemetry + + def test_otel_logger_matches_own_dedup(self): + otel_logger = _make_otel_logger(InMemorySpanExporter()) + assert type(otel_logger) is OpenTelemetry + + +class TestProxyLoggerNotOverwritten(unittest.TestCase): + """Arize / Phoenix must not overwrite ``proxy_server.open_telemetry_logger``.""" + + @patch("litellm.proxy.proxy_server.open_telemetry_logger", None) + def test_arize_phoenix_does_not_set_proxy_otel_logger(self): + from litellm.proxy import proxy_server + + _make_arize_phoenix_logger(InMemorySpanExporter()) + assert proxy_server.open_telemetry_logger is None + + @patch("litellm.proxy.proxy_server.open_telemetry_logger", None) + def test_arize_does_not_set_proxy_otel_logger(self): + from litellm.proxy import proxy_server + + _make_arize_logger(InMemorySpanExporter()) + assert proxy_server.open_telemetry_logger is None + + +if __name__ == "__main__": + unittest.main() From b3fea9e983f058b98137e80237ec7614eb2c61a8 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 10 Feb 2026 14:34:00 -0500 Subject: [PATCH 07/15] check for typeddict --- litellm/integrations/arize/arize_phoenix.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index 33f858cc42..93b5be6b1b 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -106,7 +106,7 @@ class ArizePhoenixLogger(OpenTelemetry): traces to different Phoenix projects dynamically. """ standard_logging_payload = kwargs.get("standard_logging_object") - if standard_logging_payload is not None: + if isinstance(standard_logging_payload, dict): metadata = standard_logging_payload.get("metadata") if isinstance(metadata, dict): project_name = metadata.get("phoenix_project_name") @@ -114,8 +114,11 @@ class ArizePhoenixLogger(OpenTelemetry): return str(project_name) # Also check litellm_params.metadata for SDK usage - litellm_params = kwargs.get("litellm_params") or {} - metadata = litellm_params.get("metadata") or {} + litellm_params = kwargs.get("litellm_params") + if isinstance(litellm_params, dict): + metadata = litellm_params.get("metadata") or {} + else: + metadata = {} if isinstance(metadata, dict): project_name = metadata.get("phoenix_project_name") if project_name: From d1c6e25723ae75ca5a33ed8ae95f7dbb65139dcd Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 10 Feb 2026 14:42:33 -0500 Subject: [PATCH 08/15] added tests --- litellm/integrations/arize/arize_phoenix.py | 11 +- .../arize/test_arize_otel_coexistence.py | 2 - .../hooks/test_image_generation_guardrails.py | 293 ++++++++++++++++++ 3 files changed, 298 insertions(+), 8 deletions(-) create mode 100644 tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index 93b5be6b1b..5871bd5d61 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -1,11 +1,15 @@ import os from typing import TYPE_CHECKING, Any, Optional, Union +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.trace import SpanKind + from litellm._logging import verbose_logger from litellm.integrations.arize import _utils from litellm.integrations.arize._utils import ArizeOTELAttributes -from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig from litellm.integrations.opentelemetry import OpenTelemetry +from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute +from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -46,9 +50,6 @@ class ArizePhoenixLogger(OpenTelemetry): By creating our own provider we guarantee Arize Phoenix always gets its own exporter pipeline, regardless of initialisation order. """ - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.trace import SpanKind - if tracer_provider is not None: # Explicitly supplied (e.g. in tests) — honour it. self.tracer = tracer_provider.get_tracer("litellm") @@ -81,8 +82,6 @@ class ArizePhoenixLogger(OpenTelemetry): @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): - from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute - _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) # Dynamic project name: check metadata first, then fall back to env var config diff --git a/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py b/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py index 40b814e69e..a0dcf1c091 100644 --- a/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py +++ b/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py @@ -8,11 +8,9 @@ Covers the three root-cause fixes: 3. Arize loggers do NOT overwrite ``proxy_server.open_telemetry_logger``. """ -import os import unittest from unittest.mock import patch -import pytest from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter diff --git a/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py b/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py new file mode 100644 index 0000000000..4a5d901b74 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py @@ -0,0 +1,293 @@ +""" +Tests that guardrails (post_call_success_hook) fire for image generation requests. + +The /images/generations endpoint in proxy/image_endpoints/endpoints.py calls +proxy_logging_obj.post_call_success_hook after a successful image generation. +These tests verify: +1. CustomGuardrail.async_post_call_success_hook is invoked for image generation. +2. A guardrail can inspect and transform the image response. +3. A guardrail that raises blocks the response (exception propagates). +""" + +import os +import sys +from typing import Any, Optional +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ImageObject, ImageResponse + + +def _make_image_response(**kwargs) -> ImageResponse: + """Helper to build a minimal ImageResponse for tests.""" + return ImageResponse( + data=[ImageObject(url="https://example.com/img.png")], + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# 1. Hook is invoked for image generation responses +# --------------------------------------------------------------------------- + + +class TrackingGuardrail(CustomGuardrail): + """Guardrail that records whether it was called and with what args.""" + + def __init__(self): + super().__init__( + guardrail_name="tracking_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + self.called = False + self.received_data: Optional[dict] = None + self.received_response: Optional[Any] = None + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + self.called = True + self.received_data = data + self.received_response = response + return response + + +@pytest.mark.asyncio +async def test_post_call_success_hook_invoked_for_image_generation(): + """ + Verify that a default-on guardrail's async_post_call_success_hook is + called when ProxyLogging.post_call_success_hook is invoked with an + ImageResponse (the same path used by the /images/generations endpoint). + """ + guardrail = TrackingGuardrail() + image_response = _make_image_response() + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + data = {"model": "dall-e-3", "prompt": "A sunset over mountains"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.post_call_success_hook( + data=data, + response=image_response, + user_api_key_dict=user_api_key_dict, + ) + + assert guardrail.called is True, "Guardrail hook was not invoked for image generation" + assert guardrail.received_data is not None + assert guardrail.received_data["model"] == "dall-e-3" + assert isinstance(guardrail.received_response, ImageResponse) + # The response should be passed through unchanged + assert result is image_response + + +# --------------------------------------------------------------------------- +# 2. Guardrail can transform image generation response +# --------------------------------------------------------------------------- + + +class TransformingGuardrail(CustomGuardrail): + """Guardrail that replaces the image URL in the response.""" + + def __init__(self): + super().__init__( + guardrail_name="transforming_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + # Return a modified image response (e.g., watermarked URL) + return ImageResponse( + data=[ImageObject(url="https://example.com/watermarked.png")], + ) + + +@pytest.mark.asyncio +async def test_guardrail_can_transform_image_response(): + """ + Verify that a guardrail can replace the ImageResponse returned to the client. + """ + guardrail = TransformingGuardrail() + original_response = _make_image_response() + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + data = {"model": "dall-e-3", "prompt": "A sunset"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.post_call_success_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + assert result is not original_response + assert isinstance(result, ImageResponse) + assert result.data[0].url == "https://example.com/watermarked.png" + + +# --------------------------------------------------------------------------- +# 3. Guardrail that raises blocks the image response +# --------------------------------------------------------------------------- + + +class BlockingGuardrail(CustomGuardrail): + """Guardrail that raises on unsafe image prompts.""" + + def __init__(self): + super().__init__( + guardrail_name="blocking_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + raise ValueError("Image content blocked by guardrail") + + +@pytest.mark.asyncio +async def test_guardrail_exception_propagates_for_image_generation(): + """ + Verify that an exception raised in a guardrail's post_call_success_hook + propagates up (the proxy endpoint wraps this in an error response). + """ + guardrail = BlockingGuardrail() + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + data = {"model": "dall-e-3", "prompt": "Something unsafe"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + with pytest.raises(ValueError, match="Image content blocked by guardrail"): + await proxy_logging.post_call_success_hook( + data=data, + response=_make_image_response(), + user_api_key_dict=user_api_key_dict, + ) + + +# --------------------------------------------------------------------------- +# 4. Non-guardrail CustomLogger also fires for image generation +# --------------------------------------------------------------------------- + + +class TrackingLogger(CustomLogger): + """Plain CustomLogger (not a guardrail) that tracks invocations.""" + + def __init__(self): + self.called = False + self.received_response = None + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + self.called = True + self.received_response = response + return response + + +@pytest.mark.asyncio +async def test_custom_logger_post_call_success_hook_fires_for_image_generation(): + """ + Verify that a plain CustomLogger (non-guardrail) callback also has its + async_post_call_success_hook invoked for image generation responses. + """ + logger = TrackingLogger() + image_response = _make_image_response() + + with patch("litellm.callbacks", [logger]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + data = {"model": "dall-e-3", "prompt": "A cat"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.post_call_success_hook( + data=data, + response=image_response, + user_api_key_dict=user_api_key_dict, + ) + + assert logger.called is True + assert isinstance(logger.received_response, ImageResponse) + assert result is image_response + + +# --------------------------------------------------------------------------- +# 5. Guardrail with should_run_guardrail=False is skipped +# --------------------------------------------------------------------------- + + +class OptInGuardrail(CustomGuardrail): + """Guardrail that is NOT default_on, so it only runs if explicitly requested.""" + + def __init__(self): + super().__init__( + guardrail_name="opt_in_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.post_call, + ) + self.called = False + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + self.called = True + return response + + +@pytest.mark.asyncio +async def test_non_default_guardrail_skipped_for_image_generation(): + """ + Verify that a guardrail with default_on=False is NOT invoked for image + generation unless the request explicitly enables it. + """ + guardrail = OptInGuardrail() + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + # No guardrails key in data -> should_run_guardrail returns False + data = {"model": "dall-e-3", "prompt": "A sunset"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + await proxy_logging.post_call_success_hook( + data=data, + response=_make_image_response(), + user_api_key_dict=user_api_key_dict, + ) + + assert guardrail.called is False, "Opt-in guardrail should not fire without explicit request" From 2b9b5302ef591630bee5c98d85fe926fb6692e3e Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 10 Feb 2026 15:04:06 -0500 Subject: [PATCH 09/15] add test for dynamic project name in metadata --- .../integrations/arize/test_arize_phoenix.py | 61 ++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/integrations/arize/test_arize_phoenix.py b/tests/test_litellm/integrations/arize/test_arize_phoenix.py index aa227fbff5..129b35fb06 100644 --- a/tests/test_litellm/integrations/arize/test_arize_phoenix.py +++ b/tests/test_litellm/integrations/arize/test_arize_phoenix.py @@ -1,5 +1,5 @@ import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -7,6 +7,7 @@ from litellm.integrations.arize.arize_phoenix import ( ArizePhoenixConfig, ArizePhoenixLogger, ) +from litellm.integrations.arize._utils import ArizeOTELAttributes class TestArizePhoenixConfig(unittest.TestCase): @@ -195,5 +196,63 @@ def test_get_arize_phoenix_config_expection_on_missing_api_key(monkeypatch, env_ +# --------------------------------------------------------------------------- +# Dynamic project naming from metadata +# --------------------------------------------------------------------------- + + +class TestGetDynamicProjectName: + """Tests for _get_dynamic_project_name extraction logic.""" + + def test_extracts_from_standard_logging_object_metadata(self): + kwargs = { + "standard_logging_object": { + "metadata": {"phoenix_project_name": "my-project"}, + } + } + assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) == "my-project" + + def test_extracts_from_litellm_params_metadata(self): + kwargs = { + "litellm_params": { + "metadata": {"phoenix_project_name": "sdk-project"}, + } + } + assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) == "sdk-project" + + def test_returns_none_when_no_metadata(self): + assert ArizePhoenixLogger._get_dynamic_project_name({}) is None + + def test_non_dict_standard_logging_object_does_not_raise(self): + """isinstance(dict) guard prevents AttributeError on non-dict payloads.""" + kwargs = {"standard_logging_object": "not-a-dict"} + assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) is None + + +class TestDynamicProjectNameOnSpan: + """set_arize_phoenix_attributes sets openinference.project.name on the span.""" + + @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-fallback"}, clear=False) + @patch("litellm.integrations.arize._utils.set_attributes") + def test_dynamic_name_sets_span_attribute(self, _mock_set_attrs): + span = MagicMock() + kwargs = { + "standard_logging_object": { + "metadata": {"phoenix_project_name": "dynamic-proj"}, + } + } + ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj=None) + + span.set_attribute.assert_called_once_with("openinference.project.name", "dynamic-proj") + + @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-project"}, clear=False) + @patch("litellm.integrations.arize._utils.set_attributes") + def test_falls_back_to_env_var_when_no_dynamic_name(self, _mock_set_attrs): + span = MagicMock() + ArizePhoenixLogger.set_arize_phoenix_attributes(span, {}, response_obj=None) + + span.set_attribute.assert_called_once_with("openinference.project.name", "env-project") + + if __name__ == "__main__": unittest.main() From 70a1bf92e24b186bae88e092feef113223db8694 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 10 Feb 2026 01:46:38 -0500 Subject: [PATCH 10/15] fix(bedrock): skip AssumeRole when ECS/EC2 already running as target IAM role MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When aws_role_name is configured but the environment (ECS task role, EC2 instance profile) is already running as that role, AssumeRole is unnecessary and can fail with AccessDenied. This adds same-role detection for ECS/EC2 (extending existing IRSA support) and a fallback to ambient credentials when AssumeRole returns AccessDenied. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- litellm/llms/bedrock/base_aws_llm.py | 98 +++++++-- litellm/utils.py | 6 + .../llms/bedrock/test_base_aws_llm.py | 196 ++++++++++++++++-- 3 files changed, 270 insertions(+), 30 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1de1c40c43..506b8811d3 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -211,25 +211,13 @@ class BaseAWSLLM: aws_external_id=aws_external_id, ) elif aws_role_name is not None: - # Check if we're in IRSA and trying to assume the same role we already have - current_role_arn = os.getenv("AWS_ROLE_ARN") - web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") - - # In IRSA environments, we should skip role assumption if we're already running as the target role - # This is true when: - # 1. We have AWS_ROLE_ARN set (current role) - # 2. We have AWS_WEB_IDENTITY_TOKEN_FILE set (IRSA environment) - # 3. The current role matches the requested role - if ( - current_role_arn - and web_identity_token_file - and current_role_arn == aws_role_name - ): + # Check if we're already running as the target role and can skip assumption + # This handles IRSA (EKS), ECS task roles, and EC2 instance profiles + if self._is_already_running_as_role(aws_role_name): verbose_logger.debug( - "Using IRSA same-role optimization: calling _auth_with_env_vars" + "Already running as target role %s, using ambient credentials", + aws_role_name, ) - # We're already running as this role via IRSA, no need to assume it again - # Use the default boto3 credentials (which will use the IRSA credentials) credentials, _cache_ttl = self._auth_with_env_vars() else: verbose_logger.debug( @@ -553,6 +541,65 @@ class BaseAWSLLM: aws_region_name = "us-west-2" return aws_region_name + def _is_already_running_as_role(self, aws_role_name: str) -> bool: + """ + Check if the current environment is already running as the target IAM role. + + This handles multiple AWS environments: + - IRSA (EKS): AWS_ROLE_ARN + AWS_WEB_IDENTITY_TOKEN_FILE are set + - ECS task roles: Uses sts:GetCallerIdentity to check current role ARN + - EC2 instance profiles: Uses sts:GetCallerIdentity to check current role ARN + + Returns True if the current identity matches the target role, meaning + we can skip sts:AssumeRole and use ambient credentials directly. + """ + # Fast path: IRSA environment check (no API call needed) + current_role_arn = os.getenv("AWS_ROLE_ARN") + web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") + if current_role_arn and web_identity_token_file: + return current_role_arn == aws_role_name + + # For ECS/EC2: call sts:GetCallerIdentity to check if already running as the role + try: + import boto3 + + with tracer.trace("boto3.client(sts).get_caller_identity"): + sts_client = boto3.client("sts") + identity = sts_client.get_caller_identity() + caller_arn = identity.get("Arn", "") + + # The caller ARN for an ECS task role looks like: + # arn:aws:sts::123456789012:assumed-role/MyRole/session-name + # The target role ARN looks like: + # arn:aws:iam::123456789012:role/MyRole + # We need to compare the role name portion + if ":assumed-role/" in caller_arn: + # Extract role name from assumed-role ARN + # Format: arn:aws:sts::ACCOUNT:assumed-role/ROLE_NAME/SESSION + caller_role_name = caller_arn.split(":assumed-role/")[1].split("/")[0] + + # Extract role name from target role ARN + # Format: arn:aws:iam::ACCOUNT:role/ROLE_NAME or + # arn:aws:iam::ACCOUNT:role/path/ROLE_NAME + if ":role/" in aws_role_name: + target_role_name = aws_role_name.split(":role/")[-1].split("/")[-1] + else: + target_role_name = aws_role_name + + if caller_role_name == target_role_name: + verbose_logger.debug( + "Current identity already matches target role: %s", + aws_role_name, + ) + return True + + except Exception as e: + verbose_logger.debug( + "Could not determine current role identity: %s", str(e) + ) + + return False + @tracer.wrap() def _auth_with_web_identity_token( self, @@ -867,7 +914,22 @@ class BaseAWSLLM: if aws_external_id is not None: assume_role_params["ExternalId"] = aws_external_id - sts_response = sts_client.assume_role(**assume_role_params) + try: + sts_response = sts_client.assume_role(**assume_role_params) + except Exception as e: + error_str = str(e) + # If AssumeRole fails because the caller already IS the role + # (e.g., ECS task role, root account, or same-role scenario), + # fall back to using ambient credentials directly + if "AccessDenied" in error_str: + verbose_logger.warning( + "AssumeRole failed for %s (%s). " + "Falling back to ambient credentials (boto3 default chain).", + aws_role_name, + error_str, + ) + return self._auth_with_env_vars() + raise # Extract the credentials from the response and convert to Session Credentials sts_credentials = sts_response["Credentials"] diff --git a/litellm/utils.py b/litellm/utils.py index ed0d6ee930..733baeadd8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6140,6 +6140,12 @@ def validate_environment( # noqa: PLR0915 if ( "AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ + ) or ( + # IAM role, profile, or web identity auth don't require access keys + "AWS_ROLE_ARN" in os.environ + or "AWS_PROFILE" in os.environ + or "AWS_WEB_IDENTITY_TOKEN_FILE" in os.environ + or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" in os.environ # ECS task role ): keys_in_environment = True else: diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 77eda43251..b57c20fb7c 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -853,29 +853,66 @@ def test_role_assumption_ttl_calculation(): assert 3500 <= ttl <= 3600 # Allow some variance for test execution time -def test_role_assumption_error_handling(): +def test_role_assumption_access_denied_falls_back_to_env_vars(): """ - Test that role assumption errors are properly propagated. + Test that when AssumeRole fails with AccessDenied, we fall back to ambient credentials. + This handles ECS task roles, root accounts, and same-role scenarios where + AssumeRole is unnecessary because the caller already has the role's permissions. """ base_aws_llm = BaseAWSLLM() - - # Mock the boto3 STS client to raise an exception + + # Mock the boto3 STS client to raise AccessDenied mock_sts_client = MagicMock() - mock_sts_client.assume_role.side_effect = Exception("AccessDenied: User is not authorized to perform sts:AssumeRole") - + mock_sts_client.assume_role.side_effect = Exception( + "An error occurred (AccessDenied) when calling the AssumeRole operation: " + "Roles may not be assumed by root accounts." + ) + + # Mock _auth_with_env_vars to return fallback credentials + mock_creds = MagicMock() + mock_creds.access_key = "fallback-access-key" + mock_creds.secret_key = "fallback-secret-key" + + with patch("boto3.client", return_value=mock_sts_client): + with patch.object( + base_aws_llm, "_auth_with_env_vars", return_value=(mock_creds, None) + ) as mock_env_auth: + credentials, ttl = base_aws_llm._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole", + aws_session_name="error-test-session", + ) + + # Should have fallen back to env vars + mock_env_auth.assert_called_once() + assert credentials.access_key == "fallback-access-key" + + +def test_role_assumption_non_access_denied_error_propagated(): + """ + Test that non-AccessDenied errors from AssumeRole are still propagated. + """ + base_aws_llm = BaseAWSLLM() + + # Mock the boto3 STS client to raise a non-AccessDenied exception + mock_sts_client = MagicMock() + mock_sts_client.assume_role.side_effect = Exception( + "An error occurred (MalformedPolicyDocument) when calling the AssumeRole operation" + ) + with patch("boto3.client", return_value=mock_sts_client): - - # Should raise the exception with pytest.raises(Exception) as exc_info: base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, aws_session_token=None, - aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole", - aws_session_name="error-test-session" + aws_role_name="arn:aws:iam::1111111111111:role/BadPolicyRole", + aws_session_name="error-test-session", ) - - assert "AccessDenied" in str(exc_info.value) + + assert "MalformedPolicyDocument" in str(exc_info.value) def test_multiple_role_assumptions_in_sequence(): @@ -1195,3 +1232,138 @@ def test_converse_handler_external_id_extraction(): assert hasattr(mock_get_credentials, 'called_kwargs') assert "aws_external_id" in mock_get_credentials.called_kwargs assert mock_get_credentials.called_kwargs["aws_external_id"] == "TestExternalID123" + + +def test_is_already_running_as_role_irsa_same_role(): + """Test IRSA fast path: when AWS_ROLE_ARN matches target role.""" + base_aws_llm = BaseAWSLLM() + + with patch.dict(os.environ, { + "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole", + "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token", + }): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/MyRole" + ) is True + + +def test_is_already_running_as_role_irsa_different_role(): + """Test IRSA fast path: when AWS_ROLE_ARN does NOT match target role.""" + base_aws_llm = BaseAWSLLM() + + with patch.dict(os.environ, { + "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole", + "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token", + }): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::999999999999:role/OtherRole" + ) is False + + +def test_is_already_running_as_role_ecs_task_role(): + """Test ECS/EC2 path: GetCallerIdentity shows assumed-role matching target.""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + # Ensure no IRSA env vars + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/MyEcsTaskRole" + ) is True + + +def test_is_already_running_as_role_ecs_different_role(): + """Test ECS/EC2 path: GetCallerIdentity shows a different role.""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::999999999999:role/DifferentRole" + ) is False + + +def test_is_already_running_as_role_ecs_role_with_path(): + """Test ECS path with role that has a path prefix (e.g., /service-role/MyRole).""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + # Role ARN with path + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/service-role/MyEcsTaskRole" + ) is True + + +def test_is_already_running_as_role_get_caller_identity_fails(): + """Test that when GetCallerIdentity fails, we return False (don't crash).""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.side_effect = Exception("No credentials found") + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/SomeRole" + ) is False + + +def test_get_credentials_ecs_same_role_skips_assume_role(): + """ + End-to-end test: when running on ECS with the same role as aws_role_name, + get_credentials should use ambient credentials and NOT call AssumeRole. + """ + base_aws_llm = BaseAWSLLM() + + mock_creds = MagicMock() + mock_creds.access_key = "ecs-access-key" + mock_creds.secret_key = "ecs-secret-key" + mock_creds.token = "ecs-session-token" + + with patch.object( + base_aws_llm, + "_is_already_running_as_role", + return_value=True, + ): + with patch.object( + base_aws_llm, + "_auth_with_env_vars", + return_value=(mock_creds, None), + ) as mock_env_auth: + with patch.object( + base_aws_llm, + "_auth_with_aws_role", + ) as mock_role_auth: + credentials = base_aws_llm.get_credentials( + aws_role_name="arn:aws:iam::123456789012:role/MyEcsTaskRole", + aws_region_name="us-east-1", + ) + + # Should use env vars, NOT role assumption + mock_env_auth.assert_called_once() + mock_role_auth.assert_not_called() + assert credentials.access_key == "ecs-access-key" From 6fdfdd27b599858a4f62cf07a26b41ada97d4b46 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 10 Feb 2026 02:09:52 -0500 Subject: [PATCH 11/15] fix(bedrock): address review - cross-account, SSL verify, narrow fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. Cross-account false match: Added _parse_arn_account_and_role_name() helper that compares partition + account ID + role name (not just role name) to prevent same-name-different-account false matches. 2. SSL verify: _is_already_running_as_role() now passes ssl_verify to the STS client via self._get_ssl_verify(), consistent with all other boto3 client creation in this module. 3. Overbroad AccessDenied fallback: The catch in _auth_with_aws_role now only falls back to ambient credentials when _is_already_running_as_role positively confirms the caller is the target role. Genuine trust-policy or permission misconfigurations are re-raised. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- litellm/llms/bedrock/base_aws_llm.py | 113 +++++++++--- .../llms/bedrock/test_base_aws_llm.py | 174 ++++++++++++++++-- 2 files changed, 244 insertions(+), 43 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 506b8811d3..304c707fa0 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -213,7 +213,7 @@ class BaseAWSLLM: elif aws_role_name is not None: # Check if we're already running as the target role and can skip assumption # This handles IRSA (EKS), ECS task roles, and EC2 instance profiles - if self._is_already_running_as_role(aws_role_name): + if self._is_already_running_as_role(aws_role_name, ssl_verify=ssl_verify): verbose_logger.debug( "Already running as target role %s, using ambient credentials", aws_role_name, @@ -541,7 +541,49 @@ class BaseAWSLLM: aws_region_name = "us-west-2" return aws_region_name - def _is_already_running_as_role(self, aws_role_name: str) -> bool: + @staticmethod + def _parse_arn_account_and_role_name( + arn: str, + ) -> Optional[Tuple[str, str, str]]: + """ + Parse an ARN and return (partition, account_id, role_name). + + Handles: + - arn:aws:iam::123456789012:role/MyRole + - arn:aws:iam::123456789012:role/path/to/MyRole + - arn:aws:sts::123456789012:assumed-role/MyRole/session-name + + Returns None if the ARN cannot be parsed. + """ + # ARN format: arn:PARTITION:SERVICE:REGION:ACCOUNT:RESOURCE + parts = arn.split(":") + if len(parts) < 6 or parts[0] != "arn": + return None + + partition = parts[1] # e.g. "aws", "aws-cn", "aws-us-gov" + account_id = parts[4] + resource = ":".join(parts[5:]) # rejoin in case resource contains colons + + if resource.startswith("role/"): + # arn:aws:iam::ACCOUNT:role/[path/]ROLE_NAME + role_name = resource.split("/")[-1] + elif resource.startswith("assumed-role/"): + # arn:aws:sts::ACCOUNT:assumed-role/ROLE_NAME/SESSION + role_parts = resource.split("/") + if len(role_parts) >= 2: + role_name = role_parts[1] + else: + return None + else: + return None + + return partition, account_id, role_name + + def _is_already_running_as_role( + self, + aws_role_name: str, + ssl_verify: Optional[Union[bool, str]] = None, + ) -> bool: """ Check if the current environment is already running as the target IAM role. @@ -550,9 +592,18 @@ class BaseAWSLLM: - ECS task roles: Uses sts:GetCallerIdentity to check current role ARN - EC2 instance profiles: Uses sts:GetCallerIdentity to check current role ARN + Compares partition, account ID, and role name to avoid cross-account + false matches. + Returns True if the current identity matches the target role, meaning we can skip sts:AssumeRole and use ambient credentials directly. """ + target_parsed = self._parse_arn_account_and_role_name(aws_role_name) + if target_parsed is None: + return False + + target_partition, target_account, target_role = target_parsed + # Fast path: IRSA environment check (no API call needed) current_role_arn = os.getenv("AWS_ROLE_ARN") web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") @@ -564,29 +615,20 @@ class BaseAWSLLM: import boto3 with tracer.trace("boto3.client(sts).get_caller_identity"): - sts_client = boto3.client("sts") + sts_client = boto3.client( + "sts", verify=self._get_ssl_verify(ssl_verify) + ) identity = sts_client.get_caller_identity() caller_arn = identity.get("Arn", "") - # The caller ARN for an ECS task role looks like: - # arn:aws:sts::123456789012:assumed-role/MyRole/session-name - # The target role ARN looks like: - # arn:aws:iam::123456789012:role/MyRole - # We need to compare the role name portion - if ":assumed-role/" in caller_arn: - # Extract role name from assumed-role ARN - # Format: arn:aws:sts::ACCOUNT:assumed-role/ROLE_NAME/SESSION - caller_role_name = caller_arn.split(":assumed-role/")[1].split("/")[0] - - # Extract role name from target role ARN - # Format: arn:aws:iam::ACCOUNT:role/ROLE_NAME or - # arn:aws:iam::ACCOUNT:role/path/ROLE_NAME - if ":role/" in aws_role_name: - target_role_name = aws_role_name.split(":role/")[-1].split("/")[-1] - else: - target_role_name = aws_role_name - - if caller_role_name == target_role_name: + caller_parsed = self._parse_arn_account_and_role_name(caller_arn) + if caller_parsed is not None: + caller_partition, caller_account, caller_role = caller_parsed + if ( + caller_partition == target_partition + and caller_account == target_account + and caller_role == target_role + ): verbose_logger.debug( "Current identity already matches target role: %s", aws_role_name, @@ -918,17 +960,30 @@ class BaseAWSLLM: sts_response = sts_client.assume_role(**assume_role_params) except Exception as e: error_str = str(e) - # If AssumeRole fails because the caller already IS the role - # (e.g., ECS task role, root account, or same-role scenario), - # fall back to using ambient credentials directly if "AccessDenied" in error_str: - verbose_logger.warning( - "AssumeRole failed for %s (%s). " - "Falling back to ambient credentials (boto3 default chain).", + # Only fall back to ambient credentials if we can positively + # confirm the caller is already the target role (same account, + # partition, and role name). This avoids silently using the + # wrong identity when there is a genuine trust-policy or + # permission misconfiguration. + if self._is_already_running_as_role( + aws_role_name, ssl_verify=ssl_verify + ): + verbose_logger.warning( + "AssumeRole failed for %s (%s). " + "Caller is already running as this role; " + "falling back to ambient credentials.", + aws_role_name, + error_str, + ) + return self._auth_with_env_vars() + # Genuine permission error — re-raise + verbose_logger.error( + "AssumeRole AccessDenied for %s and caller is NOT " + "the same role. Re-raising. Error: %s", aws_role_name, error_str, ) - return self._auth_with_env_vars() raise # Extract the credentials from the response and convert to Session Credentials diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index b57c20fb7c..cf9fee6bac 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -853,11 +853,10 @@ def test_role_assumption_ttl_calculation(): assert 3500 <= ttl <= 3600 # Allow some variance for test execution time -def test_role_assumption_access_denied_falls_back_to_env_vars(): +def test_role_assumption_access_denied_falls_back_when_same_role(): """ - Test that when AssumeRole fails with AccessDenied, we fall back to ambient credentials. - This handles ECS task roles, root accounts, and same-role scenarios where - AssumeRole is unnecessary because the caller already has the role's permissions. + Test that when AssumeRole fails with AccessDenied AND the caller is confirmed + to already be running as the target role, we fall back to ambient credentials. """ base_aws_llm = BaseAWSLLM() @@ -877,17 +876,51 @@ def test_role_assumption_access_denied_falls_back_to_env_vars(): with patch.object( base_aws_llm, "_auth_with_env_vars", return_value=(mock_creds, None) ) as mock_env_auth: - credentials, ttl = base_aws_llm._auth_with_aws_role( - aws_access_key_id=None, - aws_secret_access_key=None, - aws_session_token=None, - aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole", - aws_session_name="error-test-session", - ) + # _is_already_running_as_role returns True => fallback allowed + with patch.object( + base_aws_llm, "_is_already_running_as_role", return_value=True + ): + credentials, ttl = base_aws_llm._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole", + aws_session_name="error-test-session", + ) - # Should have fallen back to env vars - mock_env_auth.assert_called_once() - assert credentials.access_key == "fallback-access-key" + # Should have fallen back to env vars + mock_env_auth.assert_called_once() + assert credentials.access_key == "fallback-access-key" + + +def test_role_assumption_access_denied_raises_when_different_role(): + """ + Test that when AssumeRole fails with AccessDenied but the caller is NOT + the same role, the error is re-raised (genuine permission failure). + """ + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.assume_role.side_effect = Exception( + "An error occurred (AccessDenied) when calling the AssumeRole operation: " + "User is not authorized to perform sts:AssumeRole" + ) + + with patch("boto3.client", return_value=mock_sts_client): + # _is_already_running_as_role returns False => do NOT fallback + with patch.object( + base_aws_llm, "_is_already_running_as_role", return_value=False + ): + with pytest.raises(Exception) as exc_info: + base_aws_llm._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole", + aws_session_name="error-test-session", + ) + + assert "AccessDenied" in str(exc_info.value) def test_role_assumption_non_access_denied_error_propagated(): @@ -1367,3 +1400,116 @@ def test_get_credentials_ecs_same_role_skips_assume_role(): mock_env_auth.assert_called_once() mock_role_auth.assert_not_called() assert credentials.access_key == "ecs-access-key" + + +def test_parse_arn_account_and_role_name(): + """Test the ARN parser helper for various ARN formats.""" + parse = BaseAWSLLM._parse_arn_account_and_role_name + + # Standard IAM role ARN + assert parse("arn:aws:iam::123456789012:role/MyRole") == ( + "aws", "123456789012", "MyRole" + ) + + # IAM role ARN with path + assert parse("arn:aws:iam::123456789012:role/service-role/MyRole") == ( + "aws", "123456789012", "MyRole" + ) + + # Assumed-role ARN (from GetCallerIdentity) + assert parse("arn:aws:sts::123456789012:assumed-role/MyRole/session-id") == ( + "aws", "123456789012", "MyRole" + ) + + # China partition + assert parse("arn:aws-cn:iam::123456789012:role/MyRole") == ( + "aws-cn", "123456789012", "MyRole" + ) + + # GovCloud partition + assert parse("arn:aws-us-gov:iam::123456789012:role/MyRole") == ( + "aws-us-gov", "123456789012", "MyRole" + ) + + # Invalid ARNs + assert parse("not-an-arn") is None + assert parse("arn:aws:iam::123456789012:user/MyUser") is None + assert parse("") is None + + +def test_is_already_running_as_role_cross_account_same_name(): + """ + Test that same role NAME in different accounts does NOT match. + This is the cross-account false-match prevention. + """ + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + # Caller is in account 111111111111 + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::111111111111:assumed-role/MyRole/session-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + # Target is same role name but in account 222222222222 + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::222222222222:role/MyRole" + ) is False + + +def test_is_already_running_as_role_cross_partition(): + """ + Test that same role name + account but different partition does NOT match. + """ + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyRole/session-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + # Same account and role but aws-cn partition + assert base_aws_llm._is_already_running_as_role( + "arn:aws-cn:iam::123456789012:role/MyRole" + ) is False + + +def test_is_already_running_as_role_invalid_target_arn(): + """ + Test that an unparseable target ARN returns False immediately. + """ + base_aws_llm = BaseAWSLLM() + + # Should return False without making any API calls + assert base_aws_llm._is_already_running_as_role("not-a-valid-arn") is False + + +def test_is_already_running_as_role_ssl_verify_passed(): + """ + Test that ssl_verify parameter is correctly passed to the STS client. + """ + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyRole/session-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client: + base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/MyRole", + ssl_verify="/path/to/ca-bundle.crt", + ) + mock_boto3_client.assert_called_once_with( + "sts", verify="/path/to/ca-bundle.crt" + ) From f32f8ebf70029b069b8069dd9acb02945daabec8 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 10 Feb 2026 02:22:21 -0500 Subject: [PATCH 12/15] fix(bedrock): accept AWS_CONTAINER_CREDENTIALS_FULL_URI in validate_environment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ECS/Fargate supports both RELATIVE_URI and FULL_URI credential delivery. Only RELATIVE_URI was checked, causing false "missing keys" reports for FULL_URI setups even though boto3 can authenticate fine. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- litellm/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/utils.py b/litellm/utils.py index 733baeadd8..9ae1e9cbb2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6146,6 +6146,7 @@ def validate_environment( # noqa: PLR0915 or "AWS_PROFILE" in os.environ or "AWS_WEB_IDENTITY_TOKEN_FILE" in os.environ or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" in os.environ # ECS task role + or "AWS_CONTAINER_CREDENTIALS_FULL_URI" in os.environ # ECS/Fargate full URI credential delivery ): keys_in_environment = True else: From e2fa31edcef5625ec375d7412374344f7139f6a7 Mon Sep 17 00:00:00 2001 From: shin-bot-litellm Date: Tue, 10 Feb 2026 22:08:44 -0800 Subject: [PATCH 13/15] feat(ui): add license expiration display to usage indicator (#20763) * feat(ui): add license expiration display to usage indicator - Add getLicenseInfo() function to networking.tsx that calls /health/license - Display license expiration as human-readable 'X days remaining' or 'Expires in X months' - Show warning styling (yellow) if license expires in < 30 days - Show error styling (red) if license is expired - Fetch license info in parallel with usage data for efficiency - Include license type display in expanded card view - Compact UI suitable for sidebar widget * fix: timezone mismatch in license expiration calculation Addresses Greptile review feedback - forces UTC midnight for expiration date and normalizes current date to local midnight to prevent off-by-one errors in days remaining calculation. --------- Co-authored-by: Shin --- .../src/components/UsageIndicator.tsx | 164 +++++++++++++++--- .../src/components/networking.tsx | 42 +++++ 2 files changed, 186 insertions(+), 20 deletions(-) diff --git a/ui/litellm-dashboard/src/components/UsageIndicator.tsx b/ui/litellm-dashboard/src/components/UsageIndicator.tsx index 3976e4d3d6..47da9f78bd 100644 --- a/ui/litellm-dashboard/src/components/UsageIndicator.tsx +++ b/ui/litellm-dashboard/src/components/UsageIndicator.tsx @@ -1,8 +1,8 @@ import { useDisableUsageIndicator } from "@/app/(dashboard)/hooks/useDisableUsageIndicator"; import { Badge } from "@tremor/react"; -import { AlertTriangle, ChevronDown, ChevronUp, Loader2, Minus, TrendingUp, UserCheck, Users } from "lucide-react"; +import { AlertTriangle, Calendar, ChevronDown, ChevronUp, Loader2, Minus, TrendingUp, UserCheck, Users } from "lucide-react"; import { useEffect, useState } from "react"; -import { getRemainingUsers } from "./networking"; +import { getRemainingUsers, getLicenseInfo, LicenseInfo } from "./networking"; // Simple utility function to combine class names const cn = (...classes: (string | boolean | undefined)[]) => { @@ -23,11 +23,35 @@ interface UsageData { total_teams_remaining: number | null; } +// Calculate days until expiration +const getDaysUntilExpiration = (expirationDate: string | null): number | null => { + if (!expirationDate) return null; + const expDate = new Date(expirationDate + 'T00:00:00Z'); // Force UTC midnight + const now = new Date(); + now.setHours(0, 0, 0, 0); // Normalize to local midnight + const diffTime = expDate.getTime() - now.getTime(); + const diffDays = Math.ceil(diffTime / (1000 * 60 * 60 * 24)); + return diffDays; +}; + +// Format expiration for display +const formatExpirationDisplay = (daysRemaining: number | null): string => { + if (daysRemaining === null) return "No expiration"; + if (daysRemaining < 0) return "Expired"; + if (daysRemaining === 0) return "Expires today"; + if (daysRemaining === 1) return "1 day remaining"; + if (daysRemaining < 30) return `${daysRemaining} days remaining`; + if (daysRemaining < 60) return "1 month remaining"; + const months = Math.floor(daysRemaining / 30); + return `${months} months remaining`; +}; + export default function UsageIndicator({ accessToken, width = 220 }: UsageIndicatorProps) { const disableUsageIndicator = useDisableUsageIndicator(); const [isExpanded, setIsExpanded] = useState(false); const [isMinimized, setIsMinimized] = useState(false); const [data, setData] = useState(null); + const [licenseInfo, setLicenseInfo] = useState(null); const [isLoading, setIsLoading] = useState(false); const [error, setError] = useState(null); @@ -39,8 +63,12 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica setError(null); try { - const result = await getRemainingUsers(accessToken); - setData(result); + const [usageResult, licenseResult] = await Promise.all([ + getRemainingUsers(accessToken), + getLicenseInfo(accessToken).catch(() => null), // Don't fail if license endpoint unavailable + ]); + setData(usageResult); + setLicenseInfo(licenseResult); } catch (err) { console.error("Failed to fetch usage data:", err); setError("Failed to load usage data"); @@ -52,6 +80,13 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica fetchData(); }, [accessToken]); + // Calculate license expiration metrics + const daysUntilExpiration = licenseInfo?.expiration_date + ? getDaysUntilExpiration(licenseInfo.expiration_date) + : null; + const isLicenseExpired = daysUntilExpiration !== null && daysUntilExpiration < 0; + const isLicenseExpiringSoon = daysUntilExpiration !== null && daysUntilExpiration >= 0 && daysUntilExpiration < 30; + // Calculate derived values from data const getUsageMetrics = (data: UsageData | null) => { if (!data) { @@ -106,35 +141,38 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica const { isOverLimit, isNearLimit, usagePercentage, userMetrics, teamMetrics } = getUsageMetrics(data); + // Include license status in overall status + const hasAnyIssue = isOverLimit || isNearLimit || isLicenseExpired || isLicenseExpiringSoon; + const hasError = isOverLimit || isLicenseExpired; + const hasWarning = (isNearLimit || isLicenseExpiringSoon) && !hasError; + const getStatusColor = () => { - if (isOverLimit) return "red"; - if (isNearLimit) return "yellow"; + if (hasError) return "red"; + if (hasWarning) return "yellow"; return "green"; }; const getStatusIcon = () => { - if (isOverLimit) return ; - if (isNearLimit) return ; + if (hasError) return ; + if (hasWarning) return ; return null; }; // Minimized view - just a small restore button const MinimizedView = () => { - const hasIssues = isOverLimit || isNearLimit; - return (
@@ -198,13 +245,13 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica onClick={() => setIsExpanded(!isExpanded)} className={cn( "flex items-center gap-3 text-left hover:bg-gray-50 rounded-md px-0 py-1 transition-colors flex-1 min-w-0", - isOverLimit && "text-red-600", - isNearLimit && "text-yellow-600", + hasError && "text-red-600", + hasWarning && "text-yellow-600", )} > Usage Status - {(isOverLimit || isNearLimit) && ( + {hasAnyIssue && ( {getStatusIcon()} @@ -229,6 +276,28 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica {/* Expanded details - simple and compact */} {isExpanded && (
+ {/* License expiration section */} + {licenseInfo?.has_license && licenseInfo.expiration_date && ( +
+
+ + License +
+
+ {isLicenseExpired ? ( + + ) : isLicenseExpiringSoon ? ( + + ) : null} + {formatExpirationDisplay(daysUntilExpiration)} +
+
+ )} + {/* Users section */} {data.total_users !== null && (
@@ -323,7 +392,6 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica // Optimized CardStyleView for 220px width const CardStyleView = () => { if (isMinimized) { - const hasIssues = isOverLimit || isNearLimit; return ( @@ -416,6 +496,50 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica {/* Compact stats optimized for 220px */}
+ {/* License expiration section */} + {licenseInfo?.has_license && licenseInfo.expiration_date && ( +
+
+ + License + + {isLicenseExpired ? "Expired" : isLicenseExpiringSoon ? "Expiring soon" : "OK"} + +
+
+ Status: + + {formatExpirationDisplay(daysUntilExpiration)} + +
+ {licenseInfo.license_type && ( +
+ Type: + {licenseInfo.license_type} +
+ )} +
+ )} + {/* Users section */} {data.total_users !== null && (
=> { + try { + const url = proxyBaseUrl ? `${proxyBaseUrl}/health/license` : `/health/license`; + + const response = await fetch(url, { + method: "GET", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + }, + }); + + if (!response.ok) { + // if 404 - return null (endpoint not available) + if (response.status === 404) { + return null; + } + const errorData = await response.text(); + handleError(errorData); + throw new Error("Network response was not ok"); + } + + const data = await response.json(); + return data; + } catch (error) { + console.error("Failed to fetch license info:", error); + throw error; + } +}; + export const updatePassThroughEndpoint = async ( accessToken: string, endpointPath: string, From d3426d55f9a6407865cadf1540ecfe019f2a99d7 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Feb 2026 16:09:49 +0530 Subject: [PATCH 14/15] Fix: litellm import error --- litellm/integrations/arize/arize_phoenix.py | 23 +++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index 5871bd5d61..050fd0f3b7 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -1,35 +1,41 @@ import os from typing import TYPE_CHECKING, Any, Optional, Union -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.trace import SpanKind - from litellm._logging import verbose_logger from litellm.integrations.arize import _utils from litellm.integrations.arize._utils import ArizeOTELAttributes -from litellm.integrations.opentelemetry import OpenTelemetry -from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig if TYPE_CHECKING: + from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import Span as _Span + from opentelemetry.trace import SpanKind + from litellm.integrations.opentelemetry import OpenTelemetry as _OpenTelemetry from litellm.integrations.opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig from litellm.types.integrations.arize import Protocol as _Protocol Protocol = _Protocol OpenTelemetryConfig = _OpenTelemetryConfig Span = Union[_Span, Any] + OpenTelemetry = _OpenTelemetry else: Protocol = Any OpenTelemetryConfig = Any Span = Any + TracerProvider = Any + SpanKind = Any + # Import OpenTelemetry at runtime + try: + from litellm.integrations.opentelemetry import OpenTelemetry + except ImportError: + OpenTelemetry = None # type: ignore ARIZE_HOSTED_PHOENIX_ENDPOINT = "https://otlp.arize.com/v1/traces" -class ArizePhoenixLogger(OpenTelemetry): +class ArizePhoenixLogger(OpenTelemetry): # type: ignore """ Arize Phoenix logger that sends traces to a Phoenix endpoint. @@ -50,6 +56,9 @@ class ArizePhoenixLogger(OpenTelemetry): By creating our own provider we guarantee Arize Phoenix always gets its own exporter pipeline, regardless of initialisation order. """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import SpanKind + if tracer_provider is not None: # Explicitly supplied (e.g. in tests) — honour it. self.tracer = tracer_provider.get_tracer("litellm") @@ -82,6 +91,8 @@ class ArizePhoenixLogger(OpenTelemetry): @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): + from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute + _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) # Dynamic project name: check metadata first, then fall back to env var config From 375ebb333e2f91666852bd0e6a017e51e30e7afe Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Feb 2026 16:13:20 +0530 Subject: [PATCH 15/15] Fix: phoenix tests issues --- litellm/integrations/arize/arize_phoenix.py | 74 +++++++++++++++++++ .../test_logging_redaction_e2e_test.py | 9 +-- .../test_otel_logging.py | 8 ++ 3 files changed, 85 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index 050fd0f3b7..1b038c098f 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -136,6 +136,80 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore return None + def _handle_success(self, kwargs, response_obj, start_time, end_time): + """ + Override to prevent creating duplicate litellm_request spans when a proxy parent span exists. + + ArizePhoenixLogger should reuse the proxy parent span instead of creating a new litellm_request span, + to maintain a shallow span hierarchy as expected by Arize Phoenix. + """ + from opentelemetry.trace import Status, StatusCode + from litellm.secret_managers.main import get_secret_bool + from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME + + verbose_logger.debug( + "ArizePhoenixLogger: Logging kwargs: %s, OTEL config settings=%s", + kwargs, + self.config, + ) + ctx, parent_span = self._get_span_context(kwargs) + + # ArizePhoenixLogger NEVER creates a litellm_request span when a proxy parent span exists + # This is different from the base OpenTelemetry behavior which respects USE_OTEL_LITELLM_REQUEST_SPAN + should_create_primary_span = parent_span is None or ( + parent_span.name != LITELLM_PROXY_REQUEST_SPAN_NAME + and get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN") + ) + + if should_create_primary_span: + # Create a new litellm_request span + span = self._start_primary_span( + kwargs, response_obj, start_time, end_time, ctx + ) + # Raw-request sub-span (if enabled) - child of litellm_request span + self._maybe_log_raw_request( + kwargs, response_obj, start_time, end_time, span + ) + # Ensure proxy-request parent span is annotated with the actual operation kind + if ( + parent_span is not None + and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME + ): + self.set_attributes(parent_span, kwargs, response_obj) + else: + # Do not create primary span (keep hierarchy shallow when parent exists) + span = None + # Only set attributes if the span is still recording (not closed) + # Note: parent_span is guaranteed to be not None here + if parent_span.is_recording(): + parent_span.set_status(Status(StatusCode.OK)) + self.set_attributes(parent_span, kwargs, response_obj) + # Raw-request as direct child of parent_span + self._maybe_log_raw_request( + kwargs, response_obj, start_time, end_time, parent_span + ) + + # 3. Guardrail span + self._create_guardrail_span(kwargs=kwargs, context=ctx) + + # 4. Metrics & cost recording + self._record_metrics(kwargs, response_obj, start_time, end_time) + + # 5. Semantic logs. + if self.config.enable_events: + log_span = span if span is not None else parent_span + if log_span is not None: + self._emit_semantic_logs(kwargs, response_obj, log_span) + + # 6. Do NOT end parent span - it should be managed by its creator + # External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM + # However, proxy-created spans should be closed here + if ( + parent_span is not None + and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME + ): + parent_span.end(end_time=self._to_ns(end_time)) + @staticmethod def get_arize_phoenix_config() -> ArizePhoenixConfig: """ diff --git a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py index b261daab6a..e70d08b900 100644 --- a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py +++ b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py @@ -452,16 +452,13 @@ async def test_redaction_with_metadata_completion_api(): litellm.callbacks = [test_custom_logger] # When metadata is passed, the system uses get_metadata_variable_name_from_kwargs - # to determine which field to check + # to determine which field to check. No headers means redaction should happen + # based on the global setting (litellm.turn_off_message_logging = True) response = await litellm.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "hi"}], mock_response="hello", - metadata={ - "headers": { - "litellm-disable-message-redaction": "true" - } - } + metadata={} ) await asyncio.sleep(1) diff --git a/tests/logging_callback_tests/test_otel_logging.py b/tests/logging_callback_tests/test_otel_logging.py index a0c78305e6..f0511e7d1e 100644 --- a/tests/logging_callback_tests/test_otel_logging.py +++ b/tests/logging_callback_tests/test_otel_logging.py @@ -263,10 +263,18 @@ async def test_arize_phoenix_adds_openinference_kind_and_avoids_duplicate_litell Ensure Arize Phoenix spans include OpenInference span kind and do not create a duplicate litellm_request span when a proxy parent span is already active. """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor exporter.clear() litellm.logging_callback_manager._reset_all_callbacks() + # Set up a global TracerProvider so we can create valid spans + # This simulates the proxy server's TracerProvider + global_provider = TracerProvider() + global_provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace.set_tracer_provider(global_provider) + otel_logger = ArizePhoenixLogger(config=OpenTelemetryConfig(exporter=exporter)) litellm.callbacks = [otel_logger] litellm.success_callback = []