Merge remote-tracking branch 'origin/main' into litellm_remove_lit002_dict_ban

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	litellm/proxy/management_endpoints/team_endpoints.py
#	litellm/proxy/management_helpers/team_metadata_validation.py

litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs

@@ -5,6 +5,7 @@ use serde_json::Value;
 
 #[derive(Deserialize)]
 struct Input {
+    path: String,
     model_alias: String,
     provider_model: String,
     api_base: String,
@@ -21,7 +22,8 @@ async fn main() {
         Ok(input) => input,
         Err(error) => fail(error),
     };
-    let result = litellm_ai_gateway::trace_parity::traced_messages_request(
+    let result = litellm_ai_gateway::trace_parity::traced_request(
+        input.path,
         input.model_alias,
         input.provider_model,
         input.api_base,

litellm-rust/crates/ai-gateway/src/trace_parity.rs

@@ -29,14 +29,15 @@ pub struct TracedGatewayResponse {
     pub trace: Vec<litellm_core::observability::FunctionTraceEvent>,
 }
 
-pub async fn traced_messages_request(
+pub async fn traced_request(
+    path: String,
     model_alias: String,
     provider_model: String,
     api_base: String,
     body: Value,
 ) -> TracedGatewayResponse {
     let trace = litellm_core::observability::FunctionTrace::default();
-    let result = messages_request(model_alias, provider_model, api_base, body)
+    let result = request(path, model_alias, provider_model, api_base, body)
         .with_subscriber(trace.dispatcher())
         .await;
     let events = trace.events();
@@ -54,7 +55,8 @@ pub async fn traced_messages_request(
     }
 }
 
-pub async fn messages_request(
+pub async fn request(
+    path: String,
     model_alias: String,
     provider_model: String,
     api_base: String,
@@ -75,7 +77,7 @@ pub async fn messages_request(
     };
     let request = Request::builder()
         .method("POST")
-        .uri("/v1/messages")
+        .uri(path)
         .header(AUTHORIZATION, "Bearer trace-master-key")
         .header(CONTENT_TYPE, "application/json")
         .body(Body::from(body.to_string()))

litellm/constants.py

@@ -317,6 +317,8 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
 BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
 BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
 BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
+REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
+REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
 
 # SSL/TLS cipher configuration for faster handshakes
 # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones

litellm/integrations/arize/arize_phoenix_prompt_manager.py

@@ -379,7 +379,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
         """Reload prompts from Arize Phoenix."""
         if self.prompt_id:
             self._prompt_manager = None  # Reset to force reload
-            self.prompt_manager  # This will trigger reload
+            _ = self.prompt_manager  # access triggers lazy reload
 
     def should_run_prompt_management(
         self,

litellm/integrations/bitbucket/bitbucket_prompt_manager.py

@@ -406,7 +406,7 @@ class BitBucketPromptManager(CustomPromptManagement):
         """Reload prompts from BitBucket."""
         if self.prompt_id:
             self._prompt_manager = None  # Reset to force reload
-            self.prompt_manager  # This will trigger reload
+            _ = self.prompt_manager  # access triggers lazy reload
 
     def should_run_prompt_management(
         self,

litellm/integrations/weights_biases.py

@@ -4,18 +4,10 @@ imported_openAIResponse = True
 try:
     import io
     import logging
-    import sys
-    from typing import Any, TypeVar
+    from typing import Any, Literal, Protocol, TypeVar
 
     from wandb.sdk.data_types import trace_tree
 
-    if sys.version_info >= (3, 8):
-        from typing import Literal, Protocol
-    else:
-        from typing import Literal
-
-        from typing_extensions import Protocol
-
     logger: Final = logging.getLogger(__name__)
 
     K = TypeVar("K", bound=str)

litellm/litellm_core_utils/coroutine_checker.py

@@ -36,7 +36,7 @@ class CoroutineChecker:
         target = callback
         if not inspect.isfunction(target) and not inspect.ismethod(target):
             try:
-                call_attr: Final = getattr(target, "__call__", None)
+                call_attr: Final = getattr(target, "__call__", None)  # noqa: B004  # value unwrap so iscoroutinefunction sees through functors
                 if call_attr is not None:
                     target = call_attr
             except Exception:

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -1757,7 +1757,7 @@ def convert_to_anthropic_tool_invoke(
     anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, object]]] = []
 
     for tool in tool_calls:
-        if not get_attribute_or_key(tool, "type") == "function":
+        if get_attribute_or_key(tool, "type") != "function":
             continue
 
         tool_id = cast(str, get_attribute_or_key(tool, "id"))

litellm/litellm_core_utils/realtime_streaming.py

@@ -10,6 +10,7 @@ from typing_extensions import ReadOnly
 
 import litellm
 from litellm._logging import redact_internal_details_from_client_message, verbose_logger
+from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
 from litellm.types.llms.openai import (
@@ -35,9 +36,6 @@ else:
     CLIENT_CONNECTION_CLASS = Any
 
 
-REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
-
-
 @dataclass(frozen=True, slots=True)
 class BackendClose:
     code: int
@@ -1153,6 +1151,7 @@ class RealTimeStreaming:
         self._logging_worker.ensure_initialized_and_enqueue(
             self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True)
         )
+        self.logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True
 
     @staticmethod
     def _detect_beta_header(websocket: ScopedWebSocket) -> bool:

litellm/litellm_core_utils/token_counter.py

@@ -454,7 +454,7 @@ def token_counter(
         params: Final = _MessageCountParams(model, custom_tokenizer)
         num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count)
         if count_response_tokens is False:
-            includes_system_message: Final = any([message.get("role", None) == "system" for message in new_messages])
+            includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages)
             num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message)
 
     else:

litellm/llms/azure_ai/common_utils.py

@@ -144,10 +144,13 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
     def get_api_key(api_key: str | None = None) -> str | None:
         return api_key or litellm.api_key or get_secret_str("AZURE_AI_API_KEY")
 
+    @staticmethod
+    def get_api_version(api_version: str | None = None) -> str | None:
+        return api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
+
     @property
-    def api_version(self, api_version: str | None = None) -> str | None:
-        api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
-        return api_version
+    def api_version(self) -> str | None:
+        return AzureFoundryModelInfo.get_api_version()
 
     def get_token_counter(self) -> BaseTokenCounter | None:
         """

litellm/llms/bedrock/realtime/handler.py

@@ -20,6 +20,7 @@ from litellm.constants import (
     BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
     BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
     BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
+    REALTIME_SESSION_SUCCESS_LOGGED_KEY,
 )
 from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@@ -346,6 +347,7 @@ class BedrockRealtime(BaseAWSLLM):
                     prefer_async_handlers=True,
                 )
             )
+            logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
 
         if outcome.provider_failure is None:
             return
@@ -382,7 +384,7 @@ class BedrockRealtime(BaseAWSLLM):
         )
         bedrock_task: Final = asyncio.create_task(collect_logged_events())
 
-        await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_EXCEPTION)
+        await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_COMPLETED)
         client_disconnected: Final = (
             client_task.done() and not client_task.cancelled() and client_task.exception() is None
         )

litellm/proxy/_types.py

@@ -4418,6 +4418,9 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
     access_group_mcp_server_ids: list[str] | None = None
     access_group_agent_ids: list[str] | None = None
     access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None
+    # Parent org's model ceiling, reported only to callers who can manage the team.
+    # None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
+    organization_models: list[str] | None = None
 
 
 class TeamInfoResponseObject(TypedDict):

litellm/proxy/auth/auth_checks.py

@@ -328,9 +328,23 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads)
 last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100)
 db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL  # refresh every 5s
 
+_TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000
+_team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX)
+
+
+class _TeamMembershipCacheMiss:
+    __slots__ = ()
+
+
+_TEAM_MEMBERSHIP_CACHE_MISS: Final = _TeamMembershipCacheMiss()
+
 all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value
 
 
+def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None:
+    return result if isinstance(result, LiteLLM_TeamMembership) else None
+
+
 def _log_budget_lookup_failure(entity: str, error: Exception) -> None:
     """
     Log a warning when budget lookup fails; cache will not be populated.
@@ -888,6 +902,22 @@ async def common_checks(
         and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route))
     )
 
+    membership_user_id: Final = (
+        valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None
+    )
+    team_membership_loaded: Final = team_object is not None and membership_user_id is not None
+    loaded_team_membership: Final = (
+        await get_team_membership(
+            user_id=membership_user_id,
+            team_id=team_object.team_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+        if team_object is not None and membership_user_id is not None
+        else None
+    )
+
     unpriced_models: Final = (
         _unpriced_models_in_request(model=_model, llm_router=llm_router)
         if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route)
@@ -937,6 +967,8 @@ async def common_checks(
                 prisma_client=prisma_client,
                 user_api_key_cache=user_api_key_cache,
                 proxy_logging_obj=proxy_logging_obj,
+                team_membership=loaded_team_membership,
+                team_membership_loaded=team_membership_loaded,
             )
 
     # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
@@ -988,6 +1020,8 @@ async def common_checks(
         prisma_client=prisma_client,
         user_api_key_cache=user_api_key_cache,
         proxy_logging_obj=proxy_logging_obj,
+        team_membership=loaded_team_membership,
+        team_membership_loaded=team_membership_loaded,
     )
 
     # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body.
@@ -1097,6 +1131,8 @@ async def common_checks(
                     prisma_client=prisma_client,
                     user_api_key_cache=user_api_key_cache,
                     proxy_logging_obj=proxy_logging_obj,
+                    team_membership=loaded_team_membership,
+                    team_membership_loaded=team_membership_loaded,
                 ),
                 _check_end_user_budget(end_user_obj=end_user_object, route=route)
                 if end_user_object is not None and end_user_object.litellm_budget_table is not None
@@ -2142,7 +2178,76 @@ async def get_tag_object(
     return tag_objects.get(tag_name)
 
 
+def _membership_from_cached_payload(
+    cached: object,
+) -> LiteLLM_TeamMembership | None | _TeamMembershipCacheMiss:
+    if cached is None:
+        return _TEAM_MEMBERSHIP_CACHE_MISS
+    if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
+        return None
+    cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
+    return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS
+
+
 @log_db_metrics
+async def _fetch_team_membership_from_db(
+    user_id: str,
+    team_id: str,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: Span | None = None,
+    proxy_logging_obj: ProxyLogging | None = None,
+) -> LiteLLM_TeamMembership | None:
+    _ = parent_otel_span, proxy_logging_obj
+    response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique(
+        where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
+        include={"litellm_budget_table": True},
+    )
+    membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict())
+    _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
+    if membership is None:
+        await user_api_key_cache.async_set_cache(
+            key=_key,
+            value=NO_TEAM_MEMBERSHIP_SENTINEL,
+            ttl=get_management_object_ttl(user_api_key_cache),
+        )
+    else:
+        await user_api_key_cache.async_set_cache(
+            key=_key,
+            value=membership,
+            model_type=LiteLLM_TeamMembership,
+        )
+    return membership
+
+
+async def _load_team_membership_on_cache_miss(
+    user_id: str,
+    team_id: str,
+    cache_key: str,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: Span | None,
+    proxy_logging_obj: ProxyLogging | None,
+) -> LiteLLM_TeamMembership | None:
+    try:
+        redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
+        redis_membership: Final = _membership_from_cached_payload(redis_cached)
+        if not isinstance(redis_membership, _TeamMembershipCacheMiss):
+            return redis_membership
+
+        return await _fetch_team_membership_from_db(
+            user_id=user_id,
+            team_id=team_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=parent_otel_span,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    except Exception:
+        verbose_proxy_logger.exception("Error getting team membership")
+        return None
+
+
 async def get_team_membership(
     user_id: str,
     team_id: str,
@@ -2156,54 +2261,42 @@ async def get_team_membership(
 
     Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership).
     """
-    from litellm.proxy._types import LiteLLM_TeamMembership
-
-    if prisma_client is None:
-        raise Exception("No db connected")
-
     if user_id is None or team_id is None:
         return None
 
     _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
 
-    # check if in cache
-    cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key)
-    if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
-        return None
-    cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
-    if cached_membership_obj is not None:
-        return cached_membership_obj
+    l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True)
+    l1_membership: Final = _membership_from_cached_payload(l1_cached)
+    if not isinstance(l1_membership, _TeamMembershipCacheMiss):
+        return l1_membership
 
-    # else, check db
-    try:
-        response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique(
-            where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
-            include={"litellm_budget_table": True},
+    inflight: Final[object] = _team_membership_inflight.get(_key)
+    if isinstance(inflight, asyncio.Task):
+        return _membership_from_shared_load(await asyncio.shield(inflight))
+
+    if prisma_client is None:
+        raise Exception("No db connected")
+
+    task: Final = asyncio.ensure_future(
+        _load_team_membership_on_cache_miss(
+            user_id=user_id,
+            team_id=team_id,
+            cache_key=_key,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=parent_otel_span,
+            proxy_logging_obj=proxy_logging_obj,
         )
+    )
+    _team_membership_inflight[_key] = task
 
-        if response is None:
-            await user_api_key_cache.async_set_cache(
-                key=_key,
-                value=NO_TEAM_MEMBERSHIP_SENTINEL,
-                ttl=get_management_object_ttl(user_api_key_cache),
-            )
-            return None
+    def _clear_inflight(_done: object) -> None:
+        if _team_membership_inflight.get(_key) is task:
+            _team_membership_inflight.pop(_key, None)
 
-        _response: Final = LiteLLM_TeamMembership.model_validate(response.dict())
-        await user_api_key_cache.async_set_cache(
-            key=_key,
-            value=_response,
-            model_type=LiteLLM_TeamMembership,
-        )
-
-        return _response
-    except Exception:
-        verbose_proxy_logger.exception(
-            "Error getting team membership for user_id: %s, team_id: %s",
-            user_id,
-            team_id,
-        )
-        return None
+    task.add_done_callback(_clear_inflight)
+    return _membership_from_shared_load(await asyncio.shield(task))
 
 
 def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool:
@@ -2662,6 +2755,12 @@ async def invalidate_team_member_spend_state(
         publish_auth_cache_invalidation,
     )
 
+    inflight: Final[object] = _team_membership_inflight.pop(
+        team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None
+    )
+    if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task():
+        await asyncio.wait((inflight,))
+
     if new_spend is not None:
         from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
 
@@ -4116,18 +4215,21 @@ async def _team_member_granted_models(
     prisma_client: PrismaClient,
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
+    team_membership: LiteLLM_TeamMembership | None = None,
+    team_membership_loaded: bool = False,
 ) -> Sequence[str]:
     """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team."""
     if team_object is None or valid_token.user_id is None:
         return ()
 
-    team_membership: Final = await get_team_membership(
-        user_id=valid_token.user_id,
-        team_id=team_object.team_id,
-        prisma_client=prisma_client,
-        user_api_key_cache=user_api_key_cache,
-        proxy_logging_obj=proxy_logging_obj,
-    )
+    if not team_membership_loaded:
+        team_membership = await get_team_membership(
+            user_id=valid_token.user_id,
+            team_id=team_object.team_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
     return () if team_membership is None else _member_allowed_models(team_membership)
 
 
@@ -4163,6 +4265,8 @@ async def _granted_model_lists(
     prisma_client: PrismaClient,
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
+    team_membership: LiteLLM_TeamMembership | None = None,
+    team_membership_loaded: bool = False,
 ) -> tuple[Sequence[str], ...]:
     """One model allowlist per level that participates in authorizing the request."""
     return (
@@ -4174,6 +4278,8 @@ async def _granted_model_lists(
             prisma_client=prisma_client,
             user_api_key_cache=user_api_key_cache,
             proxy_logging_obj=proxy_logging_obj,
+            team_membership=team_membership,
+            team_membership_loaded=team_membership_loaded,
         ),
         project_object.models if project_object is not None else (),
         await _org_granted_models(
@@ -4268,6 +4374,8 @@ async def collect_matched_model_access_groups(
     prisma_client: PrismaClient | None,
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
+    team_membership: LiteLLM_TeamMembership | None = None,
+    team_membership_loaded: bool = False,
 ) -> tuple[str, ...]:
     """
     The budgeted model access groups that authorized this request, sorted and deduplicated.
@@ -4313,6 +4421,8 @@ async def collect_matched_model_access_groups(
             prisma_client=prisma_client,
             user_api_key_cache=user_api_key_cache,
             proxy_logging_obj=proxy_logging_obj,
+            team_membership=team_membership,
+            team_membership_loaded=team_membership_loaded,
         )
         for granted_model in granted_models
     )
@@ -4328,6 +4438,8 @@ async def stamp_matched_model_access_groups(
     prisma_client: PrismaClient | None,
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
+    team_membership: LiteLLM_TeamMembership | None = None,
+    team_membership_loaded: bool = False,
 ) -> tuple[str, ...]:
     """Record the groups that authorized this request on its auth object, for the post-call spend
     writer and the reservation counters, and hand them back for the budget check."""
@@ -4344,6 +4456,8 @@ async def stamp_matched_model_access_groups(
             prisma_client=prisma_client,
             user_api_key_cache=user_api_key_cache,
             proxy_logging_obj=proxy_logging_obj,
+            team_membership=team_membership,
+            team_membership_loaded=team_membership_loaded,
         )
     except Exception as e:  # noqa: BLE001  # fail-safe: attribution is spend telemetry, it must never break auth
         verbose_proxy_logger.debug("model access group attribution failed: %s", e)
@@ -5146,6 +5260,8 @@ async def _check_team_member_budget(
     prisma_client: PrismaClient | None,
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
+    team_membership: LiteLLM_TeamMembership | None = None,
+    team_membership_loaded: bool = False,
 ):
     """Check if team member is over their max budget within the team."""
     if (
@@ -5154,23 +5270,25 @@ async def _check_team_member_budget(
         and valid_token is not None
         and valid_token.user_id is not None
     ):
-        team_membership: Final = await get_team_membership(
-            user_id=valid_token.user_id,
-            team_id=team_object.team_id,
-            prisma_client=prisma_client,
-            user_api_key_cache=user_api_key_cache,
-            proxy_logging_obj=proxy_logging_obj,
-        )
+        if not team_membership_loaded:
+            team_membership = await get_team_membership(
+                user_id=valid_token.user_id,
+                team_id=team_object.team_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+        loaded_membership = team_membership
 
         # Per-member override wins; otherwise fall back to the team-level
         # default configured via team.metadata["team_member_budget_id"].
         team_member_budget: float | None = None
         if (
-            team_membership is not None
-            and team_membership.litellm_budget_table is not None
-            and team_membership.litellm_budget_table.max_budget is not None
+            loaded_membership is not None
+            and loaded_membership.litellm_budget_table is not None
+            and loaded_membership.litellm_budget_table.max_budget is not None
         ):
-            team_member_budget = team_membership.litellm_budget_table.max_budget
+            team_member_budget = loaded_membership.litellm_budget_table.max_budget
         else:
             default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
             if isinstance(default_budget_id, str):
@@ -5189,7 +5307,7 @@ async def _check_team_member_budget(
                     team_member_budget = default_budget.max_budget
 
         if team_member_budget is not None:
-            team_member_spend = (team_membership.spend if team_membership is not None else 0.0) or 0.0
+            team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0
 
             # Read from cross-pod counter (Redis-first) if available
             from litellm.proxy.proxy_server import get_current_spend
@@ -5218,6 +5336,8 @@ async def _check_team_member_model_access(
     prisma_client: Optional["PrismaClient"],
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
+    team_membership: LiteLLM_TeamMembership | None = None,
+    team_membership_loaded: bool = False,
 ) -> None:
     """
     Check if a team member's per-member model scope allows access to the requested model.
@@ -5228,22 +5348,24 @@ async def _check_team_member_model_access(
     if valid_token.user_id is None or team_object.team_id is None:
         return
 
-    team_membership: Final = await get_team_membership(
-        user_id=valid_token.user_id,
-        team_id=team_object.team_id,
-        prisma_client=prisma_client,
-        user_api_key_cache=user_api_key_cache,
-        proxy_logging_obj=proxy_logging_obj,
-    )
+    if not team_membership_loaded:
+        team_membership = await get_team_membership(
+            user_id=valid_token.user_id,
+            team_id=team_object.team_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    loaded_membership = team_membership
 
     if (
-        team_membership is None
-        or team_membership.litellm_budget_table is None
-        or not team_membership.litellm_budget_table.allowed_models
+        loaded_membership is None
+        or loaded_membership.litellm_budget_table is None
+        or not loaded_membership.litellm_budget_table.allowed_models
     ):
         return  # no per-member restriction — inherit team-level check
 
-    member_allowed_models: Final[list[str]] = team_membership.litellm_budget_table.allowed_models
+    member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models
     try:
         _can_object_call_model(
             model=model,

litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py

@@ -44,7 +44,7 @@ class JavelinGuardrail(CustomGuardrail):
         application: str | None = None,
         **kwargs,
     ):
-        f"""
+        """
         Initialize the JavelinGuardrail class.
 
         This calls: {api_base}/{api_version}/guardrail/{guardrail_name}/apply

litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py

@@ -15,7 +15,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
     # We check the raw guardrail dict because LitellmParams normalizes None → False,
     # making it impossible to distinguish "not set" from "explicitly false" via litellm_params.
     _raw_default_on: Final = cast(dict[str, Any], guardrail).get("litellm_params", {}).get("default_on")
-    _default_on: Final = False if _raw_default_on is False else True
+    _default_on: Final = _raw_default_on is not False
 
     _callback: Final = MCPEndUserPermissionGuardrail(
         guardrail_name=guardrail.get("guardrail_name", ""),

litellm/proxy/management_endpoints/team_endpoints.py

@@ -252,7 +252,7 @@ class _RawTeamRow(_TeamIdRow, _ModelDumpRow, _ObjectPermissionRow, _TeamBudgetRo
     @property
     def members_with_roles(
         self,
-    ) -> Sequence[dict[str, object]] | None: ...
+    ) -> Sequence[dict[str, object]] | None: ...  # mutable-ok: prisma deserializes this JSON column into plain dicts
 
     @property
     def organization_id(self) -> str | None: ...
@@ -431,27 +431,26 @@ async def _refresh_cached_team(
     )
 
 
+async def _can_manage_team(
+    team_obj: LiteLLM_TeamTable,
+    user_api_key_dict: UserAPIKeyAuth,
+) -> bool:
+    """True for a proxy admin, an admin of this team, or an org admin for the team's organization."""
+    if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
+        return True
+
+    if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
+        return True
+
+    return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj)
+
+
 async def _verify_team_access(
     team_obj: LiteLLM_TeamTable,
     user_api_key_dict: UserAPIKeyAuth,
 ) -> None:
-    """
-    Verify the caller is authorized to manage the given team.
-
-    Access is granted if:
-    - Caller is a proxy admin, OR
-    - Caller is an org admin for the team's organization, OR
-    - Caller is a team admin of this team
-
-    Raises HTTPException(403) otherwise.
-    """
-    if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
-        return
-
-    if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
-        return
-
-    if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
+    """Raise HTTPException(403) unless the caller can manage the given team."""
+    if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict):
         return
 
     raise HTTPException(
@@ -2197,7 +2196,7 @@ async def update_team(
 
         if "metadata" in updated_kv:
             stored_metadata: Final[Mapping[str, JsonValue] | None] = (
-                {
+                {  # mutable-ok: the validator payload's isinstance guard requires a plain dict
                     key: value
                     for key, value in existing_team_row.metadata.items()
                     if key not in TeamMemberBudgetHandler.SYSTEM_MANAGED_METADATA_KEYS
@@ -2881,7 +2880,11 @@ async def _resolve_existing_member_user_ids(
         return frozenset()
 
     found: Final = await _user_id_rows_db(UserRepository(prisma_client)).find_many(
-        where={"user_id": {"in": sorted(requested_user_ids)}}
+        where={  # mutable-ok: Prisma query filters are dict-shaped
+            "user_id": {  # mutable-ok: Prisma query filters are dict-shaped
+                "in": sorted(requested_user_ids)
+            }
+        }
     )
     return frozenset(user.user_id for user in found or () if user.user_id is not None)
 
@@ -2935,7 +2938,7 @@ def _validate_member_user_id_provisioning(
     remaining: Final = len(unknown_user_ids) - _MAX_REPORTED_UNKNOWN_USER_IDS
     raise HTTPException(
         status_code=403,
-        detail={
+        detail={  # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape
             "error": (
                 "Only proxy admins can add a user_id that does not exist yet: {}{}. "
                 "Add the member by user_email to invite a new user, or ask a proxy admin "
@@ -2951,7 +2954,11 @@ def _members_audit_value(members: Sequence[Member]) -> str:
     The audit-log columns hold a JSON object, so the member list is nested
     under a key rather than serialized as a top-level array.
     """
-    return safe_dumps({"members_with_roles": tuple(member.model_dump() for member in members)})
+    return safe_dumps(
+        {  # mutable-ok: the audit-log JSON column rejects a top-level array, so this value must be an object
+            "members_with_roles": tuple(member.model_dump() for member in members)
+        }
+    )
 
 
 async def _create_team_member_add_audit_logs(
@@ -3638,7 +3645,7 @@ def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAu
 
 
 def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn:
-    detail: Final = {"error": message}
+    detail: Final = {"error": message}  # mutable-ok: HTTPException.detail takes a dict
     raise HTTPException(status_code=status_code, detail=detail)
 
 
@@ -3672,7 +3679,7 @@ def _validate_team_member_reset_spend_value(
 
 @router.post(
     "/team/{team_id}/member/{user_id}/reset_spend",
-    tags=["team management"],
+    tags=["team management"],  # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence
     dependencies=(Depends(user_api_key_auth),),
 )
 @management_endpoint_wrapper
@@ -3708,10 +3715,12 @@ async def reset_team_member_spend_fn(
     await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict)
     _check_not_resetting_own_spend(user_id=user_id, user_api_key_dict=user_api_key_dict)
 
-    membership_where: Final = {"user_id_team_id": {"user_id": user_id, "team_id": team_id}}
+    membership_where: Final = {  # mutable-ok: prisma client requires a plain dict where= argument
+        "user_id_team_id": {"user_id": user_id, "team_id": team_id}  # mutable-ok: same prisma where= argument
+    }
     _membership_row: Final = await _team_membership_db(prisma_client).find_unique(
         where=membership_where,
-        include={"litellm_budget_table": True},
+        include={"litellm_budget_table": True},  # mutable-ok: prisma client requires a plain dict include= argument
     )
     if _membership_row is None:
         _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.")
@@ -3722,7 +3731,7 @@ async def reset_team_member_spend_fn(
 
     await _team_membership_db(prisma_client).update(
         where=membership_where,
-        data={"spend": reset_to},
+        data={"spend": reset_to},  # mutable-ok: prisma client requires a plain dict data= argument
     )
 
     await invalidate_team_member_spend_state(
@@ -3732,7 +3741,7 @@ async def reset_team_member_spend_fn(
         new_spend=reset_to,
     )
 
-    return {
+    return {  # mutable-ok: matches this router's established untyped-response-dict convention
         "team_id": team_id,
         "user_id": user_id,
         "spend": reset_to,
@@ -4336,7 +4345,15 @@ async def _hydrate_member_user_details(
     """Attach ``user_alias`` and fill in a missing ``user_email`` from ``LiteLLM_UserTable`` in one query."""
     user_ids: Final = frozenset(m.user_id for m in members if m.user_id is not None)
     user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = (
-        await _user_db(prisma_client).find_many(where={"user_id": {"in": sorted(user_ids)}}) if user_ids else ()
+        await _user_db(prisma_client).find_many(
+            where={  # mutable-ok: Prisma query filters are dict-shaped
+                "user_id": {  # mutable-ok: Prisma query filters are dict-shaped
+                    "in": sorted(user_ids)
+                }
+            }
+        )
+        if user_ids
+        else ()
     )
     user_by_id: Final = MappingProxyType({u.user_id: u for u in user_rows})
 
@@ -4352,6 +4369,20 @@ async def _hydrate_member_user_details(
     return tuple(hydrate(m) for m in members)
 
 
+class _OrganizationModelsRow(BaseModel):
+    models: list[str] = []  # mutable-ok: pydantic field default
+
+
+class _TeamRowWithOrganization(BaseModel):
+    litellm_organization_table: _OrganizationModelsRow | None = None
+
+
+def _parent_organization_models(team_row: BaseModel) -> list[str] | None:
+    """Return the parent org's model allow-list, or None when the team has no org."""
+    organization: Final = _TeamRowWithOrganization.model_validate(team_row.model_dump()).litellm_organization_table
+    return organization.models if organization is not None else None
+
+
 async def _resolve_team_access_group_resources(
     _team_info: TeamInfoResponseObjectTeamTable,
 ) -> TeamInfoResponseObjectTeamTable:
@@ -4423,7 +4454,11 @@ async def team_info(
         try:
             team_info: BaseModel | None = await _team_db(prisma_client).find_unique(
                 where={"team_id": team_id},
-                include={"litellm_model_table": True, "object_permission": True},
+                include={
+                    "litellm_model_table": True,
+                    "object_permission": True,
+                    "litellm_organization_table": True,
+                },
             )
             if team_info is None:
                 raise Exception
@@ -4432,9 +4467,12 @@ async def team_info(
                 status_code=status.HTTP_404_NOT_FOUND,
                 detail={"message": f"Team not found, passed team id: {team_id}."},
             )
-        await validate_membership(
-            user_api_key_dict=user_api_key_dict,
-            team_table=LiteLLM_TeamTable.model_validate(team_info.model_dump()),
+        team_table: Final = LiteLLM_TeamTable.model_validate(team_info.model_dump())
+        await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table)
+        organization_models: Final[list[str] | None] = (
+            _parent_organization_models(team_info)
+            if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict)
+            else None
         )
 
         ## GET ALL KEYS ##
@@ -4493,7 +4531,12 @@ async def team_info(
             prisma_client=prisma_client,
             members=resolved_team_info.members_with_roles,
         )
-        hydrated_team_info: Final = resolved_team_info.model_copy(update={"members_with_roles": hydrated_members})
+        hydrated_team_info: Final = resolved_team_info.model_copy(
+            update={
+                "members_with_roles": hydrated_members,
+                "organization_models": organization_models,
+            }
+        )
 
         response_object: Final = TeamInfoResponseObject(
             team_id=team_id,
@@ -4746,7 +4789,7 @@ async def unblock_team(
 
 @router.get(
     "/team/metadata_schema",
-    tags=["team management"],
+    tags=["team management"],  # mutable-ok: fastapi's decorator signature types tags as a list
     dependencies=(Depends(user_api_key_auth),),
     response_model=TeamMetadataSchemaResponse,
 )
@@ -6028,13 +6071,13 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi
 def _daily_activity_error(*, status_code: int, message: str) -> HTTPException:
     """Single construction site for the `{"error": ...}` detail shape the
     /team/daily/activity endpoints have always returned."""
-    return HTTPException(status_code=status_code, detail={"error": message})
+    return HTTPException(status_code=status_code, detail={"error": message})  # mutable-ok: FastAPI JSON detail
 
 
 class _TeamDailyActivityScope(NamedTuple):
     team_ids: list[str] | None  # mutable-ok: downstream daily-activity signatures take str | list unions
     exclude_team_ids: list[str] | None  # mutable-ok: downstream daily-activity signatures take str | list unions
-    team_alias_metadata: dict[str, dict[str, object]]
+    team_alias_metadata: dict[str, dict[str, object]]  # mutable-ok: entity_metadata_field shape
     api_key_filter: str | list[str] | None  # mutable-ok: downstream daily-activity signatures take str | list unions
 
 
@@ -6337,7 +6380,7 @@ class _TeamUserSpendDbRow(TypedDict):
 @router.get(
     "/team/spend/by_user",
     response_model=TeamUserSpendResponse,
-    tags=["team management"],
+    tags=["team management"],  # mutable-ok: fastapi route tags must be a list
 )
 async def get_team_spend_by_user(
     user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],

litellm/proxy/management_helpers/team_metadata_validation.py

@@ -115,13 +115,13 @@ async def run_team_metadata_validation(
                 "error": f"custom_team_metadata_validate is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}"
             },
         )
-    if not (
-        inspect.iscoroutinefunction(validator) or inspect.iscoroutinefunction(getattr(validator, "__call__", None))
-    ):
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail={"error": "custom_team_metadata_validate must be an async function"},
-        )
+    if not inspect.iscoroutinefunction(validator):
+        validator_call: Final = getattr(validator, "__call__", None)  # noqa: B004  # value unwrap for the functor check
+        if not inspect.iscoroutinefunction(validator_call):
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail={"error": "custom_team_metadata_validate must be an async function"},
+            )
 
     try:
         raw_result: Final = await asyncio.wait_for(validator(payload), timeout=timeout_seconds)

litellm/proxy/proxy_server.py

@@ -274,6 +274,8 @@ from litellm.constants import (
     PROXY_BUDGET_RESCHEDULER_MAX_TIME,
     PROXY_BUDGET_RESCHEDULER_MIN_TIME,
     PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
+    REALTIME_SESSION_FAILURE_LOGGED_KEY,
+    REALTIME_SESSION_SUCCESS_LOGGED_KEY,
     ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
     USER_SPEND_ALERTS_JOB_ID,
     WEEKLY_SPEND_REPORT_JOB_ID,
@@ -11895,6 +11897,13 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth
     )
 
 
+async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None:
+    release_like_http_disconnect: Final = (
+        proxy_logging_obj._arelease_max_parallel_requests_on_disconnect  # pyright: ignore[reportPrivateUsage]  # shared
+    )
+    await release_like_http_disconnect(user_api_key_dict)
+
+
 async def _reject_realtime_session(
     websocket: WebSocket,
     user_api_key_dict: UserAPIKeyAuth,
@@ -11914,6 +11923,7 @@ async def _reject_realtime_session(
         await websocket.close(code=code, reason=reason)
     finally:
         await _release_realtime_budget_reservation(user_api_key_dict)
+        await _release_realtime_max_parallel_slot(user_api_key_dict)
 
 
 @app.websocket("/openai/v1/realtime")
@@ -12017,6 +12027,9 @@ async def realtime_websocket_endpoint(
             websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e)
         )
         return
+    except BaseException:
+        await _release_realtime_max_parallel_slot(user_api_key_dict)
+        raise
 
     # Phase 2: route to upstream LLM.
     try:
@@ -12046,12 +12059,10 @@ async def realtime_websocket_endpoint(
         except Exception:  # noqa: BLE001  # the lower layer may have closed the socket already; closing twice is not an error
             verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone")
     finally:
-        from litellm.litellm_core_utils.realtime_streaming import (
-            REALTIME_SESSION_SUCCESS_LOGGED_KEY,
-        )
-
         if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY):
             await _release_realtime_budget_reservation(user_api_key_dict)
+            if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY):
+                await _release_realtime_max_parallel_slot(user_api_key_dict)
 
 
 ######################################################################

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -159,7 +159,7 @@ def _get_spend_logs_metadata(
             requester_ip_address=None,
             additional_usage_values=None,
             applied_guardrails=None,
-            status=None or "success",
+            status="success",
             error_information=None,
             proxy_server_request=None,
             batch_models=None,

litellm/proxy/utils.py

@@ -4,6 +4,7 @@ import copy
 import hashlib
 import inspect
 import json
+import math
 import os
 import smtplib
 import ssl
@@ -6384,7 +6385,7 @@ class PrismaClient:
             return None
         try:
             value: Final = float(response_time_ms)
-            return value if value == value and value not in (float("inf"), float("-inf")) else None
+            return value if math.isfinite(value) else None
         except (ValueError, TypeError):
             verbose_proxy_logger.warning("Invalid response_time_ms value: %s", response_time_ms)
             return None

litellm/router.py

@@ -1532,6 +1532,18 @@ class Router:
             return False
         return sum(len(self.model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1
 
+    def team_model_has_alternatives(self, deployment_id: str) -> bool:
+        deployment: Final = self.get_deployment(model_id=deployment_id)
+        if deployment is None:
+            return False
+        team_id: Final = deployment.model_info.team_id
+        public_model_name: Final = deployment.model_info.team_public_model_name
+        if team_id is None or public_model_name is None:
+            return False
+        sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()
+        routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices])
+        return len(routable_siblings) > 1
+
     _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY})
 
     def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None:
@@ -4087,16 +4099,16 @@ class Router:
         models: Final = [m.strip() for m in model.split(",")]
 
         async def _async_completion_no_exceptions(
-            model: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any
+            model_name: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any
         ) -> ModelResponse | CustomStreamWrapper | Exception:
             """
             Wrapper around self.acompletion that catches exceptions and returns them as a result
             """
             try:
-                result = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs)
+                result = await self.acompletion(model=model_name, messages=messages, stream=stream, **kwargs)
                 return result
             except asyncio.CancelledError:
-                verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model)
+                verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model_name)
                 raise
             except Exception as e:
                 return e
@@ -4123,9 +4135,9 @@ class Router:
                 except KeyError:
                     pass
 
-        for model in models:
+        for model_name in models:
             task = asyncio.create_task(
-                _async_completion_no_exceptions(model=model, messages=messages, stream=stream, **kwargs)
+                _async_completion_no_exceptions(model_name=model_name, messages=messages, stream=stream, **kwargs)
             )
             pending_tasks.append(task)
 

litellm/router_utils/cooldown_handlers.py

@@ -343,8 +343,9 @@ def _should_cooldown_deployment(
     model_group: Final = litellm_router_instance.get_model_group(id=deployment)
     is_single_deployment_model_group = False
     if model_group is not None and len(model_group) == 1:
-        is_single_deployment_model_group = not litellm_router_instance.routing_group_has_alternatives(
-            requested_model_group
+        is_single_deployment_model_group = not (
+            litellm_router_instance.routing_group_has_alternatives(requested_model_group)
+            or litellm_router_instance.team_model_has_alternatives(deployment)
         )
 
     ## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level)

litellm/utils.py

@@ -6260,7 +6260,7 @@ def function_to_dict(input_function) -> dict:
             "enum": param_enum,
         }
 
-        parameters[param_name] = dict([(k, v) for k, v in param_dict.items() if isinstance(v, str)])
+        parameters[param_name] = {k: v for k, v in param_dict.items() if isinstance(v, str)}
 
         # Check if the parameter has no default value (i.e., it's required)
         if param.default == param.empty:

ruff.toml

@@ -4,11 +4,12 @@ lint.ignore = ["F405", "E402", "F403"]
 # That gives editors and `ruff check --fix` the diagnostic, which the gate script cannot.
 lint.extend-select = [
     "T20", "PGH004", "RUF008", "RUF009", "RUF100",
-    "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208",
-    "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PLW0133", "PYI030", "PYI041", "PYI064", "RET501",
-    "RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "TC005", "UP006", "UP007",
-    "UP008",
-    "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045",
+    "B004", "B018", "B021", "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790",
+    "PIE800", "PLC0208", "PLR0124", "PLR0402", "PLR0206", "PLR1704", "PLR1711", "PLR1730", "PLR2044",
+    "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", "RUF022", "RUF023", "RUF051", "S113",
+    "SIM114", "SIM118", "SIM201", "SIM211", "SIM222", "TC005", "UP006", "UP007", "UP008",
+    "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045",
+    "C404", "C419",
 ]
 # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip
 # `# noqa` directives that protect rules enforced elsewhere. List those codes as external

tests/proxy_behavior/auth/test_auth_object_prefetch.py

@@ -66,6 +66,7 @@ async def test_join_binds_the_membership_to_the_requested_team(prisma):
         cache = _frozen_cache()
         refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id)
         await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma)
+        assert cache.in_memory_cache.get_cache(f"org_id:{org_id}") is not None
 
         dead_db = _dead_db()
         membership = await get_team_membership(

tests/rust-python-harness/AGENTS.md

@@ -63,7 +63,7 @@ tests/rust-python-harness/
 - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr`
 - `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases
 - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses
-- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`)
+- `trace_parity/` prints every collected Python call under `litellm/` and every Rust span without comparing them; mappings only filter the separate unit-test mapping strategy. Before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`)
 - E2E and trace strategies load their registered module cases and run surface-specific execution from their folders
 - `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest
 - `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report

tests/rust-python-harness/cli/__init__.py

@@ -58,16 +58,29 @@ def _strategy_command(strategy: Strategy) -> click.Command:
                 help=runner_argument.help,
             )
         )
+    for runner_option in strategy.definition.runner_options:
+        name: Final = runner_option.option.removeprefix("--").replace("-", "_")
+        params.append(
+            click.Option(
+                (runner_option.option, name),
+                type=click.Choice(runner_option.choices),
+                help=runner_option.help,
+            )
+        )
 
     def run_strategy(
         sdk_functions: tuple[str, ...],
         surface: str | None = None,
         runner_args: tuple[str, ...] = (),
+        **runner_options: str | None,
     ) -> int:
         selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions))
         selected_surface: Final = cast(Surface | None, surface)
         cases: Final = select_cases((strategy,), selected_functions, selected_surface)
-        return run_command((strategy,), cases, runner_args)
+        option_args: Final = tuple(
+            f"--{name.replace('_', '-')}={value}" for name, value in runner_options.items() if value is not None
+        )
+        return run_command((strategy,), cases, (*runner_args, *option_args))
 
     return click.Command(
         strategy.id,

tests/rust-python-harness/cli/catalog.py

@@ -21,9 +21,7 @@ def _load_strategy_module(name: str, folder: Path, prefix: str | None) -> Module
     if prefix is not None:
         return importlib.import_module(f"{prefix}.{name}")
     module_name: Final = _synthetic_module_name(folder)
-    spec: Final = importlib.util.spec_from_file_location(
-        module_name, folder / "__init__.py"
-    )
+    spec: Final = importlib.util.spec_from_file_location(module_name, folder / "__init__.py")
     if spec is None or spec.loader is None:
         raise ValueError(f"{folder}: cannot load strategy package")
     module: Final = importlib.util.module_from_spec(spec)
@@ -59,9 +57,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy:
     if duplicates:
         raise ValueError(f"{folder}: duplicate strategy cases: {duplicates}")
     expected: Final = frozenset(
-        (surface, function)
-        for surface in (definition.surfaces or (None,))
-        for function in SDK_FUNCTIONS
+        (surface, function) for surface in (definition.surfaces or (None,)) for function in SDK_FUNCTIONS
     )
     actual: Final = frozenset(keys)
     if actual != expected:
@@ -73,8 +69,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy:
     incompatible: Final = tuple(
         (case.surface, case.sdk_function)
         for case in definition.cases
-        if case.spec.disposition is CaseDisposition.RUNNABLE
-        and not isinstance(case.spec, definition.runnable_spec)
+        if case.spec.disposition is CaseDisposition.RUNNABLE and not isinstance(case.spec, definition.runnable_spec)
     )
     if incompatible:
         raise ValueError(f"{folder}: runnable cases do not match {definition.runnable_spec.__name__}: {incompatible}")
@@ -102,14 +97,10 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy:
 def load_catalog(root: Path | None = None) -> tuple[Strategy, ...]:
     resolved: Final = STRATEGIES_ROOT if root is None else root
     prefix: Final = _STRATEGIES_PACKAGE.__name__ if resolved == STRATEGIES_ROOT else None
-    folders: Final = tuple(
-        info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg
-    )
+    folders: Final = tuple(info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg)
     if not folders:
         raise ValueError(f"No strategy packages found below {resolved}")
-    strategies: Final = tuple(
-        _load_strategy(name, resolved / name, prefix) for name in sorted(folders)
-    )
+    strategies: Final = tuple(_load_strategy(name, resolved / name, prefix) for name in sorted(folders))
     ids: Final = [strategy.id for strategy in strategies]
     if len(set(ids)) != len(ids):
         raise ValueError(f"Duplicate strategy id in {resolved}")

tests/rust-python-harness/cli/commands.py

@@ -21,8 +21,7 @@ def select_cases(
         case
         for strategy in strategies
         for case in strategy.cases
-        if (not sdk_functions or case.sdk_function in sdk_functions)
-        and (surface is None or case.surface == surface)
+        if (not sdk_functions or case.sdk_function in sdk_functions) and (surface is None or case.surface == surface)
     )
 
 
@@ -32,8 +31,7 @@ def run_command(
     runner_args: Sequence[str] = (),
 ) -> int:
     grouped: Final = {
-        strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id)
-        for strategy in strategies
+        strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) for strategy in strategies
     }
     visible: Final = tuple(strategy for strategy in strategies if grouped[strategy.id])
     runners: Final = tuple(replace(strategy, cases=grouped[strategy.id]) for strategy in visible)

tests/rust-python-harness/cli/test_cli.py

@@ -242,7 +242,7 @@ def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_titl
 def test_every_unavailable_case_finishes_and_explains_itself() -> None:
     section_titles: Final = {
         "e2e_parity": "End-to-end parity outcomes",
-        "trace_parity": "trace comparisons",
+        "trace_parity": "traces",
         "unit_tests_mapping": "Python/Rust unit-test mappings",
         "unit_tests_parity": "Python backend parity outcomes",
         "unit_tests_rust": "Native Rust unit-test outcomes",
@@ -359,6 +359,25 @@ def test_strategy_command_forwards_repeated_filters_and_runner_arguments(
     ]
 
 
+def test_trace_command_forwards_engine_and_scenario(monkeypatch: pytest.MonkeyPatch) -> None:
+    cli: Final = importlib.import_module("tests.rust-python-harness.cli")
+    captured: list[tuple[str, ...]] = []
+
+    def capture_run(
+        strategies: Sequence[Strategy],
+        cases: Sequence[HarnessCase],
+        runner_args: Sequence[str] = (),
+    ) -> int:
+        del strategies, cases
+        captured.append(tuple(runner_args))
+        return 0
+
+    monkeypatch.setattr(cli, "run_command", capture_run)
+
+    assert main(["run", "trace_parity", "--scenario", "async-mistral", "--engine", "python"]) == 0
+    assert captured == [("async-mistral", "--engine=python")]
+
+
 def test_omitted_surface_selects_every_strategy_surface(monkeypatch: pytest.MonkeyPatch) -> None:
     cli: Final = importlib.import_module("tests.rust-python-harness.cli")
     selected: list[str] = []

tests/rust-python-harness/conftest.py

@@ -19,9 +19,7 @@ def subprocess_test_environment(monkeypatch: pytest.MonkeyPatch) -> None:
 def cargo_project(tmp_path: Path) -> Callable[[str, str], Path]:
     def create(package: str, source: str) -> Path:
         manifest: Final = tmp_path / "Cargo.toml"
-        manifest.write_text(
-            f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n'
-        )
+        manifest.write_text(f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n')
         (tmp_path / "src").mkdir()
         (tmp_path / "src/lib.rs").write_text(source)
         return manifest

tests/rust-python-harness/shared/native_build.py

@@ -16,6 +16,11 @@ _RUST_ROOT: Final = "litellm-rust"
 _LOCKFILE: Final = "Cargo.lock"
 _SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"})
 _FAILURE_OUTPUT_LINES: Final = 15
+_TRACE_CHECK: Final = (
+    "from litellm.rust_bridge import get_native_bridge; "
+    "bridge = get_native_bridge(); "
+    "raise SystemExit(0 if bridge is not None and getattr(bridge, '_trace', None) is not None else 1)"
+)
 
 
 def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool:
@@ -73,6 +78,17 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]:
     return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:])
 
 
+def _installed_bridge_has_trace(repo_root: Path) -> bool:
+    completed: Final = subprocess.run(
+        (sys.executable, "-c", _TRACE_CHECK),
+        cwd=repo_root,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        check=False,
+    )
+    return completed.returncode == 0
+
+
 def trace_bridge_error() -> str | None:
     bridge: Final = get_native_bridge()
     if bridge is None:
@@ -85,7 +101,10 @@ def trace_bridge_error() -> str | None:
 def ensure_trace_bridge(repo_root: Path) -> str | None:
     native_path: Final = _native_module_path()
     native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None
-    if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)):
+    rebuild_required: Final = needs_rebuild(
+        native_mtime, _newest_source_mtime(repo_root)
+    ) or not _installed_bridge_has_trace(repo_root)
+    if rebuild_required:
         print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True)
         succeeded: Final
         output: Final

tests/rust-python-harness/shared/parity/fixtures/recording.py

@@ -191,6 +191,7 @@ class _RecordingHandler(LocalHttpHandler):
         self.end_headers()
         self.wfile.write(body)
 
+
 def _recording_provider(spec: UpstreamEndpoint) -> AbstractContextManager[_RecordingProvider]:
     return serve_in_thread(_RecordingProvider(spec))
 

tests/rust-python-harness/shared/parity/fixtures/store.py

@@ -15,6 +15,8 @@ from .cassette import deserialize_cassette, serialize_cassette
 from .recording import RecordedInteraction
 
 FIXTURE_SCHEMA_VERSION: Final = 1
+
+
 class FixtureInput(Protocol):
     def canonical_input(self) -> dict[str, object]: ...
 

tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py

@@ -45,8 +45,8 @@ class _Upstream(LocalHttpServer):
         super().__init__(("127.0.0.1", 0), _UpstreamHandler)
         self.response_status: Final = status
 
-class _UpstreamHandler(LocalHttpHandler):
 
+class _UpstreamHandler(LocalHttpHandler):
     def do_POST(self) -> None:
         length: Final = int(self.headers.get("content-length") or "0")
         self.rfile.read(length)
@@ -59,6 +59,7 @@ class _UpstreamHandler(LocalHttpHandler):
         self.end_headers()
         self.wfile.write(body)
 
+
 def _upstream(status: int = 200) -> AbstractContextManager[_Upstream]:
     return serve_in_thread(_Upstream(status))
 

tests/rust-python-harness/shared/parity/fixtures/test_recording.py

@@ -238,6 +238,7 @@ class _ControlledUpstreamHandler(LocalHttpHandler):
         self.end_headers()
         self.wfile.write(body)
 
+
 def _controlled_upstream(
     stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS,
 ) -> AbstractContextManager[_ControlledUpstream]:

tests/rust-python-harness/shared/reporting/models.py

@@ -180,15 +180,11 @@ class HarnessRun:
 
     @property
     def unique_checks(self) -> int:
-        return len(
-            {nodeid for result in self.results.values() for nodeid in result.collected}
-        )
+        return len({nodeid for result in self.results.values() for nodeid in result.collected})
 
     @property
     def completed_checks(self) -> int:
-        return len(
-            {nodeid for result in self.results.values() for nodeid in result.completed}
-        )
+        return len({nodeid for result in self.results.values() for nodeid in result.completed})
 
     @classmethod
     def from_cases(cls, cases: Iterable[HarnessCase]) -> HarnessRun:

tests/rust-python-harness/shared/reporting/strategy.py

@@ -67,6 +67,13 @@ class RunnerArgumentDefinition:
     metavar: str = "ARG"
 
 
+@dataclass(frozen=True, slots=True)
+class RunnerOptionDefinition:
+    option: str
+    help: str
+    choices: tuple[str, ...]
+
+
 class StrategyRunner(Protocol):
     def __call__(
         self,
@@ -90,3 +97,4 @@ class StrategyDefinition:
     render: StrategyRenderer
     surfaces: tuple[Surface, ...] = ()
     runner_argument: RunnerArgumentDefinition | None = None
+    runner_options: tuple[RunnerOptionDefinition, ...] = ()

tests/rust-python-harness/shared/test_native_build.py

@@ -89,8 +89,8 @@ def test_ensure_trace_bridge_reports_failed_rebuild(tmp_path: Final, monkeypatch
     assert "boom" in message
 
 
-def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild(
-    tmp_path: Final, monkeypatch: pytest.MonkeyPatch
+def test_ensure_trace_bridge_rebuilds_when_trace_feature_is_missing(
+    tmp_path: Final, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
 ) -> None:
     native: Final = tmp_path / "_native.abi3.so"
     native.write_bytes(b"")
@@ -105,12 +105,17 @@ def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild(
         state.rebuilt = True
         return True, ""
 
+    def fake_get_native_bridge() -> SimpleNamespace:
+        assert state.rebuilt
+        return SimpleNamespace(_trace=object())
+
     monkeypatch.setattr(native_build, "_native_module_path", lambda: native)
     monkeypatch.setattr(native_build, "_rebuild", fake_rebuild)
-    monkeypatch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=None))
+    monkeypatch.setattr(native_build, "_installed_bridge_has_trace", lambda repo_root: False)
+    monkeypatch.setattr(native_build, "get_native_bridge", fake_get_native_bridge)
 
     message: Final = native_build.ensure_trace_bridge(tmp_path)
 
-    assert message is not None
-    assert "_trace" in message
-    assert state.rebuilt is False
+    assert message is None
+    assert state.rebuilt is True
+    assert "Rebuilding native Rust bridge" in capsys.readouterr().out

tests/rust-python-harness/shared/tracing/profiler.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 import sys
 import threading
-from collections.abc import Generator, Iterator, Mapping
+import warnings
+from collections.abc import Callable, Generator, Iterator, Mapping
 from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import lru_cache
@@ -32,6 +33,7 @@ class PythonProfiler:
         self._source_root: Final = str(source_root.resolve()) + "/"
         self._seen_frames: Final[set[FrameType]] = set()
         self._event_ids: Final[dict[FrameType, int]] = {}
+        self._lock: Final = threading.Lock()
         self.events: Final[list[FunctionTraceEvent]] = []
 
     def __call__(self, frame: FrameType, event: str, _arg: object) -> None:
@@ -40,14 +42,15 @@ class PythonProfiler:
         function_name: Final = self.function_name(frame)
         if function_name is None:
             return
-        event_id: Final = len(self.events)
-        parent_id: Final = next(
-            (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids),
-            None,
-        )
-        self._seen_frames.add(frame)
-        self._event_ids[frame] = event_id
-        self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name))
+        with self._lock:
+            event_id: Final = len(self.events)
+            parent_id: Final = next(
+                (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids),
+                None,
+            )
+            self._seen_frames.add(frame)
+            self._event_ids[frame] = event_id
+            self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name))
 
     def function_name(self, frame: FrameType) -> str | None:
         code: Final = frame.f_code
@@ -137,21 +140,51 @@ def _frame_ancestors(frame: FrameType) -> Generator[FrameType]:
 
 
 @contextmanager
-def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]:
-    profiler: Final = PythonProfiler(source_root)
+def _installed_profiler(profiler: Callable[[FrameType, str, object], None], *, threads: bool) -> Generator[None]:
+    if threads and sys.version_info >= (3, 12):
+        tool_id: Final = next((slot for slot in (2, 3, 4, 0, 1, 5) if sys.monitoring.get_tool(slot) is None), None)
+        if tool_id is None:
+            raise RuntimeError("no sys.monitoring tool ID is available for Python trace collection")
+
+        def started(_code: CodeType, _offset: int) -> None:
+            profiler(sys._getframe(1), "call", None)
+
+        sys.monitoring.use_tool_id(tool_id, "litellm-python-trace")
+        try:
+            sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, started)
+            sys.monitoring.set_events(tool_id, sys.monitoring.events.PY_START)
+            yield
+        finally:
+            sys.monitoring.set_events(tool_id, 0)
+            sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, None)
+            sys.monitoring.free_tool_id(tool_id)
+        return
+    if threads:
+        warnings.warn(
+            "Python <3.12 cannot trace existing worker threads; use Python 3.12+ for complete threaded traces",
+            RuntimeWarning,
+            stacklevel=3,
+        )
     previous_thread: Final = threading.getprofile()
     if threads:
         threading.setprofile(profiler)
     previous: Final = sys.getprofile()
     sys.setprofile(profiler)
     try:
-        yield profiler
+        yield
     finally:
         sys.setprofile(previous)
         if threads:
             threading.setprofile(previous_thread)
 
 
+@contextmanager
+def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]:
+    profiler: Final = PythonProfiler(source_root)
+    with _installed_profiler(profiler, threads=threads):
+        yield profiler
+
+
 @contextmanager
 def profile_python_function_usage(
     source_root: Path,
@@ -160,14 +193,5 @@ def profile_python_function_usage(
     threads: bool = False,
 ) -> Generator[PythonFunctionUsageProfiler]:
     profiler: Final = PythonFunctionUsageProfiler(source_root, functions)
-    previous_thread: Final = threading.getprofile()
-    if threads:
-        threading.setprofile(profiler)
-    previous: Final = sys.getprofile()
-    sys.setprofile(profiler)
-    try:
+    with _installed_profiler(profiler, threads=threads):
         yield profiler
-    finally:
-        sys.setprofile(previous)
-        if threads:
-            threading.setprofile(previous_thread)

tests/rust-python-harness/shared/tracing/steps.py

@@ -76,7 +76,7 @@ def _span_for(engine: Engine, function: str, mappings: Sequence[TraceMapping]) -
 
 
 def pipeline_projection(
-    engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping]
+    engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping] | None = None
 ) -> PipelineProjection:
     raw_parents: dict[int, int | None] = {}
     projected_ids: set[int] = set()
@@ -88,7 +88,7 @@ def pipeline_projection(
         if event.parent_id is not None and event.parent_id not in raw_parents:
             raise ValueError(f"trace event {event.id} references unknown or later parent {event.parent_id}")
         raw_parents[event.id] = event.parent_id
-        span = _span_for(engine, event.function, mappings)
+        span = event.function if mappings is None else _span_for(engine, event.function, mappings)
         if span is None:
             unmatched += 1
             continue
@@ -181,12 +181,7 @@ class TraceDiff:
 
     @property
     def matches(self) -> bool:
-        return (
-            not self.python_only
-            and not self.rust_only
-            and not self.missing_mappings
-            and self.shared_order_matches
-        )
+        return not self.python_only and not self.rust_only and not self.missing_mappings and self.shared_order_matches
 
 
 def _missing_mappings(
@@ -257,9 +252,7 @@ def trace_diff(
     rust_counts: Final = Counter(rust_spans)
     python_only_counts: Final = python_counts - rust_counts
     rust_only_counts: Final = rust_counts - python_counts
-    python_only: Final = tuple(
-        span for span, count in python_only_counts.items() for _ in range(count)
-    )
+    python_only: Final = tuple(span for span, count in python_only_counts.items() for _ in range(count))
     rust_only: Final = tuple(span for span, count in rust_only_counts.items() for _ in range(count))
     first_difference: Final = _first_difference(python, rust, mappings, contract)
     return TraceDiff(

tests/rust-python-harness/shared/tracing/test_profiler.py

@@ -4,9 +4,10 @@ import asyncio
 import sys
 import threading
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
 from functools import wraps
 from pathlib import Path
-from types import FunctionType
+from types import FrameType, FunctionType
 from typing import Final, ParamSpec, TypeVar, cast
 
 import pytest
@@ -41,11 +42,12 @@ def _events_named(profiler: PythonProfiler, name: str) -> tuple[FunctionTraceEve
     return tuple(event for event in profiler.events if event.function.endswith(name))
 
 
-def test_profiler_keeps_repeated_calls() -> None:
+@pytest.mark.parametrize("threads", (False, True))
+def test_profiler_keeps_repeated_calls(threads: bool) -> None:
     def called() -> None:
         return None
 
-    with profile_python(Path(__file__).parent) as profiler:
+    with profile_python(Path(__file__).parent, threads=threads) as profiler:
         called()
         called()
 
@@ -60,14 +62,15 @@ def test_profiler_qualifies_decorated_methods_by_class() -> None:
     assert _module_qualnames(__name__)[cast(FunctionType, Decorated.call.__wrapped__).__code__] == "Decorated.call"
 
 
-def test_profiler_records_real_frame_ancestry() -> None:
+@pytest.mark.parametrize("threads", (False, True))
+def test_profiler_records_real_frame_ancestry(threads: bool) -> None:
     def called() -> None:
         return None
 
     def outer() -> None:
         called()
 
-    with profile_python(Path(__file__).parent) as profiler:
+    with profile_python(Path(__file__).parent, threads=threads) as profiler:
         outer()
 
     outer_event, called_event = (event for event in profiler.events if event.function.endswith(("outer", "called")))
@@ -84,18 +87,20 @@ def test_profiler_restores_previous_profiler_after_failure() -> None:
     assert sys.getprofile() is previous
 
 
-def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None:
+@pytest.mark.parametrize("threads", (False, True))
+def test_profiler_does_not_count_coroutine_resumption_as_another_call(threads: bool) -> None:
     async def suspended() -> None:
         await asyncio.sleep(0)
         await asyncio.sleep(0)
 
-    with profile_python(Path(__file__).parent) as profiler:
+    with profile_python(Path(__file__).parent, threads=threads) as profiler:
         asyncio.run(suspended())
 
     assert len(_events_named(profiler, "suspended")) == 1
 
 
-def test_profiler_preserves_parent_across_coroutine_suspension() -> None:
+@pytest.mark.parametrize("threads", (False, True))
+def test_profiler_preserves_parent_across_coroutine_suspension(threads: bool) -> None:
     def called() -> None:
         return None
 
@@ -103,7 +108,7 @@ def test_profiler_preserves_parent_across_coroutine_suspension() -> None:
         await asyncio.sleep(0)
         called()
 
-    with profile_python(Path(__file__).parent) as profiler:
+    with profile_python(Path(__file__).parent, threads=threads) as profiler:
         asyncio.run(suspended())
 
     suspended_event: Final = _events_named(profiler, "suspended")[0]
@@ -124,6 +129,96 @@ def test_profiler_captures_worker_threads_when_enabled() -> None:
     assert called_event.parent_id is None
 
 
+@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring")
+@pytest.mark.parametrize("prewarm", (False, True))
+def test_profiler_captures_reused_workers_without_leaking_between_sessions(prewarm: bool) -> None:
+    def called() -> None:
+        return None
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        if prewarm:
+            executor.submit(called).result(timeout=5)
+        with profile_python(Path(__file__).parent, threads=True) as first:
+            executor.submit(called).result(timeout=5)
+        executor.submit(called).result(timeout=5)
+        with profile_python(Path(__file__).parent, threads=True) as second:
+            executor.submit(called).result(timeout=5)
+        executor.submit(called).result(timeout=5)
+
+    assert len(_events_named(first, "called")) == 1
+    assert len(_events_named(second, "called")) == 1
+
+
+def test_profiler_restores_main_and_worker_hooks_after_failure() -> None:
+    previous: Final = sys.getprofile()
+    previous_thread: Final = threading.getprofile()
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        worker_previous: Final = executor.submit(sys.getprofile).result(timeout=5)
+        with pytest.raises(RuntimeError, match="stop"):
+            with profile_python(Path(__file__).parent, threads=True):
+                raise RuntimeError("stop")
+        assert executor.submit(sys.getprofile).result(timeout=5) is worker_previous
+
+    assert sys.getprofile() is previous
+    assert threading.getprofile() is previous_thread
+
+
+@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring")
+def test_function_usage_profiler_captures_reused_workers() -> None:
+    def selected() -> None:
+        return None
+
+    function: Final = f"{Path(__file__).name}:{selected.__code__.co_firstlineno} {selected.__qualname__}"
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(selected).result(timeout=5)
+        with profile_python_function_usage(Path(__file__).parent, frozenset((function,)), threads=True) as profiler:
+            executor.submit(selected).result(timeout=5)
+
+    assert profiler.called == {function}
+
+
+@pytest.mark.skipif(sys.version_info < (3, 12), reason="independent thread hooks require sys.monitoring")
+def test_threaded_profiler_preserves_custom_worker_hook_and_releases_monitoring_slot() -> None:
+    def worker_hook(_frame: FrameType, _event: str, _arg: object) -> None:
+        return None
+
+    def fail_with_profile(executor: ThreadPoolExecutor) -> None:
+        with profile_python(Path(__file__).parent, threads=True):
+            assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook
+            raise RuntimeError("stop")
+
+    tools_before: Final = tuple(sys.monitoring.get_tool(slot) for slot in range(6))
+    with ThreadPoolExecutor(max_workers=1, initializer=lambda: sys.setprofile(worker_hook)) as executor:
+        assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook
+        with pytest.raises(RuntimeError, match="stop"):
+            fail_with_profile(executor)
+        assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook
+
+    assert tuple(sys.monitoring.get_tool(slot) for slot in range(6)) == tools_before
+
+
+@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring")
+def test_threaded_profiler_keeps_concurrent_event_ids_and_parent_links() -> None:
+    def child() -> None:
+        return None
+
+    def parent() -> None:
+        child()
+
+    with ThreadPoolExecutor(max_workers=4) as executor:
+        with profile_python(Path(__file__).parent, threads=True) as profiler:
+            futures: Final = tuple(executor.submit(parent) for _ in range(200))
+            for future in futures:
+                future.result(timeout=5)
+
+    parent_ids: Final = frozenset(event.id for event in _events_named(profiler, "parent"))
+    children: Final = _events_named(profiler, "child")
+    assert len(parent_ids) == len(children) == 200
+    assert frozenset(event.parent_id for event in children) == parent_ids
+    assert tuple(event.id for event in profiler.events) == tuple(range(len(profiler.events)))
+
+
 def test_function_usage_profiler_records_only_selected_functions() -> None:
     def selected() -> None:
         return None

tests/rust-python-harness/shared/tracing/test_steps.py

@@ -40,6 +40,23 @@ def test_python_projection_collapses_unmapped_parents_and_counts_noise() -> None
     ]
 
 
+@pytest.mark.parametrize("engine", ("python", "rust"))
+def test_projection_without_mappings_keeps_every_call_and_parent(engine: Engine) -> None:
+    events: Final = (
+        event(0, "module.py:1 entry"),
+        event(1, "module.py:2 internal_helper", 0),
+        event(2, "module.py:3 nested", 1),
+        event(3, "module.py:2 internal_helper", 0),
+    )
+
+    projection: Final = pipeline_projection(engine, events)
+
+    assert projection.unmatched == 0
+    assert tuple((step.id, step.parent_id, step.span, step.raw) for step in projection.steps) == tuple(
+        (item.id, item.parent_id, item.function, item.raw) for item in events
+    )
+
+
 def test_rust_projection_keeps_unknown_spans() -> None:
     projection: Final = pipeline_projection("rust", (event(0, "route"), event(1, "new_span", 0)), MAPPINGS)
     assert [(step.span, step.parent_id) for step in projection.steps] == [("route", None), ("new_span", 0)]
@@ -146,9 +163,7 @@ def test_trace_diff_allows_reordered_concurrent_children() -> None:
 def test_trace_diff_prunes_declared_engine_only_nodes_but_requires_them() -> None:
     mappings: Final = (MAPPINGS[0], mapping(rust_span="rust_prepare"))
     python: Final = pipeline_projection("python", (event(0, "module.py:1 entry"),), mappings).steps
-    rust: Final = pipeline_projection(
-        "rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings
-    ).steps
+    rust: Final = pipeline_projection("rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings).steps
     assert trace_diff(python, rust, mappings).matches
     assert trace_diff(python, rust[:1], mappings).missing_mappings == ("rust_prepare",)
 

tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py

@@ -264,9 +264,7 @@ def _formatting_strategy() -> SearchStrategy[ReductoFormatting]:
         ),
         st.sampled_from((False, True)).map(lambda value: {"add_page_markers": value}),
         st.sampled_from((False, True)).map(lambda value: {"merge_tables": value}),
-        st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS)
-        .map(list)
-        .map(lambda value: {"include": value}),
+        st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS).map(list).map(lambda value: {"include": value}),
     )
     return values.map(ReductoFormatting.model_validate)
 

tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py

@@ -124,10 +124,7 @@ class RecordingCallback(CustomLogger):
         if isinstance(value, Mapping):
             if any(not isinstance(map_key, str) for map_key in value):
                 raise TypeError("callback kwarg mappings must use string keys")
-            return {
-                map_key: self._normalized_kwargs(map_value, map_key)
-                for map_key, map_value in value.items()
-            }
+            return {map_key: self._normalized_kwargs(map_value, map_key) for map_key, map_value in value.items()}
         if isinstance(value, (list, tuple)):
             return [self._normalized_kwargs(item) for item in value]
         raise TypeError(f"unsupported callback kwarg type: {type(value)}")

tests/rust-python-harness/strategies/trace_parity/AGENTS.md

@@ -1 +1 @@
-Maps Python profiler frames onto feature-gated Rust span names via an explicit per-case mapping (Rust span name is the identity) and compares steps, order, and nesting of both live traces against a replayed provider response.
+Prints every collected Python call under litellm/ and every feature-gated Rust span from live traces against replayed HTTP responses. The two traces are independent and are not compared. API-key and Vertex credentials scenarios exercise separate authentication paths; credentials scenarios replay the token exchange locally.

tests/rust-python-harness/strategies/trace_parity/__init__.py

@@ -7,6 +7,7 @@ from ...shared.reporting.strategy import (
     ModuleCaseSpec,
     NotImplementedCaseSpec,
     RunnerArgumentDefinition,
+    RunnerOptionDefinition,
     StrategyDefinition,
 )
 from .reporting import render_trace_results
@@ -26,13 +27,17 @@ CASES: Final[tuple[CaseDefinition, ...]] = (
         ModuleCaseSpec(
             coverage=Coverage.PARTIAL,
             module="tests.rust-python-harness.strategies.trace_parity.sdk.messages.case",
-            note="Async only until anthropic_messages_handler supports sync calls.",
+            note="Success paths are async; sync tracing captures the currently unsupported behavior.",
         ),
         surface="sdk",
     ),
     CaseDefinition(
         "responses",
-        NotImplementedCaseSpec(reason="No Responses trace-parity case is registered."),
+        ModuleCaseSpec(
+            coverage=Coverage.PARTIAL,
+            module="tests.rust-python-harness.strategies.trace_parity.sdk.responses.case",
+            note="Core create paths: native, streaming, provider error, Azure override, and chat bridge.",
+        ),
         surface="sdk",
     ),
     CaseDefinition(
@@ -70,13 +75,17 @@ CASES: Final[tuple[CaseDefinition, ...]] = (
         ModuleCaseSpec(
             coverage=Coverage.PARTIAL,
             module="tests.rust-python-harness.strategies.trace_parity.gateway.messages.case",
-            note="Non-streaming success paths only.",
+            note="Anthropic/Azure provider routes plus a fully consumed downstream streaming path.",
         ),
         surface="gateway",
     ),
     CaseDefinition(
         "responses",
-        NotImplementedCaseSpec(reason="No gateway Responses trace-parity case is registered."),
+        ModuleCaseSpec(
+            coverage=Coverage.PARTIAL,
+            module="tests.rust-python-harness.strategies.trace_parity.gateway.responses.case",
+            note="Native OpenAI non-streaming and fully consumed downstream streaming paths.",
+        ),
         surface="gateway",
     ),
     CaseDefinition(
@@ -86,7 +95,11 @@ CASES: Final[tuple[CaseDefinition, ...]] = (
     ),
     CaseDefinition(
         "chat_completions",
-        NotImplementedCaseSpec(reason="No gateway chat trace-parity case is registered."),
+        ModuleCaseSpec(
+            coverage=Coverage.PARTIAL,
+            module="tests.rust-python-harness.strategies.trace_parity.gateway.chat_completions.case",
+            note="Anthropic non-streaming and fully consumed downstream streaming paths.",
+        ),
         surface="gateway",
     ),
     CaseDefinition(
@@ -99,8 +112,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = (
 STRATEGY: Final = StrategyDefinition(
     id="trace_parity",
     order=20,
-    label="Trace parity",
-    description="Compare pipeline steps, order, and nesting between Python profiler frames and Rust spans via an explicit mapping.",
+    label="Traces",
+    description="Print Python profiler frames and Rust spans for representative pipeline scenarios.",
     directory=Path(__file__).parent,
     runnable_spec=ModuleCaseSpec,
     cases=CASES,
@@ -112,4 +125,11 @@ STRATEGY: Final = StrategyDefinition(
         metavar="NAME",
         help="run only this named trace scenario; repeat to select more than one",
     ),
+    runner_options=(
+        RunnerOptionDefinition(
+            option="--engine",
+            choices=("python", "rust"),
+            help="show only this engine's trace; omit to print both engines",
+        ),
+    ),
 )

tests/rust-python-harness/strategies/trace_parity/fixtures.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import base64
+import binascii
+import json
+import struct
+from collections.abc import Iterable, Mapping
+from typing import Final
+
+from ...shared.parity.recorded_http import (
+    HttpHeader,
+    RecordedHttpResponse,
+    RecordedHttpStreamResponse,
+    RecordedStreamChunk,
+)
+
+JSON_HEADERS: Final = (HttpHeader(name="content-type", value="application/json"),)
+SSE_HEADERS: Final = (HttpHeader(name="content-type", value="text/event-stream"),)
+AWS_EVENT_STREAM_HEADERS: Final = (HttpHeader(name="content-type", value="application/vnd.amazon.eventstream"),)
+
+
+def json_response(body: Mapping[str, object] | bytes, *, status: int = 200) -> RecordedHttpResponse:
+    encoded: Final = body if isinstance(body, bytes) else json.dumps(body).encode()
+    return RecordedHttpResponse.from_bytes(status, JSON_HEADERS, encoded)
+
+
+def sse_event(event: str, payload: Mapping[str, object]) -> bytes:
+    return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode()
+
+
+def sse_response(events: Iterable[tuple[str, Mapping[str, object]]]) -> RecordedHttpStreamResponse:
+    return RecordedHttpStreamResponse(
+        kind="http_stream",
+        status_code=200,
+        headers=SSE_HEADERS,
+        chunks=tuple(RecordedStreamChunk.from_bytes(sse_event(event, payload)) for event, payload in events),
+    )
+
+
+def _aws_string_header(name: str, value: str) -> bytes:
+    name_bytes: Final = name.encode()
+    value_bytes: Final = value.encode()
+    return (
+        struct.pack("!B", len(name_bytes))
+        + name_bytes
+        + struct.pack("!B", 7)
+        + struct.pack("!H", len(value_bytes))
+        + value_bytes
+    )
+
+
+def aws_event_stream_frame(payload: Mapping[str, object]) -> bytes:
+    event_payload: Final = json.dumps(
+        {"bytes": base64.b64encode(json.dumps(payload, separators=(",", ":")).encode()).decode()},
+        separators=(",", ":"),
+    ).encode()
+    headers: Final = (
+        _aws_string_header(":event-type", "chunk")
+        + _aws_string_header(":content-type", "application/json")
+        + _aws_string_header(":message-type", "event")
+    )
+    total_length: Final = 12 + len(headers) + len(event_payload) + 4
+    prelude: Final = struct.pack("!II", total_length, len(headers))
+    prelude_crc: Final = binascii.crc32(prelude) & 0xFFFFFFFF
+    prelude_crc_bytes: Final = struct.pack("!I", prelude_crc)
+    message_crc: Final = binascii.crc32(prelude_crc_bytes + headers + event_payload, prelude_crc) & 0xFFFFFFFF
+    return prelude + prelude_crc_bytes + headers + event_payload + struct.pack("!I", message_crc)
+
+
+def aws_event_stream_response(
+    events: Iterable[Mapping[str, object]], *, corrupt_last_frame: bool = False
+) -> RecordedHttpStreamResponse:
+    frames: Final = tuple(aws_event_stream_frame(event) for event in events)
+    body: Final = (
+        b"".join((*frames[:-1], frames[-1][:-1] + bytes((frames[-1][-1] ^ 0xFF,))))
+        if corrupt_last_frame
+        else b"".join(frames)
+    )
+    return RecordedHttpStreamResponse(
+        kind="http_stream",
+        status_code=200,
+        headers=AWS_EVENT_STREAM_HEADERS,
+        chunks=(RecordedStreamChunk.from_bytes(body),),
+    )
+
+
+def anthropic_response_body(*, model: str = "claude-sonnet-5") -> dict[str, object]:
+    return {
+        "id": "msg_trace",
+        "type": "message",
+        "role": "assistant",
+        "model": model,
+        "content": [{"type": "text", "text": "hello"}],
+        "stop_reason": "end_turn",
+        "stop_sequence": None,
+        "usage": {"input_tokens": 2, "output_tokens": 3},
+    }
+
+
+def anthropic_stream_events(*, model: str = "claude-sonnet-5") -> tuple[tuple[str, Mapping[str, object]], ...]:
+    return (
+        (
+            "message_start",
+            {
+                "type": "message_start",
+                "message": {
+                    "id": "msg_trace",
+                    "type": "message",
+                    "role": "assistant",
+                    "model": model,
+                    "content": [],
+                    "stop_reason": None,
+                    "stop_sequence": None,
+                    "usage": {"input_tokens": 2, "output_tokens": 0},
+                },
+            },
+        ),
+        (
+            "content_block_start",
+            {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+        ),
+        (
+            "content_block_delta",
+            {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}},
+        ),
+        ("content_block_stop", {"type": "content_block_stop", "index": 0}),
+        (
+            "message_delta",
+            {
+                "type": "message_delta",
+                "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                "usage": {"output_tokens": 1},
+            },
+        ),
+        ("message_stop", {"type": "message_stop"}),
+    )
+
+
+def responses_body(*, model: str = "gpt-5", status: str = "completed") -> dict[str, object]:
+    return {
+        "id": "resp_trace",
+        "object": "response",
+        "created_at": 1_750_000_000,
+        "status": status,
+        "model": model,
+        "output": [
+            {
+                "type": "message",
+                "id": "msg_trace",
+                "status": "completed",
+                "role": "assistant",
+                "content": [{"type": "output_text", "text": "hello", "annotations": []}],
+            }
+        ],
+        "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5},
+    }
+
+
+def responses_stream_events(*, model: str = "gpt-5") -> tuple[tuple[str, Mapping[str, object]], ...]:
+    response: Final = responses_body(model=model)
+    return (
+        (
+            "response.created",
+            {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}},
+        ),
+        (
+            "response.output_text.delta",
+            {
+                "type": "response.output_text.delta",
+                "item_id": "msg_trace",
+                "output_index": 0,
+                "content_index": 0,
+                "delta": "hello",
+            },
+        ),
+        ("response.completed", {"type": "response.completed", "response": response}),
+    )

tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import Final
+
+from .....shared.tracing.steps import Engine, mapping
+from ...fixtures import anthropic_response_body, anthropic_stream_events, json_response, sse_response
+from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite
+
+MAPPINGS: Final = (
+    mapping(span="python_chat_gateway_route", python_frame=r"proxy_server\.py:\d+ chat_completion$"),
+    mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"),
+    mapping(span="python_chat_entrypoint", python_frame=r"main\.py:\d+ a?completion$"),
+    mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_chat_config$"),
+    mapping(rust_span="validate_environment", python_frame=r"(?<!_)validate_environment$"),
+    mapping(rust_span="transform_request", python_frame=r"AnthropicConfig\.transform_request$"),
+    mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"),
+    mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
+    mapping(rust_span="transform_response", python_frame=r"AnthropicConfig\.transform_response$"),
+    mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$|Logging\.success_handler$"),
+)
+
+STREAM_MAPPINGS: Final = (
+    mapping(span="python_stream_wrapper", python_frame=r"CustomStreamWrapper\.__init__$"),
+    mapping(span="python_stream_next", python_frame=r"CustomStreamWrapper\.__anext__$"),
+    mapping(span="python_stream_chunk", python_frame=r"CustomStreamWrapper\.chunk_creator$"),
+    mapping(span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$"),
+)
+
+
+def _fixture(_engine: Engine, _base_url: str) -> RouteFixture:
+    return RouteFixture(
+        kwargs={
+            "model_alias": "trace-model",
+            "provider_model": "anthropic/claude-sonnet-5",
+            "body": {
+                "model": "trace-model",
+                "messages": [{"role": "user", "content": "hello"}],
+                "max_tokens": 16,
+            },
+        },
+        provider_responses=(json_response(anthropic_response_body()),),
+    )
+
+
+def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture:
+    fixture: Final = _fixture(engine, base_url)
+    return fixture.with_body(stream=True).derive(
+        provider_responses=(sse_response(anthropic_stream_events()),),
+    )
+
+
+TRACE_SUITE: Final = TraceSuite(
+    route=GatewayRouteSpec("chat_completions", rust_supported=False),
+    scenarios=(
+        TraceScenario(name="async-anthropic", fixture=_fixture, mappings=MAPPINGS, asynchronous=True),
+        TraceScenario(
+            name="async-anthropic-downstream-stream",
+            fixture=_stream_fixture,
+            mappings=(*MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=True,
+        ),
+    ),
+)

tests/rust-python-harness/strategies/trace_parity/gateway/execution.py

@@ -13,8 +13,8 @@ from ....shared.parity.replay import replay_server
 from ....shared.tracing.native import TraceResponsePayload, native_trace_events
 from ....shared.tracing.profiler import FunctionTraceEvent, profile_python
 from ....shared.tracing.steps import Engine, PipelineProjection, pipeline_projection
-from ..models import GatewayRouteSpec, RouteFixture, TraceExecutionFailure, TraceMode, TraceScenario
-from ..reporting import TraceComparisonArtifact
+from ..models import GatewayRouteSpec, RouteFixture, TraceEngine, TraceExecutionFailure, TraceScenario
+from ..reporting import TraceArtifact
 
 
 class _GatewayResponsePayload(BaseModel):
@@ -28,7 +28,14 @@ class _GatewayClient(Protocol):
     def post(self, url: str, *, json: object, headers: dict[str, str]) -> httpx.Response: ...
 
 
-def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]:
+_ROUTE_PATHS: Final = {
+    "messages": "/v1/messages",
+    "chat_completions": "/v1/chat/completions",
+    "responses": "/v1/responses",
+}
+
+
+def _collect_python(fixture: RouteFixture, route: GatewayRouteSpec) -> tuple[FunctionTraceEvent, ...]:
     from fastapi.testclient import TestClient
 
     import litellm
@@ -61,7 +68,7 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]:
         with profile_python(Path(litellm.__file__).parent, threads=True) as profiler:
             client: Final = cast(_GatewayClient, TestClient(proxy_server.app))
             response: Final = client.post(
-                "/v1/messages",
+                _ROUTE_PATHS[route.route],
                 json=fixture.kwargs["body"],
                 headers={"authorization": "Bearer trace-key"},
             )
@@ -76,9 +83,10 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]:
             proxy_server.app.dependency_overrides[user_api_key_auth] = old_override
 
 
-def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]:
+def _collect_rust(fixture: RouteFixture, route: GatewayRouteSpec) -> tuple[FunctionTraceEvent, ...]:
     payload: Final = json.dumps(
         {
+            "path": _ROUTE_PATHS[route.route],
             "model_alias": fixture.kwargs["model_alias"],
             "provider_model": fixture.kwargs["provider_model"],
             "api_base": fixture.kwargs["api_base"],
@@ -130,7 +138,9 @@ def _gateway_trace_binary() -> Path:
     return rust_root / "target" / "debug" / "trace-parity-gateway"
 
 
-def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure:
+def _collect(
+    route: GatewayRouteSpec, scenario: TraceScenario, engine: Engine
+) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure:
     try:
         with replay_server() as provider:
             base_fixture: Final = scenario.fixture(engine, provider.url)
@@ -140,7 +150,7 @@ def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEven
             )
             for response in fixture.provider_responses:
                 provider.enqueue_response(response)
-            events: Final = _collect_python(fixture) if engine == "python" else _collect_rust(fixture)
+            events: Final = _collect_python(fixture, route) if engine == "python" else _collect_rust(fixture, route)
             provider.take_requests(len(fixture.provider_responses))
         return events
     except Exception as error:
@@ -150,40 +160,38 @@ def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEven
 def _projections(
     python_events: tuple[FunctionTraceEvent, ...],
     rust_events: tuple[FunctionTraceEvent, ...],
-    scenario: TraceScenario,
-    mode: TraceMode,
 ) -> tuple[PipelineProjection, PipelineProjection, str | None]:
-    mappings: Final = scenario.mappings_for(mode)
     try:
         return (
-            pipeline_projection("python", python_events, mappings),
-            pipeline_projection("rust", rust_events, mappings),
+            pipeline_projection("python", python_events),
+            pipeline_projection("rust", rust_events),
             None,
         )
     except ValueError as error:
         return PipelineProjection(), PipelineProjection(), f"harness: {error}"
 
 
-def execute_gateway_trace(route: GatewayRouteSpec, scenario: TraceScenario, mode: TraceMode) -> TraceComparisonArtifact:
-    mappings: Final = scenario.mappings_for(mode)
-    python_trace: Final = _collect(scenario, "python")
-    rust_trace: Final = _collect(scenario, "rust")
+def execute_gateway_trace(
+    route: GatewayRouteSpec,
+    scenario: TraceScenario,
+    engine: TraceEngine = "both",
+) -> TraceArtifact:
+    effective_engine: Final[TraceEngine] = "python" if engine == "both" and not route.rust_supported else engine
+    python_trace: Final = _collect(route, scenario, "python") if effective_engine != "rust" else ()
+    rust_trace: Final = _collect(route, scenario, "rust") if effective_engine != "python" else ()
     collection_python_error: Final = None if isinstance(python_trace, tuple) else f"python: {python_trace.message}"
     rust_error: Final = None if isinstance(rust_trace, tuple) else f"rust: {rust_trace.message}"
     python_events: Final = python_trace if isinstance(python_trace, tuple) else ()
     rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else ()
-    python, rust, projection_error = _projections(python_events, rust_events, scenario, mode)
+    python, rust, projection_error = _projections(python_events, rust_events)
     python_error: Final = projection_error or collection_python_error
-    return TraceComparisonArtifact.from_traces(
+    return TraceArtifact.from_traces(
+        engine=effective_engine,
         surface="gateway",
         sdk_function=route.route,
         scenario=scenario.name,
-        mode=mode,
-        mappings=mappings,
-        contract=scenario.contract,
         python=python.steps,
         rust=rust.steps,
-        python_unmatched=python.unmatched,
         python_error=python_error,
         rust_error=rust_error,
     )

tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-import json
 from typing import Final
 
-from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse
 from .....shared.tracing.steps import Engine, mapping
+from ...fixtures import anthropic_response_body, anthropic_stream_events, json_response, sse_response
 from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite
 
 
@@ -49,24 +48,7 @@ def _fixture(_engine: Engine, provider: str) -> RouteFixture:
                 "max_tokens": 16,
             },
         },
-        provider_responses=(
-            RecordedHttpResponse.from_bytes(
-                200,
-                (HttpHeader(name="content-type", value="application/json"),),
-                json.dumps(
-                    {
-                        "id": "msg_trace",
-                        "type": "message",
-                        "role": "assistant",
-                        "model": "claude-sonnet-5",
-                        "content": [{"type": "text", "text": "hello"}],
-                        "stop_reason": "end_turn",
-                        "stop_sequence": None,
-                        "usage": {"input_tokens": 2, "output_tokens": 3},
-                    }
-                ).encode(),
-            ),
-        ),
+        provider_responses=(json_response(anthropic_response_body()),),
     )
 
 
@@ -78,6 +60,13 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
     return _fixture(engine, "azure_ai")
 
 
+def _stream_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _anthropic_fixture(engine, _base_url)
+    return fixture.with_body(stream=True).derive(
+        provider_responses=(sse_response(anthropic_stream_events()),),
+    )
+
+
 ANTHROPIC_MAPPINGS: Final = (
     *GATEWAY_MAPPINGS,
     mapping(
@@ -100,7 +89,22 @@ AZURE_MAPPINGS: Final = (
 TRACE_SUITE: Final = TraceSuite(
     route=GatewayRouteSpec("messages"),
     scenarios=(
-        TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)),
-        TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)),
+        TraceScenario(
+            name="async-anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, asynchronous=True
+        ),
+        TraceScenario(name="async-azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True),
+        TraceScenario(
+            name="async-anthropic-downstream-stream",
+            fixture=_stream_fixture,
+            mappings=(
+                *ANTHROPIC_MAPPINGS,
+                mapping(span="python_upstream_stream", python_frame=r"AnthropicMessagesStreamingResponse\.__anext__$"),
+                mapping(
+                    span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$"
+                ),
+                mapping(span="python_stream_callback", python_frame=r"Logging\.async_success_handler$"),
+            ),
+            asynchronous=True,
+        ),
     ),
 )

tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from typing import Final
+
+from .....shared.tracing.steps import Engine, mapping
+from ...fixtures import json_response, responses_body, responses_stream_events, sse_response
+from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite
+
+MAPPINGS: Final = (
+    mapping(
+        span="python_responses_gateway_route", python_frame=r"response_api_endpoints/endpoints\.py:\d+ responses_api$"
+    ),
+    mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"),
+    mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"),
+    mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$"),
+    mapping(rust_span="validate_environment", python_frame=r"OpenAIResponsesAPIConfig\.validate_environment$"),
+    mapping(rust_span="complete_url", python_frame=r"OpenAIResponsesAPIConfig\.get_complete_url$"),
+    mapping(rust_span="transform_request", python_frame=r"OpenAIResponsesAPIConfig\.transform_responses_api_request$"),
+    mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"),
+    mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
+    mapping(rust_span="transform_response", python_frame=r"OpenAIResponsesAPIConfig\.transform_response_api_response$"),
+    mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$|Logging\.success_handler$"),
+)
+
+STREAM_MAPPINGS: Final = (
+    mapping(span="python_stream_iterator", python_frame=r"ResponsesAPIStreamingIterator\.__init__$"),
+    mapping(span="python_stream_next", python_frame=r"ResponsesAPIStreamingIterator\.__anext__$"),
+    mapping(span="python_stream_transform", python_frame=r"OpenAIResponsesAPIConfig\.transform_streaming_response$"),
+    mapping(span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$"),
+)
+
+
+def _fixture(_engine: Engine, _base_url: str) -> RouteFixture:
+    return RouteFixture(
+        kwargs={
+            "model_alias": "trace-model",
+            "provider_model": "openai/gpt-5",
+            "body": {"model": "trace-model", "input": "hello"},
+        },
+        provider_responses=(json_response(responses_body()),),
+    )
+
+
+def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture:
+    fixture: Final = _fixture(engine, base_url)
+    return fixture.with_body(stream=True).derive(
+        provider_responses=(sse_response(responses_stream_events()),),
+    )
+
+
+TRACE_SUITE: Final = TraceSuite(
+    route=GatewayRouteSpec("responses", rust_supported=False),
+    scenarios=(
+        TraceScenario(name="async-openai", fixture=_fixture, mappings=MAPPINGS, asynchronous=True),
+        TraceScenario(
+            name="async-openai-downstream-stream",
+            fixture=_stream_fixture,
+            mappings=(*MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=True,
+        ),
+    ),
+)

tests/rust-python-harness/strategies/trace_parity/models.py

@@ -1,35 +1,61 @@
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
-from typing import Final, Literal, TypeAlias
+from typing import Final, Literal, TypeAlias, cast
 
-from ...shared.parity.recorded_http import RecordedHttpResponse
+from ...shared.parity.recorded_http import RecordedResponse
 from ...shared.reporting.models import SdkFunction
-from ...shared.tracing.steps import Engine, TraceContract, TraceMapping
+from ...shared.tracing.steps import Engine, TraceMapping
 
-TraceMode = Literal["sync", "async"]
+TraceEngine = Literal["python", "rust", "both"]
 TraceFailureSource = Literal["python", "rust", "harness"]
 
 
 @dataclass(frozen=True, slots=True)
 class RouteFixture:
     kwargs: dict[str, object]
-    provider_responses: tuple[RecordedHttpResponse, ...]
+    provider_responses: tuple[RecordedResponse, ...]
     expected_failure: bool = False
+    consume_stream: bool = False
+    environment: tuple[tuple[str, str], ...] = ()
+
+    def derive(
+        self,
+        *,
+        kwargs: Mapping[str, object] | None = None,
+        provider_responses: tuple[RecordedResponse, ...] | None = None,
+        expected_failure: bool | None = None,
+        consume_stream: bool | None = None,
+    ) -> RouteFixture:
+        return RouteFixture(
+            kwargs={**self.kwargs, **(kwargs or {})},
+            provider_responses=self.provider_responses if provider_responses is None else provider_responses,
+            expected_failure=self.expected_failure if expected_failure is None else expected_failure,
+            consume_stream=self.consume_stream if consume_stream is None else consume_stream,
+            environment=self.environment,
+        )
+
+    def with_body(self, **updates: object) -> RouteFixture:
+        raw_body: Final = self.kwargs.get("body")
+        if not isinstance(raw_body, dict):
+            raise ValueError("route fixture does not contain an object body")
+        body: Final = cast(dict[str, object], raw_body)
+        return self.derive(kwargs={"body": {**body, **updates}})
 
 
 @dataclass(frozen=True, slots=True)
 class RouteSpec:
     route: SdkFunction
     python_entrypoints: tuple[str, str]
-    rust_entrypoints: tuple[str, str]
+    rust_entrypoints: tuple[str, str] | None
     fixture: Callable[[Engine, str], RouteFixture]
 
 
 @dataclass(frozen=True, slots=True)
 class GatewayRouteSpec:
     route: SdkFunction
+    rust_supported: bool = True
 
 
 TraceRouteSpec: TypeAlias = RouteSpec | GatewayRouteSpec
@@ -40,14 +66,7 @@ class TraceScenario:
     name: str
     fixture: Callable[[Engine, str], RouteFixture]
     mappings: tuple[TraceMapping, ...]
-    modes: tuple[TraceMode, ...] = ("sync", "async")
-    contract: TraceContract = TraceContract()
-    sync_mappings: tuple[TraceMapping, ...] | None = None
-    async_mappings: tuple[TraceMapping, ...] | None = None
-
-    def mappings_for(self, mode: TraceMode) -> tuple[TraceMapping, ...]:
-        selected: Final = self.async_mappings if mode == "async" else self.sync_mappings
-        return self.mappings if selected is None else selected
+    asynchronous: bool
 
 
 @dataclass(frozen=True, slots=True)

tests/rust-python-harness/strategies/trace_parity/reporting.py

@@ -1,31 +1,24 @@
 from __future__ import annotations
 
 import os
-import re
 import sys
 from collections.abc import Sequence
-from typing import Final, Literal
+from typing import Final
 
 from pydantic import BaseModel, ConfigDict, ValidationError
 
 from ...shared.reporting.models import SURFACES, CaseResult, RunStatus, SdkFunction, Surface
 from ...shared.reporting.rendering import ReportSection
 from ...shared.reporting.strategy import NotImplementedCaseSpec, SkippedCaseSpec
-from ...shared.tracing.steps import (
-    PipelineStep,
-    TraceContract,
-    TraceDiff,
-    TraceMapping,
-    trace_depths,
-    trace_diff,
-)
+from ...shared.tracing.steps import PipelineStep, trace_depths
+from .models import TraceEngine
 
-TRACE_COMPARISON_ARTIFACT: Final = "trace_comparison"
+TRACE_ARTIFACT: Final = "trace"
 TRACE_PARITY_HINT: Final = (
     "rebuild the native bridge with the trace-parity feature, e.g. `uvx maturin develop --features trace-parity`"
 )
 
-_COLORS: Final[dict[str, str]] = {"green": "32", "yellow": "33", "red": "31", "cyan": "36"}
+_COLORS: Final[dict[str, str]] = {"yellow": "33", "red": "31", "cyan": "36"}
 _RESET: Final = "\033[0m"
 
 
@@ -47,26 +40,15 @@ class TraceEventArtifact(BaseModel):
         return PipelineStep(self.id, self.parent_id, self.span, self.raw)
 
 
-class TraceMappingArtifact(BaseModel):
-    model_config = ConfigDict(frozen=True, extra="forbid")
-
-    span: str
-    python: str | None
-    rust: str | None
-
-
-class TraceComparisonArtifact(BaseModel):
+class TraceArtifact(BaseModel):
     model_config = ConfigDict(frozen=True, extra="forbid")
 
+    engine: TraceEngine = "both"
     surface: Surface
     sdk_function: SdkFunction
     scenario: str
-    mode: Literal["sync", "async"]
-    mappings: tuple[TraceMappingArtifact, ...]
     python: tuple[TraceEventArtifact, ...]
     rust: tuple[TraceEventArtifact, ...]
-    python_unmatched: int
-    unordered_children_of: frozenset[str]
     python_error: str | None = None
     rust_error: str | None = None
 
@@ -74,41 +56,27 @@ class TraceComparisonArtifact(BaseModel):
     def from_traces(
         cls,
         *,
+        engine: TraceEngine = "both",
         surface: Surface,
         sdk_function: SdkFunction,
         scenario: str,
-        mode: Literal["sync", "async"],
-        mappings: Sequence[TraceMapping],
-        contract: TraceContract,
         python: Sequence[PipelineStep],
         rust: Sequence[PipelineStep],
-        python_unmatched: int,
         python_error: str | None = None,
         rust_error: str | None = None,
-    ) -> TraceComparisonArtifact:
+    ) -> TraceArtifact:
         return cls(
+            engine=engine,
             surface=surface,
             sdk_function=sdk_function,
             scenario=scenario,
-            mode=mode,
-            mappings=tuple(
-                TraceMappingArtifact(
-                    span=item.span,
-                    python=item.python.pattern if item.python else None,
-                    rust=item.rust,
-                )
-                for item in mappings
-            ),
             python=tuple(
                 TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw)
                 for step in python
             ),
             rust=tuple(
-                TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw)
-                for step in rust
+                TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) for step in rust
             ),
-            python_unmatched=python_unmatched,
-            unordered_children_of=contract.unordered_children_of,
             python_error=python_error,
             rust_error=rust_error,
         )
@@ -119,32 +87,9 @@ class TraceComparisonArtifact(BaseModel):
     def rust_steps(self) -> tuple[PipelineStep, ...]:
         return tuple(event.step() for event in self.rust)
 
-    def diff(self) -> TraceDiff:
-        return trace_diff(
-            self.python_steps(),
-            self.rust_steps(),
-            tuple(
-                TraceMapping(
-                    item.span,
-                    re.compile(item.python) if item.python is not None else None,
-                    item.rust,
-                )
-                for item in self.mappings
-            ),
-            TraceContract(self.unordered_children_of),
-        )
-
-    def exact_match(self) -> bool:
-        return self.diff().matches
-
     def has_errors(self) -> bool:
         return self.python_error is not None or self.rust_error is not None
 
-    def contract_matches(self) -> bool:
-        if self.has_errors():
-            return False
-        return self.diff().matches
-
 
 def _split_raw(raw: str) -> tuple[str, str]:
     location, separator, name = raw.partition(" ")
@@ -153,72 +98,28 @@ def _split_raw(raw: str) -> tuple[str, str]:
     return raw, ""
 
 
-def _python_line(index: int, step: PipelineStep, depth: int, exclusive: frozenset[str]) -> str:
+def _python_line(index: int, step: PipelineStep, depth: int) -> str:
     name: Final = _split_raw(step.raw)[0]
     location: Final = _split_raw(step.raw)[1]
     suffix: Final = f"  ({location})" if location else ""
-    marker: Final = "  [python only]" if step.span in exclusive else ""
-    return _paint(f"{index} {'  ' * depth}{name}{suffix}{marker}", "cyan")
+    return _paint(f"{index} {'  ' * depth}{name}{suffix}", "cyan")
 
 
-def _python_lines(steps: tuple[PipelineStep, ...], exclusive: frozenset[str]) -> str:
+def _python_lines(steps: tuple[PipelineStep, ...]) -> str:
     depths: Final = trace_depths(steps)
-    lines: Final = tuple(
-        _python_line(index, step, depths[step.id], exclusive) for index, step in enumerate(steps, start=1)
-    )
+    lines: Final = tuple(_python_line(index, step, depths[step.id]) for index, step in enumerate(steps, start=1))
     return f"{_paint('PYTHON', 'cyan')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)")
 
 
-def _python_references(steps: tuple[PipelineStep, ...]) -> dict[tuple[str, int], str]:
-    references: dict[tuple[str, int], str] = {}
-    occurrences: dict[str, int] = {}
-    for index, step in enumerate(steps, start=1):
-        name = _split_raw(step.raw)[0]
-        occurrence = occurrences.get(step.span, 0) + 1
-        occurrences[step.span] = occurrence
-        references[(step.span, occurrence)] = f"{index} {name}"
-    return references
-
-
-def _rust_line(
-    step: PipelineStep,
-    depth: int,
-    occurrence: int,
-    references: dict[tuple[str, int], str],
-) -> str:
-    span: Final = _paint(step.span, "yellow")
-    key: Final = (step.span, occurrence)
-    reference: Final = (
-        _paint(references[key], "cyan") if key in references else _paint("[rust only]", "yellow")
-    )
-    suffix: Final = f"#{occurrence}" if occurrence > 1 else ""
-    return f"{'  ' * depth}{span}{suffix}  ->  {reference}"
-
-
-def _rust_lines(steps: tuple[PipelineStep, ...], references: dict[tuple[str, int], str]) -> str:
+def _rust_lines(steps: tuple[PipelineStep, ...]) -> str:
     depths: Final = trace_depths(steps)
-    occurrences: dict[str, int] = {}
-    lines: list[str] = []
-    for step in steps:
-        occurrence = occurrences.get(step.span, 0) + 1
-        occurrences[step.span] = occurrence
-        lines.append(_rust_line(step, depths[step.id], occurrence, references))
+    lines: Final = tuple(
+        _paint(f"{index} {'  ' * depths[step.id]}{step.span}", "yellow") for index, step in enumerate(steps, 1)
+    )
     return f"{_paint('RUST', 'yellow')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)")
 
 
-def _state_text(state: str, *, good: bool) -> str:
-    return _paint(state, "green" if good else "red")
-
-
-def _contract_line(artifact: TraceComparisonArtifact) -> str:
-    matches: Final = artifact.contract_matches()
-    status: Final = _state_text("PASS" if matches else "FAIL", good=matches)
-    if artifact.python_error or artifact.rust_error:
-        return f"Contract: {status}"
-    return f"Contract: {status}"
-
-
-def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]:
+def _error_lines(artifact: TraceArtifact) -> tuple[str, ...]:
     lines: list[str] = []
     for engine, error in (("Python", artifact.python_error), ("Rust", artifact.rust_error)):
         if error is None:
@@ -229,68 +130,20 @@ def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]:
     return tuple(lines)
 
 
-def _unseen_mappings(
-    artifact: TraceComparisonArtifact,
-    python: tuple[PipelineStep, ...],
-    rust: tuple[PipelineStep, ...],
-) -> tuple[str, ...]:
-    return artifact.diff().missing_mappings
-
-
-def _comparison_status_lines(
-    artifact: TraceComparisonArtifact,
-    python: tuple[PipelineStep, ...],
-    rust: tuple[PipelineStep, ...],
-) -> tuple[str, ...]:
-    diff: Final = artifact.diff()
-    exact_match: Final = artifact.exact_match()
-    if artifact.has_errors():
-        return (*_error_lines(artifact), _contract_line(artifact))
-    unseen: Final = _unseen_mappings(artifact, python, rust)
-    unseen_line: Final[tuple[str, ...]] = (f"Unseen mappings: {', '.join(unseen)}",) if unseen else ()
-    drift_lines: Final[tuple[str, ...]] = (
-        (_state_text("Same steps, order, and nesting", good=True),)
-        if exact_match
-        else (
-            _paint(f"Python only: {', '.join(diff.python_only) or 'none'}", "cyan"),
-            _paint(f"Rust only: {', '.join(diff.rust_only) or 'none'}", "yellow"),
-            f"First difference: {diff.first_difference or 'none'}",
-            f"Python frames outside mapping: {artifact.python_unmatched}",
-        )
-    )
-    return (
-        f"Trace: {_state_text('MATCH' if exact_match else 'DRIFT', good=exact_match)}",
-        *drift_lines,
-        *unseen_line,
-        _contract_line(artifact),
-    )
-
-
-def _render_comparison(artifact: TraceComparisonArtifact) -> str:
-    python: Final = artifact.python_steps()
-    rust: Final = artifact.rust_steps()
-    diff: Final = artifact.diff()
-    python_exclusive: Final = frozenset(item.span for item in artifact.mappings if item.rust is None)
-    status_lines: Final = _comparison_status_lines(artifact, python, rust)
-    return "\n\n".join(
-        (
-            _python_lines(python, python_exclusive | frozenset(diff.python_only)),
-            _rust_lines(rust, _python_references(python)),
-            "\n".join(status_lines),
-        )
-    )
-
-
-def _mode(nodeid: str) -> str:
-    if "[" in nodeid:
-        return nodeid.rsplit("[", 1)[-1].removesuffix("]")
-    head, _, tail = nodeid.rpartition(":")
-    return tail if head else "unknown mode"
+def _render_trace(artifact: TraceArtifact) -> str:
+    traces: tuple[str, ...]
+    if artifact.engine == "python":
+        traces = (_python_lines(artifact.python_steps()),)
+    elif artifact.engine == "rust":
+        traces = (_rust_lines(artifact.rust_steps()),)
+    else:
+        traces = (_python_lines(artifact.python_steps()), _rust_lines(artifact.rust_steps()))
+    return "\n\n".join((*traces, *_error_lines(artifact)))
 
 
 def _scenario(nodeid: str) -> str:
     parts: Final = nodeid.split(":")
-    return parts[-2] if len(parts) >= 5 else "default"
+    return parts[-1] if len(parts) >= 4 else "default"
 
 
 def _unavailable(status: RunStatus) -> str:
@@ -299,20 +152,20 @@ def _unavailable(status: RunStatus) -> str:
 
 def _render_artifact(body: str) -> str:
     try:
-        artifact: Final = TraceComparisonArtifact.model_validate_json(body)
+        artifact: Final = TraceArtifact.model_validate_json(body)
     except ValidationError as error:
-        return f"Trace comparison artifact is invalid: {error}"
-    return _render_comparison(artifact)
+        return f"Trace artifact is invalid: {error}"
+    return _render_trace(artifact)
 
 
-def _mode_section(result: CaseResult, nodeid: str, status: RunStatus) -> str:
+def _scenario_section(result: CaseResult, nodeid: str, status: RunStatus) -> str:
     artifacts: Final = tuple(
-        artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_COMPARISON_ARTIFACT
+        artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_ARTIFACT
     )
     body: Final = (
         "\n\n".join(_render_artifact(artifact.body) for artifact in artifacts) if artifacts else _unavailable(status)
     )
-    label: Final = f"Scenario: {_scenario(nodeid)} / Mode: {_mode(nodeid)}"
+    label: Final = f"Scenario: {_scenario(nodeid)}"
     return f"{label}\n{'-' * len(label)}\n\n{body}"
 
 
@@ -321,7 +174,7 @@ def _case_block(result: CaseResult) -> str:
     outcomes: Final = tuple(result.outcomes.items()) or (
         (nodeid, RunStatus.NOT_RUN) for nodeid in sorted(result.collected)
     )
-    sections: Final = tuple(_mode_section(result, nodeid, status) for nodeid, status in outcomes)
+    sections: Final = tuple(_scenario_section(result, nodeid, status) for nodeid, status in outcomes)
     return "\n\n".join((f"{header}\n{'=' * len(header)}", *sections))
 
 
@@ -357,11 +210,11 @@ def _surface_section(surface: Surface, results: Sequence[CaseResult]) -> ReportS
         *((not_implemented,) if not_implemented else ()),
         *((skipped,) if skipped else ()),
     )
-    return ReportSection(f"{surface.upper()} trace comparisons", blocks or ("No runnable trace comparisons",))
+    return ReportSection(f"{surface.upper()} traces", blocks or ("No runnable traces",))
 
 
 def render_trace_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]:
     sections: Final = tuple(
         section for surface in SURFACES if (section := _surface_section(surface, results)) is not None
     )
-    return sections or (ReportSection("Trace comparisons", ("No trace comparisons selected",)),)
+    return sections or (ReportSection("Traces", ("No traces selected",)),)

tests/rust-python-harness/strategies/trace_parity/runner.py

@@ -4,13 +4,20 @@ import importlib
 from collections.abc import Sequence
 from pathlib import Path
 from time import monotonic
-from typing import Final
+from typing import Final, cast
 
+from ...shared.native_build import ensure_trace_bridge
 from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, ResultArtifact, RunStatus, Surface
 from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback
-from ...shared.native_build import ensure_trace_bridge
-from .models import GatewayRouteSpec, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario, TraceSuite
-from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact
+from .models import (
+    GatewayRouteSpec,
+    RouteSpec,
+    TraceEngine,
+    TraceExecutionFailure,
+    TraceScenario,
+    TraceSuite,
+)
+from .reporting import TRACE_ARTIFACT, TraceArtifact
 from .sdk.execution import execute_trace
 
 
@@ -32,15 +39,13 @@ def validate_trace_suite(suite: TraceSuite, harness_case: HarnessCase) -> str |
     names: Final = tuple(scenario.name for scenario in suite.scenarios)
     if not names or len(names) != len(set(names)) or any(not name or ":" in name for name in names):
         return "scenario names must be non-empty, unique, and colon-free"
-    invalid_modes: Final = tuple(
+    invalid_names: Final = tuple(
         scenario.name
         for scenario in suite.scenarios
-        if not scenario.modes
-        or len(scenario.modes) != len(set(scenario.modes))
-        or any(mode not in {"sync", "async"} for mode in scenario.modes)
+        if not scenario.name.startswith("async-" if scenario.asynchronous else "sync-")
     )
-    if invalid_modes:
-        return f"scenarios must use non-empty, unique sync/async modes: {', '.join(invalid_modes)}"
+    if invalid_names:
+        return f"scenario names must start with sync- or async-: {', '.join(invalid_names)}"
     surface: Final = harness_case.surface
     if surface == "sdk" and not isinstance(suite.route, RouteSpec):
         return "must use RouteSpec for the sdk surface"
@@ -57,15 +62,14 @@ def scenario_nodeids(
     trace_suite: TraceSuite,
     harness_case: HarnessCase,
     selected_scenarios: frozenset[str] = frozenset(),
-) -> tuple[tuple[TraceScenario, TraceMode, str], ...]:
+) -> tuple[tuple[TraceScenario, str], ...]:
     surface: Final = harness_case.surface
     if surface is None:
         return ()
     return tuple(
-        (scenario, mode, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}:{mode}")
+        (scenario, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}")
         for scenario in trace_suite.scenarios
         if not selected_scenarios or scenario.name in selected_scenarios
-        for mode in scenario.modes
     )
 
 
@@ -77,54 +81,49 @@ def _record_setup_failure(run: HarnessRun, case: HarnessCase, message: str, stag
     run.failures.append((nodeid, message))
 
 
-def run_trace_mode(
+def run_trace_scenario(
     run: HarnessRun,
     result: CaseResult,
     trace_suite: TraceSuite,
     scenario: TraceScenario,
-    mode: TraceMode,
     surface: Surface,
     nodeid: str,
     on_update: UpdateCallback,
+    engine: TraceEngine = "both",
 ) -> None:
     started_at: Final = monotonic()
-    comparison: Final = _execute_mode(trace_suite, scenario, mode, surface)
+    trace: Final = _execute_scenario(trace_suite, scenario, surface, engine)
     duration: Final = monotonic() - started_at
-    if isinstance(comparison, TraceExecutionFailure):
+    if isinstance(trace, TraceExecutionFailure):
         result.record(nodeid, RunStatus.ERROR, duration)
-        run.failures.append((nodeid, comparison.message))
+        run.failures.append((nodeid, trace.message))
         on_update(run)
         return
-    artifact: Final = ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json())
-    if comparison.has_errors():
+    artifact: Final = ResultArtifact(TRACE_ARTIFACT, trace.model_dump_json())
+    if trace.has_errors():
         result.record(nodeid, RunStatus.ERROR, duration, (artifact,))
-        run.failures.append(
-            (nodeid, "\n".join(error for error in (comparison.python_error, comparison.rust_error) if error))
-        )
+        run.failures.append((nodeid, "\n".join(error for error in (trace.python_error, trace.rust_error) if error)))
     else:
-        status: Final = RunStatus.PASSED if comparison.contract_matches() else RunStatus.FAILED
-        result.record(nodeid, status, duration, (artifact,))
-        if status is RunStatus.FAILED:
-            run.failures.append((nodeid, "trace contract mismatch; see the rendered comparison"))
+        result.record(nodeid, RunStatus.PASSED, duration, (artifact,))
     on_update(run)
 
 
-def _execute_mode(
+def _execute_scenario(
     trace_suite: TraceSuite,
     scenario: TraceScenario,
-    mode: TraceMode,
     surface: Surface,
-) -> TraceComparisonArtifact | TraceExecutionFailure:
+    engine: TraceEngine,
+) -> TraceArtifact | TraceExecutionFailure:
     route: Final = trace_suite.route
     if isinstance(route, GatewayRouteSpec):
         if surface != "gateway":
             return TraceExecutionFailure("harness", "gateway route cannot run on the sdk surface")
         from .gateway.execution import execute_gateway_trace
 
-        return execute_gateway_trace(route, scenario, mode)
+        return execute_gateway_trace(route, scenario, engine)
     if surface != "sdk":
         return TraceExecutionFailure("harness", "sdk route cannot run on the gateway surface")
-    return execute_trace(route, scenario, mode, surface)
+    return execute_trace(route, scenario, surface, engine)
 
 
 def _run_case(
@@ -132,6 +131,7 @@ def _run_case(
     harness_case: HarnessCase,
     selected_scenarios: frozenset[str],
     on_update: UpdateCallback,
+    engine: TraceEngine,
 ) -> None:
     result: Final = run.results[harness_case.key]
     spec: Final = harness_case.spec
@@ -146,15 +146,29 @@ def _run_case(
         on_update(run)
         return
     nodeids: Final = scenario_nodeids(trace_suite, harness_case, selected_scenarios)
-    result.collected.update(nodeid for _, _, nodeid in nodeids)
+    result.collected.update(nodeid for _, nodeid in nodeids)
     if not nodeids:
         result.status = RunStatus.SKIPPED
         on_update(run)
         return
     result.status = RunStatus.RUNNING
     on_update(run)
-    for scenario, mode, nodeid in nodeids:
-        run_trace_mode(run, result, trace_suite, scenario, mode, surface, nodeid, on_update)
+    for scenario, nodeid in nodeids:
+        run_trace_scenario(run, result, trace_suite, scenario, surface, nodeid, on_update, engine)
+
+
+def runner_selection(runner_args: Sequence[str]) -> tuple[frozenset[str], TraceEngine]:
+    engine: TraceEngine = "both"
+    scenarios: list[str] = []
+    for argument in runner_args:
+        if argument.startswith("--engine="):
+            value = argument.removeprefix("--engine=")
+            if value not in {"python", "rust"}:
+                raise ValueError(f"invalid trace engine: {value}")
+            engine = cast(TraceEngine, value)
+        else:
+            scenarios.append(argument)
+    return frozenset(scenarios), engine
 
 
 def run_trace_cases(
@@ -163,10 +177,10 @@ def run_trace_cases(
     on_update: UpdateCallback,
     runner_args: Sequence[str] = (),
 ) -> tuple[int, HarnessRun]:
-    selected_scenarios: Final = frozenset(runner_args)
+    selected_scenarios, engine = runner_selection(runner_args)
     run: Final = HarnessRun.from_cases(cases)
     runnable_cases: Final = tuple(case for case in cases if isinstance(case.spec, ModuleCaseSpec))
-    bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases else None
+    bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases and engine != "python" else None
     if bridge_error is not None:
         for harness_case in runnable_cases:
             _record_setup_failure(run, harness_case, bridge_error, "bridge")
@@ -174,7 +188,7 @@ def run_trace_cases(
         on_update(run)
         return 1, run
     for harness_case in cases:
-        _run_case(run, harness_case, selected_scenarios, on_update)
+        _run_case(run, harness_case, selected_scenarios, on_update, engine)
     run.finished_at = monotonic()
     on_update(run)
     failed: Final = any(

tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
-import json
 from typing import Final
 
-from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse
 from .....shared.tracing.steps import Engine, mapping
+from ...fixtures import (
+    anthropic_response_body,
+    anthropic_stream_events,
+    aws_event_stream_response,
+    json_response,
+    sse_response,
+)
 from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite
 
 COMMON_MAPPINGS: Final = (
@@ -24,6 +29,22 @@ COMMON_MAPPINGS: Final = (
     mapping(rust_span="execute_chat_completions_provider_call"),
     mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
     mapping(rust_span="transform_response", python_frame=r"(?<!async_)transform_response$"),
+    mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"),
+    mapping(span="python_logging_post_call", python_frame=r"Logging\.post_call$"),
+    mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$|Logging\.success_handler$"),
+)
+
+STREAM_MAPPINGS: Final = (
+    mapping(span="python_stream_wrapper", python_frame=r"CustomStreamWrapper\.__init__$"),
+    mapping(span="python_stream_next", python_frame=r"CustomStreamWrapper\.__next__$|CustomStreamWrapper\.__anext__$"),
+    mapping(span="python_stream_chunk", python_frame=r"CustomStreamWrapper\.chunk_creator$"),
+    mapping(span="python_stream_finalize", python_frame=r"CustomStreamWrapper\._finalize_completed_stream$"),
+)
+
+FAILURE_MAPPINGS: Final = (
+    mapping(span="python_exception_mapping", python_frame=r"(?<!_)exception_type$"),
+    mapping(span="python_failure_callback", python_frame=r"Logging\.failure_handler$"),
+    mapping(span="python_async_failure_callback", python_frame=r"Logging\.async_failure_handler$"),
 )
 
 SYNC_MAPPINGS: Final = (
@@ -44,41 +65,23 @@ ASYNC_MAPPINGS: Final = (
 
 
 def _anthropic_fixture(engine: Engine, _base_url: str) -> RouteFixture:
-    response: Final = json.dumps(
-        {
-            "id": "msg_trace",
-            "type": "message",
-            "role": "assistant",
-            "model": "claude-sonnet-5",
-            "content": [{"type": "text", "text": "hello"}],
-            "stop_reason": "end_turn",
-            "stop_sequence": None,
-            "usage": {"input_tokens": 2, "output_tokens": 3},
-        }
-    ).encode()
     return RouteFixture(
         kwargs={
             "model": "anthropic/claude-sonnet-5",
             "messages": [{"role": "user", "content": "hello"}],
             **({"optional_params": {"max_tokens": 16}} if engine == "rust" else {"max_tokens": 16}),
         },
-        provider_responses=(
-            RecordedHttpResponse.from_bytes(
-                200, (HttpHeader(name="content-type", value="application/json"),), response
-            ),
-        ),
+        provider_responses=(json_response(anthropic_response_body()),),
     )
 
 
 def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture:
-    response: Final = json.dumps(
-        {
-            "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}},
-            "stopReason": "end_turn",
-            "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5},
-            "metrics": {"latencyMs": 1},
-        }
-    ).encode()
+    response: Final[dict[str, object]] = {
+        "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}},
+        "stopReason": "end_turn",
+        "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5},
+        "metrics": {"latencyMs": 1},
+    }
     credentials: Final = {
         "aws_access_key_id": "test-access",
         "aws_secret_access_key": "test-secret",
@@ -94,11 +97,60 @@ def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture:
                 else {**credentials, "max_tokens": 16}
             ),
         },
+        provider_responses=(json_response(response),),
+    )
+
+
+def _anthropic_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _anthropic_fixture(engine, _base_url)
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(sse_response(anthropic_stream_events()),),
+        consume_stream=True,
+    )
+
+
+def _bedrock_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _bedrock_fixture(engine, _base_url)
+    events: Final[tuple[dict[str, object], ...]] = (
+        {"messageStart": {"role": "assistant"}},
+        {"contentBlockStart": {"contentBlockIndex": 0, "start": {}}},
+        {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "hello"}}},
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {"messageStop": {"stopReason": "end_turn"}},
+        {"metadata": {"usage": {"inputTokens": 2, "outputTokens": 1, "totalTokens": 3}}},
+    )
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(aws_event_stream_response(events),),
+        consume_stream=True,
+    )
+
+
+def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _anthropic_fixture(engine, _base_url)
+    return fixture.derive(
         provider_responses=(
-            RecordedHttpResponse.from_bytes(
-                200, (HttpHeader(name="content-type", value="application/json"),), response
+            json_response(
+                {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}},
+                status=400,
             ),
         ),
+        expected_failure=True,
+    )
+
+
+def _stream_error_fixture(engine: Engine, base_url: str) -> RouteFixture:
+    fixture: Final = _anthropic_fixture(engine, base_url)
+    events: Final = (
+        anthropic_stream_events()[0],
+        ("error", {"type": "error", "error": {"type": "overloaded_error", "message": "overloaded"}}),
+    )
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(sse_response(events),),
+        expected_failure=True,
+        consume_stream=True,
     )
 
 
@@ -132,18 +184,64 @@ TRACE_SUITE: Final = TraceSuite(
     route=SPEC,
     scenarios=(
         TraceScenario(
-            name="anthropic",
+            name="sync-anthropic",
             fixture=_anthropic_fixture,
-            mappings=COMMON_MAPPINGS,
-            sync_mappings=SYNC_MAPPINGS,
-            async_mappings=ASYNC_MAPPINGS,
+            mappings=SYNC_MAPPINGS,
+            asynchronous=False,
         ),
         TraceScenario(
-            name="bedrock",
+            name="async-anthropic",
+            fixture=_anthropic_fixture,
+            mappings=ASYNC_MAPPINGS,
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-anthropic-stream",
+            fixture=_anthropic_stream_fixture,
+            mappings=(*SYNC_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=False,
+        ),
+        TraceScenario(
+            name="async-anthropic-stream",
+            fixture=_anthropic_stream_fixture,
+            mappings=(*ASYNC_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-anthropic-provider-error",
+            fixture=_provider_error_fixture,
+            mappings=(*ASYNC_MAPPINGS, *FAILURE_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-anthropic-stream-error",
+            fixture=_stream_error_fixture,
+            mappings=(*ASYNC_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-bedrock",
             fixture=_bedrock_fixture,
-            mappings=BEDROCK_COMMON_MAPPINGS,
-            sync_mappings=BEDROCK_SYNC_MAPPINGS,
-            async_mappings=BEDROCK_ASYNC_MAPPINGS,
+            mappings=BEDROCK_SYNC_MAPPINGS,
+            asynchronous=False,
+        ),
+        TraceScenario(
+            name="async-bedrock",
+            fixture=_bedrock_fixture,
+            mappings=BEDROCK_ASYNC_MAPPINGS,
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-bedrock-event-stream",
+            fixture=_bedrock_stream_fixture,
+            mappings=(*BEDROCK_SYNC_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=False,
+        ),
+        TraceScenario(
+            name="async-bedrock-event-stream",
+            fixture=_bedrock_stream_fixture,
+            mappings=(*BEDROCK_ASYNC_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=True,
         ),
     ),
 )

tests/rust-python-harness/strategies/trace_parity/sdk/execution.py

@@ -1,18 +1,20 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Awaitable
+import os
+from collections.abc import AsyncIterable, Awaitable, Iterable
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Final, Protocol, cast
+from unittest.mock import patch
 
 from ....shared.parity.replay import replay_server
 from ....shared.reporting.models import Surface
 from ....shared.tracing.native import TraceResponsePayload, native_trace_events
 from ....shared.tracing.profiler import FunctionTraceEvent, profile_python
 from ....shared.tracing.steps import Engine, pipeline_projection
-from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario
-from ..reporting import TraceComparisonArtifact
+from ..models import RouteFixture, RouteSpec, TraceEngine, TraceExecutionFailure, TraceScenario
+from ..reporting import TraceArtifact
 
 
 class SdkCall(Protocol):
@@ -25,10 +27,20 @@ class _CollectedTrace:
     error: str | None = None
 
 
-def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object:
+def _invoke(
+    function: SdkCall,
+    kwargs: dict[str, object],
+    *,
+    asynchronous: bool,
+    consume_stream: bool = False,
+) -> object:
     async def invoke_async() -> object:
         try:
-            return await cast(Awaitable[object], function(**kwargs))
+            response: Final = await cast(Awaitable[object], function(**kwargs))
+            if consume_stream and isinstance(response, AsyncIterable):
+                stream = cast(AsyncIterable[object], response)
+                return tuple([item async for item in stream])
+            return response
         finally:
             await asyncio.sleep(0)
             from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@@ -38,7 +50,10 @@ def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool)
 
     if asynchronous:
         return asyncio.run(invoke_async())
-    return function(**kwargs)
+    response: Final = function(**kwargs)
+    if consume_stream and isinstance(response, Iterable):
+        return tuple(cast(Iterable[object], response))
+    return response
 
 
 def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCall | TraceExecutionFailure:
@@ -47,6 +62,8 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa
     from litellm.rust_bridge import get_native_bridge
 
     if engine == "rust":
+        if spec.rust_entrypoints is None:
+            return TraceExecutionFailure("rust", f"{spec.route} has no native Rust trace entrypoint")
         bridge: Final = cast(object | None, get_native_bridge())
         if bridge is None:
             return TraceExecutionFailure("rust", "native Rust bridge is required for trace parity")
@@ -62,14 +79,6 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa
     return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)]))
 
 
-def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None:
-    try:
-        _invoke(function, kwargs, asynchronous=asynchronous)
-    except Exception as error:
-        return f"{type(error).__name__}: {error}"
-    return None
-
-
 def _collect(
     function: SdkCall,
     fixture: RouteFixture,
@@ -83,14 +92,23 @@ def _collect(
         return _CollectedTrace(native_trace_events(payload), payload.error)
     import litellm
 
-    with profile_python(Path(litellm.__file__).parent, threads=True) as profiler:
-        error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous)
+    previous_suppress_debug_info: Final = litellm.suppress_debug_info
+    try:
+        if fixture.expected_failure:
+            litellm.suppress_debug_info = True
+        with profile_python(Path(litellm.__file__).parent, threads=True) as profiler:
+            error: str | None
+            try:
+                _invoke(function, kwargs, asynchronous=asynchronous, consume_stream=fixture.consume_stream)
+                error = None
+            except Exception as caught:
+                error = f"{type(caught).__name__}: {caught}"
+    finally:
+        litellm.suppress_debug_info = previous_suppress_debug_info
     return _CollectedTrace(tuple(profiler.events), error)
 
 
-def collect_trace(
-    spec: RouteSpec, engine: Engine, *, asynchronous: bool
-) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure:
+def collect_trace(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure:
     function: Final = _entrypoint(spec, engine, asynchronous=asynchronous)
     if isinstance(function, TraceExecutionFailure):
         return function
@@ -101,15 +119,18 @@ def collect_trace(
                 provider.enqueue_response(response)
             fixture: Final = RouteFixture(
                 kwargs={
-                    **base_fixture.kwargs,
                     "api_key": "test-key",
+                    **base_fixture.kwargs,
                     "api_base": provider.url,
                     **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}),
                 },
                 provider_responses=base_fixture.provider_responses,
                 expected_failure=base_fixture.expected_failure,
+                consume_stream=base_fixture.consume_stream,
+                environment=base_fixture.environment,
             )
-            collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous)
+            with patch.dict(os.environ, fixture.environment):
+                collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous)
             provider.take_requests(len(fixture.provider_responses))
     except Exception as error:
         return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}")
@@ -129,48 +150,56 @@ def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFail
 
 
 def execute_trace(
-    route: RouteSpec, scenario: TraceScenario, mode: TraceMode, surface: Surface
-) -> TraceComparisonArtifact:
-    asynchronous: Final = mode == "async"
-    mappings: Final = scenario.mappings_for(mode)
+    route: RouteSpec,
+    scenario: TraceScenario,
+    surface: Surface,
+    engine: TraceEngine = "both",
+) -> TraceArtifact:
+    effective_engine: Final[TraceEngine] = "python" if engine == "both" and route.rust_entrypoints is None else engine
     scenario_route: Final = RouteSpec(
         route=route.route,
         python_entrypoints=route.python_entrypoints,
         rust_entrypoints=route.rust_entrypoints,
         fixture=scenario.fixture,
     )
-    python_trace: Final = collect_trace(scenario_route, "python", asynchronous=asynchronous)
-    rust_trace: Final = collect_trace(scenario_route, "rust", asynchronous=asynchronous)
+    python_trace: Final = (
+        collect_trace(
+            scenario_route,
+            "python",
+            asynchronous=scenario.asynchronous,
+        )
+        if effective_engine != "rust"
+        else ()
+    )
+    rust_trace: Final = (
+        collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous)
+        if effective_engine != "python"
+        else ()
+    )
     python_error: Final = _failure_message(python_trace)
     rust_error: Final = _failure_message(rust_trace)
     python_events: Final = python_trace if isinstance(python_trace, tuple) else ()
     rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else ()
     try:
-        python: Final = pipeline_projection("python", python_events, mappings)
-        rust: Final = pipeline_projection("rust", rust_events, mappings)
+        python: Final = pipeline_projection("python", python_events)
+        rust: Final = pipeline_projection("rust", rust_events)
     except ValueError as error:
-        return TraceComparisonArtifact.from_traces(
+        return TraceArtifact.from_traces(
+            engine=effective_engine,
             surface=surface,
             sdk_function=route.route,
             scenario=scenario.name,
-            mode=mode,
-            mappings=mappings,
-            contract=scenario.contract,
             python=(),
             rust=(),
-            python_unmatched=0,
             python_error=f"harness: {error}",
         )
-    return TraceComparisonArtifact.from_traces(
+    return TraceArtifact.from_traces(
+        engine=effective_engine,
         surface=surface,
         sdk_function=route.route,
         scenario=scenario.name,
-        mode=mode,
-        mappings=mappings,
-        contract=scenario.contract,
         python=python.steps,
         rust=rust.steps,
-        python_unmatched=python.unmatched,
         python_error=python_error,
         rust_error=rust_error,
     )

tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py

@@ -1,14 +1,26 @@
 from __future__ import annotations
 
-import json
 from typing import Final
 
-from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse
 from .....shared.tracing.steps import Engine, mapping
+from ...fixtures import (
+    anthropic_response_body,
+    anthropic_stream_events,
+    aws_event_stream_response,
+    json_response,
+    sse_response,
+)
 from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite
 
 COMMON_MAPPINGS: Final = (
     mapping(rust_span="messages", python_frame=r"anthropic_interface/messages/__init__\.py:\d+ a?create$"),
+    mapping(span="python_sanitize_empty_content", python_frame=r"strip_empty_content_blocks_from_anthropic_messages$"),
+    mapping(span="python_sanitize_tool_ids", python_frame=r"sanitize_tool_use_ids_in_anthropic_messages$"),
+    mapping(
+        span="python_flatten_web_search", python_frame=r"flatten_unencrypted_web_search_results_in_anthropic_messages$"
+    ),
+    mapping(span="python_cache_control", python_frame=r"AnthropicCacheControlHook\.maybe_inject_cache_control$"),
+    mapping(span="python_pre_request_hooks", python_frame=r"_execute_pre_request_hooks$"),
     mapping(
         span="python_messages_provider_config",
         python_frame=r"ProviderConfigManager\.get_provider_anthropic_messages_config$",
@@ -30,10 +42,42 @@ COMMON_MAPPINGS: Final = (
     ),
     mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
     mapping(rust_span="transform_response", python_frame=r"(?<!async_)transform_anthropic_messages_response$"),
+    mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"),
+    mapping(span="python_logging_post_call", python_frame=r"Logging\.post_call$"),
+)
+
+SUCCESS_MAPPINGS: Final = (mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$"),)
+FAILURE_MAPPINGS: Final = (
+    mapping(span="python_failure_callback", python_frame=r"Logging\.failure_handler$"),
+    mapping(span="python_async_failure_callback", python_frame=r"Logging\.async_failure_handler$"),
+    mapping(span="python_exception_mapping", python_frame=r"(?<!_)exception_type$"),
+)
+STREAM_MAPPINGS: Final = (
+    mapping(span="python_stream_wrapper", python_frame=r"AnthropicMessagesStreamingResponse\.__init__$"),
+    mapping(span="python_stream_next", python_frame=r"AnthropicMessagesStreamingResponse\.__anext__$"),
+    mapping(
+        span="python_stream_iterator",
+        python_frame=r"BaseAnthropicMessagesStreamingIterator\.get_async_streaming_response_iterator$",
+    ),
+    mapping(span="python_stream_chunks", python_frame=r"PassThroughStreamingHandler\.chunk_processor$"),
+    mapping(
+        span="python_stream_logging",
+        python_frame=r"PassThroughStreamingHandler\._route_streaming_logging_to_handler$",
+    ),
 )
 
 ANTHROPIC_MAPPINGS: Final = (
     *COMMON_MAPPINGS,
+    *SUCCESS_MAPPINGS,
+    mapping(
+        rust_span="transform_request",
+        python_frame=r"(?<!Azure)AnthropicMessagesConfig\.transform_anthropic_messages_request$",
+    ),
+)
+
+ANTHROPIC_FAILURE_MAPPINGS: Final = (
+    *COMMON_MAPPINGS,
+    *FAILURE_MAPPINGS,
     mapping(
         rust_span="transform_request",
         python_frame=r"(?<!Azure)AnthropicMessagesConfig\.transform_anthropic_messages_request$",
@@ -42,6 +86,7 @@ ANTHROPIC_MAPPINGS: Final = (
 
 AZURE_MAPPINGS: Final = (
     *COMMON_MAPPINGS,
+    *SUCCESS_MAPPINGS,
     mapping(
         rust_span="transform_request",
         python_frame=r"AzureAnthropicMessagesConfig\.transform_anthropic_messages_request$",
@@ -52,31 +97,54 @@ AZURE_MAPPINGS: Final = (
     ),
 )
 
+BEDROCK_MAPPINGS: Final = (
+    *COMMON_MAPPINGS,
+    *SUCCESS_MAPPINGS,
+    mapping(
+        rust_span="transform_request",
+        python_frame=r"AmazonAnthropicClaudeMessagesConfig\.transform_anthropic_messages_request$",
+    ),
+    mapping(
+        span="python_anthropic_transform_request",
+        python_frame=r"(?<!Azure)AnthropicMessagesConfig\.transform_anthropic_messages_request$",
+    ),
+    mapping(
+        span="python_bedrock_provider_config",
+        python_frame=r"BedrockModelInfo\.get_bedrock_provider_config_for_messages_api$",
+    ),
+    mapping(span="python_aws_signing", python_frame=r"sign_request_off_loop_if_aws$"),
+    mapping(
+        span="python_aws_sign_request",
+        python_frame=r"AmazonAnthropicClaudeMessagesConfig\.sign_request$|BaseAWSLLM\._sign_request$",
+    ),
+)
+
+RETRY_MAPPINGS: Final = (
+    *BEDROCK_MAPPINGS,
+    mapping(
+        span="python_retry_request_transform",
+        python_frame=r"transform_anthropic_messages_request_on_http_error$",
+    ),
+    mapping(
+        span="python_strip_invalid_thinking",
+        python_frame=r"strip_thinking_blocks_from_anthropic_messages_request_dict$",
+    ),
+)
+
+MOCK_MAPPINGS: Final = (
+    *COMMON_MAPPINGS,
+    mapping(span="python_mock_response", python_frame=r"messages/utils\.py:\d+ mock_response$"),
+)
+
 
 def _fixture(engine: Engine, provider: str) -> RouteFixture:
     conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16}
-    response: Final = json.dumps(
-        {
-            "id": "msg_trace",
-            "type": "message",
-            "role": "assistant",
-            "model": "claude-sonnet-5",
-            "content": [{"type": "text", "text": "hello"}],
-            "stop_reason": "end_turn",
-            "stop_sequence": None,
-            "usage": {"input_tokens": 2, "output_tokens": 3},
-        }
-    ).encode()
     return RouteFixture(
         kwargs={
             "model": f"{provider}/claude-sonnet-5",
             **({"body": {**conversation, "model": "claude-sonnet-5"}} if engine == "rust" else conversation),
         },
-        provider_responses=(
-            RecordedHttpResponse.from_bytes(
-                200, (HttpHeader(name="content-type", value="application/json"),), response
-            ),
-        ),
+        provider_responses=(json_response(anthropic_response_body()),),
     )
 
 
@@ -88,11 +156,170 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
     return _fixture(engine, "azure_ai")
 
 
+def _bedrock_kwargs(engine: Engine) -> dict[str, object]:
+    conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16}
+    return {
+        "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+        **(
+            {"body": {**conversation, "model": "anthropic.claude-3-sonnet-20240229-v1:0"}}
+            if engine == "rust"
+            else conversation
+        ),
+        "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
+        "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+        "aws_region_name": "us-east-1",
+    }
+
+
+def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    response_fixture: Final = _fixture(engine, "anthropic")
+    return RouteFixture(kwargs=_bedrock_kwargs(engine), provider_responses=response_fixture.provider_responses)
+
+
+def _bedrock_retry_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    success_fixture: Final = _bedrock_fixture(engine, _base_url)
+    messages: Final = [
+        {"role": "user", "content": "hello"},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "thinking", "thinking": "old reasoning", "signature": ""},
+                {"type": "text", "text": "partial answer"},
+            ],
+        },
+        {"role": "user", "content": "continue"},
+    ]
+    kwargs: Final = {
+        **_bedrock_kwargs(engine),
+        **(
+            {"body": {"messages": messages, "max_tokens": 16, "model": "anthropic.claude-3-sonnet-20240229-v1:0"}}
+            if engine == "rust"
+            else {"messages": messages}
+        ),
+    }
+    return success_fixture.derive(
+        kwargs=kwargs,
+        provider_responses=(
+            json_response({"message": "messages.1.content.0: Invalid `signature` in `thinking` block"}, status=400),
+            *success_fixture.provider_responses,
+        ),
+    )
+
+
+def _mock_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _fixture(engine, "anthropic")
+    return fixture.derive(kwargs={"mock_response": "hello from mock"}, provider_responses=())
+
+
+def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _fixture(engine, "anthropic")
+    return fixture.derive(
+        provider_responses=(
+            json_response(
+                {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}},
+                status=400,
+            ),
+        ),
+        expected_failure=True,
+    )
+
+
+def _sync_unsupported_fixture(engine: Engine, base_url: str) -> RouteFixture:
+    if engine == "rust":
+        return _anthropic_fixture(engine, base_url)
+    fixture: Final = _fixture(engine, "anthropic")
+    return fixture.derive(provider_responses=(), expected_failure=True)
+
+
+def _stream_fixture_for(engine: Engine, provider: str) -> RouteFixture:
+    fixture: Final = _fixture(engine, provider)
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(sse_response(anthropic_stream_events()),),
+        consume_stream=True,
+    )
+
+
+def _stream_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    return _stream_fixture_for(engine, "anthropic")
+
+
+def _azure_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    return _stream_fixture_for(engine, "azure_ai")
+
+
+def _bedrock_stream_fixture(engine: Engine, base_url: str) -> RouteFixture:
+    fixture: Final = _bedrock_fixture(engine, base_url)
+    events: Final = tuple(payload for _, payload in anthropic_stream_events())
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(aws_event_stream_response(events),),
+        consume_stream=True,
+    )
+
+
+def _bedrock_stream_error_fixture(engine: Engine, base_url: str) -> RouteFixture:
+    fixture: Final = _bedrock_fixture(engine, base_url)
+    start: Final = anthropic_stream_events(model="anthropic.claude-3-sonnet-20240229-v1:0")[0][1]
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(aws_event_stream_response((start, {"type": "message_stop"}), corrupt_last_frame=True),),
+        expected_failure=True,
+        consume_stream=True,
+    )
+
+
 SPEC: Final = RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _anthropic_fixture)
 TRACE_SUITE: Final = TraceSuite(
     route=SPEC,
     scenarios=(
-        TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)),
-        TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)),
+        TraceScenario(
+            name="async-anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, asynchronous=True
+        ),
+        TraceScenario(name="async-azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True),
+        TraceScenario(name="async-bedrock", fixture=_bedrock_fixture, mappings=BEDROCK_MAPPINGS, asynchronous=True),
+        TraceScenario(
+            name="async-bedrock-invalid-thinking-retry",
+            fixture=_bedrock_retry_fixture,
+            mappings=RETRY_MAPPINGS,
+            asynchronous=True,
+        ),
+        TraceScenario(name="async-mock-response", fixture=_mock_fixture, mappings=MOCK_MAPPINGS, asynchronous=True),
+        TraceScenario(
+            name="async-anthropic-provider-error",
+            fixture=_provider_error_fixture,
+            mappings=ANTHROPIC_FAILURE_MAPPINGS,
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-anthropic-stream",
+            fixture=_stream_fixture,
+            mappings=(*ANTHROPIC_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-azure-ai-stream",
+            fixture=_azure_stream_fixture,
+            mappings=(*AZURE_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-bedrock-event-stream",
+            fixture=_bedrock_stream_fixture,
+            mappings=(*BEDROCK_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-bedrock-event-stream-error",
+            fixture=_bedrock_stream_error_fixture,
+            mappings=(*BEDROCK_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-unsupported",
+            fixture=_sync_unsupported_fixture,
+            mappings=ANTHROPIC_MAPPINGS,
+            asynchronous=False,
+        ),
     ),
 )

tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import json
-from typing import Final, cast
+from typing import Final
 
 from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse
 from .....shared.tracing.steps import Engine, mapping
@@ -53,6 +53,23 @@ ASYNC_MAPPINGS: Final = (
     ),
 )
 
+PUBLIC_RUST_DISPATCH_MAPPINGS: Final = (
+    mapping(span="public_sdk_entrypoint", python_frame=r"ocr/main\.py:\d+ a?ocr$"),
+    mapping(span="public_request", python_frame=r"ocr/main\.py:\d+ _public_request$"),
+    mapping(span="bind_request", python_frame=r"ocr/main\.py:\d+ _bind_request$"),
+    mapping(span="rust_ocr_enabled", python_frame=r"rust_bridge/configuration\.py:\d+ rust_ocr_enabled$"),
+    mapping(span="select_native_ocr", python_frame=r"rust_bridge/ocr_lifecycle\.py:\d+ select$"),
+    mapping(span="load_native_bridge", python_frame=r"rust_bridge/bindings\.py:\d+ NativeBinding\.load$"),
+    mapping(span="native_call_setup", python_frame=r"rust_bridge/lifecycle\.py:\d+ setup$"),
+    mapping(span="native_response", python_frame=r"rust_bridge/ocr\.py:\d+ _response$"),
+    mapping(span="native_call_finalize", python_frame=r"rust_bridge/lifecycle\.py:\d+ finalize$"),
+    mapping(
+        span="native_success_bookkeeping",
+        python_frame=r"rust_bridge/lifecycle\.py:\d+ success_bookkeeping$",
+    ),
+    *(mapping(rust_span=item.rust) for item in SYNC_MAPPINGS if item.rust is not None),
+)
+
 CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING)
 CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING)
 CALLBACK_FAILURE_SYNC_MAPPINGS: Final = (
@@ -164,23 +181,6 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
     )
 
 
-def _vertex_fixture(engine: Engine, _base_url: str) -> RouteFixture:
-    fixture: Final = _fixture(
-        engine,
-        "vertex_ai/mistral-ocr-maas",
-        {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="},
-    )
-    vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"}
-    optional_params: Final = cast(dict[str, object], fixture.kwargs.get("optional_params", {}))
-    return RouteFixture(
-        kwargs={
-            **fixture.kwargs,
-            **({"optional_params": {**optional_params, **vertex}} if engine == "rust" else vertex),
-        },
-        provider_responses=fixture.provider_responses,
-    )
-
-
 def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture:
     vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"}
     return RouteFixture(
@@ -204,6 +204,62 @@ def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture:
     )
 
 
+def _vertex_deepseek_credentials_fixture(engine: Engine, base_url: str) -> RouteFixture:
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.hazmat.primitives.asymmetric import rsa
+
+    fixture: Final = _vertex_deepseek_fixture(engine, base_url)
+    private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    credentials: Final = json.dumps(
+        {
+            "type": "service_account",
+            "project_id": "trace-project",
+            "private_key_id": "trace-key",
+            "private_key": private_key.private_bytes(
+                serialization.Encoding.PEM,
+                serialization.PrivateFormat.PKCS8,
+                serialization.NoEncryption(),
+            ).decode(),
+            "client_email": "trace@trace-project.iam.gserviceaccount.com",
+            "token_uri": f"{base_url}/token",
+        }
+    )
+    return RouteFixture(
+        kwargs={**fixture.kwargs, "api_key": None},
+        environment=(("VERTEXAI_CREDENTIALS", credentials), ("VERTEX_AI_API_KEY", "")),
+        provider_responses=(
+            RecordedHttpResponse.from_bytes(
+                200,
+                (HttpHeader(name="content-type", value="application/json"),),
+                b'{"access_token":"trace-token","token_type":"Bearer","expires_in":3600}',
+            ),
+            *fixture.provider_responses,
+        ),
+    )
+
+
+def _cohere_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    return RouteFixture(
+        kwargs={
+            "model": "cohere/parse-v5.0",
+            "document": {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="},
+            **({"optional_params": {"output_format": "blocks"}} if engine == "rust" else {"output_format": "blocks"}),
+        },
+        provider_responses=(
+            RecordedHttpResponse.from_bytes(
+                200,
+                (HttpHeader(name="content-type", value="application/json"),),
+                json.dumps(
+                    {
+                        "pages": [{"index": 0, "blocks": [{"type": "text", "text": {"content": "hello"}}]}],
+                        "meta": {"billed_units": {"pages": 1}},
+                    }
+                ).encode(),
+            ),
+        ),
+    )
+
+
 def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> RouteFixture:
     completed: Final = json.dumps(
         {
@@ -249,30 +305,6 @@ def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> Route
     )
 
 
-VERTEX_COMMON_MAPPINGS: Final = (
-    *COMMON_MAPPINGS[:7],
-    mapping(
-        rust_span="transform_ocr_request",
-        python_frame=(
-            r"VertexAIOCRConfig\.(?:async_)?transform_ocr_request$"
-            r"|MistralOCRConfig\.transform_ocr_request$"
-        ),
-    ),
-    COMMON_MAPPINGS[-1],
-)
-VERTEX_SYNC_MAPPINGS: Final = (
-    *VERTEX_COMMON_MAPPINGS,
-    mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
-    mapping(span="python_transform_ocr_response_wrapper", python_frame=r"BaseLLMHTTPHandler\._transform_ocr_response$"),
-    mapping(rust_span="transform_ocr_response", python_frame=r"MistralOCRConfig\.transform_ocr_response$"),
-)
-VERTEX_ASYNC_MAPPINGS: Final = (
-    *VERTEX_COMMON_MAPPINGS,
-    mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
-    mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
-    mapping(rust_span="transform_ocr_response", python_frame=r"MistralOCRConfig\.transform_ocr_response$"),
-)
-
 DEEPSEEK_COMMON_MAPPINGS: Final = (
     mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"),
     mapping(rust_span="prepare_ocr_call", python_frame=r"ocr/main\.py:\d+ _prepare_ocr_request$"),
@@ -352,59 +384,126 @@ DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS: Final = (
     mapping(span="python_poll_http_request", python_frame=r"AsyncHTTPHandler\.get$"),
 )
 
+COHERE_COMMON_MAPPINGS: Final = (
+    *COMMON_MAPPINGS[:7],
+    mapping(
+        rust_span="transform_ocr_request",
+        python_frame=r"CohereParseConfig\.(?:async_)?transform_ocr_request$",
+    ),
+    COMMON_MAPPINGS[-1],
+    mapping(rust_span="transform_ocr_response", python_frame=r"CohereParseConfig\.transform_ocr_response$"),
+)
+COHERE_ASYNC_MAPPINGS: Final = (
+    *COHERE_COMMON_MAPPINGS,
+    mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
+    mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
+)
 
 SPEC: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _mistral_fixture)
 TRACE_SUITE: Final = TraceSuite(
     route=SPEC,
     scenarios=(
         TraceScenario(
-            name="mistral",
+            name="sync-mistral",
             fixture=_mistral_fixture,
-            mappings=COMMON_MAPPINGS,
-            sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
-            async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=False,
         ),
         TraceScenario(
-            name="mistral-callback-success",
+            name="async-mistral",
+            fixture=_mistral_fixture,
+            mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-mistral-callback-success",
             fixture=_mistral_callback_success_fixture,
-            mappings=COMMON_MAPPINGS,
-            sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS,
-            async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS,
+            mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS,
+            asynchronous=False,
         ),
         TraceScenario(
-            name="mistral-callback-failure",
+            name="async-mistral-callback-success",
+            fixture=_mistral_callback_success_fixture,
+            mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS,
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-mistral-callback-failure",
             fixture=_mistral_callback_failure_fixture,
-            mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING),
-            sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS,
-            async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS,
+            mappings=CALLBACK_FAILURE_SYNC_MAPPINGS,
+            asynchronous=False,
         ),
         TraceScenario(
-            name="azure-ai",
+            name="async-mistral-callback-failure",
+            fixture=_mistral_callback_failure_fixture,
+            mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS,
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-azure-ai",
             fixture=_azure_fixture,
-            mappings=AZURE_COMMON_MAPPINGS,
-            sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
-            async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=False,
         ),
         TraceScenario(
-            name="azure-document-intelligence",
+            name="async-azure-ai",
+            fixture=_azure_fixture,
+            mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-azure-document-intelligence",
             fixture=_azure_document_intelligence_fixture,
-            mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS,
-            sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
-            async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=False,
         ),
         TraceScenario(
-            name="vertex-ai",
-            fixture=_vertex_fixture,
-            mappings=VERTEX_COMMON_MAPPINGS,
-            sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
-            async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            name="async-azure-document-intelligence",
+            fixture=_azure_document_intelligence_fixture,
+            mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=True,
         ),
         TraceScenario(
-            name="vertex-deepseek",
+            name="sync-vertex-deepseek",
             fixture=_vertex_deepseek_fixture,
-            mappings=DEEPSEEK_COMMON_MAPPINGS,
-            sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
-            async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=False,
+        ),
+        TraceScenario(
+            name="async-vertex-deepseek",
+            fixture=_vertex_deepseek_fixture,
+            mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-vertex-deepseek-credentials",
+            fixture=_vertex_deepseek_credentials_fixture,
+            mappings=DEEPSEEK_SYNC_MAPPINGS,
+            asynchronous=False,
+        ),
+        TraceScenario(
+            name="async-vertex-deepseek-credentials",
+            fixture=_vertex_deepseek_credentials_fixture,
+            mappings=DEEPSEEK_ASYNC_MAPPINGS,
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-cohere",
+            fixture=_cohere_fixture,
+            mappings=(*COHERE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="sync-public-rust-dispatch",
+            fixture=_mistral_fixture,
+            mappings=PUBLIC_RUST_DISPATCH_MAPPINGS,
+            asynchronous=False,
+        ),
+        TraceScenario(
+            name="async-public-rust-dispatch",
+            fixture=_mistral_fixture,
+            mappings=PUBLIC_RUST_DISPATCH_MAPPINGS,
+            asynchronous=True,
         ),
     ),
 )

tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+from typing import Final
+
+from .....shared.tracing.steps import Engine, mapping
+from ...fixtures import (
+    anthropic_response_body,
+    anthropic_stream_events,
+    json_response,
+    responses_body,
+    responses_stream_events,
+    sse_response,
+)
+from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite
+
+COMMON_MAPPINGS: Final = (
+    mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"),
+    mapping(
+        span="python_responses_provider_config",
+        python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$",
+    ),
+    mapping(rust_span="responses_provider_config"),
+    mapping(rust_span="validate_environment", python_frame=r"validate_environment$"),
+    mapping(rust_span="complete_url", python_frame=r"get_complete_url$"),
+    mapping(
+        rust_span="transform_request",
+        python_frame=r"(?<!AzureOpenAIResponsesAPIConfig\.)transform_responses_api_request$",
+    ),
+    mapping(
+        rust_span="execute_responses_provider_call",
+        python_frame=r"BaseLLMHTTPHandler\.(?:async_)?response_api_handler$",
+    ),
+    mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
+    mapping(rust_span="transform_response", python_frame=r"transform_response_api_response$"),
+    mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"),
+    mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$|Logging\.success_handler$"),
+)
+
+STREAM_MAPPINGS: Final = (
+    mapping(
+        span="python_responses_stream_iterator",
+        python_frame=r"(?:Sync)?ResponsesAPIStreamingIterator\.__init__$",
+    ),
+    mapping(
+        span="python_responses_stream_next",
+        python_frame=r"(?:Sync)?ResponsesAPIStreamingIterator\.__a?next__$",
+    ),
+    mapping(span="python_responses_stream_transform", python_frame=r"transform_streaming_response$"),
+)
+
+FAILURE_MAPPINGS: Final = (
+    mapping(span="python_exception_mapping", python_frame=r"(?<!_)exception_type$"),
+    mapping(span="python_failure_callback", python_frame=r"Logging\.failure_handler$"),
+    mapping(span="python_async_failure_callback", python_frame=r"Logging\.async_failure_handler$"),
+)
+
+AZURE_MAPPINGS: Final = (
+    *COMMON_MAPPINGS,
+    mapping(
+        span="python_azure_transform_request",
+        python_frame=r"AzureOpenAIResponsesAPIConfig\.transform_responses_api_request$",
+    ),
+)
+
+BRIDGE_MAPPINGS: Final = (
+    mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"),
+    mapping(
+        span="python_responses_chat_bridge", python_frame=r"ResponsesToCompletionBridgeHandler\.response_api_handler$"
+    ),
+    mapping(span="python_chat_completions", python_frame=r"main\.py:\d+ a?completion$"),
+    mapping(span="python_chat_transform_request", python_frame=r"AnthropicConfig\.transform_request$"),
+    mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"),
+    mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
+    mapping(span="python_chat_transform_response", python_frame=r"AnthropicConfig\.transform_response$"),
+    mapping(span="python_chat_to_responses", python_frame=r"LiteLLMResponsesTransformationHandler\..*response"),
+    mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$|Logging\.success_handler$"),
+)
+
+
+def _native_fixture(engine: Engine, provider: str) -> RouteFixture:
+    model: Final = "gpt-5"
+    return RouteFixture(
+        kwargs={
+            "model": f"{provider}/{model}",
+            "input": "hello",
+            **({"body": {"model": model, "input": "hello"}} if engine == "rust" else {}),
+        },
+        provider_responses=(json_response(responses_body(model=model)),),
+    )
+
+
+def _openai_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    return _native_fixture(engine, "openai")
+
+
+def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _native_fixture(engine, "azure")
+    return fixture.derive(kwargs={"api_version": "2025-04-01-preview"})
+
+
+def _openai_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _openai_fixture(engine, _base_url)
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(sse_response(responses_stream_events()),),
+        consume_stream=True,
+    )
+
+
+def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _openai_fixture(engine, _base_url)
+    return fixture.derive(
+        provider_responses=(
+            json_response({"error": {"message": "bad request", "type": "invalid_request_error"}}, status=400),
+        ),
+        expected_failure=True,
+    )
+
+
+def _stream_failed_fixture(engine: Engine, base_url: str) -> RouteFixture:
+    fixture: Final = _openai_fixture(engine, base_url)
+    failed_response: Final[dict[str, object]] = {
+        **responses_body(),
+        "status": "failed",
+        "output": [],
+        "error": {"message": "stream failed", "type": "server_error", "code": "server_error"},
+    }
+    events: Final = (
+        (
+            "response.created",
+            {"type": "response.created", "response": {**failed_response, "status": "in_progress", "error": None}},
+        ),
+        ("response.failed", {"type": "response.failed", "response": failed_response}),
+    )
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(sse_response(events),),
+        expected_failure=True,
+        consume_stream=True,
+    )
+
+
+def _anthropic_bridge_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    return RouteFixture(
+        kwargs={
+            "model": "anthropic/claude-sonnet-5",
+            "input": "hello",
+            "max_output_tokens": 16,
+            **({"body": {"model": "claude-sonnet-5", "input": "hello"}} if engine == "rust" else {}),
+        },
+        provider_responses=(json_response(anthropic_response_body()),),
+    )
+
+
+def _anthropic_bridge_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    fixture: Final = _anthropic_bridge_fixture(engine, _base_url)
+    return fixture.derive(
+        kwargs={"stream": True},
+        provider_responses=(sse_response(anthropic_stream_events()),),
+        consume_stream=True,
+    )
+
+
+SPEC: Final = RouteSpec("responses", ("responses", "aresponses"), None, _openai_fixture)
+TRACE_SUITE: Final = TraceSuite(
+    route=SPEC,
+    scenarios=(
+        TraceScenario(name="sync-openai", fixture=_openai_fixture, mappings=COMMON_MAPPINGS, asynchronous=False),
+        TraceScenario(name="async-openai", fixture=_openai_fixture, mappings=COMMON_MAPPINGS, asynchronous=True),
+        TraceScenario(
+            name="sync-openai-stream",
+            fixture=_openai_stream_fixture,
+            mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=False,
+        ),
+        TraceScenario(
+            name="async-openai-stream",
+            fixture=_openai_stream_fixture,
+            mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-openai-provider-error",
+            fixture=_provider_error_fixture,
+            mappings=(*COMMON_MAPPINGS, *FAILURE_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-openai-stream-failed",
+            fixture=_stream_failed_fixture,
+            mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS),
+            asynchronous=True,
+        ),
+        TraceScenario(name="async-azure", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True),
+        TraceScenario(
+            name="async-anthropic-chat-bridge",
+            fixture=_anthropic_bridge_fixture,
+            mappings=BRIDGE_MAPPINGS,
+            asynchronous=True,
+        ),
+        TraceScenario(
+            name="async-anthropic-chat-bridge-stream",
+            fixture=_anthropic_bridge_stream_fixture,
+            mappings=(
+                *BRIDGE_MAPPINGS,
+                mapping(span="python_chat_stream_wrapper", python_frame=r"CustomStreamWrapper\.__init__$"),
+                mapping(span="python_chat_stream_next", python_frame=r"CustomStreamWrapper\.__anext__$"),
+                mapping(
+                    span="python_responses_bridge_stream_iterator",
+                    python_frame=r"LiteLLMCompletionStreamingIterator\.__init__$|LiteLLMCompletionStreamingIterator\.__anext__$",
+                ),
+            ),
+            asynchronous=True,
+        ),
+    ),
+)

tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Final, cast
+
+from ..models import TraceSuite
+
+
+def _suite(module: str) -> TraceSuite:
+    loaded: Final = import_module(module)
+    candidate: Final = cast(object, getattr(loaded, "TRACE_SUITE"))
+    assert isinstance(candidate, TraceSuite)
+    return candidate
+
+
+def test_core_sdk_scenario_matrix_keeps_distinct_migration_paths() -> None:
+    chat: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case")
+    messages: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.messages.case")
+    ocr: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
+    responses: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case")
+
+    assert {(scenario.name, scenario.asynchronous) for scenario in chat.scenarios} >= {
+        ("sync-anthropic", False),
+        ("async-anthropic", True),
+        ("sync-anthropic-stream", False),
+        ("async-anthropic-stream", True),
+        ("async-anthropic-provider-error", True),
+        ("async-anthropic-stream-error", True),
+        ("sync-bedrock", False),
+        ("async-bedrock", True),
+        ("sync-bedrock-event-stream", False),
+        ("async-bedrock-event-stream", True),
+    }
+    assert {(scenario.name, scenario.asynchronous) for scenario in messages.scenarios} >= {
+        ("async-anthropic-stream", True),
+        ("async-azure-ai-stream", True),
+        ("async-bedrock-event-stream", True),
+        ("async-bedrock-event-stream-error", True),
+        ("async-bedrock-invalid-thinking-retry", True),
+        ("sync-unsupported", False),
+    }
+    assert {(scenario.name, scenario.asynchronous) for scenario in ocr.scenarios} >= {
+        ("async-cohere", True),
+        ("sync-public-rust-dispatch", False),
+        ("async-public-rust-dispatch", True),
+    }
+    assert {(scenario.name, scenario.asynchronous) for scenario in responses.scenarios} >= {
+        ("sync-openai", False),
+        ("async-openai", True),
+        ("sync-openai-stream", False),
+        ("async-openai-stream", True),
+        ("async-openai-provider-error", True),
+        ("async-openai-stream-failed", True),
+        ("async-azure", True),
+        ("async-anthropic-chat-bridge", True),
+        ("async-anthropic-chat-bridge-stream", True),
+    }
+
+
+def test_core_gateway_matrix_keeps_downstream_streams_separate() -> None:
+    modules: Final = (
+        "tests.rust-python-harness.strategies.trace_parity.gateway.chat_completions.case",
+        "tests.rust-python-harness.strategies.trace_parity.gateway.messages.case",
+        "tests.rust-python-harness.strategies.trace_parity.gateway.responses.case",
+    )
+
+    for module in modules:
+        suite = _suite(module)
+        assert any("downstream-stream" in scenario.name for scenario in suite.scenarios)

tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py

@@ -92,11 +92,16 @@ TRACE_SUITE: Final = TraceSuite(
     route=SPEC,
     scenarios=(
         TraceScenario(
-            name="bedrock",
+            name="sync-bedrock",
             fixture=_fixture,
-            mappings=MAPPINGS,
-            sync_mappings=SYNC_MAPPINGS,
-            async_mappings=ASYNC_MAPPINGS,
+            mappings=SYNC_MAPPINGS,
+            asynchronous=False,
+        ),
+        TraceScenario(
+            name="async-bedrock",
+            fixture=_fixture,
+            mappings=ASYNC_MAPPINGS,
+            asynchronous=True,
         ),
     ),
 )

tests/rust-python-harness/strategies/trace_parity/test_reporting.py

@@ -1,58 +1,46 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
 from typing import Final, Literal
 
 import pytest
 
 from ...shared.reporting.models import CaseResult, Coverage, HarnessCase, ResultArtifact, RunStatus
 from ...shared.reporting.strategy import ModuleCaseSpec, NotImplementedCaseSpec
-from ...shared.tracing.steps import PipelineStep, TraceContract, TraceMapping, mapping
+from ...shared.tracing.steps import PipelineStep
 from . import reporting
-from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact, render_trace_results
-
-MAPPINGS: Final = (
-    mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"),
-    mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$"),
-)
+from .reporting import TRACE_ARTIFACT, TraceArtifact, render_trace_results
 
 
-def _result(comparison: TraceComparisonArtifact) -> CaseResult:
+def _result(trace: TraceArtifact) -> CaseResult:
     case: Final = HarnessCase(
         strategy_id="trace_parity",
         strategy_label="Trace parity",
-        sdk_function=comparison.sdk_function,
+        sdk_function=trace.sdk_function,
         spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"),
-        surface=comparison.surface,
+        surface=trace.surface,
     )
     result: Final = CaseResult(case=case)
-    nodeid: Final = f"trace:sdk:{comparison.sdk_function}:{comparison.scenario}:{comparison.mode}"
+    nodeid: Final = f"trace:{trace.surface}:{trace.sdk_function}:{trace.scenario}"
     result.collected.add(nodeid)
-    result.record(
-        nodeid,
-        RunStatus.PASSED,
-        artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),),
-    )
+    result.record(nodeid, RunStatus.PASSED, artifacts=(ResultArtifact(TRACE_ARTIFACT, trace.model_dump_json()),))
     return result
 
 
-def _comparison(
+def _trace(
     python: tuple[PipelineStep, ...],
     rust: tuple[PipelineStep, ...],
     *,
-    mappings: Sequence[TraceMapping] = MAPPINGS,
     rust_error: str | None = None,
-) -> TraceComparisonArtifact:
-    return TraceComparisonArtifact.from_traces(
+    engine: Literal["python", "rust", "both"] = "both",
+    scenario: str = "sync-default",
+) -> TraceArtifact:
+    return TraceArtifact.from_traces(
+        engine=engine,
         surface="sdk",
         sdk_function="ocr",
-        scenario="default",
-        mode="sync",
-        mappings=mappings,
-        contract=TraceContract(),
+        scenario=scenario,
         python=python,
         rust=rust,
-        python_unmatched=796,
         rust_error=rust_error,
     )
 
@@ -67,107 +55,55 @@ def _events(*items: tuple[str, int, str | None]) -> tuple[PipelineStep, ...]:
     return tuple(steps)
 
 
-def test_renderer_shows_matching_python_and_rust_paths() -> None:
-    rust: Final = _events(("ocr", 0, None), ("http_request", 1, None))
+def test_renderer_prints_python_and_rust_traces_independently() -> None:
     python: Final = _events(
         ("ocr", 0, "ocr/main.py:88 aocr"),
-        ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"),
+        ("python_prepare", 1, "prep.py:1 python_prepare"),
     )
-
-    section: Final = render_trace_results((_result(_comparison(python, rust)),))[0]
-    report: Final = "\n\n".join(section.blocks)
-
-    assert section.title == "SDK trace comparisons"
-    assert "Case: ocr" in report
-    assert "PYTHON (2 steps)\n1 aocr  (ocr/main.py:88)\n2   AsyncHTTPHandler.post  (http_handler.py:673)" in report
-    assert "RUST (2 steps)\nocr  ->  1 aocr\n  http_request  ->  2 AsyncHTTPHandler.post" in report
-    assert "Mapping (identifier -> span)" not in report
-    assert "Trace: MATCH" in report
-    assert "Same steps, order, and nesting" in report
-    assert "Unseen mappings:" not in report
-
-
-def test_renderer_reports_mappings_that_matched_nothing() -> None:
-    events: Final = _events(("ocr", 0, None))
-
-    section: Final = render_trace_results((_result(_comparison(events, events)),))[0]
-    report: Final = "\n\n".join(section.blocks)
-
-    assert "Unseen mappings: http_request" in report
-    assert "Contract: FAIL" in report
-
-
-def test_renderer_numbers_repeated_span_occurrences() -> None:
-    mappings: Final = (MAPPINGS[0], MAPPINGS[1])
-    rust: Final = _events(("ocr", 0, None), ("http_request", 1, None), ("http_request", 1, None))
-    python: Final = _events(
-        ("ocr", 0, "ocr/main.py:88 aocr"),
-        ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"),
-        ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"),
-    )
-
-    report: Final = "\n\n".join(render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0].blocks)
-
-    assert "http_request#2" in report
-
-
-def test_renderer_accepts_declared_engine_specific_steps() -> None:
-    mappings: Final = (
-        *MAPPINGS[:1],
-        mapping(span="python_prepare", python_frame=r"python_prepare$"),
-        mapping(rust_span="rust_prepare"),
-    )
-    python: Final = _events(("ocr", 0, None), ("python_prepare", 1, "prep.py:1 python_prepare"))
     rust: Final = _events(("ocr", 0, None), ("rust_prepare", 1, None))
 
-    section: Final = render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0]
+    section: Final = render_trace_results((_result(_trace(python, rust)),))[0]
     report: Final = "\n\n".join(section.blocks)
 
-    assert "2   python_prepare  (prep.py:1)  [python only]" in report
-    assert "rust_prepare  ->  [rust only]" in report
-    assert "Trace: MATCH" in report
-    assert "Contract: PASS" in report
+    assert section.title == "SDK traces"
+    assert "PYTHON (2 steps)\n1 aocr  (ocr/main.py:88)\n2   python_prepare  (prep.py:1)" in report
+    assert "RUST (2 steps)\n1 ocr\n2   rust_prepare" in report
+    assert "python only" not in report
+    assert "rust only" not in report
+    assert " -> " not in report
+    assert "Trace: MATCH" not in report
+    assert "Trace: DRIFT" not in report
+    assert "Contract:" not in report
 
 
-def test_unavailable_check_reports_mode_from_nodeid() -> None:
-    case: Final = HarnessCase(
-        strategy_id="trace_parity",
-        strategy_label="Trace parity",
-        sdk_function="ocr",
-        spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"),
-        surface="sdk",
-    )
-    result: Final = CaseResult(case=case)
-    result.collected.add("trace:sdk:ocr:default:sync")
-    result.record("trace:sdk:ocr:default:sync", RunStatus.ERROR)
+@pytest.mark.parametrize(
+    ("engine", "present", "absent"),
+    (("python", "PYTHON (1 steps)", "RUST"), ("rust", "RUST (1 steps)", "PYTHON")),
+)
+def test_renderer_prints_only_selected_engine(engine: Literal["python", "rust"], present: str, absent: str) -> None:
+    events: Final = _events(("ocr", 0, None))
 
-    section: Final = render_trace_results((result,))[0]
-    report: Final = "\n\n".join(section.blocks)
+    report: Final = "\n\n".join(render_trace_results((_result(_trace(events, events, engine=engine)),))[0].blocks)
 
-    assert "Case: ocr" in report
-    assert "Scenario: default / Mode: sync" in report
-    assert "Trace: NOT AVAILABLE\nTest outcome: error" in report
-    assert "unknown mode" not in report
+    assert present in report
+    assert absent not in report
 
 
 def test_renderer_keeps_collected_trace_when_one_engine_errors() -> None:
-    python: Final = _events(
-        ("ocr", 0, "ocr/main.py:88 aocr"),
-        ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"),
+    python: Final = _events(("ocr", 0, "ocr/main.py:88 aocr"))
+
+    report: Final = "\n\n".join(
+        render_trace_results(
+            (_result(_trace(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),)
+        )[0].blocks
     )
 
-    section: Final = render_trace_results(
-        (_result(_comparison(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),)
-    )[0]
-    report: Final = "\n\n".join(section.blocks)
-
-    assert "PYTHON (2 steps)\n1 aocr  (ocr/main.py:88)  [python only]" in report
+    assert "PYTHON (1 steps)\n1 aocr  (ocr/main.py:88)" in report
     assert "Rust error: rust: native Rust bridge must include the trace-parity feature" in report
     assert "hint: rebuild the native bridge with the trace-parity feature" in report
-    assert "Contract: FAIL" in report
 
 
-def test_renderer_groups_all_modes_under_one_case_header() -> None:
+def test_unavailable_trace_reports_scenario_from_nodeid() -> None:
     case: Final = HarnessCase(
         strategy_id="trace_parity",
         strategy_label="Trace parity",
@@ -176,76 +112,55 @@ def test_renderer_groups_all_modes_under_one_case_header() -> None:
         surface="sdk",
     )
     result: Final = CaseResult(case=case)
-    events: Final = _events(("ocr", 0, None))
-    modes: Final[tuple[Literal["sync", "async"], ...]] = ("sync", "async")
-    for mode in modes:
-        nodeid = f"trace:sdk:ocr:default:{mode}"
-        result.collected.add(nodeid)
-        comparison = TraceComparisonArtifact.from_traces(
-            surface="sdk",
-            sdk_function="ocr",
-            scenario="default",
-            mode=mode,
-            mappings=MAPPINGS,
-            contract=TraceContract(),
-            python=events,
-            rust=events,
-            python_unmatched=0,
-        )
-        result.record(
-            nodeid,
-            RunStatus.PASSED,
-            artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),),
-        )
+    result.collected.add("trace:sdk:ocr:async-error")
+    result.record("trace:sdk:ocr:async-error", RunStatus.ERROR)
 
-    section: Final = render_trace_results((result,))[0]
+    report: Final = "\n\n".join(render_trace_results((result,))[0].blocks)
+
+    assert "Scenario: async-error" in report
+    assert "Trace: NOT AVAILABLE\nTest outcome: error" in report
+
+
+def test_renderer_groups_scenarios_under_one_case_header() -> None:
+    result: Final = _result(_trace(_events(("ocr", 0, None)), (), scenario="sync-default"))
+    async_trace: Final = _trace((), _events(("ocr", 0, None)), scenario="async-default")
+    nodeid: Final = "trace:sdk:ocr:async-default"
+    result.collected.add(nodeid)
+    result.record(nodeid, RunStatus.PASSED, artifacts=(ResultArtifact(TRACE_ARTIFACT, async_trace.model_dump_json()),))
+
+    report: Final = render_trace_results((result,))[0].blocks[0]
 
-    assert len(section.blocks) == 1
-    report: Final = section.blocks[0]
     assert report.count("Case: ocr") == 1
-    assert "Scenario: default / Mode: sync" in report
-    assert "Scenario: default / Mode: async" in report
+    assert "Scenario: sync-default" in report
+    assert "Scenario: async-default" in report
 
 
 def test_renderer_colors_every_trace_line_in_a_terminal(monkeypatch: pytest.MonkeyPatch) -> None:
-    rust: Final = _events(("ocr", 0, None), ("http_request", 1, None))
-    python: Final = _events(
-        ("ocr", 0, "ocr/main.py:88 aocr"),
-        ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"),
-    )
+    events: Final = _events(("ocr", 0, "ocr/main.py:88 aocr"))
     monkeypatch.setattr(reporting.sys.stdout, "isatty", lambda: True)
     monkeypatch.delenv("NO_COLOR", raising=False)
 
-    section: Final = render_trace_results((_result(_comparison(python, rust)),))[0]
-    report: Final = "\n\n".join(section.blocks)
+    report: Final = "\n\n".join(render_trace_results((_result(_trace(events, events)),))[0].blocks)
 
-    assert "\033[36mPYTHON\033[0m (2 steps)" in report
+    assert "\033[36mPYTHON\033[0m (1 steps)" in report
     assert "\033[36m1 aocr  (ocr/main.py:88)\033[0m" in report
-    assert "\033[33mRUST\033[0m (2 steps)" in report
-    assert "\033[33mocr\033[0m  ->  \033[36m1 aocr\033[0m" in report
-    assert "\033[33mhttp_request\033[0m  ->  \033[36m2 AsyncHTTPHandler.post\033[0m" in report
+    assert "\033[33mRUST\033[0m (1 steps)" in report
+    assert "\033[33m1 ocr\033[0m" in report
 
 
-def test_renderer_groups_cases_and_unavailable_entries_by_surface() -> None:
-    events: Final = _events(("ocr", 0, None))
-    gateway_results: Final = tuple(
-        CaseResult(
-            case=HarnessCase(
-                strategy_id="trace_parity",
-                strategy_label="Trace parity",
-                sdk_function=sdk_function,
-                spec=NotImplementedCaseSpec(reason=f"No {sdk_function} case is registered."),
-                surface="gateway",
-            ),
-            status=RunStatus.NOT_IMPLEMENTED,
-        )
-        for sdk_function in ("ocr", "messages")
+def test_renderer_groups_unavailable_entries_by_surface() -> None:
+    gateway_result: Final = CaseResult(
+        case=HarnessCase(
+            strategy_id="trace_parity",
+            strategy_label="Trace parity",
+            sdk_function="messages",
+            spec=NotImplementedCaseSpec(reason="No messages case is registered."),
+            surface="gateway",
+        ),
+        status=RunStatus.NOT_IMPLEMENTED,
     )
 
-    sections: Final = render_trace_results((_result(_comparison(events, events)), *gateway_results))
+    sections: Final = render_trace_results((_result(_trace((), ())), gateway_result))
 
-    assert tuple(section.title for section in sections) == ("SDK trace comparisons", "GATEWAY trace comparisons")
-    gateway_report: Final = "\n\n".join(sections[1].blocks)
-    assert gateway_report.count("Not implemented") == 1
-    assert "- ocr: No ocr case is registered." in gateway_report
-    assert "- messages: No messages case is registered." in gateway_report
+    assert tuple(section.title for section in sections) == ("SDK traces", "GATEWAY traces")
+    assert "- messages: No messages case is registered." in "\n\n".join(sections[1].blocks)

tests/rust-python-harness/strategies/trace_parity/test_runner.py

@@ -1,12 +1,23 @@
 from __future__ import annotations
 
-from typing import Final
+import importlib
+import os
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Final, cast
+
+import pytest
+
+import litellm
 
 from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface
 from ...shared.reporting.strategy import ModuleCaseSpec
-from ...shared.tracing.steps import Engine
+from ...shared.tracing.profiler import FunctionTraceEvent
+from ...shared.tracing.steps import Engine, PipelineStep, mapping
 from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite
-from .runner import run_trace_mode, scenario_nodeids, validate_trace_suite
+from .reporting import TraceArtifact
+from .runner import run_trace_cases, run_trace_scenario, runner_selection, scenario_nodeids, validate_trace_suite
+from .sdk.execution import SdkCall, collect_trace, execute_trace
 
 
 def _fixture(_engine: Engine, _base_url: str) -> RouteFixture:
@@ -27,46 +38,243 @@ def test_scenario_filtering_and_occurrence_node_ids() -> None:
     suite: Final = TraceSuite(
         route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture),
         scenarios=(
-            TraceScenario("one", _fixture, (), modes=("sync", "async")),
-            TraceScenario("two", _fixture, (), modes=("async",)),
+            TraceScenario("sync-one", _fixture, (), asynchronous=False),
+            TraceScenario("async-one", _fixture, (), asynchronous=True),
+            TraceScenario("async-two", _fixture, (), asynchronous=True),
         ),
     )
     case: Final = _case()
 
-    nodes: Final = scenario_nodeids(suite, case, frozenset({"two"}))
+    nodes: Final = scenario_nodeids(suite, case, frozenset({"async-two"}))
 
-    assert tuple(nodeid for _, _, nodeid in nodes) == ("trace:sdk:ocr:two:async",)
+    assert tuple(nodeid for _, nodeid in nodes) == ("trace:sdk:ocr:async-two",)
+
+
+def test_python_engine_is_separate_from_scenario_selection() -> None:
+    assert runner_selection(("mistral", "--engine=python")) == (frozenset({"mistral"}), "python")
+
+
+def test_python_engine_skips_native_bridge(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner")
+    case: Final = _case()
+    selected: list[tuple[frozenset[str], str]] = []
+
+    def reject_bridge(_repo_root: Path) -> str | None:
+        raise AssertionError("Python-only tracing must not inspect or build the native bridge")
+
+    def capture_case(
+        _run: HarnessRun,
+        _case: HarnessCase,
+        scenarios: frozenset[str],
+        _on_update: object,
+        engine: str,
+    ) -> None:
+        selected.append((scenarios, engine))
+
+    monkeypatch.setattr(runner, "ensure_trace_bridge", reject_bridge)
+    monkeypatch.setattr(runner, "_run_case", capture_case)
+
+    exit_code, _ = run_trace_cases((case,), tmp_path, lambda _: None, ("mistral", "--engine=python"))
+
+    assert exit_code == 0
+    assert selected == [(frozenset({"mistral"}), "python")]
+
+
+def test_python_trace_preserves_native_ocr_dispatch_setting(monkeypatch: pytest.MonkeyPatch) -> None:
+    execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution")
+    route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture)
+    observed: list[str | None] = []
+
+    def collect(
+        _function: SdkCall,
+        _fixture: RouteFixture,
+        _engine: Engine,
+        *,
+        asynchronous: bool,
+    ) -> SimpleNamespace:
+        observed.append(os.environ.get("LITELLM_RUST"))
+        return SimpleNamespace(
+            events=(FunctionTraceEvent(0, None, "aocr" if asynchronous else "ocr"),),
+            error=None,
+        )
+
+    monkeypatch.setattr(execution, "_collect", collect)
+
+    monkeypatch.setenv("LITELLM_RUST", "0")
+    collect_trace(route, "python", asynchronous=False)
+    monkeypatch.setenv("LITELLM_RUST", "1")
+    collect_trace(route, "python", asynchronous=True)
+
+    assert observed == ["0", "1"]
+    assert os.environ["LITELLM_RUST"] == "1"
+
+
+def test_expected_provider_failure_omits_feedback_banner(
+    capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
+) -> None:
+    loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case")
+    suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
+    scenario: Final = next(item for item in suite.scenarios if item.name == "async-openai-provider-error")
+    monkeypatch.setattr(litellm, "suppress_debug_info", False)
+    assert isinstance(suite.route, RouteSpec)
+
+    result: Final = execute_trace(suite.route, scenario, "sdk", engine="python")
+
+    assert result.python_error is None
+    assert "Give Feedback / Get Help" not in capsys.readouterr().out
+    assert litellm.suppress_debug_info is False
+
+
+@pytest.mark.parametrize("asynchronous", (False, True))
+def test_vertex_trace_keeps_unmapped_helpers_and_parents(asynchronous: bool) -> None:
+    loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
+    suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
+    name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek"
+    scenario: Final = next(item for item in suite.scenarios if item.name == name)
+    assert isinstance(suite.route, RouteSpec)
+
+    trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python")
+
+    assert trace.python_error is None
+    url: Final = next(
+        event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.get_complete_url")
+    )
+    project: Final = next(
+        event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_project")
+    )
+    location: Final = next(
+        event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_location")
+    )
+    assert project.parent_id == location.parent_id == url.id
+    assert not any(event.raw.endswith(" VertexBase.get_access_token") for event in trace.python)
+
+
+@pytest.mark.parametrize("asynchronous", (False, True))
+def test_vertex_credentials_trace_runs_real_auth_helpers(asynchronous: bool, monkeypatch: pytest.MonkeyPatch) -> None:
+    loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
+    suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
+    name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek-credentials"
+    scenario: Final = next(item for item in suite.scenarios if item.name == name)
+    monkeypatch.setenv("VERTEXAI_CREDENTIALS", "original-credentials")
+    monkeypatch.setenv("VERTEX_AI_API_KEY", "original-api-key")
+    assert isinstance(suite.route, RouteSpec)
+
+    trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python")
+
+    assert trace.python_error is None
+    validate: Final = next(
+        event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.validate_environment")
+    )
+    helpers: Final = (
+        "VertexBase.safe_get_vertex_ai_project",
+        "VertexBase.safe_get_vertex_ai_credentials",
+        "VertexBase.get_access_token",
+    )
+    assert tuple(event.raw.split(" ", 1)[1] for event in trace.python if event.parent_id == validate.id) == helpers
+    token: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.get_access_token"))
+    load: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.load_auth"))
+    refresh: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.refresh_auth"))
+    assert load.parent_id == token.id
+    assert refresh.parent_id == load.id
+    assert os.environ["VERTEXAI_CREDENTIALS"] == "original-credentials"
+    assert os.environ["VERTEX_AI_API_KEY"] == "original-api-key"
+
+
+def test_gateway_trace_keeps_calls_outside_scenario_mappings(monkeypatch: pytest.MonkeyPatch) -> None:
+    execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution")
+    events: Final = (
+        FunctionTraceEvent(0, None, "route.py:1 entry"),
+        FunctionTraceEvent(1, 0, "auth.py:2 authenticate"),
+        FunctionTraceEvent(2, 1, "auth.py:3 credentials"),
+    )
+    scenario: Final = TraceScenario(
+        "async-gateway",
+        _fixture,
+        (mapping(rust_span="entry", python_frame=r" entry$"),),
+        asynchronous=True,
+    )
+    monkeypatch.setattr(execution, "_collect", lambda *_args: events)
+
+    trace: Final = execution.execute_gateway_trace(GatewayRouteSpec("messages"), scenario, engine="python")
+
+    assert trace.python_error is None
+    assert tuple((event.id, event.parent_id, event.raw) for event in trace.python) == tuple(
+        (event.id, event.parent_id, event.raw) for event in events
+    )
+
+
+def test_default_trace_skips_unavailable_rust_sdk_entrypoint(monkeypatch: pytest.MonkeyPatch) -> None:
+    execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution")
+    route: Final = RouteSpec("responses", ("responses", "aresponses"), None, _fixture)
+    scenario: Final = TraceScenario("sync-openai", _fixture, (), asynchronous=False)
+    engines: list[Engine] = []
+
+    def collect(_route: RouteSpec, engine: Engine, *, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]:
+        engines.append(engine)
+        return (FunctionTraceEvent(0, None, "responses"),)
+
+    monkeypatch.setattr(execution, "collect_trace", collect)
+
+    trace: Final = execution.execute_trace(route, scenario, "sdk")
+
+    assert engines == ["python"]
+    assert trace.engine == "python"
+    assert trace.rust_error is None
+
+
+def test_default_trace_skips_unavailable_rust_gateway_route(monkeypatch: pytest.MonkeyPatch) -> None:
+    execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution")
+    route: Final = GatewayRouteSpec("responses", rust_supported=False)
+    scenario: Final = TraceScenario("async-openai", _fixture, (), asynchronous=True)
+    engines: list[Engine] = []
+
+    def collect(_route: GatewayRouteSpec, _scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...]:
+        engines.append(engine)
+        return (FunctionTraceEvent(0, None, "responses"),)
+
+    monkeypatch.setattr(execution, "_collect", collect)
+
+    trace: Final = execution.execute_gateway_trace(route, scenario)
+
+    assert engines == ["python"]
+    assert trace.engine == "python"
+    assert trace.rust_error is None
 
 
 def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None:
     route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture)
     duplicate: Final = TraceSuite(
         route=route,
-        scenarios=(TraceScenario("same", _fixture, ()), TraceScenario("same", _fixture, ())),
+        scenarios=(
+            TraceScenario("sync-same", _fixture, (), asynchronous=False),
+            TraceScenario("sync-same", _fixture, (), asynchronous=False),
+        ),
+    )
+    unsafe: Final = TraceSuite(
+        route=route, scenarios=(TraceScenario("sync-bad:name", _fixture, (), asynchronous=False),)
     )
-    unsafe: Final = TraceSuite(route=route, scenarios=(TraceScenario("bad:name", _fixture, ()),))
     case: Final = _case()
 
     assert validate_trace_suite(duplicate, case) is not None
     assert validate_trace_suite(unsafe, case) is not None
 
 
-def test_scenario_validation_rejects_invalid_modes_and_route_registration() -> None:
-    invalid_modes: Final = TraceSuite(
+def test_scenario_validation_rejects_invalid_names_and_route_registration() -> None:
+    invalid_name: Final = TraceSuite(
         route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture),
-        scenarios=(TraceScenario("invalid", _fixture, (), modes=("sync", "sync")),),
+        scenarios=(TraceScenario("bedrock", _fixture, (), asynchronous=True),),
     )
     wrong_function: Final = TraceSuite(
         route=RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _fixture),
-        scenarios=(TraceScenario("one", _fixture, ()),),
+        scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),),
     )
     wrong_surface: Final = TraceSuite(
         route=GatewayRouteSpec("ocr"),
-        scenarios=(TraceScenario("one", _fixture, ()),),
+        scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),),
     )
     case: Final = _case()
 
-    assert "unique sync/async modes" in (validate_trace_suite(invalid_modes, case) or "")
+    assert "start with sync- or async-" in (validate_trace_suite(invalid_name, case) or "")
     assert "does not match case function" in (validate_trace_suite(wrong_function, case) or "")
     assert "must use RouteSpec" in (validate_trace_suite(wrong_surface, case) or "")
 
@@ -77,11 +285,35 @@ def test_invalid_route_dispatch_records_harness_error() -> None:
     result: Final = run.results[case.key]
     suite: Final = TraceSuite(
         route=GatewayRouteSpec("ocr"),
-        scenarios=(TraceScenario("one", _fixture, (), modes=("sync",)),),
+        scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),),
     )
-    nodeid: Final = "trace:sdk:ocr:one:sync"
+    nodeid: Final = "trace:sdk:ocr:sync-one"
 
-    run_trace_mode(run, result, suite, suite.scenarios[0], "sync", "sdk", nodeid, lambda _: None)
+    run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", nodeid, lambda _: None)
 
     assert result.outcomes[nodeid] is RunStatus.ERROR
     assert run.failures == [(nodeid, "gateway route cannot run on the sdk surface")]
+
+
+def test_different_python_and_rust_traces_pass(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner")
+    case: Final = _case()
+    run: Final = HarnessRun.from_cases((case,))
+    result: Final = run.results[case.key]
+    suite: Final = TraceSuite(
+        route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture),
+        scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),),
+    )
+    trace: Final = TraceArtifact.from_traces(
+        surface="sdk",
+        sdk_function="ocr",
+        scenario="sync-one",
+        python=(PipelineStep(0, None, "python_step", "python.py:1 python_step"),),
+        rust=(PipelineStep(0, None, "rust_step", "rust_step"),),
+    )
+    monkeypatch.setattr(runner, "_execute_scenario", lambda *_args: trace)
+
+    run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", "trace:sdk:ocr:sync-one", lambda _: None)
+
+    assert result.outcomes["trace:sdk:ocr:sync-one"] is RunStatus.PASSED
+    assert run.failures == []

tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py

@@ -38,30 +38,28 @@ def _trace_functions(
     python_functions: Final[dict[str, PythonFunctionIdentity]] = {}
     rust_functions: Final[dict[str, RustFunctionIdentity]] = {}
     for scenario in suite.scenarios:
-        for mode in scenario.modes:
-            route: Final = RouteSpec(
-                route=suite.route.route,
-                python_entrypoints=suite.route.python_entrypoints,
-                rust_entrypoints=suite.route.rust_entrypoints,
-                fixture=scenario.fixture,
-            )
-            python_trace: Final = collect_trace(route, "python", asynchronous=mode == "async")
-            rust_trace: Final = collect_trace(route, "rust", asynchronous=mode == "async")
-            if isinstance(python_trace, TraceExecutionFailure):
-                raise ValueError(f"Python trace discovery failed for {scenario.name}/{mode}: {python_trace.message}")
-            if isinstance(rust_trace, TraceExecutionFailure):
-                raise ValueError(f"Rust trace discovery failed for {scenario.name}/{mode}: {rust_trace.message}")
-            mappings: Final = scenario.mappings_for(mode)
-            python_projection: Final = pipeline_projection("python", python_trace, mappings)
-            rust_projection: Final = pipeline_projection("rust", rust_trace, mappings)
-            for step in python_projection.steps:
-                if step.span in spec.trace_spans:
-                    function: Final = PythonFunctionIdentity.from_trace(step.raw)
-                    python_functions[function.raw] = function
-            for step in rust_projection.steps:
-                if step.span in spec.trace_spans:
-                    function: Final = RustFunctionIdentity.from_trace(step.raw)
-                    rust_functions[step.raw] = function
+        route: Final = RouteSpec(
+            route=suite.route.route,
+            python_entrypoints=suite.route.python_entrypoints,
+            rust_entrypoints=suite.route.rust_entrypoints,
+            fixture=scenario.fixture,
+        )
+        python_trace: Final = collect_trace(route, "python", asynchronous=scenario.asynchronous)
+        rust_trace: Final = collect_trace(route, "rust", asynchronous=scenario.asynchronous)
+        if isinstance(python_trace, TraceExecutionFailure):
+            raise ValueError(f"Python trace discovery failed for {scenario.name}: {python_trace.message}")
+        if isinstance(rust_trace, TraceExecutionFailure):
+            raise ValueError(f"Rust trace discovery failed for {scenario.name}: {rust_trace.message}")
+        python_projection: Final = pipeline_projection("python", python_trace, scenario.mappings)
+        rust_projection: Final = pipeline_projection("rust", rust_trace, scenario.mappings)
+        for step in python_projection.steps:
+            if step.span in spec.trace_spans:
+                function: Final = PythonFunctionIdentity.from_trace(step.raw)
+                python_functions[function.raw] = function
+        for step in rust_projection.steps:
+            if step.span in spec.trace_spans:
+                function: Final = RustFunctionIdentity.from_trace(step.raw)
+                rust_functions[step.raw] = function
     if not python_functions or not rust_functions:
         raise ValueError(f"Python trace discovery found no functions for spans: {', '.join(spec.trace_spans)}")
     return (

tests/test_litellm/litellm_core_utils/test_realtime_streaming.py

@@ -10,9 +10,9 @@ from websockets.exceptions import ConnectionClosed
 from websockets.frames import Close
 
 import litellm
+from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.litellm_core_utils.realtime_streaming import (
-    REALTIME_SESSION_SUCCESS_LOGGED_KEY,
     RealTimeStreaming,
     client_sent_openai_beta_realtime_header,
 )
@@ -3399,6 +3399,26 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker()
     assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details
 
 
+@pytest.mark.asyncio
+async def test_refused_session_stamps_the_failure_ownership_marker():
+    """LIT-6463: the enqueued failure callback releases the key's max_parallel_requests
+    slot from the logging worker, so a refusal stamps REALTIME_SESSION_FAILURE_LOGGED_KEY.
+    The proxy endpoint reads it to leave the slot to that callback instead of racing it.
+    A session that relayed frames logs a success and must not carry the failure stamp."""
+    upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
+    refused: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close))
+    session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode()
+    relayed: Final = _relay_session(
+        _client_ws_that_never_sends(), _backend_ws_closing_with(session_created, upstream_close)
+    )
+
+    await refused.run()
+    await relayed.run()
+
+    assert refused.logging.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY) is True
+    assert REALTIME_SESSION_FAILURE_LOGGED_KEY not in relayed.logging.model_call_details
+
+
 @pytest.mark.asyncio
 async def test_transformed_transcription_completion_never_sends_response_create():
     from typing import Final

tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py

@@ -8,6 +8,7 @@ from unittest.mock import MagicMock
 import pytest
 
 import litellm
+from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY
 from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.llms.bedrock.realtime.handler import BedrockRealtime
 from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig
@@ -77,6 +78,7 @@ class UnavailableBedrockStream:
 class FakeLogging:
     def __init__(self, trace_id="trace-nova-sonic"):
         self.litellm_trace_id = trace_id
+        self.model_call_details = {}
 
 
 class DisconnectingClientWS:
@@ -673,6 +675,31 @@ class TestBedrockRealtimeProviderFailurePropagation:
         await spend_dispatch["coro"]
         assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
 
+    @pytest.mark.asyncio
+    async def test_success_dispatch_stamps_the_ownership_marker_only_when_spend_was_logged(
+        self, stub_aws_sdk_client, spend_dispatch
+    ):
+        stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)]
+        await BedrockRealtime().async_realtime(
+            model="amazon.nova-sonic-v1:0",
+            websocket=ConnectedClientWS([self.SESSION_UPDATE]),
+            logging_obj=spend_dispatch["logging_obj"],
+            **self.AWS_PARAMS,
+        )
+        await spend_dispatch["coro"]
+        assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
+        assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True
+
+        idle_logging = FakeLogging()
+        stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([])]
+        await BedrockRealtime().async_realtime(
+            model="amazon.nova-sonic-v1:0",
+            websocket=ConnectedClientWS([self.SESSION_UPDATE]),
+            logging_obj=idle_logging,
+            **self.AWS_PARAMS,
+        )
+        assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in idle_logging.model_call_details
+
     @pytest.mark.asyncio
     async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client):
         stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver)
@@ -684,6 +711,25 @@ class TestBedrockRealtimeProviderFailurePropagation:
 
         assert stream.input_stream.closed
 
+    @pytest.mark.asyncio
+    async def test_client_disconnect_ends_the_session_while_bedrock_output_stays_open(self, stub_aws_sdk_client):
+        receiver = DrainedThenOpenBedrockReceiver([])
+        stream = ScriptedBedrockStream([], receiver_type=lambda _payloads: receiver)
+        stub_aws_sdk_client["streams"] = [stream]
+
+        await asyncio.wait_for(
+            BedrockRealtime().async_realtime(
+                model="amazon.nova-sonic-v1:0",
+                websocket=RealtimeClientWS(),
+                logging_obj=FakeLogging(),
+                **self.AWS_PARAMS,
+            ),
+            timeout=1,
+        )
+
+        assert receiver.drained.is_set(), "the handler must have been waiting on the open provider stream"
+        assert stream.input_stream.closed
+
     @pytest.mark.asyncio
     async def test_session_updated_is_not_sent_before_bedrock_is_ready(self, stub_aws_models):
         handler = BedrockRealtime()

tests/test_litellm/proxy/auth/test_auth_checks.py

@@ -5527,7 +5527,9 @@ async def _run_internal_user_budget_alert(
 
     with (
         patch("litellm.proxy.proxy_server.prisma_client", None),  # test-quality-ok: common_checks has no database seam
-        patch("litellm.proxy.proxy_server.get_current_spend", _get_spend),  # test-quality-ok: common_checks imports it locally
+        patch(  # test-quality-ok: common_checks imports get_current_spend locally
+            "litellm.proxy.proxy_server.get_current_spend", _get_spend
+        ),
         patch.object(slack_alerting, "send_alert", send_alert),
     ):
         error: Final = await _check_for_error()
@@ -6419,9 +6421,7 @@ async def test_get_team_membership_negative_caches_a_missing_row():
     assert first is None
     assert second is None
     mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once()
-    cached = await cache.async_get_cache(
-        key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")
-    )
+    cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1"))
     assert cached == NO_TEAM_MEMBERSHIP_SENTINEL
 
 
@@ -6454,6 +6454,314 @@ async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model()
     mock_prisma_client.db.litellm_teammembership.find_unique.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_get_team_membership_coalesces_parallel_db_fetches():
+    from litellm.proxy.auth.auth_checks import get_team_membership
+
+    started = asyncio.Event()
+    release = asyncio.Event()
+    membership_row = MagicMock()
+    membership_row.dict = lambda: {"user_id": "u-parallel", "team_id": "t-parallel", "spend": 1.0}
+
+    async def _slow_find_unique(*args, **kwargs):
+        started.set()
+        await release.wait()
+        return membership_row
+
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_slow_find_unique)
+    cache = UserApiKeyCache()
+
+    async def _load():
+        return await get_team_membership(
+            user_id="u-parallel",
+            team_id="t-parallel",
+            prisma_client=mock_prisma_client,
+            user_api_key_cache=cache,
+        )
+
+    first = asyncio.create_task(_load())
+    second = asyncio.create_task(_load())
+    await started.wait()
+    await asyncio.sleep(0)
+    release.set()
+    results = await asyncio.gather(first, second)
+
+    assert results[0] is not None and results[1] is not None
+    assert results[0].user_id == "u-parallel"
+    assert results[1].user_id == "u-parallel"
+    mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_get_team_membership_invalidation_waits_for_in_flight_load_then_evicts_it():
+    from litellm.proxy._types import LiteLLM_TeamMembership
+    from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state
+    from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
+    from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
+
+    started = asyncio.Event()
+    release_stale = asyncio.Event()
+    rows = iter(("budget-old", "budget-new"))
+
+    async def _find_unique(*args, **kwargs):
+        budget_id = next(rows)
+        row = MagicMock()
+        row.dict = lambda: {"user_id": "u-inv", "team_id": "t-inv", "spend": 1.0, "budget_id": budget_id}
+        if budget_id == "budget-old":
+            started.set()
+            await release_stale.wait()
+        return row
+
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique)
+    cache = UserApiKeyCache()
+    _key = team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv")
+
+    async def _load():
+        return await get_team_membership(
+            user_id="u-inv", team_id="t-inv", prisma_client=mock_prisma_client, user_api_key_cache=cache
+        )
+
+    stale = asyncio.create_task(_load())
+    await asyncio.wait_for(started.wait(), timeout=2)
+    invalidation = asyncio.create_task(
+        invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache)
+    )
+    for _ in range(5):
+        await asyncio.sleep(0)
+    assert not invalidation.done()
+
+    release_stale.set()
+    await asyncio.wait_for(invalidation, timeout=2)
+    stale_result = await stale
+    assert stale_result is not None and stale_result.budget_id == "budget-old"
+    assert await cache.async_get_cache(key=_key) is None
+
+    fresh_result = await _load()
+    assert fresh_result is not None and fresh_result.budget_id == "budget-new"
+    assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2
+    cached = CacheCodec.deserialize(await cache.async_get_cache(key=_key), model_type=LiteLLM_TeamMembership)
+    assert cached is not None and cached.budget_id == "budget-new"
+    again = await _load()
+    assert again is not None and again.budget_id == "budget-new"
+    assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_entry():
+    from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state
+    from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
+
+    write_started = asyncio.Event()
+    release_write = asyncio.Event()
+
+    class _SlowWriteCache(UserApiKeyCache):
+        async def async_set_cache(self, key, value, local_only=False, **kwargs):
+            write_started.set()
+            await release_write.wait()
+            return await super().async_set_cache(key, value, local_only=local_only, **kwargs)
+
+    row = MagicMock()
+    row.dict = lambda: {"user_id": "u-w", "team_id": "t-w", "spend": 1.0, "budget_id": "budget-old"}
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=row)
+    cache = _SlowWriteCache()
+
+    stale = asyncio.create_task(
+        get_team_membership(user_id="u-w", team_id="t-w", prisma_client=mock_prisma_client, user_api_key_cache=cache)
+    )
+    await asyncio.wait_for(write_started.wait(), timeout=2)
+    invalidation = asyncio.create_task(
+        invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache)
+    )
+    for _ in range(5):
+        await asyncio.sleep(0)
+    assert not invalidation.done()
+
+    release_write.set()
+    await asyncio.wait_for(invalidation, timeout=2)
+    stale_result = await stale
+
+    assert stale_result is not None and stale_result.budget_id == "budget-old"
+    assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None
+
+
+@pytest.mark.asyncio
+async def test_common_checks_calls_get_team_membership_once_per_request():
+    from fastapi import Request
+
+    from litellm.proxy.auth.auth_checks import common_checks
+
+    team = LiteLLM_TeamTable(team_id="t-once")
+    token = UserAPIKeyAuth(token="k-once", user_id="u-once", team_id="t-once", models=["gpt-4o-mini"])
+    membership = MagicMock()
+    membership.litellm_budget_table = None
+    membership.spend = 0.0
+
+    with (
+        patch(  # test-quality-ok: common_checks imports prisma_client from proxy_server
+            "litellm.proxy.proxy_server.prisma_client", MagicMock()
+        ),
+        patch(  # test-quality-ok: common_checks imports user_api_key_cache from proxy_server
+            "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()
+        ),
+        patch(  # test-quality-ok: counts membership loads; common_checks has no membership seam
+            "litellm.proxy.auth.auth_checks.get_team_membership",
+            new_callable=AsyncMock,
+            return_value=membership,
+        ) as load_membership,
+        patch(  # test-quality-ok: common_checks imports get_current_spend locally
+            "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0
+        ),
+    ):
+        result = await common_checks(
+            request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]},
+            team_object=team,
+            user_object=LiteLLM_UserTable(user_id="u-once"),
+            end_user_object=None,
+            global_proxy_spend=None,
+            general_settings={},
+            route="/chat/completions",
+            llm_router=None,
+            proxy_logging_obj=MagicMock(),
+            valid_token=token,
+            request=MagicMock(spec=Request),
+        )
+
+    assert result is True
+    assert load_membership.await_count == 1
+
+
+@pytest.mark.asyncio
+async def test_common_checks_skips_membership_load_when_no_check_reads_it():
+    from fastapi import Request
+
+    from litellm.proxy.auth.auth_checks import common_checks
+
+    team = LiteLLM_TeamTable(team_id="t-lazy")
+    token = UserAPIKeyAuth(token="k-lazy", user_id="u-lazy", team_id="t-lazy")
+
+    with (
+        patch(  # test-quality-ok: common_checks imports prisma_client from proxy_server
+            "litellm.proxy.proxy_server.prisma_client", MagicMock()
+        ),
+        patch(  # test-quality-ok: common_checks imports user_api_key_cache from proxy_server
+            "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()
+        ),
+        patch(  # test-quality-ok: counts membership loads; common_checks has no membership seam
+            "litellm.proxy.auth.auth_checks.get_team_membership",
+            new_callable=AsyncMock,
+        ) as load_membership,
+    ):
+        result = await common_checks(
+            request_body={},
+            team_object=team,
+            user_object=LiteLLM_UserTable(user_id="u-lazy"),
+            end_user_object=None,
+            global_proxy_spend=None,
+            general_settings={},
+            route="/key/info",
+            llm_router=None,
+            proxy_logging_obj=MagicMock(),
+            valid_token=token,
+            request=MagicMock(spec=Request),
+        )
+
+    assert result is True
+    load_membership.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_get_team_membership_db_error_returns_none_and_retries_next_call():
+    from litellm.proxy.auth.auth_checks import get_team_membership
+    from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
+
+    membership_row = MagicMock()
+    membership_row.dict = lambda: {"user_id": "u-fail", "team_id": "t-fail", "spend": 1.0}
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(
+        side_effect=[RuntimeError("db down"), membership_row]
+    )
+    cache = UserApiKeyCache()
+
+    failed = await get_team_membership(
+        user_id="u-fail",
+        team_id="t-fail",
+        prisma_client=mock_prisma_client,
+        user_api_key_cache=cache,
+    )
+    cached_after_failure = await cache.async_get_cache(
+        key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail")
+    )
+    recovered = await get_team_membership(
+        user_id="u-fail",
+        team_id="t-fail",
+        prisma_client=mock_prisma_client,
+        user_api_key_cache=cache,
+    )
+
+    assert failed is None
+    assert cached_after_failure is None
+    assert recovered is not None
+    assert recovered.user_id == "u-fail"
+    assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_get_team_membership_string_prisma_client_returns_none():
+    from litellm.proxy.auth.auth_checks import get_team_membership
+
+    result = await get_team_membership(
+        user_id="u-str",
+        team_id="t-str",
+        prisma_client="hello-world",
+        user_api_key_cache=UserApiKeyCache(),
+    )
+    assert result is None
+
+
+@pytest.mark.asyncio
+async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load():
+    from litellm.proxy.auth.auth_checks import get_team_membership
+
+    started = asyncio.Event()
+    release = asyncio.Event()
+    membership_row = MagicMock()
+    membership_row.dict = lambda: {"user_id": "u-shield", "team_id": "t-shield", "spend": 1.0}
+
+    async def _slow_find_unique(*args, **kwargs):
+        started.set()
+        await release.wait()
+        return membership_row
+
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_slow_find_unique)
+    cache = UserApiKeyCache()
+
+    async def _load():
+        return await get_team_membership(
+            user_id="u-shield",
+            team_id="t-shield",
+            prisma_client=mock_prisma_client,
+            user_api_key_cache=cache,
+        )
+
+    owner = asyncio.create_task(_load())
+    await started.wait()
+    waiter = asyncio.create_task(_load())
+    await asyncio.sleep(0)
+    waiter.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await waiter
+    release.set()
+    result = await owner
+
+    assert result is not None
+    assert result.user_id == "u-shield"
+    mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel():
     """
@@ -6477,10 +6785,7 @@ async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sent
     assert before is None
 
     await invalidate_team_member_spend_state(user_id="u-1", team_id="t-1", user_api_key_cache=cache)
-    assert (
-        await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1"))
-        is None
-    )
+    assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) is None
 
     after = await get_team_membership(
         user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache

tests/test_litellm/proxy/auth/test_model_access_group_budgets.py

@@ -31,6 +31,7 @@ from litellm.proxy.auth.auth_checks import (
 )
 from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key
 from litellm.proxy.common_utils.user_api_key_cache import (
+    NO_TEAM_MEMBERSHIP_SENTINEL,
     UserApiKeyCache,
     model_access_group_registry_cache_key,
     model_access_group_spend_counter_key,
@@ -95,6 +96,11 @@ async def _cache(
             ),
             model_type=LiteLLM_TeamMembership,
         )
+    else:
+        await cache.async_set_cache(
+            key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID),
+            value=NO_TEAM_MEMBERSHIP_SENTINEL,
+        )
     if org_models:
         await cache.async_set_cache(
             key=f"org_id:{ORG_ID}",
@@ -314,9 +320,7 @@ class _RecordingPrismaClient:
     def __init__(self, *rows: _MagBudgetRow) -> None:
         self.rows = {row.access_group_name: row for row in rows}
         self.batches: list[list[str]] = []
-        self.db = SimpleNamespace(
-            litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many)
-        )
+        self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many))
 
     async def _find_many(self, **kwargs):
         requested = list(kwargs["where"]["access_group_name"]["in"])
@@ -345,7 +349,9 @@ async def _enforce(
     read, seen = _spend_reader(spend_by_counter_key or {})
     # The check takes its client and cache as arguments, injected just below. get_current_spend is the
     # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter.
-    with patch("litellm.proxy.proxy_server.get_current_spend", read):  # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point
+    with patch(  # test-quality-ok: get_current_spend is lazily imported inside the budget check
+        "litellm.proxy.proxy_server.get_current_spend", read
+    ):
         await _model_access_group_max_budget_check(
             matched_model_access_groups=matched,
             prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows),
@@ -491,9 +497,7 @@ async def test_a_second_request_serves_the_budget_row_from_cache():
 async def test_a_database_error_does_not_block_the_request():
     class _FailingPrismaClient:
         def __init__(self) -> None:
-            self.db = SimpleNamespace(
-                litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom)
-            )
+            self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom))
 
         async def _boom(self, **kwargs):
             raise RuntimeError("database unavailable")
@@ -507,11 +511,15 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) ->
     read, _ = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 99.0})
 
     with (
-        # common_checks resolves all three off the proxy_server module at call time; its signature
-        # has no client, cache or spend-reader parameter to pass them through instead.
-        patch("litellm.proxy.proxy_server.prisma_client", prisma_client),  # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter
-        patch("litellm.proxy.proxy_server.user_api_key_cache", cache),  # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter
-        patch("litellm.proxy.proxy_server.get_current_spend", read),  # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point
+        patch(  # test-quality-ok: common_checks lazily imports prisma_client from proxy_server
+            "litellm.proxy.proxy_server.prisma_client", prisma_client
+        ),
+        patch(  # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server
+            "litellm.proxy.proxy_server.user_api_key_cache", cache
+        ),
+        patch(  # test-quality-ok: get_current_spend is lazily imported inside the budget check
+            "litellm.proxy.proxy_server.get_current_spend", read
+        ),
     ):
         return await common_checks(
             request_body={"model": "gpt-4o", "messages": []},
@@ -524,7 +532,9 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) ->
             llm_router=Router(model_list=MODEL_LIST),
             proxy_logging_obj=ProxyLogging(user_api_key_cache=cache),
             valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID),
-            request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")),
+            request=SimpleNamespace(
+                method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")
+            ),
             skip_budget_checks=skip_budget_checks,
         )
 

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -14366,3 +14366,123 @@ async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids
     assert exc_info.value.status_code == 400
     assert expected_error in str(exc_info.value.detail)
     mock_db_client.db.query_raw.assert_not_called()
+
+
+class _TeamRowWithOrganization(LiteLLM_TeamTable):
+    litellm_organization_table: LiteLLM_OrganizationTable | None = None
+
+
+@pytest.mark.parametrize(
+    "organization, expected_models",
+    [
+        (
+            LiteLLM_OrganizationTable(
+                organization_id="org-1",
+                budget_id="budget-1",
+                models=["all-proxy-models"],
+                created_by="admin",
+                updated_by="admin",
+            ),
+            ["all-proxy-models"],
+        ),
+        (
+            LiteLLM_OrganizationTable(
+                organization_id="org-1",
+                budget_id="budget-1",
+                models=["gpt-4o"],
+                created_by="admin",
+                updated_by="admin",
+            ),
+            ["gpt-4o"],
+        ),
+        (None, None),
+    ],
+)
+@pytest.mark.asyncio
+async def test_team_info_returns_parent_organization_models(organization, expected_models):
+    """/team/info must report the parent org's model ceiling.
+
+    A team admin who is not an org admin gets a 403 from /organization/info, so this
+    is the only read that can tell the Admin UI whether the org allows all proxy
+    models. Without it the team edit form hides the "All Proxy Models" option and a
+    team admin cannot grant their team everything on the proxy.
+    """
+    from fastapi import Request
+
+    from litellm.proxy.management_endpoints import team_endpoints
+
+    team_row = _TeamRowWithOrganization(
+        team_id="team-1",
+        organization_id="org-1" if organization is not None else None,
+        litellm_organization_table=organization,
+    )
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
+    mock_prisma.get_data = AsyncMock(return_value=[])
+
+    memberships = AsyncMock(return_value=[])
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),  # test-quality-ok: no seam on team_info
+        patch.object(team_endpoints, "get_all_team_memberships", memberships),  # test-quality-ok: no seam on team_info
+    ):
+        response = await team_endpoints.team_info(
+            http_request=MagicMock(spec=Request),
+            team_id="team-1",
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+
+    assert response["team_info"].organization_models == expected_models
+
+
+@pytest.mark.parametrize(
+    "caller, expected_models",
+    [
+        (UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.INTERNAL_USER), ["gpt-4o"]),
+        (UserAPIKeyAuth(user_id="member-1", user_role=LitellmUserRoles.INTERNAL_USER), None),
+        (UserAPIKeyAuth(team_id="team-1"), None),
+    ],
+)
+@pytest.mark.asyncio
+async def test_team_info_reports_parent_organization_models_only_to_team_managers(caller, expected_models):
+    """Plain members and team keys can read their team, but not the org's wider allow-list."""
+    from fastapi import Request
+
+    from litellm.proxy.management_endpoints import team_endpoints
+
+    team_row = _TeamRowWithOrganization(
+        team_id="team-1",
+        organization_id="org-1",
+        members_with_roles=[
+            Member(user_id="admin-1", role="admin"),
+            Member(user_id="member-1", role="user"),
+        ],
+        litellm_organization_table=LiteLLM_OrganizationTable(
+            organization_id="org-1",
+            budget_id="budget-1",
+            models=["gpt-4o"],
+            created_by="admin",
+            updated_by="admin",
+        ),
+    )
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
+    mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[])
+    mock_prisma.get_data = AsyncMock(return_value=[])
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),  # test-quality-ok: no seam on team_info
+        patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])),  # test-quality-ok: no seam on team_info
+        patch.object(  # test-quality-ok: no seam on team_info
+            team_endpoints, "_is_user_org_admin_for_team", AsyncMock(return_value=False)
+        ),
+    ):
+        response = await team_endpoints.team_info(
+            http_request=MagicMock(spec=Request),
+            team_id="team-1",
+            user_api_key_dict=caller,
+        )
+
+    assert response["team_info"].organization_models == expected_models

tests/test_litellm/proxy/test_proxy_server.py

@@ -33,6 +33,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
 from litellm.caching.dual_cache import DualCache
 from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash
 from litellm.proxy.proxy_server import app, initialize
 from litellm.utils import _invalidate_model_cost_lowercase_map
 
@@ -9999,6 +10000,7 @@ async def _lit6973_drive_realtime_session(
     reservation: dict,
     *,
     backend_logged_success: bool,
+    backend_logged_failure: bool = False,
     phase_one_exit: str | None = None,
     websocket: MagicMock | None = None,
 ) -> MagicMock:
@@ -10006,8 +10008,9 @@ async def _lit6973_drive_realtime_session(
 
     phase_one_exit picks a rejection before the relay: "model_access" makes the
     key/model check raise ProxyException, "pre_call" makes pre-call processing
-    (rate limits, guardrails) raise. Neither reaches route_request, so no success
-    log can own the reservation and the endpoint has to release it on that exit.
+    (rate limits, guardrails) raise, "pre_call_cancelled" cancels the task inside
+    pre-call processing. None reaches route_request, so no success log can own the
+    reservation and the endpoint has to release it on that exit.
 
     route_request resolves normally in both cases: the relay owns the session
     once route_request returns. A successful session enqueues its success cost
@@ -10017,7 +10020,7 @@ async def _lit6973_drive_realtime_session(
     logging object carries a real model_call_details dict so the stamp is
     observable, and the reservation has empty entries so the real release touches
     no counter store."""
-    from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY
+    from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
     from litellm.proxy import proxy_server as ps
 
     user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token")
@@ -10029,6 +10032,8 @@ async def _lit6973_drive_realtime_session(
     async def fake_llm_call() -> None:
         if backend_logged_success:
             logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
+        if backend_logged_failure:
+            logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True
 
     from litellm.proxy._types import ProxyException
 
@@ -10037,7 +10042,13 @@ async def _lit6973_drive_realtime_session(
         if phase_one_exit == "model_access"
         else None
     )
-    pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None
+    pre_call_error: Final = (
+        asyncio.CancelledError()
+        if phase_one_exit == "pre_call_cancelled"
+        else Exception("Rate limit exceeded")
+        if phase_one_exit == "pre_call"
+        else None
+    )
     pre_call: Final = AsyncMock(
         side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)
     )
@@ -10155,6 +10166,114 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c
     assert reservation["finalized"] is False
 
 
+_LIT6463_COUNTER_KEY: Final = "{api_key:hashed-token}:max_parallel_requests"
+
+
+async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
+    *,
+    backend_logged_success: bool,
+    backend_logged_failure: bool = False,
+    phase_one_exit: str | None = None,
+) -> tuple[DualCache, RequestRateLimiterStash]:
+    """Run the realtime endpoint with a real v3 limiter registered and the request's
+    stash already holding slot-1 of a two-slot counter, the state pre-call leaves
+    behind. Returns the limiter's cache and the stash so the test can read what the
+    endpoint did to the slot."""
+    from litellm.proxy import proxy_server as ps
+    from litellm.proxy.hooks.parallel_request_limiter_v3 import (
+        _PROXY_MaxParallelRequestsHandler_v3,
+        _request_stash,
+    )
+    from litellm.proxy.utils import InternalUsageCache
+
+    dual_cache: Final = DualCache()
+    await dual_cache.async_set_cache(
+        key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True
+    )
+    limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache))
+    stash: Final = RequestRateLimiterStash(
+        parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}
+    )
+    reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
+
+    stash_token: Final = _request_stash.set(stash)
+    try:
+        hooks: Final = patch.dict(  # test-quality-ok: registers the real limiter the route's release reads
+            ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}
+        )
+        expected_exit: Final = (
+            pytest.raises(asyncio.CancelledError)
+            if phase_one_exit == "pre_call_cancelled"
+            else contextlib.nullcontext()
+        )
+        with hooks, expected_exit:
+            await _lit6973_drive_realtime_session(
+                reservation,
+                backend_logged_success=backend_logged_success,
+                backend_logged_failure=backend_logged_failure,
+                phase_one_exit=phase_one_exit,
+            )
+    finally:
+        _request_stash.reset(stash_token)
+    return dual_cache, stash
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("phase_one_exit", [None, "pre_call", "pre_call_cancelled"])
+async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_parallel_slot(
+    phase_one_exit: str | None,
+):
+    """The rate limiter acquires the key's max_parallel_requests slot in pre-call and
+    only frees it from the LLM success/failure callbacks. A realtime session that ends
+    without either callback (Bedrock closes without usage events, a later pre-call hook
+    rejects the session, or the task is cancelled while still in pre-call) has to be
+    released by the route itself, or the slot stays occupied until its TTL and the key's
+    next session is refused with a 429."""
+    dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
+        backend_logged_success=False, phase_one_exit=phase_one_exit
+    )
+
+    assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {"slot-2": 2.0}
+    assert stash.parallel_slot is None
+
+
+@pytest.mark.asyncio
+async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_callback():
+    """A session that enqueued its success callback hands the slot to the limiter's
+    own success handler, which runs on the logging worker. If the route also released
+    it, the two releases would race on the same stashed acquisition and, under the
+    limiter's integer in-memory fallback, double-decrement the counter so the key
+    admits more sessions than max_parallel_requests allows. With the success stamp
+    present the route leaves the slot and the stash alone."""
+    dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
+        backend_logged_success=True
+    )
+
+    assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {
+        "slot-1": 1.0,
+        "slot-2": 2.0,
+    }
+    assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}
+
+
+@pytest.mark.asyncio
+async def test_refused_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_failure_callback():
+    """An upstream refusal before any frame enqueues the failure callback instead, and
+    the limiter's failure handler releases the slot from the logging worker just like
+    the success handler does. The route sees no success stamp, so it still settles the
+    budget reservation, but it must leave the slot to that callback or the two releases
+    race on the same acquisition."""
+    dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
+        backend_logged_success=False, backend_logged_failure=True
+    )
+
+    assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {
+        "slot-1": 1.0,
+        "slot-2": 2.0,
+    }
+    assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}
+
+
 @pytest.mark.asyncio
 async def test_release_or_invalidate_falls_back_to_invalidating_the_counters():
     """If releasing the reservation itself fails (e.g. the counter store is down),

tests/test_litellm/router_utils/test_cooldown_handlers.py

@@ -437,3 +437,67 @@ class TestRoutingGroupCooldownAlternatives:
             )
             is False
         )
+
+
+class TestTeamModelCooldownAlternatives:
+    def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()) -> litellm.Router:
+        return litellm.Router(
+            model_list=[
+                {
+                    "model_name": f"model_name_team-1_{i}",
+                    "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
+                    "model_info": {
+                        "id": f"team-deploy-{i}",
+                        "team_id": "team-1",
+                        "team_public_model_name": "team-gpt-4o-mini",
+                        "blocked": f"team-deploy-{i}" in blocked_ids,
+                    },
+                }
+                for i in range(team_deployments)
+            ]
+        )
+
+    def test_429_on_team_deployment_with_sibling_cools_down(self):
+        from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment
+
+        router = self._router(team_deployments=2)
+        assert (
+            _should_cooldown_deployment(
+                litellm_router_instance=router,
+                deployment="team-deploy-0",
+                exception_status=429,
+                original_exception=Exception("rate limited"),
+                requested_model_group="team-gpt-4o-mini",
+            )
+            is True
+        )
+
+    def test_429_on_only_team_deployment_keeps_single_deployment_exemption(self):
+        from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment
+
+        router = self._router(team_deployments=1)
+        assert (
+            _should_cooldown_deployment(
+                litellm_router_instance=router,
+                deployment="team-deploy-0",
+                exception_status=429,
+                original_exception=Exception("rate limited"),
+                requested_model_group="team-gpt-4o-mini",
+            )
+            is False
+        )
+
+    def test_429_with_only_a_blocked_sibling_keeps_single_deployment_exemption(self):
+        from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment
+
+        router = self._router(team_deployments=2, blocked_ids=frozenset({"team-deploy-1"}))
+        assert (
+            _should_cooldown_deployment(
+                litellm_router_instance=router,
+                deployment="team-deploy-0",
+                exception_status=429,
+                original_exception=Exception("rate limited"),
+                requested_model_group="team-gpt-4o-mini",
+            )
+            is False
+        )

tests/test_litellm/test_router.py

@@ -896,6 +896,46 @@ def test_arouter_test_team_model():
     assert result is not None
 
 
+def test_team_model_has_alternatives():
+    def team_deployment(
+        deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False
+    ) -> DeploymentTypedDict:
+        return {
+            "model_name": f"model_name_{team_id}_{deployment_id}",
+            "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
+            "model_info": {
+                "id": deployment_id,
+                "team_id": team_id,
+                "team_public_model_name": public_model_name,
+                "blocked": blocked,
+            },
+        }
+
+    router = litellm.Router(
+        model_list=[
+            team_deployment("team-a-1", "team-a", "shared-model"),
+            team_deployment("team-a-2", "team-a", "shared-model"),
+            team_deployment("team-a-solo", "team-a", "solo-model"),
+            team_deployment("team-b-1", "team-b", "shared-model"),
+            team_deployment("team-c-1", "team-c", "paused-sibling-model"),
+            team_deployment("team-c-paused", "team-c", "paused-sibling-model", blocked=True),
+            {
+                "model_name": "plain-model",
+                "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
+                "model_info": {"id": "plain-1"},
+            },
+        ],
+    )
+
+    assert router.team_model_has_alternatives("team-a-1") is True
+    assert router.team_model_has_alternatives("team-a-2") is True
+    assert router.team_model_has_alternatives("team-a-solo") is False
+    assert router.team_model_has_alternatives("team-b-1") is False
+    assert router.team_model_has_alternatives("team-c-1") is False
+    assert router.team_model_has_alternatives("plain-1") is False
+    assert router.team_model_has_alternatives("missing-deployment") is False
+
+
 def test_arouter_ignore_invalid_deployments():
     """
     Test that router.ignore_invalid_deployments is set to True

ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts

@@ -328,7 +328,8 @@ describe("useTeam", () => {
   });
 
   it("should return team data when query is successful", async () => {
-    (teamInfoCall as any).mockResolvedValue(mockTeams[0]);
+    // /team/info answers with an envelope; the hook is typed as the team itself.
+    (teamInfoCall as any).mockResolvedValue({ team_id: "team-1", team_info: mockTeams[0], keys: [] });
 
     const { result } = renderHook(() => useTeam("team-1"), { wrapper });
 

ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts

@@ -163,7 +163,8 @@ export const useTeam = (teamId?: string) => {
         throw new Error("Missing auth or teamId");
       }
 
-      return teamInfoCall(accessToken, teamId);
+      const { team_info } = (await teamInfoCall(accessToken, teamId)) as { team_info: Team };
+      return team_info;
     },
 
     initialData: () => {

ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.test.tsx

@@ -242,13 +242,11 @@ describe("ProjectDetail", () => {
     it("should show team information when team data is available", () => {
       mockUseTeam.mockReturnValue({
         data: {
-          team_info: {
-            team_id: "team-1",
-            team_alias: "Engineering",
-            models: ["gpt-4"],
-            spend: 50,
-            members_with_roles: [],
-          },
+          team_id: "team-1",
+          team_alias: "Engineering",
+          models: ["gpt-4"],
+          spend: 50,
+          members_with_roles: [],
         },
         isLoading: false,
       });

ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.tsx

@@ -14,16 +14,6 @@ import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 import { EditProjectModal } from "./ProjectModals/EditProjectModal";
 import { ProjectKeysSection } from "./ProjectKeysSection";
 
-interface TeamInfoShape {
-  team_id: string;
-  team_alias?: string;
-  models?: string[];
-  max_budget?: number | null;
-  budget_duration?: string | null;
-  spend?: number;
-  members_with_roles?: { user_id: string; role: string }[];
-}
-
 interface ProjectDetailProps {
   projectId: string;
   onBack: () => void;
@@ -33,10 +23,7 @@ const utilisationTone = (percent: number) => (percent >= 90 ? "over" : percent >
 
 export function ProjectDetail({ projectId, onBack }: ProjectDetailProps) {
   const { data: project, isLoading } = useProjectDetails(projectId);
-  const { data: teamData } = useTeam(project?.team_id ?? undefined);
-  // teamInfoCall returns { team_id, team_info: {...}, keys, team_memberships }
-  const teamInfo: TeamInfoShape | undefined = ((teamData as unknown as { team_info?: TeamInfoShape })?.team_info ??
-    teamData) as TeamInfoShape | undefined;
+  const { data: teamInfo } = useTeam(project?.team_id ?? undefined);
   const [isEditModalVisible, setIsEditModalVisible] = useState(false);
 
   const spend = project?.spend ?? 0;

ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx

@@ -415,6 +415,139 @@ describe("ModelSelect", () => {
     }
   });
 
+  it("should take the org model ceiling from the team when /organization/info is not readable", async () => {
+    const testCases = [
+      {
+        name: "org allows all proxy models",
+        organizationModels: ["all-proxy-models"],
+        shouldShowSentinel: true,
+        offered: ["gpt-4", "claude-3"],
+        notOffered: [] as string[],
+      },
+      {
+        name: "org places no ceiling at all",
+        organizationModels: [],
+        shouldShowSentinel: true,
+        offered: ["gpt-4", "claude-3"],
+        notOffered: [] as string[],
+      },
+      {
+        name: "org restricts the team to one model",
+        organizationModels: ["gpt-4"],
+        shouldShowSentinel: false,
+        offered: ["gpt-4"],
+        notOffered: ["claude-3"],
+      },
+    ];
+
+    for (const testCase of testCases) {
+      const user = userEvent.setup();
+      // A team admin gets a 403 from /organization/info, so the org query never resolves.
+      mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any);
+      mockUseTeam.mockReturnValue({
+        data: { team_id: "team-1", organization_models: testCase.organizationModels },
+        isLoading: false,
+      } as any);
+
+      const { unmount } = renderWithProviders(
+        <ModelSelect
+          onChange={mockOnChange}
+          context="team"
+          teamID="team-1"
+          organizationID="org-1"
+          options={{ includeSpecialOptions: true }}
+        />,
+      );
+
+      await openModelList(user);
+      if (testCase.shouldShowSentinel) {
+        expectOffered("All Proxy Models");
+      } else {
+        expectNotOffered("All Proxy Models");
+      }
+      expectOffered("No Default Models");
+      testCase.offered.forEach(expectOffered);
+      testCase.notOffered.forEach(expectNotOffered);
+
+      unmount();
+    }
+  });
+
+  it("should stay in the loading state while a list-seeded team is still fetching its org ceiling", () => {
+    mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any);
+    mockUseTeam.mockReturnValue({
+      data: { team_id: "team-1", models: [] },
+      isLoading: false,
+      isFetching: true,
+    } as any);
+
+    renderWithProviders(
+      <ModelSelect
+        onChange={mockOnChange}
+        context="team"
+        teamID="team-1"
+        organizationID="org-1"
+        options={{ includeSpecialOptions: true }}
+      />,
+    );
+
+    expect(screen.queryAllByRole("combobox")).toHaveLength(0);
+  });
+
+  it("should not hold the loading state on a background refetch once the org ceiling is known", async () => {
+    const user = userEvent.setup();
+    mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any);
+    mockUseTeam.mockReturnValue({
+      data: { team_id: "team-1", organization_models: ["all-proxy-models"] },
+      isLoading: false,
+      isFetching: true,
+    } as any);
+
+    renderWithProviders(
+      <ModelSelect
+        onChange={mockOnChange}
+        context="team"
+        teamID="team-1"
+        organizationID="org-1"
+        options={{ includeSpecialOptions: true }}
+      />,
+    );
+
+    await openModelList(user);
+    expectOffered("All Proxy Models");
+  });
+
+  it("should offer no models for an org team when neither the team nor the org reports a ceiling", async () => {
+    const testCases = [
+      { name: "/team/info withheld the ceiling", team: { team_id: "team-1", organization_models: null } },
+      { name: "/team/info failed after the list seeded the team", team: { team_id: "team-1", models: [] } },
+    ];
+
+    for (const testCase of testCases) {
+      const user = userEvent.setup();
+      mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any);
+      mockUseTeam.mockReturnValue({ data: testCase.team, isLoading: false, isFetching: false } as any);
+
+      const { unmount } = renderWithProviders(
+        <ModelSelect
+          onChange={mockOnChange}
+          context="team"
+          teamID="team-1"
+          organizationID="org-1"
+          options={{ includeSpecialOptions: true }}
+        />,
+      );
+
+      await openModelList(user);
+      expectNotOffered("All Proxy Models");
+      expectOffered("No Default Models");
+      expectNotOffered("gpt-4");
+      expectNotOffered("claude-3");
+
+      unmount();
+    }
+  });
+
   it("should use custom dataTestId when provided", async () => {
     renderWithProviders(
       <ModelSelect

ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx

@@ -19,7 +19,7 @@ import {
 } from "@/components/ui/combobox";
 import { Skeleton } from "@/components/ui/skeleton";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
-import { Organization, Team } from "../networking";
+import type { Team } from "@/components/key_team_helpers/key_list";
 import { splitWildcardModels } from "./modelUtils";
 
 const MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE = {
@@ -69,12 +69,19 @@ type ModelOptionGroup = {
 
 type FilterContextArgs = {
   allProxyModels: string[];
-  selectedTeam?: Team;
-  selectedOrganization?: Organization;
+  organizationID?: string;
+  organizationModels?: string[];
   userModels?: string[];
   options?: ModelSelectProps["options"];
 };
 
+const isUncappedModelCeiling = (organizationModels: string[]) =>
+  organizationModels.length === 0 || organizationModels.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value);
+
+// useTeam seeds from the team list, which omits organization_models; /team/info is the only source of the org ceiling.
+const isAwaitingOrganizationModels = (team: Team | undefined, isFetchingTeam: boolean) =>
+  isFetchingTeam && team !== undefined && team.organization_models === undefined;
+
 const contextFilters: Record<ModelSelectProps["context"], (args: FilterContextArgs) => string[]> = {
   user: ({ allProxyModels, userModels, options }) => {
     if (!userModels) return [];
@@ -82,18 +89,10 @@ const contextFilters: Record<ModelSelectProps["context"], (args: FilterContextAr
     return [];
   },
 
-  team: ({ allProxyModels, selectedOrganization, userModels }) => {
-    if (selectedOrganization) {
-      if (
-        selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) ||
-        selectedOrganization.models.length === 0
-      ) {
-        return allProxyModels;
-      }
-      return allProxyModels.filter((model) => selectedOrganization.models.includes(model));
-    }
-
-    return allProxyModels ?? [];
+  team: ({ allProxyModels, organizationID, organizationModels }) => {
+    if (organizationModels === undefined) return organizationID ? [] : allProxyModels;
+    if (isUncappedModelCeiling(organizationModels)) return allProxyModels;
+    return allProxyModels.filter((model) => organizationModels.includes(model));
   },
 
   organization: ({ allProxyModels }) => {
@@ -108,7 +107,7 @@ const contextFilters: Record<ModelSelectProps["context"], (args: FilterContextAr
 const filterModels = (
   allProxyModels: ProxyModel[],
   ctx: ModelSelectProps,
-  extra: { selectedTeam?: Team; selectedOrganization?: Organization; userModels?: string[] },
+  extra: { organizationModels?: string[]; userModels?: string[] },
 ): string[] => {
   const deduplicatedProxyModels = Array.from(new Map(allProxyModels.map((m) => [m.id, m])).values()).map(
     (model) => model.id,
@@ -118,7 +117,13 @@ const filterModels = (
   const filterFn = contextFilters[ctx.context];
   if (!filterFn) return [];
 
-  return filterFn({ allProxyModels: deduplicatedProxyModels, ...extra, options: ctx.options });
+  const filterArgs: FilterContextArgs = {
+    allProxyModels: deduplicatedProxyModels,
+    organizationID: ctx.organizationID,
+    ...extra,
+    options: ctx.options,
+  };
+  return filterFn(filterArgs);
 };
 
 export const ModelSelect = (props: ModelSelectProps) => {
@@ -126,16 +131,17 @@ export const ModelSelect = (props: ModelSelectProps) => {
   const { id, teamID, organizationID, options, context, dataTestId, value = [], onChange, style } = props;
   const { showAllProxyModelsOverride, includeSpecialOptions } = options || {};
   const { data: allProxyModels, isLoading: isLoadingAllProxyModels } = useAllProxyModels();
-  const { data: team, isLoading: isLoadingTeam } = useTeam(teamID);
+  const { data: team, isLoading: isLoadingTeam, isFetching: isFetchingTeam } = useTeam(teamID);
   const { data: organization, isLoading: isLoadingOrganization } = useOrganization(organizationID);
   const { data: currentUser, isLoading: isCurrentUserLoading } = useCurrentUser();
 
   const isSpecialOption = (value: string) => MODEL_SENTINEL_OPTIONS.some((sv) => sv.value === value);
   const hasSpecialOptionSelected = value.some(isSpecialOption);
-  const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization || isCurrentUserLoading;
-  const organizationHasAllProxyModels =
-    organization?.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) ||
-    organization?.models.length === 0;
+  const isTeamPending = isLoadingTeam || isAwaitingOrganizationModels(team, isFetchingTeam);
+  const isLoading = isLoadingAllProxyModels || isTeamPending || isLoadingOrganization || isCurrentUserLoading;
+  // The org's ceiling rides on /team/info, which a team admin may read; /organization/info 403s for them.
+  const organizationModels = team?.organization_models ?? organization?.models;
+  const organizationHasAllProxyModels = organizationModels !== undefined && isUncappedModelCeiling(organizationModels);
   const shouldShowAllProxyModels =
     showAllProxyModelsOverride || (organizationHasAllProxyModels && includeSpecialOptions) || context === "global";
 
@@ -159,8 +165,7 @@ export const ModelSelect = (props: ModelSelectProps) => {
   };
 
   const filteredModels = filterModels(allProxyModels?.data ?? [], props, {
-    selectedTeam: team,
-    selectedOrganization: organization,
+    organizationModels,
     userModels: currentUser?.models,
   });
 

ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx

@@ -27,6 +27,8 @@ export interface Team {
   access_group_models?: string[];
   access_group_mcp_server_ids?: string[];
   access_group_agent_ids?: string[];
+  // Parent org's model ceiling. undefined = no org / not loaded; [] or ["all-proxy-models"] = no ceiling.
+  organization_models?: string[] | null;
 }
 
 export interface KeyResponse {

ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx

@@ -1,12 +1,15 @@
 import React from "react";
-import { Control } from "react-hook-form";
+import { Control, UseFormReturn } from "react-hook-form";
 import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { Textarea } from "@/components/ui/textarea";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { CircleHelp } from "lucide-react";
 import { FormField } from "@/components/shared/form/FormField";
+import { toast } from "@/lib/toast";
 import AgentSelector from "../agent_management/AgentSelector";
 import NumericalInput from "../shared/numerical_input";
 import SkillSelector from "../skills/SkillSelector";
+import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers";
 import { AgentsAndGroups, KeyEditFormValues } from "./keyEditFormValues";
 
 export const labelWithHint = (label: React.ReactNode, hint: string): React.ReactNode => (
@@ -85,6 +88,42 @@ export const KeyAgentAndSkillFields = ({
   </>
 );
 
+type KeyEditForm = Pick<
+  UseFormReturn<KeyEditFormValues, unknown, KeyEditFormValues>,
+  "control" | "getValues" | "setValue"
+>;
+
+export const moveMetadataTagsToTagsField = (form: KeyEditForm): void => {
+  const moved = moveTagsOutOfMetadataJson(form.getValues("metadata"), form.getValues("tags"));
+  if (moved === null) return;
+  form.setValue("metadata", moved.metadata, { shouldDirty: true });
+  form.setValue("tags", moved.tags, { shouldDirty: true });
+  if (moved.movedTags.length > 0) {
+    toast.info(`Moved ${moved.movedTags.join(", ")} from metadata to the Tags field`);
+  }
+};
+
+export const KeyMetadataField = ({ form }: { form: KeyEditForm }) => (
+  <FormField
+    control={form.control}
+    name="metadata"
+    label="Metadata"
+    description="Tags are managed by the Tags field above. A tags array typed here is moved to that field."
+  >
+    {(field) => (
+      <Textarea
+        {...field}
+        value={(field.value as string | undefined) ?? ""}
+        rows={10}
+        onBlur={() => {
+          field.onBlur();
+          moveMetadataTagsToTagsField(form);
+        }}
+      />
+    )}
+  </FormField>
+);
+
 export const KeyBudgetNumberField = ({
   control,
   name,

ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.test.ts

@@ -0,0 +1,50 @@
+import { describe, expect, it } from "vitest";
+
+import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers";
+
+describe("moveTagsOutOfMetadataJson", () => {
+  it("moves a tags array out of the JSON and appends it to the current tags", () => {
+    expect(moveTagsOutOfMetadataJson('{"tags": ["pilot-tag"], "env": "non-prod"}', ["ui-tag"])).toEqual({
+      metadata: '{\n  "env": "non-prod"\n}',
+      tags: ["ui-tag", "pilot-tag"],
+      movedTags: ["pilot-tag"],
+    });
+  });
+
+  it("drops duplicates and non-string entries without reporting them as moved", () => {
+    expect(moveTagsOutOfMetadataJson('{"tags": ["a", "a", "ui-tag", 7, null]}', ["ui-tag"])).toEqual({
+      metadata: "{}",
+      tags: ["ui-tag", "a"],
+      movedTags: ["a"],
+    });
+  });
+
+  it("trims whitespace and drops blank entries the same way the Tags control does", () => {
+    expect(moveTagsOutOfMetadataJson('{"tags": [" a ", "a", "   ", "", " ui-tag"]}', ["ui-tag"])).toEqual({
+      metadata: "{}",
+      tags: ["ui-tag", "a"],
+      movedTags: ["a"],
+    });
+  });
+
+  it("strips an empty tags array while moving nothing", () => {
+    expect(moveTagsOutOfMetadataJson('{"tags": []}', undefined)).toEqual({
+      metadata: "{}",
+      tags: [],
+      movedTags: [],
+    });
+  });
+
+  it("is a no-op when there is no tags array to move", () => {
+    expect(moveTagsOutOfMetadataJson('{"env": "prod"}', ["a"])).toBeNull();
+    expect(moveTagsOutOfMetadataJson('{"tags": "not-an-array"}', ["a"])).toBeNull();
+    expect(moveTagsOutOfMetadataJson("", ["a"])).toBeNull();
+    expect(moveTagsOutOfMetadataJson(undefined, ["a"])).toBeNull();
+  });
+
+  it("leaves invalid or non-object JSON alone so the save path can report it", () => {
+    expect(moveTagsOutOfMetadataJson('{"tags": [', ["a"])).toBeNull();
+    expect(moveTagsOutOfMetadataJson('["tags"]', ["a"])).toBeNull();
+    expect(moveTagsOutOfMetadataJson("null", ["a"])).toBeNull();
+  });
+});

ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts

@@ -34,6 +34,39 @@ export const modelSentinelOptions = (
   return teamLoaded ? [{ value: "all-team-models", label: "All Team Models" }] : [];
 };
 
+export type MovedMetadataTags = {
+  metadata: string;
+  tags: string[];
+  movedTags: string[];
+};
+
+const parseJsonObject = (text: string): Record<string, unknown> | null => {
+  try {
+    const parsed: unknown = JSON.parse(text);
+    return parsed !== null && typeof parsed === "object" && !Array.isArray(parsed)
+      ? (parsed as Record<string, unknown>)
+      : null;
+  } catch {
+    return null;
+  }
+};
+
+export const moveTagsOutOfMetadataJson = (
+  metadataJson: string | undefined,
+  currentTags: readonly string[] | undefined,
+): MovedMetadataTags | null => {
+  const parsed = metadataJson === undefined ? null : parseJsonObject(metadataJson);
+  if (parsed === null) return null;
+  const { tags: metadataTags, ...rest } = parsed;
+  if (!Array.isArray(metadataTags)) return null;
+  const existing = currentTags ?? [];
+  const movedTags = metadataTags
+    .filter((tag: unknown): tag is string => typeof tag === "string")
+    .map((tag) => tag.trim())
+    .filter((tag, index, all) => tag.length > 0 && !existing.includes(tag) && all.indexOf(tag) === index);
+  return { metadata: JSON.stringify(rest, null, 2), tags: [...existing, ...movedTags], movedTags };
+};
+
 export const currentValuePlaceholder = (
   premiumUser: boolean,
   current: unknown,

ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx

@@ -2180,6 +2180,32 @@ describe("KeyEditView", () => {
       expect(onSubmitMock.mock.calls[0][0].tags).toEqual(["test-tag", "typed-tag"]);
     });
 
+    it("moves a tags array typed into the metadata JSON into the Tags control on blur", async () => {
+      renderForPayload(vi.fn().mockResolvedValue(undefined));
+      await screen.findByRole("button", { name: /save changes/i });
+
+      const metadata = screen.getByLabelText("Metadata");
+      fireEvent.change(metadata, { target: { value: '{"tags": ["pilot-tag"], "env": "non-prod"}' } });
+      fireEvent.blur(metadata);
+
+      expect(await screen.findByText("pilot-tag")).toBeInTheDocument();
+      expect(metadata).toHaveValue('{\n  "env": "non-prod"\n}');
+    });
+
+    it("carries a tags array typed into the metadata JSON into the payload even without a blur", async () => {
+      const onSubmitMock = vi.fn().mockResolvedValue(undefined);
+      renderForPayload(onSubmitMock);
+      await screen.findByRole("button", { name: /save changes/i });
+
+      fireEvent.change(screen.getByLabelText("Metadata"), { target: { value: '{"tags": ["pilot-tag"]}' } });
+      fireEvent.submit(screen.getByRole("button", { name: /save changes/i }));
+
+      await waitFor(() => {
+        expect(onSubmitMock).toHaveBeenCalled();
+      });
+      expect(onSubmitMock.mock.calls[0][0]).toMatchObject({ tags: ["test-tag", "pilot-tag"], metadata: "{}" });
+    });
+
     const pickFromCombobox = async (inputLabel: RegExp | string, optionName: RegExp | string) => {
       await userEvent.click(screen.getByLabelText(inputLabel));
       await userEvent.click(await screen.findByRole("option", { name: optionName }));

ui/litellm-dashboard/src/components/templates/key_edit_view.tsx

@@ -31,7 +31,14 @@ import {
   modelSentinelOptions,
   parseAllowedRoutes,
 } from "./keyEditFieldNormalizers";
-import { KeyAgentAndSkillFields, KeyBudgetNumberField, KeyTypeSelect, labelWithHint } from "./KeyEditViewControls";
+import {
+  KeyAgentAndSkillFields,
+  KeyBudgetNumberField,
+  KeyMetadataField,
+  KeyTypeSelect,
+  labelWithHint,
+  moveMetadataTagsToTagsField,
+} from "./KeyEditViewControls";
 import {
   KeyEditFormValues,
   keyEditFormSchema,
@@ -341,9 +348,12 @@ export function KeyEditView({
   return (
     <TooltipProvider>
       <form
-        onSubmit={form.handleSubmit((values) =>
-          handleSubmit(toSubmittedValues(values, { canViewPolicies, canViewPrompts })),
-        )}
+        onSubmit={(event) => {
+          moveMetadataTagsToTagsField(form);
+          return form.handleSubmit((values) =>
+            handleSubmit(toSubmittedValues(values, { canViewPolicies, canViewPrompts })),
+          )(event);
+        }}
       >
         <FieldGroup>
           <FormField control={form.control} name="key_alias" label="Key Alias">
@@ -843,9 +853,7 @@ export function KeyEditView({
             )}
           </FormField>
 
-          <FormField control={form.control} name="metadata" label="Metadata">
-            {(field) => <Textarea {...field} value={(field.value as string | undefined) ?? ""} rows={10} />}
-          </FormField>
+          <KeyMetadataField form={form} />
 
           <div className="mb-4">
             <FormField control={form.control} name="duration">