From d086f8574dcaac12c494fa9d37d60cab3b5a5ab7 Mon Sep 17 00:00:00 2001
From: yuneng <yuneng@berri.ai>
Date: Sun, 20 Sep 2026 13:46:52 +0000
Subject: [PATCH] test: add sync and async custom_llm_provider bridge
 propagation tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 ...t_responses_bridge_provider_propagation.py | 52 +++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py
index 09ef1889818..f2e36137a19 100644
--- a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py
+++ b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py
@@ -40,6 +40,58 @@ def _bedrock_mantle_kwargs() -> dict:
     }
 
 
+def _openai_kwargs() -> dict:
+    messages = [{"role": "user", "content": "hi"}]
+    logging_obj = LiteLLMLogging(
+        litellm_call_id="test-call",
+        call_type="completion",
+        model="gpt-5.5",
+        messages=messages,
+        function_id="fn-id",
+        stream=False,
+        start_time=datetime.now(),
+    )
+    return {
+        "model": "gpt-5.5",
+        "custom_llm_provider": "openai",
+        "messages": messages,
+        "optional_params": {},
+        "litellm_params": {},
+        "headers": {},
+        "model_response": ModelResponse(),
+        "logging_obj": logging_obj,
+    }
+
+
+def test_completion_forwards_custom_llm_provider_to_responses():
+    bridge = ResponsesToCompletionBridgeHandler()
+    cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5")
+
+    with patch("litellm.responses", return_value=cached) as fake_responses:
+        result = bridge.completion(**_openai_kwargs())
+
+    assert result is cached
+    assert fake_responses.call_args.kwargs["custom_llm_provider"] == "openai"
+
+
+@pytest.mark.asyncio
+async def test_acompletion_forwards_custom_llm_provider_to_aresponses():
+    bridge = ResponsesToCompletionBridgeHandler()
+    cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5")
+
+    async def _fake_aresponses(**kwargs):
+        _fake_aresponses.kwargs = kwargs
+        return cached
+
+    _fake_aresponses.kwargs = {}
+
+    with patch("litellm.aresponses", _fake_aresponses):
+        result = await bridge.acompletion(**_openai_kwargs())
+
+    assert result is cached
+    assert _fake_aresponses.kwargs["custom_llm_provider"] == "openai"
+
+
 @pytest.mark.asyncio
 async def test_acompletion_forwards_aws_region_name_to_aresponses():
     bridge = ResponsesToCompletionBridgeHandler()