sse.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. from typing import Annotated, Any
  2. from annotated_doc import Doc
  3. from pydantic import AfterValidator, BaseModel, Field, model_validator
  4. from starlette.responses import StreamingResponse
  5. # Canonical SSE event schema matching the OpenAPI 3.2 spec
  6. # (Section 4.14.4 "Special Considerations for Server-Sent Events")
  7. _SSE_EVENT_SCHEMA: dict[str, Any] = {
  8. "type": "object",
  9. "properties": {
  10. "data": {"type": "string"},
  11. "event": {"type": "string"},
  12. "id": {"type": "string"},
  13. "retry": {"type": "integer", "minimum": 0},
  14. },
  15. }
  16. class EventSourceResponse(StreamingResponse):
  17. """Streaming response with `text/event-stream` media type.
  18. Use as `response_class=EventSourceResponse` on a *path operation* that uses `yield`
  19. to enable Server Sent Events (SSE) responses.
  20. Works with **any HTTP method** (`GET`, `POST`, etc.), which makes it compatible
  21. with protocols like MCP that stream SSE over `POST`.
  22. The actual encoding logic lives in the FastAPI routing layer. This class
  23. serves mainly as a marker and sets the correct `Content-Type`.
  24. """
  25. media_type = "text/event-stream"
  26. def _check_id_no_null(v: str | None) -> str | None:
  27. if v is not None and "\0" in v:
  28. raise ValueError("SSE 'id' must not contain null characters")
  29. return v
  30. class ServerSentEvent(BaseModel):
  31. """Represents a single Server-Sent Event.
  32. When `yield`ed from a *path operation function* that uses
  33. `response_class=EventSourceResponse`, each `ServerSentEvent` is encoded
  34. into the [SSE wire format](https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream)
  35. (`text/event-stream`).
  36. If you yield a plain object (dict, Pydantic model, etc.) instead, it is
  37. automatically JSON-encoded and sent as the `data:` field.
  38. All `data` values **including plain strings** are JSON-serialized.
  39. For example, `data="hello"` produces `data: "hello"` on the wire (with
  40. quotes).
  41. """
  42. data: Annotated[
  43. Any,
  44. Doc(
  45. """
  46. The event payload.
  47. Can be any JSON-serializable value: a Pydantic model, dict, list,
  48. string, number, etc. It is **always** serialized to JSON: strings
  49. are quoted (`"hello"` becomes `data: "hello"` on the wire).
  50. Mutually exclusive with `raw_data`.
  51. """
  52. ),
  53. ] = None
  54. raw_data: Annotated[
  55. str | None,
  56. Doc(
  57. """
  58. Raw string to send as the `data:` field **without** JSON encoding.
  59. Use this when you need to send pre-formatted text, HTML fragments,
  60. CSV lines, or any non-JSON payload. The string is placed directly
  61. into the `data:` field as-is.
  62. Mutually exclusive with `data`.
  63. """
  64. ),
  65. ] = None
  66. event: Annotated[
  67. str | None,
  68. Doc(
  69. """
  70. Optional event type name.
  71. Maps to `addEventListener(event, ...)` on the browser. When omitted,
  72. the browser dispatches on the generic `message` event.
  73. """
  74. ),
  75. ] = None
  76. id: Annotated[
  77. str | None,
  78. AfterValidator(_check_id_no_null),
  79. Doc(
  80. """
  81. Optional event ID.
  82. The browser sends this value back as the `Last-Event-ID` header on
  83. automatic reconnection. **Must not contain null (`\\0`) characters.**
  84. """
  85. ),
  86. ] = None
  87. retry: Annotated[
  88. int | None,
  89. Field(ge=0),
  90. Doc(
  91. """
  92. Optional reconnection time in **milliseconds**.
  93. Tells the browser how long to wait before reconnecting after the
  94. connection is lost. Must be a non-negative integer.
  95. """
  96. ),
  97. ] = None
  98. comment: Annotated[
  99. str | None,
  100. Doc(
  101. """
  102. Optional comment line(s).
  103. Comment lines start with `:` in the SSE wire format and are ignored by
  104. `EventSource` clients. Useful for keep-alive pings to prevent
  105. proxy/load-balancer timeouts.
  106. """
  107. ),
  108. ] = None
  109. @model_validator(mode="after")
  110. def _check_data_exclusive(self) -> "ServerSentEvent":
  111. if self.data is not None and self.raw_data is not None:
  112. raise ValueError(
  113. "Cannot set both 'data' and 'raw_data' on the same "
  114. "ServerSentEvent. Use 'data' for JSON-serialized payloads "
  115. "or 'raw_data' for pre-formatted strings."
  116. )
  117. return self
  118. def format_sse_event(
  119. *,
  120. data_str: Annotated[
  121. str | None,
  122. Doc(
  123. """
  124. Pre-serialized data string to use as the `data:` field.
  125. """
  126. ),
  127. ] = None,
  128. event: Annotated[
  129. str | None,
  130. Doc(
  131. """
  132. Optional event type name (`event:` field).
  133. """
  134. ),
  135. ] = None,
  136. id: Annotated[
  137. str | None,
  138. Doc(
  139. """
  140. Optional event ID (`id:` field).
  141. """
  142. ),
  143. ] = None,
  144. retry: Annotated[
  145. int | None,
  146. Doc(
  147. """
  148. Optional reconnection time in milliseconds (`retry:` field).
  149. """
  150. ),
  151. ] = None,
  152. comment: Annotated[
  153. str | None,
  154. Doc(
  155. """
  156. Optional comment line(s) (`:` prefix).
  157. """
  158. ),
  159. ] = None,
  160. ) -> bytes:
  161. """Build SSE wire-format bytes from **pre-serialized** data.
  162. The result always ends with `\n\n` (the event terminator).
  163. """
  164. lines: list[str] = []
  165. if comment is not None:
  166. for line in comment.splitlines():
  167. lines.append(f": {line}")
  168. if event is not None:
  169. lines.append(f"event: {event}")
  170. if data_str is not None:
  171. for line in data_str.splitlines():
  172. lines.append(f"data: {line}")
  173. if id is not None:
  174. lines.append(f"id: {id}")
  175. if retry is not None:
  176. lines.append(f"retry: {retry}")
  177. lines.append("")
  178. lines.append("")
  179. return "\n".join(lines).encode("utf-8")
  180. # Keep-alive comment, per the SSE spec recommendation
  181. KEEPALIVE_COMMENT = b": ping\n\n"
  182. # Seconds between keep-alive pings when a generator is idle.
  183. # Private but importable so tests can monkeypatch it.
  184. _PING_INTERVAL: float = 15.0