phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-28 19:08:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: Additional streaming error handling

Browse source 
This expands the `test_sse` test suite and fixes some edge cases with
bugs in our SSE error handling to ensure streaming clients always get
a proper error response.

First, we handle the case where a client disconnects before we
actually start streaming the response back. Previously we only handled
the case where a client disconnected as we were streaming the
response, but there was an edge case where a client disconnecting
before we streamed any response back did not trigger our logic to
cleanly handle that disconnect.

Second, we handle the case where an error is thrown from the server
before the actual async generator gets created from the provider. This
happens in scenarios like the newly merged OpenAI API input
validation, where we eagerly raise validation errors before returning
the async generator object that streams the responses back.

Tested via:

```
python -m pytest -s -v tests/unit/server/test_sse.py
```

Both test cases failed before, and passed afterwards. The test cases
were written based on me experimenting with actual clients that would
do bad things like randomly disconnect or send invalid input in
streaming mode and I hit these two cases, where things were
misbehaving in our error handling.

Signed-off-by: Ben Browning <bbrownin@redhat.com>
This commit is contained in:
Ben Browning 2025-04-23 13:04:16 -04:00
parent fa5dfee07b
commit cfa4b61a01
2 changed files with 41 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
llama_stack/distribution/server/server.py
View file

@ -166,14 +166,16 @@ async def maybe_await(value):
async def sse_generator(event_gen_coroutine):
event_gen = await event_gen_coroutine
event_gen = None
try:
event_gen = await event_gen_coroutine
async for item in event_gen:
yield create_sse_event(item)
await asyncio.sleep(0.01)
except asyncio.CancelledError:
logger.info("Generator cancelled")
await event_gen.aclose()
if event_gen:
await event_gen.aclose()
except Exception as e:
logger.exception("Error in sse_generator")
yield create_sse_event(