scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 98

fix(proxy): record streamed /v1/responses container ownership before the response.completed frame (#43140)

Browse source 
* test(e2e): cover Azure code_interpreter container files by native id with a service-account key

* test(e2e): require the code_interpreter tool, skip at collection, and scope the container call timeout

* fix(e2e): fail the containers suite when the Azure credentials are missing instead of skipping

* fix(proxy): record streamed /v1/responses container ownership before the response.completed frame

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-25 13:45:43 -07:00 • committed by GitHub
parent 6b7688869e
commit 601c75a475
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 144 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

87
litellm/proxy/common_request_processing.py
View file

@ -2781,15 +2781,6 @@ class ProxyBaseLLMRequestProcessing:
request=request,
)
if route_type == "aresponses":
# Streaming /v1/responses returns here without
# reaching the non-streaming ownership tail below.
# Wrap the SSE generator so container ownership is
# written once the upstream iterator finishes
# assembling ``completed_response`` — otherwise
# code-interpreter containers created during the
# stream stay unregistered and follow-up file API
# calls 403. Covers the background-polling path
# too, which loops ``body_iterator`` end-to-end.
selected_data_generator = (
ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
original_stream_response=response,
@ -3011,50 +3002,50 @@ class ProxyBaseLLMRequestProcessing:
wrapped_generator: Any,
user_api_key_dict: UserAPIKeyAuth,
):
"""Forward SSE chunks, then record container ownership at stream end.
"""Forward SSE chunks and record container ownership before the terminal chunk goes out.
Streaming ``/v1/responses`` short-circuits out of
``base_process_llm_request`` before the non-streaming ownership
tail runs, so without this wrap the
``LiteLLM_ManagedObjectTable`` row for any container created
during the stream is never written and follow-up file API calls
return 403.
tail runs. The OpenAI SDK closes the connection at ``data: [DONE]``
and starlette cancels the body task on disconnect, so a write that
waits for the generator to finish never lands. The iterator sets
``completed_response`` before it hands over its terminal chunk, so
the ``LiteLLM_ManagedObjectTable`` row is written the moment it
appears, ahead of the chunk carrying ``response.completed``.
"""
try:
async for chunk in wrapped_generator:
async for chunk in wrapped_generator:
completed_obj = ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
original_stream_response
)
if completed_obj is None:
yield chunk
finally:
try:
completed_obj: Final = ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
original_stream_response
)
if completed_obj is not None:
await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed(
response=completed_obj,
user_api_key_dict=user_api_key_dict,
)
else:
# Silent skip caused #30210: the proxy's Router wrapper
# of the responses streaming iterator wasn't propagating
# ``completed_response``, so this hook recorded nothing
# and follow-up /v1/containers/<id>/files calls 403'd
# for non-admin keys with no proxy-side hint. Log a
# warning so future regressions of the same shape
# surface in operator logs.
verbose_proxy_logger.warning(
"Container ownership recording skipped on streaming "
"/v1/responses: no completed_response on stream "
"iterator %s. If this stream created any tool "
"container (e.g. code_interpreter), follow-up "
"/v1/containers/<id>/files calls will 403 for "
"non-admin keys.",
type(original_stream_response).__name__,
)
except Exception as e:
verbose_proxy_logger.exception(
"Container ownership recording failed after streaming responses call: %s",
e,
)
continue
await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed(
response=completed_obj,
user_api_key_dict=user_api_key_dict,
)
yield chunk
async for remaining_chunk in wrapped_generator:
yield remaining_chunk
return
late_completed_obj: Final = ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
original_stream_response
)
if late_completed_obj is not None:
await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed(
response=late_completed_obj,
user_api_key_dict=user_api_key_dict,
)
return
verbose_proxy_logger.warning(
"Container ownership recording skipped on streaming "
"/v1/responses: no completed_response on stream "
"iterator %s. If this stream created any tool "
"container (e.g. code_interpreter), follow-up "
"/v1/containers/<id>/files calls will 403 for "
"non-admin keys.",
type(original_stream_response).__name__,
)
async def base_passthrough_process_llm_request(
self,

1
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -85,6 +85,7 @@
- {id: llm.responses.azure_openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Responses w/ Azure OpenAI (smoke)"}
- {id: llm.responses.azure_openai.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Responses tool calls w/ Azure OpenAI"}
- {id: llm.responses.azure_openai.code_interpreter.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: code_interpreter, streaming: nonstream, assertions: [works], source: "llm_translation/test_containers_e2e.py", rationale: "An implicit code_interpreter container on an Azure deployment that carries its own api_base must serve GET /v1/containers/{id}/files/{fid}/content by its native cntr_ id to a team service-account key, the customer's shape (#27921, #28990)", fail_before_fix: proven}
- {id: llm.responses.azure_openai.code_interpreter.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: code_interpreter, streaming: stream, assertions: [works], source: "llm_translation/test_containers_e2e.py", rationale: "A container created by a streamed /v1/responses code_interpreter call must serve /v1/containers/{id}/files to the same service-account key right after the OpenAI SDK closes at [DONE]; the ownership row used to be written after the stream and the disconnect cancelled it (LIT-8612)", fail_before_fix: proven}
- {id: llm.chat_completions.together_ai.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together reasoning surfaces as reasoning_content (LIT-5960)"}
- {id: llm.chat_completions.together_ai.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together reasoning deltas stream as reasoning_content"}
- {id: llm.chat_completions.together_ai.thinking.nonstream.template_kwargs_forwarded, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [template_kwargs_forwarded], source: "llm_translation/test_together_ai_e2e.py", rationale: "chat_template_kwargs reaches Together and turns thinking off"}

3
tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md
View file

@ -48,7 +48,7 @@ most likely to silently break and the one a mock can't prove works.
|----------|---------------|-----------|------------|-------------|--------|
| Chat | live (spend suite) | live (spend suite) | gap | live | partial |
| Embeddings | live (spend suite) | n/a | n/a | live | covered |
| Responses (Azure code_interpreter container files) | live | gap | live | gap | partial |
| Responses (Azure code_interpreter container files) | live | live | live | gap | partial |
| Image / audio / rerank / realtime | - | - | - | - | gap |
## This suite's files
@ -63,6 +63,7 @@ most likely to silently break and the one a mock can't prove works.
| `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost |
| `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost |
| `test_service_account_key_reads_container_file_by_native_id` | azure responses code_interpreter, non-stream, native container id, service-account key |
| `test_service_account_key_reads_container_file_created_by_a_streamed_response` | azure responses code_interpreter, stream, native container id, service-account key, upload right after `[DONE]` |
Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is
added at runtime instead of declared in the gateway config: the test POSTs `/model/new`

43
tests/e2e/llm_translation/test_containers_e2e.py
View file

@ -31,10 +31,11 @@ A proxy whose env carries ``AZURE_API_BASE`` for the same resource masks the
second regression, since the global-credential fallback then reaches the
container anyway.
The streaming variant is not here: a streamed ``/v1/responses`` writes the
container ownership row only after the ``[DONE]`` frame, and the OpenAI SDK
closes the connection at ``[DONE]``, so the write is cancelled and every
follow-up container call 403s (LIT-8612). That cell comes with its fix.
The streaming cell repeats the flow with ``stream=True`` and uploads right
after the last event. The OpenAI SDK closes the connection at ``[DONE]``, so an
ownership row written after the stream is cancelled with the body task and every
follow-up container call 403s (LIT-8612); the row has to land before the
``response.completed`` frame goes out.
"""
from __future__ import annotations
@ -52,7 +53,7 @@ from lifecycle import ResourceManager
from management.management_client import ManagementClient, build_client
from models import KeyGenerateBody, KeyGenerateResponse, LiteLLMParamsBody, TeamNewBody, UserNewBody
from openai import OpenAI
from openai.types.responses import Response, ResponseCodeInterpreterToolCall
from openai.types.responses import Response, ResponseCodeInterpreterToolCall, ResponseCompletedEvent
from openai.types.responses.tool_param import CodeInterpreter
from proxy_client import ProxyClient
from sdk_clients import NO_PROXY_CACHE, SdkClients
@ -120,6 +121,24 @@ def _response_with_code_interpreter(client: OpenAI, model: str) -> Response:
)
def _streamed_response_with_code_interpreter(client: OpenAI, model: str) -> Response:
events: Final = tuple(
client.with_options(timeout=CODE_INTERPRETER_TIMEOUT).responses.create(
model=model,
input=PROMPT,
tools=[CODE_INTERPRETER],
tool_choice="required",
stream=True,
extra_body=NO_PROXY_CACHE,
)
)
assert events, "responses stream returned no events"
assert isinstance(events[-1], ResponseCompletedEvent), (
f"responses stream did not terminate with response.completed: {events[-1].type}"
)
return events[-1].response
def _container_id(response: Response) -> str:
calls: Final = tuple(item for item in response.output if isinstance(item, ResponseCodeInterpreterToolCall))
assert calls, f"no code_interpreter_call in the responses output: {response.output!r}"
@ -165,3 +184,17 @@ class TestAzureContainerFiles:
f"container id is not the provider's own id: {native_id}"
)
_assert_file_round_trip(client, native_id, marker)
@pytest.mark.covers("llm.responses.azure_openai.code_interpreter.stream.works")
def test_service_account_key_reads_container_file_created_by_a_streamed_response(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
marker: Final = unique_marker()
model: Final = _register_two_azure_deployments(proxy, resources, marker)
key: Final = _service_account_key(proxy, resources, build_client(proxy), marker, model)
client: Final = sdk.openai(key)
native_id: Final = _native_container_id(
_container_id(_streamed_response_with_code_interpreter(client, model))
)
resources.defer(lambda: client.containers.delete(native_id, extra_query=AZURE_PROVIDER_QUERY))
_assert_file_round_trip(client, native_id, marker)

64
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -9899,3 +9899,67 @@ class TestErrorLogCarriesCallId:
record: Final = caplog.records[-1]
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
class TestStreamingContainerOwnershipRecordedBeforeDone:
"""Regression for LIT-8612: the OpenAI SDK closes the connection at
``data: [DONE]`` and starlette cancels the body task, so an ownership row
written after the SSE generator is exhausted never lands. The row must be
written before the chunk carrying ``response.completed`` is handed to the
client."""
CHUNKS: Final = (
'data: {"type":"response.created"}\n\n',
'data: {"type":"response.output_text.delta"}\n\n',
'data: {"type":"response.completed"}\n\n',
"data: [DONE]\n\n",
)
TERMINAL_INDEX: Final = 2
@staticmethod
def _completed_event() -> SimpleNamespace:
return SimpleNamespace(
type="response.completed",
response=SimpleNamespace(
id="resp_lit8612",
output=[SimpleNamespace(type="code_interpreter_call", container_id="cntr_lit8612")],
),
)
async def _sse(self, stream: SimpleNamespace, populate_at: int) -> AsyncGenerator[str, None]:
for index, chunk in enumerate(self.CHUNKS):
if index == populate_at:
stream.completed_response = self._completed_event()
yield chunk
if populate_at == len(self.CHUNKS):
stream.completed_response = self._completed_event()
async def _await_counts_per_chunk(self, populate_at: int) -> tuple[tuple[tuple[str, int], ...], AsyncMock]:
stream: Final = SimpleNamespace(completed_response=None, _hidden_params={"custom_llm_provider": "azure"})
recorder: Final = AsyncMock(return_value=None)
with patch(
"litellm.proxy.container_endpoints.ownership.record_container_owners_from_responses_response", recorder
):
wrapped: Final = ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
original_stream_response=stream,
wrapped_generator=self._sse(stream, populate_at),
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", team_id="team-1"),
)
observed: Final = tuple([(chunk, recorder.await_count) async for chunk in wrapped])
return observed, recorder
async def test_row_is_written_before_the_terminal_chunk_reaches_the_client(self) -> None:
observed, recorder = await self._await_counts_per_chunk(populate_at=self.TERMINAL_INDEX)
assert tuple(chunk for chunk, _ in observed) == self.CHUNKS
assert tuple(count for _, count in observed) == (0, 0, 1, 1)
recorder.assert_awaited_once()
assert recorder.await_args.kwargs["response"].output[0].container_id == "cntr_lit8612"
assert recorder.await_args.kwargs["user_api_key_dict"].team_id == "team-1"
async def test_row_is_still_written_when_the_iterator_completes_only_at_exhaustion(self) -> None:
observed, recorder = await self._await_counts_per_chunk(populate_at=len(self.CHUNKS))
assert tuple(chunk for chunk, _ in observed) == self.CHUNKS
assert tuple(count for _, count in observed) == (0, 0, 0, 0)
recorder.assert_awaited_once()