Skip to content

Commit eabb5b2

Browse files
cursoragentyassin-berriai
cursoragent
and
yassin-berriai
committed
fix(containers): record ownership for streaming /v1/responses too
Streaming /v1/responses returns through the select_data_generator branch in base_process_llm_request and bypasses the non-streaming ownership tail, so code-interpreter containers created mid-stream were never written to LiteLLM_ManagedObjectTable. Follow-up file API calls would then 403. Wrap the SSE generator so container ownership is recorded once the upstream iterator finishes assembling completed_response. Also covers the background-polling path, which loops body_iterator end-to-end. Co-authored-by: Yassin Kortam <[email protected]>
1 parent 8bede7c commit eabb5b2

2 files changed

Lines changed: 169 additions & 0 deletions

File tree

‎litellm/proxy/common_request_processing.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1368,6 +1368,21 @@ async def _on_deferred_stream_complete(
13681368
user_api_key_dict=user_api_key_dict,
13691369
request_data=self.data,
13701370
)
1371+
if route_type == "aresponses":
1372+
# Streaming /v1/responses returns here without
1373+
# reaching the non-streaming ownership tail below.
1374+
# Wrap the SSE generator so container ownership is
1375+
# written once the upstream iterator finishes
1376+
# assembling ``completed_response`` — otherwise
1377+
# code-interpreter containers created during the
1378+
# stream stay unregistered and follow-up file API
1379+
# calls 403. Covers the background-polling path
1380+
# too, which loops ``body_iterator`` end-to-end.
1381+
selected_data_generator = ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
1382+
original_stream_response=response,
1383+
wrapped_generator=selected_data_generator,
1384+
user_api_key_dict=user_api_key_dict,
1385+
)
13711386
return await create_response(
13721387
generator=selected_data_generator,
13731388
media_type="text/event-stream",
@@ -1515,6 +1530,61 @@ async def _record_container_owners_from_responses_if_needed(
15151530
e,
15161531
)
15171532

1533+
@staticmethod
1534+
def _extract_completed_responses_response(stream_response: Any) -> Any:
1535+
"""Pull the assembled ``ResponsesAPIResponse`` off a streaming iterator.
1536+
1537+
``ResponsesAPIStreamingIterator`` stores the terminal stream event
1538+
(``response.completed`` / ``response.incomplete`` / ``response.failed``)
1539+
in ``completed_response``; the actual response body hangs off
1540+
that event's ``.response`` attribute. Some iterators store the
1541+
``ResponsesAPIResponse`` directly. Handle both shapes so the
1542+
container-ownership recording path can walk ``.output`` either way.
1543+
"""
1544+
completed = getattr(stream_response, "completed_response", None)
1545+
if completed is None:
1546+
return None
1547+
response_obj = getattr(completed, "response", None)
1548+
if response_obj is not None:
1549+
return response_obj
1550+
return completed
1551+
1552+
@staticmethod
1553+
async def _wrap_responses_stream_for_container_ownership(
1554+
original_stream_response: Any,
1555+
wrapped_generator: Any,
1556+
user_api_key_dict: UserAPIKeyAuth,
1557+
):
1558+
"""Forward SSE chunks, then record container ownership at stream end.
1559+
1560+
Streaming ``/v1/responses`` short-circuits out of
1561+
``base_process_llm_request`` before the non-streaming ownership
1562+
tail runs, so without this wrap the
1563+
``LiteLLM_ManagedObjectTable`` row for any container created
1564+
during the stream is never written and follow-up file API calls
1565+
return 403.
1566+
"""
1567+
try:
1568+
async for chunk in wrapped_generator:
1569+
yield chunk
1570+
finally:
1571+
try:
1572+
completed_obj = (
1573+
ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
1574+
original_stream_response
1575+
)
1576+
)
1577+
if completed_obj is not None:
1578+
await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed(
1579+
response=completed_obj,
1580+
user_api_key_dict=user_api_key_dict,
1581+
)
1582+
except Exception as e:
1583+
verbose_proxy_logger.exception(
1584+
"Container ownership recording failed after streaming responses call: %s",
1585+
e,
1586+
)
1587+
15181588
async def base_passthrough_process_llm_request(
15191589
self,
15201590
request: Request,

‎tests/test_litellm/containers/test_container_proxy_ownership.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1008,3 +1008,102 @@ async def test_service_account_can_access_container_after_responses_tracking(
10081008
)
10091009
assert original_id == "cntr_native"
10101010
assert provider == "azure"
1011+
1012+
1013+
@pytest.mark.asyncio
1014+
async def test_should_record_container_ownership_after_streaming_responses_finish(
1015+
monkeypatch,
1016+
):
1017+
"""Streaming /v1/responses calls return through the
1018+
``select_data_generator`` branch and never reach the non-streaming
1019+
container-ownership tail. The wrapper must read
1020+
``completed_response`` off the upstream iterator once iteration
1021+
finishes and write the row, otherwise code-interpreter containers
1022+
created during the stream stay unregistered and follow-up file API
1023+
calls 403.
1024+
"""
1025+
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
1026+
1027+
encoded_container_id = (
1028+
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR"
1029+
"lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl"
1030+
)
1031+
response_body = SimpleNamespace(
1032+
output=[
1033+
SimpleNamespace(
1034+
type="code_interpreter_call",
1035+
container_id=encoded_container_id,
1036+
code_interpreter_call=None,
1037+
)
1038+
]
1039+
)
1040+
stream_response = SimpleNamespace(
1041+
completed_response=SimpleNamespace(response=response_body),
1042+
_hidden_params={"custom_llm_provider": "azure"},
1043+
)
1044+
1045+
async def fake_sse_generator():
1046+
yield "data: chunk-1\n\n"
1047+
yield "data: chunk-2\n\n"
1048+
1049+
table = AsyncMock()
1050+
table.find_unique.return_value = None
1051+
prisma_client = SimpleNamespace(
1052+
db=SimpleNamespace(litellm_managedobjecttable=table)
1053+
)
1054+
monkeypatch.setattr(
1055+
ownership,
1056+
"_get_prisma_client",
1057+
AsyncMock(return_value=prisma_client),
1058+
)
1059+
auth = UserAPIKeyAuth(team_id="team-1")
1060+
1061+
wrapped = (
1062+
ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
1063+
original_stream_response=stream_response,
1064+
wrapped_generator=fake_sse_generator(),
1065+
user_api_key_dict=auth,
1066+
)
1067+
)
1068+
1069+
chunks = [chunk async for chunk in wrapped]
1070+
assert chunks == ["data: chunk-1\n\n", "data: chunk-2\n\n"]
1071+
1072+
table.create.assert_awaited_once()
1073+
created_data = table.create.await_args.kwargs["data"]
1074+
assert created_data["created_by"] == "team:team-1"
1075+
assert created_data["unified_object_id"] == encoded_container_id
1076+
1077+
1078+
@pytest.mark.asyncio
1079+
async def test_streaming_ownership_wrap_no_op_when_stream_did_not_complete(
1080+
monkeypatch,
1081+
):
1082+
"""If the stream errored before ``response.completed``,
1083+
``completed_response`` is ``None`` — we must skip the ownership
1084+
write rather than crash the response generator."""
1085+
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
1086+
1087+
stream_response = SimpleNamespace(completed_response=None)
1088+
1089+
async def fake_sse_generator():
1090+
yield "data: chunk-1\n\n"
1091+
1092+
record = AsyncMock()
1093+
monkeypatch.setattr(
1094+
ownership,
1095+
"record_container_owners_from_responses_response",
1096+
record,
1097+
)
1098+
1099+
wrapped = (
1100+
ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
1101+
original_stream_response=stream_response,
1102+
wrapped_generator=fake_sse_generator(),
1103+
user_api_key_dict=UserAPIKeyAuth(user_id="user-1"),
1104+
)
1105+
)
1106+
chunks = [chunk async for chunk in wrapped]
1107+
1108+
assert chunks == ["data: chunk-1\n\n"]
1109+
record.assert_not_awaited()

0 commit comments

Comments
 (0)