Python: Add AG-UI FastAPI SSE keepalive support#6980
Conversation
Key decisions: add keepalive_seconds as endpoint-owned FastAPI registration configuration with default 15, accept None as the explicit off switch, validate that non-None values are greater than zero during route registration, and keep agent/workflow runner constructors unchanged. Declare sse-starlette>=3.4.5,<4 as a direct AG-UI dependency without changing the existing StreamingResponse path in this slice. Files changed: packages/ag-ui/agent_framework_ag_ui/_endpoint.py adds validation and the public endpoint parameter; packages/ag-ui/tests/ag_ui/test_endpoint.py covers default, supported runner shapes, endpoint ownership, and invalid intervals; packages/ag-ui/pyproject.toml and uv.lock add the direct sse-starlette dependency metadata. Verification: uv run pytest focused keepalive endpoint tests -q; uv run poe test -P ag-ui; uv run poe syntax -P ag-ui -C; uv run poe pyright -P ag-ui; uv run poe validate-dependency-bounds-test -P ag-ui; git diff --check; git diff --cached --check. Also ran validate-dependency-bounds-project --mode both --package ag-ui --dependency sse-starlette; it completed but broadened the lower bound, so the issue-required >=3.4.5,<4 contract was restored and re-locked. Notes: uv run poe typing -P ag-ui and uv run poe check -P ag-ui currently fail in mypy before checking project files because .venv/lib/python3.13/site-packages/numpy/__init__.pyi uses type-statement syntax while the test mypy profile targets Python 3.11. Local issue file was moved to issues/done/ but not staged.
Key decisions: switch only enabled AG-UI FastAPI endpoint keepalive responses to EventSourceResponse, keep encoded AG-UI SSE frames as bytes on that path to avoid double encoding, and emit the fixed static SSE comment ': keepalive' while preserving existing SSE headers. Files changed: packages/ag-ui/agent_framework_ag_ui/_endpoint.py adds the EventSourceResponse enabled path and static comment factory; packages/ag-ui/tests/ag_ui/test_endpoint.py adds an endpoint test for a long output-silent gap, keepalive comments, headers, valid data frames, and no data: data: double encoding. Verification: uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_keepalive_enabled_emits_static_comment_during_silent_gap -q; focused endpoint pytest selection; uv run poe test -P ag-ui; uv run poe syntax -P ag-ui -C; uv run poe pyright -P ag-ui; git diff --check; git diff --cached --check. Notes: uv run poe check -P ag-ui still fails in the test-typing mypy phase before project files are checked because .venv/lib/python3.13/site-packages/numpy/__init__.pyi uses type-statement syntax while the mypy test profile targets Python 3.11. Local PRD/Ralph/context artifacts were not staged.
Key decisions: cover keepalive_seconds=None at the FastAPI endpoint seam and assert it preserves the legacy StreamingResponse SSE shape without emitting transport keepalive comments. Files changed: packages/ag-ui/tests/ag_ui/test_endpoint.py adds disabled keepalive endpoint coverage for headers, valid AG-UI data frames, no keepalive comments, and no data: data: double encoding. Verification: uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_keepalive_disabled_preserves_streaming_response_shape packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_keepalive_enabled_emits_static_comment_during_silent_gap -q; uv run poe test -P ag-ui; uv run poe syntax -P ag-ui -C; uv run poe pyright -P ag-ui; uv run poe test-typing -P ag-ui --checker pyright; git diff --check. Notes: no production code changes were needed because the endpoint already branches to the existing StreamingResponse path when keepalive_seconds=None. Local PRD/Ralph/context artifacts were not staged.
Key decisions: document keepalive_seconds at the FastAPI endpoint seam as a default-enabled transport keepalive with None as the off switch, and record that SSE keepalive emits comments without changing AG-UI events or adding protocol heartbeat events. Files changed: packages/ag-ui/agent_framework_ag_ui/_endpoint.py expands the public endpoint docstring; packages/ag-ui/AGENTS.md records endpoint-owned keepalive guidance; packages/ag-ui/tests/ag_ui/test_endpoint.py adds a public docstring regression. Verification: uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_add_endpoint_docstring_describes_keepalive_transport_behavior -q failed before the doc update; focused keepalive endpoint tests passed; uv run poe test -P ag-ui; uv run poe syntax -P ag-ui -C; uv run poe pyright -P ag-ui; uv run poe test-typing -P ag-ui --checker pyright; uv run python scripts/check_md_code_blocks.py packages/ag-ui/AGENTS.md; git diff --check. Notes: no standalone docs page was added. Local issue bookkeeping was moved to issues/done but not staged; local PRD and Ralph/context artifacts remain unstaged.
giles17
added
documentation
There was a problem hiding this comment.
Pull request overview
This PR improves the robustness of the Python AG-UI FastAPI SSE transport by adding endpoint-owned keepalive comments to prevent idle-timeout disconnects during output-silent portions of long runs, without introducing new AG-UI protocol events or changing runner behavior.
Changes:
- Adds a
keepalive_secondsparameter toadd_agent_framework_fastapi_endpoint(default15), withNonepreserving the existing non-keepaliveStreamingResponsepath and validation rejecting non-positive intervals. - Switches keepalive-enabled responses to
sse-starlette’sEventSourceResponse(ping comments) while continuing to stream already-encoded AG-UI SSE frames unchanged. - Updates package dependencies (adds
sse-starlette, tightens FastAPI lower bound) and adds tests/documentation to codify the transport boundary and behavior.
Reviewed changes
Copilot reviewed 4 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py | Adds keepalive_seconds and uses EventSourceResponse ping comments for keepalive-enabled SSE streams while preserving the existing streaming path when disabled. |
| python/packages/ag-ui/tests/ag_ui/test_endpoint.py | Adds unit tests verifying keepalive defaults, docstring guidance, endpoint-only ownership, and comment-only keepalive behavior without altering AG-UI frames. |
| python/packages/ag-ui/pyproject.toml | Adds sse-starlette runtime dependency and raises FastAPI lower bound to a compatible range. |
| python/packages/ag-ui/AGENTS.md | Documents keepalive as endpoint-owned SSE comment behavior (not protocol events / not runner configuration). |
| python/uv.lock | Locks the new sse-starlette dependency and updates the FastAPI version range accordingly. |
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Automated Code Review
Reviewers: 5 | Confidence: 90%
✓ Correctness
The implementation is correct. The keepalive feature properly uses sse-starlette's EventSourceResponse for ping-enabled streams while preserving the existing StreamingResponse path when disabled. The generator correctly yields pre-encoded bytes that pass through ensure_bytes unchanged, avoiding double-framing. Validation correctly rejects non-positive intervals. One minor typing concern:
cast(int, keepalive_seconds)relies on sse-starlette's property setter accepting floats despite the init signature declaringping: Optional[int].
✓ Security Reliability
The keepalive implementation is well-designed with proper input validation, clear separation of transport and protocol concerns, and comprehensive tests. The only notable finding is the use of
typing.cast(int, keepalive_seconds)which does not perform runtime conversion — it passes a float to a parameter typed asint. This works today because sse-starlette internally usesasyncio.sleep(self.ping)which accepts floats, but it's a type lie that relies on an undocumented implementation detail of the library.
✓ Test Coverage
The PR has comprehensive test coverage for the new keepalive feature. Tests cover: the enabled path verifying SSE comments are emitted during silent gaps, the disabled path verifying no comments appear, validation rejecting non-positive intervals (parametrized with 0, -1, -0.5), default value verification, API surface boundary checks (endpoint-owned, not runner-owned), all runner shapes (raw/wrapped agent/workflow), and docstring contract tests. The only minor gap is no explicit test for the error paths (encode_error/stream_error) when keepalive is enabled, but these paths are pre-existing and the
prepare_framewrapper is trivially correct. No blocking issues found.
✓ Failure Modes
The PR is well-designed with no significant failure modes. The key design choices — encoding to bytes to bypass sse-starlette's
ensure_bytesstr-wrapping path, validation rejecting non-positive values, and the None fallback to StreamingResponse — are all correct. Thecast(int, keepalive_seconds)on line 241 is a type-only no-op at runtime;sse-starlette'sping_intervalproperty setter acceptsUnion[int, float]and the value flows intoanyio.sleep()which handles floats correctly. No silent failures, lost errors, or operational issues found.
✗ Design Approach
The keepalive framing itself looks consistent with the new tests, but the current integration pulls in
se-starlettetoo early. Because_endpoint.pyimportsEventSourceResponseat module load time, the package now triggersse-starlette's global Uvicorn shutdown monkey-patch even for callers that setkeepalive_seconds=None, which undercuts the PR's stated disabled-path compatibility and the new AGENTS guidance that keepalive stay endpoint-owned transport behavior.
Automated review by moonbox3's agents
Motivation & Context
Long output-silent AG-UI FastAPI streams can be killed by clients or intermediaries when no bytes are sent before their idle timeout. This change adds endpoint-owned SSE keepalive comments so the transport remains active without introducing protocol-level heartbeat events or changing runner behavior.
Description & Review Guide
keepalive_secondstoadd_agent_framework_fastapi_endpoint, defaulting to 15 seconds.StreamingResponsepath whenkeepalive_seconds=Noneand rejects non-positive keepalive intervals.sse-starletteEventSourceResponsefor keepalive-enabled streams while yielding encoded AG-UI frames unchanged.>=0.121.0so the declaredsse-starlettedependency has a satisfiable Starlette range.keepalive_seconds=None.Related Issue
Fixes #6941
An existing open PR also references this issue: #6965 (
kartikmadan11:fix/6941-sse-keepalive). This PR is frommoonbox3:6941-fixand includes the dependency-bound correction needed to keepsse-starlette>=3.4.5satisfiable with the declared FastAPI range.Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.