Skip to content

Commit 157e7a0

Browse files
Sameerlitecursoragentyassin-berriai
authored
fix(containers): record ownership for service-account keys + fix Prisma Json serialization (BerriAI#28990)
* fix(containers): record ownership for service-account keys + fix Prisma Json field serialization - Track containers created implicitly via /v1/responses by extracting container IDs from the response output and calling record_container_owner for each one, so subsequent file-API calls from the same service account pass ownership checks. - Fix DataError: Prisma Python requires Json fields to be JSON strings; serialize file_object with json.dumps() before insert/update in LiteLLM_ManagedObjectTable. - Add collect_container_ids_from_responses_response utility to responses/utils.py that walks all output item shapes (code_interpreter_call, message annotations). - Tests: two new cases covering the responses-tracking path and the end-to-end record-then-assert flow for service accounts with team scope. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(containers): swallow all exceptions in ownership hook; tighten file_object_json type to str Co-authored-by: Cursor <cursoragent@cursor.com> * fix(containers): parse file_object JSON string in existing ownership test Co-authored-by: Cursor <cursoragent@cursor.com> * fix: container ownership recording bugs - Remove unreachable _aresponses_websocket from route_type set in base_process_llm_request; the WebSocket endpoint never flows through base_process_llm_request, so this branch was dead code that gave a false impression of coverage. - Drop the HTTPException re-raise in record_container_owners_from_responses_response so per-container failures (including HTTP 403/500 from conflicting ownership rows) no longer abort the batch and skip recording for the remaining container IDs in the same response. Co-authored-by: Yassin Kortam <yassin@berri.ai> * 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 <yassin@berri.ai> --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Yassin Kortam <yassin@berri.ai>
1 parent 9cac047 commit 157e7a0

4 files changed

Lines changed: 444 additions & 4 deletions

File tree

‎litellm/proxy/common_request_processing.py‎

Lines changed: 100 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",
@@ -1483,8 +1498,93 @@ async def _on_deferred_stream_complete(
14831498

14841499
await check_response_size_is_safe(response=response)
14851500

1501+
if route_type in {"aresponses", "aget_responses"}:
1502+
await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed(
1503+
response=response,
1504+
user_api_key_dict=user_api_key_dict,
1505+
)
1506+
14861507
return response
14871508

1509+
@staticmethod
1510+
async def _record_container_owners_from_responses_if_needed(
1511+
response: Any,
1512+
user_api_key_dict: UserAPIKeyAuth,
1513+
) -> None:
1514+
"""Register code-interpreter containers so follow-up file APIs pass ownership checks."""
1515+
from litellm.proxy.container_endpoints.ownership import (
1516+
record_container_owners_from_responses_response,
1517+
)
1518+
1519+
if response is None:
1520+
return
1521+
1522+
try:
1523+
await record_container_owners_from_responses_response(
1524+
response=response,
1525+
user_api_key_dict=user_api_key_dict,
1526+
)
1527+
except Exception as e:
1528+
verbose_proxy_logger.exception(
1529+
"Container ownership recording failed after responses call: %s",
1530+
e,
1531+
)
1532+
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+
14881588
async def base_passthrough_process_llm_request(
14891589
self,
14901590
request: Request,

‎litellm/proxy/container_endpoints/ownership.py‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,58 @@ async def _get_prisma_client():
117117
return prisma_client
118118

119119

120+
def _custom_llm_provider_from_responses_response(
121+
response: Any,
122+
default: str = "openai",
123+
) -> str:
124+
hidden_params: Dict[str, Any] = {}
125+
if isinstance(response, dict):
126+
hidden_params = response.get("_hidden_params") or {}
127+
else:
128+
hidden_params = getattr(response, "_hidden_params", None) or {}
129+
130+
provider = hidden_params.get("custom_llm_provider")
131+
if isinstance(provider, str) and provider:
132+
return provider
133+
return default
134+
135+
136+
async def record_container_owners_from_responses_response(
137+
response: Any,
138+
user_api_key_dict: UserAPIKeyAuth,
139+
custom_llm_provider: Optional[str] = None,
140+
) -> None:
141+
"""Track containers created implicitly by code interpreter in /v1/responses."""
142+
container_ids = (
143+
ResponsesAPIRequestUtils.collect_container_ids_from_responses_response(response)
144+
)
145+
if not container_ids:
146+
return
147+
148+
resolved_provider = (
149+
custom_llm_provider or _custom_llm_provider_from_responses_response(response)
150+
)
151+
152+
for container_id in container_ids:
153+
try:
154+
await record_container_owner(
155+
response={"id": container_id, "object": "container"},
156+
user_api_key_dict=user_api_key_dict,
157+
custom_llm_provider=resolved_provider,
158+
)
159+
except Exception as e:
160+
# Per-container errors (including ``HTTPException`` from
161+
# conflicting/forbidden ownership rows) must not abort the
162+
# batch — other containers in the same response should still
163+
# get recorded so their follow-up file API calls don't 403.
164+
verbose_proxy_logger.exception(
165+
"Failed to record container ownership from responses output "
166+
"for container_id=%s: %s",
167+
container_id,
168+
e,
169+
)
170+
171+
120172
async def record_container_owner(
121173
response: Any,
122174
user_api_key_dict: UserAPIKeyAuth,
@@ -151,6 +203,8 @@ async def record_container_owner(
151203
file_object = _dump_response(response)
152204
file_object["custom_llm_provider"] = resolved_provider
153205
file_object["provider_container_id"] = original_container_id
206+
# Prisma Python requires Json fields to be serialized as a JSON string.
207+
file_object_json: str = json.dumps(file_object)
154208

155209
prisma_client = await _get_prisma_client()
156210
if prisma_client is None:
@@ -172,7 +226,7 @@ async def record_container_owner(
172226
where={"model_object_id": model_object_id},
173227
data={
174228
"unified_object_id": container_id,
175-
"file_object": file_object,
229+
"file_object": file_object_json,
176230
"updated_by": owner,
177231
},
178232
)
@@ -181,7 +235,7 @@ async def record_container_owner(
181235
data={
182236
"unified_object_id": container_id,
183237
"model_object_id": model_object_id,
184-
"file_object": file_object,
238+
"file_object": file_object_json,
185239
"file_purpose": CONTAINER_OBJECT_PURPOSE,
186240
"created_by": owner,
187241
"updated_by": owner,

‎litellm/responses/utils.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -738,6 +738,98 @@ def _maybe_encode(container_id: str) -> Optional[str]:
738738
model_id,
739739
)
740740

741+
@staticmethod
742+
def _collect_container_ids_from_annotations(
743+
annotations: Any,
744+
collected: set[str],
745+
) -> None:
746+
if not annotations or not isinstance(annotations, list):
747+
return
748+
for ann in annotations:
749+
ResponsesAPIRequestUtils._collect_container_ids_from_output_item(
750+
ann, collected
751+
)
752+
753+
@staticmethod
754+
def _collect_container_ids_from_message_content(
755+
content: Any,
756+
collected: set[str],
757+
) -> None:
758+
if not content:
759+
return
760+
if isinstance(content, list):
761+
for part in content:
762+
if isinstance(part, dict):
763+
ResponsesAPIRequestUtils._collect_container_ids_from_annotations(
764+
part.get("annotations"),
765+
collected,
766+
)
767+
else:
768+
ResponsesAPIRequestUtils._collect_container_ids_from_annotations(
769+
getattr(part, "annotations", None),
770+
collected,
771+
)
772+
773+
@staticmethod
774+
def _collect_container_ids_from_output_item(
775+
item: Any,
776+
collected: set[str],
777+
) -> None:
778+
"""Collect managed or raw ``container_id`` values from one output item."""
779+
if item is None:
780+
return
781+
782+
if isinstance(item, dict):
783+
cid = item.get("container_id")
784+
if isinstance(cid, str) and cid:
785+
collected.add(cid)
786+
nested = item.get("code_interpreter_call")
787+
if isinstance(nested, dict):
788+
nc = nested.get("container_id")
789+
if isinstance(nc, str) and nc:
790+
collected.add(nc)
791+
if item.get("type") == "message":
792+
ResponsesAPIRequestUtils._collect_container_ids_from_message_content(
793+
item.get("content"),
794+
collected,
795+
)
796+
return
797+
798+
cid_attr = getattr(item, "container_id", None)
799+
if isinstance(cid_attr, str) and cid_attr:
800+
collected.add(cid_attr)
801+
802+
nested_obj = getattr(item, "code_interpreter_call", None)
803+
if nested_obj is not None:
804+
ResponsesAPIRequestUtils._collect_container_ids_from_output_item(
805+
nested_obj, collected
806+
)
807+
808+
if getattr(item, "type", None) == "message":
809+
ResponsesAPIRequestUtils._collect_container_ids_from_message_content(
810+
getattr(item, "content", None),
811+
collected,
812+
)
813+
814+
@staticmethod
815+
def collect_container_ids_from_responses_response(response: Any) -> list[str]:
816+
"""Return unique container IDs referenced in a Responses API payload."""
817+
if response is None:
818+
return []
819+
820+
if isinstance(response, dict):
821+
output = response.get("output", [])
822+
else:
823+
output = getattr(response, "output", []) or []
824+
825+
collected: set[str] = set()
826+
if output:
827+
for item in output:
828+
ResponsesAPIRequestUtils._collect_container_ids_from_output_item(
829+
item, collected
830+
)
831+
return list(collected)
832+
741833
@staticmethod
742834
def _update_container_ids_in_response(
743835
responses_api_response: Union[ResponsesAPIResponse, Dict[str, Any]],

0 commit comments

Comments
 (0)