diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index bdfddb48522..1b659770181 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -21,6 +21,7 @@ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-11-20": "mcp-client-2025-11-20", "mcp-client-2025-04-04": "mcp-client-2025-04-04", + "message-threads-2026-08-12": "message-threads-2026-08-12", "mcp-servers-2025-12-04": null, "oauth-2025-04-20": "oauth-2025-04-20", "output-128k-2025-02-19": "output-128k-2025-02-19", @@ -56,7 +57,6 @@ "mcp-client-2026-09-15": "mcp-client-2026-09-15", "mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22", "message-batches-2024-09-24": "message-batches-2024-09-24", - "message-threads-2026-08-12": "message-threads-2026-08-12", "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", "model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26", "output-300k-2026-03-24": "output-300k-2026-03-24", @@ -98,6 +98,7 @@ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-11-20": "mcp-client-2025-11-20", "mcp-client-2025-04-04": "mcp-client-2025-04-04", + "message-threads-2026-08-12": "message-threads-2026-08-12", "mcp-servers-2025-12-04": null, "output-128k-2025-02-19": "output-128k-2025-02-19", "structured-output-2024-03-01": null, @@ -132,7 +133,6 @@ "mcp-client-2026-09-15": "mcp-client-2026-09-15", "mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22", "message-batches-2024-09-24": "message-batches-2024-09-24", - "message-threads-2026-08-12": "message-threads-2026-08-12", "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01", "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", @@ -177,6 +177,7 @@ "interleaved-thinking-2025-05-14": null, "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, + "message-threads-2026-08-12": null, "mcp-servers-2025-12-04": null, "output-128k-2025-02-19": null, "structured-output-2024-03-01": null, @@ -212,6 +213,7 @@ "interleaved-thinking-2025-05-14": null, "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, + "message-threads-2026-08-12": null, "mcp-servers-2025-12-04": null, "output-128k-2025-02-19": null, "structured-output-2024-03-01": null, @@ -249,6 +251,7 @@ "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-04-04": null, + "message-threads-2026-08-12": null, "mcp-client-2025-11-20": null, "mcp-servers-2025-12-04": null, "output-128k-2025-02-19": "output-128k-2025-02-19", @@ -286,6 +289,7 @@ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, + "message-threads-2026-08-12": null, "mcp-servers-2025-12-04": null, "output-128k-2025-02-19": "output-128k-2025-02-19", "structured-output-2024-03-01": null, @@ -321,7 +325,6 @@ "mcp-client-2026-09-15": null, "mcp-tunnels-2026-06-22": null, "message-batches-2024-09-24": null, - "message-threads-2026-08-12": null, "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01", "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", @@ -386,6 +389,7 @@ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, + "message-threads-2026-08-12": null, "mcp-servers-2025-12-04": null, "oauth-2025-04-20": null, "output-128k-2025-02-19": "output-128k-2025-02-19", diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index b4c95fafff3..058ac4dcf4c 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -5,15 +5,18 @@ Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthrop """ import json +from collections.abc import Mapping from typing import Final -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from .exceptions import AnthropicErrorDetail, AnthropicErrorResponse, AnthropicErrorType _PARSED_ERROR_BODY: Final = TypeAdapter(object) +_ERROR_DETAILS: Final = TypeAdapter(Mapping[str, object]) +_JSON_DECODER: Final = json.JSONDecoder() # HTTP status code -> Anthropic error type # Source: https://platform.claude.com/docs/en/api/errors @@ -166,12 +169,41 @@ class AnthropicExceptionMapping: else: message = raw_message - return AnthropicExceptionMapping.create_error_response( + response: Final = AnthropicExceptionMapping.create_error_response( status_code=status_code, message=message, request_id=request_id, error_type=error_type, ) + details: Final = AnthropicExceptionMapping.embedded_error_details(raw_message) + if details is None: + return response + with_details: Final[AnthropicErrorResponse] = {**response, "error": {**response["error"], "details": details}} + return with_details + + @staticmethod + def embedded_error_details(raw_message: str) -> Mapping[str, object] | None: + """ + ``error.details`` of the Anthropic error body quoted inside ``raw_message``, when there is one. + + litellm wraps an upstream error as ``" - {upstream json}\n\nLiteLLM: ..."``. Anthropic + puts machine-readable codes such as ``thread_not_found`` under ``error.details``, and the SDKs + branch on them, so the code is lifted next to the wrapped message instead of staying only inside it. + """ + start: Final = raw_message.find("{") + if start < 0: + return None + try: + embedded, _ = _JSON_DECODER.raw_decode(raw_message, start) + except ValueError: + return None + try: + envelope: Final = _ERROR_DETAILS.validate_python(embedded, strict=True) + error: Final = _ERROR_DETAILS.validate_python(envelope.get("error"), strict=True) + details: Final = _ERROR_DETAILS.validate_python(error.get("details"), strict=True) + except ValidationError: + return None + return details if envelope.get("type") == "error" else None class AnthropicErrorSseFrame(str): diff --git a/litellm/anthropic_interface/exceptions/exceptions.py b/litellm/anthropic_interface/exceptions/exceptions.py index 18b3ca37a3b..71d4de0029d 100644 --- a/litellm/anthropic_interface/exceptions/exceptions.py +++ b/litellm/anthropic_interface/exceptions/exceptions.py @@ -27,6 +27,7 @@ class AnthropicErrorDetail(TypedDict): type: AnthropicErrorType message: str + details: NotRequired[ReadOnly[Mapping[str, object]]] provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]] litellm_error_type: NotRequired[ReadOnly[str]] litellm_call_id: NotRequired[ReadOnly[str]] diff --git a/litellm/constants.py b/litellm/constants.py index 931ac4f2b81..5d0eab30b90 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2374,6 +2374,7 @@ ADMIN_ONLY_BATCH_LIMIT_METADATA_KEYS: Final = ( MAX_FILE_DOWNLOADS_PER_MINUTE_KEY, ) FILE_USAGE_MAX_TRACKED_COUNTERS: Final = 20_000 +MESSAGE_THREAD_MAX_TRACKED_OWNERS: Final = 20_000 # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 59ffb6fef04..08c3e0b48a8 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -15,10 +15,11 @@ from typing import Final, cast from pydantic import TypeAdapter, ValidationError import litellm +from litellm._logging import verbose_logger from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.constants import PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.reasoning_effort_capability import ( @@ -29,6 +30,7 @@ from litellm.litellm_core_utils.reasoning_effort_capability import ( resolve_deployment_reasoning_effort, ) from litellm.llms.anthropic.common_utils import ( + AnthropicError, AnthropicModelInfo, strip_provider_specific_fields_from_anthropic_messages, tools_marked_for_a_tool_config_point, @@ -58,7 +60,14 @@ from ..safeguards import ( with_safeguard_results, ) from .interceptors import get_messages_interceptors -from .utils import AnthropicMessagesRequestUtils, mock_response, prepare_native_messages +from .utils import ( + AnthropicMessagesRequestUtils, + continues_message_thread, + message_thread_not_found_body, + message_thread_unsupported_body, + mock_response, + prepare_native_messages, +) __all__ = ("anthropic_messages", "anthropic_messages_handler") @@ -769,6 +778,32 @@ def anthropic_messages_handler( anthropic_messages_provider_config: Final = _native_messages_config( model, custom_llm_provider, kwargs.get("model_info"), api_base ) + thread: Final[object] = kwargs.get("thread") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # **kwargs of the handler is untyped + if thread is not None and ( + anthropic_messages_provider_config is None or not anthropic_messages_provider_config.supports_message_threads() + ): + if continues_message_thread(thread): + raise AnthropicError( + status_code=404, + message=message_thread_not_found_body( + f"message threads are not supported for provider {custom_llm_provider}, " + "so this thread cannot be continued. Send the full conversation to start over" + ), + ) + if not drops_unsupported_params(litellm_params.drop_params): + raise AnthropicError( + status_code=400, + message=message_thread_unsupported_body( + f"message threads are not supported for provider {custom_llm_provider}. " + "Send the request without thread, or set drop_params to serve it without a server-side thread" + ), + ) + verbose_logger.warning( + "Dropping thread for provider %s: message threads are not supported there, " + "so the request is served without a server-side thread", + custom_llm_provider, + ) + kwargs.pop("thread") # pyright: ignore[reportUnknownMemberType] # **kwargs of the handler is untyped if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. if kwargs.get("compaction") is None and _should_route_to_responses_api( diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index aa0693d6f03..b0558292271 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -59,6 +59,7 @@ from .mid_conversation_system import ( convert_mid_conversation_system_turns, is_system_role_message, ) +from .utils import continues_message_thread DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -137,6 +138,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): "output_config", "reasoning_effort", "safeguards", + "thread", # TODO: Add Anthropic `metadata` support # "metadata", ] @@ -144,6 +146,11 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return self._resolved_provider != "anthropic" + def supports_message_threads(self) -> bool: + return is_beta_header_supported_by_bundled_config( + ANTHROPIC_BETA_HEADER_VALUES.MESSAGE_THREADS_2026_08_12.value, self._resolved_provider + ) + def supports_safeguards(self) -> bool: """ Whether the upstream platform answers the ``safeguards`` body field itself. A platform @@ -771,10 +778,13 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # Auto-strip advisor blocks from history if advisor tool is absent. # Prevents Anthropic 400: advisor_tool_result in history requires advisor tool. + # A thread continue turn omits tools by contract and the stored tools still apply, so it keeps them. tools: Final = anthropic_messages_optional_request_params.get("tools") _tools: Final = tools or [] _has_advisor: Final = any(isinstance(t, dict) and t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE for t in _tools) - if not _has_advisor: + if not _has_advisor and not continues_message_thread( + anthropic_messages_optional_request_params.get("thread") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # the request params arrive as a bare dict + ): messages = strip_advisor_blocks_from_messages(messages) gated: Final = ( @@ -797,7 +807,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ), max_tokens=max_tokens, model=model, - **anthropic_messages_optional_request_params, + **anthropic_messages_optional_request_params, # pyright: ignore[reportUnknownArgumentType] # the request params arrive as a bare dict ) request: Final = ( {**dict(anthropic_messages_request), "thinking": gated.thinking} @@ -877,6 +887,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): drop_params: Whether the deployment or request set drop_params """ beta_values: Final[set[str]] = set() + threaded: Final = optional_params.get("thread") is not None # pyright: ignore[reportUnknownMemberType] # the optional params arrive as a bare dict + scanned_messages: Final = ( + () + if continues_message_thread(optional_params.get("thread")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # the optional params arrive as a bare dict + else messages + ) existing_beta: Final = tuple( piece.strip() @@ -887,7 +903,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ) beta_values.update(existing_beta) - if requires_native_compaction_beta(custom_llm_provider, optional_params, messages): + if requires_native_compaction_beta(custom_llm_provider, optional_params, scanned_messages): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value) beta_values.update(context_management_betas(optional_params.get("context_management"))) @@ -911,9 +927,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value) - if messages_carry_output_config(messages): + if messages_carry_output_config(scanned_messages): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value) + if threaded: + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.MESSAGE_THREADS_2026_08_12.value) + tools: Final = optional_params.get("tools") if any(isinstance(tool, dict) and tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE for tool in tools or ()): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value) @@ -924,7 +943,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): thinking_betas: Final = AnthropicModelInfo().thinking_betas(optional_params.get("thinking")) tool_change_betas: Final = ( (ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,) - if AnthropicModelInfo().is_mid_conversation_tool_change_used(messages) + if AnthropicModelInfo().is_mid_conversation_tool_change_used(scanned_messages) else () ) all_beta_values: Final = beta_values.union(thinking_betas, tool_change_betas) diff --git a/litellm/llms/anthropic/messages/native/utils.py b/litellm/llms/anthropic/messages/native/utils.py index 8f33d000c7c..dfeedf1c191 100644 --- a/litellm/llms/anthropic/messages/native/utils.py +++ b/litellm/llms/anthropic/messages/native/utils.py @@ -1,3 +1,4 @@ +import json from collections.abc import Iterable, Mapping, Sequence from functools import lru_cache from typing import TYPE_CHECKING, Final, cast, get_type_hints @@ -13,7 +14,10 @@ from litellm.llms.anthropic.common_utils import ( ) from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled from litellm.types.llms.anthropic import ( + MESSAGE_THREAD_NOT_FOUND_ERROR_CODE, + MESSAGE_THREAD_UNSUPPORTED_ERROR_CODE, AnthropicMessagesRequestOptionalParams, + AnthropicMessageThread, AnthropicStopDetails, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -279,6 +283,39 @@ def prepare_native_messages( ) +_THREAD_ADAPTER: Final = TypeAdapter(AnthropicMessageThread) + + +def continues_message_thread(thread: object) -> bool: + """True when ``thread`` is a well-formed ``{"type": "continue", ...}`` body field.""" + try: + return _THREAD_ADAPTER.validate_python(thread)["type"] == "continue" + except ValidationError: + return False + + +def continued_message_id(thread: object) -> str | None: + try: + validated: Final = _THREAD_ADAPTER.validate_python(thread) + except ValidationError: + return None + return validated.get("previous_message_id") if validated["type"] == "continue" else None + + +def _message_thread_error_body(error_type: str, error_code: str, message: str) -> str: + return json.dumps( + {"type": "error", "error": {"type": error_type, "message": message, "details": {"error_code": error_code}}} + ) + + +def message_thread_not_found_body(message: str) -> str: + return _message_thread_error_body("not_found_error", MESSAGE_THREAD_NOT_FOUND_ERROR_CODE, message) + + +def message_thread_unsupported_body(message: str) -> str: + return _message_thread_error_body("invalid_request_error", MESSAGE_THREAD_UNSUPPORTED_ERROR_CODE, message) + + @lru_cache(maxsize=1) def anthropic_messages_optional_param_keys() -> frozenset[str]: """ diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index 487e7aca663..3e999467ad5 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -145,6 +145,9 @@ class BaseAnthropicMessagesConfig(ABC): def supports_safeguards(self) -> bool: return True + def supports_message_threads(self) -> bool: + return False + def serves_api_base(self, model: str, api_base: str | None) -> bool: return True diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index a5d825a39d7..cea57e3505d 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -37,6 +37,9 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): def is_native_messages_route(self) -> bool: return False + def supports_message_threads(self) -> bool: + return False + def hoists_wrapped_defer_loading(self, model: str) -> bool: return False diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index 138050ccc9a..49cf8760033 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -16,6 +16,7 @@ from .max_iterations_limiter import ( # noqa: F401 # backwards-compatible pack PROXY_MaxIterationsHandler, _PROXY_MaxIterationsHandler, # pyright: ignore[reportPrivateUsage] # backwards-compatible package export ) +from .message_thread_ownership import MessageThreadOwnership from .parallel_request_limiter import ( # noqa: F401 # backwards-compatible package export PROXY_MaxParallelRequestsHandler, _PROXY_MaxParallelRequestsHandler, # pyright: ignore[reportPrivateUsage] # backwards-compatible package export @@ -39,6 +40,7 @@ PROXY_HOOKS: Final = { "parallel_request_limiter": PROXY_MaxParallelRequestsHandler_v3, "cache_control_check": PROXY_CacheControlCheck, "responses_id_security": ResponsesIDSecurity, + "message_thread_ownership": MessageThreadOwnership, "litellm_skills": SkillsInjectionHook, "max_iterations_limiter": PROXY_MaxIterationsHandler, "max_budget_per_session_limiter": PROXY_MaxBudgetPerSessionHandler, diff --git a/litellm/proxy/hooks/message_thread_ownership.py b/litellm/proxy/hooks/message_thread_ownership.py new file mode 100644 index 00000000000..0f2d9fd1933 --- /dev/null +++ b/litellm/proxy/hooks/message_thread_ownership.py @@ -0,0 +1,319 @@ +import json +import re +from collections.abc import AsyncGenerator, AsyncIterable, Mapping +from typing import TYPE_CHECKING, Final, Protocol + +import httpx +from fastapi import HTTPException +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm._internal_context import with_service_target +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers +from litellm.llms.anthropic.messages.native.utils import continued_message_id, message_thread_not_found_body +from litellm.llms.custom_httpx.sse_lines import split_sse_lines_keepends +from litellm.types.utils import CallTypesLiteral, LLMResponseTypes + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import InternalUsageCache + +_PASS_THROUGH_CALL_TYPE: Final = "pass_through_endpoint" +_ANTHROPIC_BETA_HEADER: Final = "anthropic-beta" +_OWNER_KEY_PREFIX: Final = "anthropic_message_thread_owner:" +_OWNER_TTL_SECONDS: Final = 3600 +_SSE_DATA_FIELD: Final = "data:" +_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n") +_MAX_SCANNED_BYTES: Final = 65536 +_UNOWNED_THREAD_MESSAGE: Final = ( + "message thread not found for this key. Send the full conversation to start a new thread" +) + +_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_MAPPINGS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +class _FirstClaimStore(Protocol): + async def async_set_cache(self, key: str, value: object, *, nx: bool, ttl: int) -> object: ... + + async def async_get_cache(self, key: str) -> object: ... + + +class _ThreadBinding(BaseModel): + model_config = ConfigDict(frozen=True) + + owner: str + anthropic_beta: str | None = None + + +_BINDING_ADAPTER: Final = TypeAdapter(_ThreadBinding) + + +def _as_mapping(value: object) -> Mapping[str, object] | None: + try: + return _MAPPING_ADAPTER.validate_python(value, strict=True) + except ValidationError: + return None + + +def _as_binding(value: object) -> _ThreadBinding | None: + try: + return _BINDING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _anthropic_beta(headers: Mapping[str, object]) -> str | None: + return next( + ( + value + for name, value in headers.items() + if name.lower() == _ANTHROPIC_BETA_HEADER and isinstance(value, str) and value + ), + None, + ) + + +def _sent_anthropic_beta(request_data: Mapping[str, object]) -> str | None: + logging_obj: Final = request_data.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return None + response: Final[object] = logging_obj.model_call_details.get("httpx_response") + if not isinstance(response, httpx.Response): + return None + try: + sent_headers: Final = response.request.headers + except RuntimeError: + return None + return sent_headers.get(_ANTHROPIC_BETA_HEADER) or None + + +def _with_anthropic_beta(extra_headers: object, anthropic_beta: str) -> Mapping[str, object]: + headers: Final = _as_mapping(extra_headers) or {} + return { + **{name: value for name, value in headers.items() if name.lower() != _ANTHROPIC_BETA_HEADER}, + _ANTHROPIC_BETA_HEADER: merge_anthropic_beta_headers(_anthropic_beta(headers), anthropic_beta), + } + + +def _as_mappings(value: object) -> tuple[Mapping[str, object], ...]: + try: + return _MAPPINGS_ADAPTER.validate_python(value) + except ValidationError: + return () + + +def _client_headers_with_anthropic_beta(scoped: Mapping[str, object], anthropic_beta: str) -> Mapping[str, object]: + client_headers: Final = _as_mapping(scoped.get("extra_headers")) + if client_headers is None or _anthropic_beta(client_headers) is None: + return scoped + return {**scoped, "extra_headers": dict(_with_anthropic_beta(client_headers, anthropic_beta))} + + +def _provider_headers_with_anthropic_beta( + provider_specific_header: object, anthropic_beta: str +) -> Mapping[str, object]: + single: Final = _as_mapping(provider_specific_header) + if single is not None: + return {"provider_specific_header": _client_headers_with_anthropic_beta(single, anthropic_beta)} + scoped_headers: Final = _as_mappings(provider_specific_header) + if not scoped_headers: + return {} + return { + "provider_specific_header": tuple( + _client_headers_with_anthropic_beta(scoped, anthropic_beta) for scoped in scoped_headers + ) + } + + +def _owner_key(message_id: str) -> str: + return f"{_OWNER_KEY_PREFIX}{message_id}" + + +def _owner(user_api_key_dict: "UserAPIKeyAuth") -> str: + if user_api_key_dict.token: + return f"key:{user_api_key_dict.token}" + return "identity:" + json.dumps( + [ + user_api_key_dict.user_id, + user_api_key_dict.team_id, + user_api_key_dict.end_user_id, + user_api_key_dict.org_id, + user_api_key_dict.agent_id, + ] + ) + + +def _threads(data: Mapping[str, object]) -> tuple[object, object]: + extra_body: Final = _as_mapping(data.get("extra_body")) or {} + return data.get("thread"), extra_body.get("thread") + + +def _sends_thread(data: Mapping[str, object]) -> bool: + return any(thread is not None for thread in _threads(data)) + + +def _continued_message_ids(data: Mapping[str, object]) -> tuple[str, ...]: + continued: Final = (continued_message_id(thread) for thread in _threads(data)) + return tuple(message_id for message_id in continued if message_id is not None) + + +def _may_continue(binding: _ThreadBinding | None, owner: str, pass_through: bool) -> bool: + if binding is None: + return pass_through + return binding.owner == owner + + +def _response_message_id(response: object) -> str | None: + message: Final = _as_mapping(response) + if message is None or message.get("type") != "message": + return None + message_id: Final = message.get("id") + return message_id if isinstance(message_id, str) and message_id else None + + +def _message_start_id(event: Mapping[str, object]) -> str | None: + return _response_message_id(event.get("message")) if event.get("type") == "message_start" else None + + +def _frame_event(frame: bytes) -> Mapping[str, object] | None: + data_lines: Final = tuple( + line.strip()[len(_SSE_DATA_FIELD) :].strip() + for line in split_sse_lines_keepends(frame.decode("utf-8", errors="ignore")) + if line.strip().startswith(_SSE_DATA_FIELD) + ) + payload: Final = "".join(data_lines) + if not payload or payload == "[DONE]": + return None + try: + return _MAPPING_ADAPTER.validate_json(payload) + except ValidationError: + return None + + +class _StreamedMessageIdScanner: + def __init__(self, finished: bool) -> None: + self._head = b"" + self.finished = finished + + def scan(self, chunk: object) -> str | None: + if isinstance(chunk, (bytes, bytearray)): + return self._scan_encoded(bytes(chunk)) + if isinstance(chunk, str): + return self._scan_encoded(chunk.encode("utf-8")) + event: Final = _as_mapping(chunk) + if event is None: + self.finished = True + return None + return self._settled(event) + + def _settled(self, event: Mapping[str, object]) -> str | None: + if event.get("type") == "ping": + return None + self.finished = True + return _message_start_id(event) + + def _scan_encoded(self, data: bytes) -> str | None: + head: Final = self._head + data + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(head)) + frames: Final = tuple(head[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries)) + events: Final = (event for frame in frames if (event := _frame_event(frame)) is not None) + decider: Final = next((event for event in events if event.get("type") != "ping"), None) + if decider is not None: + return self._settled(decider) + self._head = head[boundaries[-1] :] if boundaries else head + if len(self._head) > _MAX_SCANNED_BYTES: + self.finished = True + return None + + +class MessageThreadOwnership(CustomLogger): + def __init__(self, message_thread_owner_cache: "InternalUsageCache") -> None: + self._owners: Final = message_thread_owner_cache + + async def async_pre_call_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + cache: "DualCache", + data: Mapping[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object] | None: # mutable-ok: the proxy swaps in a pre-call hook's result only when it is a dict + previous_message_ids: Final = _continued_message_ids(data) + if not previous_message_ids: + return None + bindings: Final = tuple( + [await self._binding(message_id, user_api_key_dict) for message_id in previous_message_ids] + ) + owner: Final = _owner(user_api_key_dict) + pass_through: Final = call_type == _PASS_THROUGH_CALL_TYPE + if not all(_may_continue(binding, owner, pass_through) for binding in bindings): + raise HTTPException(status_code=404, detail=message_thread_not_found_body(_UNOWNED_THREAD_MESSAGE)) + anthropic_beta: Final = next( + (binding.anthropic_beta for binding in bindings if binding is not None and binding.anthropic_beta), None + ) + if anthropic_beta is None or pass_through: + return None + return { + **data, + "extra_headers": _with_anthropic_beta(data.get("extra_headers"), anthropic_beta), + **_provider_headers_with_anthropic_beta(data.get("provider_specific_header"), anthropic_beta), + } + + async def async_post_call_success_hook( + self, + data: Mapping[str, object], + user_api_key_dict: "UserAPIKeyAuth", + response: LLMResponseTypes, + ) -> LLMResponseTypes: + message_id: Final = _response_message_id(response) if _sends_thread(data) else None + if message_id is not None: + await self._record_owner(message_id, user_api_key_dict, _sent_anthropic_beta(data)) + return response + + async def async_post_call_streaming_iterator_hook( # pyright: ignore[reportIncompatibleMethodOverride] # the proxy streams bytes, str and dict chunks here, not the base's ModelResponseStream + self, + user_api_key_dict: "UserAPIKeyAuth", + response: AsyncIterable[object], + request_data: Mapping[str, object], + ) -> AsyncGenerator[object, None]: + scanner: Final = _StreamedMessageIdScanner(finished=not _sends_thread(request_data)) + async for chunk in response: + yield chunk + if not scanner.finished and (message_id := scanner.scan(chunk)) is not None: + await self._record_owner(message_id, user_api_key_dict, _sent_anthropic_beta(request_data)) + + @with_service_target("message_thread_owners") + async def _binding(self, message_id: str, user_api_key_dict: "UserAPIKeyAuth") -> _ThreadBinding | None: + cached: Final[object] = await self._owners.async_get_cache( + key=_owner_key(message_id), litellm_parent_otel_span=user_api_key_dict.parent_otel_span + ) + return _as_binding(cached) + + @with_service_target("message_thread_owners") + async def _record_owner( + self, message_id: str, user_api_key_dict: "UserAPIKeyAuth", anthropic_beta: str | None + ) -> None: + if await self._binding(message_id, user_api_key_dict) is not None: + return + key: Final = _owner_key(message_id) + claimed: Final = _ThreadBinding(owner=_owner(user_api_key_dict), anthropic_beta=anthropic_beta).model_dump() + await self._owners.async_set_cache( + key=key, + value=await _first_claim(self._owners.dual_cache.redis_cache, key, claimed), + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + local_only=True, + ttl=_OWNER_TTL_SECONDS, + ) + + +async def _first_claim(store: _FirstClaimStore | None, key: str, claimed: Mapping[str, object]) -> object: + if store is None: + return claimed + try: + await store.async_set_cache(key, claimed, nx=True, ttl=_OWNER_TTL_SECONDS) + stored: Final = await store.async_get_cache(key) + except Exception: # noqa: BLE001 # a Redis outage keeps the binding in this worker, as DualCache's writes do + return claimed + return claimed if stored is None else stored diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4468a724d2e..bf00bdfb5fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -65,6 +65,7 @@ from litellm.constants import ( LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, MAX_TEAM_LIST_LIMIT, + MESSAGE_THREAD_MAX_TRACKED_OWNERS, PROXY_REJECTED_BEFORE_ROUTING_KEY, REDACTED_BY_LITELLM, REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT, @@ -1963,6 +1964,9 @@ class ProxyLogging: self.file_usage_cache: Final = InternalUsageCache( dual_cache=DualCache(in_memory_cache=InMemoryCache(max_size_in_memory=FILE_USAGE_MAX_TRACKED_COUNTERS)) ) + self.message_thread_owner_cache: Final = InternalUsageCache( + dual_cache=DualCache(in_memory_cache=InMemoryCache(max_size_in_memory=MESSAGE_THREAD_MAX_TRACKED_OWNERS)) + ) self.max_parallel_request_limiter = PROXY_MaxParallelRequestsHandler(self.internal_usage_cache) self.cache_control_check = PROXY_CacheControlCheck() self.alerting: list[str] | None = None @@ -2122,6 +2126,7 @@ class ProxyLogging: if redis_cache is not None: self.internal_usage_cache.dual_cache.redis_cache = redis_cache self.file_usage_cache.dual_cache.redis_cache = redis_cache + self.message_thread_owner_cache.dual_cache.redis_cache = redis_cache self.db_spend_update_writer.redis_update_buffer.redis_cache = redis_cache self.db_spend_update_writer.pod_lock_manager.redis_cache = redis_cache @@ -2141,6 +2146,7 @@ class ProxyLogging: continue available_args: Mapping[str, object] = { "internal_usage_cache": self.internal_usage_cache, + "message_thread_owner_cache": self.message_thread_owner_cache, "prisma_client": prisma_client, } passed_in_args = {name: value for name, value in available_args.items() if name in expected_args} diff --git a/litellm/router_rust/surface/health.py b/litellm/router_rust/surface/health.py index 80f7e148185..f2ac72a262a 100644 --- a/litellm/router_rust/surface/health.py +++ b/litellm/router_rust/surface/health.py @@ -279,6 +279,7 @@ class RustHealthCallbacks: is_caller_scoped_auth_failure, is_caller_timeout_408, is_failure_outside_selected_deployment, + is_message_thread_not_found, is_response_lookup_not_found, is_video_lookup_not_found, ) @@ -290,6 +291,7 @@ class RustHealthCallbacks: or (isinstance(exception, Exception) and is_background_response_poll_not_found(exception, litellm_params)) or is_video_lookup_not_found(exception, logged) or is_response_lookup_not_found(exception, logged) + or is_message_thread_not_found(exception) or is_caller_timeout_408(logged, status) # pyright: ignore[reportArgumentType] # the status as raised or is_caller_scoped_auth_failure(self._deployment_of(deployment_id), status) # pyright: ignore[reportArgumentType] # see above ) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 897f31840d3..190ef43a3e3 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -30,6 +30,7 @@ from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials from litellm.router_deprecated import renamed_router_param from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback +from litellm.types.llms.anthropic import MESSAGE_THREAD_NOT_FOUND_ERROR_CODE from litellm.types.router import Deployment from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -106,6 +107,12 @@ def is_video_lookup_not_found(exception: object, failure_kwargs: Mapping[str, ob return getattr(exception, "status_code", None) == 404 and failure_kwargs.get("call_type") in VIDEO_ID_CALL_TYPES +def is_message_thread_not_found(exception: object) -> bool: + """Whether a message-thread continue failed with Anthropic's 404 ``thread_not_found``: the caller's + ``previous_message_id`` is stale and the SDK replays the conversation, the deployment is healthy.""" + return getattr(exception, "status_code", None) == 404 and MESSAGE_THREAD_NOT_FOUND_ERROR_CODE in str(exception) + + # The calls on an existing response id. ``compact_responses`` is left out: it sends a model and input, # so its 404 can mean a missing model rather than a missing response. RESPONSE_ID_CALL_TYPES: Final = frozenset( diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 3abe64df3c6..244883859f0 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -448,6 +448,14 @@ class AnthropicCompaction(TypedDict, total=False): instructions: ReadOnly[str] +class AnthropicMessageThread(TypedDict, total=False): + """``thread`` body field of the message-threads beta: ``{"type": "create"}`` opens a server-side + thread, ``{"type": "continue", "previous_message_id": "msg_..."}`` appends only the new messages to it.""" + + type: Required[ReadOnly[Literal["create", "continue"]]] + previous_message_id: ReadOnly[str] + + class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): max_tokens: int | None metadata: AnthropicMetadata | dict | None @@ -472,6 +480,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): reasoning_effort: str | None safeguards: ReadOnly[list[dict[str, object]] | None] service_tier: ReadOnly[str | None] + thread: ReadOnly[AnthropicMessageThread | None] class AnthropicMessagesRequest(AnthropicMessagesRequestOptionalParams, total=False): @@ -883,6 +892,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): PER_TURN_CONTROL_2026_07_01 = "per-turn-control-2026-07-01" DANGEROUS_TOOL_USE_2026_09_03 = "dangerous-tool-use-2026-09-03" MCP_CLIENT_2025_04_04 = "mcp-client-2025-04-04" + MESSAGE_THREADS_2026_08_12 = "message-threads-2026-08-12" # Tool search beta header constant (for Anthropic direct API and Microsoft Foundry) @@ -910,6 +920,9 @@ ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24" ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER: Final = "mid-conversation-output-config-2026-07-01" ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER: Final = "thinking-display-updates-2026-08-18" + +MESSAGE_THREAD_NOT_FOUND_ERROR_CODE: Final = "thread_not_found" +MESSAGE_THREAD_UNSUPPORTED_ERROR_CODE: Final = "thread_unsupported_request" ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER: Final = "thinking-binding-controls-2026-08-01" ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER: Final = "mid-conversation-tool-changes-2026-07-01" diff --git a/tests/integration/messages_endpoint/test_message_threads_wire.py b/tests/integration/messages_endpoint/test_message_threads_wire.py new file mode 100644 index 00000000000..2ba3e627ea0 --- /dev/null +++ b/tests/integration/messages_endpoint/test_message_threads_wire.py @@ -0,0 +1,413 @@ +import json +import re +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.shares_worker_state + +ANTHROPIC_MODEL: Final = "threads-anthropic" +BRIDGED_MODEL: Final = "threads-bridged" +BRIDGED_DROP_PARAMS_MODEL: Final = "threads-bridged-drop-params" +PASSTHROUGH_MODEL: Final = "threads-messages-passthrough" +MESSAGE_THREADS_BETA: Final = "message-threads-2026-08-12" +EXPIRED_THREAD_MARKER: Final = "thread-expired" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_NONCE: Final = re.compile(rb"nonce-([0-9a-f]{32})") + + +def nonce() -> str: + return f"nonce-{uuid.uuid4().hex}" + + +def thread_not_found_reply() -> Reply: + return Reply( + status=404, + body=json.dumps( + { + "type": "error", + "error": { + "type": "not_found_error", + "message": "No thread state was found for the requested `previous_message_id`.", + "details": {"error_code": "thread_not_found"}, + }, + "request_id": "req_upstream", + } + ).encode(), + ) + + +def anthropic_message(identity: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{identity}", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": f"answered {identity}"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 3}, + } + + +def sse_event(event: str, data: Mapping[str, JsonValue]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +def streamed_anthropic_reply(identity: str) -> Reply: + message_start: Final = sse_event( + "message_start", {"type": "message_start", "message": {**anthropic_message(identity), "content": []}} + ) + return Reply( + chunks=( + message_start[:40], + message_start[40:], + sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + ), + sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + sse_event( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}, + ), + sse_event("message_stop", {"type": "message_stop"}), + ), + content_type="text/event-stream", + ) + + +def responses_reply(identity: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-5.5", + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": f"answered {identity}", "annotations": []}], + } + ], + "usage": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15}, + } + ).encode() + ) + + +def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(body=b'{"object":"list","data":[]}') + match: Final = _NONCE.search(request.body) + assert match is not None, request.body[:200] + identity: Final = match.group(1).decode() + if request.target.startswith("/openai"): + return responses_reply(identity) + if EXPIRED_THREAD_MARKER.encode() in request.body: + return thread_not_found_reply() + if b'"stream": true' in request.body or b'"stream":true' in request.body: + return streamed_anthropic_reply(identity) + return Reply(body=json.dumps(anthropic_message(identity)).encode()) + + +def write_config(directory: Path, wire_url: str) -> Path: + path: Final = directory / f"threads-{uuid.uuid4().hex}.json" + path.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": ANTHROPIC_MODEL, + "litellm_params": { + "model": "anthropic/claude-opus-4-8", + "api_base": f"{wire_url}/anthropic", + "api_key": "synthetic-anthropic-key", + }, + }, + { + "model_name": BRIDGED_MODEL, + "litellm_params": { + "model": "openai/gpt-5.5", + "api_base": f"{wire_url}/openai", + "api_key": "synthetic-openai-key", + }, + }, + { + "model_name": BRIDGED_DROP_PARAMS_MODEL, + "litellm_params": { + "model": "openai/gpt-5.5", + "api_base": f"{wire_url}/openai", + "api_key": "synthetic-openai-key", + "drop_params": True, + }, + }, + { + "model_name": PASSTHROUGH_MODEL, + "litellm_params": { + "model": "hosted_vllm/qwen3-coder", + "api_base": f"{wire_url}/vllm", + "api_key": "synthetic-vllm-key", + }, + "model_info": {"supported_endpoints": ["/v1/messages"]}, + }, + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + } + ) + ) + return path + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + wire: Wire + + def send( + self, + model: str, + request_nonce: str, + thread: Mapping[str, str], + *, + key: str | None = None, + text: str = "hi", + **fields: JsonValue, + ) -> tuple[int, bytes]: + response: Final = self.gateway.client.post( + "/v1/messages", + json={ + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": f"{text} {request_nonce}"}], + "thread": dict(thread), + **fields, + }, + headers={"Authorization": f"Bearer {self.gateway.key if key is None else key}"}, + timeout=120, + ) + return response.status_code, response.content + + def create(self, *, key: str | None = None, **fields: JsonValue) -> str: + request_nonce: Final = nonce() + status, raw = self.send(ANTHROPIC_MODEL, request_nonce, {"type": "create"}, key=key, **fields) + assert status == 200, raw[:500] + self.received(request_nonce) + return f"msg_{request_nonce.removeprefix('nonce-')}" + + def received(self, request_nonce: str) -> tuple[Request, ...]: + return tuple(request for request in self.wire.drain() if request_nonce.encode() in request.body) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("threads") + with ( + gateway_from_environment() as shared, + wire_server(respond) as wire, + owned_proxy( + shared, + directory, + { + "LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "true", + "ANTHROPIC_API_BASE": f"{wire.url}/anthropic", + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + }, + config=write_config(directory, wire.url), + workers=1, + ) as gateway, + ): + yield Rig(gateway, wire) + + +def beta_header(headers: Mapping[str, str]) -> frozenset[str]: + return frozenset( + piece for key, value in headers.items() if key.lower() == "anthropic-beta" for piece in value.split(",") + ) # comprehension-ok: one header line may carry several comma-separated betas + + +def test_a_create_turn_reaches_anthropic_with_the_thread_field_and_its_beta(rig: Rig) -> None: + request_nonce: Final = nonce() + status, raw = rig.send(ANTHROPIC_MODEL, request_nonce, {"type": "create"}, system="be terse") + assert status == 200, raw[:500] + assert _JSON_OBJECT.validate_json(raw)["id"] == f"msg_{request_nonce.removeprefix('nonce-')}" + (upstream,) = rig.received(request_nonce) + body: Final = _JSON_OBJECT.validate_json(upstream.body) + assert body["thread"] == {"type": "create"} + assert body["system"] == "be terse" + assert MESSAGE_THREADS_BETA in beta_header(upstream.headers), upstream.headers + + +def test_a_continue_turn_reaches_anthropic_with_only_the_new_messages_and_the_same_beta(rig: Rig) -> None: + request_nonce: Final = nonce() + thread: Final = {"type": "continue", "previous_message_id": rig.create()} + status, raw = rig.send(ANTHROPIC_MODEL, request_nonce, thread) + assert status == 200, raw[:500] + (upstream,) = rig.received(request_nonce) + body: Final = _JSON_OBJECT.validate_json(upstream.body) + assert body["thread"] == thread + assert body["messages"] == [{"role": "user", "content": f"hi {request_nonce}"}] + assert "system" not in body and "tools" not in body + assert MESSAGE_THREADS_BETA in beta_header(upstream.headers), upstream.headers + + +def assert_thread_not_found(status: int, raw: bytes) -> None: + assert status == 404, raw[:500] + error: Final = _JSON_OBJECT.validate_json(raw) + detail: Final = error["error"] + assert isinstance(detail, dict), error + assert detail["type"] == "not_found_error" + assert detail["details"] == {"error_code": "thread_not_found"} + + +@pytest.mark.parametrize("model", [BRIDGED_MODEL, PASSTHROUGH_MODEL], ids=["translated", "messages-passthrough"]) +def test_a_create_turn_on_a_leg_without_threads_is_thread_unsupported_and_never_reaches_the_provider( + rig: Rig, model: str +) -> None: + request_nonce: Final = nonce() + status, raw = rig.send(model, request_nonce, {"type": "create"}) + assert status == 400, raw[:500] + detail: Final = _JSON_OBJECT.validate_json(raw)["error"] + assert isinstance(detail, dict), raw[:500] + assert detail["type"] == "invalid_request_error" + assert detail["details"] == {"error_code": "thread_unsupported_request"} + assert rig.received(request_nonce) == () + + +def test_a_create_turn_on_a_translated_leg_is_served_without_the_thread_under_drop_params(rig: Rig) -> None: + request_nonce: Final = nonce() + status, raw = rig.send(BRIDGED_DROP_PARAMS_MODEL, request_nonce, {"type": "create"}) + assert status == 200, raw[:500] + (upstream,) = rig.received(request_nonce) + assert upstream.target.startswith("/openai"), upstream.target + assert "thread" not in _JSON_OBJECT.validate_json(upstream.body) + + +@pytest.mark.parametrize("model", [BRIDGED_MODEL, PASSTHROUGH_MODEL], ids=["translated", "messages-passthrough"]) +def test_a_continue_turn_on_a_leg_without_threads_is_thread_not_found_and_never_reaches_the_provider( + rig: Rig, model: str +) -> None: + request_nonce: Final = nonce() + status, raw = rig.send(model, request_nonce, {"type": "continue", "previous_message_id": rig.create()}) + assert_thread_not_found(status, raw) + assert rig.received(request_nonce) == () + + +def test_another_key_cannot_continue_a_thread_and_never_reaches_anthropic(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + owner: Final = scenario.key() + intruder: Final = scenario.key() + thread: Final = {"type": "continue", "previous_message_id": rig.create(key=owner)} + intruder_nonce: Final = nonce() + assert_thread_not_found(*rig.send(ANTHROPIC_MODEL, intruder_nonce, thread, key=intruder)) + assert rig.received(intruder_nonce) == () + owner_nonce: Final = nonce() + status, raw = rig.send(ANTHROPIC_MODEL, owner_nonce, thread, key=owner) + assert status == 200, raw[:500] + assert len(rig.received(owner_nonce)) == 1 + + +def test_another_key_cannot_continue_a_thread_through_chat_completions(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + owner: Final = scenario.key() + intruder: Final = scenario.key() + thread: Final = {"type": "continue", "previous_message_id": rig.create(key=owner)} + intruder_nonce: Final = nonce() + response: Final = rig.gateway.client.post( + "/v1/chat/completions", + json={ + "model": ANTHROPIC_MODEL, + "max_tokens": 64, + "messages": [{"role": "user", "content": f"hi {intruder_nonce}"}], + "thread": thread, + }, + headers={"Authorization": f"Bearer {intruder}"}, + timeout=120, + ) + assert response.status_code == 404, response.content[:500] + assert rig.received(intruder_nonce) == () + + +ADVISOR_TOOL: Final = {"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-8"} +ADVISOR_BETA: Final = "advisor-tool-2026-03-01" + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "streamed"]) +def test_a_continue_turn_resends_the_betas_its_create_turn_took_from_its_tools(rig: Rig, stream: bool) -> None: + create_nonce: Final = nonce() + status, raw = rig.send(ANTHROPIC_MODEL, create_nonce, {"type": "create"}, tools=[ADVISOR_TOOL], stream=stream) + assert status == 200, raw[:500] + (created,) = rig.received(create_nonce) + create_betas: Final = beta_header(created.headers) + assert {ADVISOR_BETA, MESSAGE_THREADS_BETA} <= create_betas, created.headers + continue_nonce: Final = nonce() + thread: Final = {"type": "continue", "previous_message_id": f"msg_{create_nonce.removeprefix('nonce-')}"} + status, raw = rig.send(ANTHROPIC_MODEL, continue_nonce, thread) + assert status == 200, raw[:500] + (continued,) = rig.received(continue_nonce) + assert "tools" not in _JSON_OBJECT.validate_json(continued.body) + assert beta_header(continued.headers) == create_betas, continued.headers + + +def pass_through_continue(rig: Rig, request_nonce: str, previous_message_id: str, key: str) -> int: + return rig.gateway.client.post( + "/anthropic/v1/messages", + json={ + "model": "claude-opus-4-8", + "max_tokens": 64, + "messages": [{"role": "user", "content": f"hi {request_nonce}"}], + "thread": {"type": "continue", "previous_message_id": previous_message_id}, + }, + headers={"Authorization": f"Bearer {key}"}, + timeout=120, + ).status_code + + +def test_another_key_cannot_continue_a_thread_through_the_anthropic_pass_through_route(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + owner: Final = scenario.key() + intruder: Final = scenario.key() + previous_message_id: Final = rig.create(key=owner) + intruder_nonce: Final = nonce() + assert pass_through_continue(rig, intruder_nonce, previous_message_id, intruder) == 404 + assert rig.received(intruder_nonce) == () + owner_nonce: Final = nonce() + assert pass_through_continue(rig, owner_nonce, previous_message_id, owner) == 200 + assert len(rig.received(owner_nonce)) == 1 + + +def test_a_streamed_create_turn_can_be_continued_by_its_key(rig: Rig) -> None: + previous_message_id: Final = rig.create(stream=True) + request_nonce: Final = nonce() + status, raw = rig.send( + ANTHROPIC_MODEL, request_nonce, {"type": "continue", "previous_message_id": previous_message_id} + ) + assert status == 200, raw[:500] + assert len(rig.received(request_nonce)) == 1 + + +def test_a_thread_not_found_from_anthropic_keeps_its_error_code_for_the_sdk_to_replay_on(rig: Rig) -> None: + request_nonce: Final = nonce() + thread: Final = {"type": "continue", "previous_message_id": rig.create()} + assert_thread_not_found(*rig.send(ANTHROPIC_MODEL, request_nonce, thread, text=EXPIRED_THREAD_MARKER)) + assert len(rig.received(request_nonce)) == 1 diff --git a/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py index 210ab6717aa..0e0e1c8c53e 100644 --- a/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py +++ b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py @@ -258,3 +258,48 @@ class TestAnthropicErrorSseFrame: "type": "error", "error": {"type": "api_error", "message": "upstream unavailable"}, } + + +THREAD_NOT_FOUND_UPSTREAM: Final = json.dumps( + { + "type": "error", + "error": { + "type": "not_found_error", + "message": "No thread state was found for the requested `previous_message_id`.", + "details": {"error_code": "thread_not_found"}, + }, + "request_id": "req_upstream", + } +) + + +class TestEmbeddedErrorDetails: + def test_a_wrapped_upstream_error_keeps_its_details_next_to_the_wrapped_message(self) -> None: + wrapped: Final = ( + f"litellm.NotFoundError: AnthropicException - {THREAD_NOT_FOUND_UPSTREAM}\n\n" + "LiteLLM: model group 'haiku' failed with the error above. No fallback was attempted." + ) + response: Final = AnthropicExceptionMapping.transform_to_anthropic_error(status_code=404, raw_message=wrapped) + assert response["error"]["details"] == {"error_code": "thread_not_found"} + assert response["error"]["type"] == "not_found_error" + assert response["error"]["message"] == wrapped + + def test_the_error_event_of_a_failed_stream_carries_the_same_details(self) -> None: + wrapped: Final = f"litellm.NotFoundError: AnthropicException - {THREAD_NOT_FOUND_UPSTREAM}" + frame: Final = anthropic_error_sse_frame(404, wrapped) + assert frame.error_response["error"]["details"] == {"error_code": "thread_not_found"} + + @pytest.mark.parametrize( + "raw_message", + ( + "plain text with a { brace and no json", + 'litellm.BadRequestError: OpenAIException - {"error": {"message": "bad", "type": "invalid_request_error"}}', + 'AnthropicException - {"type": "error", "error": {"type": "api_error", "message": "no details here"}}', + ), + ids=("no_json", "openai_shape", "anthropic_without_details"), + ) + def test_a_message_without_anthropic_details_gets_no_details_key(self, raw_message: str) -> None: + response: Final = AnthropicExceptionMapping.transform_to_anthropic_error( + status_code=400, raw_message=raw_message + ) + assert "details" not in response["error"] diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index f6300b2e70c..ff150af324d 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -2801,6 +2801,85 @@ async def test_anthropic_messages_strict_tools_rejection_is_sent_again_without_s assert _JSON_BODY.validate_python(response)["id"] == "msg_strict_retry" +@pytest.mark.parametrize( + "model", + ("openai/gpt-5.5", "vertex_ai/claude-opus-4-8", "bedrock/anthropic.claude-opus-4-8-v1:0"), + ids=("chat_bridge", "vertex_native", "bedrock_native"), +) +def test_a_thread_continue_on_a_leg_without_message_threads_is_thread_not_found_before_any_upstream_call( + model: str, +) -> None: + from litellm.llms.anthropic.common_utils import AnthropicError + from litellm.llms.anthropic.messages.native.handler import anthropic_messages_handler + + with pytest.raises(AnthropicError) as raised: + anthropic_messages_handler( + max_tokens=64, + messages=[{"role": "user", "content": "hi"}], + model=model, + api_key="synthetic-key", + thread={"type": "continue", "previous_message_id": "msg_01previous"}, + ) + assert raised.value.status_code == 404 + error: Final = _JSON_BODY.validate_json(raised.value.message)["error"] + assert isinstance(error, dict) + assert error["type"] == "not_found_error" + assert error["details"] == {"error_code": "thread_not_found"} + + +@pytest.mark.parametrize( + "model", + ("openai/gpt-5.5", "vertex_ai/claude-opus-4-8", "bedrock/anthropic.claude-opus-4-8-v1:0"), + ids=("chat_bridge", "vertex_native", "bedrock_native"), +) +def test_a_thread_create_on_a_leg_without_message_threads_is_thread_unsupported_before_any_upstream_call( + model: str, +) -> None: + from litellm.llms.anthropic.common_utils import AnthropicError + from litellm.llms.anthropic.messages.native.handler import anthropic_messages_handler + + with pytest.raises(AnthropicError) as raised: + anthropic_messages_handler( + max_tokens=64, + messages=[{"role": "user", "content": "hi"}], + model=model, + api_key="synthetic-key", + thread={"type": "create"}, + ) + assert raised.value.status_code == 400 + error: Final = _JSON_BODY.validate_json(raised.value.message)["error"] + assert isinstance(error, dict) + assert error["type"] == "invalid_request_error" + assert error["details"] == {"error_code": "thread_unsupported_request"} + + +@pytest.mark.asyncio +async def test_a_thread_create_on_a_native_leg_without_message_threads_is_served_without_the_thread_under_drop_params( + local_beta_headers_config: None, +) -> None: + from litellm.llms.anthropic.messages.native import handler + + captured: Final[dict[str, object]] = {} + response: Final = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="bedrock/us.anthropic.claude-sonnet-5", + custom_llm_provider="bedrock", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=_upstream_answering_with([], captured), + thread={"type": "create"}, + drop_params=True, + ) + + body: Final = captured["body"] + assert isinstance(body, dict) + assert "thread" not in body + assert body["messages"] == [{"role": "user", "content": "hi"}] + assert _JSON_BODY.validate_python(response)["id"] == "msg_1" + + @pytest.mark.asyncio @pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) @pytest.mark.parametrize("stream", [False, True], ids=["complete", "stream"]) diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py index 8f86888728f..97d5b0550d4 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py @@ -17,6 +17,7 @@ from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.anthropic.wif import get_anthropic_wif_token +from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.llms.deepseek.messages.transformation import DeepSeekAnthropicMessagesConfig from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig @@ -723,6 +724,135 @@ def test_a_first_party_messages_request_never_fetches_the_remote_beta_headers_co assert not remote_config.called +MESSAGE_THREADS_BETA: Final = "message-threads-2026-08-12" + + +@pytest.mark.parametrize( + "thread", + ({"type": "create"}, {"type": "continue", "previous_message_id": "msg_01previous"}), + ids=("create", "continue"), +) +def test_messages_request_forwards_the_thread_field_to_anthropic(thread: dict[str, str]) -> None: + request: Final = AnthropicMessagesConfig().transform_anthropic_messages_request( # pyright: ignore[reportUnknownMemberType] # bare-dict signature + model="claude-opus-4-8", + messages=[{"role": "user", "content": "what did I just say"}], + anthropic_messages_optional_request_params={"max_tokens": 64, "thread": dict(thread)}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert request["thread"] == thread + assert request["messages"] == [{"role": "user", "content": "what did I just say"}] + + +def test_messages_headers_add_the_message_threads_beta_next_to_the_callers_betas() -> None: + headers: Final = AnthropicMessagesConfig()._update_headers_with_anthropic_beta( # pyright: ignore[reportUnknownMemberType] # bare-dict signature + headers={"anthropic-beta": "context-1m-2025-08-07"}, + optional_params={"thread": {"type": "continue", "previous_message_id": "msg_01previous"}}, + messages=[], + ) + assert headers["anthropic-beta"] == f"context-1m-2025-08-07,{MESSAGE_THREADS_BETA}" + + +def test_messages_headers_add_no_message_threads_beta_without_a_thread() -> None: + headers: Final = AnthropicMessagesConfig()._update_headers_with_anthropic_beta( # pyright: ignore[reportUnknownMemberType] # bare-dict signature + headers={}, optional_params={"max_tokens": 64}, messages=[] + ) + assert MESSAGE_THREADS_BETA not in headers.get("anthropic-beta", "") # pyright: ignore[reportUnknownMemberType] # bare-dict result + + +TOOL_SEARCH_TOOLS: Final = ( + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + {"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-8"}, +) +TURN_WITH_OUTPUT_CONFIG: Final = ( + {"role": "user", "content": "be brief", "output_config": {"effort": "low"}}, + { + "role": "system", + "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "tool_name": "lookup_record"}}], + }, +) + + +def test_a_thread_create_turn_takes_betas_from_its_tools_and_messages() -> None: + headers: Final = AnthropicMessagesConfig()._update_headers_with_anthropic_beta( # pyright: ignore[reportUnknownMemberType] # bare-dict signature + headers={"anthropic-beta": "context-1m-2025-08-07"}, + optional_params={"thread": {"type": "create"}, "system": "be terse", "tools": list(TOOL_SEARCH_TOOLS)}, + messages=list(TURN_WITH_OUTPUT_CONFIG), + ) + assert set(headers["anthropic-beta"].split(",")) == { # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # bare-dict result + "context-1m-2025-08-07", + MESSAGE_THREADS_BETA, + get_tool_search_beta_header("anthropic"), + "advisor-tool-2026-03-01", + "per-turn-control-2026-07-01", + "mid-conversation-tool-changes-2026-07-01", + } + + +def test_a_thread_continue_turn_takes_no_betas_from_its_new_messages() -> None: + headers: Final = AnthropicMessagesConfig()._update_headers_with_anthropic_beta( # pyright: ignore[reportUnknownMemberType] # bare-dict signature + headers={"anthropic-beta": "context-1m-2025-08-07"}, + optional_params={"thread": {"type": "continue", "previous_message_id": "msg_01p"}}, + messages=list(TURN_WITH_OUTPUT_CONFIG), + ) + assert headers["anthropic-beta"] == f"context-1m-2025-08-07,{MESSAGE_THREADS_BETA}" + + +def test_a_request_without_a_thread_still_takes_betas_from_its_tools_and_messages() -> None: + headers: Final = AnthropicMessagesConfig()._update_headers_with_anthropic_beta( # pyright: ignore[reportUnknownMemberType] # bare-dict signature + headers={}, optional_params={"tools": list(TOOL_SEARCH_TOOLS)}, messages=list(TURN_WITH_OUTPUT_CONFIG) + ) + assert set(headers["anthropic-beta"].split(",")) >= { # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # bare-dict result + get_tool_search_beta_header("anthropic"), + "advisor-tool-2026-03-01", + "per-turn-control-2026-07-01", + "mid-conversation-tool-changes-2026-07-01", + } + + +def _advisor_turn() -> list[dict[str, object]]: + return [ + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "advisor", "input": {}}, + { + "type": "advisor_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [{"type": "text", "text": "ok"}], + }, + {"type": "text", "text": "done"}, + ], + }, + {"role": "user", "content": "and now?"}, + ] + + +def test_a_thread_continue_turn_keeps_advisor_blocks_the_stored_tools_still_explain() -> None: + request: Final = AnthropicMessagesConfig().transform_anthropic_messages_request( # pyright: ignore[reportUnknownMemberType] # bare-dict signature + model="claude-opus-4-8", + messages=_advisor_turn(), + anthropic_messages_optional_request_params={ + "max_tokens": 64, + "thread": {"type": "continue", "previous_message_id": "msg_01previous"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert request["messages"][0]["content"] == _advisor_turn()[0]["content"] + + +def test_a_request_without_a_thread_still_strips_advisor_blocks_when_the_tool_is_absent() -> None: + request: Final = AnthropicMessagesConfig().transform_anthropic_messages_request( # pyright: ignore[reportUnknownMemberType] # bare-dict signature + model="claude-opus-4-8", + messages=_advisor_turn(), + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert request["messages"][0]["content"] == [{"type": "text", "text": "done"}] + + @pytest.mark.parametrize( "model, expected", [ diff --git a/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py b/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py index 8104c0cef56..910234fa3f1 100644 --- a/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py +++ b/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py @@ -137,6 +137,14 @@ def test_sampling_param_raises_clean_400_without_drop_params(monkeypatch): ) +def test_the_thread_field_survives_optional_param_filtering(): + thread: Final = {"type": "continue", "previous_message_id": "msg_01previous"} + result: Final = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + {"max_tokens": 64, "thread": thread} + ) + assert result == {"max_tokens": 64, "thread": thread} + + @pytest.mark.parametrize( ("model", "param", "value"), [ diff --git a/tests/unit/llms/bedrock/test_claude_platform_provider.py b/tests/unit/llms/bedrock/test_claude_platform_provider.py index 9142fe45752..b1c170a6f6f 100644 --- a/tests/unit/llms/bedrock/test_claude_platform_provider.py +++ b/tests/unit/llms/bedrock/test_claude_platform_provider.py @@ -927,3 +927,30 @@ def test_sigv4_no_duplicate_content_type_when_caller_sets_lowercase(): "Duplicate keys produce 'application/json, application/json' in the " "SigV4 canonical string and cause a 401." ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "thread", + ({"type": "create"}, {"type": "continue", "previous_message_id": "msg_01previous"}), + ids=("create", "continue"), +) +async def test_anthropic_messages_bedrock_claude_platform_forwards_a_message_thread(thread: dict[str, str]) -> None: + import litellm + + requests: list[dict] = [] + + await litellm.anthropic_messages( + model="bedrock/claude_platform/claude-sonnet-4-6", + messages=[{"role": "user", "content": "what did I just say"}], + max_tokens=10, + api_base="https://aws-external-anthropic.us-west-2.api.aws", + api_key="fake-platform-key", + workspace_id="wrkspc_test", + thread=dict(thread), + client=_async_gateway_client(requests, lambda url, headers, data: _anthropic_response(url)), + ) + + assert len(requests) == 1 + assert requests[0]["body"]["thread"] == thread + assert requests[0]["headers"]["anthropic-beta"] == "message-threads-2026-08-12" diff --git a/tests/unit/proxy/hooks/test_message_thread_ownership.py b/tests/unit/proxy/hooks/test_message_thread_ownership.py new file mode 100644 index 00000000000..e596e2958b9 --- /dev/null +++ b/tests/unit/proxy/hooks/test_message_thread_ownership.py @@ -0,0 +1,558 @@ +import asyncio +import datetime +import json +from collections.abc import AsyncIterator, Sequence +from typing import Final + +import httpx +import pytest +from fastapi import HTTPException +from pydantic import JsonValue, TypeAdapter + +from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.message_thread_ownership import MessageThreadOwnership +from litellm.types.utils import CallTypesLiteral +from litellm.proxy.utils import InternalUsageCache, ProxyLogging + +OWNER: Final = UserAPIKeyAuth(token="hashed-owner-key") +INTRUDER: Final = UserAPIKeyAuth(token="hashed-intruder-key") +CREATE: Final = {"type": "create"} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def continue_from(message_id: str) -> dict[str, object]: + return {"thread": {"type": "continue", "previous_message_id": message_id}, "messages": []} + + +def anthropic_message(message_id: str) -> dict[str, object]: + return {"id": message_id, "type": "message", "role": "assistant", "content": []} + + +def sse_frame(event: dict[str, object]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def message_start(message_id: str) -> dict[str, object]: + return {"type": "message_start", "message": anthropic_message(message_id)} + + +def sent_with(anthropic_beta: str, **data: object) -> dict[str, object]: + logging_obj: Final = Logging( + model="claude-opus-4-8", + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id="call", + function_id="function", + ) + sent: Final = httpx.Request("POST", "https://api.anthropic.com/v1/messages", headers={"anthropic-beta": anthropic_beta}) + logging_obj.model_call_details["httpx_response"] = httpx.Response(200, request=sent) + logging_obj.post_call(original_response="{}", additional_args={"complete_input_dict": {}}) + return {**data, "litellm_logging_obj": logging_obj} + + +def hook() -> MessageThreadOwnership: + return MessageThreadOwnership(message_thread_owner_cache=InternalUsageCache(DualCache())) + + +async def chunks(items: Sequence[object]) -> AsyncIterator[object]: + for item in items: + yield item + + +async def served( + ownership: MessageThreadOwnership, items: Sequence[object], request_data: dict[str, object] +) -> list[object]: + return [ + chunk + async for chunk in ownership.async_post_call_streaming_iterator_hook( + user_api_key_dict=OWNER, response=chunks(items), request_data=request_data + ) + ] + + +async def rejection( + ownership: MessageThreadOwnership, key: UserAPIKeyAuth, message_id: str, call_type: CallTypesLiteral +) -> HTTPException | None: + try: + await ownership.async_pre_call_hook( + user_api_key_dict=key, cache=DualCache(), data=continue_from(message_id), call_type=call_type + ) + except HTTPException as rejected: + return rejected + return None + + +async def allows( + ownership: MessageThreadOwnership, + key: UserAPIKeyAuth, + message_id: str, + call_type: CallTypesLiteral = "anthropic_messages", +) -> bool: + rejected: Final = await rejection(ownership, key, message_id, call_type) + if rejected is None: + return True + assert rejected.status_code == 404 + error: Final = _JSON_OBJECT.validate_json(rejected.detail)["error"] + assert isinstance(error, dict) + assert error["type"] == "not_found_error" + assert error["details"] == {"error_code": "thread_not_found"} + return False + + +@pytest.mark.asyncio +async def test_the_proxy_starts_with_a_hook_that_keeps_each_thread_with_its_key() -> None: + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging._add_proxy_hooks() + await proxy_logging.post_call_success_hook( + data={"thread": CREATE}, response=anthropic_message("msg_started"), user_api_key_dict=OWNER + ) + + with pytest.raises(HTTPException) as rejected: + await proxy_logging.pre_call_hook( + user_api_key_dict=INTRUDER, data=continue_from("msg_started"), call_type="anthropic_messages" + ) + assert rejected.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_only_the_key_that_created_a_thread_can_continue_it() -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=OWNER, response=anthropic_message("msg_created") + ) + + assert await allows(ownership, OWNER, "msg_created") + assert not await allows(ownership, INTRUDER, "msg_created") + + +@pytest.mark.asyncio +async def test_without_a_key_only_the_signed_in_identity_that_created_a_thread_can_continue_it() -> None: + ownership: Final = hook() + alice: Final = UserAPIKeyAuth(user_id="alice", team_id="team-1") + bob: Final = UserAPIKeyAuth(user_id="bob", team_id="team-1") + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=alice, response=anthropic_message("msg_signed_in") + ) + + assert await allows(ownership, alice, "msg_signed_in") + assert not await allows(ownership, bob, "msg_signed_in") + + +@pytest.mark.asyncio +async def test_without_a_key_one_managed_agent_cannot_continue_another_agents_thread() -> None: + ownership: Final = hook() + planner: Final = UserAPIKeyAuth(agent_id="agent-planner") + coder: Final = UserAPIKeyAuth(agent_id="agent-coder") + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=planner, response=anthropic_message("msg_agent") + ) + + assert await allows(ownership, planner, "msg_agent") + assert not await allows(ownership, coder, "msg_agent") + + +@pytest.mark.parametrize("call_type", ("anthropic_messages", "acompletion", "aresponses")) +@pytest.mark.asyncio +async def test_a_continue_naming_a_message_this_proxy_never_served_in_a_thread_is_thread_not_found( + call_type: CallTypesLiteral, +) -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data={}, user_api_key_dict=OWNER, response=anthropic_message("msg_unthreaded") + ) + + assert not await allows(ownership, OWNER, "msg_unthreaded", call_type) + assert not await allows(ownership, OWNER, "msg_unknown", call_type) + assert not await allows(ownership, INTRUDER, "msg_unknown", call_type) + + +@pytest.mark.asyncio +async def test_each_continued_turn_hands_its_own_message_to_the_same_key() -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=OWNER, response=anthropic_message("msg_first") + ) + await ownership.async_post_call_success_hook( + data=continue_from("msg_first"), user_api_key_dict=OWNER, response=anthropic_message("msg_second") + ) + + assert await allows(ownership, OWNER, "msg_second") + assert not await allows(ownership, INTRUDER, "msg_second") + + +@pytest.mark.asyncio +async def test_a_reply_served_again_to_another_key_stays_with_the_key_that_got_it_first() -> None: + ownership: Final = hook() + for key in (OWNER, INTRUDER): + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=key, response=anthropic_message("msg_cached") + ) + + assert await allows(ownership, OWNER, "msg_cached") + assert not await allows(ownership, INTRUDER, "msg_cached") + + +@pytest.mark.asyncio +async def test_a_worker_that_writes_after_another_worker_claimed_the_reply_keeps_the_first_owner( + router_redis: tuple[RedisCache, object], +) -> None: + redis_cache, _ = router_redis + late_write_started: Final = asyncio.Event() + owner_stored: Final = asyncio.Event() + + class WritesAfterTheOwner(RedisCache): + async def async_set_cache( + self, + key: str | None, + value: object, + **kwargs: object, # kwargs-ok: RedisCache write signature + ) -> None: + late_write_started.set() + await owner_stored.wait() + await super().async_set_cache(key, value, **kwargs) + + first: Final = MessageThreadOwnership( + message_thread_owner_cache=InternalUsageCache(DualCache(redis_cache=redis_cache)) + ) + late: Final = MessageThreadOwnership( + message_thread_owner_cache=InternalUsageCache( + DualCache(redis_cache=WritesAfterTheOwner(url="redis://router-unit.example:6379/0")) + ) + ) + + async def owner_claims() -> None: + await late_write_started.wait() + await first.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=OWNER, response=anthropic_message("msg_raced") + ) + owner_stored.set() + + await asyncio.gather( + late.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=INTRUDER, response=anthropic_message("msg_raced") + ), + owner_claims(), + ) + + for worker in (first, late): + assert await allows(worker, OWNER, "msg_raced") + assert not await allows(worker, INTRUDER, "msg_raced") + + +@pytest.mark.asyncio +async def test_every_redis_command_for_a_thread_binding_is_traced_under_the_message_thread_owners_family( + router_redis: tuple[RedisCache, asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]], +) -> None: + redis_cache, commands = router_redis + creating_worker: Final = MessageThreadOwnership( + message_thread_owner_cache=InternalUsageCache(DualCache(redis_cache=redis_cache)) + ) + continuing_worker: Final = MessageThreadOwnership( + message_thread_owner_cache=InternalUsageCache(DualCache(redis_cache=redis_cache)) + ) + await creating_worker.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=OWNER, response=anthropic_message("msg_traced") + ) + + assert await allows(continuing_worker, OWNER, "msg_traced") + sent: Final = tuple(commands.get_nowait() for _ in range(commands.qsize())) + thread_commands: Final = tuple((str(args[0]), target) for args, target, _ in sent if args[0] != "CLIENT") + assert {command for command, _ in thread_commands} == {"GET", "SET"} + assert {target for _, target in thread_commands} == {"message_thread_owners"} + + +@pytest.mark.parametrize("call_type", ("acompletion", "aresponses")) +@pytest.mark.asyncio +async def test_a_thread_continue_on_another_route_stays_with_the_key_that_created_the_thread( + call_type: CallTypesLiteral, +) -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=OWNER, response=anthropic_message("msg_cross_route") + ) + + assert await allows(ownership, OWNER, "msg_cross_route", call_type) + assert not await allows(ownership, INTRUDER, "msg_cross_route", call_type) + + +@pytest.mark.parametrize("call_type", ("anthropic_messages", "acompletion", "aresponses")) +@pytest.mark.asyncio +async def test_a_continue_sent_in_extra_body_stays_with_the_key_that_created_the_thread( + call_type: CallTypesLiteral, +) -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=OWNER, response=anthropic_message("msg_extra_body") + ) + in_extra_body: Final = {"messages": [], "extra_body": {"thread": continue_from("msg_extra_body")["thread"]}} + + assert ( + await ownership.async_pre_call_hook( + user_api_key_dict=OWNER, cache=DualCache(), data=in_extra_body, call_type=call_type + ) + is None + ) + with pytest.raises(HTTPException) as rejected: + await ownership.async_pre_call_hook( + user_api_key_dict=INTRUDER, cache=DualCache(), data=in_extra_body, call_type=call_type + ) + assert rejected.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_an_owned_top_level_continue_does_not_let_another_keys_extra_body_continue_through() -> None: + ownership: Final = hook() + for key, message_id in ((OWNER, "msg_owner"), (INTRUDER, "msg_intruder")): + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=key, response=anthropic_message(message_id) + ) + mixed: Final = {**continue_from("msg_intruder"), "extra_body": {"thread": continue_from("msg_owner")["thread"]}} + + with pytest.raises(HTTPException) as rejected: + await ownership.async_pre_call_hook( + user_api_key_dict=INTRUDER, cache=DualCache(), data=mixed, call_type="acompletion" + ) + assert rejected.value.status_code == 404 + + +@pytest.mark.parametrize( + "create", + ( + pytest.param( + lambda ownership, data: ownership.async_post_call_success_hook( + data=data, user_api_key_dict=OWNER, response=anthropic_message("msg_body_create") + ), + id="reply", + ), + pytest.param( + lambda ownership, data: served( + ownership, (sse_frame(message_start("msg_body_create")), sse_frame({"type": "message_stop"})), data + ), + id="stream", + ), + ), +) +@pytest.mark.asyncio +async def test_a_thread_created_in_extra_body_stays_with_the_key_that_created_it(create) -> None: + ownership: Final = hook() + await create(ownership, {"messages": [], "extra_body": {"thread": CREATE}}) + + assert await allows(ownership, OWNER, "msg_body_create", "acompletion") + assert not await allows(ownership, INTRUDER, "msg_body_create", "acompletion") + + +@pytest.mark.asyncio +async def test_a_pass_through_continue_of_a_message_the_proxy_never_bound_is_left_to_the_upstream() -> None: + assert await allows(hook(), OWNER, "msg_unknown", "pass_through_endpoint") + assert await allows(hook(), INTRUDER, "msg_unknown", "pass_through_endpoint") + + +@pytest.mark.asyncio +async def test_a_pass_through_continue_of_another_keys_thread_is_thread_not_found() -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=OWNER, response=anthropic_message("msg_bound") + ) + + assert await allows(ownership, OWNER, "msg_bound", "pass_through_endpoint") + assert not await allows(ownership, INTRUDER, "msg_bound", "pass_through_endpoint") + + +@pytest.mark.asyncio +async def test_a_continue_resends_the_betas_its_create_turn_sent_next_to_its_own_extra_headers() -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data=sent_with("advisor-tool-2026-03-01,message-threads-2026-08-12", thread=CREATE), + user_api_key_dict=OWNER, + response=anthropic_message("msg_advisor"), + ) + continued: Final = await ownership.async_pre_call_hook( + user_api_key_dict=OWNER, + cache=DualCache(), + data={ + **continue_from("msg_advisor"), + "extra_headers": {"Anthropic-Beta": "context-1m-2025-08-07", "x-trace": "abc"}, + }, + call_type="anthropic_messages", + ) + + assert continued is not None + assert continued["extra_headers"] == { + "x-trace": "abc", + "anthropic-beta": "advisor-tool-2026-03-01,context-1m-2025-08-07,message-threads-2026-08-12", + } + assert continued["thread"] == continue_from("msg_advisor")["thread"] + + +@pytest.mark.asyncio +async def test_the_create_turns_betas_reach_every_later_turn_of_the_thread() -> None: + ownership: Final = hook() + items: Final = (sse_frame(message_start("msg_streamed_first")), sse_frame({"type": "message_stop"})) + await served(ownership, items, sent_with("advisor-tool-2026-03-01,message-threads-2026-08-12", thread=CREATE)) + second: Final = await ownership.async_pre_call_hook( + user_api_key_dict=OWNER, cache=DualCache(), data=continue_from("msg_streamed_first"), call_type="anthropic_messages" + ) + assert second is not None + second_headers: Final = second["extra_headers"] + assert isinstance(second_headers, dict) + await ownership.async_post_call_success_hook( + data=sent_with(str(second_headers["anthropic-beta"]), **continue_from("msg_streamed_first")), + user_api_key_dict=OWNER, + response=anthropic_message("msg_second_turn"), + ) + third: Final = await ownership.async_pre_call_hook( + user_api_key_dict=OWNER, cache=DualCache(), data=continue_from("msg_second_turn"), call_type="anthropic_messages" + ) + + assert third is not None + assert third["extra_headers"] == {"anthropic-beta": "advisor-tool-2026-03-01,message-threads-2026-08-12"} + + +@pytest.mark.parametrize("call_type", ("acompletion", "aresponses")) +@pytest.mark.asyncio +async def test_a_continue_on_chat_or_responses_resends_the_betas_its_create_turn_sent( + call_type: CallTypesLiteral, +) -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data=sent_with("advisor-tool-2026-03-01,message-threads-2026-08-12", thread=CREATE), + user_api_key_dict=OWNER, + response=anthropic_message("msg_other_route"), + ) + continued: Final = await ownership.async_pre_call_hook( + user_api_key_dict=OWNER, cache=DualCache(), data=continue_from("msg_other_route"), call_type=call_type + ) + + assert continued is not None + assert continued["extra_headers"] == {"anthropic-beta": "advisor-tool-2026-03-01,message-threads-2026-08-12"} + + +CLIENT_BETAS: Final = {"custom_llm_provider": "anthropic,azure_ai", "extra_headers": {"anthropic-beta": "context-1m-2025-08-07"}} +CLIENT_VERSION: Final = {"custom_llm_provider": "anthropic", "extra_headers": {"anthropic-version": "2023-06-01"}} +CLIENT_BETAS_WITH_THE_CREATES: Final = { + "custom_llm_provider": "anthropic,azure_ai", + "extra_headers": {"anthropic-beta": "advisor-tool-2026-03-01,context-1m-2025-08-07,message-threads-2026-08-12"}, +} + + +@pytest.mark.parametrize( + ("client_headers", "expected"), + ( + (CLIENT_BETAS, CLIENT_BETAS_WITH_THE_CREATES), + ([CLIENT_BETAS, CLIENT_VERSION], (CLIENT_BETAS_WITH_THE_CREATES, CLIENT_VERSION)), + (CLIENT_VERSION, CLIENT_VERSION), + ), +) +@pytest.mark.parametrize("call_type", ("acompletion", "aresponses", "anthropic_messages")) +@pytest.mark.asyncio +async def test_a_continue_that_sends_its_own_betas_keeps_the_create_turns_betas_next_to_them( + call_type: CallTypesLiteral, client_headers: object, expected: object +) -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data=sent_with("advisor-tool-2026-03-01,message-threads-2026-08-12", thread=CREATE), + user_api_key_dict=OWNER, + response=anthropic_message("msg_client_betas"), + ) + continued: Final = await ownership.async_pre_call_hook( + user_api_key_dict=OWNER, + cache=DualCache(), + data={**continue_from("msg_client_betas"), "provider_specific_header": client_headers}, + call_type=call_type, + ) + + assert continued is not None + assert continued["provider_specific_header"] == expected + assert continued["extra_headers"] == {"anthropic-beta": "advisor-tool-2026-03-01,message-threads-2026-08-12"} + + +@pytest.mark.asyncio +async def test_a_pass_through_continue_is_checked_but_left_unchanged() -> None: + ownership: Final = hook() + await ownership.async_post_call_success_hook( + data=sent_with("advisor-tool-2026-03-01", thread=CREATE), + user_api_key_dict=OWNER, + response=anthropic_message("msg_other_route"), + ) + + assert ( + await ownership.async_pre_call_hook( + user_api_key_dict=OWNER, + cache=DualCache(), + data=continue_from("msg_other_route"), + call_type="pass_through_endpoint", + ) + is None + ) + + +@pytest.mark.asyncio +async def test_thread_bindings_never_push_the_proxys_rate_limit_counters_out_of_memory() -> None: + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + ownership: Final = MessageThreadOwnership(message_thread_owner_cache=proxy_logging.message_thread_owner_cache) + await proxy_logging.internal_usage_cache.async_set_cache( + key="rpm-counter", value=1, litellm_parent_otel_span=None, ttl=60 + ) + for turn in range(250): + await ownership.async_post_call_success_hook( + data={"thread": CREATE}, user_api_key_dict=OWNER, response=anthropic_message(f"msg_{turn}") + ) + + assert await proxy_logging.internal_usage_cache.async_get_cache(key="rpm-counter", litellm_parent_otel_span=None) == 1 + assert await allows(ownership, OWNER, "msg_0") + + +@pytest.mark.parametrize( + "stream", + ( + pytest.param(lambda frame: (frame,), id="whole_frame"), + pytest.param(lambda frame: (frame[:25], frame[25:60], frame[60:]), id="torn_frame"), + pytest.param(lambda frame: (frame.decode()[:30], frame.decode()[30:]), id="torn_text_frame"), + pytest.param( + lambda frame: (sse_frame({"type": "ping"}) + frame[:30], frame[30:]), id="after_a_ping_in_one_chunk" + ), + ), +) +@pytest.mark.asyncio +async def test_a_streamed_thread_turn_hands_its_message_to_its_key_and_streams_unchanged(stream) -> None: + ownership: Final = hook() + items: Final = (*stream(sse_frame(message_start("msg_streamed"))), sse_frame({"type": "message_stop"})) + + assert await served(ownership, items, {"thread": CREATE}) == list(items) + assert await allows(ownership, OWNER, "msg_streamed") + assert not await allows(ownership, INTRUDER, "msg_streamed") + + +@pytest.mark.asyncio +async def test_a_streamed_thread_turn_of_event_dicts_hands_its_message_to_its_key() -> None: + ownership: Final = hook() + items: Final = (message_start("msg_dict_stream"), {"type": "message_stop"}) + + assert await served(ownership, items, {"thread": CREATE}) == list(items) + assert await allows(ownership, OWNER, "msg_dict_stream") + + +@pytest.mark.asyncio +async def test_a_stream_without_a_thread_hands_out_nothing() -> None: + ownership: Final = hook() + items: Final = (sse_frame(message_start("msg_plain_stream")), sse_frame({"type": "message_stop"})) + + assert await served(ownership, items, {}) == list(items) + assert not await allows(ownership, OWNER, "msg_plain_stream") + + +@pytest.mark.asyncio +async def test_a_message_start_after_another_event_is_not_read_as_the_threads_message() -> None: + ownership: Final = hook() + items: Final = ( + sse_frame({"type": "content_block_start", "index": 0}), + sse_frame(message_start("msg_late")), + ) + + assert await served(ownership, items, {"thread": CREATE}) == list(items) + assert not await allows(ownership, OWNER, "msg_late") diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index 26ee9f04d01..ff800b24912 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -454,6 +454,7 @@ def test_every_pre_call_customlogger_is_deliberately_classified(): "_PROXY_BatchRedisRequests", "_PROXY_SensitiveDataRoutingHandler", "ResponsesIDSecurity", + "MessageThreadOwnership", "SkillsInjectionHook", "_PROXY_LiteLLMManagedFiles", "_PROXY_LiteLLMManagedVectorStores", diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index e16a8e01d81..3a61ca48f26 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -3411,6 +3411,75 @@ def test_the_refusal_names_each_cooled_deployment_some_cooldown_state_routes_the assert {deployment["model_id"] for deployment in refusal.value.cooldown_deployments or ()} == reachable +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +class TestMessageThreadNotFoundCooldown: + THREAD_NOT_FOUND_BODY: Final = ( + '{"type":"error","error":{"type":"not_found_error","message":"No thread state was found for the ' + 'requested `previous_message_id`.","details":{"error_code":"thread_not_found"}}}' + ) + + @staticmethod + def _not_found(message: str) -> Exception: + return litellm.NotFoundError( + message=f"AnthropicException - {message}", + model="claude-opus-4-8", + llm_provider="anthropic", + ) + + @pytest.mark.asyncio + async def test_a_stale_thread_404_leaves_the_deployment_out_of_cooldown(self) -> None: + model_id: Final = "anthropic-threads" + router: Final = _make_router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-opus-4-8", "api_key": "synthetic-key"}, + "model_info": {"id": model_id}, + } + ], + ) + + result: Final = router.deployment_callback_on_failure( + kwargs={ + "exception": self._not_found(self.THREAD_NOT_FOUND_BODY), + "litellm_params": {"model_info": {"id": model_id}}, + }, + completion_response=None, + start_time=0, + end_time=1, + ) + + assert result is False + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [] + + @pytest.mark.asyncio + async def test_a_404_without_the_thread_code_still_cools_the_deployment_down(self) -> None: + model_id: Final = "anthropic-missing-model" + router: Final = _make_router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-opus-4-8", "api_key": "synthetic-key"}, + "model_info": {"id": model_id}, + } + ], + ) + + router.deployment_callback_on_failure( + kwargs={ + "exception": self._not_found( + '{"type":"error","error":{"type":"not_found_error","message":"model: x"}}' + ), + "litellm_params": {"model_info": {"id": model_id}}, + }, + completion_response=None, + start_time=0, + end_time=1, + ) + + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [model_id] + + @pytest.mark.parametrize( ("status_code", "call_type", "expected"), [