address review: clear a failed build via a done-callback, not the waiter

The previous cleanup only ran inside a caller's own except handler, so a
build that failed after its only caller had already been cancelled left
the failed task cached with nothing left to clear it. Move the cleanup
onto the task itself as a done-callback, which fires whether or not
anyone is still awaiting it, so the next request always gets a fresh
attempt instead of replaying the stale failure.

Adds a regression test for exactly that ordering (cancel the only
caller, let the build fail unobserved, then confirm the next request
builds successfully); it fails against the previous except-based
cleanup, which left the task cached.

litellm/router_strategy/auto_router/auto_router.py

@@ -138,6 +138,17 @@ class AutoRouter(CustomLogger):
         self.routelayer = routelayer
         return routelayer
 
+    def _clear_build_task_on_failure(self, build_task: "asyncio.Task[SemanticRouter]") -> None:
+        """Done-callback: drop a failed build so the next caller gets a fresh attempt.
+
+        Runs whether or not any caller is still awaiting `build_task` (that's the point:
+        a caller cancelled via `cancel_on_disconnect` mid-build must not leave a later
+        failure cached with nothing left to clear it), and `not build_task.cancelled()`
+        guards `.exception()`, which raises on a cancelled task instead of returning one.
+        """
+        if build_task is self._routelayer_build_task and not build_task.cancelled() and build_task.exception():
+            self._routelayer_build_task = None
+
     async def _ensure_routelayer(self) -> "SemanticRouter":
         """Return the cached route layer, building it once under a lock if needed.
 
@@ -158,18 +169,9 @@ class AutoRouter(CustomLogger):
             build_task = self._routelayer_build_task
             if build_task is None:
                 build_task = asyncio.ensure_future(asyncio.to_thread(self._build_routelayer))
+                build_task.add_done_callback(self._clear_build_task_on_failure)
                 self._routelayer_build_task = build_task
-        try:
-            return await asyncio.shield(build_task)
-        except Exception:
-            # Only a real build failure (not this caller's own cancellation, which
-            # asyncio.shield turns into a CancelledError here while the task keeps
-            # running for everyone else) clears the slot, so the next call retries
-            # a fresh build instead of replaying the same failure forever.
-            async with self._routelayer_lock:
-                if self._routelayer_build_task is build_task:
-                    self._routelayer_build_task = None
-            raise
+        return await asyncio.shield(build_task)
 
     @staticmethod
     def _extract_text_from_messages(messages: list[dict[str, Any]]) -> str:

tests/test_litellm/router_strategy/test_auto_router.py

@@ -721,3 +721,62 @@ class TestAutoRouterColdStartDoesNotBlockTheEventLoop:
 
         assert result is not None
         assert len(embedding_router.embedding_call_threads) == _EMBEDDING_CALLS_PER_ROUTELAYER_BUILD
+
+    @pytest.mark.asyncio
+    async def test_should_clear_a_failed_build_even_with_no_caller_left_to_observe_it(self):
+        """Regression: a build that fails after its only caller was already cancelled must
+        still clear the slot, so the next request gets a fresh attempt instead of replaying
+        the same stale failure forever."""
+        import threading
+
+        class FailsOnFirstAttemptEmbeddingRouter(ThreadTrackingEmbeddingRouter):
+            def __init__(self) -> None:
+                super().__init__()
+                self.started = threading.Event()
+                self.release = threading.Event()
+                self.attempts = 0
+
+            def embedding(self, input: list[str], model: str, **kwargs: Any) -> Any:
+                self.attempts += 1
+                attempt = self.attempts
+                self.started.set()
+                self.release.wait(timeout=5)
+                if attempt == 1:
+                    raise ValueError("boom")
+                return super().embedding(input, model, **kwargs)
+
+        embedding_router: Final = FailsOnFirstAttemptEmbeddingRouter()
+        auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router)
+
+        first_call: Final = asyncio.ensure_future(
+            auto_router.async_pre_routing_hook(
+                model="my-auto-router",
+                request_kwargs={},
+                messages=[{"role": "user", "content": "fix this stack trace"}],
+            )
+        )
+        while not embedding_router.started.is_set():
+            await asyncio.sleep(0.01)
+        first_call.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await first_call
+
+        # Nobody awaits the build now. Let the first attempt fail on its own.
+        embedding_router.release.set()
+        build_task = auto_router._routelayer_build_task
+        assert build_task is not None
+        while not build_task.done():
+            await asyncio.sleep(0.01)
+        await asyncio.sleep(0.01)  # let the done-callback (scheduled via call_soon) run
+
+        assert auto_router._routelayer_build_task is None
+
+        embedding_router.started.clear()
+        embedding_router.release.clear()
+        result: Final = await auto_router.async_pre_routing_hook(
+            model="my-auto-router",
+            request_kwargs={},
+            messages=[{"role": "user", "content": "fix this stack trace"}],
+        )
+
+        assert result is not None