Browse Source

🐛 Fix error responses for exceptions raised before the first yield in SSE endpoints

pull/15917/head
serverestaa 3 weeks ago
parent
commit
59e23d2262
  1. 12
      docs/en/docs/tutorial/server-sent-events.md
  2. 92
      fastapi/routing.py
  3. 184
      tests/test_sse.py

12
docs/en/docs/tutorial/server-sent-events.md

@ -113,8 +113,18 @@ This is useful for protocols like [MCP](https://modelcontextprotocol.io) that st
FastAPI implements some SSE best practices out of the box.
* Send a **"keep alive" `ping` comment** every 15 seconds when there hasn't been any message, to prevent some proxies from closing the connection, as suggested in the [HTML specification: Server-Sent Events](https://html.spec.whatwg.org/multipage/server-sent-events.html#authoring-notes).
* Send a **"keep alive" `ping` comment** every 15 seconds when there hasn't been any message, to prevent some proxies from closing the connection, as suggested in the [HTML specification: Server-Sent Events](https://html.spec.whatwg.org/multipage/server-sent-events.html#authoring-notes). The pings start with the response, once the **first event** has been produced (see the note below).
* Set the `Cache-Control: no-cache` header to **prevent caching** of the stream.
* Set a special header `X-Accel-Buffering: no` to **prevent buffering** in some proxies like Nginx.
* Handle exceptions raised **before the first event**, for example an `HTTPException` from an authorization check, with the regular exception handlers, sending a proper error response instead of an empty `200` stream.
* Stop waiting for the first event if the **client disconnects**, closing the generator and releasing its resources (e.g. dependencies with `yield`).
You don't have to do anything about it, it works out of the box. 🤓
/// note
The response starts once your *path operation function* yields its **first event**, this also applies to the keep alive pings.
If your code could take a long time before producing the first event, you can yield an early `ServerSentEvent(comment="connected")` to start the response (and the pings) right away.
///

92
fastapi/routing.py

@ -38,6 +38,7 @@ from typing import (
)
import anyio
import anyio.to_thread
from annotated_doc import Doc
from anyio.abc import ObjectReceiveStream
from fastapi import params
@ -364,6 +365,18 @@ def _build_response_args(
return response_args
class _ClientDisconnectResponse(Response):
# Returned when the client disconnects before the response has
# started (e.g. while waiting for the first SSE item). Nothing has
# been sent yet and there's no one left to send a response to, so
# send nothing, the same way Starlette's StreamingResponse handles
# a disconnect before the response starts. Background tasks still
# run, as they would after a disconnect during a streaming response.
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if self.background is not None:
await self.background()
def get_request_handler(
dependant: Dependant,
body_field: ModelField | None = None,
@ -547,6 +560,83 @@ def get_request_handler(
else:
sse_aiter = iterate_in_threadpool(gen)
async def _sse_first_item() -> Any:
if dependant.is_async_gen_callable:
return await sse_aiter.__anext__()
def _next_or_stop() -> Any:
try:
return next(gen) # ty: ignore[invalid-argument-type]
except StopIteration:
raise StopAsyncIteration() from None
# A worker thread blocked inside a sync generator
# can't be cancelled, so abandon it on cancellation
# (client disconnect) instead of blocking the
# request task until the generator produces its
# first item.
return await anyio.to_thread.run_sync(
_next_or_stop, abandon_on_cancel=True
)
# Pull the first item before starting the streaming
# response, so that exceptions raised before the first
# yield (e.g. an HTTPException from auth checks)
# propagate to the regular exception handlers instead
# of aborting an already started 200 response.
#
# While waiting, also listen for client disconnect (the
# streaming response is not running yet, so Starlette's
# own disconnect listener isn't active): otherwise a
# disconnected client would leave the generator and its
# dependencies alive until the first yield.
first_item: Any = None
first_item_received = False
first_item_exc: Exception | None = None
client_disconnected = False
async with anyio.create_task_group() as first_item_tg:
async def _cancel_on_disconnect() -> None:
nonlocal client_disconnected
while True:
message = await request.receive()
if message["type"] == "http.disconnect":
client_disconnected = True
first_item_tg.cancel_scope.cancel()
return
first_item_tg.start_soon(_cancel_on_disconnect)
try:
first_item = await _sse_first_item()
first_item_received = True
except StopAsyncIteration:
pass
except Exception as exc:
first_item_exc = exc
finally:
first_item_tg.cancel_scope.cancel()
if first_item_exc is not None:
raise first_item_exc
if client_disconnected:
# Close the generator so its cleanup runs promptly
# (cancelling `__anext__` normally already unwound
# an async generator, making this a no-op).
if dependant.is_async_gen_callable:
await gen.aclose()
elif first_item_received: # pragma: no cover
# The first item arrived in the same moment as
# the disconnect; the worker thread is done, so
# the sync generator can be closed. Otherwise
# it's still running in the abandoned thread and
# is finalized when that thread finishes (best
# effort, a blocked thread can't be cancelled).
await anyio.to_thread.run_sync(gen.close)
return _ClientDisconnectResponse(
background=solved_result.background_tasks
)
@asynccontextmanager
async def _sse_producer_cm() -> AsyncIterator[
ObjectReceiveStream[bytes]
@ -571,6 +661,8 @@ def get_request_handler(
async def _producer() -> None:
async with send_stream:
if first_item_received:
await send_stream.send(_serialize_sse_item(first_item))
async for raw_item in sse_aiter:
await send_stream.send(_serialize_sse_item(raw_item))

184
tests/test_sse.py

@ -1,14 +1,18 @@
import asyncio
import threading
import time
from collections.abc import AsyncIterable, Iterable
import anyio
import anyio.to_thread
import fastapi.routing
import pytest
from fastapi import APIRouter, FastAPI
from fastapi import APIRouter, BackgroundTasks, Depends, FastAPI, HTTPException
from fastapi.responses import EventSourceResponse
from fastapi.sse import ServerSentEvent
from fastapi.testclient import TestClient
from pydantic import BaseModel
from starlette.types import Message, Scope
class Item(BaseModel):
@ -87,6 +91,30 @@ async def sse_items_raw():
yield ServerSentEvent(raw_data="cpu,87.3,1709145600", event="csv")
@app.get("/items/stream-early-error", response_class=EventSourceResponse)
async def sse_items_early_error() -> AsyncIterable[Item]:
raise HTTPException(status_code=403, detail="Not authorized")
yield items[0] # pragma: no cover
@app.get("/items/stream-early-error-sync", response_class=EventSourceResponse)
def sse_items_early_error_sync() -> Iterable[Item]:
raise HTTPException(status_code=403, detail="Not authorized")
yield items[0] # pragma: no cover
@app.get("/items/stream-empty", response_class=EventSourceResponse)
async def sse_items_empty() -> AsyncIterable[Item]:
return
yield items[0] # pragma: no cover
@app.get("/items/stream-empty-sync", response_class=EventSourceResponse)
def sse_items_empty_sync() -> Iterable[Item]:
return
yield items[0] # pragma: no cover
router = APIRouter()
@ -264,6 +292,37 @@ def test_data_and_raw_data_mutually_exclusive():
ServerSentEvent(data="json", raw_data="raw")
def test_http_exception_before_first_yield_async(client: TestClient):
"""An HTTPException raised before the first yield is handled by the
regular exception handlers instead of aborting a 200 stream."""
response = client.get("/items/stream-early-error")
assert response.status_code == 403, response.text
assert response.headers["content-type"] == "application/json"
assert response.json() == {"detail": "Not authorized"}
def test_http_exception_before_first_yield_sync(client: TestClient):
response = client.get("/items/stream-early-error-sync")
assert response.status_code == 403, response.text
assert response.headers["content-type"] == "application/json"
assert response.json() == {"detail": "Not authorized"}
def test_empty_stream(client: TestClient):
"""A generator that finishes without yielding sends an empty stream."""
response = client.get("/items/stream-empty")
assert response.status_code == 200
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
assert response.text == ""
def test_empty_stream_sync(client: TestClient):
response = client.get("/items/stream-empty-sync")
assert response.status_code == 200
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
assert response.text == ""
def test_sse_on_router_included_in_app(client: TestClient):
response = client.get("/api/events")
assert response.status_code == 200
@ -325,3 +384,126 @@ def test_no_keepalive_when_fast(client: TestClient):
assert response.status_code == 200
# KEEPALIVE_COMMENT is ": ping\n\n".
assert ": ping\n" not in response.text
# Client disconnect before the first item
def _http_scope(path: str) -> Scope:
return {
"type": "http",
"asgi": {"version": "3.0", "spec_version": "2.0"},
"http_version": "1.1",
"method": "GET",
"path": path,
"query_string": b"",
"root_path": "",
"headers": [],
"server": ("test", 80),
}
@pytest.mark.anyio
async def test_disconnect_before_first_yield_async() -> None:
"""A client disconnect while an async generator is still waiting to
produce its first item cancels the wait, runs cleanup (the generator's
finally, dependencies with yield, background tasks) within a bounded
time, and sends no response."""
disconnect_app = FastAPI()
generator_started = anyio.Event()
generator_cleaned_up = anyio.Event()
dependency_cleaned_up = anyio.Event()
background_ran = anyio.Event()
async def dep() -> AsyncIterable[str]:
try:
yield "resource"
finally:
dependency_cleaned_up.set()
@disconnect_app.get("/stream", response_class=EventSourceResponse)
async def stream(
background_tasks: BackgroundTasks, res: str = Depends(dep)
) -> AsyncIterable[Item]:
background_tasks.add_task(background_ran.set)
try:
generator_started.set()
await anyio.sleep(3600)
finally:
generator_cleaned_up.set()
yield items[0] # pragma: no cover
request_sent = False
async def receive() -> Message:
nonlocal request_sent
if not request_sent:
request_sent = True
return {"type": "http.request", "body": b"", "more_body": False}
# Disconnect once the generator is waiting for its first item.
await generator_started.wait()
return {"type": "http.disconnect"}
sent_messages: list[Message] = []
async def send(message: Message) -> None:
sent_messages.append(message) # pragma: no cover
with anyio.fail_after(5):
await disconnect_app(_http_scope("/stream"), receive, send)
assert generator_cleaned_up.is_set()
assert dependency_cleaned_up.is_set()
assert background_ran.is_set()
# Nothing was sent: the client was already gone.
assert sent_messages == []
@pytest.mark.anyio
async def test_disconnect_before_first_yield_sync() -> None:
"""A worker thread blocked inside a sync generator can't be cancelled,
but a client disconnect must still complete the request task within a
bounded time (abandoning the thread) and run dependency cleanup."""
disconnect_app = FastAPI()
generator_started = threading.Event()
release_generator = threading.Event()
generator_resumed = threading.Event()
dependency_cleaned_up = anyio.Event()
async def dep() -> AsyncIterable[str]:
try:
yield "resource"
finally:
dependency_cleaned_up.set()
@disconnect_app.get("/stream", response_class=EventSourceResponse)
def stream(res: str = Depends(dep)) -> Iterable[Item]:
generator_started.set()
release_generator.wait()
generator_resumed.set()
yield items[0]
request_sent = False
async def receive() -> Message:
nonlocal request_sent
if not request_sent:
request_sent = True
return {"type": "http.request", "body": b"", "more_body": False}
# Disconnect once the generator is blocked in the worker thread.
await anyio.to_thread.run_sync(generator_started.wait)
return {"type": "http.disconnect"}
sent_messages: list[Message] = []
async def send(message: Message) -> None:
sent_messages.append(message) # pragma: no cover
with anyio.fail_after(5):
await disconnect_app(_http_scope("/stream"), receive, send)
assert dependency_cleaned_up.is_set()
assert sent_messages == []
# Let the abandoned worker thread finish.
release_generator.set()
await anyio.to_thread.run_sync(generator_resumed.wait)

Loading…
Cancel
Save