fix(router): harden fusion accounting and contracts

This commit is contained in:
moe-berri 2026-09-03 18:07:26 -07:00
parent 123f486d20
commit dd40897f24
9 changed files with 470 additions and 220 deletions

View file

@ -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.

View file

@ -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. "

View file

@ -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

View file

@ -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)

View file

@ -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:

View file

@ -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)

View file

@ -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,

View file

@ -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:

View file

@ -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",