diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 41198887cab..33c68e95fd1 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -753,6 +753,7 @@ jobs: tests/unit/responses/test_metadata_codex_callback.py tests/unit/responses/test_no_duplicate_spend_logs.py tests/unit/responses/test_null_test_fix.py + tests/unit/responses/test_regression_issue_28553.py tests/unit/responses/test_responses_api_bridge_flag.py tests/unit/responses/test_responses_api_lifecycle.py tests/unit/responses/test_responses_api_request_body.py diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 5dc9679fe78..447c18afc88 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -12,12 +12,14 @@ A2A Protocol Format: import json from collections.abc import Sequence +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Optional from pydantic import ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.llms.a2a.common_utils import serialize_a2a_data_part from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -37,6 +39,8 @@ class _A2ATextPart(TypedDict, total=False): kind: ReadOnly[str] text: ReadOnly[str] + data: ReadOnly[object] + parts: ReadOnly[Sequence["_A2ATextPart"]] class A2AGuardrailHandler(BaseTranslation): @@ -48,8 +52,10 @@ class A2AGuardrailHandler(BaseTranslation): 2. Process output responses (post-call hook) - extracts text from A2A response parts A2A Message Format: - - Input: params.message.parts[].text (where kind == "text") - - Output: result.message.parts[].text or result.artifacts[].parts[].text + - Input: params.message.parts[].text (where kind == "text") or + params.message.parts[].data (where kind == "data") + - Output: result.message.parts[].text or result.artifacts[].parts[].text, + and the "data" equivalents of both """ async def process_input_messages( @@ -80,20 +86,28 @@ class A2AGuardrailHandler(BaseTranslation): verbose_proxy_logger.debug("A2A: No parts in message, skipping guardrail") return data - texts_to_check: Final[list[str]] = [] - text_part_indices: Final[list[int]] = [] # Track which parts contain text + def _scan_input_part(part_idx: int, part: dict[str, Any]) -> tuple[str, int, str] | None: + kind: Final = part.get("kind") + if kind == "text": + text: Final = part.get("text", "") + return (text, part_idx, "text") if text else None + if kind == "data": + part_data: Final = part.get("data") + return (serialize_a2a_data_part(part_data), part_idx, "data") if part_data is not None else None + return None - # Step 1: Extract text from all text parts - for part_idx, part in enumerate(parts): - if part.get("kind") == "text": - text = part.get("text", "") - if text: - texts_to_check.append(text) - text_part_indices.append(part_idx) + # Extract text from all text parts, and serialized data from all data parts + scanned: Final = tuple( + entry for part_idx, part in enumerate(parts) if (entry := _scan_input_part(part_idx, part)) is not None + ) + texts_to_check: Final = tuple(text for text, _, _ in scanned) + # Track which parts contain scannable content, and which field to write + # the guardrailed value back to ("text" or "data") + part_mappings: Final = tuple((part_idx, field) for _, part_idx, field in scanned) # Step 2: Apply guardrail to all texts in batch if texts_to_check: - inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + inputs: Final = GenericGuardrailAPIInputs(texts=list(texts_to_check)) # Pass the structured A2A message to guardrails inputs["structured_messages"] = [message] @@ -113,9 +127,9 @@ class A2AGuardrailHandler(BaseTranslation): guardrailed_texts: Final = guardrailed_inputs.get("texts", []) # Step 3: Apply guardrailed text back to original parts - if guardrailed_texts and len(guardrailed_texts) == len(text_part_indices): - for task_idx, part_idx in enumerate(text_part_indices): - parts[part_idx]["text"] = guardrailed_texts[task_idx] + if guardrailed_texts and len(guardrailed_texts) == len(part_mappings): + for task_idx, (part_idx, field) in enumerate(part_mappings): + parts[part_idx][field] = guardrailed_texts[task_idx] verbose_proxy_logger.debug("A2A: Processed input message: %s", message) @@ -163,18 +177,12 @@ class A2AGuardrailHandler(BaseTranslation): verbose_proxy_logger.debug("A2A: No result in response, skipping guardrail") return response - # Find all text-containing parts in the response - texts_to_check: Final[list[str]] = [] - # Each mapping is (path_to_parts_list, part_index) - # path_to_parts_list is a tuple of keys to navigate to the parts list - task_mappings: Final[list[tuple[tuple[str, ...], int]]] = [] - - # Extract texts from all possible locations - self._extract_texts_from_result( - result=result, - texts_to_check=texts_to_check, - task_mappings=task_mappings, - ) + # Find all text-containing parts in the response. Each scanned entry is + # (text, path_to_parts_list, part_index, field); path_to_parts_list is a + # tuple of keys to navigate to the parts list. + scanned: Final = self._extract_texts_from_result(result=result) + texts_to_check: Final = tuple(text for text, _, _, _ in scanned) + task_mappings: Final = tuple((path, part_idx, field) for _, path, part_idx, field in scanned) if not texts_to_check: verbose_proxy_logger.debug("A2A: No text content in response") @@ -195,7 +203,7 @@ class A2AGuardrailHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata - inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + inputs: Final = GenericGuardrailAPIInputs(texts=list(texts_to_check)) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -208,11 +216,12 @@ class A2AGuardrailHandler(BaseTranslation): # Step 3: Apply guardrailed text back to original response if guardrailed_texts and len(guardrailed_texts) == len(task_mappings): - for task_idx, (path, part_idx) in enumerate(task_mappings): + for task_idx, (path, part_idx, field) in enumerate(task_mappings): self._apply_text_to_path( result=result, path=path, part_idx=part_idx, + field=field, text=guardrailed_texts[task_idx], ) @@ -284,31 +293,27 @@ class A2AGuardrailHandler(BaseTranslation): result = obj.get("result", {}) if not isinstance(result, dict): continue - texts_in_chunk: list[str] = [] - mappings: list[tuple[tuple[str, ...], int]] = [] - self._extract_texts_from_result( - result=result, - texts_to_check=texts_in_chunk, - task_mappings=mappings, - ) - if not mappings: + scanned = self._extract_texts_from_result(result=result) + if not scanned: continue if orig_i == first_chunk_with_text: # Put full guardrailed text in first text part; clear others - for task_idx, (path, part_idx) in enumerate(mappings): + for task_idx, (_, path, part_idx, field) in enumerate(scanned): text = guardrailed_text if task_idx == 0 else "" self._apply_text_to_path( result=result, path=path, part_idx=part_idx, + field=field, text=text, ) else: - for path, part_idx in mappings: + for _, path, part_idx, field in scanned: self._apply_text_to_path( result=result, path=path, part_idx=part_idx, + field=field, text="", ) @@ -365,9 +370,7 @@ class A2AGuardrailHandler(BaseTranslation): def _extract_texts_from_result( self, result: dict[str, Any], - texts_to_check: list[str], - task_mappings: list[tuple[tuple[str, ...], int]], - ) -> None: + ) -> tuple[tuple[str, tuple[str, ...], int, str], ...]: """ Extract text from all possible locations in an A2A result. @@ -377,80 +380,100 @@ class A2AGuardrailHandler(BaseTranslation): 3. Task with artifacts: {"artifacts": [{"parts": [...]}]} 4. Task with status message: {"status": {"message": {"parts": [...]}}} 5. Streaming artifact-update: {"artifact": {"parts": [...]}} + + Returns a tuple of (text, path_to_parts_list, part_index, field) entries. """ # Case 1: Direct parts in result (direct message) - if "parts" in result: - self._extract_texts_from_parts( - parts=result["parts"], - path=("parts",), - texts_to_check=texts_to_check, - task_mappings=task_mappings, - ) + case1: Final = ( + self._extract_texts_from_parts(parts=result["parts"], path=("parts",)) if "parts" in result else () + ) # Case 2: Nested message message: Final = result.get("message") - if message and isinstance(message, dict) and "parts" in message: - self._extract_texts_from_parts( - parts=message["parts"], - path=("message", "parts"), - texts_to_check=texts_to_check, - task_mappings=task_mappings, - ) + case2: Final = ( + self._extract_texts_from_parts(parts=message["parts"], path=("message", "parts")) + if message and isinstance(message, dict) and "parts" in message + else () + ) # Case 3: Streaming artifact-update (singular artifact) artifact: Final = result.get("artifact") - if artifact and isinstance(artifact, dict) and "parts" in artifact: - self._extract_texts_from_parts( - parts=artifact["parts"], - path=("artifact", "parts"), - texts_to_check=texts_to_check, - task_mappings=task_mappings, - ) + case3: Final = ( + self._extract_texts_from_parts(parts=artifact["parts"], path=("artifact", "parts")) + if artifact and isinstance(artifact, dict) and "parts" in artifact + else () + ) # Case 4: Task with status message status: Final = result.get("status", {}) - if isinstance(status, dict): - status_message: Final = status.get("message") - if status_message and isinstance(status_message, dict) and "parts" in status_message: - self._extract_texts_from_parts( - parts=status_message["parts"], - path=("status", "message", "parts"), - texts_to_check=texts_to_check, - task_mappings=task_mappings, - ) + status_message: Final = status.get("message") if isinstance(status, dict) else None + case4: Final = ( + self._extract_texts_from_parts(parts=status_message["parts"], path=("status", "message", "parts")) + if isinstance(status_message, dict) and "parts" in status_message + else () + ) # Case 5: Task with artifacts (plural, array) artifacts: Final = result.get("artifacts", []) - if artifacts and isinstance(artifacts, list): - for artifact_idx, art in enumerate(artifacts): - if isinstance(art, dict) and "parts" in art: - self._extract_texts_from_parts( - parts=art["parts"], - path=("artifacts", str(artifact_idx), "parts"), - texts_to_check=texts_to_check, - task_mappings=task_mappings, - ) + case5: Final = ( + tuple( + chain.from_iterable( + self._extract_texts_from_parts(parts=art["parts"], path=("artifacts", str(artifact_idx), "parts")) + for artifact_idx, art in enumerate(artifacts) + if isinstance(art, dict) and "parts" in art + ) + ) + if artifacts and isinstance(artifacts, list) + else () + ) + + return case1 + case2 + case3 + case4 + case5 def _extract_texts_from_parts( self, parts: Sequence[_A2ATextPart], path: tuple[str, ...], - texts_to_check: list[str], - task_mappings: list[tuple[tuple[str, ...], int]], - ) -> None: - """Extract text from message parts.""" - for part_idx, part in enumerate(parts): - if part.get("kind") == "text": - text = part.get("text", "") - if text: - texts_to_check.append(text) - task_mappings.append((path, part_idx)) + depth: int = 0, + max_depth: int = 10, + ) -> tuple[tuple[str, tuple[str, ...], int, str], ...]: + """ + Extract text from message parts, serialized data from data parts, and + recurse into any part that itself carries a nested "parts" list. + + Mirrors `extract_text_from_a2a_message`'s handling (including the + recursion depth guard) so the two stay in sync. Returns a tuple of + (text, path_to_parts_list, part_index, field) entries. + """ + if depth >= max_depth: + return () + + def _scan(part_idx: int, part: _A2ATextPart) -> tuple[tuple[str, tuple[str, ...], int, str], ...]: + kind: Final = part.get("kind") + if kind == "text": + text: Final = part.get("text", "") + return ((text, path, part_idx, "text"),) if text else () + if kind == "data": + part_data: Final = part.get("data") + if part_data is None: + return () + return ((serialize_a2a_data_part(part_data), path, part_idx, "data"),) + if "parts" in part: + return self._extract_texts_from_parts( + parts=part["parts"], + path=path + (str(part_idx), "parts"), + depth=depth + 1, + max_depth=max_depth, + ) + return () + + return tuple(chain.from_iterable(_scan(part_idx, part) for part_idx, part in enumerate(parts))) def _apply_text_to_path( self, result: dict[str | int, Any], path: tuple[str, ...], part_idx: int, + field: str, text: str, ) -> None: """Apply guardrailed text back to the specified path in the result.""" @@ -463,5 +486,5 @@ class A2AGuardrailHandler(BaseTranslation): else: current = current[key] - # Update the text in the part - current[part_idx]["text"] = text + # Update the guardrailed value in the part + current[part_idx][field] = text diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index be3b5fea601..3319cfb9262 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -413,6 +413,7 @@ class A2AConfig(BaseConfig): role: Final = message.get("role", "user") return { + "kind": "message", "role": role, "parts": [{"kind": "text", "text": str(content)}], "messageId": str(uuid.uuid4()), diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index f6ff6d3bc47..39c6a0f1fd0 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -2,6 +2,7 @@ Common utilities for A2A (Agent-to-Agent) Protocol """ +import json from collections.abc import Awaitable, Callable, Mapping from typing import Any, Final @@ -67,6 +68,19 @@ def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str: return "\n".join(conversation_parts) +def serialize_a2a_data_part(data: Any) -> str: + """ + Serialize an A2A ``data``-kind part's payload to text. + + Used both to build the flattened completion text shown to callers and to + extract guardrail-scannable text, so the two stay in sync. + """ + try: + return json.dumps(data, ensure_ascii=False) + except (TypeError, ValueError): + return str(data) + + def extract_text_from_a2a_message(message: dict[str, Any], depth: int = 0, max_depth: int = 10) -> str: """ Extract text content from A2A message parts. @@ -86,8 +100,13 @@ def extract_text_from_a2a_message(message: dict[str, Any], depth: int = 0, max_d text_parts: Final[list[str]] = [] for part in parts: - if part.get("kind") == "text": + kind = part.get("kind") + if kind == "text": text_parts.append(part.get("text", "")) + elif kind == "data": + data = part.get("data") + if data is not None: + text_parts.append(serialize_a2a_data_part(data)) # Handle nested parts if they exist elif "parts" in part: nested_text = extract_text_from_a2a_message(part, depth + 1, max_depth) diff --git a/ruff-strict.toml b/ruff-strict.toml index e9db4d95cc8..323d9d7ce1d 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -35,6 +35,11 @@ external = [ # hook signatures inherit `Any` for `response` from CustomLogger, so narrowing them here # would break the override rather than describe it. "litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py" = ["ANN401"] +# `serialize_a2a_data_part`'s `data` param is the payload of an A2A `data`-kind part, +# which the A2A protocol defines as arbitrary JSON-serializable content with no fixed +# shape. `object` was not viable: json.dumps requires the narrower JSON-value types, +# so typing it `object` just moves the Any into a cast at every call site instead. +"litellm/llms/a2a/common_utils.py" = ["ANN401"] [lint.mccabe] max-complexity = 15 diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index f4ca7fa2965..fe63d8b5453 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -48,6 +48,7 @@ IGNORE_FUNCTIONS = [ "_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation. "_basic_json_schema_validate", # max depth set. "extract_text_from_a2a_message", # max depth set (default 10) to prevent infinite recursion in A2A message parsing. + "_extract_texts_from_parts", # max depth set (default 10), mirrors extract_text_from_a2a_message's recursion guard. "_convert_to_json_serializable_dict", # max depth set (default 20) and circular reference protection to prevent infinite recursion. "dict", # max depth set. _LiteLLMParamsDictView.dict() calls builtin dict(), not itself. "_read_image_bytes", # max depth set. diff --git a/tests/unit/a2a_protocol/test_regression_issue_28577.py b/tests/unit/a2a_protocol/test_regression_issue_28577.py new file mode 100644 index 00000000000..ae8b670083d --- /dev/null +++ b/tests/unit/a2a_protocol/test_regression_issue_28577.py @@ -0,0 +1,63 @@ +from litellm.llms.a2a.chat.transformation import A2AConfig +from litellm.llms.a2a.common_utils import extract_text_from_a2a_message + + +def test_regression_issue_28577_a2a_discriminator(): + """ + Test that A2A transformation adds the mandatory 'kind': 'message' discriminator. + Fixes Bug 1 in #28577. + """ + config = A2AConfig() + messages = [{"role": "user", "content": "ping"}] + + # transform_request creates the A2A JSON-RPC payload + request_data = config.transform_request( + model="a2a/demo", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + # Check Bug 1: message.kind missing + a2a_message = request_data["params"]["message"] + assert a2a_message["kind"] == "message" + assert a2a_message["role"] == "user" + assert "parts" in a2a_message + + +def test_regression_issue_28577_a2a_data_serialization(): + """ + Test that A2A common_utils handle kind: 'data' parts by serializing them. + Fixes Bug 2 in #28577. + """ + message_with_data = { + "kind": "message", + "role": "assistant", + "parts": [{"kind": "data", "data": {"result": {"msg": "pong"}}}], + "messageId": "msg-123", + } + + text = extract_text_from_a2a_message(message_with_data) + assert '"result": {"msg": "pong"}' in text + + +def test_regression_issue_28577_a2a_blocking_param(): + """ + Test that A2A requests include configuration.blocking: True. + Fixes Bug 3 in #28577 (async task unblocking). + """ + config = A2AConfig() + messages = [{"role": "user", "content": "ping"}] + + request_data = config.transform_request( + model="a2a/demo", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + # Check Bug 3 fix: configuration.blocking = True + assert "configuration" in request_data["params"] + assert request_data["params"]["configuration"]["blocking"] is True diff --git a/tests/unit/llms/a2a/chat/guardrail_translation/test_handler.py b/tests/unit/llms/a2a/chat/guardrail_translation/test_handler.py new file mode 100644 index 00000000000..150538f1d2c --- /dev/null +++ b/tests/unit/llms/a2a/chat/guardrail_translation/test_handler.py @@ -0,0 +1,167 @@ +""" +Unit tests for A2A Protocol Guardrail Translation Handler + +Regression coverage for the "data"-kind part guardrail bypass: A2A responses +can carry structured content in `kind: "data"` parts, which +`extract_text_from_a2a_message` (used to build the completion text callers +see) folds into the final text, but the guardrail handler previously only +inspected `kind: "text"` parts, so guarded output checks were skipped for +that content path. +""" + +from typing import Any, Literal, Optional + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler +from litellm.types.utils import GenericGuardrailAPIInputs + + +class MockGuardrail(CustomGuardrail): + """Mock guardrail that uppercases text so we can assert exactly what was scanned and where the result landed.""" + + def __init__(self, guardrail_name: str = "test"): + super().__init__(guardrail_name=guardrail_name) + self.last_inputs: Optional[GenericGuardrailAPIInputs] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.last_inputs = inputs + texts = inputs.get("texts", []) + return {"texts": [text.upper() for text in texts]} + + +@pytest.mark.asyncio +async def test_process_output_response_scans_data_parts(): + """A `kind: data` part in the output must be sent to the guardrail and the + guardrailed value written back into `data`, not silently skipped.""" + handler = A2AGuardrailHandler() + guardrail = MockGuardrail() + + response = { + "result": { + "kind": "message", + "parts": [ + {"kind": "text", "text": "hello"}, + {"kind": "data", "data": {"secret": "leak-me"}}, + ], + } + } + + result = await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + ) + + # The data part's serialized content must have reached the guardrail. + assert guardrail.last_inputs is not None + scanned_texts = guardrail.last_inputs["texts"] + assert any("leak-me" in t for t in scanned_texts) + + # The guardrailed (uppercased) value must be written back into "data", + # and the part must remain a "data" part, not be silently dropped or + # converted into an unguarded pass-through. + data_part = result["result"]["parts"][1] + assert data_part["kind"] == "data" + assert "LEAK-ME" in data_part["data"] + + # The text part must still be guardrailed as before (no regression). + text_part = result["result"]["parts"][0] + assert text_part["text"] == "HELLO" + + +@pytest.mark.asyncio +async def test_process_output_response_data_only_still_scanned(): + """A response with ONLY a data part (no text parts at all) must not be + skipped as "no text content in response".""" + handler = A2AGuardrailHandler() + guardrail = MockGuardrail() + + response = { + "result": { + "kind": "message", + "parts": [{"kind": "data", "data": {"result": {"msg": "pong"}}}], + } + } + + result = await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + ) + + assert guardrail.last_inputs is not None + assert guardrail.last_inputs["texts"] + assert "PONG" in result["result"]["parts"][0]["data"] + + +@pytest.mark.asyncio +async def test_process_output_response_scans_nested_parts(): + """A part that itself carries a nested "parts" list (grouping sub-parts) + must be recursed into, matching extract_text_from_a2a_message's own + recursion, instead of being silently skipped as neither text nor data.""" + handler = A2AGuardrailHandler() + guardrail = MockGuardrail() + + response = { + "result": { + "kind": "message", + "parts": [ + { + "kind": "group", + "parts": [ + {"kind": "text", "text": "hello"}, + {"kind": "data", "data": {"secret": "leak-me"}}, + ], + }, + ], + } + } + + result = await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + ) + + assert guardrail.last_inputs is not None + scanned_texts = guardrail.last_inputs["texts"] + assert "hello" in scanned_texts + assert any("leak-me" in t for t in scanned_texts) + + nested_parts = result["result"]["parts"][0]["parts"] + assert nested_parts[0]["text"] == "HELLO" + assert "LEAK-ME" in nested_parts[1]["data"] + + +@pytest.mark.asyncio +async def test_process_input_messages_scans_data_parts(): + """The same bypass existed on the request/input side of the handler.""" + handler = A2AGuardrailHandler() + guardrail = MockGuardrail() + + data = { + "params": { + "message": { + "kind": "message", + "role": "user", + "parts": [{"kind": "data", "data": {"secret": "leak-me"}}], + } + } + } + + result = await handler.process_input_messages( + data=data, + guardrail_to_apply=guardrail, + ) + + assert guardrail.last_inputs is not None + assert any("leak-me" in t for t in guardrail.last_inputs["texts"]) + + data_part = result["params"]["message"]["parts"][0] + assert data_part["kind"] == "data" + assert "LEAK-ME" in data_part["data"] diff --git a/tests/unit/responses/test_regression_issue_28553.py b/tests/unit/responses/test_regression_issue_28553.py new file mode 100644 index 00000000000..5ff7f9aeadc --- /dev/null +++ b/tests/unit/responses/test_regression_issue_28553.py @@ -0,0 +1,23 @@ +import pytest + +from litellm.proxy.common_request_processing import _stream_usage_tracking_updates + + +@pytest.mark.parametrize( + "route_type,expect_injected", + [ + ("acompletion", True), + ("atext_completion", False), + ("aresponses", False), + ("aget_responses", False), + ("arealtime", False), + ], +) +def test_stream_usage_only_injected_for_whitelisted_route_types(route_type, expect_injected): + updates = _stream_usage_tracking_updates( + data={"stream": True, "model": "gpt-4o"}, + general_settings={}, + route_type=route_type, + supports_stream_options=lambda: True, + ) + assert ("stream_options" in updates) is expect_injected