Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_health_check_db_storm

litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql

@@ -0,0 +1,15 @@
+-- DropForeignKey
+DO $$
+BEGIN
+    IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
+        ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey";
+    END IF;
+END $$;
+
+-- AddForeignKey
+DO $$
+BEGIN
+    IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
+        ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE;
+    END IF;
+END $$;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
     updated_at        DateTime @default(now()) @updatedAt
     updated_by        String?
 
-    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
+    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
 
     @@unique([jwt_claim_name, jwt_claim_value])
     @@index([jwt_claim_name, jwt_claim_value, is_active])

litellm/a2a_protocol/providers/bedrock_agentcore/handler.py

@@ -13,6 +13,7 @@ from litellm._logging import verbose_logger
 from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
     BedrockAgentCoreA2ATransformation,
 )
+from litellm.llms.bedrock.base_aws_llm import run_aws_signing
 from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
 from litellm.types.llms.custom_http import httpxSpecialProvider
 
@@ -45,7 +46,8 @@ class BedrockAgentCoreA2AHandler:
         Returns:
             A2A JSON-RPC response dict from the AgentCore agent
         """
-        url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
+        url, headers, body = await run_aws_signing(
+            BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
             request_id=request_id,
             params=params,
             litellm_params=litellm_params,
@@ -91,7 +93,8 @@ class BedrockAgentCoreA2AHandler:
         Yields:
             A2A streaming response events from the AgentCore agent
         """
-        url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
+        url, headers, body = await run_aws_signing(
+            BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
             request_id=request_id,
             params=params,
             litellm_params=litellm_params,

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -25,7 +25,7 @@ import litellm
 from litellm import ModelResponse
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
-    responses_reasoning_item_from_thinking_blocks,
+    responses_reasoning_items_from_thinking_blocks,
 )
 from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
 from litellm.llms.base_llm.bridges.completion_transformation import (
@@ -129,8 +129,8 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]:
         return stored
     raw_blocks: Final = msg.get("thinking_blocks") or ()
     blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks)  # cast-ok: untyped client json
-    from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks)
-    return [] if from_thinking is None else [dict(from_thinking)]  # mutable-ok: API message payload
+    replayed: Final = responses_reasoning_items_from_thinking_blocks(blocks)
+    return [dict(item) for item in replayed]  # mutable-ok: API message payload
 
 
 def _build_reasoning_item(
@@ -227,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
     provider_specific_fields: Mapping[str, object]
 
 
-def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
+def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
     """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
     completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
     string payload in ``input`` rather than ``arguments``; both map to
@@ -755,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     # Tool calls accumulate into the single trailing tool_calls choice
                     # like the typed branches above; a choice per call would hide every
                     # call after choices[0] from chat clients
-                    accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index))
+                    accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
                     tool_call_index += 1
                 elif handle_raw_dict_callback is not None:
                     choice, index = handle_raw_dict_callback(item=raw_item, index=index)
@@ -1409,7 +1409,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
             # New output item added
             output_item = parsed_chunk.get("item", {})
             if output_item.get("type") in ("function_call", "custom_tool_call"):
-                converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
+                converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
                 provider_specific_fields: Final = converted.get("provider_specific_fields")
 
                 function_chunk: Final = ChatCompletionToolCallFunctionChunk(
@@ -1484,7 +1484,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                                 index=0,
                                 delta=Delta(
                                     tool_calls=(
-                                        _tool_call_dict_from_output_item(
+                                        tool_call_dict_from_output_item(
                                             output_item, parsed_chunk.get("output_index", 0)
                                         ),
                                     )

litellm/constants.py

@@ -398,6 +398,18 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range(
     minimum=1,
     maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS,
 )
+TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range(
+    "TOKEN_COUNTER_MAX_EXACT_CHARS",
+    default=4_000_000,
+    minimum=1,
+    maximum=1_000_000_000,
+)
+TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range(
+    "TOKEN_COUNTER_MAX_CONCURRENT_COUNTS",
+    default=4,
+    minimum=1,
+    maximum=256,
+)
 MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
 MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
 OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))
@@ -570,6 +582,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float(
 LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100)
 LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000)
 LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0
+AWS_SIGNING_MAX_THREADS: Final = 16
 DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
     "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
 )

litellm/integrations/s3_v2.py

@@ -24,7 +24,7 @@ from litellm.integrations.s3 import (
 from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
-from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
 from litellm.llms.custom_httpx.http_handler import (
     _get_httpx_client,
     get_async_httpx_client,
@@ -366,7 +366,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
             # Sign the request
             aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers)
             aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name)
-            S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
+            await run_aws_signing(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request)
 
             # Prepare the signed headers
             signed_headers: Final = dict(aws_request.headers.items())
@@ -597,7 +597,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
 
             # Sign the request
             aws_request: Final = AWSRequest(method="GET", url=url, headers=headers)
-            S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request)
+            await run_aws_signing(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request)
 
             # Prepare the signed headers
             signed_headers: Final = dict(aws_request.headers.items())

litellm/integrations/sqs.py

@@ -22,7 +22,7 @@ from litellm.constants import (
     SQS_SEND_MESSAGE_ACTION,
 )
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
-from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
@@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM):
                 data=prepped.body,
                 headers=prepped.headers,
             )
-            SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth(aws_request)
+            await run_aws_signing(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request)
 
             signed_headers: Final = dict(aws_request.headers.items())
 

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -6,7 +6,7 @@ import io
 import json
 import mimetypes
 import re
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Iterator, Mapping, Sequence
 from itertools import groupby
 from os import PathLike
 from pathlib import Path
@@ -1823,14 +1823,11 @@ def _extract_reasoning_content(message: dict) -> tuple[str | None, str | None]:
     return None, message_content
 
 
-def _readable_thinking_text(
-    block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
-) -> str:
+def _readable_thinking_text(block: Mapping[str, object]) -> str:
     """The text a chat model can read back, empty for redacted blocks and malformed ones."""
     if block.get("type") != "thinking":
         return ""
-    thinking: Final = cast(ChatCompletionThinkingBlock, block).get("thinking")  # cast-ok: narrowed by the type tag
-    return str(thinking or "")
+    return str(block.get("thinking") or "")
 
 
 def reasoning_content_from_thinking_blocks(
@@ -1843,24 +1840,125 @@ def reasoning_content_from_thinking_blocks(
     return "\n".join(text for block in thinking_blocks if (text := _readable_thinking_text(block)))
 
 
-def responses_reasoning_item_from_thinking_blocks(
-    thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock],
-) -> ChatCompletionReasoningItem | None:
-    """Build a Responses API `reasoning` input item from Anthropic thinking blocks.
+ENCRYPTED_REASONING_SIGNATURE_PREFIX: Final = "litellm_encrypted_reasoning:"
 
-    The item carries no `id`: the Responses API rejects an empty one and 404s on any id it
-    did not mint itself, while an item without an id is always accepted.
+
+def encrypted_reasoning_signature(encrypted_content: str) -> str:
+    """The opaque value a Responses API reasoning item's `encrypted_content` travels in.
+
+    Anthropic clients echo a thinking block's `signature` and a redacted block's `data`
+    back verbatim, so either field can carry the encrypted reasoning across turns; the
+    prefix tells the two apart from a signature Anthropic minted.
     """
+    return f"{ENCRYPTED_REASONING_SIGNATURE_PREFIX}{encrypted_content}"
+
+
+def _carries_encrypted_reasoning(signature: object) -> bool:
+    return isinstance(signature, str) and signature.startswith(ENCRYPTED_REASONING_SIGNATURE_PREFIX)
+
+
+def encrypted_content_from_signature(signature: object) -> str | None:
+    if not isinstance(signature, str) or not _carries_encrypted_reasoning(signature):
+        return None
+    return signature.removeprefix(ENCRYPTED_REASONING_SIGNATURE_PREFIX) or None
+
+
+def _encrypted_reasoning_field(block: Mapping[str, object]) -> object:
+    match block.get("type"):
+        case "thinking":
+            return block.get("signature")
+        case "redacted_thinking":
+            return block.get("data")
+        case _:
+            return None
+
+
+def encrypted_content_of_block(block: Mapping[str, object]) -> str | None:
+    return encrypted_content_from_signature(_encrypted_reasoning_field(block))
+
+
+def is_encrypted_reasoning_block(block: object) -> bool:
+    """A thinking or redacted_thinking block carrying Responses API encrypted reasoning.
+
+    Only the Responses API that minted the content can read it back, so an Anthropic
+    backend has to drop such a block rather than fail signature verification on it.
+    """
+    if not isinstance(block, Mapping):
+        return False
+    mapping: Final = cast(Mapping[str, object], block)  # cast-ok: narrowed by isinstance
+    return _carries_encrypted_reasoning(_encrypted_reasoning_field(mapping))
+
+
+def strip_encrypted_reasoning_from_messages(messages: object) -> None:
+    """Drop the bridge-tagged reasoning blocks a routed deployment cannot decrypt from
+    Anthropic-shaped history.
+
+    The whole block goes, the way #40280 drops undecryptable Responses ``input`` items: a
+    provider that did not mint the block rejects it signed (a foreign signature) and unsigned
+    (a missing signature) alike, so keeping its text as an unsigned thinking block only moves
+    the 400 from the router to the provider.
+
+    Mutates the content lists in place: the router's fallback snapshot shares these
+    message objects, so a rebound list would replay the stripped blocks on the fallback hop.
+    """
+    if not isinstance(messages, list):
+        return
+    for content in _anthropic_content_lists(cast(list[object], messages)):  # cast-ok: untyped client json
+        _strip_encrypted_reasoning_from_blocks(content)
+
+
+def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]:
+    return (
+        cast(list[object], content)  # cast-ok: narrowed by isinstance
+        for message in messages
+        if isinstance(message, Mapping)
+        for content in (cast(Mapping[str, object], message).get("content"),)  # cast-ok: narrowed by isinstance
+        if isinstance(content, list)
+    )
+
+
+def _strip_encrypted_reasoning_from_blocks(content: object) -> None:
+    blocks: Final = cast(list[object], content)  # cast-ok: narrowed by the caller's isinstance
+    kept: Final = tuple(block for block in blocks if not is_encrypted_reasoning_block(block))
+    blocks[:] = kept  # rebind-ok: shared with fallback snapshot
+
+
+def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str:
+    index, block = indexed_block
+    return f"encrypted:{index}" if is_encrypted_reasoning_block(block) else "summary"
+
+
+def _reasoning_item_from_block_group(group: tuple[Mapping[str, object], ...]) -> ChatCompletionReasoningItem | None:
     summary: Final[list[ChatCompletionReasoningSummaryTextBlock]] = [  # mutable-ok: API message payload
         ChatCompletionReasoningSummaryTextBlock(type="summary_text", text=text)
-        for block in thinking_blocks
+        for block in group
         if (text := _readable_thinking_text(block))
     ]
+    encrypted_content: Final = encrypted_content_of_block(group[0])
+    if encrypted_content is not None:
+        return ChatCompletionReasoningItem(type="reasoning", summary=summary, encrypted_content=encrypted_content)
     if not summary:
         return None
     return ChatCompletionReasoningItem(type="reasoning", summary=summary)
 
 
+def responses_reasoning_items_from_thinking_blocks(
+    thinking_blocks: Iterable[Mapping[str, object]],
+) -> tuple[ChatCompletionReasoningItem, ...]:
+    """Build Responses API `reasoning` input items from Anthropic thinking blocks.
+
+    A block carrying encrypted reasoning replays the item it came from byte for byte;
+    a run of plain thinking blocks collapses into one summary-only item. No item carries
+    an `id`: the Responses API 404s on any id it did not mint itself and rejects an empty
+    one, while an item without an id is always accepted.
+    """
+    return tuple(
+        item
+        for _, group in groupby(enumerate(thinking_blocks), key=_reasoning_replay_group_key)
+        if (item := _reasoning_item_from_block_group(tuple(block for _, block in group))) is not None
+    )
+
+
 def _parse_content_for_reasoning(
     message_text: str | None,
 ) -> tuple[str | None, str | None]:

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -46,6 +46,7 @@ from litellm.types.utils import GenericImageParsingChunk
 from .common_utils import (
     convert_content_list_to_str,
     infer_content_type_from_url_and_content,
+    is_encrypted_reasoning_block,
     is_non_content_values_set,
     parse_tool_call_arguments,
 )
@@ -2299,13 +2300,16 @@ def sanitize_messages_for_tool_calling(
 
 
 def _is_unsignable_thinking_block(block: object) -> bool:
-    """A `thinking` block that Anthropic cannot accept on input.
+    """A thinking block that Anthropic cannot accept on input.
 
     Anthropic verifies the thinking signature cryptographically, so a block whose
     signature is null, empty, or missing (e.g. from an open-source reasoning model)
-    is rejected with a 400 and must be dropped rather than blanked or repaired.
-    `redacted_thinking` blocks carry no signature and are always kept.
+    is rejected with a 400 and must be dropped rather than blanked or repaired, and
+    so is a block whose signature or data carries another provider's encrypted
+    reasoning. A `redacted_thinking` block Anthropic minted is always kept.
     """
+    if is_encrypted_reasoning_block(block):
+        return True
     if not isinstance(block, dict) or block.get("type") != "thinking":
         return False
     signature: Final = block.get("signature")

litellm/litellm_core_utils/token_counter.py

@@ -3,11 +3,15 @@
 import base64
 import io
 import struct
-from collections.abc import Callable, Iterable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
 from typing import Final, Literal, cast
 
+import anyio
+import anyio.lowlevel
 import httpx
 import tiktoken
+from tokenizers import Tokenizer
+from typing_extensions import ParamSpec, TypeVar
 
 import litellm
 from litellm import verbose_logger
@@ -21,7 +25,10 @@ from litellm.constants import (
     MAX_TILE_HEIGHT,
     MAX_TILE_WIDTH,
     TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS,
+    TOKEN_COUNTER_MAX_CONCURRENT_COUNTS,
+    TOKEN_COUNTER_MAX_EXACT_CHARS,
 )
+from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.default_encoding import encoding as default_encoding
 from litellm.litellm_core_utils.url_utils import safe_get
 from litellm.llms.custom_httpx.http_handler import _get_httpx_client
@@ -317,6 +324,32 @@ TokenCounterFunction = Callable[[str], int]
 Type for a function that counts tokens in a string.
 """
 
+EXTRAPOLATION_SAMPLES: Final = 16
+T_ParamSpec: Final = ParamSpec("T_ParamSpec")
+T_Retval = TypeVar("T_Retval")
+_COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter")
+
+
+def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter:
+    existing: Final = _COUNT_OFFLOAD_LIMITER.get(None)
+    if existing is not None:
+        return existing
+    created: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS)
+    _COUNT_OFFLOAD_LIMITER.set(created)
+    return created
+
+
+def offload_token_count(
+    function: Callable[T_ParamSpec, T_Retval],
+) -> Callable[T_ParamSpec, Awaitable[T_Retval]]:
+    async def offloaded(
+        *args: T_ParamSpec.args,
+        **kwargs: T_ParamSpec.kwargs,  # kwargs-ok: ParamSpec keeps the wrapped function's own keyword contract
+    ) -> T_Retval:
+        return await asyncify(function, limiter=_count_offload_limiter_for_this_loop())(*args, **kwargs)
+
+    return offloaded
+
 
 def _get_tiktoken_count_function(
     encode_length: Callable[[str], int],
@@ -538,9 +571,40 @@ def _count_extra(
     return num_tokens
 
 
+def _get_extrapolating_count_function(
+    count_exactly: TokenCounterFunction,
+    max_exact_chars: int = TOKEN_COUNTER_MAX_EXACT_CHARS,
+) -> TokenCounterFunction:
+    def count_tokens(text: str) -> int:
+        if len(text) <= max_exact_chars:
+            return count_exactly(text)
+        samples: Final = _evenly_spaced_samples(text, max_exact_chars)
+        sampled_chars: Final = sum(len(sample) for sample in samples)
+        return round(sum(count_exactly(sample) for sample in samples) * len(text) / sampled_chars)
+
+    return count_tokens
+
+
+def _evenly_spaced_samples(text: str, total_chars: int) -> tuple[str, ...]:
+    sample_count: Final = min(EXTRAPOLATION_SAMPLES, total_chars)
+    sample_chars: Final = total_chars // sample_count
+    last_start: Final = len(text) - sample_chars
+    return tuple(
+        text[start : start + sample_chars]
+        for start in (last_start * index // max(sample_count - 1, 1) for index in range(sample_count))
+    )
+
+
 def _get_count_function(
     model: str | None,
     custom_tokenizer: dict | SelectTokenizerResponse | None = None,
+) -> TokenCounterFunction:
+    return _get_extrapolating_count_function(_get_exact_count_function(model, custom_tokenizer))
+
+
+def _get_exact_count_function(
+    model: str | None,
+    custom_tokenizer: dict | SelectTokenizerResponse | None = None,
 ) -> TokenCounterFunction:
     """
     Get the function to count tokens based on the model and custom tokenizer."""
@@ -549,10 +613,10 @@ def _get_count_function(
     if model is not None or custom_tokenizer is not None:
         tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model)
         if tokenizer_json["type"] == "huggingface_tokenizer":
+            tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"]
 
             def count_tokens(text: str) -> int:
-                enc: Final = tokenizer_json["tokenizer"].encode(text)
-                return len(enc.ids)
+                return len(tokenizer.encode_batch_fast([text])[0])
 
             return count_tokens
         elif tokenizer_json["type"] == "openai_tokenizer":

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -213,11 +213,17 @@ class AnthropicMessagesHandler(BaseTranslation):
     """
 
     delivers_ended_stream_rewrites = True
+    assembles_streamed_response = True
 
     def __init__(self):
         super().__init__()
         self.adapter = LiteLLMAnthropicMessagesAdapter()
 
+    def post_call_hook_response(self, response: object) -> object:
+        if not isinstance(response, ModelResponse):
+            return response
+        return self.adapter.translate_openai_response_to_anthropic(response)
+
     @staticmethod
     def _build_streaming_usage_response(
         responses_so_far: Sequence[object],

litellm/llms/anthropic/common_utils.py

@@ -21,6 +21,7 @@ from litellm.constants import (
 )
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_file_ids_from_messages,
+    is_encrypted_reasoning_block,
 )
 from litellm.litellm_core_utils.prompt_templates.factory import (
     THOUGHT_SIGNATURE_SEPARATOR,
@@ -1201,6 +1202,32 @@ def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[A
     return out
 
 
+def _without_encrypted_reasoning_blocks(message: dict) -> dict | None:  # mutable-ok: Anthropic message payload shape
+    if not isinstance(message, Mapping):
+        return message
+    content: Final = message.get("content")
+    if not isinstance(content, list):
+        return message
+    kept: Final = [b for b in content if not is_encrypted_reasoning_block(b)]  # mutable-ok: API message payload
+    if len(kept) == len(content):
+        return message
+    if not kept:
+        return None
+    return {**message, "content": kept}  # mutable-ok: API message payload
+
+
+def strip_encrypted_reasoning_blocks_from_anthropic_messages(
+    messages: Sequence[dict],  # mutable-ok: Anthropic message payload shape
+) -> list[dict]:  # mutable-ok: AnthropicMessagesRequest.messages is typed list[dict]
+    """
+    Drop thinking / redacted_thinking blocks that carry another provider's encrypted
+    reasoning (a turn the Responses API bridge served) before the request reaches
+    Anthropic, which cannot verify them. Anthropic's own signed blocks are kept.
+    """
+    stripped: Final = (_without_encrypted_reasoning_blocks(m) for m in messages)
+    return [m for m in stripped if m is not None]  # mutable-ok: API message payload
+
+
 def strip_thinking_blocks_from_anthropic_messages_request_dict(
     data: dict[str, Any],
 ) -> None:

litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py

@@ -113,6 +113,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import (
 from litellm.llms.anthropic.common_utils import (
     is_empty_unsigned_thinking_block,
     normalize_anthropic_tool_use_id,
+    strip_encrypted_reasoning_blocks_from_anthropic_messages,
 )
 from litellm.llms.anthropic.experimental_pass_through.context_management import (
     PolyfillResult,
@@ -417,7 +418,8 @@ class LiteLLMAnthropicMessagesAdapter:
         model: str | None = None,
     ) -> list:
         new_messages: Final[list[AllMessageValues]] = []
-        for m in messages:
+        replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)
+        for m in replayable_messages:
             user_message: ChatCompletionUserMessage | None = None
             tool_message_list: list[ChatCompletionToolMessage] = []
             new_user_content_list: list[ChatCompletionTextObject | ChatCompletionImageObject] = []

litellm/llms/anthropic/experimental_pass_through/messages/transformation.py

@@ -25,6 +25,7 @@ from ...common_utils import (
     AnthropicModelInfo,
     optionally_handle_anthropic_oauth,
     strip_advisor_blocks_from_messages,
+    strip_encrypted_reasoning_blocks_from_anthropic_messages,
 )
 
 DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01"
@@ -613,7 +614,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
             messages = strip_advisor_blocks_from_messages(messages)
 
         anthropic_messages_request: Final[AnthropicMessagesRequest] = AnthropicMessagesRequest(
-            messages=messages,
+            messages=strip_encrypted_reasoning_blocks_from_anthropic_messages(messages),
             max_tokens=max_tokens,
             model=model,
             **anthropic_messages_optional_request_params,

litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py

@@ -19,6 +19,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,
 )
 from litellm.types.llms.openai import ResponsesAPIResponse
+from litellm.utils import ProviderConfigManager
 
 from ..utils import litellm_logging_obj_from_kwargs, local_model_name
 from .streaming_iterator import AnthropicResponsesStreamWrapper
@@ -34,6 +35,15 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str,
     return extra_kwargs or {}
 
 
+def _provider_returns_encrypted_reasoning(model: str, custom_llm_provider: object) -> bool:
+    provider: Final = (
+        custom_llm_provider if isinstance(custom_llm_provider, str) else litellm.get_llm_provider(model=model)[1]
+    )
+    provider_model: Final = local_model_name(model, provider)
+    responses_config: Final = ProviderConfigManager.get_provider_responses_api_config(provider, provider_model)
+    return responses_config is not None and "include" in responses_config.get_supported_openai_params(provider_model)
+
+
 def _build_responses_kwargs(
     *,
     max_tokens: int,
@@ -85,8 +95,13 @@ def _build_responses_kwargs(
         request_data["output_format"] = output_format
 
     anthropic_request: Final = AnthropicMessagesRequest(**request_data)
-    responses_kwargs: Final = _ADAPTER.translate_request(anthropic_request)
     forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs)
+    responses_kwargs: Final = _ADAPTER.translate_request(
+        anthropic_request,
+        include_encrypted_reasoning=_provider_returns_encrypted_reasoning(
+            model, forwarded_kwargs.get("custom_llm_provider")
+        ),
+    )
 
     # Normalize reasoning effort based on model capabilities
     # (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported)
@@ -111,7 +126,7 @@ def _build_responses_kwargs(
         responses_kwargs["stream"] = True
 
     # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
-    excluded: Final = {"anthropic_messages"}
+    excluded: Final = frozenset(("anthropic_messages",))
     for key, value in forwarded_kwargs.items():
         if key == "litellm_logging_obj" and value is not None:
             from litellm.litellm_core_utils.litellm_logging import (
@@ -132,6 +147,14 @@ def _build_responses_kwargs(
     if explicit_prompt_cache_key is not None:
         responses_kwargs["prompt_cache_key"] = explicit_prompt_cache_key
 
+    deployment_include: Final = forwarded_kwargs.get("include")
+    bridge_include: Final = responses_kwargs.get("include")
+    if isinstance(deployment_include, list) and isinstance(bridge_include, list):
+        responses_kwargs["include"] = [
+            *bridge_include,
+            *(item for item in deployment_include if item not in bridge_include),
+        ]
+
     return responses_kwargs
 
 

litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py

@@ -9,13 +9,19 @@ from typing import TYPE_CHECKING, Any, Final
 
 from litellm import verbose_logger
 from litellm._uuid import uuid
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    encrypted_reasoning_signature,
+)
 from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
     refusal_stop_details,
     responses_output_refusal_text,
 )
 from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
 
-from .transformation import LiteLLMAnthropicToResponsesAPIAdapter
+from .transformation import (
+    REASONING_SUMMARY_PART_SEPARATOR,
+    LiteLLMAnthropicToResponsesAPIAdapter,
+)
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
@@ -29,9 +35,10 @@ class AnthropicResponsesStreamWrapper:
       response.created                   -> message_start
       response.output_item.added         -> content_block_start (if message/function_call)
       response.output_text.delta         -> content_block_delta (text_delta)
+      response.reasoning_summary_part.added -> content_block_delta (thinking_delta separator)
       response.reasoning_summary_text.delta -> content_block_delta (thinking_delta)
       response.function_call_arguments.delta -> content_block_delta (input_json_delta)
-      response.output_item.done          -> content_block_stop
+      response.output_item.done          -> content_block_delta (signature_delta) + content_block_stop
       response.completed                 -> message_delta + message_stop
     """
 
@@ -94,6 +101,38 @@ class AnthropicResponsesStreamWrapper:
         )
         return block_idx
 
+    @staticmethod
+    def _field(source: object, name: str) -> object:
+        return source.get(name) if isinstance(source, dict) else getattr(source, name, None)
+
+    def _close_reasoning_item(self, item: object, item_id: str | None) -> None:
+        block_idx: Final = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index
+        encrypted_content: Final = self._field(item, "encrypted_content")
+        signature: Final = (
+            encrypted_reasoning_signature(encrypted_content)
+            if isinstance(encrypted_content, str) and encrypted_content
+            else None
+        )
+        if block_idx < 0 and signature is None:
+            return
+        if block_idx < 0:
+            redacted_idx: Final = self._open_block(
+                item_id,
+                {"type": "redacted_thinking", "data": signature},  # mutable-ok: API message payload
+            )
+            stop: Final = {"type": "content_block_stop", "index": redacted_idx}  # mutable-ok: API message payload
+            self._chunk_queue.append(stop)
+            return
+        if signature is not None:
+            self._chunk_queue.append(
+                {  # mutable-ok: API message payload
+                    "type": "content_block_delta",
+                    "index": block_idx,
+                    "delta": {"type": "signature_delta", "signature": signature},  # mutable-ok: API message payload
+                }
+            )
+        self._chunk_queue.append({"type": "content_block_stop", "index": block_idx})  # mutable-ok: API message payload
+
     def _process_event(self, event: object) -> None:
         """Convert one Responses API event into zero or more Anthropic chunks queued for emission."""
         event_type = getattr(event, "type", None)
@@ -175,6 +214,26 @@ class AnthropicResponsesStreamWrapper:
             )
             return
 
+        if event_type == "response.reasoning_summary_part.added":
+            part_item_id: Final = self._field(event, "item_id")
+            summary_index: Final = self._field(event, "summary_index")
+            part_block_idx: Final = (
+                self._item_id_to_block_index.get(part_item_id, -1) if isinstance(part_item_id, str) else -1
+            )
+            if part_block_idx < 0 or not isinstance(summary_index, int) or summary_index == 0:
+                return
+            self._chunk_queue.append(
+                {  # mutable-ok: API message payload
+                    "type": "content_block_delta",
+                    "index": part_block_idx,
+                    "delta": {  # mutable-ok: API message payload
+                        "type": "thinking_delta",
+                        "thinking": REASONING_SUMMARY_PART_SEPARATOR,
+                    },
+                }
+            )
+            return
+
         # ---- reasoning summary text delta ----
         if event_type == "response.reasoning_summary_text.delta":
             item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None)
@@ -220,6 +279,9 @@ class AnthropicResponsesStreamWrapper:
             item_id = (
                 getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None
             )
+            if self._field(item, "type") == "reasoning":
+                self._close_reasoning_item(item, item_id)
+                return
             block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index
             if block_idx < 0:
                 return

litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py

@@ -13,7 +13,8 @@ from typing import Any, Final, cast
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     TOOL_RESULT_IMAGE_BOUNDARY,
     TOOL_RESULT_IMAGE_PLACEHOLDER,
-    responses_reasoning_item_from_thinking_blocks,
+    encrypted_reasoning_signature,
+    responses_reasoning_items_from_thinking_blocks,
     with_prompt_cache_breakpoint,
 )
 from litellm.litellm_core_utils.reasoning_effort_utils import (
@@ -33,6 +34,7 @@ from litellm.types.llms.anthropic import (
     AnthropicFinishReason,
     AnthropicMessagesRequest,
     AnthropicMessagesToolChoice,
+    AnthropicResponseContentBlockRedactedThinking,
     AnthropicResponseContentBlockText,
     AnthropicResponseContentBlockThinking,
     AnthropicResponseContentBlockToolUse,
@@ -43,11 +45,13 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicUsage,
 )
 from litellm.types.llms.openai import (
-    ChatCompletionThinkingBlock,
     ResponseAPIUsage,
     ResponsesAPIResponse,
 )
 
+REASONING_SUMMARY_PART_SEPARATOR: Final = "\n\n"
+RESPONSES_INCLUDE_ENCRYPTED_REASONING: Final = "reasoning.encrypted_content"
+
 
 class LiteLLMAnthropicToResponsesAPIAdapter:
     """
@@ -163,49 +167,55 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
         return str(getattr(part, "text", None) or "")
 
     @classmethod
-    def _thinking_blocks_from_reasoning_item(
+    def _thinking_block_from_reasoning_item(
         cls,
         summary: Iterable[object],
-    ) -> tuple[dict[str, Any], ...]:  # mutable-ok: API message payload
-        """Anthropic thinking blocks for one Responses reasoning item.
+        encrypted_content: object,
+    ) -> dict[str, Any] | None:  # mutable-ok: API message payload
+        """The one Anthropic block for a Responses reasoning item.
 
-        The signature stays empty: only Anthropic can sign a thinking block, and a stand-in
-        value would be replayed as a real one and rejected by every backend that verifies it.
+        The item's encrypted reasoning rides the block's opaque field (`signature`, or
+        `data` when there is no summary text) so the client echoes it back and the next
+        turn replays the very item OpenAI produced; without it the signature stays empty,
+        since only Anthropic can sign a thinking block.
         """
-        return tuple(
-            AnthropicResponseContentBlockThinking(
-                type="thinking",
-                thinking=text,
-                signature=None,
-            ).model_dump()
-            for part in summary
-            if (text := cls._summary_part_text(part))
+        text: Final = REASONING_SUMMARY_PART_SEPARATOR.join(
+            part_text for part in summary if (part_text := cls._summary_part_text(part))
         )
+        if not isinstance(encrypted_content, str) or not encrypted_content:
+            if not text:
+                return None
+            return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=None).model_dump()
+        signature: Final = encrypted_reasoning_signature(encrypted_content)
+        if not text:
+            return AnthropicResponseContentBlockRedactedThinking(type="redacted_thinking", data=signature).model_dump()
+        return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=signature).model_dump()
 
     @staticmethod
     def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str:
         """Group a run of consecutive thinking blocks together; keep every other block alone."""
         index, block = indexed_block
-        return "thinking" if block.get("type") == "thinking" else f"block:{index}"
+        return "thinking" if block.get("type") in ("thinking", "redacted_thinking") else f"block:{index}"
 
     @classmethod
-    def _assistant_group_to_input_item(
+    def _assistant_group_to_input_items(
         cls, group: tuple[Mapping[str, object], ...]
-    ) -> dict[str, Any] | None:  # mutable-ok: API message payload
+    ) -> tuple[dict[str, Any], ...]:  # mutable-ok: API message payload
         first: Final = group[0]
         btype: Final = first.get("type")
-        if btype == "thinking":
-            blocks: Final = cast(tuple[ChatCompletionThinkingBlock, ...], group)  # cast-ok: untrusted client payload
-            reasoning_item: Final = responses_reasoning_item_from_thinking_blocks(blocks)
-            return None if reasoning_item is None else dict(reasoning_item)  # mutable-ok: API message payload
+        if btype in ("thinking", "redacted_thinking"):
+            replayed: Final = responses_reasoning_items_from_thinking_blocks(group)
+            return tuple(dict(item) for item in replayed)  # mutable-ok: API message payload
         if btype == "tool_use":
-            return {  # mutable-ok: API message payload
-                "type": "function_call",
-                "call_id": first.get("id", ""),
-                "name": first.get("name", ""),
-                "arguments": json.dumps(first.get("input", {})),  # mutable-ok: API message payload
-            }
-        return None
+            return (
+                {  # mutable-ok: API message payload
+                    "type": "function_call",
+                    "call_id": first.get("id", ""),
+                    "name": first.get("name", ""),
+                    "arguments": json.dumps(first.get("input", {})),  # mutable-ok: API message payload
+                },
+            )
+        return ()
 
     def translate_messages_to_responses_input(
         self,
@@ -362,7 +372,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
                     input_items.extend(
                         item
                         for _, group in groupby(enumerate(blocks), key=self._assistant_block_group_key)
-                        if (item := self._assistant_group_to_input_item(tuple(block for _, block in group))) is not None
+                        for item in self._assistant_group_to_input_items(tuple(block for _, block in group))
                     )
                     asst_parts: list[dict[str, Any]] = [  # mutable-ok: API message payload
                         {"type": "output_text", "text": block.get("text", "")}  # mutable-ok: API message payload
@@ -495,10 +505,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
     def translate_request(
         self,
         anthropic_request: AnthropicMessagesRequest,
+        include_encrypted_reasoning: bool = True,
     ) -> dict[str, Any]:
         """
         Translate a full Anthropic /v1/messages request dict to
         litellm.responses() / litellm.aresponses() kwargs.
+
+        ``include_encrypted_reasoning`` asks the provider for ``reasoning.encrypted_content``
+        on every call, so a reasoning model's items can be replayed intact next turn even
+        when the client sent no ``thinking`` block; pass False for a provider whose
+        Responses API rejects ``include``.
         """
         model: Final[str] = anthropic_request["model"]
         messages_list: Final = cast(
@@ -528,6 +544,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
             "model": model,
             "input": input_items,
         }
+        if include_encrypted_reasoning:
+            responses_kwargs["include"] = [RESPONSES_INCLUDE_ENCRYPTED_REASONING]  # mutable-ok: API request payload
 
         if system and not developer_parts:
             if isinstance(system, str):
@@ -634,7 +652,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
 
         for item in response.output:
             if isinstance(item, ResponseReasoningItem):
-                content.extend(self._thinking_blocks_from_reasoning_item(item.summary))
+                reasoning_block = self._thinking_block_from_reasoning_item(item.summary, item.encrypted_content)
+                if reasoning_block is not None:
+                    content.append(reasoning_block)
 
             elif isinstance(item, ResponseOutputMessage):
                 for part in item.content:
@@ -684,11 +704,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
                                     ).model_dump()
                                 )
                 elif item_type == "reasoning":
-                    content.extend(
-                        self._thinking_blocks_from_reasoning_item(
-                            cast(Iterable[object], item.get("summary") or ()),  # cast-ok: untyped provider json
-                        )
+                    reasoning_block = self._thinking_block_from_reasoning_item(
+                        cast(Iterable[object], item.get("summary") or ()),  # cast-ok: untyped provider json
+                        item.get("encrypted_content"),
                     )
+                    if reasoning_block is not None:
+                        content.append(reasoning_block)
                 elif item_type == "function_call":
                     try:
                         input_data = json.loads(item.get("arguments", "{}"))

litellm/llms/azure/common_utils.py

@@ -14,6 +14,7 @@ from typing_extensions import ReadOnly, TypedDict
 import litellm
 from litellm._logging import verbose_logger
 from litellm.caching.caching import DualCache
+from litellm.constants import DEFAULT_MAX_RETRIES
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.openai.common_utils import BaseOpenAILLM
 from litellm.secret_managers.get_azure_ad_token_provider import (
@@ -582,7 +583,8 @@ class BaseAzureLLM(BaseOpenAILLM):
         if scope is None:
             scope = "https://cognitiveservices.azure.com/.default"
 
-        max_retries: Final = litellm_params.get("max_retries")
+        configured_max_retries: Final = litellm_params.get("max_retries")
+        max_retries: Final = DEFAULT_MAX_RETRIES if configured_max_retries is None else configured_max_retries
         timeout: Final = litellm_params.get("timeout")
         if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret:
             verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth")
@@ -642,8 +644,7 @@ class BaseAzureLLM(BaseOpenAILLM):
         else:
             azure_client_params["http_client"] = self._get_sync_http_client()
 
-        if max_retries is not None:
-            azure_client_params["max_retries"] = max_retries
+        azure_client_params["max_retries"] = max_retries
         if timeout is not None:
             azure_client_params["timeout"] = timeout
 

litellm/llms/azure_ai/image_edit/__init__.py

@@ -23,7 +23,7 @@ def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig:
     """
     Get the appropriate image edit config for an Azure AI model.
 
-    - MAI models use /mai/v1/images/edits with multipart form data and size
+    - MAI models use /mai/v1/images/edits with multipart form data
     - FLUX 2 models use JSON with base64 image
     - FLUX 1 models use multipart/form-data
     """

litellm/llms/azure_ai/image_edit/mai_transformation.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Final, cast
+from typing import TYPE_CHECKING, Any, Final
 
 import httpx
 from httpx._types import RequestFiles
@@ -13,7 +13,6 @@ from litellm.llms.azure_ai.image_generation.mai_transformation import (
 from litellm.llms.openai.common_utils import OpenAIError
 from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
 from litellm.secret_managers.main import get_secret_str
-from litellm.types.images.main import ImageEditOptionalRequestParams
 from litellm.types.llms.openai import FileTypes
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import ImageResponse
@@ -26,65 +25,8 @@ if TYPE_CHECKING:
 class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig):
     """Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5)."""
 
-    DEFAULT_SIZE = "1024x1024"
-
     def get_supported_openai_params(self, model: str) -> list:
-        return ["prompt", "image", "model", "n", "size"]
-
-    def map_openai_params(
-        self,
-        image_edit_optional_params: ImageEditOptionalRequestParams,
-        model: str,
-        drop_params: bool,
-    ) -> dict:
-        optional_params: Final[dict[str, Any]] = {}
-        supported_params: Final = self.get_supported_openai_params(model)
-
-        for key, value in dict(image_edit_optional_params).items():
-            if value is None or key in optional_params:
-                continue
-
-            if key in supported_params:
-                if key == "size" and value:
-                    size_param = cast(str, value)
-                    self._validate_size_param(size_param)
-                    optional_params[key] = size_param
-                else:
-                    optional_params[key] = value
-            elif not drop_params:
-                raise ValueError(
-                    f"Parameter {key} is not supported for model {model}. "
-                    f"Supported parameters are {supported_params}. "
-                    f"Set drop_params=True to drop unsupported parameters."
-                )
-
-        if "size" not in optional_params:
-            optional_params["size"] = self.DEFAULT_SIZE
-
-        return optional_params
-
-    def _validate_size_param(self, size: str) -> None:
-        known_sizes: Final = {
-            "1024x1024",
-            "1792x1024",
-            "1024x1792",
-            "512x512",
-            "256x256",
-        }
-
-        if size in known_sizes:
-            return
-
-        if "x" in size:
-            try:
-                tuple(map(int, size.lower().split("x", 1)))
-                return
-            except ValueError:
-                raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').")
-
-        raise ValueError(
-            f"Unsupported size value: '{size}'. Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
-        )
+        return ["prompt", "image", "model", "n"]
 
     def validate_environment(
         self,

litellm/llms/azure_ai/image_generation/mai_transformation.py

@@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Final
 
 import httpx
 
+from litellm.exceptions import UnsupportedParamsError
 from litellm.llms.base_llm.image_generation.transformation import (
     BaseImageGenerationConfig,
 )
@@ -21,6 +22,10 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
     DEFAULT_WIDTH = 1024
     DEFAULT_HEIGHT = 1024
 
+    MAX_IMAGES_PER_REQUEST: Final = 1
+    MIN_DIMENSION_PX: Final = 768
+    MAX_TOTAL_PX: Final = 1_056_768
+
     @staticmethod
     def get_mai_image_generation_url(
         api_base: str | None,
@@ -145,16 +150,27 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
 
             if k in supported_params:
                 if k == "size" and v:
-                    self._map_size_param(v, optional_params)
+                    self._map_size_param(v, optional_params, model)
+                elif k == "n" and v is not None and self._image_count(v, model) != self.MAX_IMAGES_PER_REQUEST:
+                    if not drop_params:
+                        raise self._unsupported(
+                            model,
+                            f"n={v} is not supported for model {model}. The Azure AI MAI image "
+                            f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per "
+                            "request and ignores any count, so a larger value would silently "
+                            "return fewer images than requested. Send one request per image, or "
+                            "set drop_params=True to drop n.",
+                        )
                 else:
                     optional_params[k] = v
             elif k in ("width", "height"):
                 optional_params[k] = v
             elif not drop_params:
-                raise ValueError(
+                raise self._unsupported(
+                    model,
                     f"Parameter {k} is not supported for model {model}. "
                     f"Supported parameters are {supported_params} and width/height. "
-                    f"Set drop_params=True to drop unsupported parameters."
+                    f"Set drop_params=True to drop unsupported parameters.",
                 )
 
         if "width" not in optional_params:
@@ -165,7 +181,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
         optional_params.pop("size", None)
         return optional_params
 
-    def _map_size_param(self, size: str, optional_params: dict) -> None:
+    @staticmethod
+    def _unsupported(model: str, message: str) -> UnsupportedParamsError:
+        return UnsupportedParamsError(message=message, llm_provider="azure_ai", model=model)
+
+    def _image_count(self, n: object, model: str) -> int:
+        if isinstance(n, int):
+            return n
+        try:
+            return int(str(n))
+        except ValueError:
+            raise self._unsupported(model, f"n={n!r} is not a whole number of images for model {model}.")
+
+    def _map_size_param(self, size: str, optional_params: dict, model: str) -> None:
         size_mapping: Final = {
             "1024x1024": (1024, 1024),
             "1792x1024": (1792, 1024),
@@ -176,19 +204,36 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
 
         if size in size_mapping:
             width, height = size_mapping[size]
-            optional_params["width"] = width
-            optional_params["height"] = height
         elif "x" in size:
             try:
                 width, height = map(int, size.lower().split("x"))
-                optional_params["width"] = width
-                optional_params["height"] = height
             except ValueError:
-                raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').")
+                raise self._unsupported(
+                    model, f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')."
+                )
         else:
-            raise ValueError(
+            raise self._unsupported(
+                model,
                 f"Unsupported size value: '{size}'. "
-                f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
+                f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string.",
+            )
+
+        self._validate_dimensions(model=model, size=size, width=width, height=height)
+        optional_params["width"] = width
+        optional_params["height"] = height
+
+    def _validate_dimensions(self, model: str, size: str, width: int, height: int) -> None:
+        if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX:
+            raise self._unsupported(
+                model,
+                f"Unsupported size value: '{size}'. Azure AI MAI image models require width and "
+                f"height of at least {self.MIN_DIMENSION_PX} pixels.",
+            )
+        if width * height > self.MAX_TOTAL_PX:
+            raise self._unsupported(
+                model,
+                f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most "
+                f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height}).",
             )
 
     def transform_image_generation_response(

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -61,6 +61,20 @@ class BaseTranslation(ABC):
     on every other translation are undeliverable: the pipeline executor
     discards them and releases the original chunks."""
 
+    assembles_streamed_response: ClassVar[bool] = False
+    """Whether ``process_output_streaming_response`` stores the assembled response of an
+    ended stream under ``request_data["response"]`` before scanning it, the way the chat,
+    Responses, and Messages translations do. A streaming pipeline runs a guardrail that only
+    has the legacy post-call hook against that response, so on a translation without it such
+    a guardrail keeps running on its own."""
+
+    def post_call_hook_response(self, response: object) -> object:
+        """The ``response`` this endpoint's non-streaming post-call hooks receive, derived from
+        the object the translation stores under ``request_data["response"]`` while scanning an
+        ended stream. Chat and Responses scan that shape already; a translation that scans a
+        different one (Messages scans an OpenAI-shaped ModelResponse) overrides this."""
+        return response
+
     @staticmethod
     def transform_user_api_key_dict_to_metadata(
         user_api_key_dict: Any | None,

litellm/llms/bedrock/base_aws_llm.py

@@ -1,13 +1,17 @@
+import asyncio
 import base64
+import contextvars
 import hashlib
 import json
 import os
 import re
 import urllib.parse
 from collections.abc import Callable, Mapping
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
+from functools import partial
 from threading import Lock
-from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload
+from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload
 
 import httpx
 from pydantic import BaseModel, ValidationError
@@ -16,6 +20,7 @@ from litellm._logging import verbose_logger
 from litellm.caching.caching import DualCache
 from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.constants import (
+    AWS_SIGNING_MAX_THREADS,
     BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
     BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES,
     BEDROCK_IAM_CACHE_MAX_ENTRIES,
@@ -80,7 +85,11 @@ class AwsAuthError(Exception):
         super().__init__(self.message)  # Call the base class constructor with the parameters it needs
 
 
-class BaseAWSLLM:
+class SignsRequestsWithAWS:
+    pass
+
+
+class BaseAWSLLM(SignsRequestsWithAWS):
     # Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request).
     # Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static
     # access-key + secret + region use ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); ambient
@@ -1668,3 +1677,52 @@ class BaseAWSLLM:
             request_headers_dict["Authorization"] = incoming_authorization
 
         return request_headers_dict, request.body
+
+
+def sign_aws_json_post(
+    get_credentials: Callable[[], Credentials],
+    service_name: str,
+    aws_region_name: str | None,
+    url: str,
+    body: str,
+    headers: Mapping[str, str],
+) -> AWSPreparedRequest:
+    try:
+        from botocore.auth import SigV4Auth
+        from botocore.awsrequest import AWSRequest
+    except ImportError:
+        raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.")
+
+    aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers)
+    SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request)
+    return aws_request.prepare()
+
+
+_SignParams = ParamSpec("_SignParams")
+_SignedRequest = TypeVar("_SignedRequest")
+
+AWS_SIGNING_EXECUTOR: Final = ThreadPoolExecutor(max_workers=AWS_SIGNING_MAX_THREADS, thread_name_prefix="aws-signing")
+
+
+async def run_aws_signing(
+    sign: Callable[_SignParams, _SignedRequest],
+    /,
+    *args: _SignParams.args,
+    **kwargs: _SignParams.kwargs,  # kwargs-ok: ParamSpec forwarding keeps the wrapped signing signature
+) -> _SignedRequest:
+    context: Final = contextvars.copy_context()
+    return await asyncio.get_running_loop().run_in_executor(
+        AWS_SIGNING_EXECUTOR, partial(context.run, sign, *args, **kwargs)
+    )
+
+
+async def sign_request_off_loop_if_aws(
+    provider_config: object,
+    sign_request: Callable[_SignParams, _SignedRequest],
+    /,
+    *args: _SignParams.args,
+    **kwargs: _SignParams.kwargs,  # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature
+) -> _SignedRequest:
+    if isinstance(provider_config, SignsRequestsWithAWS):
+        return await run_aws_signing(sign_request, *args, **kwargs)
+    return sign_request(*args, **kwargs)

litellm/llms/bedrock/chat/converse_handler.py

@@ -21,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
 from litellm.types.utils import ModelResponse
 from litellm.utils import CustomStreamWrapper
 
-from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
+from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing
 from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text
 from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
 
@@ -136,7 +136,8 @@ class BedrockConverseLLM(BaseAWSLLM):
         )
         data: Final = json.dumps(request_data)
 
-        prepped: Final = self.get_request_headers(
+        prepped: Final = await run_aws_signing(
+            self.get_request_headers,
             credentials=credentials,
             aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
             extra_headers=headers,
@@ -206,7 +207,8 @@ class BedrockConverseLLM(BaseAWSLLM):
         )
         data: Final = json.dumps(request_data)
 
-        prepped: Final = self.get_request_headers(
+        prepped: Final = await run_aws_signing(
+            self.get_request_headers,
             credentials=credentials,
             aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
             extra_headers=headers,

litellm/llms/bedrock/chat/converse_transformation.py

@@ -4,6 +4,7 @@ Translating between OpenAI's `/chat/completion` format and Amazon's `/converse`
 
 import copy
 import json
+import re
 import time
 import types
 from collections.abc import Mapping
@@ -293,6 +294,10 @@ class AmazonConverseConfig(BaseConfig):
                     llm_provider="bedrock",
                 )
 
+    @staticmethod
+    def _is_openai_gpt_reasoning_model(model: str) -> bool:
+        return re.search(r"openai\.gpt-\d", model) is not None
+
     def _is_nova_2_model(self, model: str) -> bool:
         """
         Check if the model is a Nova 2 model that supports reasoningConfig.
@@ -422,15 +427,15 @@ class AmazonConverseConfig(BaseConfig):
         """
         Handle the reasoning_effort parameter based on the model type.
 
-        - GPT-OSS models: passed through unchanged via additionalModelRequestFields.
-        - OpenAI GPT-5.x models: mapped to ``reasoning.effort`` via additionalModelRequestFields.
+        - GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields.
+        - OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields.
         - Nova 2 models: transformed to reasoningConfig.
         - Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on
           adaptive Claude 4.6 / 4.7).
         """
-        if "gpt-oss" in model:
+        if "gpt-oss" in model or "deepseek" in model:
             optional_params["reasoning_effort"] = reasoning_effort
-        elif "openai.gpt-5" in model:
+        elif self._is_openai_gpt_reasoning_model(model):
             reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort}
             optional_params["reasoning"] = reasoning
         elif self._is_nova_2_model(model):
@@ -509,6 +514,36 @@ class AmazonConverseConfig(BaseConfig):
                 )
                 thinking["budget_tokens"] = BEDROCK_MIN_THINKING_BUDGET_TOKENS
 
+    def _is_deepseek_model(self, model: str, base_model: str) -> bool:
+        return "deepseek" in model or "deepseek" in base_model
+
+    def _is_deepseek_r1_model(self, model: str, base_model: str) -> bool:
+        return "deepseek.r1" in model or "deepseek.r1" in base_model
+
+    def _model_accepts_anthropic_thinking_param(self, model: str, base_model: str) -> bool:
+        """Whether the model accepts the Anthropic-shaped ``thinking`` request field.
+
+        Only Claude reasoning models accept it. DeepSeek advertises ``supports_reasoning`` but reasons
+        natively: R1 returns a 400 when the field is sent and V3 silently ignores it.
+        """
+        if self._is_deepseek_model(model=model, base_model=base_model):
+            return False
+        return (
+            "claude-3-7" in model
+            or "claude-sonnet-4" in model
+            or "claude-opus-4" in model
+            or supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)
+            or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider)
+        )
+
+    def _model_rejects_reasoning_effort_param(self, model: str, base_model: str) -> bool:
+        """Whether the model returns a 400 for every ``reasoning_effort`` shape on Converse.
+
+        DeepSeek R1 always reasons and rejects any reasoning request field. DeepSeek V3 accepts a raw
+        ``reasoning_effort`` like gpt-oss does, and every other model maps it to a shape it accepts.
+        """
+        return self._is_deepseek_r1_model(model=model, base_model=base_model)
+
     def get_supported_openai_params(self, model: str) -> list[str]:
         from litellm.utils import supports_function_calling
 
@@ -564,23 +599,20 @@ class AmazonConverseConfig(BaseConfig):
             # only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
             supported_params.append("tool_choice")
 
-        if "gpt-oss" in model or "openai.gpt-5" in model or "openai.gpt-5" in base_model:
+        if (
+            "gpt-oss" in model
+            or self._is_openai_gpt_reasoning_model(model)
+            or self._is_openai_gpt_reasoning_model(base_model)
+        ):
             supported_params.append("reasoning_effort")
+        elif self._is_deepseek_model(model=model, base_model=base_model):
+            if not self._is_deepseek_r1_model(model=model, base_model=base_model):
+                supported_params.append("reasoning_effort")
         elif self._is_nova_2_model(model):
             # Nova 2 models support reasoning_effort (transformed to reasoningConfig)
             # These models use a different reasoning structure than Anthropic's thinking parameter
             supported_params.append("reasoning_effort")
-        elif (
-            "claude-3-7" in model
-            or "claude-sonnet-4" in model
-            or "claude-opus-4" in model
-            or "deepseek.r1" in model
-            or supports_reasoning(
-                model=model,
-                custom_llm_provider=self.custom_llm_provider,
-            )
-            or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider)
-        ):
+        elif self._model_accepts_anthropic_thinking_param(model=model, base_model=base_model):
             supported_params.append("thinking")
             supported_params.append("reasoning_effort")
             supported_params.append("output_config")
@@ -872,6 +904,11 @@ class AmazonConverseConfig(BaseConfig):
         drop_params: bool,
     ) -> dict:
         is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params)
+        base_model: Final = BedrockModelInfo.get_base_model(model)
+        drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model)
+        drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param(
+            model=model, base_model=base_model
+        )
 
         for param, value in non_default_params.items():
             if param == "response_format" and isinstance(value, dict):
@@ -920,7 +957,12 @@ class AmazonConverseConfig(BaseConfig):
                 optional_params["_parallel_tool_use_config"] = {
                     "tool_choice": {"type": "auto", "disable_parallel_tool_use": not value}
                 }
-            if param == "thinking" and "openai.gpt-5" not in model:
+            if param == "thinking" and drop_thinking_param:
+                verbose_logger.debug(
+                    "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.",
+                    model,
+                )
+            elif param == "thinking" and not self._is_openai_gpt_reasoning_model(model):
                 if (
                     isinstance(value, dict)
                     and value.get("type") == "adaptive"
@@ -946,6 +988,11 @@ class AmazonConverseConfig(BaseConfig):
                     AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
                         model=model, optional_params=optional_params, custom_llm_provider="bedrock"
                     )
+            elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_effort_param:
+                verbose_logger.debug(
+                    "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it always reasons and rejects it.",
+                    model,
+                )
             elif param == "reasoning_effort" and isinstance(value, str):
                 self._handle_reasoning_effort_parameter(
                     model=model, reasoning_effort=value, optional_params=optional_params
@@ -1805,6 +1852,7 @@ class AmazonConverseConfig(BaseConfig):
             data=request_data,
             messages=messages,
             encoding=encoding,
+            json_mode=json_mode,
         )
 
     def _transform_reasoning_content(self, reasoning_content_blocks: list[BedrockConverseReasoningContentBlock]) -> str:
@@ -2237,6 +2285,7 @@ class AmazonConverseConfig(BaseConfig):
         data: dict | str,
         messages: list,
         encoding,
+        json_mode: bool | None = None,
     ) -> ModelResponse:
         ## LOGGING
         if logging_obj is not None:
@@ -2247,7 +2296,9 @@ class AmazonConverseConfig(BaseConfig):
                 additional_args={"complete_input_dict": data},
             )
 
-        json_mode: Final[bool | None] = optional_params.get("json_mode", None)
+        resolved_json_mode: Final[bool | None] = (
+            json_mode if json_mode is not None else optional_params.get("json_mode", None)
+        )
         ## RESPONSE OBJECT
         try:
             completion_response: Final = ConverseResponseBlock(**response.json())
@@ -2339,7 +2390,7 @@ class AmazonConverseConfig(BaseConfig):
             chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks)
         chat_completion_message["content"] = content_str
         filtered_tools: Final = self._filter_json_mode_tools(
-            json_mode=json_mode,
+            json_mode=resolved_json_mode,
             tools=tools,
             chat_completion_message=chat_completion_message,
         )
@@ -2363,7 +2414,7 @@ class AmazonConverseConfig(BaseConfig):
         # When json_mode filtered out all synthetic tool calls the response
         # is plain content, not a pending tool invocation. Fix finish_reason
         # so callers (e.g. OpenAI SDK) don't misinterpret it.
-        if json_mode and not filtered_tools and tools:
+        if resolved_json_mode and not filtered_tools and tools:
             initial_finish_reason = "stop"
 
         (

litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py

@@ -340,6 +340,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
                     optional_params=optional_params,
                     litellm_params=litellm_params,
                     encoding=encoding,
+                    json_mode=json_mode,
                 )
             elif provider == "twelvelabs":
                 return litellm.AmazonTwelveLabsPegasusConfig().transform_response(

litellm/llms/bedrock/count_tokens/handler.py

@@ -10,9 +10,10 @@ import httpx
 
 import litellm
 from litellm._logging import verbose_logger
+from litellm.llms.bedrock.base_aws_llm import run_aws_signing
 from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig
-from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
 
 
 class BedrockCountTokensHandler(BedrockCountTokensConfig):
@@ -27,6 +28,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
         request_data: dict[str, Any],
         litellm_params: dict[str, Any],
         resolved_model: str,
+        client: AsyncHTTPHandler | None = None,
     ) -> dict[str, Any]:
         """
         Handle a CountTokens request using existing LiteLLM patterns.
@@ -75,7 +77,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
             # Extract api_key for bearer token auth if provided
             api_key: Final = litellm_params.get("api_key", None)
             headers: Final = {"Content-Type": "application/json"}
-            signed_headers, signed_body = self._sign_request(
+            signed_headers, signed_body = await run_aws_signing(
+                self._sign_request,
                 service_name="bedrock",
                 headers=headers,
                 optional_params=litellm_params,
@@ -85,7 +88,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
                 api_key=api_key,
             )
 
-            async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
+            async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
 
             response: Final = await async_client.post(
                 endpoint_url,

litellm/llms/bedrock/embed/embedding.py

@@ -5,7 +5,7 @@ Handles embedding calls to Bedrock's `/invoke` endpoint
 import copy
 import json
 import urllib.parse
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import TYPE_CHECKING, Final, get_args, overload
 
 import httpx
@@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import (
 )
 from litellm.types.utils import EmbeddingResponse, LlmProviders
 
-from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
+from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing
 from ..common_utils import BedrockError
 from .amazon_nova_transformation import AmazonNovaEmbeddingConfig
 from .amazon_titan_g1_transformation import AmazonTitanG1Config
@@ -41,6 +41,20 @@ if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 
 
+def _sign_get_request(
+    credentials: Credentials, url: str, headers: Mapping[str, str], aws_region_name: str
+) -> AWSPreparedRequest:
+    try:
+        from botocore.auth import SigV4Auth
+        from botocore.awsrequest import AWSRequest
+    except ImportError:
+        raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
+
+    request: Final = AWSRequest(method="GET", url=url, data=None, headers=headers)
+    SigV4Auth(credentials, "bedrock", aws_region_name).add_auth(request)
+    return request.prepare()
+
+
 class BedrockEmbedding(BaseAWSLLM):
     @overload
     def _load_credentials(
@@ -342,7 +356,8 @@ class BedrockEmbedding(BaseAWSLLM):
             if extra_headers is not None:
                 headers = {"Content-Type": "application/json", **extra_headers}
 
-            prepped = self.get_request_headers(
+            prepped = await run_aws_signing(
+                self.get_request_headers,
                 credentials=credentials,
                 aws_region_name=aws_region_name,
                 extra_headers=extra_headers,
@@ -600,9 +615,6 @@ class BedrockEmbedding(BaseAWSLLM):
             dict: Status response from AWS Bedrock
         """
 
-        # Get AWS credentials using the same method as other Bedrock methods
-        credentials, _ = self._load_credentials(kwargs)
-
         # Get the runtime endpoint
         endpoint_url, _ = self.get_runtime_endpoint(
             api_base=None,
@@ -619,27 +631,13 @@ class BedrockEmbedding(BaseAWSLLM):
         # Prepare headers for GET request
         headers: Final = {"Content-Type": "application/json"}
 
-        # Use AWSRequest directly for GET requests (get_request_headers hardcodes POST)
-        try:
-            from botocore.auth import SigV4Auth
-            from botocore.awsrequest import AWSRequest
-        except ImportError:
-            raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
+        def sign_status_request() -> AWSPreparedRequest:
+            credentials, _ = self._load_credentials(kwargs)
+            return _sign_get_request(
+                credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name
+            )
 
-        # Create AWSRequest with GET method and encoded URL
-        request: Final = AWSRequest(
-            method="GET",
-            url=status_url,
-            data=None,  # GET request, no body
-            headers=headers,
-        )
-
-        # Sign the request - SigV4Auth will create canonical string from request URL
-        sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name)
-        sigv4.add_auth(request)
-
-        # Prepare the request
-        prepped: Final = request.prepare()
+        prepped: Final = await run_aws_signing(sign_status_request)
 
         # LOGGING
         if logging_obj is not None:

litellm/llms/bedrock/realtime/handler.py

@@ -21,7 +21,7 @@ from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeE
 from litellm.types.llms.openai import OpenAIRealtimeEvents
 from litellm.types.realtime import RealtimeResponseTransformInput
 
-from ..base_aws_llm import BaseAWSLLM
+from ..base_aws_llm import BaseAWSLLM, run_aws_signing
 from ..common_utils import BedrockError
 from .transformation import BedrockRealtimeConfig
 
@@ -149,7 +149,8 @@ class BedrockRealtime(BaseAWSLLM):
 
         verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model)
 
-        credentials: Final = self.get_credentials(
+        credentials: Final = await run_aws_signing(
+            self.get_credentials,
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
             aws_session_token=aws_session_token,
@@ -169,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM):
                     "or configure credentials in the environment"
                 ),
             )
-        frozen_credentials: Final = credentials.get_frozen_credentials()
+        frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials)
 
         # Initialize Bedrock client with aws_sdk_bedrock_runtime
         config: Final = Config(

litellm/llms/bedrock_mantle/common_utils.py

@@ -23,7 +23,7 @@ from botocore.exceptions import (
     ProfileNotFound,
 )
 
-from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS
 from litellm.secret_managers.main import get_secret_str
 
 BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
@@ -55,7 +55,7 @@ def resolve_mantle_region(params: Mapping[str, object]) -> str:
     )
 
 
-class BedrockMantleAuthMixin:
+class BedrockMantleAuthMixin(SignsRequestsWithAWS):
     _aws_signer: BaseAWSLLM
 
     @staticmethod

litellm/llms/custom_httpx/llm_http_handler.py

@@ -77,6 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import (
     BaseVectorStoreFilesConfig,
 )
 from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
+from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws
 from litellm.llms.custom_httpx.container_handler import raise_for_error_status
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
@@ -637,7 +638,12 @@ class BaseLLMHTTPHandler:
                         headers=request_headers,
                     ),
                 )
-                return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed))
+                signed_request: Final = await (
+                    run_aws_signing(sign_and_log, transformed)
+                    if isinstance(provider_config, SignsRequestsWithAWS)
+                    else asyncio.to_thread(sign_and_log, transformed)
+                )
+                return await dispatch_async(*signed_request)
 
             return transform_then_dispatch()
 
@@ -1973,7 +1979,9 @@ class BaseLLMHTTPHandler:
             api_key=api_key,
         )
 
-        signed_headers, signed_json_body = provider_config.sign_request(
+        signed_headers, signed_json_body = await sign_request_off_loop_if_aws(
+            provider_config,
+            provider_config.sign_request,
             headers=headers,
             optional_params=optional_params,
             request_data=data,
@@ -2074,7 +2082,9 @@ class BaseLLMHTTPHandler:
                         max_attempts,
                     )
                     provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body)
-                    headers, signed_json_body = provider_config.sign_request(
+                    headers, signed_json_body = await sign_request_off_loop_if_aws(
+                        provider_config,
+                        provider_config.sign_request,
                         headers=headers,
                         optional_params=optional_params_dict,
                         request_data=request_body,
@@ -2234,7 +2244,9 @@ class BaseLLMHTTPHandler:
             stream=stream,
         )
 
-        headers, signed_json_body = anthropic_messages_provider_config.sign_request(
+        headers, signed_json_body = await sign_request_off_loop_if_aws(
+            anthropic_messages_provider_config,
+            anthropic_messages_provider_config.sign_request,
             headers=headers,
             optional_params=dict(litellm_params),  # dynamic aws_* params are passed under litellm_params
             request_data=request_body,
@@ -2910,7 +2922,9 @@ class BaseLLMHTTPHandler:
                 fake_stream=fake_stream,
             )
 
-        headers, signed_body = responses_api_provider_config.sign_request(
+        headers, signed_body = await sign_request_off_loop_if_aws(
+            responses_api_provider_config,
+            responses_api_provider_config.sign_request,
             headers=headers,
             optional_params=dict(litellm_params),
             request_data=data,
@@ -4618,7 +4632,9 @@ class BaseLLMHTTPHandler:
         )
         data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
 
-        headers, signed_body = responses_api_provider_config.sign_request(
+        headers, signed_body = await sign_request_off_loop_if_aws(
+            responses_api_provider_config,
+            responses_api_provider_config.sign_request,
             headers=headers,
             optional_params=dict(litellm_params),
             request_data=data,
@@ -9845,7 +9861,9 @@ class BaseLLMHTTPHandler:
             )
         all_optional_params: Final[dict[str, object]] = dict(litellm_params)
         all_optional_params.update(vector_store_search_optional_params or {})
-        headers, signed_json_body = vector_store_provider_config.sign_request(
+        headers, signed_json_body = await sign_request_off_loop_if_aws(
+            vector_store_provider_config,
+            vector_store_provider_config.sign_request,
             headers=headers,
             optional_params=all_optional_params,
             request_data=request_body,

litellm/llms/databricks/chat/transformation.py

@@ -250,8 +250,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
         litellm_params: dict,
         stream: bool | None = None,
     ) -> str:
-        api_base = self._get_api_base(api_base)
-        complete_url: Final = f"{api_base}/chat/completions"
+        use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2
+        api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway)
+        url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base
+        complete_url: Final = f"{url_base}/chat/completions"
         return complete_url
 
     def get_supported_openai_params(self, model: str | None = None) -> list:

litellm/llms/databricks/common_utils.py

@@ -177,19 +177,13 @@ class DatabricksBase:
         # Default: just litellm
         return f"litellm/{version}"
 
-    def _get_api_base(self, api_base: str | None) -> str:
-        """
-        Get the Databricks API base URL.
-
-        If not provided, attempts to get it from the Databricks SDK.
-        """
+    def _get_api_base(self, api_base: str | None, use_ai_gateway: bool = False) -> str:
         if api_base is None:
             try:
                 from databricks.sdk import WorkspaceClient
 
                 databricks_client: Final = WorkspaceClient()
                 api_base = f"{databricks_client.config.host}/serving-endpoints"
-                return api_base
             except ImportError:
                 raise DatabricksException(
                     status_code=400,
@@ -198,6 +192,18 @@ class DatabricksBase:
                         "or install the databricks-sdk Python library."
                     ),
                 )
+
+        if not use_ai_gateway:
+            return api_base
+
+        normalized_api_base: Final = api_base.rstrip("/")
+        if normalized_api_base.endswith("/ai-gateway/mlflow/v1"):
+            return normalized_api_base
+        if normalized_api_base.endswith("/serving-endpoints"):
+            return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1"
+        api_base_parts: Final = urlsplit(normalized_api_base)
+        if api_base_parts.path in ("", "/"):
+            return f"{normalized_api_base}/ai-gateway/mlflow/v1"
         return api_base
 
     def _get_oauth_m2m_token(

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -81,6 +81,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     """
 
     delivers_ended_stream_rewrites = True
+    assembles_streamed_response = True
 
     def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
         """

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -37,14 +37,14 @@ from itertools import accumulate, chain, repeat
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
 
-from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.completion_extras.litellm_responses_transformation.transformation import (
     LiteLLMResponsesTransformationHandler,
     OpenAiResponsesToChatCompletionStreamIterator,
+    tool_call_dict_from_output_item,
 )
 from litellm.llms.base_llm.guardrail_translation.base_translation import (
     BaseTranslation,
@@ -84,7 +84,6 @@ from litellm.types.llms.openai import (
 )
 from litellm.types.responses.main import (
     GenericResponseOutputItem,
-    OutputFunctionToolCall,
     OutputText,
 )
 from litellm.types.utils import GenericGuardrailAPIInputs
@@ -106,6 +105,19 @@ class _ToolCallShape(NamedTuple):
     arguments: str
 
 
+class _ToolCallFunctionFields(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    name: str | None = None
+    arguments: str = ""
+
+
+class _ToolCallFields(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    function: _ToolCallFunctionFields
+
+
 def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]:
     return tuple(
         _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", ""))
@@ -113,6 +125,47 @@ def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tupl
     )
 
 
+def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None:
+    payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
+    try:
+        fields: Final = _ToolCallFields.model_validate(payload)
+    except ValidationError:
+        return None
+    return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments)
+
+
+def _post_guardrail_tool_call_shapes(
+    returned_tool_calls: Sequence[object] | None,
+    pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+    guardrail_name: str | None,
+) -> tuple[_ToolCallShape, ...]:
+    if not pre_guardrail_tool_calls:
+        return pre_guardrail_tool_calls
+    if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls):
+        verbose_proxy_logger.warning(
+            "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, "
+            "leaving the tool call output items unchanged",
+            guardrail_name,
+            "no" if returned_tool_calls is None else len(returned_tool_calls),
+            len(pre_guardrail_tool_calls),
+        )
+        return pre_guardrail_tool_calls
+    returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls)
+    validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None)
+    if len(validated_shapes) != len(returned_shapes):
+        verbose_proxy_logger.warning(
+            "OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, "
+            "leaving the tool call output items unchanged",
+            guardrail_name,
+        )
+        return pre_guardrail_tool_calls
+    return validated_shapes
+
+
+def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape:
+    return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments)
+
+
 class ResponseOutputEnvelope(TypedDict, total=False):
     """Dict form of a Responses API response, as far as guardrail write-back reads it."""
 
@@ -140,8 +193,18 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
 )
 
 
-_FUNCTION_CALL_ARGUMENT_EVENT_TYPES: Final = frozenset(
-    {"response.function_call_arguments.delta", "response.function_call_arguments.done"}
+_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
+_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
+    {"function_call": "arguments", "custom_tool_call": "input"}
+)
+_TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: Final = frozenset(
+    {"response.function_call_arguments.delta", "response.custom_tool_call_input.delta"}
+)
+_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
+    {"response.function_call_arguments.done": "arguments", "response.custom_tool_call_input.done": "input"}
+)
+_TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | frozenset(
+    _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS
 )
 _OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
 _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
@@ -180,8 +243,20 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp
     return {**item, field: converted_value}  # mutable-ok: request input items must stay JSON-plain dicts
 
 
-def _is_function_call_item(item: object) -> bool:
-    return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call")
+def _is_tool_call_item(item: object) -> bool:
+    return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES
+
+
+def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None:
+    if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES:
+        return None
+    if isinstance(item, Mapping):
+        return cast("Mapping[str, object]", item)  # cast-ok: output items are str-keyed JSON objects
+    return item.model_dump() if isinstance(item, BaseModel) else None
+
+
+def _is_tool_call_output_item(item: object) -> bool:
+    return _tool_call_output_item_mapping(item) is not None
 
 
 def _last_message_role(messages: Sequence[object]) -> str | None:
@@ -205,7 +280,7 @@ def _provenance_unit_bounds(
     start_indexes: Final = tuple(
         index
         for index in range(len(raw_input))
-        if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
+        if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
     )
     return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input))))
 
@@ -357,6 +432,7 @@ class OpenAIResponsesHandler(BaseTranslation):
     """
 
     delivers_ended_stream_rewrites = True
+    assembles_streamed_response = True
 
     def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
         """
@@ -603,7 +679,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             - response.output is a list of output items
             - Each output item can be:
               * GenericResponseOutputItem with a content list of OutputText objects
-              * ResponseFunctionToolCall with tool call data
+              * ResponseFunctionToolCall or CustomToolCallOutputItem with tool call data
             - Each OutputText object has a text field
         """
 
@@ -668,6 +744,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             if response_model:
                 inputs["model"] = response_model
 
+            pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
             guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
                 inputs=inputs,
                 request_data=request_data,
@@ -676,6 +753,11 @@ class OpenAIResponsesHandler(BaseTranslation):
             )
 
             guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
+            post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
+                returned_tool_calls=guardrailed_inputs.get("tool_calls"),
+                pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+                guardrail_name=guardrail_to_apply.guardrail_name,
+            )
 
             # Step 3: Map guardrail responses back to original response structure
             await self._apply_guardrail_responses_to_output(
@@ -683,6 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation):
                 responses=guardrailed_texts,
                 task_mappings=task_mappings,
             )
+            self._write_tool_call_rewrites_to_output(
+                tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)),
+                pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+                post_guardrail_tool_calls=post_guardrail_tool_calls,
+            )
 
         verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response)
 
@@ -779,11 +866,10 @@ class OpenAIResponsesHandler(BaseTranslation):
                 )
 
                 guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
-                returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
-                post_guardrail_tool_calls: Final = _tool_call_shapes(
-                    returned_tool_calls
-                    if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
-                    else tool_calls_to_check
+                post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
+                    returned_tool_calls=guardrailed_inputs.get("tool_calls"),
+                    pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+                    guardrail_name=guardrail_to_apply.guardrail_name,
                 )
 
                 # Write guardrailed texts back into the output items in-place.
@@ -933,11 +1019,12 @@ class OpenAIResponsesHandler(BaseTranslation):
         guardrail_name: str,
     ) -> None:
         """Write ended-stream guardrail tool-call rewrites into the completed
-        envelope's ``function_call`` items and sync the earlier stream events,
-        keyed by ``call_id``. The guardrail sees the envelope's function calls
-        in output order, which is how a rewritten call finds its ``call_id``;
-        the stream events find their call through the ``call_id`` on
-        ``output_item`` events and the ``item_id`` on argument events, since an
+        envelope's ``function_call`` and ``custom_tool_call`` items and sync the
+        earlier stream events, keyed by ``call_id``. The guardrail sees the
+        envelope's tool calls in output order, which is how a rewritten call
+        finds its ``call_id``; the stream events find their call through the
+        ``call_id`` on ``output_item`` events and the ``item_id`` on argument
+        and custom-input events, since an
         event's ``output_index`` need not match the envelope's (the chat bridge
         numbers tool calls from 1 while the envelope lists them after the
         message). A rewrite whose calls do not line up with the envelope, or
@@ -945,32 +1032,30 @@ class OpenAIResponsesHandler(BaseTranslation):
         pipeline executor discards it and releases the original events."""
         if post_guardrail_tool_calls == pre_guardrail_tool_calls:
             return
-        function_call_items: Final = tuple(
-            output_item for output_item in outputs if stream_item_field(output_item, "type") == "function_call"
-        )
+        tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item))
         call_ids: Final = tuple(
             call_id
-            for output_item in function_call_items
+            for output_item in tool_call_items
             if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id
         )
         stream_events: Final = responses_so_far[:-1]
-        call_id_by_item_id: Final = self._function_call_ids_by_item_id(stream_events)
+        call_id_by_item_id: Final = self._tool_call_ids_by_item_id(stream_events)
         event_call_ids: Final = tuple(
-            self._function_call_event_call_id(event, call_id_by_item_id) for event in stream_events
+            self._tool_call_event_call_id(event, call_id_by_item_id) for event in stream_events
         )
         rewrites_by_call_id: Final = MappingProxyType(
             {
-                call_id: after
+                call_id: _tool_call_rewrite(before, after)
                 for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls)
                 if after != before
             }
         )
         unresolved_argument_event: Final = any(
-            call_id is None and stream_item_field(event, "type") in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES
+            call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
             for event, call_id in zip(stream_events, event_call_ids)
         )
         if (
-            len(call_ids) != len(function_call_items)
+            len(call_ids) != len(tool_call_items)
             or len(frozenset(call_ids)) != len(call_ids)
             or len(call_ids) != len(post_guardrail_tool_calls)
             or unresolved_argument_event
@@ -981,10 +1066,10 @@ class OpenAIResponsesHandler(BaseTranslation):
             raise UndeliverableStreamRewrite(guardrail_name)
         for output_item, rewrite in (
             (output_item, rewrites_by_call_id[call_id])
-            for output_item, call_id in zip(function_call_items, call_ids)
+            for output_item, call_id in zip(tool_call_items, call_ids)
             if call_id in rewrites_by_call_id
         ):
-            self._write_function_call_item(output_item, rewrite.name, rewrite.arguments)
+            self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
         delta_replacements: Final = MappingProxyType(
             {call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()}
         )
@@ -992,16 +1077,18 @@ class OpenAIResponsesHandler(BaseTranslation):
             if call_id not in rewrites_by_call_id:
                 continue
             match stream_item_field(event, "type"):
-                case "response.function_call_arguments.delta":
+                case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES:
                     self._write_event_field(event, "delta", next(delta_replacements[call_id]))
-                case "response.function_call_arguments.done":
-                    self._write_event_field(event, "arguments", rewrites_by_call_id[call_id].arguments)
+                case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS:
+                    self._write_event_field(
+                        event, _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS[event_type], rewrites_by_call_id[call_id].arguments
+                    )
                 case "response.output_item.added":
-                    self._write_function_call_item(
+                    self._write_tool_call_item(
                         stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None
                     )
                 case "response.output_item.done":
-                    self._write_function_call_item(
+                    self._write_tool_call_item(
                         stream_item_field(event, "item"),
                         rewrites_by_call_id[call_id].name,
                         rewrites_by_call_id[call_id].arguments,
@@ -1009,8 +1096,23 @@ class OpenAIResponsesHandler(BaseTranslation):
                 case _:
                     pass
 
+    def _write_tool_call_rewrites_to_output(
+        self,
+        tool_call_items: Sequence[object],
+        pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+        post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+    ) -> None:
+        if len(tool_call_items) != len(post_guardrail_tool_calls):
+            return
+        for output_item, rewrite in (
+            (output_item, _tool_call_rewrite(before, after))
+            for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls)
+            if after != before
+        ):
+            self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
+
     @staticmethod
-    def _function_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
+    def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
         items: Final = tuple(
             stream_item_field(event, "item")
             for event in stream_events
@@ -1020,32 +1122,35 @@ class OpenAIResponsesHandler(BaseTranslation):
             {
                 item_id: call_id
                 for item in items
-                if stream_item_field(item, "type") == "function_call"
+                if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
                 and isinstance(item_id := stream_item_field(item, "id"), str)
                 and isinstance(call_id := stream_item_field(item, "call_id"), str)
             }
         )
 
     @staticmethod
-    def _function_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
+    def _tool_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
         event_type: Final = stream_item_field(event, "type")
-        if event_type in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES:
+        if event_type in _TOOL_CALL_PAYLOAD_EVENT_TYPES:
             item_id: Final = stream_item_field(event, "item_id")
             return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None
         if event_type not in _OUTPUT_ITEM_EVENT_TYPES:
             return None
         item: Final = stream_item_field(event, "item")
         call_id: Final = stream_item_field(item, "call_id")
-        return call_id if stream_item_field(item, "type") == "function_call" and isinstance(call_id, str) else None
+        return (
+            call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None
+        )
 
     @staticmethod
-    def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None:
+    def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None:
         if item is None:
             return
         if name is not None:
             OpenAIResponsesHandler._write_event_field(item, "name", name)
-        if arguments is not None:
-            OpenAIResponsesHandler._write_event_field(item, "arguments", arguments)
+        item_type: Final = stream_item_field(item, "type")
+        if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS:
+            OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload)
 
     def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
         """
@@ -1073,7 +1178,7 @@ class OpenAIResponsesHandler(BaseTranslation):
     def _completed_response_scan_key(response: object) -> StreamingScanKey:
         output_items: Final = stream_item_items(response, "output")
         message_items: Final = tuple(
-            item for item in output_items if stream_item_field(item, "type") != "function_call"
+            item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES
         )
         return StreamingScanKey(
             texts=tuple(
@@ -1085,7 +1190,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             tool_calls=tuple(
                 stream_item_fingerprint(item)
                 for item in output_items
-                if stream_item_field(item, "type") == "function_call"
+                if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
             ),
             stream_ended=True,
         )
@@ -1196,34 +1301,10 @@ class OpenAIResponsesHandler(BaseTranslation):
         Override this method to customize text/image/tool extraction logic.
         """
 
-        # Check if this is a tool call (OutputFunctionToolCall)
-        if isinstance(output_item, OutputFunctionToolCall) or (
-            isinstance(output_item, BaseModel)
-            and hasattr(output_item, "type")
-            and getattr(output_item, "type") == "function_call"
-        ):
+        tool_call_item: Final = _tool_call_output_item_mapping(output_item)
+        if tool_call_item is not None:
             if tool_calls_to_check is not None:
-                tool_call_dict = (
-                    LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
-                        tool_call_item=output_item,
-                        index=output_idx,
-                    )
-                )
-                tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
-            return
-        elif isinstance(output_item, dict) and output_item.get("type") == "function_call":
-            # Handle dict representation of tool call
-            if tool_calls_to_check is not None:
-                # Convert dict to ResponseFunctionToolCall for processing
-                try:
-                    tool_call_obj: Final = ResponseFunctionToolCall(**output_item)
-                    tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
-                        tool_call_item=tool_call_obj,
-                        index=output_idx,
-                    )
-                    tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
-                except Exception:
-                    pass
+                tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx))
             return
 
         # Handle both GenericResponseOutputItem and dict

litellm/main.py

@@ -5398,6 +5398,14 @@ def completion(
         if dynamic_api_key is not None:
             api_key = dynamic_api_key
         # check if user passed in any of the OpenAI optional params
+        bridges_to_responses_api: Final = (
+            responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge
+        )
+        allowed_openai_params: Final[list[str] | None] = (
+            [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"]
+            if bridges_to_responses_api
+            else kwargs.get("allowed_openai_params")
+        )
         optional_param_args: Final = {
             "functions": functions,
             "function_call": function_call,
@@ -5442,7 +5450,7 @@ def completion(
             "service_tier": service_tier,
             "store": store,
             "prompt_cache_key": prompt_cache_key,
-            "allowed_openai_params": kwargs.get("allowed_openai_params"),
+            "allowed_openai_params": allowed_openai_params,
             "base_model": base_model,
         }
         optional_params = get_optional_params(**optional_param_args, **non_default_params)

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -1234,7 +1234,7 @@ if MCP_AVAILABLE:
         return client_id, client_secret, scopes
 
     _STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset(
-        (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization)
+        (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization, MCPAuth.token)
     )
 
     @dataclass(frozen=True, slots=True)
@@ -1243,6 +1243,17 @@ if MCP_AVAILABLE:
         mcp_auth_header: str | None
         oauth2_headers: dict[str, str] | None
 
+    def _preview_origin(url: str | None) -> tuple[str, str, int | None] | None:
+        if not url:
+            return None
+        try:
+            parsed: Final = httpx.URL(url)
+        except httpx.InvalidURL:
+            return None
+        if parsed.scheme not in ("http", "https") or not parsed.host:
+            return None
+        return parsed.scheme, parsed.host, parsed.port
+
     def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest:
         """
         Resolve the credentials a not-yet-saved server config carries for a preview call.
@@ -1255,7 +1266,19 @@ if MCP_AVAILABLE:
             MCPRequestHandler,
         )
 
-        request: Final = _inherit_credentials_from_existing_server(new_mcp_server_request)
+        saved_server: Final = (
+            global_mcp_server_manager.get_mcp_server_by_id(new_mcp_server_request.server_id)
+            if new_mcp_server_request.server_id
+            else None
+        )
+        saved_origin: Final = _preview_origin(saved_server.url) if saved_server else None
+        preview_origin: Final = _preview_origin(new_mcp_server_request.url)
+        may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or (
+            saved_origin is not None and saved_origin == preview_origin
+        )
+        request: Final = (
+            _inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request
+        )
         mcp_auth_header: Final = (
             request.credentials.get("auth_value")
             if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict)
@@ -1318,8 +1341,15 @@ if MCP_AVAILABLE:
             if _oauth2_flow == "client_credentials" and not request.token_url:
                 _oauth2_flow = None
 
+            # Static previews inherit credentials before this step, but must not resolve back to
+            # the saved record during client creation and discard the edited connection settings.
+            preview_server_id: Final = (
+                ""
+                if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES or request.auth_type in (None, MCPAuth.none)
+                else request.server_id or ""
+            )
             server_model: Final = MCPServer(
-                server_id=request.server_id or "",
+                server_id=preview_server_id,
                 name=request.alias or request.server_name or "",
                 url=request.url,
                 transport=request.transport,

litellm/proxy/client/cli/README.md

@@ -508,7 +508,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_
 
 ### Route Every Claude Code Session Through the Proxy
 
-`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
+`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
 
 Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you.
 
@@ -532,12 +532,28 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi
 lite --base-url https://your-proxy.example.com login --config-claude
 ```
 
-It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
+It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
 
 Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`.
 
 Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops.
 
+#### Configuring Claude Code Once, With a Virtual Key or Your Login
+
+`lite configure claude` wires Claude Code up persistently and `lite unconfigure claude` puts things back. It is what `lite login --config-claude` does, plus a pinned model and an undo, and it also takes a long-lived virtual key when that is what you have:
+
+```bash
+curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
+lite --base-url https://your-proxy.example.com configure claude --api-key sk-... --model claude-auto
+claude
+```
+
+With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (the ones whose id contains `claude` or `anthropic`) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window and sends no thinking parameters for it, so either name the group like a Claude model id or append `[1m]` to opt into the 1M window. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control
+
+Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt
+
+What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Like `--config-claude`, both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any login prompt or request
+
 ### QA Complexity-Based Auto-Routing Against Your Real Proxy
 
 `lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session.
@@ -584,7 +600,7 @@ An interactive wizard. It runs the same model-group discovery as above, splits t
 
 The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/<model-name>` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key.
 
-You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
+You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets the top-level `model` and `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, `ANTHROPIC_DEFAULT_OPUS_MODEL` and `ANTHROPIC_DEFAULT_FABLE_MODEL` to `autorouter` in `~/.claude/settings.json` (and `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when missing, like every other wiring), so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
 
 You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first.
 

litellm/proxy/client/cli/commands/agents.py

@@ -4,6 +4,7 @@ import subprocess
 import sys
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
+from pathlib import Path
 from types import MappingProxyType
 from typing import Final, TypeAlias
 
@@ -12,10 +13,12 @@ import requests
 from pydantic import BaseModel, TypeAdapter, ValidationError
 
 from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login
+from .claude_settings import claude_settings_path, lite_api_key_helper_configured
 from .cmd_quoting import quote_for_cmd
 from .pi import (
     LITELLM_PROXY_API_KEY_ENV,
     PI_PROVIDER_NAME,
+    ListingFailure,
     PiSyncError,
     fetch_model_ids,
     fetch_model_limits,
@@ -83,6 +86,8 @@ def build_agent_env(
     base_url: str,
     api_key: str,
     profiles: frozenset[str],
+    *,
+    export_anthropic_token: bool = True,
 ) -> dict[str, str]:
     """Return a copy of base_env wired to route the agent through the proxy.
 
@@ -97,12 +102,19 @@ def build_agent_env(
     proxy's /v1/models; likewise left alone when already set.
     pi ignores both base URL variables and instead resolves $LITELLM_PROXY_API_KEY
     from its synced models.json provider entry.
+
+    With export_anthropic_token=False the bearer is left out (and any inherited
+    one dropped) so Claude Code asks its configured apiKeyHelper instead; Claude
+    Code prefers ANTHROPIC_AUTH_TOKEN over the helper and warns when both are set.
     """
     env: Final = dict(base_env)
     root: Final = base_url.rstrip("/")
     if PROFILE_ANTHROPIC in profiles:
         env[ANTHROPIC_BASE_URL_ENV] = root
-        env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key
+        if export_anthropic_token:
+            env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key
+        else:
+            env.pop(ANTHROPIC_AUTH_TOKEN_ENV, None)
         env.pop(ANTHROPIC_API_KEY_ENV, None)
         if ENABLE_TOOL_SEARCH_ENV not in env:
             env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE
@@ -165,7 +177,9 @@ def prepare_pi(
     """
     ids: Final = fetch_model_ids(base_url, api_key, get=get)
     if isinstance(ids, PiSyncError):
-        raise AgentRunError(ids.message)
+        raise AgentRunError(
+            f"{ids.message} pi would have nothing to run." if ids.kind is ListingFailure.EMPTY else ids.message
+        )
     limits: Final = fetch_model_limits(base_url, api_key, get=get)
     path: Final = models_json_path(base_env)
     error: Final = sync_models_json(path, base_url, ids, limits)
@@ -460,6 +474,7 @@ def run_agent(
     launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off,
     reattach_terminal: Callable[[], None] | None = None,
     preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS),
+    export_anthropic_token: bool = True,
 ) -> None:
     """Validate, wire the environment, and hand off to the agent.
 
@@ -491,7 +506,9 @@ def run_agent(
 
     env: Final = MappingProxyType(
         {
-            **build_agent_env(env_before_sync, base_url, api_key, profiles),
+            **build_agent_env(
+                env_before_sync, base_url, api_key, profiles, export_anthropic_token=export_anthropic_token
+            ),
             **(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced),
         }
     )
@@ -529,14 +546,26 @@ def resolve_api_key(ctx: click.Context) -> str:
 _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy."
 
 
+def _helper_supplies_token(
+    ctx_obj: CliContextObj, base_url: str, profiles: frozenset[str], settings_path: Path
+) -> bool:
+    if PROFILE_ANTHROPIC not in profiles or not ctx_obj.get("api_key_from_token_file"):
+        return False
+    return lite_api_key_helper_configured(base_url, settings_path)
+
+
 def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None:
     ctx_obj: Final[CliContextObj] = ctx.obj
     base_url: Final = ctx_obj["base_url"]
     started_interactive: Final = _is_interactive()
     api_key: Final = resolve_api_key(ctx)
 
-    display_name, _ = agent_profile(binary)
+    display_name, profiles = agent_profile(binary)
+    settings_path: Final = claude_settings_path(os.environ)
+    helper_supplies_token: Final = _helper_supplies_token(ctx_obj, base_url, profiles, settings_path)
     click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}")
+    if helper_supplies_token:
+        click.echo(f"litellm: {display_name} reads its key from the apiKeyHelper in {settings_path}")
 
     try:
         run_agent(
@@ -545,6 +574,7 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify
             [binary, *args],
             skip_verify=skip_verify,
             reattach_terminal=(_restore_controlling_terminal if started_interactive else None),
+            export_anthropic_token=not helper_supplies_token,
         )
     except AgentRunError as e:
         raise click.ClickException(str(e))

litellm/proxy/client/cli/commands/auth.py

@@ -1,3 +1,4 @@
+import os
 import sys
 import time
 import webbrowser
@@ -40,10 +41,16 @@ from litellm.litellm_core_utils.cli_token_utils import (
 )
 
 from .claude_settings import (
-    CLAUDE_SETTINGS_PATH,
-    SETTINGS_FILE_OWNERS,
+    STARTING_MODEL_ROLE,
+    ApiKeyHelper,
     ClaudeSettingsError,
-    write_claude_settings,
+    KeepModel,
+    claude_settings_path,
+    configure_claude_settings,
+    configure_state_path,
+    refuse_while_owned,
+    resolve_api_key_helper,
+    settings_file_owners,
 )
 from .pkce_login import (
     Http,
@@ -778,13 +785,24 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None:
 
 
 def _configure_claude_code(base_url: str) -> None:
-    """Point Claude Code at base_url by patching ~/.claude/settings.json."""
+    """Point Claude Code at base_url by patching the settings.json it reads, undoable with `lite unconfigure claude`."""
+    settings_path: Final = claude_settings_path(os.environ)
     try:
-        write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
+        configure_claude_settings(
+            base_url,
+            ApiKeyHelper(resolve_api_key_helper(base_url)),
+            KeepModel(),
+            settings_path,
+            configure_state_path(settings_path),
+            settings_file_owners(settings_path),
+        )
     except ClaudeSettingsError as e:
         raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}")
-    click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.")
-    click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.")
+    click.echo(f"\nConfigured Claude Code: {settings_path} now routes through {base_url.rstrip('/')}.")
+    click.echo(
+        "Your other Claude Code settings were left untouched. Restart Claude Code to pick this up. "
+        f"Undo with `lite unconfigure claude`; `lite configure claude --model` sets {STARTING_MODEL_ROLE}."
+    )
 
 
 def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: SecretSave) -> None:
@@ -853,6 +871,12 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
 
     ctx_obj: Final[CliContextObj] = ctx.obj
     base_url: Final = ctx_obj["base_url"]
+    if config_claude:
+        settings_path: Final = claude_settings_path(os.environ)
+        try:
+            refuse_while_owned(settings_path, settings_file_owners(settings_path))
+        except ClaudeSettingsError as e:
+            raise click.ClickException(f"Cannot configure Claude Code, so not logging in: {e}")
 
     try:
         if pkce:

litellm/proxy/client/cli/commands/autoroute/commands.py

@@ -14,11 +14,13 @@ from ..claude_settings import (
     AUTOROUTE_BACKUP_PATH,
     CLAUDE_SETTINGS_PATH,
     ClaudeSettingsError,
+    StaticToken,
     load_json_or_empty,
+    merge_claude_settings,
 )
 from ..up import BackupRecord as ClaudeBackupRecord
 from ..up import restore_claude_settings, write_backup
-from .config import master_key_from_config
+from .config import AUTOROUTER_MODEL_NAME, master_key_from_config
 from .process import (
     CONFIG_PATH,
     DEFAULT_AUTOROUTE_PORT,
@@ -37,7 +39,6 @@ from .process import (
     terminate,
     write_pid_record,
 )
-from .settings import merge_claude_settings_static_token
 from .wizard import run_configure_wizard
 
 _GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
@@ -156,7 +157,9 @@ def up(port: int) -> None:
             ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None),
             AUTOROUTE_BACKUP_PATH,
         )
-        merged: Final = merge_claude_settings_static_token(original_settings, base_url, master_key)
+        merged: Final = merge_claude_settings(
+            original_settings, base_url, StaticToken(master_key), AUTOROUTER_MODEL_NAME, AUTOROUTER_MODEL_NAME
+        )
         CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
         with secure_create(CLAUDE_SETTINGS_PATH) as f:
             json.dump(merged, f, indent=2)

litellm/proxy/client/cli/commands/autoroute/settings.py

@@ -1,51 +0,0 @@
-from typing import Final
-
-from pydantic import JsonValue
-
-from .config import AUTOROUTER_MODEL_NAME
-
-ENV_KEY: Final = "env"
-API_KEY_HELPER_KEY: Final = "apiKeyHelper"
-ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
-ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
-ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
-ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
-ENABLE_TOOL_SEARCH_VALUE: Final = "true"
-# Force every one of Claude Code's own model tiers to request the auto-router by name.
-# Router's auto-router registry is keyed by the literal requested model string
-# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*"
-# model_name can never work as a catch-all -- these overrides are what actually makes
-# Claude Code send "autorouter" regardless of /model or its own version-specific defaults.
-ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = (
-    "ANTHROPIC_DEFAULT_SONNET_MODEL",
-    "ANTHROPIC_DEFAULT_HAIKU_MODEL",
-    "ANTHROPIC_DEFAULT_OPUS_MODEL",
-)
-
-
-def merge_claude_settings_static_token(
-    settings: dict[str, JsonValue], base_url: str, auth_token: str
-) -> dict[str, JsonValue]:
-    """Return a new settings dict wired to a local ephemeral proxy with a static token.
-
-    Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real
-    remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key is the
-    locally persisted autoroute master key, so a plain env var is simpler and correct. Any
-    existing apiKeyHelper is cleared so it can't fight with the static token.
-    """
-    raw_env: Final = settings.get(ENV_KEY, {})
-    base_env: Final = raw_env if isinstance(raw_env, dict) else {}
-    env: Final[dict[str, JsonValue]] = {
-        ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
-        **base_env,
-        ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
-        ANTHROPIC_AUTH_TOKEN_KEY: auth_token,
-        **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS},
-    }
-    env.pop(ANTHROPIC_API_KEY_KEY, None)
-    merged: Final[dict[str, JsonValue]] = {**settings, ENV_KEY: env}
-    merged.pop(API_KEY_HELPER_KEY, None)
-    return merged
-
-
-__all__ = ["merge_claude_settings_static_token"]

litellm/proxy/client/cli/commands/claude_settings.py

@@ -1,37 +1,71 @@
 """Shared handling of Claude Code's ~/.claude/settings.json.
 
-`lite up` patches this file temporarily and restores it on exit; `lite login
---config-claude` patches it persistently. Both need the same merge and the same
-apiKeyHelper command, and `up` already imports from `auth`, so the shared parts
-live here rather than in either command module.
+`lite up` and `lite autoroute up` patch this file temporarily and restore it on
+exit; `lite login --config-claude` and `lite configure claude` patch it
+persistently and record how to undo it. All of them need the same merge and the
+same apiKeyHelper command, and `up` already imports from `auth`, so the shared
+parts live here rather than in any one command module.
 """
 
+import hashlib
+import json
 import shlex
 import shutil
 import sys
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
+from functools import reduce
+from itertools import chain
 from pathlib import Path
-from typing import Final
+from types import MappingProxyType
+from typing import Final, TypeAlias
 
-from pydantic import JsonValue, TypeAdapter, ValidationError
+from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
 
-from litellm.litellm_core_utils.private_json import write_private_json
+from litellm.litellm_core_utils.private_json import (
+    commit_staged_json,
+    discard_staged_json,
+    ensure_private_dir,
+    stage_private_json,
+)
 
 from .cmd_quoting import quote_for_cmd
 
 ENV_KEY: Final = "env"
 API_KEY_HELPER_KEY: Final = "apiKeyHelper"
+MODEL_KEY: Final = "model"
 ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
+ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
 ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
 ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
 ENABLE_TOOL_SEARCH_VALUE: Final = "true"
 ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
 ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1"
+ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = (
+    "ANTHROPIC_DEFAULT_SONNET_MODEL",
+    "ANTHROPIC_DEFAULT_HAIKU_MODEL",
+    "ANTHROPIC_DEFAULT_OPUS_MODEL",
+    "ANTHROPIC_DEFAULT_FABLE_MODEL",
+)
+OWNED_ENV_KEYS: Final = (
+    ENABLE_TOOL_SEARCH_KEY,
+    ENABLE_GATEWAY_MODEL_DISCOVERY_KEY,
+    ANTHROPIC_BASE_URL_KEY,
+    ANTHROPIC_AUTH_TOKEN_KEY,
+    ANTHROPIC_API_KEY_KEY,
+)
+OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY)
+OWNED_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in OWNED_ENV_KEYS), *OWNED_TOP_LEVEL_KEYS)
+_CREDENTIAL_ENV_KEYS: Final = frozenset((ANTHROPIC_API_KEY_KEY, ANTHROPIC_AUTH_TOKEN_KEY))
+_CREDENTIAL_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in sorted(_CREDENTIAL_ENV_KEYS)), API_KEY_HELPER_KEY)
+_BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}"
+STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts on"
 
 CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
+CLAUDE_CONFIG_DIR_ENV: Final = "CLAUDE_CONFIG_DIR"
 BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
 AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json"
+CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json"
 
 
 @dataclass(frozen=True, slots=True)
@@ -55,6 +89,129 @@ class ClaudeSettingsError(Exception):
     """Raised for any user-actionable failure while reading or writing Claude Code settings."""
 
 
+def claude_settings_path(environ: Mapping[str, str]) -> Path:
+    """The settings.json Claude Code reads: under CLAUDE_CONFIG_DIR when set, else ~/.claude/settings.json."""
+    config_dir: Final = environ.get(CLAUDE_CONFIG_DIR_ENV, "")
+    if not config_dir:
+        return CLAUDE_SETTINGS_PATH
+    return Path(config_dir).expanduser() / "settings.json"
+
+
+def _is_default_settings_file(settings_path: Path) -> bool:
+    return settings_path.resolve() == CLAUDE_SETTINGS_PATH.resolve()
+
+
+def settings_file_owners(settings_path: Path) -> tuple[SettingsFileOwner, ...]:
+    """The commands whose backups guard settings_path: `lite up` and `lite autoroute up` only ever manage the default file."""
+    return SETTINGS_FILE_OWNERS if _is_default_settings_file(settings_path) else ()
+
+
+def configure_state_path(settings_path: Path) -> Path:
+    """The receipt describing settings_path: the default file keeps CONFIGURE_STATE_PATH, and any other file
+    (a CLAUDE_CONFIG_DIR) gets its own beside it, keyed by its resolved path, so two settings files never
+    share one undo record."""
+    if _is_default_settings_file(settings_path):
+        return CONFIGURE_STATE_PATH
+    digest: Final = hashlib.sha256(str(settings_path.resolve()).encode()).hexdigest()
+    return CONFIGURE_STATE_PATH.parent / CONFIGURE_STATE_PATH.stem / f"{digest}.json"
+
+
+@dataclass(frozen=True, slots=True)
+class StaticToken:
+    """A long-lived virtual key, written into env.ANTHROPIC_AUTH_TOKEN."""
+
+    token: str
+
+
+@dataclass(frozen=True, slots=True)
+class ApiKeyHelper:
+    """A `lite auth print-token` command Claude Code runs per request, so a login renews in place."""
+
+    command: str
+
+
+ClaudeCredential: TypeAlias = StaticToken | ApiKeyHelper
+
+
+@dataclass(frozen=True, slots=True)
+class KeepModel:
+    """Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login)."""
+
+
+@dataclass(frozen=True, slots=True)
+class UnpinModel:
+    """Let go of a `model` an earlier configure pinned; one the user set themselves stays."""
+
+
+@dataclass(frozen=True, slots=True)
+class StartOn:
+    """Pin the top-level `model`, the row Claude Code starts on."""
+
+    model: str
+
+
+ModelChoice: TypeAlias = KeepModel | UnpinModel | StartOn
+
+
+class OwnedValue(BaseModel):
+    """What one key held at a moment in time; `present=False` is an absent key, not a null one."""
+
+    model_config = ConfigDict(frozen=True)
+
+    present: bool
+    value: JsonValue = None
+
+
+class ConfigureReceipt(BaseModel):
+    """What `lite configure claude` found and what it owns, keyed by dotted path (`env.X` or a top-level key).
+
+    Ownership moves only by a write: `written` fingerprints the keys some configure changed, at the
+    value it wrote; a repeat configure refreshes a fingerprint only for a key its merge changed and
+    carries the earlier one otherwise, so a key the user edited in between stops matching and is left
+    alone. `previous` is what each key held before configure took it over; a repeat keeps the earlier
+    snapshot while the key still holds our value and snapshots afresh otherwise, so whatever the
+    repeat displaces is what comes back. `endpoints` is the ANTHROPIC_BASE_URL each credential slot
+    was captured beside, so a credential is only ever put back next to the server it was issued for.
+    No fingerprint is a second copy of a token.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    file_existed: bool
+    env_present: bool
+    env_was_object: bool
+    previous: Mapping[str, OwnedValue]
+    written: Mapping[str, str]
+    endpoints: Mapping[str, OwnedValue]
+
+
+@dataclass(frozen=True, slots=True)
+class WithheldCredential:
+    """A credential left removed: captured beside `endpoint`, while the restored file points elsewhere."""
+
+    key: str
+    endpoint: str
+
+
+@dataclass(frozen=True, slots=True)
+class UnconfigureOutcome:
+    """Keys whose value unconfigure changed back, keys the user changed since and so were left as they
+    are, credentials withheld (the receipt is kept for them, so a later unconfigure can finish once the
+    URL points back), and whether no settings file remains."""
+
+    restored: tuple[str, ...]
+    kept: tuple[str, ...]
+    withheld: tuple[WithheldCredential, ...] = ()
+    file_removed: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class _Claim:
+    previous: OwnedValue
+    written: str | None
+    endpoint: OwnedValue | None
+
+
 def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
     try:
         content: Final = path.read_bytes() if path.exists() else b""
@@ -70,29 +227,104 @@ def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
         )
 
 
-def merge_claude_settings(
-    settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
-) -> dict[str, JsonValue]:
-    """Return a new settings dict wired to route Claude Code through the proxy.
+def _env_object(settings: Mapping[str, JsonValue], path: Path) -> Mapping[str, JsonValue]:
+    raw_env: Final = settings.get(ENV_KEY)
+    if raw_env is None:
+        return MappingProxyType({})
+    if not isinstance(raw_env, dict):
+        raise ClaudeSettingsError(
+            f'{path} has a non-object "{ENV_KEY}" value, which this would discard. Fix or remove it, then retry.'
+        )
+    return raw_env
 
-    Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
-    stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
-    token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH
-    defaults to true because Claude Code turns tool search off when
-    ANTHROPIC_BASE_URL is not a first-party Anthropic host, and
-    CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker
-    is filled from the proxy's /v1/models; existing values of both are left
-    alone. Every other key is preserved untouched.
+
+def refuse_while_owned(settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
+    """Refuse while `lite up` or `lite autoroute up` holds a backup it will restore over any write; a
+    purely local check, so commands run it before any login prompt or request."""
+    for owner in owners:
+        if owner.backup_path.exists():
+            raise ClaudeSettingsError(
+                f"`{owner.start_command}` is currently managing {settings_path} (backup at "
+                f"{owner.backup_path}) and will restore it when it stops. "
+                f"Run `{owner.stop_command}` first, then retry."
+            )
+
+
+def _write_target(settings_path: Path) -> Path:
+    """Write through a symlinked settings.json rather than replacing the link, which would silently
+    detach a file symlinked into a dotfiles repo."""
+    try:
+        return settings_path.resolve() if settings_path.is_symlink() else settings_path
+    except OSError as e:
+        raise ClaudeSettingsError(f"Could not resolve {settings_path}: {e}") from e
+
+
+def _stage(path: Path, document: Mapping[str, object]) -> str:
+    try:
+        return stage_private_json(str(path), document)
+    except OSError as e:
+        raise ClaudeSettingsError(f"Could not write {path}: {e}") from e
+
+
+def _land(
+    path: Path,
+    staged: str | None,
+    also_discard: Sequence[str | None] = (),
+    commit: Callable[[str, str], None] = commit_staged_json,
+) -> None:
+    """Commit a staged file to `path`, or remove `path` when nothing is staged for it. The one place a
+    filesystem error becomes a ClaudeSettingsError; on failure the operation's other staged files are
+    discarded, so no temp file holding a token is left behind."""
+    try:
+        if staged is None:
+            path.unlink(missing_ok=True)
+        else:
+            commit(staged, str(path))
+    except OSError as e:
+        for other in also_discard:
+            if other is not None:
+                discard_staged_json(other)
+        raise ClaudeSettingsError(f"Could not {'remove' if staged is None else 'write'} {path}: {e}") from e
+
+
+def merge_claude_settings(
+    settings: Mapping[str, JsonValue],
+    base_url: str,
+    credential: ClaudeCredential,
+    default_model: str | None = None,
+    tier_model: str | None = None,
+) -> Mapping[str, JsonValue]:
+    """Return a new settings mapping wired to route Claude Code through the proxy.
+
+    A StaticToken lands in env.ANTHROPIC_AUTH_TOKEN, an ApiKeyHelper in the top-level apiKeyHelper;
+    the other credential slots are removed either way, since Claude Code given two credentials may
+    send the wrong one. ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their
+    defaults only when missing. `default_model` is the top-level `model`, the row Claude Code starts
+    on; `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one
+    group. Apart from those tier keys, exactly OWNED_PATHS are touched.
     """
     raw_env: Final = settings.get(ENV_KEY, {})
-    base_env: Final = raw_env if isinstance(raw_env, dict) else {}
-    env: Final = {
-        ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
-        ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE,
-        **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY},
-        ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
-    }
-    return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
+    current_env: Final = raw_env if isinstance(raw_env, dict) else {}
+    env: Final = dict(  # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
+        chain(
+            (
+                (ENABLE_TOOL_SEARCH_KEY, ENABLE_TOOL_SEARCH_VALUE),
+                (ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE),
+            ),
+            ((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS),
+            ((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")),),
+            ((ANTHROPIC_AUTH_TOKEN_KEY, credential.token),) if isinstance(credential, StaticToken) else (),
+            ((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None),
+        )
+    )
+    return dict(  # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
+        chain(
+            ((key, value) for key, value in settings.items() if key not in (API_KEY_HELPER_KEY, ENV_KEY)),
+            ((ENV_KEY, env),),
+            ((API_KEY_HELPER_KEY, credential.command),) if isinstance(credential, ApiKeyHelper) else (),
+            ((MODEL_KEY, default_model),) if default_model is not None else (),
+        )
+    )
 
 
 def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
@@ -121,56 +353,273 @@ def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
     return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token"))
 
 
-def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
-    """Persistently point Claude Code at base_url, preserving every unrelated setting.
+def lite_api_key_helper_configured(base_url: str, settings_path: Path) -> bool:
+    """Whether settings_path already carries the apiKeyHelper `lite login --config-claude` writes for base_url.
 
-    Refuses while any owner holds a backup: each restores its backup when it
-    stops, which would silently undo this write.
+    Only an exact match counts: a helper for another proxy, a hand-written one, or
+    settings that cannot be read leave the caller on the env-token path.
     """
-    for owner in owners:
-        if owner.backup_path.exists():
-            raise ClaudeSettingsError(
-                f"`{owner.start_command}` is currently managing {settings_path} (backup at "
-                f"{owner.backup_path}) and will restore it when it stops. "
-                f"Run `{owner.stop_command}` first, then retry."
-            )
-    normalized_base_url: Final = base_url.rstrip("/")
-    api_key_helper: Final = resolve_api_key_helper(normalized_base_url)
-    existing: Final = load_json_or_empty(settings_path)
-    raw_env: Final = existing.get(ENV_KEY)
-    if raw_env is not None and not isinstance(raw_env, dict):
-        raise ClaudeSettingsError(
-            f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. '
-            "Fix or remove it, then retry."
-        )
-    merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper)
-    # os.replace() swaps the symlink itself for a regular file, silently detaching a
-    # settings.json that is symlinked into a dotfiles repo. There is no backup to undo
-    # that here, unlike `lite up`, so write through to the link's target instead.
-    target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path
     try:
-        write_private_json(str(target), merged)
+        configured_helper: Final = load_json_or_empty(settings_path).get(API_KEY_HELPER_KEY)
+        return configured_helper == resolve_api_key_helper(base_url.rstrip("/"))
+    except ClaudeSettingsError:
+        return False
+
+
+def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue:
+    return OwnedValue(present=key in container, value=container.get(key))
+
+
+def _fingerprint(owned: OwnedValue) -> str:
+    return hashlib.sha256(json.dumps(owned.model_dump(mode="json"), sort_keys=True).encode()).hexdigest()
+
+
+def _env(settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    raw_env: Final = settings.get(ENV_KEY)
+    return raw_env if isinstance(raw_env, dict) else MappingProxyType({})
+
+
+def _lookup(settings: Mapping[str, JsonValue], path: str) -> OwnedValue:
+    section, _, key = path.rpartition(".")
+    return _owned(_env(settings) if section else settings, key)
+
+
+def _with_key(container: Mapping[str, JsonValue], key: str, owned: OwnedValue) -> Mapping[str, JsonValue]:
+    return dict(  # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
+        chain(((k, v) for k, v in container.items() if k != key), ((key, owned.value),) if owned.present else ())
+    )
+
+
+def _with(settings: Mapping[str, JsonValue], path: str, owned: OwnedValue) -> Mapping[str, JsonValue]:
+    """`settings` with the key at `path` set (or removed when `owned` is absent); nothing else changes."""
+    section, _, key = path.rpartition(".")
+    if not section:
+        return _with_key(settings, key, owned)
+    return _with_key(settings, section, OwnedValue(present=True, value=_with_key(_env(settings), key, owned)))
+
+
+def _with_all(settings: Mapping[str, JsonValue], updates: Mapping[str, OwnedValue]) -> Mapping[str, JsonValue]:
+    return reduce(lambda acc, item: _with(acc, *item), updates.items(), settings)
+
+
+def _ours(settings: Mapping[str, JsonValue], path: str, receipt: ConfigureReceipt) -> bool:
+    """Whether the key still holds what a configure wrote (a key no configure ever changed is never ours)."""
+    return receipt.written.get(path) == _fingerprint(_lookup(settings, path))
+
+
+def _claim(
+    path: str,
+    current: Mapping[str, JsonValue],
+    merged: Mapping[str, JsonValue],
+    earlier: ConfigureReceipt | None,
+    url_now: OwnedValue,
+) -> _Claim:
+    """What this configure records for one key; see ConfigureReceipt for the rules."""
+    before, after = _lookup(current, path), _lookup(merged, path)
+    carried: Final = earlier if earlier is not None and _ours(current, path, earlier) else None
+    return _Claim(
+        previous=before if carried is None else carried.previous.get(path, before),
+        written=_fingerprint(after) if before != after else (None if earlier is None else earlier.written.get(path)),
+        endpoint=None
+        if path not in _CREDENTIAL_PATHS
+        else (url_now if carried is None else carried.endpoints.get(path, url_now)),
+    )
+
+
+def _receipt(
+    current: Mapping[str, JsonValue],
+    merged: Mapping[str, JsonValue],
+    earlier: ConfigureReceipt | None,
+    file_exists: bool,
+) -> ConfigureReceipt:
+    url_now: Final = _lookup(current, _BASE_URL_PATH)
+    claims: Final = MappingProxyType({path: _claim(path, current, merged, earlier, url_now) for path in OWNED_PATHS})
+    return ConfigureReceipt(
+        file_existed=file_exists if earlier is None else earlier.file_existed,
+        env_present=ENV_KEY in current if earlier is None else earlier.env_present,
+        env_was_object=isinstance(current.get(ENV_KEY), dict) if earlier is None else earlier.env_was_object,
+        previous=MappingProxyType({path: claim.previous for path, claim in claims.items()}),
+        written=MappingProxyType({path: claim.written for path, claim in claims.items() if claim.written is not None}),
+        endpoints=MappingProxyType(
+            {path: claim.endpoint for path, claim in claims.items() if claim.endpoint is not None}
+        ),
+    )
+
+
+def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None:
+    if not state_path.exists():
+        return None
+    try:
+        return ConfigureReceipt.model_validate_json(state_path.read_bytes())
+    except (OSError, ValidationError) as e:
+        raise ClaudeSettingsError(
+            f"{state_path} is not a readable `lite configure claude` receipt ({e}). "
+            "Remove it and edit Claude Code's settings by hand if they still point at the proxy."
+        ) from e
+
+
+def configure_claude_settings(
+    base_url: str,
+    credential: ClaudeCredential,
+    model: ModelChoice,
+    settings_path: Path,
+    state_path: Path,
+    owners: Sequence[SettingsFileOwner],
+    commit: Callable[[str, str], None] = commit_staged_json,
+) -> None:
+    """Persistently route Claude Code through base_url, recording how to undo it.
+
+    Both files are staged before either is committed, so a full disk or a read-only directory fails
+    before anything changes. The two commits are still two renames: a receipt rename that fails
+    discards the staged settings, and a settings rename that fails after the receipt landed puts the
+    earlier receipt back (or removes the new one), so the receipt on disk never describes settings
+    that were not written. `model`: StartOn pins the starting model, UnpinModel lets go of a pin an
+    earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login).
+    """
+    refuse_while_owned(settings_path, owners)
+    current: Final = load_json_or_empty(settings_path)
+    _env_object(current, settings_path)
+    earlier: Final = read_configure_receipt(state_path)
+    existing: Final = (
+        _with(current, MODEL_KEY, earlier.previous[MODEL_KEY])
+        if isinstance(model, UnpinModel) and earlier is not None and _ours(current, MODEL_KEY, earlier)
+        else current
+    )
+    merged: Final = merge_claude_settings(
+        existing, base_url, credential, model.model if isinstance(model, StartOn) else None
+    )
+    receipt: Final = _receipt(current, merged, earlier, settings_path.exists())
+    target: Final = _write_target(settings_path)
+    try:
+        ensure_private_dir(state_path.parent)
     except OSError as e:
-        raise ClaudeSettingsError(f"Could not write {target}: {e}") from e
+        raise ClaudeSettingsError(f"Could not write {state_path}: {e}") from e
+    staged_receipt: Final = _stage(state_path, receipt.model_dump(mode="json"))
+    try:
+        staged_settings: Final = _stage(target, merged)
+    except ClaudeSettingsError:
+        discard_staged_json(staged_receipt)
+        raise
+    _land(state_path, staged_receipt, (staged_settings,), commit)
+    try:
+        _land(target, staged_settings, commit=commit)
+    except ClaudeSettingsError as settings_error:
+        try:
+            _land(state_path, None if earlier is None else _stage(state_path, earlier.model_dump(mode="json")))
+        except ClaudeSettingsError as receipt_error:
+            raise ClaudeSettingsError(
+                f"{settings_error} The receipt at {state_path} now describes settings that were not written and "
+                f"could not be put back either ({receipt_error}); remove it before retrying."
+            ) from settings_error
+        raise
+
+
+def _endpoint_text(endpoint: OwnedValue) -> str:
+    if not endpoint.present:
+        return f"no {ANTHROPIC_BASE_URL_KEY} (Anthropic's default endpoint)"
+    return endpoint.value if isinstance(endpoint.value, str) else json.dumps(endpoint.value)
+
+
+def unconfigure_claude_settings(
+    settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner]
+) -> UnconfigureOutcome:
+    """Undo `lite configure claude`: put back every key still holding what configure wrote, leave the
+    rest alone, and withhold a credential the restored file would send to a different server than it
+    was issued for (the receipt stays, owning only those slots, so a later unconfigure can finish)."""
+    refuse_while_owned(settings_path, owners)
+    receipt: Final = read_configure_receipt(state_path)
+    if receipt is None:
+        raise ClaudeSettingsError(
+            f"Claude Code is not configured by `lite configure claude` (no receipt at {state_path}); nothing to undo."
+        )
+    current: Final = load_json_or_empty(settings_path)
+    _env_object(current, settings_path)
+    ours: Final = tuple(path for path in receipt.written if _ours(current, path, receipt))
+    kept: Final = tuple(path for path in receipt.written if path not in ours and _lookup(current, path).present)
+    put_back: Final = _with_all(current, MappingProxyType({path: receipt.previous[path] for path in ours}))
+    url_after: Final = _lookup(put_back, _BASE_URL_PATH)
+    withheld: Final = tuple(
+        WithheldCredential(path, _endpoint_text(receipt.endpoints[path]))
+        for path in _CREDENTIAL_PATHS
+        if path in ours and receipt.previous[path].present and receipt.endpoints[path] != url_after
+    )
+    absent: Final = OwnedValue(present=False)
+    trimmed: Final = _with_all(put_back, MappingProxyType({item.key: absent for item in withheld}))
+    settings: Final = (
+        trimmed
+        if _env(trimmed) or receipt.env_was_object
+        else _with_key(trimmed, ENV_KEY, OwnedValue(present=receipt.env_present, value=None))
+    )
+    target: Final = _write_target(settings_path)
+    file_removed: Final = not settings and not (receipt.file_existed and target.exists())
+    kept_receipt: Final = (  # mutable-ok: pydantic serializes the update as given and rejects a mappingproxy
+        receipt.model_copy(update={"written": {item.key: _fingerprint(absent) for item in withheld}})
+        if withheld
+        else None
+    )
+    staged_settings: Final = None if file_removed else _stage(target, settings)
+    try:
+        staged_receipt: Final = (
+            None if kept_receipt is None else _stage(state_path, kept_receipt.model_dump(mode="json"))
+        )
+    except ClaudeSettingsError:
+        if staged_settings is not None:
+            discard_staged_json(staged_settings)
+        raise
+    _land(target, staged_settings, (staged_receipt,))
+    _land(state_path, staged_receipt)
+    return UnconfigureOutcome(
+        restored=tuple(path for path in ours if _lookup(current, path) != _lookup(settings, path)),
+        kept=kept,
+        withheld=withheld,
+        file_removed=file_removed,
+    )
 
 
 __all__ = (
     "ANTHROPIC_API_KEY_KEY",
+    "ANTHROPIC_AUTH_TOKEN_KEY",
     "ANTHROPIC_BASE_URL_KEY",
+    "ANTHROPIC_DEFAULT_MODEL_ENV_KEYS",
     "API_KEY_HELPER_KEY",
     "AUTOROUTE_BACKUP_PATH",
     "BACKUP_PATH",
+    "CLAUDE_CONFIG_DIR_ENV",
     "CLAUDE_SETTINGS_PATH",
+    "CONFIGURE_STATE_PATH",
     "ENABLE_GATEWAY_MODEL_DISCOVERY_KEY",
     "ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE",
     "ENABLE_TOOL_SEARCH_KEY",
     "ENABLE_TOOL_SEARCH_VALUE",
     "ENV_KEY",
+    "MODEL_KEY",
+    "OWNED_ENV_KEYS",
+    "OWNED_PATHS",
+    "OWNED_TOP_LEVEL_KEYS",
     "SETTINGS_FILE_OWNERS",
+    "STARTING_MODEL_ROLE",
+    "ApiKeyHelper",
+    "ClaudeCredential",
     "ClaudeSettingsError",
+    "ConfigureReceipt",
+    "KeepModel",
+    "ModelChoice",
+    "OwnedValue",
     "SettingsFileOwner",
+    "StartOn",
+    "StaticToken",
+    "UnconfigureOutcome",
+    "UnpinModel",
+    "WithheldCredential",
+    "claude_settings_path",
+    "configure_claude_settings",
+    "configure_state_path",
+    "lite_api_key_helper_configured",
     "load_json_or_empty",
     "merge_claude_settings",
+    "read_configure_receipt",
+    "refuse_while_owned",
     "resolve_api_key_helper",
-    "write_claude_settings",
+    "settings_file_owners",
+    "unconfigure_claude_settings",
 )

litellm/proxy/client/cli/commands/configure.py

@@ -0,0 +1,262 @@
+"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable."""
+
+import os
+import re
+import sys
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Final
+
+import click
+from InquirerPy import inquirer
+from InquirerPy.base.control import Choice
+
+from .auth import CliContextObj, context_secret_vault, get_stored_api_key
+from .claude_settings import (
+    STARTING_MODEL_ROLE,
+    ApiKeyHelper,
+    ClaudeCredential,
+    ClaudeSettingsError,
+    ModelChoice,
+    StartOn,
+    StaticToken,
+    UnconfigureOutcome,
+    UnpinModel,
+    claude_settings_path,
+    configure_claude_settings,
+    configure_state_path,
+    refuse_while_owned,
+    resolve_api_key_helper,
+    settings_file_owners,
+    unconfigure_claude_settings,
+)
+from .pi import ListingFailure, PiSyncError, fetch_model_ids
+from .up import ensure_fresh_login
+
+_LISTED_MODELS_SHOWN: Final = 20
+_CLAUDE_TARGET: Final = "claude"
+_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),)
+_KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default"
+_CLAUDE_CODE_PICKER_FILTER: Final = re.compile(r"claude|anthropic", re.IGNORECASE)
+_MODEL_OPTION_HELP: Final = (
+    f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, "
+    "Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude "
+    "Code's sub-agent or background tiers; `lite autoroute up` is the mode that does."
+)
+
+
+def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, str]:
+    """The credential to write and the key to check the proxy with.
+
+    An explicit key (--api-key, `lite --api-key`, LITELLM_PROXY_API_KEY) is long-lived and goes
+    into settings.json as a static token. Without one, the stored `lite login` credential is used
+    the way `lite login --config-claude` uses it, through apiKeyHelper, since it expires within a
+    day and renews in place there; a missing or stale login is refreshed first, as `lite up` does.
+    """
+    ctx_obj: Final[CliContextObj] = ctx.obj
+    explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key"))
+    if explicit:
+        return StaticToken(explicit), explicit
+    base_url: Final = ctx_obj["base_url"]
+    ensure_fresh_login(ctx)
+    stored: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx))
+    if not stored:
+        raise ClaudeSettingsError("Login did not produce a usable token.")
+    return ApiKeyHelper(resolve_api_key_helper(base_url)), stored
+
+
+def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, tuple[str, ...]]:
+    """Every configure path begins the same way: the local ownership check first, so a `lite up`
+    session is refused before any login prompt or request, then the credential, then the listing."""
+    settings_path: Final = claude_settings_path(os.environ)
+    try:
+        refuse_while_owned(settings_path, settings_file_owners(settings_path))
+        credential, key = resolve_credential(ctx, api_key)
+    except ClaudeSettingsError as e:
+        raise click.ClickException(str(e))
+    return credential, _listed_models(ctx.obj["base_url"], key)
+
+
+def _listing_error(base_url: str, error: PiSyncError) -> str:
+    """The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question."""
+    if error.kind is ListingFailure.REJECTED:
+        return f"LiteLLM rejected your key (HTTP {error.status}). Run `lite login` to refresh it, or pass a valid --api-key."
+    if error.kind is ListingFailure.UNREACHABLE:
+        return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?"
+    if error.kind is ListingFailure.EMPTY:
+        return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model."
+    return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy."
+
+
+def _listed_models(base_url: str, key: str) -> tuple[str, ...]:
+    listed: Final = fetch_model_ids(base_url, key)
+    if isinstance(listed, PiSyncError):
+        raise click.ClickException(_listing_error(base_url, listed))
+    return listed
+
+
+def _model_choice(model: str | None) -> ModelChoice:
+    return StartOn(model) if model is not None else UnpinModel()
+
+
+def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listed: Sequence[str], model: str | None) -> None:
+    ctx_obj: Final[CliContextObj] = ctx.obj
+    base_url: Final = ctx_obj["base_url"]
+    if model is not None and model not in listed:
+        shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN])
+        more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else ""
+        raise click.ClickException(
+            f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}."
+        )
+    settings_path: Final = claude_settings_path(os.environ)
+    try:
+        configure_claude_settings(
+            base_url,
+            credential,
+            _model_choice(model),
+            settings_path,
+            configure_state_path(settings_path),
+            settings_file_owners(settings_path),
+        )
+    except ClaudeSettingsError as e:
+        raise click.ClickException(str(e))
+    in_picker: Final = sum(1 for listed_model in listed if _CLAUDE_CODE_PICKER_FILTER.search(listed_model))
+    click.echo(f"Configured Claude Code: {settings_path} now routes through {base_url}.")
+    click.echo(
+        "Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN."
+        if isinstance(credential, StaticToken)
+        else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it."
+    )
+    click.echo(
+        f"Starting model: {model} ({STARTING_MODEL_ROLE}); switch any time with /model."
+        if model is not None
+        else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or "
+        "pass --model to start on a proxy model."
+    )
+    click.echo(
+        f"/model will list {in_picker} of the proxy's {len(listed)} models (Claude Code shows only ids containing "
+        "'claude' or 'anthropic')."
+    )
+    click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.")
+    if isinstance(credential, StaticToken) and settings_path.is_symlink():
+        click.echo(
+            f"Note: {settings_path} is a symlink to {settings_path.resolve()}, so your key now lives in "
+            "that file; keep it out of version control.",
+            err=True,
+        )
+
+
+def _pick_targets() -> tuple[str, ...]:
+    picked: Final = inquirer.checkbox(
+        message="Which agents should route through LiteLLM?",
+        choices=[Choice(value, name=label, enabled=True) for value, label in _TARGETS],
+        validate=lambda chosen: len(chosen) > 0,
+        invalid_message="Pick at least one.",
+    ).execute()
+    return tuple(str(value) for value in picked)
+
+
+def _pick_model(listed: Sequence[str]) -> str | None:
+    picked: Final = inquirer.fuzzy(
+        message="Model Claude Code starts on (type to filter; /model switches any time):",
+        choices=[_KEEP_DEFAULT_MODEL, *listed],
+    ).execute()
+    return None if picked == _KEEP_DEFAULT_MODEL else str(picked)
+
+
+def interactive_configure(
+    ctx: click.Context,
+    pick_targets: Callable[[], tuple[str, ...]] = _pick_targets,
+    pick_model: Callable[[Sequence[str]], str | None] = _pick_model,
+) -> None:
+    """`lite configure` with no agent named: ask which agents to wire and which model to pin."""
+    targets: Final = pick_targets()
+    if _CLAUDE_TARGET not in targets:
+        return
+    credential, listed = _start(ctx, None)
+    _apply_claude(ctx, credential, listed, pick_model(listed))
+
+
+@click.group(name="configure", invoke_without_command=True)
+@click.pass_context
+def configure_group(ctx: click.Context) -> None:
+    """Persistently route a coding agent through your LiteLLM proxy.
+
+    With no agent named, asks which agents to wire and which proxy model to pin.
+    """
+    if ctx.invoked_subcommand is not None:
+        return
+    if not sys.stdin.isatty():
+        raise click.ClickException(
+            "`lite configure` asks questions, so it needs a terminal. Non-interactively, run "
+            "`lite configure claude --api-key <key> --model <model>`."
+        )
+    interactive_configure(ctx)
+
+
+@click.group(name="unconfigure")
+def unconfigure_group() -> None:
+    """Undo `lite configure` for a coding agent."""
+
+
+@configure_group.command(name="claude")
+@click.option(
+    "--api-key",
+    "api_key",
+    default=None,
+    help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / "
+    "LITELLM_PROXY_API_KEY value; with neither, your `lite login` credential is used through apiKeyHelper.",
+)
+@click.option("--model", default=None, help=_MODEL_OPTION_HELP)
+@click.pass_context
+def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None:
+    """Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`.
+
+    Patches ~/.claude/settings.json in place: the proxy URL, your credential (a virtual key as a
+    static token, or your `lite login` through apiKeyHelper), and gateway model discovery so
+    /model lists the proxy's models; --model picks the one Claude Code starts on. Every other
+    setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back.
+    Assumes the proxy is already running.
+    """
+    credential, listed = _start(ctx, api_key)
+    _apply_claude(ctx, credential, listed, model)
+
+
+@unconfigure_group.command(name="claude")
+def unconfigure_claude() -> None:
+    """Return Claude Code's settings to what they were before `lite configure claude`.
+
+    Also undoes `lite login --config-claude`. Only keys still holding what configure wrote are
+    put back; anything you changed since is left as it is and named in the output.
+    """
+    settings_path: Final = claude_settings_path(os.environ)
+    state_path: Final = configure_state_path(settings_path)
+    try:
+        outcome: Final = unconfigure_claude_settings(settings_path, state_path, settings_file_owners(settings_path))
+    except ClaudeSettingsError as e:
+        raise click.ClickException(str(e))
+    _report_unconfigure(settings_path, state_path, outcome)
+
+
+def _report_unconfigure(settings_path: Path, state_path: Path, outcome: UnconfigureOutcome) -> None:
+    """Say what unconfigure did, naming only keys whose value it changed."""
+    if outcome.file_removed:
+        click.echo(
+            f"No settings file remains at {settings_path}; it held nothing but `lite configure claude`'s own keys."
+        )
+    elif outcome.restored:
+        click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.")
+    else:
+        click.echo(f"Nothing in {settings_path} was still ours to restore.")
+    if outcome.kept:
+        click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.")
+    if outcome.withheld:
+        click.echo(
+            "Left removed, since the file now points at a different server than they were issued for: "
+            + "; ".join(f"{item.key} (captured with {item.endpoint})" for item in outcome.withheld)
+            + f". They stay in {state_path}: point env.ANTHROPIC_BASE_URL back and run `lite unconfigure claude` "
+            "again to put them back, or delete that file to drop them."
+        )
+
+
+__all__ = ("configure_group", "interactive_configure", "resolve_credential", "unconfigure_group")

litellm/proxy/client/cli/commands/pi.py

@@ -10,6 +10,7 @@ import os
 import tempfile
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
+from enum import StrEnum
 from pathlib import Path
 from types import MappingProxyType
 from typing import Final
@@ -20,11 +21,28 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
 PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR"
 PI_PROVIDER_NAME: Final = "litellm"
 LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
+_REJECTED_STATUSES: Final = frozenset((401, 403))
+
+
+class ListingFailure(StrEnum):
+    """Why a proxy could not be listed, decided once where the HTTP outcome is classified.
+
+    `unreachable` means no response at all; the other kinds prove the proxy answered, so callers
+    must not suggest checking whether it is running.
+    """
+
+    UNREACHABLE = "unreachable"
+    REJECTED = "rejected"
+    BAD_BODY = "bad_body"
+    EMPTY = "empty"
+    OTHER = "other"
 
 
 @dataclass(frozen=True, slots=True)
 class PiSyncError:
     message: str
+    status: int | None = None
+    kind: ListingFailure | None = None
 
 
 @dataclass(frozen=True, slots=True)
@@ -65,16 +83,20 @@ def fetch_model_ids(
             timeout=10,
         )
     except requests.RequestException as e:
-        return PiSyncError(f"Could not list models from the proxy: {e}")
+        return PiSyncError(f"Could not list models from the proxy: {e}", kind=ListingFailure.UNREACHABLE)
     if resp.status_code != 200:
-        return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.")
+        return PiSyncError(
+            f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot list models.",
+            resp.status_code,
+            ListingFailure.REJECTED if resp.status_code in _REJECTED_STATUSES else ListingFailure.OTHER,
+        )
     try:
         listing: Final = _ModelList.model_validate(resp.json())
     except (ValueError, ValidationError) as e:
-        return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}")
+        return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY)
     ids: Final = tuple(dict.fromkeys(model.id for model in listing.data))
     if not ids:
-        return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.")
+        return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY)
     return ids
 
 
@@ -200,6 +222,7 @@ __all__ = (
     "LITELLM_PROXY_API_KEY_ENV",
     "PI_CONFIG_DIR_ENV",
     "PI_PROVIDER_NAME",
+    "ListingFailure",
     "ModelLimits",
     "PiSyncError",
     "fetch_model_ids",

litellm/proxy/client/cli/commands/up.py

@@ -23,6 +23,7 @@ from .auth import CliContextObj, context_secret_vault, get_stored_api_key, load_
 from .claude_settings import (
     BACKUP_PATH,
     CLAUDE_SETTINGS_PATH,
+    ApiKeyHelper,
     ClaudeSettingsError,
     load_json_or_empty,
     merge_claude_settings,
@@ -123,7 +124,7 @@ def _stored_login_is_pkce(vault: SecretVault) -> bool:
     return token_data is not None and token_data.get("refresh_token") is not None
 
 
-def _ensure_fresh_login(ctx: click.Context) -> None:
+def ensure_fresh_login(ctx: click.Context) -> None:
     ctx_obj: Final[CliContextObj] = ctx.obj
     base_url: Final = ctx_obj["base_url"].rstrip("/")
     vault: Final = context_secret_vault(ctx)
@@ -141,7 +142,7 @@ def _ensure_fresh_login(ctx: click.Context) -> None:
     click.echo("No fresh LiteLLM login found for this proxy; starting login...")
     ctx.invoke(login, pkce=pkce)
     if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault):
-        raise UpError("Login did not produce a usable token; cannot start `lite up`.")
+        raise UpError("Login did not produce a usable token.")
 
 
 def _restore_and_report() -> None:
@@ -169,7 +170,7 @@ def up(ctx: click.Context) -> None:
     base_url: Final = ctx.obj["base_url"]
 
     try:
-        _ensure_fresh_login(ctx)
+        ensure_fresh_login(ctx)
         api_key: Final = resolve_api_key(ctx)
         verify_proxy_key(base_url, api_key)
 
@@ -190,7 +191,7 @@ def up(ctx: click.Context) -> None:
         )
 
         CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True)
-        merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper)
+        merged: Final = merge_claude_settings(original_settings, base_url, ApiKeyHelper(api_key_helper))
         with open(CLAUDE_SETTINGS_PATH, "w") as f:
             json.dump(merged, f, indent=2)
     except (AgentRunError, ClaudeSettingsError) as e:

litellm/proxy/client/cli/main.py

@@ -13,6 +13,7 @@ from .commands.auth import auth_group, context_secret_vault, get_stored_api_key,
 from .commands.autoroute.commands import autoroute_group
 from .commands.chat import chat
 from .commands.config import config_commands, get_config_value, hidden_command_names
+from .commands.configure import configure_group, unconfigure_group
 from .commands.credentials import credentials
 from .commands.debug import debug
 from .commands.encryption import encryption
@@ -162,6 +163,9 @@ cli.add_command(model_groups)
 # Add the autoroute command group (QA auto-routing against your real proxy)
 cli.add_command(autoroute_group, name="autoroute")
 cli.add_command(config_commands)
+# Add configure/unconfigure (persistently wire a coding agent to the proxy with a virtual key)
+cli.add_command(configure_group)
+cli.add_command(unconfigure_group)
 
 
 if __name__ == "__main__":

litellm/proxy/common_request_processing.py

@@ -3300,9 +3300,10 @@ class ProxyBaseLLMRequestProcessing:
         has completed.
 
         Guardrails routed through unified_guardrail are skipped, since they already ran
-        via its streaming iterator.  Guardrails that override
-        async_post_call_success_hook directly run here, including those that implement
-        apply_guardrail but keep their native lifecycle hooks.
+        via its streaming iterator, and so are guardrails a post_call policy pipeline
+        manages, since the pipeline ran them against the buffered stream.  Guardrails
+        that override async_post_call_success_hook directly run here, including those
+        that implement apply_guardrail but keep their native lifecycle hooks.
 
         This is audit-only — content has already been delivered to the client.
 
@@ -3312,12 +3313,18 @@ class ProxyBaseLLMRequestProcessing:
         _response = assembled_response
         try:
             from litellm.proxy.proxy_server import llm_router as _global_llm_router
-            from litellm.proxy.utils import _check_and_merge_model_level_guardrails
+            from litellm.proxy.utils import (
+                _check_and_merge_model_level_guardrails,
+                stream_gated_guardrail_names,
+            )
 
             guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router)
+            stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict)
             for cb in litellm.callbacks:
                 if not isinstance(cb, CustomGuardrail):
                     continue
+                if cb.guardrail_name in stream_gated:
+                    continue
                 if not cb.should_run_guardrail(
                     data=guardrail_data,
                     event_type=GuardrailEventHooks.post_call,

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -44,7 +44,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM
 from litellm.llms.base_llm.guardrail_translation.utils import (
     effective_scan_only_tool_results_for_guardrail,
 )
-from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token
+from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
@@ -917,7 +917,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
                 source,
             )
             return BedrockGuardrailResponse()
-        credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key))
+        credentials, aws_region_name = await run_aws_signing(
+            self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
+        )
         allow_chunking: Final = not self._content_uses_contextual_grounding(content)
 
         completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = []  # mutable-ok: billed-chunk usage accumulator
@@ -1178,7 +1180,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             **base_request_data,
             "content": content,
         }  # mutable-ok: outbound JSON request body
-        prepared_request: Final = self._prepare_request(
+        prepared_request: Final = await run_aws_signing(
+            self._prepare_request,
             credentials=credentials,
             data=bedrock_request_data,
             optional_params=self.optional_params,
@@ -1875,10 +1878,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             return BedrockGuardrailResponse()
 
         api_key: Final[str | None] = request_data.get("api_key") if request_data else None
-        credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key))
+        credentials, aws_region_name = await run_aws_signing(
+            self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
+        )
         body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
 
-        prepared_request: Final = self._prepare_request(
+        prepared_request: Final = await run_aws_signing(
+            self._prepare_request,
             credentials=credentials,
             data=body,
             optional_params=self.optional_params,

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -31,6 +31,7 @@ from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_str_from_messages,
 )
+from litellm.litellm_core_utils.token_counter import offload_token_count
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.auth_utils import (
     ESTIMATED_OUTPUT_TOKENS_FIELD,
@@ -3307,7 +3308,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             min_configured_tpm_limit=min_configured_otpm_limit,
             call_type=call_type,
         )
-        raw_estimated_input_tokens: Final = self._estimate_precise_input_tokens(
+        raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)(
             data=data, model=requested_model, call_type=call_type
         )
         estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1)

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -4559,6 +4559,23 @@ async def delete_verification_tokens(
                 litellm_changed_by=litellm_changed_by,
             )
 
+            # Snapshot before the delete: the FK cascade drops the mapping rows, but their
+            # cached jwt_key_mapping entries still resolve to the now-dead token (LIT-5380).
+            jwt_mapping_cache_keys: Final[tuple[str, ...]] = tuple(
+                cache_key
+                for keys_for_token in await asyncio.gather(
+                    *(
+                        get_jwt_key_mapping_cache_keys_for_token(
+                            hashed_token=key.token,
+                            prisma_client=prisma_client,
+                        )
+                        for key in authorized_keys
+                        if key.token is not None
+                    )
+                )
+                for cache_key in keys_for_token
+            )
+
             if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
                 deleted_tokens = await prisma_client.delete_data(tokens=tokens)
                 if deleted_tokens is not None and len(deleted_tokens) != len(tokens):
@@ -4571,6 +4588,8 @@ async def delete_verification_tokens(
                 if len(deleted_tokens) != len(tokens):
                     failed_tokens = [token for token in tokens if token not in deleted_tokens]
 
+            await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
+
         else:
             raise Exception("DB not connected. prisma_client is None")
     except Exception as e:

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -15,6 +15,7 @@ import os
 import re
 from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
 from dataclasses import dataclass
+from functools import partial
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast
 
@@ -1099,13 +1100,6 @@ async def bedrock_proxy_route(
     """
     create_request_copy(request)
 
-    try:
-        from botocore.auth import SigV4Auth
-        from botocore.awsrequest import AWSRequest
-        from botocore.credentials import Credentials
-    except ImportError:
-        raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
-
     aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME")
     if not _is_bedrock_agent_runtime_route(endpoint=endpoint):
         return await bedrock_llm_proxy_route(
@@ -1136,20 +1130,24 @@ async def bedrock_proxy_route(
     )
 
     # Add or update query parameters
+    from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
     from litellm.llms.bedrock.chat import BedrockConverseLLM
 
     bedrock_llm: Final = BedrockConverseLLM()
-    credentials: Final[Credentials] = bedrock_llm.get_credentials()
-    sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name)
-    headers: Final = {"Content-Type": "application/json"}
     # Assuming the body contains JSON data, parse it
     try:
         data: Final = await _json_request_body(request)
     except Exception as e:
         raise HTTPException(status_code=400, detail={"error": e})
-    _request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers)
-    sigv4.add_auth(_request)
-    prepped: Final = _request.prepare()
+    prepped: Final = await run_aws_signing(
+        sign_aws_json_post,
+        get_credentials=bedrock_llm.get_credentials,
+        service_name="bedrock",
+        aws_region_name=aws_region_name,
+        url=str(updated_url),
+        body=json.dumps(data),
+        headers=MappingProxyType({"Content-Type": "application/json"}),
+    )
 
     ## check for streaming
     is_streaming_request = False
@@ -1207,13 +1205,6 @@ async def comprehend_medical_proxy_route(
 
     [Docs](https://docs.litellm.ai/docs/pass_through/comprehend_medical)
     """
-    try:
-        from botocore.auth import SigV4Auth
-        from botocore.awsrequest import AWSRequest
-        from botocore.credentials import Credentials
-    except ImportError:
-        raise ImportError("Missing boto3 to call comprehendmedical. Run 'pip install boto3'.")
-
     from .llm_provider_handlers.comprehend_medical_passthrough_logging_handler import (
         COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS,
     )
@@ -1244,20 +1235,23 @@ async def comprehend_medical_proxy_route(
     if "stream" in data:
         raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member")
 
-    from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+    from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post
 
-    credentials: Final[Credentials] = BaseAWSLLM().get_credentials(aws_region_name=aws_region_name)
-    sigv4: Final = SigV4Auth(credentials, "comprehendmedical", aws_region_name)
-    headers: Final = MappingProxyType(
-        {
-            "Content-Type": "application/x-amz-json-1.1",
-            "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}",
-        }
-    )
     target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/"
-    _request: Final = AWSRequest(method="POST", url=target_url, data=json.dumps(data), headers=headers)
-    sigv4.add_auth(_request)
-    prepped: Final = _request.prepare()
+    prepped: Final = await run_aws_signing(
+        sign_aws_json_post,
+        get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name),
+        service_name="comprehendmedical",
+        aws_region_name=aws_region_name,
+        url=target_url,
+        body=json.dumps(data),
+        headers=MappingProxyType(
+            {
+                "Content-Type": "application/x-amz-json-1.1",
+                "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}",
+            }
+        ),
+    )
 
     endpoint_func: Final = create_pass_through_route(
         endpoint=operation,

litellm/proxy/policy_engine/pipeline_executor.py

@@ -70,6 +70,10 @@ def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
     return None if texts is None else tuple(texts)
 
 
+def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]:
+    return tuple(texts or ())
+
+
 def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None:
     return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls)
 
@@ -133,6 +137,94 @@ class _StreamRewriteObserver(CustomGuardrail):
         return outputs
 
 
+class _ScannedTextRecorder(CustomGuardrail):
+    def __init__(self, guardrail_name: str) -> None:
+        super().__init__(guardrail_name=guardrail_name)
+        self.inputs: GenericGuardrailAPIInputs | None = None
+
+    @_logged_by_inner_guardrail
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,  # mutable-ok: matches CustomGuardrail.apply_guardrail
+        input_type: Literal["request", "response"],
+        logging_obj: "LiteLLMLoggingObj | None" = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.inputs = inputs
+        return inputs
+
+
+class _LegacyHookStreamAdapter(CustomGuardrail):
+    """Runs a guardrail that only implements the legacy post-call hook (no unified
+    ``apply_guardrail``, or ``use_native_lifecycle_hooks``) as a streaming pipeline step. The
+    endpoint translation hands it the texts it scanned plus the assembled response under
+    ``request_data["response"]``; the hook gets that response in the shape its route gives
+    non-streaming hooks, an exception it raises ends the stream through the executor's
+    fail/error classification, and the response it hands back, or the one it changed in place
+    and returned ``None`` for, is re-scanned by the same translation so its texts reach the
+    client through the translation's ended-stream write-back. A
+    replacement whose scanned texts do not line up with the originals, or whose tool calls
+    differ from them, is undeliverable, so the executor releases the original chunks. A stream
+    that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as
+    long as the hook left the tool calls alone."""
+
+    def __init__(
+        self,
+        inner: CustomGuardrail,
+        endpoint_translation: "BaseTranslation",
+        user_api_key_dict: "UserAPIKeyAuth",
+    ) -> None:
+        super().__init__(guardrail_name=inner.guardrail_name)
+        self.inner: Final = inner
+        self.endpoint_translation: Final = endpoint_translation
+        self.user_api_key_dict: Final = user_api_key_dict
+
+    def structured_messages_cover_full_request(self) -> bool:
+        return self.inner.structured_messages_cover_full_request()
+
+    @_logged_by_inner_guardrail
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,  # mutable-ok: matches CustomGuardrail.apply_guardrail
+        input_type: Literal["request", "response"],
+        logging_obj: "LiteLLMLoggingObj | None" = None,
+    ) -> GenericGuardrailAPIInputs:
+        hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response"))
+        replacement: Final = await self.inner.async_post_call_success_hook(
+            data=request_data,
+            user_api_key_dict=self.user_api_key_dict,
+            response=hooked,
+        )
+        rewrite: Final = hooked if replacement is None else replacement
+        if rewrite is None:
+            return inputs
+        rescanned: Final = await self._rescan(rewrite, logging_obj)
+        if rescanned is None:
+            raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
+        rewritten: Final = rescanned.get("texts")
+        if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))):
+            raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
+        if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")):
+            raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
+        if not rewritten:
+            return inputs
+        rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
+        return rewritten_inputs
+
+    async def _rescan(
+        self, response: object, logging_obj: "LiteLLMLoggingObj | None"
+    ) -> GenericGuardrailAPIInputs | None:
+        recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown")
+        await self.endpoint_translation.process_output_response(
+            response=response,
+            guardrail_to_apply=recorder,
+            litellm_logging_obj=logging_obj,
+            user_api_key_dict=self.user_api_key_dict,
+        )
+        return recorder.inputs
+
+
 def _prepare_hook_input(
     step: PipelineStep,
     callback: CustomGuardrail,
@@ -300,18 +392,29 @@ class PipelineExecutor:
         endpoint_translation: "BaseTranslation",
         streaming_chunks: list[object],  # mutable-ok: shared buffered-stream chunks the translation rewrites in place
         hook_input: dict[str, object],  # mutable-ok: same request-payload shape as data
-        user_api_key_dict: "UserAPIKeyAuth | None",
+        user_api_key_dict: "UserAPIKeyAuth",
         litellm_logging_obj: "LiteLLMLoggingObj | None",
     ) -> None:
         """Run one streaming post_call step through the endpoint translation, delivering
         text and tool-call rewrites on translations that support ended-stream write-back. A
-        rewrite that cannot reach the client yet (one on a translation without write-back, or
-        one the translation refused with ``UndeliverableStreamRewrite``) is discarded: the
-        buffered chunks go back to the originals and the step passes, so the client gets the
-        stream the merge base sent."""
-        observer: Final = _StreamRewriteObserver(callback)
+        guardrail without the unified interface runs its legacy post-call hook against the
+        assembled response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the
+        client yet (one on a translation without write-back, one that drops or adds a tool call,
+        or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is
+        discarded: the buffered chunks go back to the originals and the step passes, so the
+        client gets the stream the merge base sent, and the guardrail stays out of the
+        applied-guardrails header since its output never reached the client. The response an
+        earlier step's translation stored under ``request_data["response"]`` is dropped first,
+        so this step's hook sees the stream as the steps before it left it."""
+        scanner: Final = (
+            callback
+            if PipelineExecutor.supports_unified_execution(callback)
+            else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict)
+        )
+        observer: Final = _StreamRewriteObserver(scanner)
         deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites
         originals: Final = copy.deepcopy(streaming_chunks)
+        hook_input.pop("response", None)  # rebind-ok: an earlier step's stored response goes so this step's is stored
         try:
             if deliver_rewrites:
                 await endpoint_translation.process_output_streaming_response(
@@ -332,11 +435,12 @@ class PipelineExecutor:
                 )
         except UndeliverableStreamRewrite:
             _release_original_chunks(step.guardrail, streaming_chunks, originals)
-        else:
-            if observer.changed_tool_call_count or (
-                not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
-            ):
-                _release_original_chunks(step.guardrail, streaming_chunks, originals)
+            return
+        if observer.changed_tool_call_count or (
+            not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
+        ):
+            _release_original_chunks(step.guardrail, streaming_chunks, originals)
+            return
         if not callback.records_own_guardrail_information:
             add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)
 
@@ -396,11 +500,11 @@ class PipelineExecutor:
                     if isinstance(response, dict):
                         callback.mark_pre_call_hook_ran(response)
             elif mode == "post_call" and streaming_chunks is not None:
-                if not use_unified or endpoint_translation is None:
+                if endpoint_translation is None:
                     return (
                         "error",
                         None,
-                        f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
+                        f"Guardrail '{step.guardrail}' cannot run on a stream without an endpoint translation",
                         None,
                     )
                 await PipelineExecutor._run_streaming_step(
@@ -456,10 +560,22 @@ class PipelineExecutor:
 
     @staticmethod
     def supports_unified_execution(callback: CustomGuardrail) -> bool:
-        """Whether this guardrail runs through the unified apply_guardrail path,
-        the interface streaming pipeline execution requires."""
+        """Whether this guardrail runs through the unified apply_guardrail path."""
         return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
 
+    @staticmethod
+    def supports_streaming_execution(callback: CustomGuardrail) -> bool:
+        """Whether a streaming pipeline step can run this guardrail against the buffered
+        stream: through the unified path, or through its post-call hook on the assembled
+        response when that hook is its only streaming path. A guardrail with its own
+        streaming iterator hook, or with neither hook, keeps running on its own."""
+        callback_type: Final = type(callback)
+        return PipelineExecutor.supports_unified_execution(callback) or (
+            callback_type.async_post_call_success_hook is not CustomLogger.async_post_call_success_hook
+            and callback_type.async_post_call_streaming_iterator_hook
+            is CustomLogger.async_post_call_streaming_iterator_hook
+        )
+
     @staticmethod
     def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
         """Look up an initialized guardrail callback by name from litellm.callbacks."""

litellm/proxy/proxy_server.py

@@ -73,11 +73,13 @@ from litellm.constants import (
     LITELLM_UI_SESSION_DURATION,
     RUNTIME_UPDATABLE_ROUTER_SETTINGS,
 )
+from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.litellm_logging import (
     _init_custom_logger_compatible_class,
 )
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
+from litellm.litellm_core_utils.token_counter import offload_token_count
 from litellm.proxy._types import (
     UI_TEAM_ID,
     CallbackDelete,
@@ -283,7 +285,6 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
 from litellm.litellm_core_utils.agentic_loop_settings import (
     validated_max_agentic_loops,
 )
-from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type
 from litellm.litellm_core_utils.core_helpers import (
     _get_parent_otel_span_from_kwargs,
@@ -12867,7 +12868,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
             CustomHuggingfaceTokenizer | None,
             model_info.get("custom_tokenizer", None),
         )
-    _tokenizer_used: Final = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer)
+    _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)(
+        model=model_to_use, custom_tokenizer=custom_tokenizer
+    )
 
     tokenizer_used: Final = str(_tokenizer_used["type"])
     system_message: Final = _system_message(system)
@@ -12880,7 +12883,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
     counted_tools: Final = cast(  # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats
         list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None
     )
-    total_tokens: Final = await asyncify(litellm.token_counter)(
+    total_tokens: Final = await offload_token_count(litellm.token_counter)(
         model=model_to_use,
         text=prompt,
         messages=counted_messages,

litellm/proxy/response_polling/background_streaming.py

@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias
 
 from fastapi import Request, Response
 from fastapi.responses import StreamingResponse
+from starlette.types import Message
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
@@ -74,6 +75,20 @@ class _StreamEventParser:
     parse: Callable[[str], _StreamEvent] = staticmethod(json.loads)
 
 
+async def _never_receive() -> Message:
+    await asyncio.Event().wait()
+    raise AssertionError("unreachable")
+
+
+def detach_request_from_client(request: Request) -> Request:
+    """Same scope (headers, parsed body, auth) but a receive() that never yields http.disconnect.
+
+    The polling client closes its connection right after getting the polling id, so the
+    upstream call must not be cancelled by the client-disconnect guards.
+    """
+    return Request(request.scope, _never_receive)
+
+
 async def background_streaming_task(
     polling_id: str,
     data: dict[str, object],
@@ -123,7 +138,7 @@ async def background_streaming_task(
         # Pre-call checks (rate limits, guardrails, budget) were already run
         # before polling ID creation, so skip them here to avoid double-counting.
         response: Final[StreamingResponse] = await processor.base_process_llm_request(
-            request=request,
+            request=detach_request_from_client(request),
             fastapi_response=fastapi_response,
             user_api_key_dict=user_api_key_dict,
             route_type="aresponses",

litellm/proxy/schema.prisma

@@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
     updated_at        DateTime @default(now()) @updatedAt
     updated_by        String?
 
-    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
+    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
 
     @@unique([jwt_claim_name, jwt_claim_value])
     @@index([jwt_claim_name, jwt_claim_value, is_active])

litellm/proxy/utils.py

@@ -101,6 +101,7 @@ from litellm.litellm_core_utils.core_helpers import (
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
+from litellm.litellm_core_utils.token_counter import offload_token_count
 from litellm.llms import load_guardrail_translation_mappings
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.proxy._types import (
@@ -465,7 +466,7 @@ def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipe
     return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps)
 
 
-def _pipeline_managed_guardrail_names(
+def pipeline_managed_guardrail_names(
     data: Mapping[str, object], mode: Literal["pre_call", "post_call"]
 ) -> frozenset[str]:
     return _pipeline_step_guardrail_names(
@@ -528,9 +529,17 @@ def _merge_pipeline_metadata_writes(
         _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
 
 
-def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
+def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool:
     callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
-    return callback is not None and PipelineExecutor.supports_unified_execution(callback)
+    if callback is None:
+        return False
+    if PipelineExecutor.supports_unified_execution(callback):
+        return True
+    return (
+        translation is not None
+        and type(translation).assembles_streamed_response
+        and PipelineExecutor.supports_streaming_execution(callback)
+    )
 
 
 def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
@@ -587,7 +596,7 @@ def _withdraw_deferred_claims(
     outside_by_policy: Final = MappingProxyType(
         {policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred}
     )
-    running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(
+    running_elsewhere: Final = pipeline_managed_guardrail_names(data, "pre_call").union(
         _guardrails_run_standalone_pre_call(data), *outside_by_policy.values()
     )
     withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside)
@@ -662,37 +671,51 @@ def _body_selected_deferrals(
     return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed)
 
 
-def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool:
-    unsupported: Final = tuple(
+def _pipeline_unsupported_streaming_guardrails(
+    pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
+) -> tuple[str, ...]:
+    return tuple(
         dict.fromkeys(
-            step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail)
+            step.guardrail
+            for step in pipeline.steps
+            if not _pipeline_step_supports_streaming(step.guardrail, translation)
         )
     )
+
+
+def _pipeline_is_streamable(
+    policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
+) -> bool:
+    unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation)
     if not unsupported:
         return True
     verbose_proxy_logger.warning(
-        "Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, "
-        "which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s",
+        "Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they "
+        "need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a "
+        "route whose translation assembles the streamed response. The stream skips the pipeline and its "
+        "guardrails run on their own: %s",
         policy_name,
         ", ".join(unsupported),
     )
     return False
 
 
-def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool:
-    return resolve_endpoint_translation(user_api_key_dict, None) is not None
+def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None":
+    resolved: Final = resolve_endpoint_translation(user_api_key_dict, None)
+    return None if resolved is None else resolved[1]
 
 
-def _stream_gated_guardrail_names(
+def stream_gated_guardrail_names(
     request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
 ) -> frozenset[str]:
-    if not _route_supports_streaming_pipelines(user_api_key_dict):
+    translation: Final = _streaming_pipeline_translation(user_api_key_dict)
+    if translation is None:
         return frozenset()
     return _pipeline_step_guardrail_names(
         tuple(
             (policy_name, pipeline)
             for policy_name, pipeline in _post_call_pipelines(request_data)
-            if all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps)
+            if not _pipeline_unsupported_streaming_guardrails(pipeline, translation)
         )
     )
 
@@ -704,16 +727,19 @@ def _streamable_post_call_pipelines(
     The post_call pipelines a streaming response can be gated through.
 
     Streaming pipelines scan the buffered stream through the endpoint guardrail
-    translation of the request route, so every step's guardrail needs the
-    unified apply_guardrail interface and the route needs a translation. A
-    pipeline that cannot be run that way yet is left out and its guardrails
-    run on the stream on their own, the way they did before pipelines ran on
-    streams at all, with a warning naming the pipeline.
+    translation of the request route, so every step's guardrail needs either the
+    unified apply_guardrail interface or, on a route whose translation assembles
+    the streamed response, a post-call hook that is its only streaming path, and
+    the route needs a translation. A pipeline that
+    cannot be run that way yet is left out and its guardrails run on the stream
+    on their own, the way they did before pipelines ran on streams at all, with
+    a warning naming the pipeline.
     """
     post_call_pipelines: Final = _post_call_pipelines(request_data)
     if not post_call_pipelines:
         return ()
-    if not _route_supports_streaming_pipelines(user_api_key_dict):
+    translation: Final = _streaming_pipeline_translation(user_api_key_dict)
+    if translation is None:
         verbose_proxy_logger.warning(
             "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
             "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run "
@@ -725,7 +751,7 @@ def _streamable_post_call_pipelines(
     return tuple(
         (policy_name, pipeline)
         for policy_name, pipeline in post_call_pipelines
-        if _pipeline_is_streamable(policy_name, pipeline)
+        if _pipeline_is_streamable(policy_name, pipeline, translation)
     )
 
 
@@ -2115,7 +2141,7 @@ class ProxyLogging:
             )
 
             # Get pipeline-managed guardrails to skip in normal loop
-            pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call")
+            pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call")
 
             caps: Final = ProxyLogging._callback_capabilities()
             # Skip the per-request callback walk entirely when nothing in
@@ -2880,7 +2906,7 @@ class ProxyLogging:
                 original_exception=original_exception,
             )
 
-        request_data.update(_failure_fields_to_lift(request_data))
+        request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data))
 
         # Remove before callbacks iterate — not serialisable
         request_data.pop("litellm_logging_obj", None)
@@ -3119,7 +3145,7 @@ class ProxyLogging:
         if pipeline_response is not None:
             response = pipeline_response  # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
 
-        pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call")
+        pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call")
         guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
         try:
             # Merge model-level guardrails before checking which guardrails to run
@@ -3435,7 +3461,7 @@ class ProxyLogging:
             _cached_guardrail_data: dict | None = None
             _guardrail_data_computed = False
             pipeline_gated: Final = (
-                _stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
+                stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
             )
 
             for callback in litellm.callbacks:

litellm/responses/litellm_completion_transformation/streaming_iterator.py

@@ -437,14 +437,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
             response_created_event_data["temperature"] = self.responses_api_request["temperature"]
         if "text" in self.responses_api_request:
             response_created_event_data["text"] = self.responses_api_request["text"]
-        if "tool_choice" in self.responses_api_request:
-            # Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format
-            response_created_event_data["tool_choice"] = (
-                LiteLLMCompletionResponsesConfig._transform_tool_choice(self.responses_api_request["tool_choice"])
-                or "auto"
+        response_created_event_data["tool_choice"] = (
+            LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response(
+                self.responses_api_request.get("tool_choice")
             )
-        else:
-            response_created_event_data["tool_choice"] = "auto"
+        )
         if "tools" in self.responses_api_request:
             response_created_event_data["tools"] = self.responses_api_request["tools"]
         else:

litellm/responses/litellm_completion_transformation/transformation.py

@@ -27,8 +27,10 @@ from openai.types.chat.chat_completion_named_tool_choice_param import (
 )
 from openai.types.responses import ResponseFunctionToolCall
 from openai.types.responses.response_create_params import ResponseInputParam
+from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam
+from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam
 from openai.types.responses.tool_param import FunctionToolParam
-from pydantic import TypeAdapter
+from pydantic import TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_logger
@@ -68,6 +70,7 @@ from litellm.types.llms.openai import (
     ResponsesAPIOptionalRequestParams,
     ResponsesAPIResponse,
     ResponsesAPIStatus,
+    ToolChoice,
     ValidChatCompletionMessageContentTypes,
     ValidChatCompletionMessageContentTypesLiteral,
 )
@@ -126,6 +129,7 @@ _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
 _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
 _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]])
 _TEXT_ADAPTER: Final = TypeAdapter(str)
+_RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice)
 
 
 @runtime_checkable
@@ -267,6 +271,27 @@ class LiteLLMCompletionResponsesConfig:
         # Return as-is for unknown formats
         return tool_choice
 
+    @staticmethod
+    def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice:
+        if tool_choice is None:
+            return "auto"
+        try:
+            return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(tool_choice)
+        except ValidationError:
+            return LiteLLMCompletionResponsesConfig._chat_tool_choice_as_responses_api_tool_choice(tool_choice)
+
+    @staticmethod
+    def _chat_tool_choice_as_responses_api_tool_choice(tool_choice: object) -> ToolChoice:
+        match tool_choice, LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice):
+            case {"type": "custom"}, {"function": {"name": str(custom_name)}}:
+                return ToolChoiceCustomParam(type="custom", name=custom_name)
+            case _, {"type": "function", "function": {"name": str(function_name)}}:
+                return ToolChoiceFunctionParam(type="function", name=function_name)
+            case _, "none" | "auto" | "required" as normalized:
+                return normalized
+            case _, _:
+                return "auto"
+
     @staticmethod
     def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool:
         """
@@ -2263,7 +2288,9 @@ class LiteLLMCompletionResponsesConfig:
             ),
             parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False),
             temperature=getattr(chat_completion_response, "temperature", 0),
-            tool_choice=getattr(chat_completion_response, "tool_choice", "auto"),
+            tool_choice=LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response(
+                responses_api_request.get("tool_choice")
+            ),
             tools=getattr(chat_completion_response, "tools", []),
             top_p=getattr(chat_completion_response, "top_p", None),
             max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None),

litellm/responses/streaming_iterator.py

@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti
 
 import httpx
 from openai._streaming import SSEDecoder
+from pydantic import BaseModel, ValidationError
 from typing_extensions import TypeIs
 
 import litellm
@@ -438,18 +439,7 @@ class BaseResponsesAPIStreamingIterator:
         if self._persist_completed_response_before_logging:
             self._persist_completed_response_to_cache(is_async=is_async)
 
-        # Create a copy for logging to avoid modifying the response object that will be returned to the user
-        # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
-        # to chat completion format (prompt_tokens/completion_tokens) for internal logging
-        # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
-        # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
-        logging_response = self.completed_response
-        if self.completed_response is not None and hasattr(self.completed_response, "model_dump"):
-            try:
-                logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump())
-            except Exception:
-                # Fallback to original if serialization fails
-                pass
+        logging_response: Final[object] = _logging_copy(self.completed_response)
         self._restore_provider_response_headers(logging_response)
 
         end_time: Final = datetime.now()
@@ -488,10 +478,10 @@ class BaseResponsesAPIStreamingIterator:
     def _restore_provider_response_headers(self, logging_response: object) -> None:
         """Re-apply the provider's response headers to the copy handed to logging callbacks.
 
-        ``model_validate(model_dump())`` above drops pydantic private attributes, so the
+        ``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the
         ``_hidden_params`` the provider transform set on the nested response are lost. Returns early
-        when that copy fell back to the original event, so logging-only state never lands on the
-        object the caller is iterating.
+        when the event was not a pydantic model and logging got the original, so logging-only state
+        never lands on the object the caller is iterating.
         """
         if logging_response is self.completed_response:
             return
@@ -544,7 +534,7 @@ class BaseResponsesAPIStreamingIterator:
     def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
         if response_obj is None or self.logging_obj is None:
             return
-        usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
+        usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None))
         if usage_obj is None:
             return
         try:
@@ -1293,14 +1283,46 @@ def _add_text_like_part_events(
         )
 
 
+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
+    deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow
+    copies of the event and its nested response still keep the caller's ``usage`` attribute separate."""
+    if not isinstance(event, BaseModel):
+        return event
+    try:
+        return type(event).model_validate(event.model_dump())
+    except Exception:
+        return _detached_shallow_copy(event)
+
+
+def _detached_shallow_copy(event: BaseModel) -> BaseModel:
+    nested: Final[object] = getattr(event, "response", None)
+    if isinstance(nested, BaseModel):
+        return event.model_copy(update={"response": nested.model_copy()})
+    return event.model_copy()
+
+
+def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
+    if isinstance(usage, ResponseAPIUsage):
+        return usage
+    if not isinstance(usage, dict):
+        return None
+    try:
+        return ResponseAPIUsage.model_validate(usage)
+    except ValidationError:
+        return None
+
+
 def _stamp_responses_usage_cost(
     response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
 ) -> None:
     if response_obj is None or logging_obj is None:
         return
-    usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
+    usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None))
     if usage_obj is None:
         return
+    response_obj.usage = usage_obj  # rebind-ok: the stamped cost has to ride on the response the client receives
     if isinstance(getattr(usage_obj, "cost", None), (int, float)):
         return
     try:

litellm/router.py

@@ -67,7 +67,7 @@ from litellm.constants import (
     SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
 )
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.asyncify import asyncify, run_async_function
+from litellm.litellm_core_utils.asyncify import run_async_function
 from litellm.litellm_core_utils.core_helpers import (
     _get_parent_otel_span_from_kwargs,
     coerce_token_limit,
@@ -98,6 +98,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
     mask_credentials_in_payload,
     mask_sensitive_structure,
 )
+from litellm.litellm_core_utils.token_counter import offload_token_count
 from litellm.llms.base_llm.vector_store.transformation import (
     RouterVectorStoreEmbeddingExecutor,
     vector_store_request_metadata,
@@ -12113,7 +12114,7 @@ class Router:
         try:
             if not self._pre_call_checks_need_token_count(model, healthy_deployments):
                 return None
-            return await asyncify(self._count_pre_call_check_tokens)(
+            return await offload_token_count(self._count_pre_call_check_tokens)(
                 messages=cast(list[dict[str, str]] | None, messages),  # cast-ok: forwarded to the sync counter
                 input=cast(str | list | None, input),  # cast-ok: forwarded to the sync counter
                 request_kwargs=request_kwargs,

litellm/router_strategy/complexity_router/complexity_router.py

@@ -2568,14 +2568,14 @@ class ComplexityRouter(CustomLogger):
         """Real-tokenizer count of the resolved messages plus the out-of-band carriers, off the
         event loop; None when counting fails, and the gate then leaves the placement alone."""
         import litellm
-        from litellm.litellm_core_utils.asyncify import asyncify
+        from litellm.litellm_core_utils.token_counter import offload_token_count
 
         out_of_band: Final = self._out_of_band_request_text(request_kwargs)
         try:
-            counted: Final = await asyncify(litellm.token_counter)(
+            counted: Final = await offload_token_count(litellm.token_counter)(
                 messages=cast(list, resolved_messages)  # cast-ok: token_counter only iterates the sequence
             )
-            return counted + (await asyncify(litellm.token_counter)(text=out_of_band) if out_of_band else 0)
+            return counted + (await offload_token_count(litellm.token_counter)(text=out_of_band) if out_of_band else 0)
         except Exception as e:  # noqa: BLE001  # best-effort: an uncountable prompt must not fail the request
             verbose_router_logger.debug("ComplexityRouter: context-window token count failed. Got - %s", e)
             return None

litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py

@@ -37,7 +37,7 @@ Safe to enable globally:
 """
 
 import time
-from collections.abc import Mapping
+from collections.abc import Iterator, Mapping
 from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
 
 import httpx
@@ -48,6 +48,10 @@ from litellm.exceptions import (
     ServiceUnavailableError,
 )
 from litellm.integrations.custom_logger import CustomLogger, Span
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    encrypted_content_of_block,
+    strip_encrypted_reasoning_from_messages,
+)
 from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.router_utils.cooldown_cache import CooldownCacheValue
 from litellm.types.llms.openai import AllMessageValues
@@ -138,15 +142,48 @@ class EncryptedContentAffinityCheck(CustomLogger):
             # If no encoded ID, check if encrypted_content itself is wrapped
             encrypted_content = item.get("encrypted_content")
             if encrypted_content and isinstance(encrypted_content, str):
-                (
-                    model_id,
-                    _,
-                ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content)
+                model_id = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content)
                 if model_id:
                     return model_id
 
         return None
 
+    @staticmethod
+    def _anthropic_content_blocks(messages: object) -> Iterator[Mapping[str, object]]:
+        if not isinstance(messages, list):
+            return iter(())
+        return (
+            cast(Mapping[str, object], block)  # cast-ok: narrowed by isinstance
+            for message in cast(list[object], messages)  # cast-ok: narrowed by isinstance
+            if isinstance(message, Mapping)
+            for content in (cast(Mapping[str, object], message).get("content"),)  # cast-ok: narrowed by isinstance
+            if isinstance(content, list)
+            for block in cast(list[object], content)  # cast-ok: narrowed by isinstance
+            if isinstance(block, Mapping)
+        )
+
+    @staticmethod
+    def _model_id_from_wrapped_encrypted_content(encrypted_content: str) -> str | None:
+        model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content)
+        return model_id or None
+
+    @staticmethod
+    def _extract_model_id_from_anthropic_messages(messages: object) -> str | None:
+        return next(
+            (
+                model_id
+                for block in EncryptedContentAffinityCheck._anthropic_content_blocks(messages)
+                if (encrypted_content := encrypted_content_of_block(block)) is not None
+                if (
+                    model_id := EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(
+                        encrypted_content
+                    )
+                )
+                is not None
+            ),
+            None,
+        )
+
     @staticmethod
     def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
         for deployment in healthy_deployments:
@@ -240,8 +277,9 @@ class EncryptedContentAffinityCheck(CustomLogger):
         parent_otel_span: Span | None = None,
     ) -> list[dict]:
         """
-        If the request ``input`` contains litellm-encoded item IDs, decode the
-        embedded ``model_id`` and pin the request to that deployment. Raises
+        If the request ``input`` contains litellm-encoded item IDs, or its Anthropic
+        ``messages`` replay a bridge-tagged thinking block, decode the embedded
+        ``model_id`` and pin the request to that deployment. Raises
         ``RateLimitError`` / ``ServiceUnavailableError`` when the originating
         deployment is a member of the routed model group but currently unavailable
         and no encryption-boundary peer exists, rather than dispatching a doomed
@@ -270,12 +308,15 @@ class EncryptedContentAffinityCheck(CustomLogger):
             request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True
 
         request_input: Final = request_kwargs.get("input")
-        model_id: Final = self._extract_model_id_from_input(request_input)
+        anthropic_messages: Final = messages or request_kwargs.get("messages")
+        model_id: Final = self._extract_model_id_from_input(
+            request_input
+        ) or self._extract_model_id_from_anthropic_messages(anthropic_messages)
         if not model_id:
             return typed_healthy_deployments
 
         verbose_router_logger.debug(
-            "EncryptedContentAffinityCheck: decoded model_id=%s from input item IDs",
+            "EncryptedContentAffinityCheck: decoded model_id=%s from the request's encrypted content markers",
             model_id,
         )
 
@@ -327,6 +368,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
                 model,
             )
             ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input)
+            strip_encrypted_reasoning_from_messages(anthropic_messages)
             return typed_healthy_deployments
 
         # The origin is a member of the routed group but currently unavailable (cooled down); fail fast

litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py

@@ -21,6 +21,7 @@ import litellm
 from litellm import token_counter
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
+from litellm.litellm_core_utils.token_counter import offload_token_count
 from litellm.types.router import RouterCacheEnum, RouterErrors
 from litellm.utils import get_utc_datetime
 
@@ -466,7 +467,7 @@ async def async_io_token_pre_call_check(
 
     request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
     _model: Final = (deployment.get("litellm_params") or {}).get("model") or ""
-    estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model)
+    estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model)
     max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
 
     dt: Final = get_utc_datetime()

litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py

@@ -14,6 +14,7 @@ from litellm.integrations.anthropic_cache_control_hook import (
     AnthropicCacheControlHook,
 )
 from litellm.integrations.custom_logger import CustomLogger, Span
+from litellm.litellm_core_utils.token_counter import offload_token_count
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import CallTypes, StandardLoggingPayload
 from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt
@@ -61,7 +62,7 @@ class PromptCachingDeploymentCheck(CustomLogger):
         if request_kwargs is not None and request_kwargs.get("_target_order") is not None:
             return healthy_deployments
 
-        if messages is not None and is_prompt_caching_valid_prompt(
+        if messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)(
             messages=messages,
             model=model,
             min_token_count=_get_min_token_count_for_deployments(healthy_deployments),
@@ -139,7 +140,7 @@ class PromptCachingDeploymentCheck(CustomLogger):
             return
 
         ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider
-        if is_prompt_caching_valid_prompt(
+        if await offload_token_count(is_prompt_caching_valid_prompt)(
             model=model,
             messages=cast(list[AllMessageValues], messages),
         ):

litellm/types/router.py

@@ -525,6 +525,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
     input_cost_per_second: float | None
     output_cost_per_second: float | None
     output_cost_per_second_480p: ReadOnly[float | None]
+    output_cost_per_second_720p: ReadOnly[float | None]
     output_cost_per_second_1080p: float | None
     output_cost_per_second_4k: ReadOnly[float | None]
     num_retries: int | None

litellm/types/utils.py

@@ -318,6 +318,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
         float | None
     )  # video_generation tier: key output_cost_per_second_<resolution> (e.g. 1080p, 720p)
     output_cost_per_second_480p: ReadOnly[float | None]
+    output_cost_per_second_720p: ReadOnly[float | None]
     output_cost_per_second_4k: ReadOnly[float | None]
     ocr_cost_per_page: float | None  # for OCR models
     ocr_cost_per_credit: float | None  # for OCR models priced by credit
@@ -3522,6 +3523,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
     output_cost_per_second: float | None = None
     output_cost_per_second_1080p: float | None = None
     output_cost_per_second_480p: float | None = None
+    output_cost_per_second_720p: float | None = None
     output_cost_per_second_4k: float | None = None
     input_cost_per_pixel: float | None = None
     output_cost_per_pixel: float | None = None

litellm/utils.py

@@ -2293,15 +2293,7 @@ def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: st
     dict: A dictionary with the tokenizer and its type.
     """
 
-    try:
-        tokenizer = Tokenizer.from_pretrained(
-            identifier,
-            revision=revision,
-            auth_token=auth_token,
-        )
-    except Exception as e:
-        verbose_logger.error("Error creating pretrained tokenizer: %s. Defaulting to version without 'auth_token'.", e)
-        tokenizer = Tokenizer.from_pretrained(identifier, revision=revision)
+    tokenizer: Final = Tokenizer.from_pretrained(identifier, revision=revision, token=auth_token)
     return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
 
 
@@ -3412,7 +3404,7 @@ def get_optional_params_image_gen(
             non_default_params=non_default_params,
             optional_params=optional_params,
             model=model or "",
-            drop_params=drop_params if drop_params is not None else False,
+            drop_params=litellm.drop_params is True or drop_params is True,
         )
     elif (
         custom_llm_provider == "openai"
@@ -5913,6 +5905,7 @@ def _get_model_info_helper(
                 output_cost_per_second=_model_info.get("output_cost_per_second", None),
                 output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None),
                 output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None),
+                output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None),
                 output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None),
                 output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None),
                 output_cost_per_image=_model_info.get("output_cost_per_image", None),

model_prices_and_context_window.schema.json

@@ -478,6 +478,10 @@
           "type": "number",
           "minimum": 0
         },
+        "output_cost_per_second_720p": {
+          "type": "number",
+          "minimum": 0
+        },
         "output_cost_per_token": {
           "type": "number",
           "minimum": 0,

schema.prisma

@@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
     updated_at        DateTime @default(now()) @updatedAt
     updated_by        String?
 
-    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
+    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
 
     @@unique([jwt_claim_name, jwt_claim_value])
     @@index([jwt_claim_name, jwt_claim_value, is_active])

terraform/provider/CHANGELOG.md

@@ -16,6 +16,7 @@ longer signal it.
 
 ### Added
 
+- **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement
 - **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes
 - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it
 - **user**: New `litellm_user` resource and `litellm_user` / `litellm_users` data sources for managing internal users
@@ -38,12 +39,14 @@ longer signal it.
 
 - **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state
 - **key**: Read now unwraps the `info` envelope `/key/info` actually returns; previously reads mapped nothing back into state, so drift on a key was never detected
+- **key**: Read now picks up `model_rpm_limit`, `model_tpm_limit`, `guardrails`, `tags`, `enforced_params`, `allowed_passthrough_routes`, `rpm_limit_type`, `tpm_limit_type` and `prompts` from `info.metadata`, where the proxy actually stores them; previously they stayed empty in state, so a matching config showed a permanent phantom diff on them and out-of-band changes to them were never detected
 - **key**: Updates no longer send an empty `budget_duration`, which the proxy rejects with a 400; any update to a key without a configured `budget_duration` previously failed outright
 - **key**: A config-supplied `key` value (write-only) is now forwarded to `/key/generate`; previously it was silently dropped and the proxy generated a random key instead
 - **security**: The `litellm_key` data source and `litellm_key_block` resource normalize raw `sk-` keys to their SHA-256 token hash before building request URLs and resource IDs, so plaintext keys no longer land in reverse-proxy access logs, Terraform plan output, or state IDs
 
 ### Changed
 
+- **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten
 - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying
 
 ## [0.4.0] - 2026-08-06

terraform/provider/README.md

@@ -103,9 +103,12 @@ resource "litellm_key" "example_key" {
   permissions          = {
     can_create_keys = "true"
   }
-  model_max_budget     = {
-    "gpt-4" = 50.0
-  }
+  model_max_budget     = jsonencode({
+    "gpt-4" = {
+      budget_limit = 50.0
+      time_period  = "30d"
+    }
+  })
   model_rpm_limit      = {
     "claude-3.5-sonnet" = 30
   }

terraform/provider/docs/resources/key.md

@@ -30,9 +30,12 @@ resource "litellm_key" "example" {
   permissions          = {
     "can_create_keys" = "true"
   }
-  model_max_budget     = {
-    "gpt-4" = 50.0
-  }
+  model_max_budget     = jsonencode({
+    "gpt-4" = {
+      budget_limit = 50.0
+      time_period  = "30d"
+    }
+  })
   model_rpm_limit      = {
     "gpt-3.5-turbo" = 30
   }
@@ -73,7 +76,7 @@ The following arguments are supported:
 
 * `key_alias` - (Optional) Alias for this key. This provides a human-readable identifier for the key.
 
-* `duration` - (Optional) Duration for which this key is valid. This sets an expiration time for the key.
+* `duration` - (Optional) How long the key stays valid, e.g. "30d" or "12h". The proxy stores this as an absolute `expires` timestamp. Changing the value resets the expiry to the time of the update plus the new duration; removing it from the configuration leaves the current expiry in place.
 
 * `aliases` - (Optional) Map of model aliases. This allows you to create custom names for models when using this key.
 
@@ -81,7 +84,7 @@ The following arguments are supported:
 
 * `permissions` - (Optional) Permissions associated with this key. This defines what actions are allowed with this key.
 
-* `model_max_budget` - (Optional) Maximum budget per model. This allows setting different budget limits for each model.
+* `model_max_budget` - (Optional) JSON string of per-model budget config, e.g. `jsonencode({"gpt-4" = {budget_limit = 50.0, time_period = "30d"}})`. Each model maps to an object with `budget_limit` (or `max_budget`), `time_period` (or `budget_duration`), `tpm_limit` and `rpm_limit`.
 
 * `model_rpm_limit` - (Optional) Requests per minute limit per model. This allows setting different RPM limits for each model.
 

terraform/provider/docs/resources/team.md

@@ -14,6 +14,16 @@ resource "litellm_team" "engineering" {
 }
 ```
 
+### Team with a Custom ID
+
+```hcl
+resource "litellm_team" "platform" {
+  team_id    = "platform-team"
+  team_alias = "platform"
+  models     = ["gpt-4-proxy"]
+}
+```
+
 ### Team with Comprehensive Configuration
 
 ```hcl
@@ -92,6 +102,8 @@ resource "litellm_team" "model_dependent_team" {
 
 The following arguments are supported:
 
+* `team_id` - (Optional) A stable, human-readable ID for the team (for example `platform-team`). If omitted, the provider generates a random UUID. Changing this forces a new team to be created.
+
 * `team_alias` - (Required) A human-readable identifier for the team.
 
 * `organization_id` - (Optional) The ID of the organization this team belongs to.
@@ -152,7 +164,7 @@ The following arguments are supported:
 
 In addition to the arguments above, the following attributes are exported:
 
-* `id` - The unique identifier for the team.
+* `id` - The unique identifier for the team, equal to `team_id`.
 
 ## Import
 
@@ -162,7 +174,7 @@ Teams can be imported using the team ID:
 terraform import litellm_team.engineering <team-id>
 ```
 
-Note: The team ID is generated when the team is created and is different from the `team_alias`.
+Note: Unless `team_id` is set, the team ID is generated when the team is created and is different from the `team_alias`.
 
 ## Note on Team Members
 

terraform/provider/litellm/client.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"crypto/tls"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -19,6 +20,20 @@ type Client struct {
 	InsecureSkipVerify bool
 }
 
+type apiError struct {
+	StatusCode int
+	Body       string
+}
+
+func (e *apiError) Error() string {
+	return fmt.Sprintf("API request failed with status code %d: %s", e.StatusCode, e.Body)
+}
+
+func isNotFound(err error) bool {
+	var apiErr *apiError
+	return errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound
+}
+
 func NewClient(apiBase, apiKey string, insecureSkipVerify bool) *Client {
 	tr := &http.Transport{
 		TLSClientConfig: &tls.Config{InsecureSkipVerify: insecureSkipVerify},
@@ -57,6 +72,9 @@ func (c *Client) CreateKey(key *Key) (*Key, error) {
 
 func (c *Client) GetKey(keyID string) (*Key, error) {
 	resp, err := c.sendRequest("GET", fmt.Sprintf("/key/info?key=%s", keyID), nil)
+	if isNotFound(err) {
+		return nil, nil
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -69,32 +87,71 @@ func (c *Client) GetKey(keyID string) (*Key, error) {
 				info["key"] = k
 			}
 		}
+		hoistKeyFieldsStoredInMetadata(info)
 		return c.parseKeyResponse(info)
 	}
 
 	return c.parseKeyResponse(resp)
 }
 
+var keyFieldsStoredInMetadata = []string{
+	"model_rpm_limit",
+	"model_tpm_limit",
+	"guardrails",
+	"tags",
+	"enforced_params",
+	"allowed_passthrough_routes",
+	"rpm_limit_type",
+	"tpm_limit_type",
+	"prompts",
+}
+
+func hoistKeyFieldsStoredInMetadata(info map[string]interface{}) {
+	metadata, ok := info["metadata"].(map[string]interface{})
+	if !ok {
+		return
+	}
+	for _, field := range keyFieldsStoredInMetadata {
+		if existing, present := info[field]; present && existing != nil {
+			continue
+		}
+		if v, present := metadata[field]; present {
+			info[field] = v
+		}
+	}
+}
+
 func (c *Client) UpdateKey(key *Key) (*Key, error) {
 	// Create a new map with only the fields that can be updated
 	updateData := map[string]interface{}{
 		"key":              key.Key,
 		"team_id":          key.TeamID,
-		"metadata":         key.Metadata,
 		"key_alias":        key.KeyAlias,
 		"aliases":          key.Aliases,
 		"permissions":      key.Permissions,
 		"model_max_budget": key.ModelMaxBudget,
-		"model_rpm_limit":  key.ModelRPMLimit,
-		"model_tpm_limit":  key.ModelTPMLimit,
 		"blocked":          key.Blocked,
 	}
 
+	// The proxy keeps the stored metadata only when the field is absent, so nil means omit.
+	if key.Metadata != nil {
+		updateData["metadata"] = key.Metadata
+	}
+	if key.ModelRPMLimit != nil {
+		updateData["model_rpm_limit"] = key.ModelRPMLimit
+	}
+	if key.ModelTPMLimit != nil {
+		updateData["model_tpm_limit"] = key.ModelTPMLimit
+	}
+
 	// The proxy rejects an empty-string budget_duration with a 400, so only
 	// send it when set.
 	if key.BudgetDuration != "" {
 		updateData["budget_duration"] = key.BudgetDuration
 	}
+	if key.Duration != "" {
+		updateData["duration"] = key.Duration
+	}
 
 	// Only add pointer fields if they are explicitly set
 	if key.MaxBudget != nil {
@@ -366,7 +423,7 @@ func (c *Client) sendRequest(method, path string, body interface{}) (map[string]
 	log.Printf("Response body: %s", c.redactSensitiveData(string(bodyBytes)))
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("API request failed with status code %d: %s", resp.StatusCode, string(bodyBytes))
+		return nil, &apiError{StatusCode: resp.StatusCode, Body: string(bodyBytes)}
 	}
 
 	var result map[string]interface{}

terraform/provider/litellm/resource_key.go

@@ -2,7 +2,9 @@ package litellm
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"log"
 
 	"github.com/hashicorp/go-cty/cty"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
@@ -10,7 +12,7 @@ import (
 )
 
 func resourceKey() *schema.Resource {
-	return &schema.Resource{
+	r := &schema.Resource{
 		CreateContext: resourceKeyCreate,
 		ReadContext:   resourceKeyRead,
 		UpdateContext: resourceKeyUpdate,
@@ -18,6 +20,7 @@ func resourceKey() *schema.Resource {
 		Importer: &schema.ResourceImporter{
 			StateContext: schema.ImportStatePassthroughContext,
 		},
+		SchemaVersion: 1,
 		Schema: map[string]*schema.Schema{
 			"key": {
 				Type:      schema.TypeString,
@@ -86,8 +89,9 @@ func resourceKey() *schema.Resource {
 				Optional: true,
 			},
 			"duration": {
-				Type:     schema.TypeString,
-				Optional: true,
+				Type:        schema.TypeString,
+				Optional:    true,
+				Description: "How long the key stays valid, e.g. \"30d\" or \"12h\". Changing it resets the expiry to the time of the update plus the new duration; removing it leaves the current expiry in place",
 			},
 			"aliases": {
 				Type:     schema.TypeMap,
@@ -105,9 +109,11 @@ func resourceKey() *schema.Resource {
 				Elem:     &schema.Schema{Type: schema.TypeString},
 			},
 			"model_max_budget": {
-				Type:     schema.TypeMap,
-				Optional: true,
-				Elem:     &schema.Schema{Type: schema.TypeFloat, Computed: true},
+				Type:             schema.TypeString,
+				Optional:         true,
+				ValidateFunc:     validateKeyModelMaxBudget,
+				DiffSuppressFunc: budgetSuppressEquivalentJSON,
+				Description:      "JSON string of per-model budget config (e.g. '{\"gpt-4o-mini\": {\"budget_limit\": 50, \"time_period\": \"30d\"}}')",
 			},
 			"model_rpm_limit": {
 				Type:     schema.TypeMap,
@@ -182,6 +188,79 @@ func resourceKey() *schema.Resource {
 			},
 		},
 	}
+	r.StateUpgraders = []schema.StateUpgrader{{
+		Version: 0,
+		Type:    resourceKeyV0Type(r.Schema),
+		Upgrade: resourceKeyStateUpgradeV0,
+	}}
+	return r
+}
+
+// Schema version 0 typed model_max_budget as map(number), which the proxy
+// rejects; version 1 stores the per-model BudgetConfig objects as a JSON string.
+func resourceKeyV0Type(current map[string]*schema.Schema) cty.Type {
+	v0 := make(map[string]*schema.Schema, len(current))
+	for k, v := range current {
+		v0[k] = v
+	}
+	v0["model_max_budget"] = &schema.Schema{
+		Type:     schema.TypeMap,
+		Optional: true,
+		Elem:     &schema.Schema{Type: schema.TypeFloat},
+	}
+	return (&schema.Resource{Schema: v0}).CoreConfigSchema().ImpliedType()
+}
+
+func resourceKeyStateUpgradeV0(_ context.Context, rawState map[string]interface{}, _ interface{}) (map[string]interface{}, error) {
+	delete(rawState, "model_max_budget")
+	return rawState, nil
+}
+
+var keyModelBudgetFields = map[string]bool{
+	"budget_limit":    true,
+	"max_budget":      true,
+	"time_period":     true,
+	"budget_duration": true,
+	"tpm_limit":       true,
+	"rpm_limit":       true,
+}
+
+func validateKeyModelMaxBudget(v interface{}, k string) ([]string, []error) {
+	var parsed map[string]json.RawMessage
+	if err := json.Unmarshal([]byte(v.(string)), &parsed); err != nil || parsed == nil {
+		return nil, []error{fmt.Errorf("%q must be a JSON object keyed by model name, got %s", k, v)}
+	}
+	for model, cfg := range parsed {
+		var budget map[string]json.RawMessage
+		if err := json.Unmarshal(cfg, &budget); err != nil || len(budget) == 0 {
+			return nil, []error{fmt.Errorf("%q[%q] must be a budget object such as {\"budget_limit\": 50, \"time_period\": \"30d\"}, got %s", k, model, cfg)}
+		}
+		for field := range budget {
+			if !keyModelBudgetFields[field] {
+				return nil, []error{fmt.Errorf("%q[%q] has unknown budget field %q; supported fields are budget_limit, max_budget, time_period, budget_duration, tpm_limit, rpm_limit", k, model, field)}
+			}
+		}
+	}
+	return nil, nil
+}
+
+func parseKeyModelMaxBudget(raw string) map[string]interface{} {
+	var parsed map[string]interface{}
+	if err := json.Unmarshal([]byte(raw), &parsed); err != nil || parsed == nil {
+		return map[string]interface{}{}
+	}
+	return parsed
+}
+
+func keyModelMaxBudgetJSON(modelMaxBudget map[string]interface{}) string {
+	if len(modelMaxBudget) == 0 {
+		return ""
+	}
+	encoded, err := json.Marshal(modelMaxBudget)
+	if err != nil {
+		return ""
+	}
+	return string(encoded)
 }
 
 func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics {
@@ -219,10 +298,12 @@ func resourceKeyRead(ctx context.Context, d *schema.ResourceData, m interface{})
 	}
 
 	if key == nil {
+		log.Printf("[WARN] Key %s not found, removing from state", d.Id())
 		d.SetId("")
 		return nil
 	}
 
+	key.Metadata = declaredKeyMetadata(key.Metadata, d.Get("metadata").(map[string]interface{}))
 	mapKeyToResourceData(d, key)
 	return nil
 }
@@ -232,15 +313,75 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{
 
 	key := &Key{Key: d.Id()}
 	mapResourceDataToKey(d, key)
+	if !d.HasChange("duration") {
+		key.Duration = ""
+	}
+	key.ModelRPMLimit = changedMap(d, "model_rpm_limit")
+	key.ModelTPMLimit = changedMap(d, "model_tpm_limit")
 
-	_, err := c.UpdateKey(key)
+	metadata, err := plannedKeyMetadata(c, d)
 	if err != nil {
+		d.Partial(true)
+		return diag.FromErr(fmt.Errorf("error updating key: %s", err))
+	}
+	key.Metadata = metadata
+
+	if _, err := c.UpdateKey(key); err != nil {
 		return diag.FromErr(fmt.Errorf("error updating key: %s", err))
 	}
 
 	return resourceKeyRead(ctx, d, m)
 }
 
+func changedMap(d *schema.ResourceData, name string) map[string]interface{} {
+	if !d.HasChange(name) {
+		return nil
+	}
+	return d.Get(name).(map[string]interface{})
+}
+
+func plannedKeyMetadata(c *Client, d *schema.ResourceData) (map[string]interface{}, error) {
+	if !d.HasChange("metadata") {
+		return nil, nil
+	}
+	current, err := c.GetKey(d.Id())
+	if err != nil {
+		return nil, err
+	}
+	if current == nil {
+		return nil, fmt.Errorf("key %s no longer exists", d.Id())
+	}
+	oldDeclared, newDeclared := d.GetChange("metadata")
+	return mergeKeyMetadata(current.Metadata, oldDeclared.(map[string]interface{}), newDeclared.(map[string]interface{})), nil
+}
+
+func declaredKeyMetadata(server, declared map[string]interface{}) map[string]interface{} {
+	if server == nil {
+		return nil
+	}
+	result := make(map[string]interface{}, len(declared))
+	for k := range declared {
+		if v, ok := server[k]; ok {
+			result[k] = v
+		}
+	}
+	return result
+}
+
+func mergeKeyMetadata(server, oldDeclared, newDeclared map[string]interface{}) map[string]interface{} {
+	result := make(map[string]interface{}, len(server)+len(newDeclared))
+	for k, v := range server {
+		result[k] = v
+	}
+	for k := range oldDeclared {
+		delete(result, k)
+	}
+	for k, v := range newDeclared {
+		result[k] = v
+	}
+	return result
+}
+
 func resourceKeyDelete(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics {
 	c := m.(*Client)
 
@@ -285,7 +426,7 @@ func mapResourceDataToKey(d *schema.ResourceData, key *Key) {
 	key.Aliases = d.Get("aliases").(map[string]interface{})
 	key.Config = d.Get("config").(map[string]interface{})
 	key.Permissions = d.Get("permissions").(map[string]interface{})
-	key.ModelMaxBudget = d.Get("model_max_budget").(map[string]interface{})
+	key.ModelMaxBudget = parseKeyModelMaxBudget(d.Get("model_max_budget").(string))
 	key.ModelRPMLimit = d.Get("model_rpm_limit").(map[string]interface{})
 	key.ModelTPMLimit = d.Get("model_tpm_limit").(map[string]interface{})
 	key.Guardrails = expandStringList(d.Get("guardrails").([]interface{}))
@@ -358,9 +499,7 @@ func mapKeyToResourceData(d *schema.ResourceData, key *Key) {
 	if key.Permissions != nil {
 		d.Set("permissions", key.Permissions)
 	}
-	if key.ModelMaxBudget != nil {
-		d.Set("model_max_budget", key.ModelMaxBudget)
-	}
+	d.Set("model_max_budget", keyModelMaxBudgetJSON(key.ModelMaxBudget))
 	if key.ModelRPMLimit != nil {
 		d.Set("model_rpm_limit", key.ModelRPMLimit)
 	}

terraform/provider/litellm/resource_key_test.go

@@ -6,9 +6,11 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"testing"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 )
 
 func newKeyResourceData(t *testing.T, raw map[string]interface{}) *schema.ResourceData {
@@ -193,6 +195,102 @@ func TestCreateKeySendsConfigSuppliedKey(t *testing.T) {
 	}
 }
 
+// The proxy validates each model_max_budget entry as a BudgetConfig object and
+// 500s on a bare number, so the JSON string must reach /key/generate as nested
+// objects and the proxy's response must map back to equivalent JSON in state.
+func TestCreateKeySendsModelMaxBudgetAsBudgetObjects(t *testing.T) {
+	var captured map[string]interface{}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		if r.URL.Path == "/key/generate" {
+			body, _ := io.ReadAll(r.Body)
+			json.Unmarshal(body, &captured)
+			w.Write([]byte(`{"key": "sk-test", "token_id": "hash-1"}`))
+			return
+		}
+		w.Write([]byte(`{"key": "hash-1", "info": {"model_max_budget": {"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d", "rpm_limit": 60}}}}`))
+	}))
+	defer srv.Close()
+
+	client := NewClient(srv.URL, "test-key", true)
+	d := newKeyResourceData(t, map[string]interface{}{
+		"model_max_budget": `{"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d"}}`,
+	})
+
+	if diags := resourceKeyCreate(context.Background(), d, client); diags.HasError() {
+		t.Fatalf("create returned error: %v", diags)
+	}
+
+	budgets, ok := captured["model_max_budget"].(map[string]interface{})
+	if !ok {
+		t.Fatalf("create payload model_max_budget = %v, want object", captured["model_max_budget"])
+	}
+	cfg, ok := budgets["gpt-4o-mini"].(map[string]interface{})
+	if !ok {
+		t.Fatalf("model_max_budget[gpt-4o-mini] = %v, want BudgetConfig object", budgets["gpt-4o-mini"])
+	}
+	if cfg["budget_limit"] != float64(50) || cfg["time_period"] != "30d" {
+		t.Errorf("BudgetConfig = %v, want budget_limit 50 and time_period 30d", cfg)
+	}
+
+	var state map[string]interface{}
+	if err := json.Unmarshal([]byte(d.Get("model_max_budget").(string)), &state); err != nil {
+		t.Fatalf("state model_max_budget %q is not JSON: %v", d.Get("model_max_budget"), err)
+	}
+	if got, _ := state["gpt-4o-mini"].(map[string]interface{}); got["budget_limit"] != float64(50) || got["rpm_limit"] != float64(60) {
+		t.Errorf("state model_max_budget = %v, want the BudgetConfig read back from /key/info", state)
+	}
+}
+
+// Schema version 0 stored model_max_budget as map(number); that state cannot
+// decode into the version 1 string attribute, so the upgrader must drop it.
+func TestKeyStateUpgradeV0DropsMapModelMaxBudget(t *testing.T) {
+	upgraded, err := resourceKey().StateUpgraders[0].Upgrade(context.Background(), map[string]interface{}{
+		"id":               "hash-1",
+		"key_alias":        "legacy",
+		"model_max_budget": map[string]interface{}{"gpt-4o-mini": 50.0},
+	}, nil)
+	if err != nil {
+		t.Fatalf("upgrade returned error: %v", err)
+	}
+	if _, present := upgraded["model_max_budget"]; present {
+		t.Errorf("upgraded state still carries map model_max_budget: %v", upgraded["model_max_budget"])
+	}
+	if upgraded["key_alias"] != "legacy" {
+		t.Errorf("upgrade dropped unrelated attribute: %v", upgraded)
+	}
+}
+
+func TestKeyModelMaxBudgetValidationRequiresBudgetObjects(t *testing.T) {
+	validate := resourceKey().Schema["model_max_budget"].ValidateFunc
+	for _, valid := range []string{
+		`{}`,
+		`{"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d"}}`,
+		`{"gpt-4o-mini": {"max_budget": 50, "rpm_limit": 60}, "gpt-4o": {"budget_duration": "1d", "tpm_limit": 1000}}`,
+	} {
+		if _, errs := validate(valid, "model_max_budget"); len(errs) != 0 {
+			t.Errorf("validate(%s) = %v, want accepted", valid, errs)
+		}
+	}
+	for _, invalid := range []string{
+		`null`,
+		`[]`,
+		`"gpt-4o-mini"`,
+		`50`,
+		`{"gpt-4o-mini": 50}`,
+		`{"gpt-4o-mini": null}`,
+		`{"gpt-4o-mini": [50]}`,
+		`{"gpt-4o-mini": {}}`,
+		`{"gpt-4o-mini": {"budget_limt": 50}}`,
+		`{"gpt-4o-mini": {"budget_limit": 50, "max_tokens": 100}}`,
+		`not json`,
+	} {
+		if _, errs := validate(invalid, "model_max_budget"); len(errs) == 0 {
+			t.Errorf("validate(%s) accepted a value that would send no per-model budget", invalid)
+		}
+	}
+}
+
 // The proxy 400s on budget_duration: "", so an unset duration must be
 // omitted from the update payload entirely.
 func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) {
@@ -221,6 +319,47 @@ func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) {
 	}
 }
 
+func TestResourceKeyUpdateFailureKeepsPriorState(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		if r.URL.Path == "/key/update" {
+			w.WriteHeader(http.StatusBadRequest)
+			w.Write([]byte(`{"error":{"message":"Invalid budget_duration 'bad'"}}`))
+			return
+		}
+		w.Write([]byte(`{"key":"hash-1","info":{"key_alias":"demo","models":["fake-model"]}}`))
+	}))
+	defer srv.Close()
+
+	res := resourceKey()
+	priorData := newKeyResourceData(t, map[string]interface{}{
+		"key_alias": "demo",
+		"models":    []interface{}{"fake-model"},
+	})
+	priorData.SetId("hash-1")
+	prior := priorData.State()
+	config := terraform.NewResourceConfigRaw(map[string]interface{}{
+		"key_alias":       "demo",
+		"models":          []interface{}{"fake-model"},
+		"budget_duration": "bad",
+	})
+	diff, err := res.Diff(context.Background(), prior, config, nil)
+	if err != nil {
+		t.Fatalf("diff failed: %v", err)
+	}
+
+	newState, diags := res.Apply(context.Background(), prior, diff, NewClient(srv.URL, "test-key", true))
+	if !diags.HasError() {
+		t.Fatal("apply succeeded, want the proxy's 400 surfaced as an error")
+	}
+	if got, ok := newState.Attributes["budget_duration"]; ok {
+		t.Errorf("failed update persisted budget_duration=%q into state, want it absent", got)
+	}
+	if newState.Attributes["key_alias"] != "demo" {
+		t.Errorf("prior key_alias lost from state: %v", newState.Attributes)
+	}
+}
+
 // /key/info nests the key's fields under "info"; GetKey must unwrap that
 // envelope or reads map nothing back into state.
 func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) {
@@ -254,3 +393,296 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) {
 		t.Errorf("RPMLimit not parsed: %+v", key.RPMLimit)
 	}
 }
+
+func TestGetKeyReadsFieldsStoredInMetadata(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.Write([]byte(`{
+			"key": "hash-1",
+			"info": {
+				"models": ["gpt-4o-mini"],
+				"metadata": {
+					"team": "core-infra",
+					"model_rpm_limit": {"gpt-4o-mini": 7},
+					"model_tpm_limit": {"gpt-4o-mini": 10000},
+					"guardrails": ["pii-guard"],
+					"tags": ["prod"],
+					"enforced_params": ["user"],
+					"allowed_passthrough_routes": ["/v1/foo"],
+					"rpm_limit_type": "guaranteed_throughput",
+					"tpm_limit_type": "dynamic",
+					"prompts": ["p1"]
+				}
+			}
+		}`))
+	}))
+	defer srv.Close()
+
+	client := NewClient(srv.URL, "test-key", true)
+	key, err := client.GetKey("hash-1")
+	if err != nil {
+		t.Fatalf("GetKey returned error: %v", err)
+	}
+	if got, ok := key.ModelRPMLimit["gpt-4o-mini"].(float64); !ok || got != 7 {
+		t.Errorf("ModelRPMLimit = %v, want gpt-4o-mini=7 read from metadata", key.ModelRPMLimit)
+	}
+	if got, ok := key.ModelTPMLimit["gpt-4o-mini"].(float64); !ok || got != 10000 {
+		t.Errorf("ModelTPMLimit = %v, want gpt-4o-mini=10000 read from metadata", key.ModelTPMLimit)
+	}
+	if len(key.Guardrails) != 1 || key.Guardrails[0] != "pii-guard" {
+		t.Errorf("Guardrails = %v, want [pii-guard]", key.Guardrails)
+	}
+	if len(key.Tags) != 1 || key.Tags[0] != "prod" {
+		t.Errorf("Tags = %v, want [prod]", key.Tags)
+	}
+	if len(key.EnforcedParams) != 1 || key.EnforcedParams[0] != "user" {
+		t.Errorf("EnforcedParams = %v, want [user]", key.EnforcedParams)
+	}
+	if len(key.AllowedPassthroughRoutes) != 1 || key.AllowedPassthroughRoutes[0] != "/v1/foo" {
+		t.Errorf("AllowedPassthroughRoutes = %v, want [/v1/foo]", key.AllowedPassthroughRoutes)
+	}
+	if key.RPMLimitType != "guaranteed_throughput" || key.TPMLimitType != "dynamic" {
+		t.Errorf("limit types = %q/%q, want guaranteed_throughput/dynamic", key.RPMLimitType, key.TPMLimitType)
+	}
+	if len(key.Prompts) != 1 || key.Prompts[0] != "p1" {
+		t.Errorf("Prompts = %v, want [p1]", key.Prompts)
+	}
+	if key.Metadata["team"] != "core-infra" {
+		t.Errorf("Metadata = %v, want team=core-infra preserved", key.Metadata)
+	}
+}
+
+func TestGetKeyPrefersTopLevelOverMetadataCopy(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.Write([]byte(`{
+			"key": "hash-1",
+			"info": {
+				"tags": ["top-level"],
+				"guardrails": null,
+				"metadata": {
+					"tags": ["from-metadata"],
+					"guardrails": ["from-metadata"]
+				}
+			}
+		}`))
+	}))
+	defer srv.Close()
+
+	client := NewClient(srv.URL, "test-key", true)
+	key, err := client.GetKey("hash-1")
+	if err != nil {
+		t.Fatalf("GetKey returned error: %v", err)
+	}
+	if len(key.Tags) != 1 || key.Tags[0] != "top-level" {
+		t.Errorf("Tags = %v, want [top-level]", key.Tags)
+	}
+	if len(key.Guardrails) != 1 || key.Guardrails[0] != "from-metadata" {
+		t.Errorf("Guardrails = %v, want [from-metadata] (null top-level must not shadow)", key.Guardrails)
+	}
+}
+
+func TestResourceKeyReadDropsMissingKeyFromState(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusNotFound)
+		w.Write([]byte(`{"error":{"message":"Key not found in database","type":"not_found_error","param":"key","code":"404"}}`))
+	}))
+	defer srv.Close()
+
+	d := newKeyResourceData(t, map[string]interface{}{"key_alias": "stale"})
+	d.SetId("deleted-out-of-band")
+
+	diags := resourceKeyRead(context.Background(), d, NewClient(srv.URL, "test-key", true))
+	if diags.HasError() {
+		t.Fatalf("read of a missing key must not error, got: %v", diags)
+	}
+	if d.Id() != "" {
+		t.Errorf("Id = %q, want empty so Terraform plans a recreate", d.Id())
+	}
+}
+
+func TestResourceKeyReadStillFailsOnNon404Errors(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(`{"error":{"message":"db down"}}`))
+	}))
+	defer srv.Close()
+
+	d := newKeyResourceData(t, map[string]interface{}{"key_alias": "live"})
+	d.SetId("still-exists")
+
+	diags := resourceKeyRead(context.Background(), d, NewClient(srv.URL, "test-key", true))
+	if !diags.HasError() {
+		t.Fatal("a 500 from /key/info must surface as an error, not be treated as a deleted key")
+	}
+	if d.Id() != "still-exists" {
+		t.Errorf("Id = %q, want unchanged on a transient error", d.Id())
+	}
+}
+
+// fakeKeyProxy serves /key/info from stored metadata and applies /key/update
+// the way the proxy does: an absent "metadata" keeps the stored map, a
+// present one replaces it wholesale.
+type fakeKeyProxy struct {
+	metadata map[string]interface{}
+	updates  []map[string]interface{}
+}
+
+func (p *fakeKeyProxy) handler() http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		switch r.URL.Path {
+		case "/key/info":
+			json.NewEncoder(w).Encode(map[string]interface{}{
+				"key":  "hash-1",
+				"info": map[string]interface{}{"key_alias": "alias-1", "models": []string{"gpt-4o-mini"}, "metadata": p.metadata},
+			})
+		case "/key/update":
+			var body map[string]interface{}
+			json.NewDecoder(r.Body).Decode(&body)
+			p.updates = append(p.updates, body)
+			if m, ok := body["metadata"].(map[string]interface{}); ok {
+				p.metadata = m
+			}
+			json.NewEncoder(w).Encode(map[string]interface{}{"key": "hash-1", "metadata": p.metadata})
+		default:
+			http.NotFound(w, r)
+		}
+	}
+}
+
+func applyKeyUpdate(t *testing.T, client *Client, stateAttrs map[string]string, config map[string]interface{}) *terraform.InstanceState {
+	t.Helper()
+	r := resourceKey()
+	state := &terraform.InstanceState{ID: "hash-1", Attributes: stateAttrs}
+	diff, err := r.Diff(context.Background(), state, terraform.NewResourceConfigRaw(config), client)
+	if err != nil {
+		t.Fatalf("Diff returned error: %v", err)
+	}
+	if diff == nil {
+		t.Fatalf("expected a non-empty diff between %v and %v", stateAttrs, config)
+	}
+	newState, diags := r.Apply(context.Background(), state, diff, client)
+	if diags.HasError() {
+		t.Fatalf("Apply returned error: %v", diags)
+	}
+	return newState
+}
+
+func TestKeyUpdateWithoutMetadataChangePreservesServerMetadata(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": float64(5)}}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	newState := applyKeyUpdate(t, client,
+		map[string]string{"key_alias": "alias-1", "max_budget": "10", "metadata.%": "1", "metadata.a": "1"},
+		map[string]interface{}{"key_alias": "alias-1", "max_budget": 20, "metadata": map[string]interface{}{"a": "1"}},
+	)
+
+	if len(proxy.updates) != 1 {
+		t.Fatalf("expected one /key/update call, got %d", len(proxy.updates))
+	}
+	for _, field := range []string{"metadata", "model_rpm_limit", "model_tpm_limit"} {
+		if _, present := proxy.updates[0][field]; present {
+			t.Errorf("unchanged %q was sent on /key/update: %v", field, proxy.updates[0][field])
+		}
+	}
+	if proxy.metadata["server_side"] != "x" {
+		t.Errorf("server-side metadata lost: %v", proxy.metadata)
+	}
+	if got := newState.Attributes["metadata.%"]; got != "1" {
+		t.Errorf("state metadata should hold only the declared entry, got %v", newState.Attributes)
+	}
+	if got := newState.Attributes["metadata.a"]; got != "1" {
+		t.Errorf("metadata.a = %q, want 1", got)
+	}
+}
+
+func TestKeyUpdateWithMetadataChangeMergesOverServerMetadata(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "b": "2", "server_side": "x"}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	applyKeyUpdate(t, client,
+		map[string]string{"key_alias": "alias-1", "metadata.%": "2", "metadata.a": "1", "metadata.b": "2"},
+		map[string]interface{}{"key_alias": "alias-1", "metadata": map[string]interface{}{"a": "2", "c": "3"}},
+	)
+
+	want := map[string]interface{}{"a": "2", "c": "3", "server_side": "x"}
+	if !reflect.DeepEqual(proxy.metadata, want) {
+		t.Errorf("metadata after update = %v, want %v", proxy.metadata, want)
+	}
+}
+
+func TestKeyUpdateSendsChangedModelLimits(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	applyKeyUpdate(t, client,
+		map[string]string{"key_alias": "alias-1", "model_rpm_limit.%": "1", "model_rpm_limit.gpt-4o-mini": "5"},
+		map[string]interface{}{"key_alias": "alias-1", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 7}},
+	)
+
+	got, ok := proxy.updates[0]["model_rpm_limit"].(map[string]interface{})
+	if !ok || got["gpt-4o-mini"] != float64(7) {
+		t.Errorf("changed model_rpm_limit not sent: %v", proxy.updates[0])
+	}
+}
+
+func TestKeyReadKeepsOnlyDeclaredMetadata(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x"}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	d := newKeyResourceData(t, map[string]interface{}{"metadata": map[string]interface{}{"a": "1"}})
+	d.SetId("hash-1")
+	if diags := resourceKeyRead(context.Background(), d, client); diags.HasError() {
+		t.Fatalf("Read returned error: %v", diags)
+	}
+
+	want := map[string]interface{}{"a": "1"}
+	if got := d.Get("metadata"); !reflect.DeepEqual(got, want) {
+		t.Errorf("metadata in state = %v, want %v", got, want)
+	}
+}
+
+func TestKeyUpdateSendsChangedDuration(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	applyKeyUpdate(t, client,
+		map[string]string{"key_alias": "alias-1", "duration": "30d"},
+		map[string]interface{}{"key_alias": "alias-1", "duration": "90d"},
+	)
+
+	if got := proxy.updates[0]["duration"]; got != "90d" {
+		t.Errorf("update payload duration = %v, want 90d", got)
+	}
+}
+
+func TestKeyUpdateOmitsUnchangedDuration(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	applyKeyUpdate(t, client,
+		map[string]string{"key_alias": "alias-1", "duration": "30d"},
+		map[string]interface{}{"key_alias": "alias-2", "duration": "30d"},
+	)
+
+	if got := proxy.updates[0]["key_alias"]; got != "alias-2" {
+		t.Fatalf("update payload key_alias = %v, want alias-2", got)
+	}
+	if v, present := proxy.updates[0]["duration"]; present {
+		t.Errorf("update payload unexpectedly contains duration = %v", v)
+	}
+}

terraform/provider/litellm/resource_key_utils.go

@@ -65,7 +65,7 @@ func buildKeyData(d *schema.ResourceData) map[string]interface{} {
 		keyData["permissions"] = v.(map[string]interface{})
 	}
 	if v, ok := d.GetOkExists("model_max_budget"); ok {
-		keyData["model_max_budget"] = v.(map[string]interface{})
+		keyData["model_max_budget"] = parseKeyModelMaxBudget(v.(string))
 	}
 	if v, ok := d.GetOkExists("model_rpm_limit"); ok {
 		keyData["model_rpm_limit"] = v.(map[string]interface{})
@@ -107,7 +107,7 @@ func setKeyResourceData(d *schema.ResourceData, key *Key) error {
 		"aliases":                key.Aliases,
 		"config":                 key.Config,
 		"permissions":            key.Permissions,
-		"model_max_budget":       key.ModelMaxBudget,
+		"model_max_budget":       keyModelMaxBudgetJSON(key.ModelMaxBudget),
 		"model_rpm_limit":        key.ModelRPMLimit,
 		"model_tpm_limit":        key.ModelTPMLimit,
 		"guardrails":             key.Guardrails,

terraform/provider/litellm/resource_team.go

@@ -31,6 +31,13 @@ func ResourceLiteLLMTeam() *schema.Resource {
 		},
 
 		Schema: map[string]*schema.Schema{
+			"team_id": {
+				Type:        schema.TypeString,
+				Optional:    true,
+				Computed:    true,
+				ForceNew:    true,
+				Description: "Unique ID for the team. Generated by the provider if not provided",
+			},
 			"team_alias": {
 				Type:     schema.TypeString,
 				Required: true,
@@ -162,7 +169,7 @@ func ResourceLiteLLMTeam() *schema.Resource {
 func resourceLiteLLMTeamCreate(d *schema.ResourceData, m interface{}) error {
 	client := m.(*Client)
 
-	teamID := uuid.New().String()
+	teamID := resolveTeamID(d)
 	teamData := buildTeamData(d, teamID)
 
 	// Throughput limit types are only accepted by /team/new, not /team/update.
@@ -214,6 +221,7 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error {
 	teamResp := infoResp.TeamInfo
 
 	// Update the state with values from the response or fall back to the data passed in during creation
+	d.Set("team_id", d.Id())
 	d.Set("team_alias", GetStringValue(teamResp.TeamAlias, d.Get("team_alias").(string)))
 	d.Set("organization_id", GetStringValue(teamResp.OrganizationID, d.Get("organization_id").(string)))
 
@@ -263,11 +271,11 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error {
 		d.Set("team_member_tpm_limit", *teamResp.TeamMemberTPMLimit)
 	}
 	d.Set("team_member_key_duration", GetStringValue(teamResp.TeamMemberKeyDuration, d.Get("team_member_key_duration").(string)))
-	if teamResp.ModelRPMLimit != nil {
-		d.Set("model_rpm_limit", teamResp.ModelRPMLimit)
+	if v := teamModelLimit(teamResp.ModelRPMLimit, teamResp.Metadata, "model_rpm_limit"); v != nil {
+		d.Set("model_rpm_limit", v)
 	}
-	if teamResp.ModelTPMLimit != nil {
-		d.Set("model_tpm_limit", teamResp.ModelTPMLimit)
+	if v := teamModelLimit(teamResp.ModelTPMLimit, teamResp.Metadata, "model_tpm_limit"); v != nil {
+		d.Set("model_tpm_limit", v)
 	}
 	if teamResp.AllowedPassthroughRoutes != nil {
 		d.Set("allowed_passthrough_routes", teamResp.AllowedPassthroughRoutes)
@@ -354,6 +362,13 @@ func resourceLiteLLMTeamDelete(d *schema.ResourceData, m interface{}) error {
 	return nil
 }
 
+func resolveTeamID(d *schema.ResourceData) string {
+	if v, ok := d.GetOk("team_id"); ok {
+		return v.(string)
+	}
+	return uuid.New().String()
+}
+
 func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} {
 	teamData := map[string]interface{}{
 		"team_id":    teamID,
@@ -364,14 +379,19 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{}
 		"organization_id", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models",
 		"blocked", "team_member_permissions", "model_aliases", "guardrails", "prompts",
 		"team_member_budget", "team_member_budget_duration", "team_member_rpm_limit",
-		"team_member_tpm_limit", "team_member_key_duration", "model_rpm_limit",
-		"model_tpm_limit", "allowed_passthrough_routes",
+		"team_member_tpm_limit", "team_member_key_duration", "allowed_passthrough_routes",
 	} {
 		if v, ok := d.GetOk(key); ok {
 			teamData[key] = v
 		}
 	}
 
+	for _, key := range []string{"model_rpm_limit", "model_tpm_limit"} {
+		if v, ok := d.GetOk(key); ok || d.HasChange(key) {
+			teamData[key] = v
+		}
+	}
+
 	if v, ok := d.GetOk("soft_budget"); ok {
 		teamData["soft_budget"] = v
 	} else if d.HasChange("soft_budget") {
@@ -404,6 +424,14 @@ func buildTeamMetadata(d *schema.ResourceData) map[string]interface{} {
 	return metadata
 }
 
+func teamModelLimit(topLevel, metadata map[string]interface{}, key string) map[string]interface{} {
+	if topLevel != nil {
+		return topLevel
+	}
+	nested, _ := metadata[key].(map[string]interface{})
+	return nested
+}
+
 func splitTeamMetadata(raw map[string]interface{}) (map[string]string, []string, []string) {
 	metadata := map[string]string{}
 	var tags, alertEmails []string

terraform/provider/litellm/resource_team_test.go

@@ -9,6 +9,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 )
@@ -85,6 +86,87 @@ func TestTeamCreateSendsSoftBudgetTagsAndAlertEmails(t *testing.T) {
 	}
 }
 
+func TestTeamCreateSendsConfiguredTeamID(t *testing.T) {
+	var captured map[string]interface{}
+	srv := newTeamTestServer(t, &captured, `{"team_id":"platform-team","team_info":{"team_id":"platform-team","team_alias":"platform"},"keys":[],"team_memberships":[]}`)
+	defer srv.Close()
+
+	d := newTeamResourceData(t, map[string]interface{}{
+		"team_id":    "platform-team",
+		"team_alias": "platform",
+	})
+
+	if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil {
+		t.Fatalf("create failed: %v", err)
+	}
+
+	if got := captured["team_id"]; got != "platform-team" {
+		t.Fatalf("payload team_id = %v, want platform-team", got)
+	}
+	if got := d.Id(); got != "platform-team" {
+		t.Fatalf("resource id = %q, want platform-team", got)
+	}
+	if got := d.Get("team_id"); got != "platform-team" {
+		t.Fatalf("state team_id = %v, want platform-team", got)
+	}
+}
+
+func TestTeamCreateGeneratesTeamIDWhenUnset(t *testing.T) {
+	var captured map[string]interface{}
+	srv := newTeamTestServer(t, &captured, `{"team_id":"x","team_info":{"team_alias":"eng"},"keys":[],"team_memberships":[]}`)
+	defer srv.Close()
+
+	d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"})
+
+	if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil {
+		t.Fatalf("create failed: %v", err)
+	}
+
+	sent, _ := captured["team_id"].(string)
+	if _, err := uuid.Parse(sent); err != nil {
+		t.Fatalf("payload team_id = %q, want a generated UUID: %v", sent, err)
+	}
+	if d.Id() != sent || d.Get("team_id") != sent {
+		t.Fatalf("id = %q, state team_id = %v, want both to equal the sent id %q", d.Id(), d.Get("team_id"), sent)
+	}
+}
+
+func TestTeamReadSetsTeamIDFromResourceID(t *testing.T) {
+	var captured map[string]interface{}
+	srv := newTeamTestServer(t, &captured, `{"team_id":"imported-team","team_info":{"team_id":"imported-team","team_alias":"imported"},"keys":[],"team_memberships":[]}`)
+	defer srv.Close()
+
+	d := newTeamResourceData(t, map[string]interface{}{})
+	d.SetId("imported-team")
+
+	if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil {
+		t.Fatalf("read failed: %v", err)
+	}
+	if got := d.Get("team_id"); got != "imported-team" {
+		t.Fatalf("team_id = %v, want imported-team", got)
+	}
+}
+
+func TestTeamIDChangeForcesReplacement(t *testing.T) {
+	res := ResourceLiteLLMTeam()
+	priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{
+		"team_id":    "old-team",
+		"team_alias": "eng",
+	})
+	priorData.SetId("old-team")
+	config := terraform.NewResourceConfigRaw(map[string]interface{}{
+		"team_id":    "new-team",
+		"team_alias": "eng",
+	})
+	diff, err := res.Diff(context.Background(), priorData.State(), config, nil)
+	if err != nil {
+		t.Fatalf("diff failed: %v", err)
+	}
+	if diff == nil || !diff.RequiresNew() {
+		t.Fatalf("changing team_id must force replacement, diff = %+v", diff)
+	}
+}
+
 func TestTeamReadMapsTeamInfoEnvelope(t *testing.T) {
 	var captured map[string]interface{}
 	srv := newTeamTestServer(t, &captured, teamInfoWithSoftBudget)
@@ -250,6 +332,77 @@ func TestTeamReadMapsNewFields(t *testing.T) {
 	}
 }
 
+func TestTeamReadMapsPerModelLimitsFromMetadata(t *testing.T) {
+	var captured map[string]interface{}
+	srv := newTeamTestServer(t, &captured, `{
+		"team_id": "team-1",
+		"team_info": {
+			"team_id": "team-1",
+			"team_alias": "eng",
+			"model_rpm_limit": null,
+			"model_tpm_limit": null,
+			"metadata": {
+				"department": "eng",
+				"model_rpm_limit": {"gpt-4o-mini": 250},
+				"model_tpm_limit": {"gpt-4o-mini": 5000}
+			}
+		}
+	}`)
+	defer srv.Close()
+
+	d := newTeamResourceData(t, map[string]interface{}{
+		"team_alias":      "eng",
+		"model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 100},
+	})
+	d.SetId("team-1")
+
+	if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil {
+		t.Fatalf("read returned error: %v", err)
+	}
+	if got := d.Get("model_rpm_limit"); !reflect.DeepEqual(got, map[string]interface{}{"gpt-4o-mini": 250}) {
+		t.Errorf("model_rpm_limit = %v, want server value 250", got)
+	}
+	if got := d.Get("model_tpm_limit"); !reflect.DeepEqual(got, map[string]interface{}{"gpt-4o-mini": 5000}) {
+		t.Errorf("model_tpm_limit = %v, want server value 5000", got)
+	}
+	if got := d.Get("metadata"); !reflect.DeepEqual(got, map[string]interface{}{"department": "eng"}) {
+		t.Errorf("metadata = %v, want per-model limits kept out of the string map", got)
+	}
+}
+
+func TestTeamUpdateClearsRemovedPerModelLimits(t *testing.T) {
+	var captured map[string]interface{}
+	srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`)
+	defer srv.Close()
+
+	res := ResourceLiteLLMTeam()
+	priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{
+		"team_alias":      "eng",
+		"model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 100},
+		"model_tpm_limit": map[string]interface{}{"gpt-4o-mini": 5000},
+	})
+	priorData.SetId("team-1")
+	prior := priorData.State()
+	config := terraform.NewResourceConfigRaw(map[string]interface{}{"team_alias": "eng"})
+	diff, err := res.Diff(context.Background(), prior, config, nil)
+	if err != nil {
+		t.Fatalf("diff failed: %v", err)
+	}
+	d, err := schema.InternalMap(res.Schema).Data(prior, diff)
+	if err != nil {
+		t.Fatalf("data failed: %v", err)
+	}
+
+	if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil {
+		t.Fatalf("update failed: %v", err)
+	}
+	for _, k := range []string{"model_rpm_limit", "model_tpm_limit"} {
+		if got, ok := captured[k]; !ok || !reflect.DeepEqual(got, map[string]interface{}{}) {
+			t.Errorf("payload %s = %v (present=%v), want explicit empty map", k, got, ok)
+		}
+	}
+}
+
 // rpm_limit_type / tpm_limit_type are accepted by /team/new but not
 // /team/update, so create must send them and update must not.
 func TestTeamLimitTypesSentOnCreateOnly(t *testing.T) {

tests/e2e/CLAUDE.md

@@ -181,13 +181,15 @@ quota_management.<behavior>.<variant>.<assertion>
                                | team_multi_window | fallback | spend_counter
               <spend_tracking> chat_completions | stream | messages_bridge | embeddings
                                | cache_hit | key_rollup | concurrent_burst | tags | end_user
-                               | per_model | failure | spend_calculate | pagination
+                               | per_model | failure | spend_calculate | pagination | key_attribution
   assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm
               | blocks_then_resets | resets_windows_independently | alerts_without_blocking
               | isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys
               | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost
               | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows
-              | writes_failure_row | returns_cost | keeps_total
+              | writes_failure_row | returns_cost | keeps_total | joins_key | reports_alias_and_email
+              | health_rows_keep_service_account | retrieve_batch_cost_joins_retrieving_key
+              | poller_batch_cost_joins_creating_key
   e.g.  quota_management.ratelimit.rpm.blocks_over_limit           exercised_on=[chat_completions, messages]
         quota_management.budget.key.blocks_over_limit              exercised_on=[chat_completions]
 ```

tests/e2e/coverage_registry/quota_management.yaml

@@ -58,3 +58,8 @@
 - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"}
 - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"}
 - {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"}
+- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, batch create, and a replayed callback log carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"}
+- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"}
+- {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"}
+- {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"}
+- {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"}

tests/e2e/models.py

@@ -84,6 +84,7 @@ class KeyGenerateBody(BaseModel):
 
 class KeyGenerateResponse(BaseModel):
     key: str
+    token: str | None = None
     key_alias: str | None = None
     models: list[str] = []
     max_budget: float | None = None
@@ -672,6 +673,7 @@ class GuardrailRunRecord(BaseModel):
 
 
 class SpendLogMetadata(BaseModel):
+    user_api_key_alias: str | None = None
     applied_guardrails: list[str] | None = None
     guardrail_information: list[GuardrailRunRecord] | None = None
 

tests/e2e/quota_management/spend_tracking/conftest.py

@@ -36,6 +36,7 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = (
     ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"),
     ("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"),
     ("openai-responses-codex", "openai/gpt-5.3-codex", "OPENAI_API_KEY"),
+    ("openai-gpt-4o-mini", "openai/gpt-4o-mini", "OPENAI_API_KEY"),
 )
 
 

tests/e2e/quota_management/spend_tracking/spend_e2e_client.py

@@ -15,9 +15,12 @@ import time
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
+from typing import Final
 
 from e2e_config import unique_marker
 from e2e_http import (
+    FileUploadForm,
+    Headers,
     NoBody,
     ProbeResult,
     Result,
@@ -35,6 +38,8 @@ from models import (
     DateRangeParams,
     EmbedBody,
     EmbedResponse,
+    KeyGenerateBody,
+    KeyGenerateResponse,
     OpenAPISchema,
     SpendCalculateBody,
     SpendCalculateResponse,
@@ -43,13 +48,27 @@ from models import (
     SpendLogsPageParams,
     SpendTagsResponse,
     TagSpend,
+    UserDeleteBody,
+    UserDeleteResponse,
+    UserNewBody,
+    UserNewResponse,
+    UserRole,
 )
-from proxy_client import ProxyClient
+from proxy_client import Converged, ProxyClient, await_converged
+from pydantic import BaseModel, Field
 
 __all__ = [
+    "BatchCreateBody",
+    "CallbackLogMetadata",
+    "CallbackLogPayload",
+    "BatchObject",
+    "DailyActivityKeyBreakdown",
+    "FileObject",
     "ProbeResult",
+    "ResponseIdentity",
     "SpendClient",
     "SpendLogRow",
+    "StreamingResponse",
     "build_client",
     "is_ok",
     "unique_marker",
@@ -57,6 +76,139 @@ __all__ = [
 ]
 
 
+class GeminiApiKeyHeaders(Headers):
+    x_goog_api_key: str = Field(serialization_alias="x-goog-api-key")
+    content_type: str = Field(default="application/json", serialization_alias="Content-Type")
+
+
+class GeminiPart(BaseModel):
+    text: str
+
+
+class GeminiContent(BaseModel):
+    parts: list[GeminiPart]
+
+
+class GeminiGenerationConfig(BaseModel):
+    maxOutputTokens: int
+
+
+class GeminiGenerateBody(BaseModel):
+    contents: list[GeminiContent]
+    generationConfig: GeminiGenerationConfig
+
+
+class ResponsesBody(BaseModel):
+    model: str
+    input: str
+    cache: dict[str, bool] | None = {"no-cache": True}
+
+
+class QueuedChatBody(ChatBody):
+    priority: int = 0
+
+
+class ResponseIdentity(BaseModel):
+    id: str | None = None
+
+
+class HealthParams(BaseModel):
+    model: str
+
+
+class ModelQuery(BaseModel):
+    model: str
+
+
+class FileObject(BaseModel):
+    id: str
+
+
+class BatchCreateBody(BaseModel):
+    input_file_id: str
+    endpoint: str = "/v1/chat/completions"
+    completion_window: str = "24h"
+    model: str
+    metadata: dict[str, str]
+
+
+class BatchObject(BaseModel):
+    id: str
+    status: str
+
+
+class ProviderQuery(BaseModel):
+    provider: str
+
+
+class CallbackLogMetadata(BaseModel):
+    user_api_key_hash: str
+    user_api_key_alias: str
+    user_api_key_user_id: str
+
+
+class CallbackLogPayload(BaseModel):
+    id: str
+    litellm_call_id: str
+    model: str
+    call_type: str = "acompletion"
+    start_time: float = Field(serialization_alias="startTime")
+    end_time: float = Field(serialization_alias="endTime")
+    response_cost: float
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    metadata: CallbackLogMetadata
+
+
+class CallbackLogRecord(BaseModel):
+    status: str = "success"
+    standard_logging_payload: CallbackLogPayload
+
+
+class CallbackLogsRequest(BaseModel):
+    records: list[CallbackLogRecord]
+
+
+class CallbackLogsResponse(BaseModel):
+    processed: int
+    failed: int
+
+
+class DailyActivityParams(BaseModel):
+    start_date: str
+    end_date: str
+    api_key: str
+
+
+class DailyActivityKeyMetadata(BaseModel):
+    key_alias: str | None = None
+    team_id: str | None = None
+    user_email: str | None = None
+
+
+class DailyActivityKeyMetrics(BaseModel):
+    api_requests: int = 0
+
+
+class DailyActivityKeyBreakdown(BaseModel):
+    metrics: DailyActivityKeyMetrics
+    metadata: DailyActivityKeyMetadata
+
+
+class DailyActivityBreakdown(BaseModel):
+    api_keys: dict[str, DailyActivityKeyBreakdown] = {}
+
+
+class DailyActivityRow(BaseModel):
+    date: str
+    breakdown: DailyActivityBreakdown
+
+
+class DailyActivityResponse(BaseModel):
+    results: list[DailyActivityRow] = []
+
+
 def _chat_body(
     model: str,
     content: str,
@@ -207,6 +359,166 @@ class SpendClient:
     def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult:
         return self.proxy.transport.probe(path, params=params)
 
+    def create_user(self, *, email: str, role: UserRole, user_id: str) -> str:
+        return unwrap(
+            self.proxy.transport.post(
+                "/user/new",
+                headers=self.proxy.transport.master,
+                json=UserNewBody(user_email=email, user_role=role, user_id=user_id),
+                response_type=UserNewResponse,
+            )
+        ).user_id
+
+    def delete_user(self, user_id: str) -> None:
+        _ = unwrap(
+            self.proxy.transport.post(
+                "/user/delete",
+                headers=self.proxy.transport.master,
+                json=UserDeleteBody(user_ids=[user_id]),
+                response_type=UserDeleteResponse,
+            )
+        )
+
+    def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse:
+        return unwrap(
+            self.proxy.transport.post(
+                "/key/generate",
+                headers=self.proxy.transport.master,
+                json=body,
+                response_type=KeyGenerateResponse,
+            )
+        )
+
+    def send_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/chat/completions",
+            headers=self.proxy.transport.bearer(key),
+            json=_chat_body(model, content, max_tokens=max_tokens),
+        )
+
+    def send_queued_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/queue/chat/completions",
+            headers=self.proxy.transport.bearer(key),
+            json=QueuedChatBody(
+                model=model,
+                messages=[ChatMessage(role="user", content=content)],
+                max_tokens=max_tokens,
+            ),
+        )
+
+    def send_messages(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/v1/messages",
+            headers=self.proxy.transport.bearer(key),
+            json=AnthropicMessagesBody(
+                model=model,
+                messages=[ChatMessage(role="user", content=content)],
+                max_tokens=max_tokens,
+            ),
+        )
+
+    def send_responses(self, key: str, model: str, content: str) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/v1/responses",
+            headers=self.proxy.transport.bearer(key),
+            json=ResponsesBody(model=model, input=content),
+        )
+
+    def send_embed(self, key: str, model: str, content: str) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/embeddings",
+            headers=self.proxy.transport.bearer(key),
+            json=EmbedBody(model=model, input=content),
+        )
+
+    def send_gemini_generate(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
+        return self.proxy.transport.send(
+            f"/gemini/v1beta/models/{model}:generateContent",
+            headers=GeminiApiKeyHeaders(x_goog_api_key=key),
+            json=GeminiGenerateBody(
+                contents=[GeminiContent(parts=[GeminiPart(text=content)])],
+                generationConfig=GeminiGenerationConfig(maxOutputTokens=max_tokens),
+            ),
+        )
+
+    def upload_batch_file(self, key: str, model: str, content: bytes) -> FileObject:
+        return unwrap(
+            self.proxy.transport.upload(
+                "/v1/files",
+                headers=self.proxy.transport.bearer(key),
+                form=FileUploadForm(purpose="batch"),
+                filename="key_attribution.jsonl",
+                content=content,
+                params=ModelQuery(model=model),
+                response_type=FileObject,
+            )
+        )
+
+    def create_batch(self, key: str, body: BatchCreateBody) -> BatchObject:
+        return unwrap(
+            self.proxy.transport.post(
+                "/v1/batches",
+                headers=self.proxy.transport.bearer(key),
+                json=body,
+                response_type=BatchObject,
+            )
+        )
+
+    def retrieve_batch(self, key: str, batch_id: str, *, provider: str) -> BatchObject:
+        return unwrap(
+            self.proxy.transport.get(
+                f"/v1/batches/{batch_id}",
+                headers=self.proxy.transport.bearer(key),
+                params=ProviderQuery(provider=provider),
+                response_type=BatchObject,
+            )
+        )
+
+    def replay_callback_log(self, key: str, payload: CallbackLogPayload) -> CallbackLogsResponse:
+        return unwrap(
+            self.proxy.transport.post(
+                "/v1/rust_control_plane/logs",
+                headers=self.proxy.transport.bearer(key),
+                json=CallbackLogsRequest(records=[CallbackLogRecord(standard_logging_payload=payload)]),
+                response_type=CallbackLogsResponse,
+            )
+        )
+
+    def health(self, model: str) -> ProbeResult:
+        return self.proxy.transport.probe("/health", params=HealthParams(model=model))
+
+    def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None:
+        response: Final = unwrap(
+            self.proxy.transport.get(
+                "/user/daily/activity",
+                headers=self.proxy.transport.master,
+                params=DailyActivityParams(
+                    start_date=start.strftime("%Y-%m-%d"),
+                    end_date=end.strftime("%Y-%m-%d"),
+                    api_key=token,
+                ),
+                response_type=DailyActivityResponse,
+            )
+        )
+        return next(
+            (row.breakdown.api_keys[token] for row in response.results if token in row.breakdown.api_keys),
+            None,
+        )
+
+    def poll_daily_activity_for_key(
+        self, token: str, *, start: datetime, end: datetime, min_requests: int
+    ) -> DailyActivityKeyBreakdown | None:
+        outcome: Final = await_converged(
+            lambda: self.daily_activity_for_key(token, start=start, end=end),
+            converged=lambda found: found is not None and found.metrics.api_requests >= min_requests,
+            timeout=self.proxy.poll_timeout,
+            interval=self.proxy.poll_interval,
+            now=time.monotonic,
+            sleep=time.sleep,
+        )
+        return outcome.result if isinstance(outcome, Converged) else outcome.last_result
+
     def openapi(self) -> OpenAPISchema:
         return unwrap(
             self.proxy.transport.get(

tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py

@@ -0,0 +1,405 @@
+"""Every spend row a live proxy writes joins its virtual key (MAT-180).
+
+One virtual key with an alias, owned by a user with an email, drives every spend
+write path a key can reach: /chat/completions, /queue/chat/completions,
+/v1/messages, /v1/responses, /embeddings, the Gemini native passthrough, a batch
+input file upload, a batch create, and a replayed callback log (POST
+/v1/rust_control_plane/logs, the writer an external gateway feeds). Each row those calls write must carry
+`api_key` equal to the key's LiteLLM_VerificationToken.token (the sha256 hash
+/key/generate returns as `token`), which is the join /spend/logs?api_key= and
+/user/daily/activity rely on to report key_alias and user_email. A row keyed by a
+re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a
+key-hash-* row with no alias and no email in the customer's usage exports.
+
+The health-check service account writes rows too; those must stay keyed by the
+literal service-account name, never by a hash of it. A batch's cost row is
+written by the retrieve that first sees the batch in a terminal state, so the
+batch the run creates is one OpenAI fails at validation within seconds (its one
+line targets /v1/embeddings under a /v1/chat/completions batch), and the test
+retrieves it by its raw provider id with the same key until it is failed. A raw
+id is never owned by the CheckBatchCost poller, so that retrieve prices the batch
+inline against the retrieving key and its {provider_batch_id}_batch_cost row
+must join the key's token with its alias. A completed batch with a positive
+cost is out of a single run's reach (OpenAI's completion window is 24h, and a
+stack booted fresh per run lists no earlier run's batches), so the poller's own
+row is not asserted here.
+
+/spend/logs carries no email field, so the email assertion lives on
+/user/daily/activity alone; /spend/logs is held to the alias in metadata.
+"""
+
+import base64
+import time
+from collections.abc import Iterator
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Final
+
+import pytest
+from models import KeyGenerateBody
+from proxy_client import Converged, await_converged
+from pydantic import BaseModel
+from spend_e2e_client import (
+    BatchCreateBody,
+    BatchObject,
+    CallbackLogMetadata,
+    CallbackLogPayload,
+    DailyActivityKeyBreakdown,
+    ResponseIdentity,
+    SpendClient,
+    SpendLogRow,
+    StreamingResponse,
+    unique_marker,
+)
+
+pytestmark = pytest.mark.e2e
+
+CHAT_MODEL: Final = "gemini-2.5-flash"
+MESSAGES_MODEL: Final = "claude-haiku-4-5"
+RESPONSES_MODEL: Final = "openai-responses-codex"
+EMBED_MODEL: Final = "openai-text-embedding-3-small"
+BATCH_MODEL: Final = "openai-gpt-4o-mini"
+BATCH_BACKEND_MODEL: Final = "gpt-4o-mini"
+BATCH_PROVIDER: Final = "openai"
+HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check"
+BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"})
+FAILED_BATCH_POLL_SECONDS: Final = 120.0
+FAILED_BATCH_POLL_INTERVAL_SECONDS: Final = 5.0
+MAX_TOKENS: Final = 8
+REPLAY_RESPONSE_COST: Final = 0.0001
+REPLAY_PROMPT_TOKENS: Final = 5
+REPLAY_COMPLETION_TOKENS: Final = 1
+WRITE_PATHS: Final = (
+    "chat_completions",
+    "queue_chat_completions",
+    "messages",
+    "responses",
+    "embeddings",
+    "gemini_passthrough",
+    "batch_file_upload",
+    "batch_create",
+    "callback_replay",
+)
+
+
+class EmbeddingLineBody(BaseModel):
+    model: str
+    input: str
+
+
+class EmbeddingLine(BaseModel):
+    custom_id: str
+    method: str = "POST"
+    url: str = "/v1/embeddings"
+    body: EmbeddingLineBody
+
+
+@dataclass(frozen=True, slots=True)
+class AttributedKey:
+    key: str
+    token: str
+    alias: str
+    email: str
+    user_id: str
+
+
+@dataclass(frozen=True, slots=True)
+class WritePath:
+    name: str
+    request_id: str
+
+
+@dataclass(frozen=True, slots=True)
+class DrivenKey:
+    identity: AttributedKey
+    paths: tuple[WritePath, ...]
+    started_at: datetime
+
+
+def _body_id(name: str, sent: StreamingResponse) -> WritePath:
+    assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}"
+    response_id: Final = ResponseIdentity.model_validate_json(sent.body).id
+    assert response_id, f"{name} answered without a response id: {sent.body[:300]}"
+    return WritePath(name=name, request_id=response_id)
+
+
+def _call_id(name: str, sent: StreamingResponse) -> WritePath:
+    assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}"
+    assert sent.call_id, f"{name} answered without an x-litellm-call-id header"
+    return WritePath(name=name, request_id=sent.call_id)
+
+
+def _endpoint_mismatched_jsonl(marker: str) -> bytes:
+    line: Final = EmbeddingLine(custom_id=marker, body=EmbeddingLineBody(model=BATCH_BACKEND_MODEL, input=marker))
+    return f"{line.model_dump_json()}\n".encode()
+
+
+def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> tuple[WritePath, WritePath]:
+    uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _endpoint_mismatched_jsonl(marker))
+    created: Final = client.create_batch(
+        identity.key,
+        BatchCreateBody(
+            input_file_id=uploaded.id,
+            model=BATCH_MODEL,
+            metadata={"run": marker},
+        ),
+    )
+    return (
+        WritePath(name="batch_file_upload", request_id=uploaded.id),
+        WritePath(name="batch_create", request_id=created.id),
+    )
+
+
+def _drive_callback_replay(client: SpendClient, identity: AttributedKey, marker: str) -> WritePath:
+    request_id: Final = f"callback-replay-{marker}"
+    finished_at: Final = time.time()
+    replayed: Final = client.replay_callback_log(
+        identity.key,
+        CallbackLogPayload(
+            id=request_id,
+            litellm_call_id=request_id,
+            model=CHAT_MODEL,
+            start_time=finished_at - 1,
+            end_time=finished_at,
+            response_cost=REPLAY_RESPONSE_COST,
+            prompt_tokens=REPLAY_PROMPT_TOKENS,
+            completion_tokens=REPLAY_COMPLETION_TOKENS,
+            total_tokens=REPLAY_PROMPT_TOKENS + REPLAY_COMPLETION_TOKENS,
+            metadata=CallbackLogMetadata(
+                user_api_key_hash=identity.token,
+                user_api_key_alias=identity.alias,
+                user_api_key_user_id=identity.user_id,
+            ),
+        ),
+    )
+    assert replayed.processed == 1 and replayed.failed == 0, f"callback replay rejected the payload: {replayed}"
+    return WritePath(name="callback_replay", request_id=request_id)
+
+
+def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tuple[WritePath, ...]:
+    marker: Final = unique_marker()
+    prompt: Final = f"Reply with the word ok. {marker}"
+    key: Final = identity.key
+    return (
+        _body_id("chat_completions", client.send_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
+        _body_id("queue_chat_completions", client.send_queued_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
+        _body_id("messages", client.send_messages(key, MESSAGES_MODEL, prompt, max_tokens=MAX_TOKENS)),
+        _body_id("responses", client.send_responses(key, RESPONSES_MODEL, prompt)),
+        _call_id("embeddings", client.send_embed(key, EMBED_MODEL, prompt)),
+        _call_id("gemini_passthrough", client.send_gemini_generate(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
+        *_drive_batch(client, identity, marker),
+        _drive_callback_replay(client, identity, marker),
+    )
+
+
+def _provider_batch_id(unified_batch_id: str) -> str:
+    encoded: Final = unified_batch_id.removeprefix("batch_")
+    decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode()
+    return decoded.removeprefix("litellm:").split(";", 1)[0]
+
+
+def _driven_batch_id(driven: DrivenKey) -> str:
+    return next(path.request_id for path in driven.paths if path.name == "batch_create")
+
+
+def _await_terminal_batch(client: SpendClient, key: str, provider_batch_id: str) -> BatchObject:
+    outcome: Final = await_converged(
+        lambda: client.retrieve_batch(key, provider_batch_id, provider=BATCH_PROVIDER),
+        converged=lambda batch: batch.status in BATCH_TERMINAL_STATUSES,
+        timeout=FAILED_BATCH_POLL_SECONDS,
+        interval=FAILED_BATCH_POLL_INTERVAL_SECONDS,
+        now=time.monotonic,
+        sleep=time.sleep,
+    )
+    return outcome.result if isinstance(outcome, Converged) else outcome.last_result
+
+
+def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]:
+    return [
+        row
+        for row in client.proxy.spend_logs_window(
+            start=started_at - timedelta(minutes=1), end=datetime.now(timezone.utc) + timedelta(minutes=1)
+        )
+        if HEALTH_SERVICE_ACCOUNT in (row.request_tags or [])
+    ]
+
+
+def _health_rows_since(client: SpendClient, started_at: datetime) -> list[SpendLogRow]:
+    outcome: Final = await_converged(
+        lambda: _health_rows_between(client, started_at),
+        converged=lambda rows: bool(rows),
+        timeout=client.proxy.poll_timeout,
+        interval=client.proxy.poll_interval,
+        now=time.monotonic,
+        sleep=time.sleep,
+    )
+    return outcome.result if isinstance(outcome, Converged) else outcome.last_result
+
+
+class TestKeyAttribution:
+    @pytest.fixture(scope="class")
+    def driven(self, client: SpendClient) -> Iterator[DrivenKey]:
+        marker: Final = unique_marker()
+        user_id: Final = client.create_user(
+            email=f"key-attribution-{marker}@example.com",
+            role="proxy_admin",
+            user_id=f"key-attribution-{marker}",
+        )
+        record: Final = client.generate_key_record(
+            KeyGenerateBody(models=[], user_id=user_id, key_alias=f"key-attribution-{marker}")
+        )
+        assert record.token, "/key/generate answered without the key's token hash"
+        assert record.key_alias, "/key/generate dropped the key alias"
+        identity: Final = AttributedKey(
+            key=record.key,
+            token=record.token,
+            alias=record.key_alias,
+            email=f"key-attribution-{marker}@example.com",
+            user_id=user_id,
+        )
+        started_at: Final = datetime.now(timezone.utc)
+        try:
+            yield DrivenKey(
+                identity=identity,
+                paths=_drive_every_write_path(client, identity),
+                started_at=started_at,
+            )
+        finally:
+            client.proxy.delete_key(identity.key)
+            client.delete_user(identity.user_id)
+
+    @pytest.mark.covers(
+        "quota_management.spend_tracking.key_attribution.joins_key",
+        exercised_on=[
+            "chat_completions",
+            "messages",
+            "responses",
+            "embeddings",
+            "batches",
+            "files",
+            "google_native",
+            "rust_control_plane",
+        ],
+    )
+    def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None:
+        assert tuple(path.name for path in driven.paths) == WRITE_PATHS
+        found: Final = tuple((path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths)
+        unwritten: Final = [path.name for path, rows in found if not rows]
+        assert not unwritten, f"write paths that produced no spend row within the poll window: {unwritten}"
+        unjoined: Final = [
+            (path.name, row.call_type, row.api_key)
+            for path, rows in found
+            for row in rows
+            if row.api_key != driven.identity.token
+        ]
+        assert not unjoined, (
+            "spend rows whose api_key does not join LiteLLM_VerificationToken.token "
+            f"{driven.identity.token}: {unjoined}"
+        )
+        unaliased: Final = [
+            (path.name, row.call_type, row.metadata.user_api_key_alias if row.metadata else None)
+            for path, rows in found
+            for row in rows
+            if row.metadata is None or row.metadata.user_api_key_alias != driven.identity.alias
+        ]
+        assert not unaliased, f"spend rows written without key alias {driven.identity.alias!r}: {unaliased}"
+
+    @pytest.mark.covers(
+        "quota_management.spend_tracking.key_attribution.reports_alias_and_email",
+        exercised_on=[
+            "chat_completions",
+            "messages",
+            "responses",
+            "embeddings",
+            "batches",
+            "files",
+            "google_native",
+            "rust_control_plane",
+        ],
+    )
+    def test_spend_logs_by_key_return_every_row_with_the_alias(self, client: SpendClient, driven: DrivenKey) -> None:
+        expected_ids: Final = frozenset(path.request_id for path in driven.paths)
+        rows: Final = client.poll_logs_for_key(
+            driven.identity.key,
+            min_rows=len(driven.paths),
+            predicate=lambda found: expected_ids <= frozenset(row.request_id or "" for row in found),
+        )
+        missing: Final = expected_ids - frozenset(row.request_id or "" for row in rows)
+        assert not missing, (
+            f"/spend/logs?api_key= does not return {len(missing)} of {len(expected_ids)} rows for the key: "
+            f"{sorted(path.name for path in driven.paths if path.request_id in missing)}"
+        )
+        aliases: Final = frozenset(row.metadata.user_api_key_alias if row.metadata else None for row in rows)
+        assert aliases == {driven.identity.alias}, f"/spend/logs rows carry aliases {sorted(map(str, aliases))}"
+
+    @pytest.mark.covers(
+        "quota_management.spend_tracking.key_attribution.reports_alias_and_email",
+        exercised_on=[
+            "chat_completions",
+            "messages",
+            "responses",
+            "embeddings",
+            "batches",
+            "files",
+            "google_native",
+            "rust_control_plane",
+        ],
+    )
+    def test_user_daily_activity_reports_alias_and_email(self, client: SpendClient, driven: DrivenKey) -> None:
+        breakdown: Final[DailyActivityKeyBreakdown | None] = client.poll_daily_activity_for_key(
+            driven.identity.token,
+            start=driven.started_at - timedelta(days=1),
+            end=datetime.now(timezone.utc) + timedelta(days=1),
+            min_requests=len(driven.paths),
+        )
+        assert breakdown is not None, (
+            f"/user/daily/activity?api_key={driven.identity.token} has no api_keys breakdown: "
+            "the key's rows did not aggregate under its token"
+        )
+        assert breakdown.metrics.api_requests >= len(driven.paths), (
+            f"/user/daily/activity counts {breakdown.metrics.api_requests} requests for the key, "
+            f"expected at least {len(driven.paths)}"
+        )
+        assert breakdown.metadata.key_alias == driven.identity.alias, f"key_alias={breakdown.metadata.key_alias!r}"
+        assert breakdown.metadata.user_email == driven.identity.email, f"user_email={breakdown.metadata.user_email!r}"
+
+    @pytest.mark.covers(
+        "quota_management.spend_tracking.key_attribution.health_rows_keep_service_account",
+        exercised_on=["chat_completions"],
+    )
+    def test_health_check_rows_keep_the_service_account_key(self, client: SpendClient) -> None:
+        started_at: Final = datetime.now(timezone.utc)
+        probe: Final = client.health(CHAT_MODEL)
+        assert probe.healthy, f"/health?model={CHAT_MODEL} answered {probe.status_code}: {probe.body[:300]}"
+        rows: Final = _health_rows_since(client, started_at)
+        assert rows, f"/health?model={CHAT_MODEL} wrote no {HEALTH_SERVICE_ACCOUNT}-tagged spend row"
+        rehashed: Final = [(row.request_id, row.api_key) for row in rows if row.api_key != HEALTH_SERVICE_ACCOUNT]
+        assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}"
+
+    @pytest.mark.covers(
+        "quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key",
+        exercised_on=["batches"],
+    )
+    def test_terminal_batch_cost_row_joins_the_retrieving_key(self, client: SpendClient, driven: DrivenKey) -> None:
+        provider_batch_id: Final = _provider_batch_id(_driven_batch_id(driven))
+        fetched: Final = _await_terminal_batch(client, driven.identity.key, provider_batch_id)
+        assert fetched.status == "failed", (
+            f"endpoint-mismatched batch {provider_batch_id} is {fetched.status!r} after "
+            f"{FAILED_BATCH_POLL_SECONDS:.0f}s, so its terminal cost row cannot be asserted"
+        )
+        cost_request_id: Final = f"{provider_batch_id}_batch_cost"
+        rows: Final = client.proxy.poll_logs_for_request_id(cost_request_id)
+        assert rows, f"retrieving failed batch {provider_batch_id} wrote no cost row under {cost_request_id}"
+        call_types: Final = tuple(sorted({row.call_type or "" for row in rows}))
+        assert call_types == ("aretrieve_batch",), f"cost rows under {cost_request_id} carry call types {call_types}"
+        unjoined: Final = [
+            (row.call_type, row.api_key, row.metadata.user_api_key_alias if row.metadata else None)
+            for row in rows
+            if row.api_key != driven.identity.token
+            or row.metadata is None
+            or row.metadata.user_api_key_alias != driven.identity.alias
+        ]
+        assert not unjoined, (
+            f"batch cost rows that do not join the retrieving key's token {driven.identity.token} "
+            f"with alias {driven.identity.alias!r}: {unjoined}"
+        )

tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py

@@ -2,6 +2,7 @@ import glob
 import os
 import re
 import sys
+from pathlib import Path
 
 import pytest
 
@@ -870,3 +871,69 @@ class TestMigrateDeployAttemptAccounting:
             harness.run()
         assert len(harness.deploy_calls) == 1
         assert harness.resolved == []
+
+
+class TestJWTKeyMappingCascade:
+    """Regression tests for issue #33702.
+
+    A virtual key referenced by a LiteLLM_JWTKeyMapping row could not be deleted
+    because LiteLLM_JWTKeyMapping_token_fkey was created ON DELETE RESTRICT, so
+    deleting the key (Admin UI, /key/delete, team delete, ...) raised a foreign
+    key violation. The mapping must be removed automatically when its key is
+    deleted, which the FK now enforces via ON DELETE CASCADE.
+    """
+
+    _FK_NAME = "LiteLLM_JWTKeyMapping_token_fkey"
+
+    def _effective_on_delete(self):
+        """Replay every migration in order and return the last ON DELETE action
+        declared for the JWT key mapping FK."""
+        action = None
+        for _migration_name, sql in _get_all_migrations():
+            for match in re.finditer(
+                rf'ADD\s+CONSTRAINT\s+"{re.escape(self._FK_NAME)}".*?'
+                r"ON\s+DELETE\s+(CASCADE|RESTRICT|SET\s+NULL|NO\s+ACTION|SET\s+DEFAULT)",
+                sql,
+                re.IGNORECASE | re.DOTALL,
+            ):
+                action = re.sub(r"\s+", " ", match.group(1).upper())
+        return action
+
+    def test_fk_effective_on_delete_is_cascade(self):
+        """The final FK definition across all migrations must cascade deletes."""
+        assert self._effective_on_delete() == "CASCADE", (
+            f"{self._FK_NAME} must end up ON DELETE CASCADE so deleting a "
+            "virtual key removes its JWT key mapping (issue #33702)"
+        )
+
+    def test_schema_declares_cascade_on_relation(self):
+        """schema.prisma must declare onDelete: Cascade on the mapping relation
+        so the generated client and DB agree."""
+        schema_paths = glob.glob(
+            os.path.abspath(
+                os.path.join(
+                    os.path.dirname(__file__), "../../**/schema.prisma"
+                )
+            ),
+            recursive=True,
+        )
+        declaring = tuple(
+            (path, schema)
+            for path, schema in ((p, Path(p).read_text()) for p in schema_paths)
+            if "model LiteLLM_JWTKeyMapping" in schema
+        )
+        assert declaring, "No schema.prisma declaring LiteLLM_JWTKeyMapping found"
+        for path, schema in declaring:
+            match = re.search(
+                r"litellm_verification_token\s+LiteLLM_VerificationToken\s+@relation\(([^)]*)\)",
+                schema,
+            )
+            assert match is not None, (
+                f"{path} declares LiteLLM_JWTKeyMapping but its verification token "
+                "relation could not be parsed, so this test cannot vouch for it "
+                "(issue #33702)"
+            )
+            assert "onDelete: Cascade" in match.group(1), (
+                f"{path} must declare onDelete: Cascade on the JWT key mapping "
+                "relation (issue #33702)"
+            )

tests/proxy_unit_tests/test_custom_tokenizer_bug.py

@@ -23,9 +23,9 @@ from litellm.proxy.proxy_server import token_counter
 
 def _fake_hf_tokenizer(num_tokens: int) -> MagicMock:
     encoding = MagicMock()
-    encoding.ids = list(range(num_tokens))
+    encoding.__len__.return_value = num_tokens
     tokenizer = MagicMock()
-    tokenizer.encode.return_value = encoding
+    tokenizer.encode_batch_fast.return_value = [encoding]
     return tokenizer
 
 
@@ -68,13 +68,11 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch):
             )
         )
 
-    mock_tokenizer_cls.from_pretrained.assert_called_once_with(
-        "my-org/custom-tokenizer", revision="v2", auth_token=None
-    )
+    mock_tokenizer_cls.from_pretrained.assert_called_once_with("my-org/custom-tokenizer", revision="v2", token=None)
     assert response.tokenizer_type == "huggingface_tokenizer"
     assert response.request_model == "my-embedding-model"
     assert response.model_used == "self-hosted-embedder"
-    assert response.total_tokens > 0
+    assert response.total_tokens >= 7
 
 
 @pytest.mark.asyncio

tests/proxy_unit_tests/test_response_polling_handler.py

@@ -14,13 +14,14 @@ These tests ensure the polling handler correctly manages response state
 following the OpenAI Response API format.
 """
 
+import asyncio
 import json
 from datetime import datetime, timezone
 from typing import Any, Dict, Optional
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-
+from fastapi import Request
 
 from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
 
@@ -1414,7 +1415,7 @@ def _make_background_streaming_kwargs(
         polling_id=polling_id,
         data={"model": "gpt-4o", "stream": False, "background": True},
         polling_handler=polling_handler,
-        request=Mock(),
+        request=Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}),
         fastapi_response=Mock(),
         user_api_key_dict=Mock(),
         general_settings={},
@@ -1663,6 +1664,63 @@ class TestBackgroundStreamingTerminalEvents:
         final_call = handler.update_state.call_args_list[-1]
         assert final_call.kwargs["status"] == "completed"
 
+    @pytest.mark.asyncio
+    async def test_polling_client_disconnect_does_not_cancel_upstream_call(self):
+        """The polling client hangs up right after getting its polling id. The detached task
+        must still stream the upstream response through the client-disconnect guards."""
+        from litellm.proxy.common_request_processing import create_response
+        from litellm.proxy.response_polling.background_streaming import (
+            background_streaming_task,
+        )
+
+        async def client_already_left():
+            return {"type": "http.disconnect"}
+
+        async def slow_upstream_stream():
+            await asyncio.sleep(0.05)
+            for event in (
+                {"type": "response.in_progress"},
+                {
+                    "type": "response.completed",
+                    "response": {
+                        "id": "resp_123",
+                        "status": "completed",
+                        "usage": {"input_tokens": 13, "output_tokens": 10},
+                        "model": "gpt-4o",
+                        "output": [{"id": "item_1", "type": "message"}],
+                    },
+                },
+            ):
+                yield f"data: {json.dumps(event)}\n\n"
+
+        async def upstream_call_behind_disconnect_guard(**kwargs):
+            return await create_response(
+                slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"]
+            )
+
+        handler = AsyncMock(spec=ResponsePollingHandler)
+        kwargs = _make_background_streaming_kwargs("poll_7", handler)
+        kwargs["request"] = Request(
+            {
+                "type": "http",
+                "method": "POST",
+                "path": "/v1/responses",
+                "headers": [(b"x-litellm-call-id", b"call-123")],
+                "query_string": b"",
+            },
+            client_already_left,
+        )
+
+        with patch(  # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests
+            "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing"
+        ) as MockProcessor:
+            MockProcessor.return_value.base_process_llm_request = upstream_call_behind_disconnect_guard
+            await background_streaming_task(**kwargs)
+
+        final_call = handler.update_state.call_args_list[-1]
+        assert final_call.kwargs["status"] == "completed"
+        assert final_call.kwargs["usage"] == {"input_tokens": 13, "output_tokens": 10}
+
 
 class TestEdgeCases:
     """Test edge cases and error scenarios"""

tests/test_litellm/litellm_core_utils/event_loop_lag.py

@@ -0,0 +1,40 @@
+import asyncio
+import time
+from collections.abc import Awaitable, Callable
+from typing import Final, TypeVar
+
+import litellm
+
+T = TypeVar("T")
+
+
+def warm_tokenizer(model: str) -> None:
+    litellm.token_counter(model=model, text="load the tokenizer before anything is timed")
+
+
+async def loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]:
+    async def wake_lag() -> float:
+        started: Final = time.perf_counter()
+        await asyncio.sleep(0.001)
+        return time.perf_counter() - started - 0.001
+
+    return tuple([await wake_lag() for _ in iter(until.is_set, True)])
+
+
+async def timed_with_loop_lags(run: Callable[[], Awaitable[T]]) -> tuple[T, float, tuple[float, ...]]:
+    finished: Final = asyncio.Event()
+
+    async def timed() -> tuple[T, float]:
+        await asyncio.sleep(0)
+        started: Final = time.perf_counter()
+        try:
+            return await run(), time.perf_counter() - started
+        finally:
+            finished.set()
+
+    (result, took), lags = await asyncio.gather(timed(), loop_wake_lags(finished))
+    return result, took, lags
+
+
+def assert_loop_stayed_free(took: float, lags: tuple[float, ...]) -> None:
+    assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count"

tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py

@@ -1,5 +1,6 @@
-import os
+import json
 from collections.abc import Mapping, Sequence
+from pathlib import Path
 
 import pytest
 
@@ -11,8 +12,6 @@ from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebS
 from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams
 
 
-
-
 def test_web_search_cost_low():
     web_search_options = WebSearchOptions(search_context_size="low")
     model_info = litellm.get_model_info("gpt-4o-search-preview")
@@ -683,12 +682,13 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model(
 
 
 def _openai_responses_with_web_search_calls(model, num_calls):
-    from litellm.types.llms.openai import ResponsesAPIResponse
     from openai.types.responses.response_function_web_search import (
         ActionSearch,
         ResponseFunctionWebSearch,
     )
 
+    from litellm.types.llms.openai import ResponsesAPIResponse
+
     output = [
         ResponseFunctionWebSearch(
             id=f"ws_{i}",
@@ -859,11 +859,62 @@ def test_dated_search_preview_entries_carry_search_pricing(local_model_cost_map)
         custom_llm_provider="openai",
         standard_built_in_tools_params=None,
     )
-    assert cost == pytest.approx(0.035), (
-        f"dated search-preview id must bill the $0.035 search fee, got ${cost}"
+    assert cost == pytest.approx(0.025), (
+        f"dated search-preview id must bill the $0.025 search fee, got ${cost}"
     )
 
 
+@pytest.mark.parametrize(
+    "web_search_options",
+    [
+        None,
+        WebSearchOptions(search_context_size="low"),
+        WebSearchOptions(search_context_size="medium"),
+        WebSearchOptions(search_context_size="high"),
+    ],
+)
+def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias(
+    web_search_options: WebSearchOptions | None, local_model_cost_map: None
+) -> None:
+    alias_info = litellm.get_model_info("gpt-4o-mini")
+    snapshot_info = litellm.get_model_info("gpt-4o-mini-2024-07-18")
+
+    assert not snapshot_info["supports_web_search"]
+    assert not alias_info["supports_web_search"]
+
+    snapshot_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search(
+        web_search_options=web_search_options, model_info=snapshot_info
+    )
+    alias_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search(
+        web_search_options=web_search_options, model_info=alias_info
+    )
+
+    assert snapshot_cost == alias_cost == 0.025
+
+
+def test_gpt_4o_mini_web_search_price_matches_in_both_cost_maps():
+    repo_root = Path(__file__).parents[4]
+    cost_maps = tuple(
+        json.loads((repo_root / path).read_text(encoding="utf-8"))
+        for path in (
+            "model_prices_and_context_window.json",
+            "litellm/model_prices_and_context_window_backup.json",
+        )
+    )
+    canonical, backup = cost_maps
+    expected_search_price = {
+        "search_context_size_low": 0.025,
+        "search_context_size_medium": 0.025,
+        "search_context_size_high": 0.025,
+    }
+    for model_name in ("gpt-4o-mini", "gpt-4o-mini-2024-07-18"):
+        canonical_entry = canonical[model_name]
+        backup_entry = backup[model_name]
+        assert canonical_entry["search_context_cost_per_query"] == expected_search_price
+        assert backup_entry["search_context_cost_per_query"] == expected_search_price
+        assert canonical_entry == backup_entry
+
+
 # Note: File search integration test removed due to complex annotation detection logic
 # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage
 

tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py

@@ -1,3 +1,4 @@
+import copy
 import functools
 import json
 import os
@@ -7,14 +8,20 @@ import pytest
 
 
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    ENCRYPTED_REASONING_SIGNATURE_PREFIX,
     TOOL_RESULT_IMAGE_BOUNDARY,
     TOOL_RESULT_IMAGE_PLACEHOLDER,
     add_system_prompt_to_messages,
+    encrypted_content_from_signature,
+    encrypted_reasoning_signature,
     get_file_ids_from_messages,
     get_format_from_file_id,
     handle_any_messages_to_chat_completion_str_messages_conversion,
     hoist_images_from_tool_messages,
+    is_encrypted_reasoning_block,
+    responses_reasoning_items_from_thinking_blocks,
     split_concatenated_json_objects,
+    strip_encrypted_reasoning_from_messages,
     update_messages_with_model_file_ids,
 )
 
@@ -1554,3 +1561,117 @@ class TestRequestContainsImageContent:
         for _ in range(50):
             nested = {"type": "tool_result", "content": [nested]}
         assert request_contains_image_content([{"role": "user", "content": [nested]}]) is False
+
+
+class TestEncryptedReasoningReplay:
+    """Regression for https://github.com/BerriAI/litellm/issues/40288."""
+
+    def test_signature_round_trips_the_encrypted_content(self):
+        assert encrypted_content_from_signature(encrypted_reasoning_signature("gAAAA_bytes")) == "gAAAA_bytes"
+
+    @pytest.mark.parametrize(
+        "signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7]
+    )
+    def test_anything_else_is_not_encrypted_content(self, signature):
+        assert encrypted_content_from_signature(signature) is None
+
+    def test_encrypted_thinking_block_replays_its_own_item(self):
+        items = responses_reasoning_items_from_thinking_blocks(
+            [{"type": "thinking", "thinking": "Plan.", "signature": encrypted_reasoning_signature("gAAAA_1")}]
+        )
+        assert items == (
+            {
+                "type": "reasoning",
+                "summary": [{"type": "summary_text", "text": "Plan."}],
+                "encrypted_content": "gAAAA_1",
+            },
+        )
+
+    def test_encrypted_redacted_block_replays_with_an_empty_summary(self):
+        items = responses_reasoning_items_from_thinking_blocks(
+            [{"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_1")}]
+        )
+        assert items == ({"type": "reasoning", "summary": [], "encrypted_content": "gAAAA_1"},)
+
+    def test_plain_blocks_collapse_into_one_summary_item_around_encrypted_ones(self):
+        items = responses_reasoning_items_from_thinking_blocks(
+            [
+                {"type": "thinking", "thinking": "A.", "signature": None},
+                {"type": "thinking", "thinking": "B.", "signature": ""},
+                {"type": "thinking", "thinking": "C.", "signature": encrypted_reasoning_signature("gAAAA_c")},
+                {"type": "redacted_thinking", "data": "anthropic-minted-opaque-data"},
+                {"type": "thinking", "thinking": "D."},
+            ]
+        )
+        assert items == (
+            {
+                "type": "reasoning",
+                "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}],
+            },
+            {"type": "reasoning", "summary": [{"type": "summary_text", "text": "C."}], "encrypted_content": "gAAAA_c"},
+            {"type": "reasoning", "summary": [{"type": "summary_text", "text": "D."}]},
+        )
+        assert all("id" not in item for item in items)
+
+    def test_blocks_without_text_or_encrypted_content_produce_nothing(self):
+        assert responses_reasoning_items_from_thinking_blocks([{"type": "thinking", "thinking": ""}]) == ()
+        assert responses_reasoning_items_from_thinking_blocks([]) == ()
+
+    @pytest.mark.parametrize(
+        ("block", "expected"),
+        [
+            ({"type": "thinking", "thinking": "x", "signature": encrypted_reasoning_signature("g")}, True),
+            ({"type": "redacted_thinking", "data": encrypted_reasoning_signature("g")}, True),
+            ({"type": "thinking", "thinking": "x", "signature": ENCRYPTED_REASONING_SIGNATURE_PREFIX}, True),
+            ({"type": "redacted_thinking", "data": ENCRYPTED_REASONING_SIGNATURE_PREFIX}, True),
+            ({"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}, False),
+            ({"type": "redacted_thinking", "data": "EmwKAhgBEgy"}, False),
+            ({"type": "text", "text": encrypted_reasoning_signature("g")}, False),
+            ("not a block", False),
+        ],
+    )
+    def test_is_encrypted_reasoning_block(self, block, expected):
+        assert is_encrypted_reasoning_block(block) is expected
+
+    def test_strip_drops_every_bridge_tagged_block_and_leaves_no_unsigned_thinking_behind(self):
+        assistant_content = [
+            {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"},
+            {"type": "thinking", "thinking": "packed by the bridge", "signature": encrypted_reasoning_signature("g1")},
+            {"type": "redacted_thinking", "data": encrypted_reasoning_signature("g2")},
+            {"type": "thinking", "thinking": "", "signature": encrypted_reasoning_signature("g3")},
+            {"type": "text", "text": "answer"},
+        ]
+        messages = [
+            {"role": "user", "content": "question"},
+            {"role": "assistant", "content": assistant_content},
+            {"role": "user", "content": [{"type": "text", "text": "follow-up"}]},
+        ]
+
+        strip_encrypted_reasoning_from_messages(messages)
+
+        assert messages[1]["content"] is assistant_content
+        assert assistant_content == [
+            {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"},
+            {"type": "text", "text": "answer"},
+        ]
+        assert all(block["signature"] for block in assistant_content if block["type"] == "thinking")
+        assert messages[0] == {"role": "user", "content": "question"}
+        assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]}
+
+    @pytest.mark.parametrize(
+        "messages",
+        [
+            "not a list",
+            None,
+            [{"role": "user", "content": None}],
+            [{"role": "user", "content": "plain string"}],
+            ["not a message"],
+            [{"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}]}],
+        ],
+    )
+    def test_strip_leaves_history_without_bridge_reasoning_untouched(self, messages):
+        before = copy.deepcopy(messages)
+
+        strip_encrypted_reasoning_from_messages(messages)
+
+        assert messages == before

tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py

@@ -191,8 +191,16 @@ def test_bedrock_converse_assistant_with_empty_thinking_block_and_tool_calls():
         {"type": "thinking", "thinking": "oss reasoning", "signature": None},
         {"type": "thinking", "thinking": "oss reasoning", "signature": ""},
         {"type": "thinking", "thinking": "oss reasoning"},
+        {"type": "thinking", "thinking": "openai reasoning", "signature": "litellm_encrypted_reasoning:gAAAA"},
+        {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:gAAAA"},
+    ],
+    ids=[
+        "null_signature",
+        "empty_signature",
+        "missing_signature",
+        "encrypted_reasoning_signature",
+        "encrypted_reasoning_redacted_data",
     ],
-    ids=["null_signature", "empty_signature", "missing_signature"],
 )
 def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block):
     """Open-source reasoning models (DeepSeek-R1, Qwen, etc.) emit thinking blocks
@@ -219,7 +227,7 @@ def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block):
     assistant = next(m for m in result if m["role"] == "assistant")
     content = assistant["content"]
     assert all(
-        block.get("type") != "thinking" for block in content
+        block.get("type") not in ("thinking", "redacted_thinking") for block in content
     ), f"unsignable thinking block must be dropped, got {content!r}"
     assert any(
         block.get("type") == "text" and block.get("text") == "2+2 equals 4."

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -1,10 +1,15 @@
 #### What this tests ####
 #    This tests litellm.token_counter.token_counter() function
+import asyncio
 import importlib
+import threading
 import time
 import traceback
+from concurrent.futures import Future, wait
+from typing import Final
 from unittest.mock import MagicMock
 
+import anyio.to_thread
 import pytest
 import tiktoken
 
@@ -14,9 +19,21 @@ import litellm
 from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens
 from litellm import token_counter as token_counter_old
 import litellm.constants
-from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function
+from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+from litellm.litellm_core_utils.asyncify import asyncify
+from litellm.litellm_core_utils.token_counter import (
+    _get_exact_count_function,
+    _get_extrapolating_count_function,
+    _get_tiktoken_count_function,
+    offload_token_count,
+)
 from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
 from tests.large_text import text
+from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+    assert_loop_stayed_free,
+    timed_with_loop_lags,
+    warm_tokenizer,
+)
 from tests.test_litellm.litellm_core_utils.messages_with_counts import (
     MESSAGES_TEXT,
     MESSAGES_WITH_IMAGES,
@@ -120,6 +137,135 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch):
         importlib.reload(litellm.constants)
 
 
+async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free():
+    warm_tokenizer("claude-fable-5")
+
+    tokens, took, lags = await timed_with_loop_lags(
+        lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100)
+    )
+
+    assert tokens > 0
+    assert_loop_stayed_free(took, lags)
+
+
+@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500])
+def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int):
+    count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk))
+    front_heavy: Final = "a" * 1_000 + "b" * 4_000
+    exact: Final = 1_000 + len(front_heavy)
+
+    estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy)
+
+    assert abs(estimate - exact) <= exact // 100
+    assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars
+
+
+def test_count_at_or_below_the_cap_is_exact():
+    count_exactly: Final = MagicMock(side_effect=len)
+
+    assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000
+    assert count_exactly.call_args_list == [(("a" * 5_000,),)]
+
+
+class _SlowEncoder:
+    def __init__(self) -> None:
+        self._lock: Final = threading.Lock()
+        self.in_flight = 0
+        self.peak_in_flight = 0
+
+    def encode_batch_fast(self, texts: list[str]) -> list[list[int]]:
+        with self._lock:
+            self.in_flight += 1
+            self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
+        time.sleep(0.1)
+        with self._lock:
+            self.in_flight -= 1
+        return [[0] * len(text) for text in texts]
+
+
+@pytest.mark.asyncio
+async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool():
+    encoder: Final = _SlowEncoder()
+    count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder})
+    shared_pool: Final = anyio.to_thread.current_default_thread_limiter()
+    burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+
+    async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]:
+        if counting.done():
+            return ()
+        await asyncio.sleep(0.01)
+        return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting))
+
+    counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst))))
+    borrowed: Final = await shared_pool_borrowed_until_done(counting)
+
+    assert await counting == [3] * burst
+    assert len(borrowed) > 1 and max(borrowed) == 0
+    assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+
+
+def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None:
+    def slow_count(counted: str) -> int:
+        time.sleep(0.1)
+        return len(counted)
+
+    result.set_result(asyncio.run(offload_token_count(slow_count)(text)))
+
+
+def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process():
+    loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+    results: Final = tuple(Future[int]() for _ in range(loops))
+    threads: Final = tuple(
+        threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True)
+        for size, result in enumerate(results, start=1)
+    )
+    for thread in threads:
+        thread.start()
+
+    _, pending = wait(results, timeout=5)
+
+    assert not pending
+    assert tuple(result.result() for result in results) == tuple(range(1, loops + 1))
+
+
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [("8", 8), ("0", 4), ("not-an-int", 4)],
+)
+def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int):
+    monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured)
+    try:
+        assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected
+    finally:
+        monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS")
+        importlib.reload(litellm.constants)
+
+
+def test_token_counter_applies_the_default_cap():
+    max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS
+    prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars]
+    over_the_cap: Final = prose + "a" * 200_000
+    exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap)
+
+    estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap)
+
+    assert estimate != exact
+    assert abs(estimate - exact) <= exact // 100
+
+
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)],
+)
+def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int):
+    monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured)
+    try:
+        assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected
+    finally:
+        monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS")
+        importlib.reload(litellm.constants)
+
+
 def test_token_counter_with_prefix():
     messages = [
         {"role": "user", "content": "Who won the world cup in 2022?"},

tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py

@@ -2270,3 +2270,29 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
         assert open_key == StreamingScanKey(texts=("hi",))
         assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
         assert ended_key != open_key
+
+
+class TestAnthropicMessagesHandlerPostCallHookResponse:
+    def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self):
+        from litellm.types.utils import Choices, Message, ModelResponse, Usage
+
+        assembled = ModelResponse(
+            id="msg_1",
+            model="claude",
+            choices=[Choices(message=Message(role="assistant", content="hello world"), finish_reason="stop")],
+            usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
+        )
+
+        hook_response = AnthropicMessagesHandler().post_call_hook_response(assembled)
+
+        assert hook_response["type"] == "message"
+        assert hook_response["role"] == "assistant"
+        assert hook_response["content"] == [{"type": "text", "text": "hello world"}]
+        assert hook_response["stop_reason"] == "end_turn"
+        assert hook_response["usage"]["input_tokens"] == 1
+        assert hook_response["usage"]["output_tokens"] == 2
+
+    def test_anything_else_reaches_the_hook_untouched(self):
+        native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]}
+
+        assert AnthropicMessagesHandler().post_call_hook_response(native) is native

tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py

@@ -10,6 +10,7 @@ import litellm
 
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     TOOL_RESULT_IMAGE_PLACEHOLDER,
+    encrypted_reasoning_signature,
 )
 from litellm.litellm_core_utils.prompt_templates.factory import (
     THOUGHT_SIGNATURE_SEPARATOR,
@@ -423,6 +424,43 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks():
     assert result[1]["tool_calls"][0]["id"] == "toolu_01234"
 
 
+def test_translate_anthropic_messages_to_openai_drops_bridge_encrypted_reasoning_blocks():
+    """A session that moves from an OpenAI reasoning model to a chat provider replays reasoning only OpenAI can read.
+
+    Gemini rejects the whole request when such a block reaches it as a thought_signature, so the
+    adapter drops those blocks and keeps the provider-signed ones.
+    """
+
+    anthropic_messages = [
+        AnthropicMessagesUserMessageParam(
+            role="user",
+            content=[{"type": "text", "text": "Who drinks water?"}],
+        ),
+        AnthopicMessagesAssistantMessageParam(
+            role="assistant",
+            content=[
+                {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")},
+                {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")},
+                {"type": "text", "text": "The Norwegian."},
+            ],
+        ),
+        AnthopicMessagesAssistantMessageParam(
+            role="assistant",
+            content=[
+                {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_signed"},
+                {"type": "text", "text": "Still the Norwegian."},
+            ],
+        ),
+    ]
+
+    result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=anthropic_messages)
+
+    assert [m["role"] for m in result] == ["user", "assistant", "assistant"]
+    assert not result[1].get("thinking_blocks")
+    assert result[1]["content"] == "The Norwegian."
+    assert [b["signature"] for b in result[2]["thinking_blocks"]] == ["EqQBCkYIAxgCIkA_signed"]
+
+
 def test_translate_anthropic_messages_to_openai_sets_reasoning_content():
     """Reasoning-aware chat providers read reasoning_content, so thinking text must land there.
 

tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py

@@ -0,0 +1,50 @@
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    encrypted_reasoning_signature,
+)
+from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
+    AnthropicMessagesConfig,
+)
+
+
+def _transform(messages):
+    return AnthropicMessagesConfig().transform_anthropic_messages_request(
+        model="claude-sonnet-4-5",
+        messages=messages,
+        anthropic_messages_optional_request_params={"max_tokens": 1024},
+        litellm_params={},
+        headers={},
+    )
+
+
+def test_reasoning_replayed_from_the_responses_bridge_never_reaches_anthropic():
+    """Claude Code resumed on a Claude model echoes the thinking blocks a gpt turn produced."""
+    messages = [
+        {"role": "user", "content": "Solve it."},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")},
+                {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")},
+                {"type": "text", "text": "The answer."},
+            ],
+        },
+        {"role": "user", "content": "And the next one?"},
+    ]
+    request = _transform(messages)
+    assert request["messages"][1]["content"] == [{"type": "text", "text": "The answer."}]
+    assert len(messages[1]["content"]) == 3
+
+
+def test_anthropic_signed_thinking_blocks_are_forwarded_untouched():
+    messages = [
+        {"role": "user", "content": "Solve it."},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "thinking", "thinking": "plan", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"},
+                {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"},
+                {"type": "text", "text": "The answer."},
+            ],
+        },
+    ]
+    assert _transform(messages)["messages"] == messages

tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py

@@ -67,6 +67,39 @@ def test_build_responses_kwargs_prefers_explicit_prompt_cache_key_over_derived()
     assert responses_kwargs["prompt_cache_key"] == "explicit-key"
 
 
+def test_build_responses_kwargs_asks_openai_for_encrypted_reasoning_without_thinking():
+    responses_kwargs = _build_responses_kwargs(
+        max_tokens=1024,
+        messages=MESSAGES,
+        model="openai/gpt-5.6-luna",
+        extra_kwargs={"custom_llm_provider": "openai"},
+    )
+    assert responses_kwargs["include"] == ["reasoning.encrypted_content"]
+    assert "reasoning" not in responses_kwargs
+
+
+def test_build_responses_kwargs_skips_include_for_a_responses_provider_that_rejects_it():
+    responses_kwargs = _build_responses_kwargs(
+        max_tokens=1024,
+        messages=MESSAGES,
+        model="perplexity/sonar",
+        thinking={"type": "enabled", "budget_tokens": 4096},
+        extra_kwargs={"custom_llm_provider": "perplexity"},
+    )
+    assert "include" not in responses_kwargs
+    assert "reasoning" in responses_kwargs
+
+
+def test_build_responses_kwargs_keeps_the_deployment_include_next_to_encrypted_reasoning():
+    responses_kwargs = _build_responses_kwargs(
+        max_tokens=1024,
+        messages=MESSAGES,
+        model="openai/gpt-5.6-luna",
+        extra_kwargs={"custom_llm_provider": "openai", "include": ["file_search_call.results"]},
+    )
+    assert responses_kwargs["include"] == ["reasoning.encrypted_content", "file_search_call.results"]
+
+
 def test_build_responses_kwargs_without_metadata_sets_no_prompt_cache_key():
     responses_kwargs = _build_responses_kwargs(
         max_tokens=1024,