From d2fbaff2c948e5beb0fe02623640e1fc9c21cf69 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:05:28 -0700 Subject: [PATCH 01/11] fix(proxy): record estimated input tokens in spend logs for dispatched failed requests Failure rows in the spend log only carried token counts when a broken stream stashed recovered partial usage; non-stream requests that reached the provider and then failed (timeouts, provider 4xx/5xx) logged 0/0/0 even though the provider billed the input tokens. Estimate the input side in post_call_failure_hook with the same tokenizer fallback interrupted streams use, gated to requests that were actually dispatched (first_api_call_start_time set and no litellm_no_upstream_llm_call marker), and pin response_cost to 0.0 so failed requests never bill spend. Recovered partial-stream usage still wins over the estimate. --- litellm/proxy/utils.py | 67 +++++++-- tests/test_litellm/proxy/test_proxy_utils.py | 139 +++++++++++++++++++ 2 files changed, 197 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2ad7180bd5f..3457ae0f352 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -40,7 +40,7 @@ from litellm.proxy._types import ( from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse -from litellm.types.utils import CallTypes, CallTypesLiteral, ModelInfo +from litellm.types.utils import CallTypes, CallTypesLiteral, ModelInfo, Usage try: from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( @@ -403,6 +403,52 @@ def _exception_changes_request_flow(exc: BaseException) -> bool: return isinstance(exc, (SensitiveDataRouteException, ModifyResponseException)) +def _count_request_input_tokens(model: str, request_input: object) -> int: + if isinstance(request_input, str): + return litellm.token_counter(model=model, text=request_input) + if not isinstance(request_input, list) or not request_input: + return 0 + text_entries: Final = tuple(entry for entry in request_input if isinstance(entry, str)) + if len(text_entries) == len(request_input): + return litellm.token_counter(model=model, text="".join(text_entries)) + return litellm.token_counter(model=model, messages=request_input) + + +def _estimate_dispatched_failure_usage(model: str, request_input: object) -> Usage | None: + """A request that failed after dispatch consumed provider-billed input + tokens, but no provider usage ever came back. Estimate the input side with + the same tokenizer fallback interrupted streams use, so the spend log's + failure row records what was sent instead of zero.""" + try: + input_tokens: Final = _count_request_input_tokens(model=model, request_input=request_input) + except Exception: + return None + if input_tokens <= 0: + return None + return Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens) + + +def _failure_usage_to_lift(model_call_details: Mapping[str, object], dispatched: bool) -> tuple[object, object] | None: + """A stream that broke mid-flight still billed the provider for the chunks + already delivered; the streaming handler stashes that recovered usage and + cost in model_call_details, so prefer it. Otherwise a request that was + dispatched to a provider and failed without upstream usage gets an + estimated input-side Usage with zero cost. Returns the + (combined_usage_object, response_cost) pair to lift, or None.""" + recovered_usage: Final = model_call_details.get("combined_usage_object") + if recovered_usage is not None: + return recovered_usage, model_call_details.get("response_cost") + if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL): + return None + estimated_usage: Final = _estimate_dispatched_failure_usage( + model=str(model_call_details.get("model") or ""), + request_input=model_call_details.get("messages"), + ) + if estimated_usage is None: + return None + return estimated_usage, 0.0 + + @dataclass(frozen=True) class _CallbackCapabilities: """Cached per-hook capability flags derived from ``litellm.callbacks``. @@ -2190,15 +2236,18 @@ class ProxyLogging: if _first_handoff is not None: request_data["first_api_call_start_time"] = _first_handoff - # A stream that broke mid-flight still billed the provider for the - # chunks already delivered; the streaming handler stashes that - # recovered usage and cost here. Lift them onto request_data so the + # Lift recovered partial-stream usage, or an estimated input-side + # usage for a dispatched failure, onto request_data so the # failure-path spend callbacks (which run after the logging object - # is popped) record the real partial spend instead of zero. - _recovered_usage: Final = _model_call_details.get("combined_usage_object") - if _recovered_usage is not None: - request_data["combined_usage_object"] = _recovered_usage - request_data["response_cost"] = _model_call_details.get("response_cost") + # is popped) record real token counts instead of zero. + _usage_to_lift: Final = _failure_usage_to_lift( + model_call_details=_model_call_details, + dispatched=_first_handoff is not None, + ) + if _usage_to_lift is not None: + _lifted_usage, _lifted_cost = _usage_to_lift + request_data["combined_usage_object"] = _lifted_usage + request_data["response_cost"] = _lifted_cost # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 1504c3c3103..70baf157ab9 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -478,6 +478,145 @@ class TestPostCallFailureHookLiftsRecoveredPartialSpend: assert "response_cost" not in request_data +class TestPostCallFailureHookEstimatesDispatchedInputTokens: + """A non-stream request that failed after dispatch (timeout, provider + error) consumed provider-billed input tokens but recovered no usage. + post_call_failure_hook must estimate the input side onto request_data so + the spend log's failure row records what was sent instead of zero, while + never charging spend for the failure (LIT-5690). + """ + + async def _run(self, request_data): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import UserAPIKeyAuth + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("boom"), + user_api_key_dict=UserAPIKeyAuth(), + ) + + def _logging_obj(self, model_call_details): + logging_obj = MagicMock() + logging_obj.model_call_details = model_call_details + return logging_obj + + @pytest.mark.asyncio + async def test_dispatched_failure_estimates_input_tokens_with_zero_cost(self): + from datetime import datetime + + from litellm.types.utils import Usage + + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "count these input tokens please"}], + } + ), + "metadata": {}, + "response_cost": 123.0, + } + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + assert estimated.prompt_tokens > 0 + assert estimated.completion_tokens == 0 + assert estimated.total_tokens == estimated.prompt_tokens + assert request_data["response_cost"] == 0.0 + + @pytest.mark.asyncio + async def test_failure_before_dispatch_stays_zero(self): + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "never dispatched"}], + } + ), + "metadata": {}, + } + await self._run(request_data) + + assert "combined_usage_object" not in request_data + assert "response_cost" not in request_data + + @pytest.mark.asyncio + async def test_proxy_only_error_never_dispatched_stays_zero(self): + from datetime import datetime + + from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL + + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": "no-such-model", + "messages": [{"role": "user", "content": "hi"}], + LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True, + } + ), + "metadata": {}, + } + await self._run(request_data) + + assert "combined_usage_object" not in request_data + assert "response_cost" not in request_data + + @pytest.mark.asyncio + async def test_recovered_partial_usage_wins_over_estimate(self): + from datetime import datetime + + from litellm.types.utils import Usage + + recovered_usage = Usage(prompt_tokens=30, completion_tokens=7, total_tokens=37) + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "mid-stream failure"}], + "combined_usage_object": recovered_usage, + "response_cost": 3.5e-05, + } + ), + "metadata": {}, + } + await self._run(request_data) + + assert request_data["combined_usage_object"] is recovered_usage + assert request_data["response_cost"] == 3.5e-05 + + @pytest.mark.asyncio + async def test_dispatched_failure_with_text_completion_prompt(self): + from datetime import datetime + + from litellm.types.utils import Usage + + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": "gpt-3.5-turbo", + "messages": "a plain text-completion prompt string", + } + ), + "metadata": {}, + } + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + assert estimated.prompt_tokens > 0 + assert estimated.completion_tokens == 0 + + from typing import cast import litellm From 803113c63af3c543473c42c91bb1846974f782ab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:21:47 -0700 Subject: [PATCH 02/11] fix(proxy): estimate failed-request input tokens on /v1/messages and count system prompts The Anthropic messages endpoint's exception handler passed the raw request body dict to the failure hook, but request setup had already replaced the processor's dict with one carrying the logging object, so failure rows for /v1/messages never lifted recovered or estimated usage. Pass the processor's dict instead. The input-side estimate only counted the messages list, missing the Anthropic top-level system prompt (string or text-block list) and the Responses API instructions field, which live in optional_params. Count them too. --- .../proxy/anthropic_endpoints/endpoints.py | 4 +- litellm/proxy/utils.py | 42 ++++++++++-- .../anthropic_endpoints/test_endpoints.py | 35 ++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 68 +++++++++++++++++++ 4 files changed, 140 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index a48ef0f08bb..f742965ade2 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -179,7 +179,7 @@ async def anthropic_response( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data=data, + request_data=base_llm_response_processor.data, ) body: Final = AnthropicExceptionMapping.transform_to_anthropic_error( status_code=e.status_code, @@ -189,7 +189,7 @@ async def anthropic_response( return JSONResponse(status_code=e.status_code, content=body) except Exception as e: await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data + user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data ) verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3457ae0f352..5bfc1c2d1e1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -403,24 +403,45 @@ def _exception_changes_request_flow(exc: BaseException) -> bool: return isinstance(exc, (SensitiveDataRouteException, ModifyResponseException)) -def _count_request_input_tokens(model: str, request_input: object) -> int: +def _prompt_block_text(block: object) -> str: + if isinstance(block, str): + return block + if not isinstance(block, dict): + return "" + block_text: Final = block.get("text") + return block_text if isinstance(block_text, str) else "" + + +def _system_prompt_text(system_input: object) -> str: + if isinstance(system_input, str): + return system_input + if not isinstance(system_input, list): + return "" + return "".join(_prompt_block_text(block) for block in system_input) + + +def _count_request_input_tokens(model: str, request_input: object, system_input: object) -> int: + system_text: Final = _system_prompt_text(system_input) + system_tokens: Final = litellm.token_counter(model=model, text=system_text) if system_text else 0 if isinstance(request_input, str): - return litellm.token_counter(model=model, text=request_input) + return system_tokens + litellm.token_counter(model=model, text=request_input) if not isinstance(request_input, list) or not request_input: - return 0 + return system_tokens text_entries: Final = tuple(entry for entry in request_input if isinstance(entry, str)) if len(text_entries) == len(request_input): - return litellm.token_counter(model=model, text="".join(text_entries)) - return litellm.token_counter(model=model, messages=request_input) + return system_tokens + litellm.token_counter(model=model, text="".join(text_entries)) + return system_tokens + litellm.token_counter(model=model, messages=request_input) -def _estimate_dispatched_failure_usage(model: str, request_input: object) -> Usage | None: +def _estimate_dispatched_failure_usage(model: str, request_input: object, system_input: object) -> Usage | None: """A request that failed after dispatch consumed provider-billed input tokens, but no provider usage ever came back. Estimate the input side with the same tokenizer fallback interrupted streams use, so the spend log's failure row records what was sent instead of zero.""" try: - input_tokens: Final = _count_request_input_tokens(model=model, request_input=request_input) + input_tokens: Final = _count_request_input_tokens( + model=model, request_input=request_input, system_input=system_input + ) except Exception: return None if input_tokens <= 0: @@ -440,9 +461,16 @@ def _failure_usage_to_lift(model_call_details: Mapping[str, object], dispatched: return recovered_usage, model_call_details.get("response_cost") if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL): return None + optional_params: Final = model_call_details.get("optional_params") + system_input: Final = ( + (optional_params.get("system") or optional_params.get("instructions")) + if isinstance(optional_params, dict) + else None + ) estimated_usage: Final = _estimate_dispatched_failure_usage( model=str(model_call_details.get("model") or ""), request_input=model_call_details.get("messages"), + system_input=system_input, ) if estimated_usage is None: return None diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index 0a427df0cb7..9a90daeccb7 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -164,6 +164,41 @@ class TestProxyExceptionPassthrough: mock_logging.post_call_failure_hook.assert_awaited_once() +class TestFailureHookRequestData: + @pytest.mark.asyncio + async def test_failure_hook_gets_post_setup_data_with_logging_obj(self): + """Request setup replaces the processor's data dict (adding the logging + object the failure hook needs to lift token usage from); the exception + handler must pass that replaced dict, not the raw request body dict.""" + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + + captured = {} + + async def fake_process(self, **kwargs): + self.data = {**self.data, "litellm_logging_obj": "logging-obj-sentinel"} + captured["processor_data"] = self.data + raise RuntimeError("provider timeout") + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), + patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), + patch.object(proxy_server, "proxy_logging_obj") as mock_logging, + ): + mock_logging.post_call_failure_hook = AsyncMock() + with pytest.raises(ProxyException): + await ep.anthropic_response( + fastapi_response=MagicMock(), + request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + + hook_request_data = mock_logging.post_call_failure_hook.await_args.kwargs["request_data"] + assert hook_request_data is captured["processor_data"] + assert hook_request_data["litellm_logging_obj"] == "logging-obj-sentinel" + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 70baf157ab9..0455a806c0b 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -616,6 +616,74 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: assert estimated.prompt_tokens > 0 assert estimated.completion_tokens == 0 + def _dispatched_request_data(self, messages, optional_params): + from datetime import datetime + + return { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": "gpt-3.5-turbo", + "messages": messages, + "optional_params": optional_params, + } + ), + "metadata": {}, + } + + @pytest.mark.asyncio + async def test_anthropic_system_prompt_counted_in_estimate(self): + import litellm as litellm_module + from litellm.types.utils import Usage + + system_prompt = "You are a verbose historian who narrates every fact in exhaustive detail." + messages = [{"role": "user", "content": "write a short essay"}] + request_data = self._dispatched_request_data(messages, {"system": system_prompt, "max_tokens": 100}) + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( + model="gpt-3.5-turbo", text=system_prompt + ) + assert estimated.prompt_tokens == expected + + @pytest.mark.asyncio + async def test_anthropic_system_text_blocks_counted_in_estimate(self): + import litellm as litellm_module + from litellm.types.utils import Usage + + system_blocks = [ + {"type": "text", "text": "part one of the system prompt. "}, + {"type": "text", "text": "part two of the system prompt."}, + ] + messages = [{"role": "user", "content": "write a short essay"}] + request_data = self._dispatched_request_data(messages, {"system": system_blocks}) + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( + model="gpt-3.5-turbo", text="part one of the system prompt. part two of the system prompt." + ) + assert estimated.prompt_tokens == expected + + @pytest.mark.asyncio + async def test_responses_instructions_counted_in_estimate(self): + import litellm as litellm_module + from litellm.types.utils import Usage + + instructions = "Answer every question as a meticulous archivist." + request_data = self._dispatched_request_data("summarize the archive", {"instructions": instructions}) + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", text="summarize the archive" + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=instructions) + assert estimated.prompt_tokens == expected + from typing import cast From 3a4d3a01af14cb9a1058ff57466869b0735f23d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:36:57 -0700 Subject: [PATCH 03/11] fix(proxy): only estimate failed-request input tokens for call types whose input is countable --- litellm/proxy/utils.py | 31 ++++++++++++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 28 +++++++++++++++++- 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5bfc1c2d1e1..1465e76d01f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -449,6 +449,35 @@ def _estimate_dispatched_failure_usage(model: str, request_input: object, system return Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens) +_INPUT_ESTIMABLE_CALL_TYPES: Final = frozenset( + call_type.value + for call_type in ( + CallTypes.completion, + CallTypes.acompletion, + CallTypes.text_completion, + CallTypes.atext_completion, + CallTypes.anthropic_messages, + CallTypes.aanthropic_messages, + CallTypes.responses, + CallTypes.aresponses, + CallTypes.embedding, + CallTypes.aembedding, + CallTypes.moderation, + CallTypes.amoderation, + CallTypes.image_generation, + CallTypes.aimage_generation, + CallTypes.speech, + CallTypes.aspeech, + CallTypes.rerank, + CallTypes.arerank, + CallTypes.generate_content, + CallTypes.agenerate_content, + CallTypes.generate_content_stream, + CallTypes.agenerate_content_stream, + ) +) + + def _failure_usage_to_lift(model_call_details: Mapping[str, object], dispatched: bool) -> tuple[object, object] | None: """A stream that broke mid-flight still billed the provider for the chunks already delivered; the streaming handler stashes that recovered usage and @@ -461,6 +490,8 @@ def _failure_usage_to_lift(model_call_details: Mapping[str, object], dispatched: return recovered_usage, model_call_details.get("response_cost") if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL): return None + if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES: + return None optional_params: Final = model_call_details.get("optional_params") system_input: Final = ( (optional_params.get("system") or optional_params.get("instructions")) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 0455a806c0b..96e2c74e5d4 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -517,6 +517,7 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: "first_api_call_start_time": datetime.now(), "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "count these input tokens please"}], + "call_type": "acompletion", } ), "metadata": {}, @@ -582,6 +583,7 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: "first_api_call_start_time": datetime.now(), "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "mid-stream failure"}], + "call_type": "acompletion", "combined_usage_object": recovered_usage, "response_cost": 3.5e-05, } @@ -605,6 +607,7 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: "first_api_call_start_time": datetime.now(), "model": "gpt-3.5-turbo", "messages": "a plain text-completion prompt string", + "call_type": "atext_completion", } ), "metadata": {}, @@ -616,7 +619,7 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: assert estimated.prompt_tokens > 0 assert estimated.completion_tokens == 0 - def _dispatched_request_data(self, messages, optional_params): + def _dispatched_request_data(self, messages, optional_params, call_type="acompletion"): from datetime import datetime return { @@ -626,11 +629,34 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: "model": "gpt-3.5-turbo", "messages": messages, "optional_params": optional_params, + "call_type": call_type, } ), "metadata": {}, } + @pytest.mark.asyncio + async def test_embedding_string_list_input_counted_in_estimate(self): + import litellm as litellm_module + from litellm.types.utils import Usage + + embedding_input = ["first embedding text", "second embedding text"] + request_data = self._dispatched_request_data(embedding_input, {}, call_type="aembedding") + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected = litellm_module.token_counter(model="gpt-3.5-turbo", text="".join(embedding_input)) + assert estimated.prompt_tokens == expected + + @pytest.mark.asyncio + async def test_transcription_checksum_not_estimated(self): + request_data = self._dispatched_request_data("a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6", {}, call_type="atranscription") + await self._run(request_data) + + assert "combined_usage_object" not in request_data + assert "response_cost" not in request_data + @pytest.mark.asyncio async def test_anthropic_system_prompt_counted_in_estimate(self): import litellm as litellm_module From e9355a7fe9d9e72b894e1dd1fcf82977af57d5df Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:49:42 -0700 Subject: [PATCH 04/11] fix(proxy): hand the embeddings failure hook the post-setup request data --- litellm/proxy/proxy_server.py | 12 +----- tests/test_litellm/proxy/test_proxy_server.py | 43 +++++++++++++++++++ 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5d4a306a73e..b9a2be82969 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10197,11 +10197,9 @@ async def embeddings( """ global proxy_logging_obj - data: Any = {} + data: Final = await _read_request_body(request=request) + base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - # Use shared request body reading helper (same as chat/completions) - data = await _read_request_body(request=request) - ### HANDLE TOKEN ARRAY INPUT DECODING ### # This must happen BEFORE base_process_llm_request() since it modifies the input router_model_names: Final = llm_router.model_names if llm_router is not None else [] @@ -10245,10 +10243,6 @@ async def embeddings( if hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None: data["metadata"]["agent_id"] = user_api_key_dict.agent_id - # Use unified request processor (same as chat/completions and responses) - base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) - - # Process the request with all optimizations (shared sessions, network tuning, etc.) response: Final = await base_llm_response_processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, @@ -10270,8 +10264,6 @@ async def embeddings( return response except Exception as e: - # Use unified error handler - base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) raise await base_llm_response_processor._handle_llm_api_exception( e=e, user_api_key_dict=user_api_key_dict, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5545ee92e84..7fdfbea843f 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11098,3 +11098,46 @@ async def test_moderations_reraises_proxy_exception_unwrapped(): assert exc_info.value.code == "400" assert exc_info.value.param == "metadata" mock_logging.post_call_failure_hook.assert_awaited_once() + + +class TestEmbeddingsFailureHookRequestData: + @pytest.mark.asyncio + async def test_failure_hook_gets_post_setup_data_with_logging_obj(self): + """Request setup replaces the processor's data dict (adding the logging + object the failure hook needs to lift token usage from); the embeddings + exception handler must pass that replaced dict, not the raw request body + dict it was rebuilt from.""" + from litellm.proxy._types import ProxyException + + captured = {} + logging_obj_sentinel = MagicMock() + + async def fake_process(self, **kwargs): + self.data = {**self.data, "litellm_logging_obj": logging_obj_sentinel} + captured["processor_data"] = self.data + raise RuntimeError("provider timeout") + + with ( + patch.object( + proxy_server_module, + "_read_request_body", + new=AsyncMock(return_value={"model": "my-embed", "input": "hello"}), + ), + patch.object( + proxy_server_module.ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + new=fake_process, + ), + patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, + ): + mock_logging.post_call_failure_hook = AsyncMock(return_value=None) + with pytest.raises(ProxyException): + await proxy_server_module.embeddings( + request=MagicMock(), + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + + hook_request_data = mock_logging.post_call_failure_hook.await_args.kwargs["request_data"] + assert hook_request_data is captured["processor_data"] + assert hook_request_data["litellm_logging_obj"] is logging_obj_sentinel From 42ddc5c5359bb32762926577fca8fcb1e5b3836d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:18:23 -0700 Subject: [PATCH 05/11] fix(proxy): estimate image message tokens without fetching the image url --- litellm/proxy/utils.py | 4 ++- tests/test_litellm/proxy/test_proxy_utils.py | 28 ++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1465e76d01f..c45b4ad17c0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -430,7 +430,9 @@ def _count_request_input_tokens(model: str, request_input: object, system_input: text_entries: Final = tuple(entry for entry in request_input if isinstance(entry, str)) if len(text_entries) == len(request_input): return system_tokens + litellm.token_counter(model=model, text="".join(text_entries)) - return system_tokens + litellm.token_counter(model=model, messages=request_input) + return system_tokens + litellm.token_counter( + model=model, messages=request_input, use_default_image_token_count=True + ) def _estimate_dispatched_failure_usage(model: str, request_input: object, system_input: object) -> Usage | None: diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 96e2c74e5d4..b70f93054d2 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -635,6 +635,34 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: "metadata": {}, } + @pytest.mark.asyncio + async def test_image_message_estimated_without_fetching_image(self): + import litellm as litellm_module + from litellm.types.utils import Usage + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this image"}, + { + "type": "image_url", + "image_url": {"url": "http://127.0.0.1:1/unreachable.png", "detail": "high"}, + }, + ], + } + ] + request_data = self._dispatched_request_data(messages, {}) + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages, use_default_image_token_count=True + ) + assert estimated.prompt_tokens == expected + assert estimated.prompt_tokens > 0 + @pytest.mark.asyncio async def test_embedding_string_list_input_counted_in_estimate(self): import litellm as litellm_module From 138c77023a4b4b0a112f6f6f737b3fe33f16148c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:44:31 -0700 Subject: [PATCH 06/11] fix: accept bool thinking param instead of crashing with AttributeError litellm.completion(thinking=True) crashed pre-network in is_thinking_enabled with a retryable APIConnectionError ('bool' object has no attribute 'get'), so the router burned retries on a deterministic failure and proxy clients got a traceback instead of a usable response. validate_and_fix_thinking_param now coerces thinking=True to the enabled dict with the default medium budget and drops thinking=False, and the remaining dict-assuming thinking accessors (base config, bedrock converse, deepseek) guard with isinstance so raw bools can never crash a transform. --- litellm/llms/base_llm/chat/transformation.py | 13 ++++++++----- .../llms/bedrock/chat/converse_transformation.py | 5 ++++- litellm/llms/deepseek/chat/transformation.py | 4 +++- litellm/main.py | 1 - litellm/utils.py | 13 +++++++++++-- .../bedrock/chat/test_converse_transformation.py | 7 +++++++ .../chat/test_deepseek_chat_transformation.py | 5 +++++ tests/test_litellm/test_thinking_enabled.py | 2 ++ tests/test_litellm/test_utils.py | 14 ++++++++++++++ 9 files changed, 54 insertions(+), 10 deletions(-) diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 0d6d942e686..d147063df73 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -5,7 +5,7 @@ Common base config for all LLM providers import types from abc import ABC, abstractmethod from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, Union import httpx from pydantic import BaseModel @@ -90,9 +90,9 @@ class BaseConfig(ABC): return type_to_response_format_param(response_format=response_format) def is_thinking_enabled(self, non_default_params: dict) -> bool: - return (non_default_params.get("thinking") or {}).get("type") == "enabled" or non_default_params.get( - "reasoning_effort" - ) is not None + thinking: Final = non_default_params.get("thinking") + thinking_type: Final = thinking.get("type") if isinstance(thinking, dict) else None + return thinking is True or thinking_type == "enabled" or non_default_params.get("reasoning_effort") is not None def is_max_tokens_in_request(self, non_default_params: dict) -> bool: """ @@ -112,7 +112,10 @@ class BaseConfig(ABC): if is_thinking_enabled and ( "max_tokens" not in non_default_params and "max_completion_tokens" not in non_default_params ): - thinking_token_budget: Final = cast(dict, optional_params["thinking"]).get("budget_tokens", None) + thinking_value: Final = optional_params.get("thinking") + thinking_token_budget: Final = ( + thinking_value.get("budget_tokens") if isinstance(thinking_value, dict) else None + ) if thinking_token_budget is not None: optional_params["max_tokens"] = thinking_token_budget + DEFAULT_MAX_TOKENS diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index fd07999395b..cee89f42c2d 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1090,7 +1090,10 @@ class AmazonConverseConfig(BaseConfig): is_thinking_enabled: Final = self.is_thinking_enabled(optional_params) is_max_tokens_in_request: Final = self.is_max_tokens_in_request(non_default_params) if is_thinking_enabled and not is_max_tokens_in_request: - thinking_token_budget: Final = cast(dict, optional_params["thinking"]).get("budget_tokens", None) + thinking_value: Final = optional_params.get("thinking") + thinking_token_budget: Final = ( + thinking_value.get("budget_tokens") if isinstance(thinking_value, dict) else None + ) if thinking_token_budget is not None: optional_params["maxTokens"] = thinking_token_budget + DEFAULT_MAX_TOKENS diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index 24da5b79261..566c960333a 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -131,9 +131,11 @@ class DeepSeekChatConfig(OpenAIGPTConfig): - model supports reasoning (capability check) - user explicitly passed thinking={"type": "enabled"} (opt-in check) """ + thinking: Final = optional_params.get("thinking") return ( supports_reasoning(model=model, custom_llm_provider="deepseek") - and (optional_params.get("thinking") or {}).get("type") == "enabled" + and isinstance(thinking, dict) + and thinking.get("type") == "enabled" ) @staticmethod diff --git a/litellm/main.py b/litellm/main.py index cc27da830d8..f0b20eba9b6 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5007,7 +5007,6 @@ def completion( tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice) # validate optional params stop = validate_openai_optional_params(stop=stop) - # normalize camelCase thinking keys (e.g. budgetTokens -> budget_tokens) thinking = validate_and_fix_thinking_param(thinking=thinking) ######### unpacking kwargs ##################### diff --git a/litellm/utils.py b/litellm/utils.py index 1c880ee9521..a7b70c4129a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -65,6 +65,7 @@ from litellm.constants import ( DEFAULT_EMBEDDING_PARAM_VALUES, DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_TRIM_RATIO, FUNCTION_DEFINITION_TOKEN_COUNT, INITIAL_RETRY_DELAY, @@ -7638,12 +7639,20 @@ def validate_and_fix_openai_tools(tools: list | None) -> list[dict] | None: def validate_and_fix_thinking_param( - thinking: AnthropicThinkingParam | None, + thinking: AnthropicThinkingParam | bool | None, ) -> AnthropicThinkingParam | None: """ - Normalizes camelCase keys in the thinking param to snake_case. + Coerces bool thinking values (True becomes enabled with the default medium budget, False becomes None) + and normalizes camelCase keys in the thinking param to snake_case. Handles clients that send budgetTokens instead of budget_tokens. """ + if thinking is True: + return cast( + "AnthropicThinkingParam", + {"type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET}, + ) + if thinking is False: + return None if thinking is None or not isinstance(thinking, dict): return thinking normalized: Final = dict(thinking) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 2509f6480d5..ff3d38eb374 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -6043,3 +6043,10 @@ def test_streaming_usage_chunk_is_transformed(): assert chunk.usage.prompt_tokens == 11 assert chunk.usage.completion_tokens == 4 assert chunk.usage.total_tokens == 15 + + +def test_update_optional_params_with_thinking_tokens_bool_thinking_does_not_crash(): + config = AmazonConverseConfig() + optional_params = {"thinking": True} + config.update_optional_params_with_thinking_tokens(non_default_params={"thinking": True}, optional_params=optional_params) + assert "maxTokens" not in optional_params diff --git a/tests/test_litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/test_litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py index ec51e5d303d..d5783e3567f 100644 --- a/tests/test_litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/test_litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -101,3 +101,8 @@ async def test_async_transform_request_strips_unsupported_tools_from_body(): assert [tool["type"] for tool in body["tools"]] == ["function"] assert body["tools"][0]["function"]["name"] == "shell" + + +def test_thinking_mode_active_bool_thinking_returns_false_without_crashing(): + config = DeepSeekChatConfig() + assert config._thinking_mode_active(model="deepseek-reasoner", optional_params={"thinking": True}) is False diff --git a/tests/test_litellm/test_thinking_enabled.py b/tests/test_litellm/test_thinking_enabled.py index 8ba406c395a..38c4534e45c 100644 --- a/tests/test_litellm/test_thinking_enabled.py +++ b/tests/test_litellm/test_thinking_enabled.py @@ -60,6 +60,8 @@ class TestIsThinkingEnabled: ({"reasoning_effort": "medium"}, True), # both thinking enabled and reasoning_effort returns True ({"thinking": {"type": "enabled"}, "reasoning_effort": "high"}, True), + # thinking=True (bool) should not crash, returns True + ({"thinking": True}, True), # falsy thinking values should not crash ({"thinking": False}, False), ({"thinking": 0}, False), diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index afdfdf170ac..efccdc4a986 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -3766,6 +3766,20 @@ class TestValidateAndFixThinkingParam: assert "budgetTokens" in thinking assert "budget_tokens" not in thinking + def test_bool_true_maps_to_enabled_with_default_budget(self): + from litellm.constants import DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET + from litellm.utils import validate_and_fix_thinking_param + + assert validate_and_fix_thinking_param(thinking=True) == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + } + + def test_bool_false_returns_none(self): + from litellm.utils import validate_and_fix_thinking_param + + assert validate_and_fix_thinking_param(thinking=False) is None + def test_deepseek_v4_models_in_cost_map(): """ From f8c8b41bf5c30cb3aebf7030c028820d97c5265b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:52:31 -0700 Subject: [PATCH 07/11] fix(proxy): backfill system prompt from the request body when estimating bridged failure tokens --- litellm/proxy/utils.py | 14 +++++-- tests/test_litellm/proxy/test_proxy_utils.py | 40 ++++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c45b4ad17c0..a743526e975 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -480,12 +480,18 @@ _INPUT_ESTIMABLE_CALL_TYPES: Final = frozenset( ) -def _failure_usage_to_lift(model_call_details: Mapping[str, object], dispatched: bool) -> tuple[object, object] | None: +def _failure_usage_to_lift( + model_call_details: Mapping[str, object], + request_body: Mapping[str, object], + dispatched: bool, +) -> tuple[object, object] | None: """A stream that broke mid-flight still billed the provider for the chunks already delivered; the streaming handler stashes that recovered usage and cost in model_call_details, so prefer it. Otherwise a request that was dispatched to a provider and failed without upstream usage gets an - estimated input-side Usage with zero cost. Returns the + estimated input-side Usage with zero cost. The raw request body backfills + the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a + chat-completions provider) without filling optional_params. Returns the (combined_usage_object, response_cost) pair to lift, or None.""" recovered_usage: Final = model_call_details.get("combined_usage_object") if recovered_usage is not None: @@ -495,11 +501,12 @@ def _failure_usage_to_lift(model_call_details: Mapping[str, object], dispatched: if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES: return None optional_params: Final = model_call_details.get("optional_params") - system_input: Final = ( + dispatched_system: Final = ( (optional_params.get("system") or optional_params.get("instructions")) if isinstance(optional_params, dict) else None ) + system_input: Final = dispatched_system or request_body.get("system") or request_body.get("instructions") estimated_usage: Final = _estimate_dispatched_failure_usage( model=str(model_call_details.get("model") or ""), request_input=model_call_details.get("messages"), @@ -2303,6 +2310,7 @@ class ProxyLogging: # is popped) record real token counts instead of zero. _usage_to_lift: Final = _failure_usage_to_lift( model_call_details=_model_call_details, + request_body=request_data, dispatched=_first_handoff is not None, ) if _usage_to_lift is not None: diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index b70f93054d2..d6cf0e30139 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -738,6 +738,46 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=instructions) assert estimated.prompt_tokens == expected + @pytest.mark.asyncio + async def test_request_body_system_counted_when_optional_params_empty(self): + import litellm as litellm_module + from litellm.types.utils import Usage + + system_prompt = "You are a meticulous cartographer who labels every landmark." + messages = [{"role": "user", "content": "draw me a map"}] + request_data = { + **self._dispatched_request_data(messages, {}, call_type="aanthropic_messages"), + "system": system_prompt, + } + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( + model="gpt-3.5-turbo", text=system_prompt + ) + assert estimated.prompt_tokens == expected + + @pytest.mark.asyncio + async def test_optional_params_system_wins_over_request_body_system(self): + import litellm as litellm_module + from litellm.types.utils import Usage + + dispatched_system = "short dispatched system prompt" + messages = [{"role": "user", "content": "hello"}] + request_data = { + **self._dispatched_request_data(messages, {"system": dispatched_system}), + "system": "a much longer request body system prompt that must not be double counted here", + } + await self._run(request_data) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( + model="gpt-3.5-turbo", text=dispatched_system + ) + assert estimated.prompt_tokens == expected + from typing import cast From 54cc988a9e5f3db851d7d372b23a2893f3896707 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:55:22 -0700 Subject: [PATCH 08/11] test: drop restating comment and wrap long call in thinking tests --- .../llms/bedrock/chat/test_converse_transformation.py | 4 +++- tests/test_litellm/test_thinking_enabled.py | 1 - 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index ff3d38eb374..298360789eb 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -6048,5 +6048,7 @@ def test_streaming_usage_chunk_is_transformed(): def test_update_optional_params_with_thinking_tokens_bool_thinking_does_not_crash(): config = AmazonConverseConfig() optional_params = {"thinking": True} - config.update_optional_params_with_thinking_tokens(non_default_params={"thinking": True}, optional_params=optional_params) + config.update_optional_params_with_thinking_tokens( + non_default_params={"thinking": True}, optional_params=optional_params + ) assert "maxTokens" not in optional_params diff --git a/tests/test_litellm/test_thinking_enabled.py b/tests/test_litellm/test_thinking_enabled.py index 38c4534e45c..744b258e617 100644 --- a/tests/test_litellm/test_thinking_enabled.py +++ b/tests/test_litellm/test_thinking_enabled.py @@ -60,7 +60,6 @@ class TestIsThinkingEnabled: ({"reasoning_effort": "medium"}, True), # both thinking enabled and reasoning_effort returns True ({"thinking": {"type": "enabled"}, "reasoning_effort": "high"}, True), - # thinking=True (bool) should not crash, returns True ({"thinking": True}, True), # falsy thinking values should not crash ({"thinking": False}, False), From b9a267c69378ef1bfaa7541c521934c6473e5faf Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 18 Aug 2026 22:37:04 -0700 Subject: [PATCH 09/11] refactor(ui): migrate the teams form graph off antd Form onto react-hook-form (#37417) * test(ui): pin the teams create and update payloads before the form migration The teams graph (Teams.tsx, TeamInfo.tsx and the MetadataKeyValueFields child they share) is next for the antd Form to react-hook-form migration, and its submit payload is a function of which collapsible sections the user happened to open. Nine sections across the two files use the shadcn Collapsible, none of them passes keepMounted, and Base UI unmounts the closed branch, so a closed section registers nothing and its keys never reach the request body. That matters beyond parity. /team/update reads the body with exclude_unset, so an omitted key is never written, while an explicitly null team member budget key reaches clear_team_member_budget_fields and nulls max_budget, budget_duration, rpm_limit and tpm_limit on the shared budget row. antd cannot reach that today because the field is unregistered rather than null. A port that seeds those fields or coalesces on the way into the payload would turn a save with the section never opened into a silent clear. The coverage that shipped with the team modal reached one of the four gating sections on the create side and asserted key sets rather than the request body, so a null where antd sent undefined would have passed. These cases assert both the raw payload and its JSON round trip with toStrictEqual, which is what separates absent from null from undefined, and they cover every gating section on both screens. Also pinned, because each is a live behaviour a port can quietly change: - the create path sends max_budget, tpm_limit and rpm_limit as strings, while team_member_budget arrives as a number through its normalize prop - an invalid secret manager config blocks the create with its rule message suppressed by the item's help prop, so nothing is shown to the user - the disable global guardrails switch is inert for a non premium user - a value typed into a section survives collapsing and re-expanding it Verified by adding keepMounted to all nine panels, which is the change a porter reaches for on noticing that fields go missing: 35 of 118 went red, including every one of these cases. The files were restored byte identical afterwards. No production file changes here. 145 tests pass across the three files. * refactor(ui): migrate the teams form graph off antd Form onto react-hook-form Teams.tsx and TeamInfo.tsx were the last large antd `Form` graph in the dashboard. Both now use `useZodForm` + `FormField`, with the shared `MetadataKeyValueFields` child converted to a `useFieldArray`. antd only returns the mounted registered fields from `onFinish`, so a closed collapsible contributed no keys at all. react-hook-form keeps unmounted values in the store (and `shouldUnregister: true` would lose them on re-expand), so both forms project the submitted values through the currently mounted section list before handing them to the existing payload builders. Closed sections therefore still produce absent keys rather than nulls, which matters at /team/update where an explicit null clears the shared budget row. Widgets that had no shadcn equivalent are replaced with the existing shared ones: SearchSelect for the organization pickers, MultiSelect for default member models, TagsInput for guardrails/policies, and a new GuardrailsSelect for the grouped global/other guardrail dropdown. * refactor(ui): forward the field ref to NumericalInput in the teams forms staging turned NumericalInput into a forwardRef, so the teams graph can stop dropping the react-hook-form ref on the floor. * test(ui): pin the capability gate, required rules and guardrail kill switch A mutation run over the ported teams forms found five survivors the payload cases did not reach: the viewPolicies gate on both forms, the team name rule on both forms, the guardrail kill switch resync, and the number coercion on a typed model rate limit. Six cases close them. --- .../components/ModelSelect/ModelSelect.tsx | 5 +- .../src/components/Teams.test.tsx | 258 +++ ui/litellm-dashboard/src/components/Teams.tsx | 1247 +++++++------- .../MetadataKeyValueFields.test.tsx | 20 +- .../MetadataKeyValueFields.tsx | 131 +- .../components/shared/form/LabelWithHint.tsx | 34 + .../src/components/team/GuardrailsSelect.tsx | 133 ++ .../src/components/team/TeamInfo.test.tsx | 316 +++- .../src/components/team/TeamInfo.tsx | 1471 ++++++++++------- 9 files changed, 2362 insertions(+), 1253 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/shared/form/LabelWithHint.tsx create mode 100644 ui/litellm-dashboard/src/components/team/GuardrailsSelect.tsx diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx index d3ede5f4318..c3936cc456d 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx @@ -40,6 +40,7 @@ export const MODEL_SENTINEL_OPTIONS = [ const MAX_VISIBLE_MODEL_CHIPS = 5; export interface ModelSelectProps { + id?: string; teamID?: string; organizationID?: string; options?: { @@ -122,7 +123,7 @@ const filterModels = ( export const ModelSelect = (props: ModelSelectProps) => { const anchor = useComboboxAnchor(); - const { teamID, organizationID, options, context, dataTestId, value = [], onChange, style } = props; + const { id, teamID, organizationID, options, context, dataTestId, value = [], onChange, style } = props; const { showAllProxyModelsOverride, includeSpecialOptions } = options || {}; const { data: allProxyModels, isLoading: isLoadingAllProxyModels } = useAllProxyModels(); const { data: team, isLoading: isLoadingTeam } = useTeam(teamID); @@ -256,7 +257,7 @@ export const ModelSelect = (props: ModelSelectProps) => { )} - + No models found diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 0ef4357e73f..d0d41837dce 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -1223,3 +1223,261 @@ describe("Teams - which fields reach the create payload depends on the open sect expect(payload.team_id).toBe("tid-kept"); }); }); + +describe("Teams - the exact bytes the create call sends", () => { + beforeEach(() => { + vi.clearAllMocks(); + can.mockReturnValue(true); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(getPoliciesList).mockResolvedValue({ policies: [] }); + vi.mocked(getDefaultTeamSettings).mockResolvedValue({ values: {} }); + vi.mocked(teamCreateCall).mockResolvedValue({ team_id: "new-team-1" }); + vi.mocked(useTeamMetadataSchema).mockReturnValue({ data: [], isLoading: false } as any); + mockUseOrganizations.mockReturnValue({ data: null }); + }); + + const openCreateModal = async (options?: { premiumUser?: boolean }) => { + renderWithQueryClient( + , + ); + act(() => { + fireEvent.click(screen.getAllByRole("button", { name: /create team/i })[0]); + }); + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Byte Contract Team" } }); + }; + + const submit = async () => { + const buttons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(buttons[buttons.length - 1]); + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalled(); + }); + return vi.mocked(teamCreateCall).mock.calls[0][1] as Record; + }; + + const wireBody = (payload: Record) => JSON.parse(JSON.stringify(payload)) as Record; + + const openSection = async (title: string, mountedProbe: RegExp | string) => { + fireEvent.click(screen.getByText(title)); + await waitFor(() => { + expect(screen.getAllByText(mountedProbe).length).toBeGreaterThan(0); + }); + }; + + it("sends three keys and nothing else when every section is left closed", async () => { + await openCreateModal(); + + const payload = await submit(); + + expect(payload).toStrictEqual({ + team_alias: "Byte Contract Team", + organization_id: null, + models: ["no-default-models"], + max_budget: undefined, + budget_duration: undefined, + tpm_limit: undefined, + rpm_limit: undefined, + metadata: undefined, + }); + expect(wireBody(payload)).toStrictEqual({ + team_alias: "Byte Contract Team", + organization_id: null, + models: ["no-default-models"], + }); + }); + + it("keeps every newly mounted but untouched field out of the request body", async () => { + await openCreateModal(); + + await openSection("Additional Settings", /Team Member Key Duration/); + await openSection("MCP Settings", /Allowed MCP Servers/); + await openSection("Agent Settings", /Allowed Agents/); + await openSection("Search Tool Settings", /Allowed Search Tools/); + + const payload = await submit(); + + expect(payload).toStrictEqual({ + team_alias: "Byte Contract Team", + organization_id: null, + models: ["no-default-models"], + max_budget: undefined, + budget_duration: undefined, + tpm_limit: undefined, + rpm_limit: undefined, + metadata: undefined, + team_id: undefined, + team_member_budget: undefined, + team_member_key_duration: undefined, + team_member_rpm_limit: undefined, + team_member_tpm_limit: undefined, + secret_manager_settings: undefined, + guardrails: undefined, + disable_global_guardrails: undefined, + policies: undefined, + access_group_ids: undefined, + allowed_vector_store_ids: undefined, + allowed_passthrough_routes: undefined, + allowed_mcp_servers_and_groups: undefined, + mcp_tool_permissions: {}, + allowed_agents_and_groups: undefined, + object_permission_search_tools: undefined, + }); + expect(wireBody(payload)).toStrictEqual({ + team_alias: "Byte Contract Team", + organization_id: null, + models: ["no-default-models"], + mcp_tool_permissions: {}, + }); + }); + + it.each([ + ["MCP Settings", /Allowed MCP Servers/, ["allowed_mcp_servers_and_groups", "mcp_tool_permissions"]], + ["Agent Settings", /Allowed Agents/, ["allowed_agents_and_groups"]], + ["Search Tool Settings", /Allowed Search Tools/, ["object_permission_search_tools"]], + ])("registers %s fields only while that one section is open", async (title, probe, keys) => { + await openCreateModal(); + + const closedPayload = await submit(); + for (const key of keys as string[]) { + expect(closedPayload).not.toHaveProperty(key); + } + }); + + it("carries every typed value to the payload at the type antd sends today", async () => { + await openCreateModal(); + + fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "150.75" } }); + fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "900" } }); + fireEvent.change(screen.getByLabelText("Requests per minute Limit (RPM)"), { target: { value: "800" } }); + + await openSection("Additional Settings", /Team Member Key Duration/); + + fireEvent.change(screen.getByLabelText("Team ID"), { target: { value: "tid-1" } }); + fireEvent.change(screen.getByLabelText("Team Member Budget (USD)"), { target: { value: "12.5" } }); + fireEvent.change(screen.getByLabelText(/Team Member Key Duration/), { target: { value: "30d" } }); + fireEvent.change(screen.getByLabelText("Team Member RPM Limit"), { target: { value: "7" } }); + fireEvent.change(screen.getByLabelText("Team Member TPM Limit"), { target: { value: "8" } }); + fireEvent.change(screen.getByLabelText("Secret Manager Settings"), { + target: { value: '{"namespace":"admin"}' }, + }); + + const payload = await submit(); + + expect(payload.max_budget).toBe("150.75"); + expect(payload.tpm_limit).toBe("900"); + expect(payload.rpm_limit).toBe("800"); + expect(payload.team_id).toBe("tid-1"); + expect(payload.team_member_budget).toBe(12.5); + expect(payload.team_member_key_duration).toBe("30d"); + expect(payload.team_member_rpm_limit).toBe("7"); + expect(payload.team_member_tpm_limit).toBe("8"); + expect(payload.secret_manager_settings).toStrictEqual({ namespace: "admin" }); + }); + + it("blocks the create on an invalid secret manager config, with the rule message suppressed by help", async () => { + await openCreateModal(); + await openSection("Additional Settings", /Team Member Key Duration/); + + fireEvent.change(screen.getByLabelText("Secret Manager Settings"), { target: { value: " " } }); + + const buttons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(buttons[buttons.length - 1]); + + await waitFor(() => { + expect(screen.getByLabelText("Secret Manager Settings")).toHaveAttribute("aria-invalid", "true"); + }); + expect(teamCreateCall).not.toHaveBeenCalled(); + expect(screen.queryByText("Please enter valid JSON")).not.toBeInTheDocument(); + }); + + it("turns the disable-global-guardrails switch into a boolean for a premium user", async () => { + await openCreateModal({ premiumUser: true }); + await openSection("Additional Settings", /Team Member Key Duration/); + + const switches = screen.getAllByRole("switch"); + fireEvent.click(switches[switches.length - 1]); + + const payload = await submit(); + + expect(payload.disable_global_guardrails).toBe(true); + }); + + it("leaves the disable-global-guardrails switch inert for a non-premium user", async () => { + await openCreateModal(); + await openSection("Additional Settings", /Team Member Key Duration/); + + const switches = screen.getAllByRole("switch"); + fireEvent.click(switches[switches.length - 1]); + + const payload = await submit(); + + expect(payload.disable_global_guardrails).toBeUndefined(); + }); + + it.each([ + ["MCP Settings", /Allowed MCP Servers/, ["allowed_mcp_servers_and_groups", "mcp_tool_permissions"]], + ["Agent Settings", /Allowed Agents/, ["allowed_agents_and_groups"]], + ["Search Tool Settings", /Allowed Search Tools/, ["object_permission_search_tools"]], + ])("adds the %s keys as soon as that one section is opened", async (title, probe, keys) => { + await openCreateModal(); + + await openSection(title as string, probe as RegExp); + const payload = await submit(); + + for (const key of keys as string[]) { + expect(payload).toHaveProperty(key); + } + }); + + it("leaves policies out of the request body for a caller without the viewPolicies capability", async () => { + can.mockReturnValue(false); + + await openCreateModal(); + await openSection("Additional Settings", /Team Member Key Duration/); + + const payload = await submit(); + + expect(payload).toStrictEqual({ + team_alias: "Byte Contract Team", + organization_id: null, + models: ["no-default-models"], + max_budget: undefined, + budget_duration: undefined, + tpm_limit: undefined, + rpm_limit: undefined, + metadata: undefined, + team_id: undefined, + team_member_budget: undefined, + team_member_key_duration: undefined, + team_member_rpm_limit: undefined, + team_member_tpm_limit: undefined, + secret_manager_settings: undefined, + guardrails: undefined, + disable_global_guardrails: undefined, + access_group_ids: undefined, + allowed_vector_store_ids: undefined, + allowed_passthrough_routes: undefined, + }); + }); + + it("blocks the create on an empty team name and names the rule", async () => { + renderWithQueryClient(); + act(() => { + fireEvent.click(screen.getAllByRole("button", { name: /create team/i })[0]); + }); + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + const buttons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(buttons[buttons.length - 1]); + + expect(await screen.findByText("Please input a team name")).toBeInTheDocument(); + expect(teamCreateCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index b1f96f51028..94e84c6e721 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -4,12 +4,21 @@ import AvailableTeamsPanel from "@/components/team/AvailableTeamsPanel"; import TeamInfoView from "@/components/team/TeamInfo"; import TeamSSOSettings from "@/components/TeamSSOSettings"; import { isProxyAdminRole } from "@/utils/roles"; -import { InfoCircleOutlined } from "@ant-design/icons"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { Input as UIInput } from "@/components/ui/input"; -import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, theme, Tooltip, Typography } from "antd"; +import { Switch } from "@/components/ui/switch"; +import { Textarea } from "@/components/ui/textarea"; +import { TooltipProvider } from "@/components/ui/tooltip"; +import { Field, FieldDescription, FieldGroup, FieldLabel } from "@/components/shared/form/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { labelWithDocsHint, labelWithHint } from "@/components/shared/form/LabelWithHint"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { TagsInput } from "@/app/(dashboard)/guardrails/_components/content_filter/TagsInput"; +import { Layout, Modal, Tabs, theme } from "antd"; import { ChevronDown, Plus, Users } from "lucide-react"; -import React, { useEffect, useState } from "react"; +import React, { useEffect, useMemo, useState } from "react"; +import { z } from "zod/v4"; import { useQuery, useQueryClient } from "@tanstack/react-query"; import { PageHeader } from "@/components/shared/PageHeader"; import { Button as UIButton } from "@/components/ui/button"; @@ -17,7 +26,10 @@ import { teamsTableKeys } from "@/app/(dashboard)/hooks/teams/useTeams"; import { parseAsString, useQueryState } from "nuqs"; import { TeamsTable } from "./TeamsPage/TeamsTable"; import AccessGroupSelector from "./common_components/AccessGroupSelector"; -import MetadataKeyValueFields, { metadataPairsToObject } from "./common_components/MetadataKeyValueFields"; +import MetadataKeyValueFields, { + metadataPairsSchema, + metadataPairsToObject, +} from "./common_components/MetadataKeyValueFields"; import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import PassThroughRoutesSelector from "./common_components/PassThroughRoutesSelector"; import AgentSelector from "./agent_management/AgentSelector"; @@ -51,6 +63,102 @@ import { teamCreateCall } from "./networking"; import { normalizeTeamModelSelection } from "./team/teamModelAccess"; import { ModelSelect } from "./ModelSelect/ModelSelect"; +const SUPPRESSED_BY_DESCRIPTION = ""; + +const numericInputSchema = z.union([z.string(), z.number()]).optional(); + +const teamCreateFieldsSchema = z.object({ + team_alias: z.string().min(1, "Please input a team name"), + organization_id: z.string().nullish(), + models: z.array(z.string()).optional(), + max_budget: numericInputSchema, + budget_duration: z.string().nullish(), + tpm_limit: numericInputSchema, + rpm_limit: numericInputSchema, + metadata: metadataPairsSchema.optional(), + team_id: z.string().optional(), + team_member_budget: z.number().optional(), + team_member_key_duration: z.string().optional(), + team_member_rpm_limit: numericInputSchema, + team_member_tpm_limit: numericInputSchema, + secret_manager_settings: z.string().optional(), + guardrails: z.array(z.string()).optional(), + disable_global_guardrails: z.boolean().optional(), + policies: z.array(z.string()).optional(), + access_group_ids: z.array(z.string()).optional(), + allowed_vector_store_ids: z.array(z.string()).optional(), + allowed_passthrough_routes: z.array(z.string()).optional(), + allowed_mcp_servers_and_groups: z + .object({ + servers: z.array(z.string()), + accessGroups: z.array(z.string()), + toolsets: z.array(z.string()).optional(), + }) + .optional(), + mcp_tool_permissions: z.record(z.string(), z.array(z.string())).optional(), + allowed_agents_and_groups: z.object({ agents: z.array(z.string()), accessGroups: z.array(z.string()) }).optional(), + object_permission_search_tools: z.array(z.string()).optional(), +}); + +type TeamCreateFormValues = z.infer; + +const EMPTY_TEAM_CREATE_VALUES: TeamCreateFormValues = { + team_alias: "", + organization_id: null, + models: [], + max_budget: undefined, + budget_duration: undefined, + tpm_limit: undefined, + rpm_limit: undefined, + metadata: [], + team_id: undefined, + team_member_budget: undefined, + team_member_key_duration: undefined, + team_member_rpm_limit: undefined, + team_member_tpm_limit: undefined, + secret_manager_settings: undefined, + guardrails: undefined, + disable_global_guardrails: undefined, + policies: undefined, + access_group_ids: undefined, + allowed_vector_store_ids: undefined, + allowed_passthrough_routes: undefined, + allowed_mcp_servers_and_groups: undefined, + mcp_tool_permissions: {}, + allowed_agents_and_groups: undefined, + object_permission_search_tools: undefined, +}; + +const ADDITIONAL_SETTINGS_FIELDS = [ + "team_id", + "team_member_budget", + "team_member_key_duration", + "team_member_rpm_limit", + "team_member_tpm_limit", + "secret_manager_settings", + "guardrails", + "disable_global_guardrails", + "policies", + "access_group_ids", + "allowed_vector_store_ids", + "allowed_passthrough_routes", +] as const; +const MCP_SETTINGS_FIELDS = ["allowed_mcp_servers_and_groups", "mcp_tool_permissions"] as const; +const AGENT_SETTINGS_FIELDS = ["allowed_agents_and_groups"] as const; +const SEARCH_TOOL_SETTINGS_FIELDS = ["object_permission_search_tools"] as const; + +const isParsableJson = (value: string | undefined): boolean => { + if (!value) { + return true; + } + try { + JSON.parse(value); + return true; + } catch { + return false; + } +}; + const canCreateOrManageTeams = ( userRole: string | null, userID: string | null, @@ -101,7 +209,29 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [currentOrg] = useState(null); const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState(null); - const [form] = Form.useForm(); + const isOrgAdmin = userRole !== "Admin"; + const [additionalSettingsOpen, setAdditionalSettingsOpen] = useState(false); + const [mcpSettingsOpen, setMcpSettingsOpen] = useState(false); + const [agentSettingsOpen, setAgentSettingsOpen] = useState(false); + const [searchToolSettingsOpen, setSearchToolSettingsOpen] = useState(false); + + const teamCreateSchema = useMemo( + () => + teamCreateFieldsSchema.superRefine((values, ctx) => { + if (isOrgAdmin && !values.organization_id) { + ctx.addIssue({ code: "custom", message: SUPPRESSED_BY_DESCRIPTION, path: ["organization_id"] }); + } + if (additionalSettingsOpen && !isParsableJson(values.secret_manager_settings)) { + ctx.addIssue({ code: "custom", message: SUPPRESSED_BY_DESCRIPTION, path: ["secret_manager_settings"] }); + } + }), + [isOrgAdmin, additionalSettingsOpen], + ); + + const form = useZodForm(teamCreateSchema, { defaultValues: EMPTY_TEAM_CREATE_VALUES }); + const watchedOrganizationId = form.watch("organization_id"); + const watchedMcpSelection = form.watch("allowed_mcp_servers_and_groups"); + const watchedToolPermissions = form.watch("mcp_tool_permissions"); const [selectedTeam, setSelectedTeam] = useState(null); const [selectedTeamId, setSelectedTeamId] = useQueryState("team", parseAsString.withOptions({ history: "push" })); @@ -134,27 +264,26 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser : "n/a"; useEffect(() => { - form.setFieldValue("models", []); + form.setValue("models", []); }, [currentOrgForCreateTeam, userModels]); // Handle organization preselection when modal opens useEffect(() => { if (isTeamModalVisible) { const adminOrgs = getAdminOrganizations(userRole, userID, organizations); - const isOrgAdmin = userRole !== "Admin"; // Org admins must scope a team to an org, so with exactly one we preselect it. // Proxy admins can create org-less teams, so the field stays optional regardless of org count. if (isOrgAdmin && adminOrgs.length === 1) { const org = adminOrgs[0]; - form.setFieldValue("organization_id", org.organization_id); + form.setValue("organization_id", org.organization_id); setCurrentOrgForCreateTeam(org); } else { - form.setFieldValue("organization_id", currentOrg?.organization_id || null); + form.setValue("organization_id", currentOrg?.organization_id || null); setCurrentOrgForCreateTeam(currentOrg); } } - }, [isTeamModalVisible, userRole, userID, organizations, currentOrg]); + }, [isTeamModalVisible, isOrgAdmin, userRole, userID, organizations, currentOrg]); // Add this useEffect to fetch guardrails useEffect(() => { @@ -190,22 +319,26 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser if (canViewPolicies) fetchPolicies(); }, [accessToken, canViewPolicies]); - const handleOk = () => { - setIsTeamModalVisible(false); - form.resetFields(); + const resetCreateForm = () => { + form.reset(EMPTY_TEAM_CREATE_VALUES); + setAdditionalSettingsOpen(false); + setMcpSettingsOpen(false); + setAgentSettingsOpen(false); + setSearchToolSettingsOpen(false); setLoggingSettings([]); setModelAliases({}); setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); }; + const handleOk = () => { + setIsTeamModalVisible(false); + resetCreateForm(); + }; + const handleCancel = () => { setIsTeamModalVisible(false); - form.resetFields(); - setLoggingSettings([]); - setModelAliases({}); - setRouterSettings(null); - setRouterSettingsKey((prev) => prev + 1); + resetCreateForm(); }; const handleDelete = async (team: Team) => { @@ -378,11 +511,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser await teamCreateCall(accessToken, { ...formValues, models: normalizeTeamModelSelection(formValues.models) }); toast.success("Team created"); await refreshTeams(); - form.resetFields(); - setLoggingSettings([]); - setModelAliases({}); - setRouterSettings(null); - setRouterSettingsKey((prev) => prev + 1); + resetCreateForm(); setIsTeamModalVisible(false); } } catch (error) { @@ -391,6 +520,19 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser } }; + const mountedCreateValues = (values: TeamCreateFormValues): Record => { + const unmounted = new Set([ + ...(additionalSettingsOpen ? [] : ADDITIONAL_SETTINGS_FIELDS), + ...(additionalSettingsOpen && canViewPolicies ? [] : ["policies"]), + ...(mcpSettingsOpen ? [] : MCP_SETTINGS_FIELDS), + ...(agentSettingsOpen ? [] : AGENT_SETTINGS_FIELDS), + ...(searchToolSettingsOpen ? [] : SEARCH_TOOL_SETTINGS_FIELDS), + ]); + return Object.fromEntries(Object.entries(values).filter(([key]) => !unmounted.has(key))); + }; + + const onCreateSubmit = (values: TeamCreateFormValues) => handleCreate(mountedCreateValues(values)); + const is_team_admin = (team: any) => { if (team == null || team.members_with_roles == null) { return false; @@ -405,7 +547,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }; const { token } = theme.useToken(); - const { Text } = Typography; const { Content } = Layout; const tabItems = [ @@ -531,556 +672,538 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser onCancel={handleCancel} destroyOnHidden > -
- <> - - - - {(() => { - const adminOrgs = getAdminOrganizations(userRole, userID, organizations); - const isOrgAdmin = userRole !== "Admin"; - const isSingleOrg = adminOrgs.length === 1; - const hasNoOrgs = adminOrgs.length === 0; - - return ( - <> - - Organization{" "} - - Organizations can have multiple teams. Learn more about{" "} - e.stopPropagation()} - > - user management hierarchy - - - } - > - - - - } - name="organization_id" - initialValue={currentOrg ? currentOrg.organization_id : null} - className="mt-8" - rules={ - isOrgAdmin - ? [ - { - required: true, - message: "Please select an organization", - }, - ] - : [] - } - help={ - isOrgAdmin && isSingleOrg - ? "You can only create teams within this organization" - : isOrgAdmin - ? "required" - : "" - } - > - - - - {/* Show message when org admin needs to select organization */} - {isOrgAdmin && !isSingleOrg && adminOrgs.length > 1 && ( -
- - Please select an organization to create a team for. You can only create teams within - organizations where you are an admin. - -
- )} - - ); - })()} - - Models{" "} - - - - - } - name="models" - > - form.setFieldValue("models", values)} - organizationID={form.getFieldValue("organization_id")} - options={{ - includeSpecialOptions: true, - showAllProxyModelsOverride: !form.getFieldValue("organization_id"), - }} - context="team" - dataTestId="create-team-models-select" - /> - - - - - - - - - - - - - - - - - - - - - Additional Settings - - - - - { - e.target.value = e.target.value.trim(); - }} - /> - - (value ? Number(value) : undefined)} - tooltip="This is the individual budget for a user in the team." - > - - - - - - - - - - - - { - if (!value) { - return Promise.resolve(); - } - try { - JSON.parse(value); - return Promise.resolve(); - } catch (error) { - return Promise.reject(new Error("Please enter valid JSON")); - } - }, - }, - ]} - > - - - - Guardrails{" "} - - e.stopPropagation()} - > - - - - - } - name="guardrails" - className="mt-8" - help="Select existing guardrails or enter new ones" - > - ({ - value: name, - label: name, - }))} - /> - + + + + + {({ ref, value, ...field }) => ( + )} - - Access Groups{" "} - - - - - } - name="access_group_ids" - className="mt-8" - help="Select access groups to assign to this team" - > - - - - Allowed Vector Stores{" "} - - - - - } - name="allowed_vector_store_ids" - className="mt-8" - help="Select vector stores this team can access. Leave empty for access to all vector stores" - > - form.setFieldValue("allowed_vector_store_ids", values)} - value={form.getFieldValue("allowed_vector_store_ids")} - accessToken={accessToken || ""} - placeholder="Select vector stores (optional)" - /> - - - - - - + + {(() => { + const adminOrgs = getAdminOrganizations(userRole, userID, organizations); + const isSingleOrg = adminOrgs.length === 1; + const hasNoOrgs = adminOrgs.length === 0; - - - MCP Settings - - - - - Allowed MCP Servers{" "} - - - - - } - name="allowed_mcp_servers_and_groups" - className="mt-4" - help="Select MCP servers or access groups this team can access" - > - form.setFieldValue("allowed_mcp_servers_and_groups", val)} - value={form.getFieldValue("allowed_mcp_servers_and_groups")} - accessToken={accessToken || ""} - placeholder="Select MCP servers or access groups (optional)" - allowAllProxyMcpServers={isProxyAdminRole(userRole || "")} + return ( + <> + + {({ id, value, onChange }) => ( + ({ + value: org.organization_id ?? "", + label: org.organization_alias ?? "", + sublabel: org.organization_id ?? "", + }))} + disabled={isOrgAdmin && isSingleOrg} + allowClear={!isOrgAdmin} + placeholder={hasNoOrgs ? "No organizations available" : "Search or select an Organization"} + emptyText="No organizations available" + onValueChange={(next) => { + onChange(next === "" ? null : next); + setCurrentOrgForCreateTeam(adminOrgs.find((org) => org.organization_id === next) ?? null); + }} + /> + )} + + + {isOrgAdmin && !isSingleOrg && adminOrgs.length > 1 && ( +
+ + Please select an organization to create a team for. You can only create teams within + organizations where you are an admin. + +
+ )} + + ); + })()} + + {({ id, value, onChange }) => ( + - + )} + - {/* Hidden field to register mcp_tool_permissions with the form */} - + + {({ ref, value, ...field }) => ( + + )} + + + {({ id, value, onChange }) => ( + + )} + + + {({ ref, value, ...field }) => ( + + )} + + + {({ ref, value, ...field }) => ( + + )} + + + Metadata + + + Values are saved as text. Enter JSON for typed values, e.g. 3, true, or {'{"region": "us"}'}. + + - - prevValues.allowed_mcp_servers_and_groups !== currentValues.allowed_mcp_servers_and_groups || - prevValues.mcp_tool_permissions !== currentValues.mcp_tool_permissions - } - > - {() => ( -
- + + Additional Settings + + + + + + {({ ref, value, ...field }) => } + + + {({ ref, value, onChange, ...field }) => ( + ) => + onChange(event.target.value ? Number(event.target.value) : undefined) + } + step={0.01} + precision={2} + width={200} + /> + )} + + + {({ ref, value, ...field }) => ( + + )} + + + {({ ref, value, ...field }) => ( + + )} + + + {({ ref, value, ...field }) => ( + + )} + + + {({ ref, value, ...field }) => ( +