From dd40897f24d672bee3e0ef5500d8504befa8c5b4 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 3 Sep 2026 18:07:26 -0700 Subject: [PATCH] fix(router): harden fusion accounting and contracts --- cookbook/fusion_models.md | 1 + litellm/constants.py | 3 + litellm/fusion_router.py | 88 ++++++-- litellm/litellm_core_utils/fusion_budget.py | 120 +++++++++++ .../proxy/hooks/proxy_track_cost_callback.py | 38 +++- .../streaming_iterator.py | 5 + litellm/router.py | 105 ++------- .../hooks/test_proxy_track_cost_callback.py | 202 ++++++++++++------ tests/test_litellm/test_fusion_router.py | 128 ++++++----- 9 files changed, 470 insertions(+), 220 deletions(-) create mode 100644 litellm/litellm_core_utils/fusion_budget.py diff --git a/cookbook/fusion_models.md b/cookbook/fusion_models.md index 0a95b1bd770..1d641ab6b20 100644 --- a/cookbook/fusion_models.md +++ b/cookbook/fusion_models.md @@ -49,6 +49,7 @@ The outer model must support function calling. Panel and analyst models only nee - The outer model is the only hard health dependency. Panel failures degrade into tool-result data, and analyst failure degrades to raw responses. - Initial outer, panel, analyst, continuation, and search calls are marked separately in spend logs. They inherit the caller identity and remain part of one logical Fusion request. +- Client-visible `usage` describes the outer response returned to that client. Hidden panel, analyst, and search usage remains in its separately tagged spend-log rows; budget reconciliation includes the cost of every hidden call rather than merging heterogeneous model tokens into one public token count. - Admission control reserves the worst-case model-call cost. Hidden calls accumulate against that shared reservation, and the direct initial response or final continuation reconciles it once. This keeps concurrent requests from spending the same remaining budget while Fusion is still running. - Chat-completion streaming is buffered until LiteLLM knows whether the private tool was invoked. A direct response is replayed as a normal stream; a Fusion invocation suppresses the private tool-call stream and exposes only the final outer-model stream. - A request-level `tool_choice: required` is considered satisfied when Fusion runs. The continuation changes it to `auto` when client tools exist, or removes it when they do not, so the outer model can finish instead of being forced into a second tool call. diff --git a/litellm/constants.py b/litellm/constants.py index 1f62347eebd..47c57c66c33 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1460,6 +1460,9 @@ FUSION_BUDGET_ACCUMULATED_COST_KEY: Final = "_fusion_accumulated_actual_cost" FUSION_BUDGET_ACCUMULATED_CALL_IDS_KEY: Final = "_fusion_accumulated_call_ids" FUSION_BUDGET_ACTIVE_KEY: Final = "_fusion_logical_request" FUSION_BUDGET_CONTINUATION_STARTED_KEY: Final = "_fusion_continuation_started" +FUSION_BUDGET_PENDING_CALL_IDS_KEY: Final = "_fusion_pending_cost_call_ids" +FUSION_BUDGET_UNPRICED_CALL_IDS_KEY: Final = "_fusion_unpriced_cost_call_ids" +FUSION_BUDGET_CALL_ID_METADATA_KEY: Final = "user_api_key_fusion_budget_call_id" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/fusion_router.py b/litellm/fusion_router.py index 8082dd31cae..af7638b0032 100644 --- a/litellm/fusion_router.py +++ b/litellm/fusion_router.py @@ -14,6 +14,11 @@ from litellm.constants import ( FUSION_BUDGET_CONTINUATION_STARTED_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.fusion_budget import ( + complete_fusion_budget_call, + register_fusion_budget_call, + wait_for_fusion_budget_calls, +) from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata from litellm.router_utils.auto_router_model_naming import StrategyRouterDependency from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage @@ -890,17 +895,25 @@ class FusionRouter: else: try: metadata: Final = _fusion_call_metadata(request_kwargs, FUSION_RESEARCH_CALL_ORIGIN) - result = await self._search( # rebind-ok: orchestration branch state - model=self.config.search_tool_name, - query=query, - # Search routing stores its internal metadata in the newer - # bucket. Passing this as plain `metadata` would let the - # router create a second bucket and hide the Fusion origin - # from spend reconciliation. - litellm_metadata=metadata, - max_tokens_per_page=1024, - _fusion_proxy_auth_required=isinstance(request_kwargs.get("proxy_server_request"), Mapping), - ) + register_fusion_budget_call(metadata) + try: + result = await self._search( # rebind-ok: orchestration branch state + model=self.config.search_tool_name, + query=query, + # Search routing stores its internal metadata in the newer + # bucket. Passing this as plain `metadata` would let the + # router create a second bucket and hide the Fusion origin + # from spend reconciliation. + litellm_metadata=metadata, + max_tokens_per_page=1024, + _fusion_proxy_auth_required=isinstance(request_kwargs.get("proxy_server_request"), Mapping), + ) + except asyncio.CancelledError: + complete_fusion_budget_call(metadata, cost_known=False) + raise + except Exception: + complete_fusion_budget_call(metadata, cost_known=True) + raise if isinstance(result, BaseModel): result = result.model_dump() # rebind-ok: orchestration branch state except Exception as exc: # noqa: BLE001 # provider search failures become advisory tool results @@ -930,6 +943,16 @@ class FusionRouter: ) # rebind-ok: orchestration branch state while True: call_kwargs = dict(kwargs) # mutable-ok: local provider payload + metadata = _optional_object_mapping( # rebind-ok: each loop dispatches a new provider call + call_kwargs.get("metadata") + ) + call_metadata = ( # rebind-ok: each loop dispatches a new provider call + dict(metadata) # mutable-ok: each provider call owns its synchronization token + if metadata is not None + else {} # mutable-ok: provider metadata requires a native mapping + ) + register_fusion_budget_call(call_metadata) + call_kwargs["metadata"] = call_metadata if remaining_searches > 0 and self._search is not None: call_kwargs["tools"] = [ # mutable-ok: local provider payload _research_tool() @@ -941,7 +964,14 @@ class FusionRouter: "model": model, "messages": current_messages, } # mutable-ok: local provider payload - response = await self._completion(model=model, messages=current_messages, stream=False, **call_kwargs) + try: + response = await self._completion(model=model, messages=current_messages, stream=False, **call_kwargs) + except asyncio.CancelledError: + complete_fusion_budget_call(call_metadata, cost_known=False) + raise + except Exception: + complete_fusion_budget_call(call_metadata, cost_known=True) + raise if not isinstance(response, ModelResponse): return response search_calls = _research_tool_calls(response) @@ -981,13 +1011,20 @@ class FusionRouter: ) -> tuple[ModelResponse, FusionReplayStream | None]: kwargs: Final = _outer_kwargs(request_kwargs) kwargs.pop("litellm_metadata", None) - kwargs["metadata"] = _fusion_call_metadata(request_kwargs, FUSION_INITIAL_CALL_ORIGIN) + metadata: Final = _fusion_call_metadata(request_kwargs, FUSION_INITIAL_CALL_ORIGIN) + register_fusion_budget_call(metadata) + kwargs["metadata"] = metadata client_tools: Final = _client_tools(request_kwargs.get("tools")) kwargs["tools"] = [ # mutable-ok: local provider payload *client_tools, _fusion_tool(), ] # mutable-ok: local provider payload if self.config.invocation == "required": + # A forced private deliberation still honors the caller's ban on + # executable tools. Some providers can emit more than the named + # forced tool in one response, so do not expose client schemas here. + if kwargs.get("tool_choice") == "none": + kwargs["tools"] = [_fusion_tool()] # mutable-ok: local provider payload kwargs["tool_choice"] = { # mutable-ok: local provider payload "type": "function", "function": { # mutable-ok: function schema requires a native mapping @@ -1003,13 +1040,20 @@ class FusionRouter: kwargs["tool_choice"] = "auto" elif kwargs.get("tool_choice") is None: kwargs["tool_choice"] = "auto" - response: Final = await self._completion( - model=self.config.outer_model, - messages=messages, - stream=stream, - _fusion_depth=1, - **kwargs, - ) + try: + response: Final = await self._completion( + model=self.config.outer_model, + messages=messages, + stream=stream, + _fusion_depth=1, + **kwargs, + ) + except asyncio.CancelledError: + complete_fusion_budget_call(metadata, cost_known=False) + raise + except Exception: + complete_fusion_budget_call(metadata, cost_known=True) + raise if isinstance(response, ModelResponse): sanitized_response, _ = _without_mixed_fusion_tool_call(response) return sanitized_response, None @@ -1216,6 +1260,10 @@ class FusionRouter: final_kwargs["metadata"] = final_metadata reservation: Final = final_metadata.get(_BUDGET_RESERVATION_METADATA_KEY) if isinstance(reservation, dict): + # Do not remove this barrier: success cost callbacks run on the + # background logging worker. Finalizing the continuation before all + # registered hidden calls report would omit late panel/search spend. + await wait_for_fusion_budget_calls(final_metadata) # Cancellation accounting can now distinguish an in-flight final # outer call from cancellation while the private panel was running. reservation[FUSION_BUDGET_CONTINUATION_STARTED_KEY] = True diff --git a/litellm/litellm_core_utils/fusion_budget.py b/litellm/litellm_core_utils/fusion_budget.py new file mode 100644 index 00000000000..06550dc4b9d --- /dev/null +++ b/litellm/litellm_core_utils/fusion_budget.py @@ -0,0 +1,120 @@ +"""Synchronize Fusion's hidden provider calls with its shared budget reservation.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping +from typing import Final +from uuid import uuid4 + +from litellm.constants import ( + FUSION_BUDGET_ACTIVE_KEY, + FUSION_BUDGET_CALL_ID_METADATA_KEY, + FUSION_BUDGET_PENDING_CALL_IDS_KEY, + FUSION_BUDGET_UNPRICED_CALL_IDS_KEY, +) + +_BUDGET_RESERVATION_METADATA_KEY: Final = "user_api_key_budget_reservation" +_CALLBACK_POLL_INTERVAL_SECONDS: Final = 0.01 +_CALLBACK_WAIT_TIMEOUT_SECONDS: Final = 5.0 + + +def _fusion_budget_reservation( + metadata: Mapping[str, object], +) -> dict[str, object] | None: # mutable-ok: callers coordinate through one shared request ledger + reservation: Final = metadata.get(_BUDGET_RESERVATION_METADATA_KEY) + if not isinstance(reservation, dict) or reservation.get(FUSION_BUDGET_ACTIVE_KEY) is not True: + return None + return reservation + + +def register_fusion_budget_call(metadata: dict[str, object]) -> None: # mutable-ok: request metadata owns token + """Register one hidden logical call before it is dispatched to a provider.""" + reservation: Final = _fusion_budget_reservation(metadata) + if reservation is None: + return + token: Final = uuid4().hex + pending: Final = reservation.setdefault( + FUSION_BUDGET_PENDING_CALL_IDS_KEY, + [], # mutable-ok: request-scoped pending-call ledger + ) + if not isinstance(pending, list): + return + pending.append(token) + metadata[FUSION_BUDGET_CALL_ID_METADATA_KEY] = token # rebind-ok: caller passes mutable request metadata + + +def complete_fusion_budget_call( + metadata: Mapping[str, object], + *, + cost_known: bool, +) -> None: + """Finish a registered call once its cost was recorded, or mark it conservatively unknown.""" + reservation: Final = _fusion_budget_reservation(metadata) + token: Final = metadata.get(FUSION_BUDGET_CALL_ID_METADATA_KEY) + if reservation is None or not isinstance(token, str): + return + pending: Final = reservation.get(FUSION_BUDGET_PENDING_CALL_IDS_KEY) + if not isinstance(pending, list): + return + try: + pending.remove(token) + except ValueError: + return + if cost_known: + return + unpriced: Final = reservation.setdefault( + FUSION_BUDGET_UNPRICED_CALL_IDS_KEY, + [], # mutable-ok: request-scoped conservative-cost ledger + ) + if isinstance(unpriced, list): + unpriced.append(token) + + +async def wait_for_fusion_budget_calls( + metadata: Mapping[str, object], + *, + timeout_seconds: float = _CALLBACK_WAIT_TIMEOUT_SECONDS, +) -> None: + """Wait for registered hidden costs, falling back to the reserved maximum on timeout.""" + reservation: Final = _fusion_budget_reservation(metadata) + if reservation is None: + return + loop: Final = asyncio.get_running_loop() + deadline: Final = loop.time() + timeout_seconds + while True: + pending = reservation.get( # rebind-ok: poll the shared ledger until every callback reports + FUSION_BUDGET_PENDING_CALL_IDS_KEY + ) + if not isinstance(pending, list) or not pending: + return + if loop.time() >= deadline: + unresolved = tuple( # rebind-ok: timeout snapshot is local to this polling iteration + token for token in pending if isinstance(token, str) + ) + pending.clear() + unpriced = reservation.setdefault( # rebind-ok: timeout ledger is read only in this iteration + FUSION_BUDGET_UNPRICED_CALL_IDS_KEY, + [], # mutable-ok: timeout converts unresolved calls to a conservative charge + ) + if isinstance(unpriced, list): + unpriced.extend(token for token in unresolved if token not in unpriced) + return + await asyncio.sleep(_CALLBACK_POLL_INTERVAL_SECONDS) + + +def fusion_budget_reconciliation_cost( + budget_reservation: Mapping[str, object], + known_cost: float, +) -> float: + """Use actual cost normally and the pre-call maximum if any hidden cost stayed unknown.""" + unpriced: Final = budget_reservation.get(FUSION_BUDGET_UNPRICED_CALL_IDS_KEY) + if not isinstance(unpriced, list) or not unpriced: + return known_cost + reserved_cost_value: Final = budget_reservation.get("reserved_cost") + reserved_cost: Final = ( + float(reserved_cost_value) + if isinstance(reserved_cost_value, (int, float)) and not isinstance(reserved_cost_value, bool) + else 0.0 + ) + return max(known_cost, reserved_cost) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 4bdeb46fb1c..ae96e34120c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -18,6 +18,11 @@ from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_litellm_metadata_from_kwargs, ) +from litellm.litellm_core_utils.fusion_budget import ( + complete_fusion_budget_call, + fusion_budget_reconciliation_cost, + wait_for_fusion_budget_calls, +) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.proxy._types import UserAPIKeyAuth @@ -337,6 +342,7 @@ class _ProxyDBLogger(CustomLogger): ) verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") + metadata: dict = {} # mutable-ok: populated from the callback payload before final token cleanup try: verbose_proxy_logger.debug( "kwargs stream: %s + complete streaming response: %s", @@ -346,7 +352,9 @@ class _ProxyDBLogger(CustomLogger): parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs=kwargs) litellm_params: Final = kwargs.get("litellm_params", {}) or {} end_user_id: Final = get_end_user_id_for_cost_tracking(litellm_params) - metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) + metadata = get_litellm_metadata_from_kwargs( # rebind-ok: callback payload supplies request metadata + kwargs=kwargs + ) # Only fetch key details when user_id wasn't already populated (e.g. direct MCP REST calls). # Avoids a cache/DB lookup on every normal LLM request. if metadata.get("user_api_key") and not metadata.get("user_api_key_user_id"): @@ -356,6 +364,14 @@ class _ProxyDBLogger(CustomLogger): ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) + if ( + budget_reservation is not None + and budget_reservation.get(FUSION_BUDGET_ACTIVE_KEY) is True + and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == "fusion_continuation" + ): + # Defense in depth for continuations dispatched outside the + # FusionRouter helper or while its local wait timed out. + await wait_for_fusion_budget_calls(metadata) user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) @@ -393,8 +409,20 @@ class _ProxyDBLogger(CustomLogger): if defer_fusion_reconciliation and budget_reservation is not None else True ) + # Completing the token immediately after accumulation is + # sufficient for the final-call barrier; persistence and alerts + # can continue without delaying the model orchestration. + complete_fusion_budget_call(metadata, cost_known=True) + known_fusion_cost: Final = float(response_cost) + ( + float(budget_reservation.get(FUSION_BUDGET_ACCUMULATED_COST_KEY) or 0.0) + if budget_reservation is not None + else 0.0 + ) budget_counter_response_cost: Final = ( - float(response_cost) + float(budget_reservation.get(FUSION_BUDGET_ACCUMULATED_COST_KEY) or 0.0) + fusion_budget_reconciliation_cost( + budget_reservation=budget_reservation, + known_cost=known_fusion_cost, + ) if budget_reservation is not None and budget_reservation.get(FUSION_BUDGET_ACTIVE_KEY) is True and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == "fusion_continuation" @@ -466,6 +494,7 @@ class _ProxyDBLogger(CustomLogger): elif budget_reservation is not None and not defer_fusion_reconciliation: await _release_budget_reservation(budget_reservation=budget_reservation) else: + complete_fusion_budget_call(metadata, cost_known=False) if _is_unbilled_interaction_response(completion_response): if BACKGROUND_INTERACTION_COST_POLLING_ENABLED and _is_unbilled_in_progress_interaction( completion_response @@ -526,6 +555,11 @@ class _ProxyDBLogger(CustomLogger): ) spend_log_error("Error in tracking cost callback - %s", str(e), exc=e) + finally: + # If metadata/cost processing itself failed, unblock the final call + # but force it to retain the reservation's conservative maximum. + # The helper is idempotent when the token already completed above. + complete_fusion_budget_call(metadata, cost_known=False) @staticmethod async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict: diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index d7f8cd8f8bd..9a2884db345 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1190,6 +1190,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): responses_api_response.output = list(self._output_with_streamed_item_ids(responses_api_response)) + # Keep the terminal snapshot consistent with response.created. The + # assembled chat chunks can contain a resolved provider model name, + # while this Responses stream represents the model the caller used. + responses_api_response.model = self.model + # Encode the response ID to match non-streaming behavior encoded_response: Final = self._with_encoded_response_id(responses_api_response) diff --git a/litellm/router.py b/litellm/router.py index 3ae9712b3f5..801081874f4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -28,7 +28,6 @@ from collections.abc import ( Generator, Iterator, Mapping, - MutableMapping, Sequence, ) from functools import lru_cache, partial @@ -77,7 +76,6 @@ 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 from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider -from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.ptu_pricing import ( PTU_COST_ATTRIBUTION_ENV_VAR, @@ -2552,10 +2550,7 @@ class Router: model=model, llm_provider="", ) - await self._authorize_fusion_dependencies( - fusion_router=fusion_router, - request_kwargs=kwargs, - ) + self._validate_fusion_proxy_context(request_kwargs=kwargs) response = ( # rebind-ok: one mutually exclusive dispatch branch assigns it await fusion_router.acompletion( messages=messages, @@ -9477,14 +9472,19 @@ class Router: search=self._fusion_asearch, ) - async def _authorize_fusion_dependencies( - self, - fusion_router: FusionRouter, - request_kwargs: MutableMapping[ # mutable-ok: request-local carrier is enriched before hidden calls dispatch - str, object - ], + @staticmethod + def _validate_fusion_proxy_context( + request_kwargs: Mapping[str, object], ) -> None: - """Apply the originating proxy caller's model access and group budgets to every hidden call.""" + """Validate proxy identity before dispatching administrator-configured dependencies. + + The public Fusion model is the authorization boundary. Its outer, panel, + analyst, and search dependencies come from proxy configuration rather + than request input, so requiring the caller to access them directly + would break the virtual-model contract. Caller identity and the Fusion + model's matched access groups remain in forwarded metadata for spend and + reservation accounting. + """ metadata_values: Final = tuple(request_kwargs.get(key) for key in ("litellm_metadata", "metadata")) raw_user_api_key_auth: Final = next( ( @@ -9499,14 +9499,10 @@ class Router: return from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth - from litellm.proxy.auth.auth_checks import can_key_call_resolved_model try: - user_api_key_auth: Final = ( - raw_user_api_key_auth - if isinstance(raw_user_api_key_auth, UserAPIKeyAuth) - else UserAPIKeyAuth.model_validate(raw_user_api_key_auth) - ) + if not isinstance(raw_user_api_key_auth, UserAPIKeyAuth): + UserAPIKeyAuth.model_validate(raw_user_api_key_auth) except ValidationError as exc: raise ProxyException( message="Fusion model authorization context is missing or invalid", @@ -9515,61 +9511,6 @@ class Router: code=403, ) from exc - dependency_models: Final = tuple( - dict.fromkeys( - ( - fusion_router.config.outer_model, - *fusion_router.config.panel_models, - fusion_router.config.resolved_analyst_model, - ) - ) - ) - matched_dependency_groups: Final = await asyncio.gather( - *( - can_key_call_resolved_model( - model=dependency_model, - llm_model_list=self.model_list, - valid_token=user_api_key_auth, - llm_router=self, - ) - for dependency_model in dependency_models - ) - ) - dependency_access_groups: Final = frozenset( - group for matched_groups in matched_dependency_groups for group in matched_groups - ) - if not dependency_access_groups: - return - - # Keep the group gating the virtual Fusion model and add every group - # gating a hidden dependency. The normal spend writer narrows this - # authorization upper bound to groups serving each provider call. - all_groups: Final = tuple( - dict.fromkeys( - ( - *(user_api_key_auth.matched_model_access_groups or ()), - *sorted(dependency_access_groups), - ) - ) - ) - user_api_key_auth.matched_model_access_groups = list( # mutable-ok: auth schema requires a list carrier - all_groups - ) - for metadata_key in ("litellm_metadata", "metadata"): - metadata = request_kwargs.get(metadata_key) - if not isinstance(metadata, Mapping): - continue - mutable_metadata = ( - metadata - if isinstance(metadata, dict) - else dict(metadata) # mutable-ok: SDK metadata boundary requires a native mapping - ) - mutable_metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list( # mutable-ok: metadata JSON requires a list - all_groups - ) - mutable_metadata["user_api_key_auth"] = user_api_key_auth - request_kwargs[metadata_key] = mutable_metadata # rebind-ok: enrich request-local metadata for dispatch - async def _fusion_asearch( # kwargs-ok: bridge preserves the Router.asearch keyword surface self, *, @@ -9577,7 +9518,7 @@ class Router: query: str, **kwargs: object, # kwargs-ok: SDK passthrough ) -> object: - """Late-bound Search API bridge with the originating caller's permissions.""" + """Late-bound Search API bridge for an administrator-configured dependency.""" metadata_values: Final = tuple(kwargs.get(key) for key in ("litellm_metadata", "metadata")) raw_user_api_key_auth: Final = next( ( @@ -9592,16 +9533,10 @@ class Router: proxy_auth_required: Final = kwargs.pop("_fusion_proxy_auth_required", False) is True if raw_user_api_key_auth is not None or proxy_auth_required: from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth - from litellm.proxy.search_endpoints.endpoints import ( - authorize_search_tool_call, - ) try: - user_api_key_auth: Final = ( - raw_user_api_key_auth - if isinstance(raw_user_api_key_auth, UserAPIKeyAuth) - else UserAPIKeyAuth.model_validate(raw_user_api_key_auth) - ) + if not isinstance(raw_user_api_key_auth, UserAPIKeyAuth): + UserAPIKeyAuth.model_validate(raw_user_api_key_auth) except ValidationError as exc: raise ProxyException( message="Fusion Search Tool authorization context is missing or invalid", @@ -9609,10 +9544,6 @@ class Router: param=None, code=403, ) from exc - await authorize_search_tool_call( - search_tool_name=model, - user_api_key_dict=user_api_key_auth, - ) return await self.asearch( model=model, query=query, diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 4f63f37c7a5..44bab5becdb 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,4 +1,3 @@ - import asyncio from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -6,7 +5,18 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm -from litellm.constants import FUSION_BUDGET_ACCUMULATED_COST_KEY, FUSION_BUDGET_ACTIVE_KEY +from litellm.constants import ( + FUSION_BUDGET_ACCUMULATED_COST_KEY, + FUSION_BUDGET_ACTIVE_KEY, + FUSION_BUDGET_PENDING_CALL_IDS_KEY, + FUSION_BUDGET_UNPRICED_CALL_IDS_KEY, +) +from litellm.litellm_core_utils.fusion_budget import ( + complete_fusion_budget_call, + fusion_budget_reconciliation_cost, + register_fusion_budget_call, + wait_for_fusion_budget_calls, +) from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( @@ -75,9 +85,7 @@ async def test_async_post_call_failure_hook(): # Check that metadata was properly updated assert "litellm_params" in call_args["kwargs"] - assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == { - "request_id": "test_request_id" - } + assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"} metadata = call_args["kwargs"]["litellm_params"]["metadata"] assert metadata["user_api_key"] == "test_api_key" assert metadata["status"] == "failure" @@ -341,9 +349,7 @@ async def test_should_continue_failure_tracking_when_budget_release_fails(): ) assert mock_invalidate_budget_reservation_counters.await_count == 1 assert ( - mock_invalidate_budget_reservation_counters.await_args.kwargs[ - "budget_reservation" - ] + mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["finalized"] is True @@ -438,36 +444,21 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): "entries": [{"counter_key": "spend:key:test_api_key"}], } + assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} - ) - is None - ) - assert ( - _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": UserAPIKeyAuth( - budget_reservation=budget_reservation - ) - } + metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)} ) == budget_reservation ) assert ( _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": dict( - UserAPIKeyAuth(budget_reservation=budget_reservation) - ) - } + metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))} ) == budget_reservation ) assert ( - _get_budget_reservation_from_metadata( - metadata={"user_api_key_budget_reservation": budget_reservation} - ) + _get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation}) is budget_reservation ) @@ -475,9 +466,7 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=Exception("db unavailable") - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable")) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -513,9 +502,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): proxy_logging_obj = MagicMock() db_exception = RuntimeError("db unavailable") - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=db_exception - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -559,12 +546,8 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re budget_reservation=budget_reservation, ) assert mock_log_exception.call_count == 2 - mock_log_exception.assert_any_call( - "Failed to release budget reservation after database update failed" - ) - mock_log_exception.assert_any_call( - "Failed to invalidate budget reservation counters after release failed" - ) + mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed") + mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed") increment_spend_counters.assert_not_awaited() @@ -749,6 +732,105 @@ async def test_fusion_hidden_costs_accumulate_then_continuation_reconciles_once( assert increment.await_args.kwargs["budget_reservation"] is reservation +@pytest.mark.asyncio +async def test_fusion_continuation_waits_for_delayed_hidden_cost_callback(): + logger = _ProxyDBLogger() + reservation = { + "reserved_cost": 1.0, + "entries": [], + "finalized": False, + FUSION_BUDGET_ACTIVE_KEY: True, + } + panel_metadata = { + "user_api_key": "hashed-key", + "user_api_key_user_id": "user-1", + "internal_call_origin": "fusion_panel", + "user_api_key_budget_reservation": reservation, + } + register_fusion_budget_call(panel_metadata) + + def kwargs_for(origin: str, response_cost: float, call_id: str, metadata: dict) -> dict: + return { + "call_type": "acompletion", + "model": "test-model", + "litellm_call_id": call_id, + "litellm_params": { + "metadata": { + "user_api_key": "hashed-key", + "user_api_key_user_id": "user-1", + "internal_call_origin": origin, + "user_api_key_budget_reservation": reservation, + **metadata, + } + }, + "standard_logging_object": { + "response_cost": response_cost, + "request_tags": [], + "metadata": {}, + }, + } + + with ( + patch( # test-quality-ok: observes the aggregate passed to the real counter boundary + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as increment, + patch( # test-quality-ok: isolates deployment-group accounting from the ordering assertion + "litellm.proxy.proxy_server.increment_fusion_model_access_group_spend_counters", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: avoids unrelated cache work in the callback race test + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: injects the callback persistence boundary + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as proxy_logging, + ): + proxy_logging.db_spend_update_writer.update_database = AsyncMock() + proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + + continuation_task = asyncio.create_task( + logger._PROXY_track_cost_callback( + kwargs=kwargs_for("fusion_continuation", 0.4, "continuation-call", {}), + completion_response=litellm.ModelResponse(choices=[]), + start_time=datetime.now(), + end_time=datetime.now(), + ) + ) + await asyncio.sleep(0) + assert not continuation_task.done() + increment.assert_not_awaited() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs_for("fusion_panel", 0.2, "panel-call", panel_metadata), + completion_response=litellm.ModelResponse(choices=[]), + start_time=datetime.now(), + end_time=datetime.now(), + ) + await asyncio.wait_for(continuation_task, timeout=1) + + increment.assert_awaited_once() + assert increment.await_args.kwargs["response_cost"] == pytest.approx(0.6) + + +@pytest.mark.asyncio +async def test_fusion_budget_callback_timeout_uses_reserved_maximum(): + reservation = { + "reserved_cost": 1.0, + FUSION_BUDGET_ACTIVE_KEY: True, + } + metadata = {"user_api_key_budget_reservation": reservation} + register_fusion_budget_call(metadata) + + await wait_for_fusion_budget_calls(metadata, timeout_seconds=0) + + assert reservation[FUSION_BUDGET_PENDING_CALL_IDS_KEY] == [] + assert len(reservation[FUSION_BUDGET_UNPRICED_CALL_IDS_KEY]) == 1 + assert fusion_budget_reconciliation_cost(reservation, known_cost=0.4) == pytest.approx(1.0) + # A callback that arrives after the timeout cannot undo the conservative marker. + complete_fusion_budget_call(metadata, cost_known=True) + assert len(reservation[FUSION_BUDGET_UNPRICED_CALL_IDS_KEY]) == 1 + + def test_mixed_fusion_and_client_tool_calls_reconcile_on_the_initial_response(): def response_with_tools(*tool_names: str) -> litellm.ModelResponse: return litellm.ModelResponse( @@ -815,9 +897,15 @@ async def test_cached_fusion_hidden_call_accumulates_zero_cost(): } with ( - patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as increment, # test-quality-ok: isolates proxy persistence while reservation state remains observable - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), # test-quality-ok: isolates proxy persistence while reservation state remains observable - patch("litellm.proxy.proxy_server.proxy_logging_obj") as proxy_logging, # test-quality-ok: injects the callback persistence boundary + patch( # test-quality-ok: isolates proxy persistence while reservation state remains observable + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as increment, + patch( # test-quality-ok: isolates proxy persistence while reservation state remains observable + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: injects the callback persistence boundary + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as proxy_logging, ): proxy_logging.db_spend_update_writer.update_database = AsyncMock() proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() @@ -863,7 +951,9 @@ async def test_unpriced_fusion_hidden_call_does_not_release_parent_reservation() } with ( - patch("litellm.proxy.proxy_server.proxy_logging_obj") as proxy_logging, # test-quality-ok: injects the callback alert boundary + patch( # test-quality-ok: injects the callback alert boundary + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as proxy_logging, patch( # test-quality-ok: verifies unpriced hidden calls cannot release the parent reservation "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", new_callable=AsyncMock, @@ -1413,10 +1503,7 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj # standard_logging_object should have been propagated from logging obj assert call_kwargs.get("standard_logging_object") is not None - assert ( - call_kwargs["standard_logging_object"]["trace_id"] - == "trace-id-from-logging-obj" - ) + assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj" # litellm_trace_id should also be propagated as a fallback assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" @@ -2003,9 +2090,7 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): "metadata": {}, "proxy_server_request": {"request_id": "rid"}, "response_cost": 3.5e-05, - "combined_usage_object": Usage( - prompt_tokens=30, completion_tokens=1, total_tokens=31 - ), + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), } with patch( @@ -2084,15 +2169,10 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): assert mock_increment.call_args.kwargs["team_id"] == "team-123" assert mock_increment.call_args.kwargs["org_id"] == "org-456" - update_kwargs = ( - mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs - ) + update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs assert update_kwargs["user_id"] == "mcp-user@example.com" assert update_kwargs["team_id"] == "team-123" - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] - == "mcp-user@example.com" - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" @pytest.mark.parametrize( @@ -2140,9 +2220,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ], ) @pytest.mark.asyncio -async def test_track_cost_callback_logs_unauthenticated_pass_through_request( - call_type, expect_spend_log -): +async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log): """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It @@ -2188,9 +2266,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( end_time=datetime.now(), ) - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( - 1 if expect_spend_log else 0 - ) + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0) class _FakeDeploymentLookup: diff --git a/tests/test_litellm/test_fusion_router.py b/tests/test_litellm/test_fusion_router.py index 20402ec5633..263b136e67d 100644 --- a/tests/test_litellm/test_fusion_router.py +++ b/tests/test_litellm/test_fusion_router.py @@ -18,7 +18,7 @@ from litellm.fusion_router import ( fusion_router_dependencies, validate_fusion_router_write, ) -from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.litellm_core_utils.fusion_budget import complete_fusion_budget_call from litellm.router import Router from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponseStream @@ -94,6 +94,11 @@ class RecordingCompletion: response = self.responses[model].popleft() if isinstance(response, Exception): raise response + metadata = kwargs.get("metadata") + if isinstance(metadata, dict): + # The fake provider has no logging worker, so model the cost + # callback reaching the request-scoped Fusion ledger. + complete_fusion_budget_call(metadata, cost_known=True) return response @@ -206,6 +211,37 @@ async def test_tool_choice_none_allows_private_deliberation_but_never_client_too assert final["tool_choice"] == "none" +@pytest.mark.asyncio +async def test_required_fusion_hides_client_tools_when_tool_choice_is_none() -> None: + completion = RecordingCompletion( + { + "outer": [_fusion_call(), _response("Final answer without a tool call")], + "panel-a": [_response("Panel A")], + "panel-b": [_response("Panel B")], + "analyst": [_response(_analysis())], + } + ) + client_tool = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object"}}, + } + + response = await _router(completion, invocation="required").acompletion( + messages=[{"role": "user", "content": "Explain the forecast without calling tools"}], + stream=False, + request_kwargs={"tools": [client_tool], "tool_choice": "none"}, + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.tool_calls is None + initial = completion.calls[0] + assert [tool["function"]["name"] for tool in initial["tools"]] == [FUSION_TOOL_NAME] + assert initial["tool_choice"] == {"type": "function", "function": {"name": FUSION_TOOL_NAME}} + final = completion.calls[-1] + assert final["tools"] == [client_tool] + assert final["tool_choice"] == "none" + + @pytest.mark.asyncio async def test_named_client_tool_choice_is_preserved_and_bypasses_private_deliberation() -> None: client_call = { @@ -680,6 +716,9 @@ async def test_configured_search_tool_is_private_to_panel_and_analyst() -> None: async def search(**kwargs: object) -> object: search_calls.append(dict(kwargs)) + metadata = kwargs.get("litellm_metadata") + if isinstance(metadata, dict): + complete_fusion_budget_call(metadata, cost_known=True) return {"results": [{"title": "Source", "url": "https://example.com", "snippet": "Evidence"}]} research_call = _response( @@ -999,7 +1038,7 @@ async def test_nested_fusion_dependency_fails_without_recursing() -> None: @pytest.mark.asyncio -async def test_proxy_fusion_authorizes_every_hidden_model(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_proxy_fusion_uses_virtual_model_as_authorization_boundary(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth import auth_checks @@ -1013,59 +1052,41 @@ async def test_proxy_fusion_authorizes_every_hidden_model(monkeypatch: pytest.Mo ) model_list[-1]["litellm_params"]["fusion_router_config"]["analyst_model"] = "analyst" router = Router(model_list=model_list) - authorize = AsyncMock(side_effect=[("outer-budget",), ("panel-budget",), ("analyst-budget",)]) + authorize = AsyncMock(side_effect=AssertionError("hidden dependencies must not be re-authorized")) monkeypatch.setattr(auth_checks, "can_key_call_resolved_model", authorize) - auth = UserAPIKeyAuth(models=["*"], matched_model_access_groups=["fusion-budget"]) + auth = UserAPIKeyAuth(models=["fusion/test"], matched_model_access_groups=["fusion-budget"]) metadata: dict[str, object] = {"user_api_key_auth": auth} - await router._authorize_fusion_dependencies( # pyright: ignore[reportPrivateUsage] - fusion_router=router.fusion_routers["fusion/test"], + router._validate_fusion_proxy_context( # pyright: ignore[reportPrivateUsage] request_kwargs={ "metadata": metadata, "proxy_server_request": {"body": {"model": "fusion/test"}}, }, ) - assert [call.kwargs["model"] for call in authorize.await_args_list] == ["outer", "panel-a", "analyst"] - assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == [ - "fusion-budget", - "analyst-budget", - "outer-budget", - "panel-budget", - ] - assert auth.matched_model_access_groups == [ - "fusion-budget", - "analyst-budget", - "outer-budget", - "panel-budget", - ] + authorize.assert_not_awaited() + assert auth.models == ["fusion/test"] + assert auth.matched_model_access_groups == ["fusion-budget"] @pytest.mark.asyncio -async def test_proxy_fusion_denies_hidden_model_before_any_provider_call(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth - from litellm.proxy.auth import auth_checks +async def test_proxy_fusion_dispatches_with_only_virtual_model_grant(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._types import UserAPIKeyAuth router = Router(model_list=_router_model_list()) - denial = ProxyException( - message="key not allowed to access model", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=403, - ) - monkeypatch.setattr(auth_checks, "can_key_call_resolved_model", AsyncMock(side_effect=denial)) fusion_completion = AsyncMock() + fusion_completion.return_value = _response("Final") monkeypatch.setattr(router.fusion_routers["fusion/test"], "acompletion", fusion_completion) - with pytest.raises(ProxyException, match="key not allowed"): - await router.acompletion( - model="fusion/test", - messages=[{"role": "user", "content": "Answer"}], - metadata={"user_api_key_auth": UserAPIKeyAuth(models=["fusion/test"])}, - proxy_server_request={"body": {"model": "fusion/test"}}, - ) + result = await router.acompletion( + model="fusion/test", + messages=[{"role": "user", "content": "Answer"}], + metadata={"user_api_key_auth": UserAPIKeyAuth(models=["fusion/test"])}, + proxy_server_request={"body": {"model": "fusion/test"}}, + ) - fusion_completion.assert_not_awaited() + assert result.choices[0].message.content == "Final" + fusion_completion.assert_awaited_once() @pytest.mark.asyncio @@ -1128,9 +1149,11 @@ async def test_invoked_fusion_closes_suppressed_initial_stream(monkeypatch: pyte @pytest.mark.asyncio -async def test_fusion_search_checks_proxy_permissions_before_router_search(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_fusion_search_uses_admin_configured_dependency_without_separate_grant( + monkeypatch: pytest.MonkeyPatch, +) -> None: from litellm.models.object_permission import LiteLLM_ObjectPermissionTable - from litellm.proxy._types import LiteLLM_TeamTable, ProxyException, UserAPIKeyAuth + from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth from litellm.proxy.auth import auth_checks router = Router(model_list=[]) @@ -1147,15 +1170,20 @@ async def test_fusion_search_checks_proxy_permissions_before_router_search(monke monkeypatch.setattr(auth_checks, "get_team_object", get_team_object) monkeypatch.setattr(router, "asearch", raw_search) - with pytest.raises(ProxyException, match="Team not allowed to access search tool"): - await router._fusion_asearch( # pyright: ignore[reportPrivateUsage] - model="restricted-search", - query="evidence", - litellm_metadata={"user_api_key_auth": user_api_key_auth.model_dump()}, - ) + result = await router._fusion_asearch( # pyright: ignore[reportPrivateUsage] + model="restricted-search", + query="evidence", + litellm_metadata={"user_api_key_auth": user_api_key_auth.model_dump()}, + _fusion_proxy_auth_required=True, + ) - get_team_object.assert_awaited_once() - raw_search.assert_not_awaited() + assert result == {"results": []} + get_team_object.assert_not_awaited() + raw_search.assert_awaited_once_with( + model="restricted-search", + query="evidence", + litellm_metadata={"user_api_key_auth": user_api_key_auth.model_dump()}, + ) @pytest.mark.asyncio @@ -1271,7 +1299,11 @@ async def test_router_responses_and_anthropic_adapters_stream_direct_outer_respo responses_stream = await router.aresponses(model="fusion/test", input="Answer", stream=True) response_events = [event async for event in responses_stream] - assert any(str(getattr(event, "type", "")).endswith("RESPONSE_COMPLETED") for event in response_events) + completed_events = [ + event for event in response_events if str(getattr(event, "type", "")).endswith("RESPONSE_COMPLETED") + ] + assert completed_events + assert all(event.response.model == "fusion/test" for event in completed_events) anthropic_stream = await router.aanthropic_messages( model="fusion/test",