From f5c1c82f8175e42491c364dca3208eb5151b815b Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 00:32:24 +0000 Subject: [PATCH 01/12] fix(responses): estimate usage from text when streamed completed event omits usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 15 +++ litellm/responses/utils.py | 16 ++++ .../responses/test_streaming_iterator.py | 96 ++++++++++++++++++- 3 files changed, 124 insertions(+), 3 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 38874768ca8..52b50167190 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -420,6 +420,21 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): self.completed_response = openai_responses_api_chunk + _response_obj: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) + if ( + _chunk_type + in ( + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + ) + and _response_obj is not None + and _response_obj.usage is None + ): + _response_obj.usage = ResponseAPILoggingUtils.estimate_usage_from_text( + self.model or "", self.request_data.get("input"), self._generated_content + ) _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 41a3ded7022..e523b3771af 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,4 +1,5 @@ import base64 +import json import re from collections.abc import Iterable, Mapping, Sequence from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload @@ -1239,3 +1240,18 @@ class ResponseAPILoggingUtils: setattr(chat_usage, "cost", response_api_usage.cost) return chat_usage + + @staticmethod + def estimate_usage_from_text(model: str, request_input: object, generated_text: str) -> ResponseAPIUsage: + input_text: Final = request_input if isinstance(request_input, str) else json.dumps(request_input, default=str) + input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped + model=model, text=input_text + ) + output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped + model=model, text=generated_text, count_response_tokens=True + ) + return ResponseAPIUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + ) diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 5e0e794d93e..c593946ac4a 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -31,8 +31,15 @@ def _sse_event(payload: dict) -> bytes: def _mock_config() -> Mock: mock_config = Mock(spec=BaseResponsesAPIConfig) - mock_responses_api_response = Mock(spec=ResponsesAPIResponse) - mock_responses_api_response.id = "resp_ttft" + mock_responses_api_response = ResponsesAPIResponse( + id="resp_ttft", + created_at=0, + status="completed", + model="gpt-4o-mini", + object="response", + output=[], + usage=ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2), + ) def _transform(model, parsed_chunk, logging_obj): evt_type = parsed_chunk.get("type") @@ -54,6 +61,8 @@ def _make_iterator( sse_events: list[bytes], logging_obj: LiteLLMLoggingObj, trailing_error: Optional[Exception] = None, + config: Mock | None = None, + request_data: dict | None = None, ) -> ResponsesAPIStreamingIterator: async def aiter_bytes(): for evt in sse_events: @@ -68,10 +77,11 @@ def _make_iterator( return ResponsesAPIStreamingIterator( response=mock_response, model="gpt-4o-mini", - responses_api_provider_config=_mock_config(), + responses_api_provider_config=config or _mock_config(), logging_obj=logging_obj, litellm_metadata={}, custom_llm_provider="openai", + request_data=request_data, ) @@ -329,6 +339,86 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params +def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock: + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + evt_type = parsed_chunk.get("type") + if evt_type == "response.completed": + completed = Mock(spec=ResponseCompletedEvent) + completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED + completed.response = response + return completed + stub = Mock() + stub.type = evt_type + if evt_type == "response.output_text.delta": + stub.delta = parsed_chunk.get("delta") + return stub + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +def _responses_api_response_without_usage() -> ResponsesAPIResponse: + return ResponsesAPIResponse( + id="resp_no_usage", + created_at=int(datetime(2025, 1, 1).timestamp()), + status="completed", + model="gpt-4o-mini", + object="response", + output=[], + usage=None, + ) + + +@pytest.mark.asyncio +async def test_completed_event_without_usage_gets_text_estimate(): + """A response.completed event carrying usage: null still bills: the + iterator estimates usage from the request input and generated text.""" + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "hello world"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "count these input tokens please"}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage is not None + assert usage.input_tokens > 0 + assert usage.output_tokens > 0 + assert usage.total_tokens == usage.input_tokens + usage.output_tokens + + +@pytest.mark.asyncio +async def test_completed_event_with_usage_is_left_untouched(): + """Provider-reported usage on response.completed wins over the estimate.""" + response = _responses_api_response_with_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "hello world"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "count these input tokens please"}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage.input_tokens == 20 + assert usage.output_tokens == 60 + assert usage.total_tokens == 80 + + def _responses_api_response_with_usage() -> ResponsesAPIResponse: return ResponsesAPIResponse( id="resp_lit6427", From de9aa48cd62e73b5b291d6c5a2a952dd35f09f00 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:02:29 +0000 Subject: [PATCH 02/12] fix(responses): count multimodal input and tool-call output in the streamed usage fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 38 ++++++++- litellm/responses/utils.py | 16 ---- .../responses/test_streaming_iterator.py | 78 ++++++++++++++++++- 3 files changed, 112 insertions(+), 20 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 52b50167190..6cb0331620f 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -32,6 +32,9 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ) from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.integrations.custom_logger import converted_stream_requested from litellm.types.llms.openai import ( @@ -257,6 +260,7 @@ class BaseResponsesAPIStreamingIterator: self._failure_handled = False # Track if failure handler has been called self._yielded_first_chunk = False self._generated_content = "" + self._generated_tool_arguments = "" self._completed_response_cached = False self._completed_response_logged = False self._completed_response_cache_hit: bool | None = None @@ -352,6 +356,10 @@ class BaseResponsesAPIStreamingIterator: _delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_delta, str): self._generated_content += _delta + elif _event_type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA: + _args_delta: Final = getattr(openai_responses_api_chunk, "delta", None) + if isinstance(_args_delta, str): + self._generated_tool_arguments += _args_delta _stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata) if _event_type in ( ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -432,8 +440,11 @@ class BaseResponsesAPIStreamingIterator: and _response_obj is not None and _response_obj.usage is None ): - _response_obj.usage = ResponseAPILoggingUtils.estimate_usage_from_text( - self.model or "", self.request_data.get("input"), self._generated_content + _response_obj.usage = _estimate_usage_from_text( + self.model or "", + self.request_data.get("input"), + self.request_data, + self._generated_content + self._generated_tool_arguments, ) _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj) @@ -1347,6 +1358,29 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None: return None +def _estimate_usage_from_text( + model: str, + request_input: object, + responses_api_request: Mapping[str, object], + generated_text: str, +) -> ResponseAPIUsage: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped + input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union + responses_api_request=dict(responses_api_request), + ) + input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped + model=model, messages=messages + ) + output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped + model=model, text=generated_text, count_response_tokens=True + ) + return ResponseAPIUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + ) + + def _stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index e523b3771af..41a3ded7022 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,5 +1,4 @@ import base64 -import json import re from collections.abc import Iterable, Mapping, Sequence from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload @@ -1240,18 +1239,3 @@ class ResponseAPILoggingUtils: setattr(chat_usage, "cost", response_api_usage.cost) return chat_usage - - @staticmethod - def estimate_usage_from_text(model: str, request_input: object, generated_text: str) -> ResponseAPIUsage: - input_text: Final = request_input if isinstance(request_input, str) else json.dumps(request_input, default=str) - input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped - model=model, text=input_text - ) - output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped - model=model, text=generated_text, count_response_tokens=True - ) - return ResponseAPIUsage( - input_tokens=input_tokens, - output_tokens=output_tokens, - total_tokens=input_tokens + output_tokens, - ) diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index c593946ac4a..c3721353d42 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -5,12 +5,13 @@ completion_start_time = end_time.""" import json from datetime import datetime -from typing import Optional +from typing import Final, Optional from unittest.mock import Mock, patch import httpx import pytest +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( @@ -351,8 +352,10 @@ def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock return completed stub = Mock() stub.type = evt_type - if evt_type == "response.output_text.delta": + if "delta" in parsed_chunk: stub.delta = parsed_chunk.get("delta") + if "item" in parsed_chunk: + stub.item = parsed_chunk.get("item") return stub mock_config.transform_streaming_response.side_effect = _transform @@ -718,3 +721,74 @@ async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_val assert isinstance(client_usage, ResponseAPIUsage) assert client_usage.input_tokens == 29 assert client_usage.cost == pytest.approx(0.0001) + + +@pytest.mark.asyncio +async def test_completed_event_without_usage_counts_tool_call_arguments(): + """A function-call-only stream still bills output tokens: streamed + function_call_arguments deltas feed the text estimate.""" + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event( + { + "type": "response.output_item.added", + "item": {"type": "function_call", "name": "get_weather", "call_id": "call_1"}, + } + ), + _sse_event( + { + "type": "response.function_call_arguments.delta", + "delta": '{"location": "San Francisco", "unit": "celsius"}', + } + ), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "what is the weather in san francisco"}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage is not None + assert usage.output_tokens > 0 + assert usage.total_tokens == usage.input_tokens + usage.output_tokens + + +@pytest.mark.asyncio +async def test_completed_event_without_usage_counts_multimodal_input_as_messages(): + """Multimodal request input is counted as chat messages, not as a JSON blob: + a huge base64 image must not inflate the estimated input tokens.""" + image_input: Final = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "what is in this image"}, + { + "type": "input_image", + "image_url": "data:image/png;base64," + "A" * 4000, + }, + ], + } + ] + json_count: Final = litellm.token_counter(model="gpt-4o-mini", text=json.dumps(image_input)) + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "it is a cat"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": image_input}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage is not None + assert usage.input_tokens < json_count / 2 From 1484fd76001fae494c31e186c079c92868ddd0c5 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:13:26 +0000 Subject: [PATCH 03/12] fix(responses): keep the usage estimate best-effort when token counting raises Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 20 +++++++++++++++- .../responses/test_streaming_iterator.py | 24 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6cb0331620f..6cf36e3e660 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -440,7 +440,7 @@ class BaseResponsesAPIStreamingIterator: and _response_obj is not None and _response_obj.usage is None ): - _response_obj.usage = _estimate_usage_from_text( + _response_obj.usage = _estimate_usage_safely( self.model or "", self.request_data.get("input"), self.request_data, @@ -1381,6 +1381,24 @@ def _estimate_usage_from_text( ) +def _estimate_usage_safely( + model: str, + request_input: object, + responses_api_request: Mapping[str, object], + generated_text: str, +) -> ResponseAPIUsage | None: + try: + return _estimate_usage_from_text( + model=model, + request_input=request_input, + responses_api_request=responses_api_request, + generated_text=generated_text, + ) + except Exception as e: + verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e) + return None + + def _stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index c3721353d42..234e63ca3be 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -792,3 +792,27 @@ async def test_completed_event_without_usage_counts_multimodal_input_as_messages usage = iterator.completed_response.response.usage assert usage is not None assert usage.input_tokens < json_count / 2 + + +@pytest.mark.asyncio +async def test_completed_event_survives_a_failing_usage_estimate(): + """A raising token_counter must not break a stream that previously completed: + the estimate is best-effort and falls back to usage None.""" + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "hello world"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "count these input tokens please"}, + ) + + with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer exploded")): + yielded: list = [] + async for chunk in iterator: + yielded.append(chunk) + + assert yielded + assert iterator.completed_response.response.usage is None From 5a9fe56aff9ac8e3136f7faeb40891718262e3a3 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:24:05 +0000 Subject: [PATCH 04/12] fix(responses): count custom-tool and MCP argument deltas in the streamed usage fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 11 +++++++- .../responses/test_streaming_iterator.py | 28 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6cf36e3e660..df3c26c1a4f 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -356,7 +356,7 @@ class BaseResponsesAPIStreamingIterator: _delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_delta, str): self._generated_content += _delta - elif _event_type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA: + elif _event_type in _TOOL_ARGUMENTS_DELTA_EVENTS: _args_delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_args_delta, str): self._generated_tool_arguments += _args_delta @@ -1358,6 +1358,15 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None: return None +_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset( + { + ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA, + ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, + } +) + + def _estimate_usage_from_text( model: str, request_input: object, diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 234e63ca3be..7bd03863717 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -816,3 +816,31 @@ async def test_completed_event_survives_a_failing_usage_estimate(): assert yielded assert iterator.completed_response.response.usage is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_delta_event_type", + ["response.custom_tool_call_input.delta", "response.mcp_call_arguments.delta"], +) +async def test_completed_event_without_usage_counts_tool_input_deltas(tool_delta_event_type): + """Custom-tool and MCP argument deltas feed the streamed usage fallback the + same way function_call_arguments deltas do.""" + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": tool_delta_event_type, "delta": '{"query": "weather in sf"}'}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "what is the weather in san francisco"}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage is not None + assert usage.output_tokens > 0 + assert usage.total_tokens == usage.input_tokens + usage.output_tokens From a6d2332f7a19822ada77c4fd2cc13b853d1f9cad Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:24:35 +0000 Subject: [PATCH 05/12] test(responses): drive the usage-estimate failure path without patching litellm Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../responses/test_streaming_iterator.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 7bd03863717..d3849e6ffce 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -17,6 +17,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.responses.streaming_iterator import ( ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator, + _estimate_usage_from_text, ) from litellm.types.llms.openai import ( ResponseAPIUsage, @@ -796,8 +797,13 @@ async def test_completed_event_without_usage_counts_multimodal_input_as_messages @pytest.mark.asyncio async def test_completed_event_survives_a_failing_usage_estimate(): - """A raising token_counter must not break a stream that previously completed: - the estimate is best-effort and falls back to usage None.""" + """A malformed request input that makes the message transformer raise must not + break a stream that previously completed: the estimate is best-effort and + falls back to usage None.""" + malformed_input: Final = [{"type": "message", "role": "user", "content": 42}] + with pytest.raises(ValueError): + _estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world") + response = _responses_api_response_without_usage() iterator = _make_iterator( sse_events=[ @@ -806,13 +812,12 @@ async def test_completed_event_survives_a_failing_usage_estimate(): ], logging_obj=_logging_obj_stub(), config=_mock_config_with_completed_response(response), - request_data={"input": "count these input tokens please"}, + request_data={"input": malformed_input}, ) - with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer exploded")): - yielded: list = [] - async for chunk in iterator: - yielded.append(chunk) + yielded: list = [] + async for chunk in iterator: + yielded.append(chunk) assert yielded assert iterator.completed_response.response.usage is None From 7ed20406d7c2f7df3ebfce179769dd6ce26a5383 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:47:45 +0000 Subject: [PATCH 06/12] test(responses): narrow the ValueError assertion to satisfy PT011 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/responses/test_streaming_iterator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index d3849e6ffce..8a76b0bedff 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -801,7 +801,7 @@ async def test_completed_event_survives_a_failing_usage_estimate(): break a stream that previously completed: the estimate is best-effort and falls back to usage None.""" malformed_input: Final = [{"type": "message", "role": "user", "content": 42}] - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid content type"): _estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world") response = _responses_api_response_without_usage() From 7680d3de863648b500fc605d0eaa0c894d0a6bf4 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 02:05:44 +0000 Subject: [PATCH 07/12] fix(responses): tolerate dict terminal responses when estimating usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 35 ++++++++-------- .../responses/test_streaming_iterator.py | 42 +++++++++++++++++++ 2 files changed, 60 insertions(+), 17 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index df3c26c1a4f..10ddf07ac7e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -428,24 +428,25 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): self.completed_response = openai_responses_api_chunk - _response_obj: Final[ResponsesAPIResponse | None] = getattr( - openai_responses_api_chunk, "response", None - ) - if ( - _chunk_type - in ( - openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, - ) - and _response_obj is not None - and _response_obj.usage is None + _response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None) + if _chunk_type in ( + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, ): - _response_obj.usage = _estimate_usage_safely( - self.model or "", - self.request_data.get("input"), - self.request_data, - self._generated_content + self._generated_tool_arguments, - ) + if isinstance(_response_obj, ResponsesAPIResponse) and _response_obj.usage is None: + _response_obj.usage = _estimate_usage_safely( + self.model or "", + self.request_data.get("input"), + self.request_data, + self._generated_content + self._generated_tool_arguments, + ) + elif isinstance(_response_obj, dict) and _response_obj.get("usage") is None: # pyright: ignore[reportUnknownMemberType] # the model_constructed terminal event leaves response as an untyped dict + _response_obj["usage"] = _estimate_usage_safely( + self.model or "", + self.request_data.get("input"), + self.request_data, + self._generated_content + self._generated_tool_arguments, + ) _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 8a76b0bedff..8cf9556761e 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -849,3 +849,45 @@ async def test_completed_event_without_usage_counts_tool_input_deltas(tool_delta assert usage is not None assert usage.output_tokens > 0 assert usage.total_tokens == usage.input_tokens + usage.output_tokens + + +@pytest.mark.asyncio +async def test_completed_event_with_a_dict_response_still_gets_the_usage_estimate(): + """transform_streaming_response can model_construct a terminal event whose + response stays a plain dict; the estimate must fill it without raising.""" + dict_response: Final = { + "id": "resp_dict", + "model": "gpt-4o-mini", + "object": "response", + "output": [], + "usage": None, + } + + def _transform(model, parsed_chunk, logging_obj): + if parsed_chunk.get("type") == "response.completed": + return ResponseCompletedEvent.model_construct(type="response.completed", response=dict_response) + stub: Final = Mock() + stub.type = parsed_chunk.get("type") + if "delta" in parsed_chunk: + stub.delta = parsed_chunk.get("delta") + return stub + + config: Final = Mock(spec=BaseResponsesAPIConfig) + config.transform_streaming_response.side_effect = _transform + iterator: Final = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "hello world"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=config, + request_data={"input": "count these input tokens please"}, + ) + + async for _ in iterator: + pass + + usage: Final = iterator.completed_response.response["usage"] + assert usage is not None + assert usage.input_tokens > 0 + assert usage.output_tokens > 0 From 7121e64db4903dc7b832b25f79f1f5f048b73a36 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 03:13:53 +0000 Subject: [PATCH 08/12] fix(responses): type the dict terminal response so the estimated usage is billed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 42 +++++++++++-------- .../responses/test_streaming_iterator.py | 15 +++++-- 2 files changed, 35 insertions(+), 22 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 10ddf07ac7e..3004337f9d1 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -429,25 +429,31 @@ class BaseResponsesAPIStreamingIterator: ): self.completed_response = openai_responses_api_chunk _response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None) - if _chunk_type in ( - openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + _typed_response: Final[ResponsesAPIResponse | None] = ( + ResponsesAPIResponse.model_construct(**_response_obj) # pyright: ignore[reportUnknownArgumentType] # the model_constructed terminal event leaves response as an untyped dict + if isinstance(_response_obj, dict) + else _response_obj + if isinstance(_response_obj, ResponsesAPIResponse) + else None + ) + if ( + _typed_response is not None + and _chunk_type + in ( + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + ) + and _typed_response.usage is None ): - if isinstance(_response_obj, ResponsesAPIResponse) and _response_obj.usage is None: - _response_obj.usage = _estimate_usage_safely( - self.model or "", - self.request_data.get("input"), - self.request_data, - self._generated_content + self._generated_tool_arguments, - ) - elif isinstance(_response_obj, dict) and _response_obj.get("usage") is None: # pyright: ignore[reportUnknownMemberType] # the model_constructed terminal event leaves response as an untyped dict - _response_obj["usage"] = _estimate_usage_safely( - self.model or "", - self.request_data.get("input"), - self.request_data, - self._generated_content + self._generated_tool_arguments, - ) - _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj) + _typed_response.usage = _estimate_usage_safely( + self.model or "", + self.request_data.get("input"), + self.request_data, + self._generated_content + self._generated_tool_arguments, + ) + if _typed_response is not None and _typed_response is not _response_obj: + openai_responses_api_chunk.response = _typed_response # pyright: ignore[reportAttributeAccessIssue] # reached only on the dict path, which only response-carrying terminal events produce + _stamp_responses_usage_cost(_typed_response, self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: self._handle_logging_failed_response() diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 8cf9556761e..30a7c4faaed 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -852,9 +852,10 @@ async def test_completed_event_without_usage_counts_tool_input_deltas(tool_delta @pytest.mark.asyncio -async def test_completed_event_with_a_dict_response_still_gets_the_usage_estimate(): +async def test_completed_event_with_a_dict_response_is_typed_and_billed(): """transform_streaming_response can model_construct a terminal event whose - response stays a plain dict; the estimate must fill it without raising.""" + response stays a plain dict; the iterator must type it so the estimated + usage reaches the cost stamping path.""" dict_response: Final = { "id": "resp_dict", "model": "gpt-4o-mini", @@ -874,12 +875,14 @@ async def test_completed_event_with_a_dict_response_still_gets_the_usage_estimat config: Final = Mock(spec=BaseResponsesAPIConfig) config.transform_streaming_response.side_effect = _transform + logging_obj: Final = _logging_obj_stub() + logging_obj._response_cost_calculator.return_value = 0.000704 iterator: Final = _make_iterator( sse_events=[ _sse_event({"type": "response.output_text.delta", "delta": "hello world"}), _sse_event({"type": "response.completed", "response": {}}), ], - logging_obj=_logging_obj_stub(), + logging_obj=logging_obj, config=config, request_data={"input": "count these input tokens please"}, ) @@ -887,7 +890,11 @@ async def test_completed_event_with_a_dict_response_still_gets_the_usage_estimat async for _ in iterator: pass - usage: Final = iterator.completed_response.response["usage"] + completed_response: Final = iterator.completed_response.response + assert isinstance(completed_response, ResponsesAPIResponse) + usage: Final = completed_response.usage assert usage is not None assert usage.input_tokens > 0 assert usage.output_tokens > 0 + assert usage.cost == pytest.approx(0.000704) + logging_obj._response_cost_calculator.assert_any_call(result=completed_response) From a6fb21c3f86cfb645b0e94c46c0ea6c8a1955d91 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 15 Sep 2026 20:26:20 -0700 Subject: [PATCH 09/12] fix(e2e): record cookie-setting provider responses and keep prompt-caching tests live The first cache-enabled litellm-e2e build (211) showed three gaps in the shared provider cache: Every OpenAI response carries Cloudflare bot-management Set-Cookie headers, and the capture rejected any response with Set-Cookie, so no OpenAI response was ever recorded (179 of 372 misses rejected). The edge already withholds Set-Cookie from the proxy, so drop it before validating and storing instead of rejecting. The provider prompt-caching tests need fresh provider state: a replayed priming response reports cache creation rather than a cache read, and the TPM test then trips the key limit. Mark both modules provider_live. TestApiBaseSeam::test_live_mode_returns_none ran inside the cache-enabled runner and saw the shared edge; isolate it from E2E_PROVIDER_CACHE. --- tests/code_coverage_tests/test_provider_cache.py | 11 +++++++++++ tests/e2e/PROVIDER_CACHE.md | 4 ++-- tests/e2e/llm_translation/test_cache_control.py | 2 +- tests/e2e/provider_cache.py | 10 ++++++---- .../ratelimit/test_tpm_excludes_cached_tokens_e2e.py | 2 +- tests/e2e/test_provider_edge.py | 3 ++- 6 files changed, 23 insertions(+), 9 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 9d89a5fe622..828227ed239 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -38,6 +38,7 @@ class Provider(ThreadingHTTPServer): delay: float = 0 stream: bool = False truncated: bool = False + cookie: str = "" class Handler(BaseHTTPRequestHandler): @@ -62,6 +63,8 @@ class Handler(BaseHTTPRequestHandler): return self.send_header("content-type", "application/json") self.send_header("content-length", str(len(server.response))) + if server.cookie: + self.send_header("set-cookie", server.cookie) self.end_headers() self.wfile.write(server.response) @@ -172,6 +175,14 @@ def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore, assert len(provider.hits) == 2 +def test_cookie_setting_success_is_reused_without_the_cookie(store: RedisResponseStore, provider: Provider) -> None: + provider.cookie = "__cf_bm=synthetic-bot-management; Path=/; HttpOnly; Secure" + with edge(CacheEdge(store, SECRET), provider) as url: + replies: Final = tuple(call(url) for _ in range(2)) + assert len(provider.hits) == 1 + assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies) + + def test_expiry_does_not_slide(store: RedisResponseStore, provider: Provider) -> None: short: Final = replace(store, lifetime_ms=250) with edge(CacheEdge(short, SECRET), provider) as url: diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index efd9eb66daa..8635c9ed9ae 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -2,7 +2,7 @@ `E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away +The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure @@ -20,7 +20,7 @@ The trusted runner receives: Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits -Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay +Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay ## Recorded response semantics diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index a18e03c982b..102b3f00698 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -44,7 +44,7 @@ from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage from passthrough_client import PassthroughClient import os -pytestmark = pytest.mark.e2e +pytestmark = [pytest.mark.e2e, pytest.mark.provider_live] BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" VERTEX_MODEL = "vertex_ai/gemini-2.5-flash" diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 50b336f5263..0c6eac75a43 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -28,6 +28,7 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationEr LIFETIME_SECONDS: Final = 86_400 MAX_REQUEST_BYTES: Final = 256 * 1024 MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024 +UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"}) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -104,8 +105,6 @@ def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool: def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES: return False - if any(name.lower() == "set-cookie" for name in headers): - return False streaming: Final = "text/event-stream" in headers.get("content-type", "").lower() if streaming: try: @@ -283,11 +282,14 @@ class CacheEdge: yield step capture.observe(step) chunks: Final = capture.chunks() if capture.eligible else () - if not capture.eligible or not successful_response(url, head.status_code, head.headers, b"".join(chunks)): + headers: Final = { + name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS + } + if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)): self.counters.increment("rejected") return response: Final = CachedResponse( - request_key=key, status_code=head.status_code, headers=head.headers, + request_key=key, status_code=head.status_code, headers=headers, chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks), ) published: Final = self.store.publish(key, lease, encode_response(self.secret, response)) diff --git a/tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py b/tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py index b0bc6b3508c..33d869ee80e 100644 --- a/tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py +++ b/tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py @@ -26,7 +26,7 @@ from models import ( ) from quota_client import QuotaClient -pytestmark = pytest.mark.e2e +pytestmark = [pytest.mark.e2e, pytest.mark.provider_live] # Anthropic prompt caching (host has ANTHROPIC_API_KEY; Bedrock was "Operation not allowed"). ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001" diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index 81be81e7b59..5d0c79f26f6 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -1254,7 +1254,8 @@ class TestHandleEdgeRequestPure: class TestApiBaseSeam: - def test_live_mode_returns_none(self, tmp_path: Path) -> None: + def test_live_mode_returns_none(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("E2E_PROVIDER_CACHE", raising=False) for mode_raw in ("live", ""): assert ( provider_edge_api_base( From 7cc07d437a3d76b5f42204adf6f4561150957a3c Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 03:47:03 +0000 Subject: [PATCH 10/12] fix(responses): build the billed terminal response immutably and guard the cache dump Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 90 +++++++++++++------ .../responses/test_streaming_iterator.py | 50 ++++++++++- 2 files changed, 112 insertions(+), 28 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 3004337f9d1..d6d3e2576f6 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -427,39 +427,45 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): - self.completed_response = openai_responses_api_chunk _response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None) - _typed_response: Final[ResponsesAPIResponse | None] = ( - ResponsesAPIResponse.model_construct(**_response_obj) # pyright: ignore[reportUnknownArgumentType] # the model_constructed terminal event leaves response as an untyped dict - if isinstance(_response_obj, dict) - else _response_obj - if isinstance(_response_obj, ResponsesAPIResponse) - else None + _estimate_wanted: Final[bool] = _chunk_type in ( + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, ) - if ( - _typed_response is not None - and _chunk_type - in ( - openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + _billed_response: Final[ResponsesAPIResponse | None] = _billed_terminal_response( + _response_obj, + ( + lambda: ( + _estimate_usage_safely( + self.model or "", + self.request_data.get("input"), + self.request_data, + self._generated_content + self._generated_tool_arguments, + ) + if _estimate_wanted + else None + ) + ), + ) + _terminal_chunk: Final = ( + openai_responses_api_chunk + if _billed_response is None or _billed_response is _response_obj + else ( + openai_responses_api_chunk.model_copy(update={"response": _billed_response}) + if issubclass(type(openai_responses_api_chunk), BaseModel) # pyright: ignore[reportUnnecessaryIsInstance] # test stubs use spec'd Mocks whose __class__ reports BaseModel but whose model_copy returns a Mock + else _replace_response(openai_responses_api_chunk, _billed_response) ) - and _typed_response.usage is None - ): - _typed_response.usage = _estimate_usage_safely( - self.model or "", - self.request_data.get("input"), - self.request_data, - self._generated_content + self._generated_tool_arguments, - ) - if _typed_response is not None and _typed_response is not _response_obj: - openai_responses_api_chunk.response = _typed_response # pyright: ignore[reportAttributeAccessIssue] # reached only on the dict path, which only response-carrying terminal events produce - _stamp_responses_usage_cost(_typed_response, self.logging_obj) + ) + self.completed_response = _terminal_chunk + _stamp_responses_usage_cost(_billed_response, self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: self._handle_logging_failed_response() else: self._handle_logging_completed_response() + return _terminal_chunk + return openai_responses_api_chunk return None @@ -688,7 +694,9 @@ class BaseResponsesAPIStreamingIterator: if cache is None: return - cached_response: Final = response_obj.model_dump_json() + cached_response: Final = _dump_json_safely(response_obj) + if cached_response is None: + return if is_async: from litellm.caching.caching_handler import create_cache_write_task @@ -1334,6 +1342,38 @@ def _add_text_like_part_events( ) +def _billed_terminal_response( + response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None +) -> ResponsesAPIResponse | None: + if isinstance(response_obj, ResponsesAPIResponse): + return ( + response_obj + if response_obj.usage is not None or estimate is None + else response_obj.model_copy(update={"usage": estimate()}) + ) + if not isinstance(response_obj, dict): + return None + usage: Final[object] = response_obj.get("usage") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a model_constructed terminal event leaves response as an untyped dict + return ResponsesAPIResponse.model_construct( + **{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # same untyped dict spread + ) + + +def _replace_response( + event: ResponsesAPIStreamingResponse, response: ResponsesAPIResponse +) -> ResponsesAPIStreamingResponse: + setattr(event, "response", response) + return event + + +def _dump_json_safely(response: BaseModel) -> str | None: + try: + return response.model_dump_json() + except Exception as exc: + verbose_logger.debug("could not serialize completed response for cache: %s", exc) + return None + + def _logging_copy(event: object) -> object: """Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 30a7c4faaed..b33cc4e93c2 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -10,6 +10,7 @@ from unittest.mock import Mock, patch import httpx import pytest +from pydantic_core import PydanticSerializationError import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -887,10 +888,11 @@ async def test_completed_event_with_a_dict_response_is_typed_and_billed(): request_data={"input": "count these input tokens please"}, ) - async for _ in iterator: - pass + yielded: Final = [chunk async for chunk in iterator] - completed_response: Final = iterator.completed_response.response + terminal_event: Final = iterator.completed_response + assert yielded[-1] is terminal_event + completed_response: Final = terminal_event.response assert isinstance(completed_response, ResponsesAPIResponse) usage: Final = completed_response.usage assert usage is not None @@ -898,3 +900,45 @@ async def test_completed_event_with_a_dict_response_is_typed_and_billed(): assert usage.output_tokens > 0 assert usage.cost == pytest.approx(0.000704) logging_obj._response_cost_calculator.assert_any_call(result=completed_response) + + +def test_billed_terminal_response_keeps_a_response_that_already_has_usage(): + from litellm.responses.streaming_iterator import _billed_terminal_response + + response: Final = _responses_api_response_with_usage() + + assert _billed_terminal_response(response, None) is response + + +def test_billed_terminal_response_copies_when_estimating_and_leaves_the_original_untouched(): + from litellm.responses.streaming_iterator import _billed_terminal_response + + response: Final = _responses_api_response_without_usage() + estimated: Final = ResponseAPIUsage(input_tokens=3, output_tokens=4, total_tokens=7) + + billed: Final = _billed_terminal_response(response, lambda: estimated) + + assert billed is not response + assert billed.usage is estimated + assert response.usage is None + + +def test_persist_completed_response_to_cache_survives_an_unserializable_response(monkeypatch): + bad_response: Final = ResponsesAPIResponse.model_construct(id="r", output=[object()], usage=None) + with pytest.raises(PydanticSerializationError): + bad_response.model_dump_json() + + logging_obj: Final = _logging_obj_stub() + caching_handler: Final = Mock() + caching_handler.request_kwargs = {"stream": True} + logging_obj._llm_caching_handler = caching_handler + iterator: Final = _make_iterator(sse_events=[], logging_obj=logging_obj) + iterator.completed_response = ResponseCompletedEvent.model_construct( + type="response.completed", response=bad_response + ) + cache: Final = Mock() + monkeypatch.setattr(litellm, "cache", cache) + + iterator._persist_completed_response_to_cache(is_async=False) + + cache.add_cache.assert_not_called() From ff878e7df06e85a0c6ee2a75f1541894929e0af2 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 03:54:04 +0000 Subject: [PATCH 11/12] refactor(responses): copy the terminal event instead of mutating stubbed chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 13 +------------ .../responses/test_streaming_iterator.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 20 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d6d3e2576f6..8d766cf1cd0 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -450,11 +450,7 @@ class BaseResponsesAPIStreamingIterator: _terminal_chunk: Final = ( openai_responses_api_chunk if _billed_response is None or _billed_response is _response_obj - else ( - openai_responses_api_chunk.model_copy(update={"response": _billed_response}) - if issubclass(type(openai_responses_api_chunk), BaseModel) # pyright: ignore[reportUnnecessaryIsInstance] # test stubs use spec'd Mocks whose __class__ reports BaseModel but whose model_copy returns a Mock - else _replace_response(openai_responses_api_chunk, _billed_response) - ) + else openai_responses_api_chunk.model_copy(update={"response": _billed_response}) ) self.completed_response = _terminal_chunk _stamp_responses_usage_cost(_billed_response, self.logging_obj) @@ -1359,13 +1355,6 @@ def _billed_terminal_response( ) -def _replace_response( - event: ResponsesAPIStreamingResponse, response: ResponsesAPIResponse -) -> ResponsesAPIStreamingResponse: - setattr(event, "response", response) - return event - - def _dump_json_safely(response: BaseModel) -> str | None: try: return response.model_dump_json() diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index b33cc4e93c2..dbf54ec3b9b 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -47,10 +47,10 @@ def _mock_config() -> Mock: def _transform(model, parsed_chunk, logging_obj): evt_type = parsed_chunk.get("type") if evt_type == "response.completed": - completed = Mock(spec=ResponseCompletedEvent) - completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED - completed.response = mock_responses_api_response - return completed + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=mock_responses_api_response, + ) stub = Mock() stub.type = evt_type return stub @@ -348,10 +348,10 @@ def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock def _transform(model, parsed_chunk, logging_obj): evt_type = parsed_chunk.get("type") if evt_type == "response.completed": - completed = Mock(spec=ResponseCompletedEvent) - completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED - completed.response = response - return completed + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=response, + ) stub = Mock() stub.type = evt_type if "delta" in parsed_chunk: From 9d640b86ed1874562e6b3bbd68242e2923f38275 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 15 Sep 2026 20:58:38 -0700 Subject: [PATCH 12/12] fix(ui): simplify Capability and Fuse routing options --- .../add_model/ComplexityRouterConfig.tsx | 42 +++++++++---------- ...ecastClassifierConfig.integration.test.tsx | 2 +- .../add_model/add_auto_router_tab.test.tsx | 18 ++++++-- .../build_complexity_router_config.test.ts | 31 ++++++++++++++ .../build_complexity_router_config.ts | 27 ++++++------ .../forecast_classifier_config.test.ts | 2 +- ...d_updated_complexity_router_config.test.ts | 37 ++++++++++++++++ .../edit_auto_router_modal.tsx | 5 ++- 8 files changed, 122 insertions(+), 42 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 4ccabff18b9..f6b50ce20bc 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -794,19 +794,15 @@ const ComplexityRouterConfig: React.FC = ({ }, ] : []), - ...(value.classifier_type !== "llm_v2" - ? [ - { - key: "adaptive", - label: Advanced: Adaptive Routing, - children: ( - - - - ), - }, - ] - : []), + { + key: "adaptive", + label: Advanced: Adaptive Routing, + children: ( + + + + ), + }, { key: "affinity", label: Advanced: Affinity, @@ -906,15 +902,17 @@ const ComplexityRouterConfig: React.FC = ({ }, ] : []), - ].map(({ key, label, children }) => ( - - - - {label} - - {children} - - ))} + ] + .filter(({ key }) => !forecast || !["adaptive", "context-window", "escalation"].includes(key)) + .map(({ key, label, children }) => ( + + + + {label} + + {children} + + ))} diff --git a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx index c249a63899a..4a574ac736d 100644 --- a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx @@ -105,7 +105,7 @@ describe("forecast classifier form", () => { expect(output).toHaveTextContent('"REASONING":["capable"]'); expect(output).toHaveTextContent('"reasoning_effort":"low","speed":"fast","max_tokens":1024'); expect(output).toHaveTextContent('"reasoning_effort":"high"'); - expect(output).toHaveTextContent('"adaptive":true'); + expect(output).toHaveTextContent('"adaptive":false'); expect(output).not.toHaveTextContent("leftover-medium"); expect(output).not.toHaveTextContent("leftover-complex"); expect(output).not.toHaveTextContent('"plan_mode_min_tier"'); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index b98636c9c5c..48903d585ff 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -213,17 +213,24 @@ describe("AddAutoRouterTab", () => { expect(screen.queryByRole("alert")).not.toBeInTheDocument(); expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled(); await user.click(screen.getByRole("button", { name: "Advanced routing options" })); - if (capability) expect(screen.getByText("Advanced: Adaptive Routing")).toBeInTheDocument(); - else expect(screen.queryByText("Advanced: Adaptive Routing")).not.toBeInTheDocument(); + for (const label of ["Adaptive Routing", "Context Window Escalation", "Escalation Keywords"]) { + expect(screen.queryByText(`Advanced: ${label}`)).not.toBeInTheDocument(); + } + expect(screen.getByText("Advanced: Stalled Task Escalation")).toBeInTheDocument(); + expect(screen.getByText("Advanced: Response Format")).toBeInTheDocument(); expect(screen.queryByText("Advanced: Classification Method")).not.toBeInTheDocument(); expect(screen.getByText("Advanced: Affinity")).toBeInTheDocument(); await user.click(screen.getByRole("button", { name: "Add Auto Router" })); await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledTimes(1)); - expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls[0][0].complexity_router_config).toMatchObject({ + const expected = { classifier_type: capability ? "capability" : "llm_v2", + adaptive: false, + enable_context_window_escalation: false, + escalation_keywords: [], tiers: { SIMPLE: ["efficient"], REASONING: ["capable"] }, classifier_llm_config: { model: "judge" }, - }); + }; + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls[0][0].complexity_router_config).toMatchObject(expected); }, ); @@ -239,6 +246,9 @@ describe("AddAutoRouterTab", () => { expect(screen.getByTestId("template-selector")).toBeInTheDocument(); expandDetailedConfiguration(); expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument(); + for (const label of ["Adaptive Routing", "Context Window Escalation", "Escalation Keywords"]) { + expect(screen.getByText(`Advanced: ${label}`)).toBeInTheDocument(); + } await user.click(screen.getByText("Advanced: Classification Method")); expect(screen.queryByRole("radio", { name: /^Capability/ })).not.toBeInTheDocument(); expect(screen.queryByRole("radio", { name: /^Fuse v2/ })).not.toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index bc30b591ea9..6e6e7a3c6cd 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -48,6 +48,37 @@ const baseParams: BuildComplexityRouterConfigParams = { }; describe("buildComplexityRouterConfig", () => { + it.each(["capability", "llm_v2", "heuristic"] as const)( + "disables the removed overrides only for forecast creates: %s", + (classifierType) => { + const forecast = classifierType !== "heuristic"; + const params = { + ...baseParams, + classifierType, + adaptive: true, + enableContextWindowEscalation: true, + contextWindowEscalationBuffer: 0.9, + }; + const config = buildComplexityRouterConfig(params); + expect(config.adaptive).toBe(!forecast); + expect(config.enable_context_window_escalation).toBe(!forecast); + expect(config.escalation_keywords).toEqual(forecast ? [] : ["LITELLM ESCALATE"]); + for (const key of [ + "adaptive_weights", + "adaptive_eligible", + "tier_distance_penalty", + "context_window_escalation_buffer", + ]) { + expect(Object.hasOwn(config, key)).toBe(!forecast); + } + if (forecast) { + const untouched = buildComplexityRouterConfig({ ...baseParams, classifierType }); + expect(untouched.enable_context_window_escalation).toBe(false); + expect(untouched.escalation_keywords).toEqual([]); + } + }, + ); + it("carries Fast and reasoning overrides independently into a new router payload", () => { const params = { speed: "fast", reasoning_effort: "high", max_tokens: 1024 }; const config = buildComplexityRouterConfig({ diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index d16be537ce6..8a377c17ad7 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -628,13 +628,11 @@ export const buildComplexityRouterConfig = ({ // An edited tier set forces the LLM classifier, so llm-only inputs must survive a classifier_type // the form never rewrote. The UI gates the same controls on this, not on the raw value. const effectiveType: ClassifierType = customTierSet ? "llm" : classifierType; + const forecast = isForecastClassifier(effectiveType); - const supportsOpeningPrompt = - !customTierSet && !isForecastClassifier(effectiveType) && usesLlmClassifier(effectiveType); + const supportsOpeningPrompt = !customTierSet && !forecast && usesLlmClassifier(effectiveType); const payload: ComplexityRouterConfigPayload = { - tiers: isForecastClassifier(effectiveType) - ? Object.fromEntries(Object.entries(tiers).filter(([, models]) => models.length > 0)) - : tiers, + tiers: forecast ? Object.fromEntries(Object.entries(tiers).filter(([, models]) => models.length > 0)) : tiers, // The backend rejects the flag beside a custom tier set. ...(!customTierSet && enableNonReasoningTier && { enable_non_reasoning_tier: true }), ...(serializedTierModelConfigs && { tier_model_configs: serializedTierModelConfigs }), @@ -645,7 +643,8 @@ export const buildComplexityRouterConfig = ({ ...classifierWireFields(effectiveType, classifierInputs), ...(effectiveType === "capability" && capabilityClassifierConfig && { capability_classifier_config: capabilityClassifierConfig }), - ...(effectiveType === "llm_v2" && { llm_v2_config: llmV2Config, adaptive: false }), + ...(effectiveType === "llm_v2" && { llm_v2_config: llmV2Config }), + ...(forecast && { adaptive: false }), // A built-in router's opening instructions. Suppressed beside a legacy whole-prompt override, // which the backend rejects as a second override of the same prompt. ...(supportsOpeningPrompt && @@ -660,7 +659,7 @@ export const buildComplexityRouterConfig = ({ modality_pin_override: modalityPinOverride ?? false, ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }), - escalation_keywords: cleanedEscalationKeywords, + escalation_keywords: forecast ? [] : cleanedEscalationKeywords, // Only written when on: the backend rejects it alongside session_affinity, user_turn mode and // a custom tier set, so an off router must not carry the key into any of those saves. ...(stallEscalationEnabled && { @@ -676,19 +675,21 @@ export const buildComplexityRouterConfig = ({ match_threshold: matchThreshold, }), ...(adaptive && - effectiveType !== "llm_v2" && { + !forecast && { adaptive: true, adaptive_weights: adaptiveWeights, ...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }), adaptive_eligible: adaptiveEligible, }), ...(returnRawModelName && { return_raw_model_name: true }), - ...(enableContextWindowEscalation !== undefined && { - enable_context_window_escalation: enableContextWindowEscalation, - }), - ...(contextWindowEscalationBuffer !== undefined && { - context_window_escalation_buffer: contextWindowEscalationBuffer, + // Omission enables the backend default, so hidden forecast controls need an explicit opt-out. + ...((forecast || enableContextWindowEscalation !== undefined) && { + enable_context_window_escalation: forecast ? false : enableContextWindowEscalation, }), + ...(!forecast && + contextWindowEscalationBuffer !== undefined && { + context_window_escalation_buffer: contextWindowEscalationBuffer, + }), ...(sessionAffinityTtlSeconds !== undefined && { session_affinity_ttl_seconds: sessionAffinityTtlSeconds, }), diff --git a/ui/litellm-dashboard/src/components/add_model/forecast_classifier_config.test.ts b/ui/litellm-dashboard/src/components/add_model/forecast_classifier_config.test.ts index d07d4a24cf0..218eebe8430 100644 --- a/ui/litellm-dashboard/src/components/add_model/forecast_classifier_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/forecast_classifier_config.test.ts @@ -264,7 +264,7 @@ describe("forecast classifier configuration", () => { }); expect(saved.tiers).toEqual({ SIMPLE: ["efficient"], MEDIUM: ["middle"], REASONING: ["capable"] }); expect(saved.tier_model_configs).toEqual(stored.tier_model_configs); - expect(saved.adaptive).toBe(true); + expect(saved.adaptive).toBe(false); expect(saved.plan_mode_min_tier).toBe("MEDIUM"); }); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 1aa6f0e3d35..4ae6efbb12d 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -46,6 +46,43 @@ const hydratedState: KeywordMatchingState = { }; describe("buildUpdatedComplexityRouterConfig keyword matching", () => { + it.each(["capability", "llm_v2", "heuristic"] as const)( + "handles enabled stored overrides when editing %s with or without keyword form state", + (classifier_type) => { + const stored = { + ...STORED, + classifier_type, + adaptive: classifier_type !== "llm_v2", + adaptive_weights: { quality: 0.6, cost: 0.4 }, + adaptive_eligible: "all", + tier_distance_penalty: 0.8, + enable_context_window_escalation: true, + context_window_escalation_buffer: 0.9, + }; + const value = hydrateComplexityRouterConfig(stored, undefined); + for (const keywordState of [undefined, hydratedState]) { + const saved = buildUpdatedComplexityRouterConfig(stored, value, undefined, keywordState); + const forecast = classifier_type !== "heuristic"; + expect(saved.adaptive).toBe(!forecast); + expect(saved.enable_context_window_escalation).toBe(!forecast); + expect(saved.escalation_keywords).toEqual(forecast ? [] : stored.escalation_keywords); + for (const key of [ + "adaptive_weights", + "adaptive_eligible", + "tier_distance_penalty", + "context_window_escalation_buffer", + ]) { + expect(Object.hasOwn(saved, key)).toBe(!forecast); + } + expect(saved.keyword_tier_rules).toEqual(STORED.keyword_tier_rules); + expect(saved.semantic_keyword_matching).toBe(true); + expect(saved.some_future_backend_key).toEqual(STORED.some_future_backend_key); + } + expect(value.enable_context_window_escalation).toBe(true); + expect(stored.escalation_keywords).toEqual(["urgent", "outage"]); + }, + ); + it("round-trips an untouched edit without changing any keyword-matching value", () => { // Opening the modal hydrates state from STORED; saving with nothing changed must be a // no-op. These keys are now MANAGED, so a hydration bug silently wipes them. diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 4a6c22689cc..e25c7f07dd7 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -3,6 +3,7 @@ import type { StoredComplexityRouterConfig } from "../add_model/build_complexity export type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; import { getForecastConfigError, + isForecastClassifier, capabilitySettingsSchema, fuseSettingsSchema, } from "../add_model/forecast_classifier_config"; @@ -71,6 +72,7 @@ import { } from "../add_model/heuristic_scoring_knobs"; import ComplexityRouterConfig, { ComplexityRouterConfigValue, + effectiveClassifierType, heuristicScoringRole, DEFAULT_ADAPTIVE_WEIGHTS, DEFAULT_SESSION_AFFINITY, @@ -305,6 +307,7 @@ export const buildUpdatedComplexityRouterConfig = ( ): Record => { const isManaged = (key: string): boolean => { if (MANAGED_COMPLEXITY_ROUTER_KEYS.has(key)) return true; + if (key === "escalation_keywords" && isForecastClassifier(effectiveClassifierType(value))) return true; if (keywordMatching !== undefined && KEYWORD_MATCHING_KEYS.has(key)) return true; return customTechnicalKeywords !== undefined && key === "custom_technical_keywords"; }; @@ -365,7 +368,7 @@ export const buildUpdatedComplexityRouterConfig = ( // Keys this call does not own stay as the stored config left them. const unowned: readonly string[] = [ - ...(keywordMatching === undefined ? KEYWORD_MATCHING_KEYS : []), + ...(keywordMatching === undefined ? [...KEYWORD_MATCHING_KEYS].filter((key) => !isManaged(key)) : []), ...(customTechnicalKeywords === undefined ? ["custom_technical_keywords"] : []), ]; return {