From 2e138511ed12dbfc6b6bf15300aea2488b674abe Mon Sep 17 00:00:00 2001
From: Yucheng He <yucheng@berri.ai>
Date: Mon, 7 Sep 2026 02:06:06 -0700
Subject: [PATCH] test(router): cover the fallback header adoption helper
 directly

The router_code_coverage gate wants every router.py function named in a
router test, and this also pins the weak-reference behavior: a wrapper
collected mid-stream must not break the generator still draining it.
---
 tests/test_litellm/test_router.py | 38 +++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 7d93758d4ad..70ebdc4766c 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -2523,6 +2523,44 @@ def test_adopt_fallback_response_headers_replaces_rather_than_merges():
     assert wrapper._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"}
 
 
+def test_adopt_fallback_response_headers_survives_a_collected_wrapper():
+    """LIT-6767: the fallback generator holds only a weak reference to its wrapper.
+
+    A client that disconnects mid-stream can drop the wrapper while the generator is
+    still draining, and the adoption call has to keep working with nothing to adopt into
+    so the deployment slot is still released.
+    """
+    import weakref
+    from unittest.mock import MagicMock
+
+    from litellm.router import FallbackAwareStreamWrapper, Router
+
+    fallback: Final = MagicMock()
+    fallback._response_headers = {"x-request-id": "req-FALLBACK"}
+    fallback._hidden_params = {
+        "model_id": "fallback-deployment",
+        "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
+    }
+
+    wrapper = FallbackAwareStreamWrapper(
+        completion_stream=iter([]),
+        model="gpt-4",
+        custom_llm_provider="openai",
+        logging_obj=MagicMock(),
+    )
+    live_ref: Final = weakref.ref(wrapper)
+    prepared: Final = Router._adopt_fallback_response_headers(live_ref, fallback)
+    assert prepared == (fallback._hidden_params, fallback._hidden_params["additional_headers"])
+    assert wrapper.fallback_headers_adopted is True
+    assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"}
+
+    dead_ref: Final = weakref.ref(wrapper)
+    del wrapper
+    assert dead_ref() is None
+    # no wrapper left to repoint, and the caller still needs the params for the chunks
+    assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared
+
+
 def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_replace():
     """LIT-6767: a fallback that carries no raw provider headers publishes none.