test(integration): look up streamed Responses spend rows by the client-facing id (#49752)

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/providers/test_bedrock_converse_missing_content_wire.py

@@ -547,7 +547,7 @@ def test_r21_responses_stream_input_item_without_content_after_an_assistant_item
         assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_ANSWERED], received
         (row,) = _success_rows(marker, 1)
         assert row["call_type"] == "aresponses" and row["end_user"] == marker, row
-        assert row["request_id"] == _inner_response_id(str(completed["response"]["id"])), (row, completed)
+        assert row["request_id"] == str(completed["response"]["id"]), (row, completed)
 
 
 def test_r22_passthrough_converse_forwards_a_message_without_content_verbatim(gateway: Gateway) -> None:

tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py

@@ -1,7 +1,6 @@
 import asyncio
 import base64
 import json
-import os
 import re
 import threading
 import uuid
@@ -26,8 +25,6 @@ from integration._support.upstream import _aws_event_frame, aws_event_stream_fra
 from integration._support.wire import Reply, Request, Wire, wire_server
 from pydantic import JsonValue, TypeAdapter
 
-from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with
-
 pytestmark: Final = pytest.mark.shares_worker_state
 
 _MODEL_ID: Final = "global.moonshotai.kimi-k3"
@@ -43,7 +40,6 @@ _REJECTION: Final = "structured output schema uses unsupported regex negative lo
 _THROTTLED: Final = "Too many requests, please wait before trying again"
 _UNKNOWN_TYPE: Final = "somethingBedrockAddedLater"
 _UNKNOWN_ONLY: Final = "none of its 1 events carried a known event type"
-_SIGNING_KEY: Final = os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt")
 _JSON_HEADERS: Final = MappingProxyType({":content-type": "application/json", ":message-type": "event"})
 _USAGE: Final[dict[str, JsonValue]] = {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}}
 _RESPONSE: Final = json.dumps(
@@ -277,17 +273,9 @@ def _responses_text(events: Iterable[dict[str, JsonValue]]) -> str:
     )
 
 
-def _inner_response_id(identity: str) -> str:
-    managed: Final = decrypt_if_encrypted_with(identity.removeprefix("resp_"), _SIGNING_KEY)
-    assert managed is not None, identity
-    issued: Final = managed.split(";", 1)[0].rsplit("response_id:", 1)[1]
-    decoded: Final = base64.b64decode(issued.removeprefix("resp_")).decode()
-    return decoded.rsplit("response_id:", 1)[1]
-
-
 def _completed_response_id(events: Iterable[dict[str, JsonValue]]) -> str:
     (completed,) = tuple(event for event in events if event.get("type") == "response.completed")
-    return _inner_response_id(str(_JSON.validate_python(completed["response"])["id"]))
+    return str(_JSON.validate_python(completed["response"])["id"])
 
 
 def _only_received(wire: Wire) -> tuple[str, dict[str, JsonValue]]:

tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py

@@ -456,7 +456,7 @@ def _inner_response_id(identity: str) -> str:
 
 def _completed_response_id(events: Iterable[dict[str, JsonValue]]) -> str:
     (completed,) = tuple(event for event in events if event.get("type") == "response.completed")
-    return _inner_response_id(str(_JSON.validate_python(completed["response"])["id"]))
+    return str(_JSON.validate_python(completed["response"])["id"])
 
 
 def _chat_text(body: Mapping[str, JsonValue]) -> str:
@@ -911,7 +911,7 @@ async def test_k06_openai_sdk_async_responses_stream_reaches_kimi_without_a_cach
     received: Final = _only_received(rig.wire)
     assert received.model == _KIMI_US and received.streaming, received
     assert received.cache_points == 0, received.body
-    _success_row(_inner_response_id(completed[0].id))
+    _success_row(completed[0].id)
 
 
 _LOCATIONS: Final = (

tests/integration/providers/test_gemini_thinking_replay_chaos.py

@@ -233,6 +233,18 @@ def _assert_answered_with_its_own_marker(served: _Served) -> None:
     assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text
 
 
+def _spend_id(served: _Served) -> str:
+    if served.call.endpoint == "responses" and served.call.stream:
+        payloads: Final = (
+            _JSON_OBJECT.validate_json(line.removeprefix("data: "))
+            for line in served.text.splitlines()
+            if line.startswith("data: {")
+        )
+        (completed,) = tuple(event for event in payloads if event.get("type") == "response.completed")
+        return str(_JSON_OBJECT.validate_python(completed["response"])["id"])
+    return _response_id(served.call.marker)
+
+
 async def test_concurrent_thinking_replays_across_endpoints_reach_gemini_signature_free(gateway: Gateway) -> None:
     calls: Final = _calls(30, ("chat", "messages", "responses"), lambda index: index % 2 == 0)
     with wire_server(_echo) as wire, gateway.scenario() as scenario:
@@ -244,7 +256,7 @@ async def test_concurrent_thinking_replays_across_endpoints_reach_gemini_signatu
         _assert_no_bleed(wire.drain(), frozenset(call.marker for call in calls))
         landed: Final = _spend_request_ids(model, 30)
         assert landed >= {
-            _response_id(call.marker) for call in calls if not (call.endpoint == "messages" and call.stream)
+            _spend_id(item) for item in served if not (item.call.endpoint == "messages" and item.call.stream)
         }
 
 

tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py

@@ -301,6 +301,18 @@ def _assert_answered_with_its_own_marker(served: _Served) -> None:
     assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text
 
 
+def _spend_id(served: _Served) -> str:
+    if served.call.endpoint == "responses" and served.call.stream:
+        payloads: Final = (
+            _JSON_OBJECT.validate_json(line.removeprefix("data: "))
+            for line in served.text.splitlines()
+            if line.startswith("data: {")
+        )
+        (completed,) = tuple(event for event in payloads if event.get("type") == "response.completed")
+        return str(_JSON_OBJECT.validate_python(completed["response"])["id"])
+    return _response_id(served.call.marker)
+
+
 async def test_a_concurrent_mixed_burst_reaches_gemini_with_each_tool_result_in_its_own_form(gateway: Gateway) -> None:
     calls: Final = _calls(
         30, ("chat", "messages", "responses"), lambda index: index % 2 == 0, lambda index: (index // 2) % 2 == 0
@@ -314,7 +326,7 @@ async def test_a_concurrent_mixed_burst_reaches_gemini_with_each_tool_result_in_
         _assert_each_in_its_own_form(wire.drain(), calls)
         landed: Final = _spend_rows_by_model(model, 30)
         assert landed >= {
-            _response_id(call.marker) for call in calls if not (call.endpoint == "messages" and call.stream)
+            _spend_id(item) for item in served if not (item.call.endpoint == "messages" and item.call.stream)
         }