llama-stack/llama_stack/distribution
Ben Browning dc46725f56
fix: properly handle streaming client disconnects (#2000)
# What does this PR do?

Previously, when a streaming client would disconnect before we were
finished streaming the entire response, an error like the below would
get raised from the `sse_generator` function in
`llama_stack/distribution/server/server.py`:

```
AttributeError: 'coroutine' object has no attribute 'aclose'. Did you mean: 'close'?
```

This was because we were calling `aclose` on a coroutine instead of the
awaited value from that coroutine. This change fixes that, so that we
save off the awaited value and then can call `aclose` on it if we
encounter an `asyncio.CancelledError`, like we see when a client
disconnects before we're finished streaming.

The other changes in here are to add a simple set of tests for the happy
path of our SSE streaming and this client disconnect path.

That unfortunately requires adding one more dependency into our unit
test section of pyproject.toml since `server.py` requires loading some
of the telemetry code for me to test this functionality.

## Test Plan

I wrote the tests in `tests/unit/server/test_sse.py` first, verified the
client disconnected test failed before my change, and that it passed
afterwards.

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

Signed-off-by: Ben Browning <bbrownin@redhat.com>
2025-04-23 15:44:28 +02:00
..
routers fix: 100% OpenAI API verification for together and fireworks (#1946) 2025-04-14 08:56:29 -07:00
server fix: properly handle streaming client disconnects (#2000) 2025-04-23 15:44:28 +02:00
store fix: handle registry errors gracefully (#1732) 2025-03-20 15:24:07 -07:00
ui feat: Hide tool output under an expander in Playground UI (#2003) 2025-04-23 15:32:12 +02:00
utils feat: add health to all providers through providers endpoint (#1418) 2025-04-14 11:59:36 +02:00
__init__.py API Updates (#73) 2024-09-17 19:51:35 -07:00
access_control.py feat: make sure agent sessions are under access control (#1737) 2025-03-21 07:31:16 -07:00
build.py feat: allow building distro with external providers (#1967) 2025-04-18 17:18:28 +02:00
build_conda_env.sh chore: remove straggler references to llama-models (#1345) 2025-03-01 14:26:03 -08:00
build_container.sh feat: allow building distro with external providers (#1967) 2025-04-18 17:18:28 +02:00
build_venv.sh chore: remove straggler references to llama-models (#1345) 2025-03-01 14:26:03 -08:00
client.py chore: move all Llama Stack types from llama-models to llama-stack (#1098) 2025-02-14 09:10:59 -08:00
common.sh fix: Fixing some small issues with the build scripts (#1132) 2025-02-19 22:20:49 -08:00
configure.py feat: add provider API for listing and inspecting provider info (#1429) 2025-03-13 15:07:21 -07:00
datatypes.py feat: allow building distro with external providers (#1967) 2025-04-18 17:18:28 +02:00
distribution.py feat: allow building distro with external providers (#1967) 2025-04-18 17:18:28 +02:00
inspect.py feat: add health to all providers through providers endpoint (#1418) 2025-04-14 11:59:36 +02:00
library_client.py feat: add health to all providers through providers endpoint (#1418) 2025-04-14 11:59:36 +02:00
providers.py feat: add health to all providers through providers endpoint (#1418) 2025-04-14 11:59:36 +02:00
request_headers.py feat(server): add attribute based access control for resources (#1703) 2025-03-19 21:28:52 -07:00
resolver.py feat: add health to all providers through providers endpoint (#1418) 2025-04-14 11:59:36 +02:00
stack.py feat: add health to all providers through providers endpoint (#1418) 2025-04-14 11:59:36 +02:00
start_stack.sh docs: Update docs and fix warning in start-stack.sh (#1937) 2025-04-11 16:26:17 -07:00