From 9ebbcc7710821fc3d245417fb3d00a06b906c114 Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:33:54 -0800 Subject: [PATCH] feat(anthropic): forward message threads on native /v1/messages (#50976) * feat(anthropic): forward message threads on native /v1/messages Keep the thread body field on native Anthropic legs, add the message-threads-2026-08-12 beta for it, reject it with a 400 on legs that cannot continue a stored thread, and relay error.details of an upstream Anthropic error so thread_not_found reaches SDK thread helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(anthropic): keep a stale thread 404 from cooling the deployment and type the thread beta ignores Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style: sort the cooldown handler imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): dedupe the message-threads beta map rows and annotate test finals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): bind message threads to the key that created them and answer unsupported continues with thread_not_found Every virtual key on a proxy shares one Anthropic credential, so forwarding thread alone let any key continue another key's thread by its message id. A new default proxy hook binds each thread reply id to its virtual key, or to the signed-in identity when there is no key, and answers a continue naming an id it did not bind with the 404 thread_not_found the SDK replays on. Legs without message threads now serve a create without the thread and answer a continue with that 404 before any upstream call, instead of answering it with no history. * fix(proxy): keep a cached message thread reply with the key that got it first A response cache hit replays the original reply id, and binding it again handed the thread to the second key, so the creator's continue got thread_not_found. * fix(anthropic): keep message threads bound on every route and stable across turns Check the thread binding on every route that carries thread except pass-through, keep bindings in their own cache so they never evict rate-limit counters, add no tool or message derived betas on thread turns, answer a create on a leg without threads with a 400 thread_unsupported_request unless drop_params is set, and forward threads on Claude Platform on AWS * fix(anthropic): reject message threads on /v1/messages passthrough deployments A deployment that forwards /v1/messages untranslated to an OpenAI-compatible server inherited the native config's thread support, so `thread` reached a server that keeps no thread state. It now answers the same 400 thread_unsupported_request on create and 404 thread_not_found on continue as a translated leg * fix(anthropic): resend a thread's create betas on continues and check bound ids on pass-through A thread create now gets the betas LiteLLM derives from its tools and messages again, so an advisor tool create that answered 400 on the previous commit answers 200. Anthropic fingerprints the beta set, so the proxy keeps the anthropic-beta header the create sent upstream in the thread binding and adds it to every continue on /v1/messages. It reads that header from the upstream request, since success logging overwrites the logged request details before the success hook runs A continue on the /anthropic pass-through route of a thread the proxy bound to another key now gets the 404 thread_not_found. Ids the proxy never bound still pass through, as before * fix(anthropic): continue unbound threads off /v1/messages and keep every anthropic-beta on chat and Responses * fix(anthropic): skip non-string anthropic-beta list entries instead of raising * test(anthropic): keep the header merge regression test clear of main's appended tests * fix(anthropic): add a thread's create betas to the client's own on continues and keep base header merging on chat and Responses * fix(anthropic): answer an unbound thread id on chat and Responses with the 404 and drop the stale beta-join tests * fix(proxy): check a thread continue sent inside extra_body against its owner Azure AI promotes extra_body into the request body, so another key could continue a thread by putting thread there, past the ownership hook that read only the top-level field * fix(proxy): write a thread's first owner atomically and bind creates sent inside extra_body Two workers that both miss a reply's binding now agree on the first key stored in Redis, since the claim uses SET NX and each worker caches the stored winner. A create sent inside extra_body now binds its reply like a top-level create. The redundant Claude Platform supports_message_threads override is removed, since the provider already resolves to the native Anthropic beta list * fix(proxy): include agent_id in the key-less message thread owner * fix(proxy): tag message thread ownership Redis calls with their key family --------- Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/anthropic_beta_headers_config.json | 10 +- .../exceptions/exception_mapping_utils.py | 36 +- .../exceptions/exceptions.py | 1 + litellm/constants.py | 1 + .../llms/anthropic/messages/native/handler.py | 39 +- .../messages/native/transformation.py | 29 +- .../llms/anthropic/messages/native/utils.py | 37 ++ .../anthropic_messages/transformation.py | 3 + .../openai_like/messages/transformation.py | 3 + litellm/proxy/hooks/__init__.py | 2 + .../proxy/hooks/message_thread_ownership.py | 319 ++++++++++ litellm/proxy/utils.py | 6 + litellm/router_rust/surface/health.py | 2 + litellm/router_utils/cooldown_handlers.py | 7 + litellm/types/llms/anthropic.py | 13 + .../test_message_threads_wire.py | 413 +++++++++++++ .../test_exception_mapping_utils.py | 45 ++ .../native/test_anthropic_messages_handler.py | 79 +++ .../test_anthropic_messages_transformation.py | 130 ++++ .../test_request_optional_param_utils.py | 8 + .../bedrock/test_claude_platform_provider.py | 27 + .../hooks/test_message_thread_ownership.py | 558 ++++++++++++++++++ .../utils/proxy_logging/test_pre_call_hook.py | 1 + .../router_utils/test_cooldown_handlers.py | 69 +++ 24 files changed, 1826 insertions(+), 12 deletions(-) create mode 100644 litellm/proxy/hooks/message_thread_ownership.py create mode 100644 tests/integration/messages_endpoint/test_message_threads_wire.py create mode 100644 tests/unit/proxy/hooks/test_message_thread_ownership.py 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"), [