From 0ecbfa376b658e06403177fac8b7d36db04abb38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:12:00 -0700 Subject: [PATCH] fix(mcp): hold round one's log for the guardrail, bill the chunk cost once, and keep round one's id When the follow-up call of a streamed chat MCP auto-execute fails, round one's success log now goes through the proxy's deferred guardrail release instead of straight to the handlers, so a post-call guardrail runs on it and the row carries it. Without a proxy release it logs directly as before The last usage chunk now takes its cost from the combined row's billed total, so a fixed cost margin is counted once in the chunk, as it is on the row. The follow-up's usage chunk is held until the follow-up ends to make that possible The combined row now keeps round one's id, the stream's first chunk id, which matches how a single-round stream is logged --- litellm/proxy/common_request_processing.py | 6 +- .../responses/mcp/chat_completions_handler.py | 11 ++- .../mcp/chat_stream_round_billing.py | 82 ++++++++++++------ .../integration/_support/hook_count_probe.py | 18 ++++ .../integration/mcp/test_mcp_llm_endpoints.py | 85 +++++++++++++++++-- .../mcp/test_chat_completions_handler.py | 3 +- 6 files changed, 166 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f5d6417fb79..e011f01c162 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4053,11 +4053,13 @@ class ProxyBaseLLMRequestProcessing: _captured_user_api_key_dict: Final = user_api_key_dict _captured_logging_obj: Final = logging_obj - async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: + async def _on_deferred_stream_complete( + assembled_response: object, cache_hit: object, logged_on: LiteLLMLoggingObj | None = None + ) -> None: await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( captured_data=_captured_data, captured_user_api_key_dict=_captured_user_api_key_dict, - captured_logging_obj=_captured_logging_obj, + captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on, assembled_response=assembled_response, cache_hit=cache_hit, ) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index f35089f2e1f..90438af1487 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -367,8 +367,10 @@ async def acompletion_with_mcp( try: return await self._next_chunk() except StopAsyncIteration: - round_billing.finish() - raise + usage_chunk: Final = round_billing.finish() + if usage_chunk is None: + raise + return usage_chunk except Exception: round_billing.fail() raise @@ -434,7 +436,6 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", chunk) - return chunk if self._round_billing is None else self._round_billing.client_chunk(chunk) except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger @@ -442,6 +443,10 @@ async def acompletion_with_mcp( verbose_logger.debug("Follow-up stream exhausted") # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration raise StopAsyncIteration + client_chunk: Final = ( + chunk if self._round_billing is None else self._round_billing.client_chunk(chunk) + ) + return await self._next_chunk() if client_chunk is None else client_chunk # If we're here and follow_up_stream is None but we expected it, log a warning if ( diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py index 7274652890e..4df79ca13ba 100644 --- a/litellm/responses/mcp/chat_stream_round_billing.py +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Generator +from collections.abc import Callable, Coroutine, Generator from contextlib import contextmanager from dataclasses import dataclass, replace from types import MappingProxyType @@ -14,6 +14,7 @@ from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import ModelResponse, ModelResponseStream, Usage Parked: TypeAlias = tuple[object, ...] +_Release: TypeAlias = Callable[..., Coroutine[object, object, None]] _PARKED_DISPATCH: Final = "_deferred_stream_complete_args" _DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" @@ -55,12 +56,31 @@ class _PricedRound: return BilledRoundsSpend(self.usage, self.price.cost or 0.0, self.price.cost_breakdown) +@dataclass(frozen=True, slots=True) +class _HeldUsageChunk: + chunk: ModelResponseStream + usage: Usage + + @dataclass(frozen=True, slots=True) class _FirstRound: priced: _PricedRound parked: Parked logged_alone: Parked banked: LiteLLMLoggingObj | None = None + usage_chunk: _HeldUsageChunk | None = None + + def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None: + if self.usage_chunk is None: + return None + merged: Final = merged_chat_usage(self.priced.usage, self.usage_chunk.usage) + usage: Final = merged if billed is None else merged.model_copy(update=MappingProxyType({"cost": _cost(billed)})) + return self.usage_chunk.chunk.model_copy(update=MappingProxyType({"usage": usage})) + + +def _cost(response: ModelResponse) -> object: + hidden_params: Final = get_hidden_params(response) + return None if hidden_params is None else hidden_params.get("response_cost") def _assembled_round(parked: Parked | None) -> _AssembledRound | None: @@ -100,27 +120,17 @@ class ChatStreamRoundBilling: self._bank_apart_from_the_follow_up() return None if self._first is None else self._first.priced.spend - def client_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + def client_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: usage: Final[object] = getattr(chunk, "usage", None) if self._first is None or not isinstance(usage, Usage): return chunk - return chunk.model_copy(update=MappingProxyType({"usage": merged_chat_usage(self._first.priced.usage, usage)})) + self._first = replace(self._first, usage_chunk=_HeldUsageChunk(chunk, usage)) + return None - def finish(self) -> None: - parked: Final = self._take_parked() + def finish(self) -> ModelResponseStream | None: first: Final = self._first - self._first = None - if parked is None: - self._disarm() - return - if first is None or parked is first.logged_alone: - self._release(parked) - return - final_round: Final = _assembled_round(parked) - if final_round is None: - self._release(first.logged_alone) - return - self._release(self._billed_for_both_rounds(first.priced, self._priced(final_round))) + billed: Final = self._release_rounds() + return None if first is None else first.client_usage_chunk(billed) def fail(self) -> None: self._bank_apart_from_the_follow_up() @@ -130,7 +140,7 @@ class ChatStreamRoundBilling: return self._first = None setattr(self._logging_obj, _PARKED_DISPATCH, None) - self._log_apart((*first.parked, first.banked)) + self._log_through_the_release((*first.parked, first.banked)) def release_abandoned(self) -> None: if not self._owns_release(): @@ -154,12 +164,31 @@ class ChatStreamRoundBilling: self._first = replace(first, logged_alone=logged_alone, banked=banked) setattr(self._logging_obj, _PARKED_DISPATCH, logged_alone) - def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> Parked: + def _release_rounds(self) -> ModelResponse | None: + parked: Final = self._take_parked() + first: Final = self._first + self._first = None + if parked is None: + self._disarm() + return None + if first is None or parked is first.logged_alone: + self._release(parked) + return None + final_round: Final = _assembled_round(parked) + if final_round is None: + self._release(first.logged_alone) + return None + billed, cache_hit = self._billed_for_both_rounds(first.priced, self._priced(final_round)) + self._release((billed, cache_hit)) + return billed + + def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> tuple[ModelResponse, bool]: first_breakdown: Final = first.price.cost_breakdown final_breakdown: Final = final.price.cost_breakdown - billed: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + rounds: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + billed: Final = rounds.model_copy(update=MappingProxyType({"id": first.response.id})) cache_hit: Final = first.cache_hit and final.cache_hit - set_hidden_params(billed, {**(get_hidden_params(billed) or {}), "cache_hit": cache_hit}) + set_hidden_params(billed, {**(get_hidden_params(rounds) or {}), "cache_hit": cache_hit}) self._logging_obj.cost_breakdown = summed_cost_breakdown( first.response, first_breakdown, final.response, final_breakdown ) @@ -192,12 +221,13 @@ class ChatStreamRoundBilling: if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): setattr(self._logging_obj, _PARKED_DISPATCH, parked) return - setattr(self._logging_obj, _DISPATCH_RELEASE, None) - self._log_apart(parked) + self._log_through_the_release(parked) - def _log_apart(self, parked: Parked) -> None: - loop: Final = asyncio.get_running_loop() - spawn_logging_task(loop, self._log_success(*parked), task_name="Chat completions MCP stream success log") + def _log_through_the_release(self, parked: Parked) -> None: + armed: Final[_Release | None] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + logged: Final = (self._log_success if armed is None else armed)(*parked) + spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log") def _disarm(self) -> None: if self._owns_release(): diff --git a/tests/integration/_support/hook_count_probe.py b/tests/integration/_support/hook_count_probe.py index 0d4056e8fb9..3cd72f1a6ba 100644 --- a/tests/integration/_support/hook_count_probe.py +++ b/tests/integration/_support/hook_count_probe.py @@ -75,6 +75,24 @@ class BlockingPostCallGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail") +class UsageRecordingPostCallGuardrail(CustomGuardrail): + """Records the token usage of the response it is handed in the request's guardrail information, so the spend row + shows which response the guardrail saw.""" + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: object, response: object + ) -> None: + usage: Final = getattr(response, "usage", None) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ + "prompt_tokens": getattr(usage, "prompt_tokens", None), + "completion_tokens": getattr(usage, "completion_tokens", None), + }, + request_data=data, + guardrail_status="success", + ) + + probe: Final = HookCountProbe() sync_only_probe: Final = SyncOnlyHookCountProbe() diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index d0fdc6935c7..c2f8845f0d5 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -45,9 +45,10 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state] -@pytest.fixture(scope="session") -def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: - directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") +@contextmanager +def _private_gateway( + directory: Path, guardrails: Sequence[Mapping[str, JsonValue]] = () +) -> Iterator[tuple[Gateway, Mapping[str, str]]]: stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = directory / "stored-responses.yaml" config.write_text( @@ -59,6 +60,7 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat **object_value(stock["litellm_settings"]), "cost_margin_config": {"hosted_vllm": {"fixed_amount": FIXED_MARGIN}}, }, + **({"guardrails": list(guardrails)} if guardrails else {}), } ) ) @@ -70,15 +72,50 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat yield private, services -@pytest.fixture -def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: - private, services = _llm_gateway +@contextmanager +def _serving(served: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + private, services = served with pytest.MonkeyPatch.context() as environment: for name, value in services.items(): environment.setenv(name, value) yield private +@pytest.fixture(scope="session") +def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-gateway")) as served: + yield served + + +@pytest.fixture +def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_llm_gateway) as private: + yield private + + +ROUND_USAGE_GUARDRAIL: Final = "mcp-llm-round-usage" + + +@pytest.fixture(scope="module") +def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + guardrail: Final[Mapping[str, JsonValue]] = { + "guardrail_name": ROUND_USAGE_GUARDRAIL, + "litellm_params": { + "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail", + "mode": "post_call", + "default_on": True, + }, + } + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), (guardrail,)) as served: + yield served + + +@pytest.fixture +def guarded_gateway(_guarded_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_guarded_llm_gateway) as private: + yield private + + Surface = Literal["chat", "responses", "messages", "messages_bridge"] SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge") ADD: Final = {"a": 2, "b": 3} @@ -973,7 +1010,7 @@ def test_streaming_chat_auto_execute_logs_the_answer_and_bills_both_rounds(gatew rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") assert len(rows) == 1, rows row: Final = rows[0] - assert row["request_id"] == chunks[-1]["id"], (row, chunks[-1]) + assert row["request_id"] == chunks[0]["id"], (row, chunks[0]) assert (row["prompt_tokens"], row["completion_tokens"]) == ( FIRST_ROUND[0] + FOLLOW_UP_ROUND[0], FIRST_ROUND[1] + FOLLOW_UP_ROUND[1], @@ -1001,6 +1038,40 @@ def test_streaming_chat_auto_execute_whose_follow_up_fails_bills_round_one_once( assert [float(str(row["spend"])) for row in failed] == [0.0], failed +def _guardrail_information(metadata: JsonValue) -> tuple[JsonValue, ...]: + entries: Final = object_value(metadata).get("guardrail_information") + return tuple(entries) if isinstance(entries, list) else () + + +def _round_usage_seen_by_the_guardrail(call_id: str) -> list[JsonValue]: + rows: Final = read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "acompletion", "success"), + ) + entries: Final = tuple(chain.from_iterable(_guardrail_information(row["metadata"]) for row in rows)) + return [ + entry.get("guardrail_response") + for entry in entries + if isinstance(entry, dict) and entry.get("guardrail_name") == ROUND_USAGE_GUARDRAIL + ] + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_guardrail_on_round_one( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert _round_usage_seen_by_the_guardrail(call_id) == [ + {"prompt_tokens": FIRST_ROUND[0], "completion_tokens": FIRST_ROUND[1]} + ] + + @pytest.mark.parametrize("include_usage", (True, False), ids=("include-usage", "no-usage")) def test_streaming_chat_auto_execute_answered_in_one_round_ends_with_the_usage_it_billed( gateway: Gateway, include_usage: bool diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ce681ef7e6f..23ff0727098 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1823,13 +1823,14 @@ def _assert_one_stream_row_bills_both_rounds_once( charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN assert len(logged) == 1 row: Final = logged[0] - assert row["id"] == chunks[-1].id + assert row["id"] == chunks[0].id assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 assert row["response_cost"] == pytest.approx(charged_once) assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) usage: Final = _streamed_usage(chunks) assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71 + 149, 19 + 5, 90 + 154) + assert usage.cost == pytest.approx(row["response_cost"]) def _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> None: