Python: Add SSE keepalive interval to AG-UI FastAPI endpoint

[Vaibhav Patel (vaibhav-patel)]

Fixes #6611.

Problem

add_agent_framework_fastapi_endpoint serves the AG-UI event stream as a bare
StreamingResponse(media_type="text/event-stream") with no application-level heartbeat.
During long silent gaps between agent events (slow server-side tools like OCR or
retrieval), idle-timeout proxies in front of the endpoint (Azure ingress, nginx,
serverless front doors) drop the healthy connection and the client sees a spurious
HTTP 500. The Connection: keep-alive header is TCP-level only and emits no application
data during an idle stream.

Fix

Wrap the event-stream generator so that when no upstream event is produced within a
configurable interval, a transport-level SSE keepalive comment (: keepalive) is written
to the wire. SSE comment lines are a protocol no-op that clients and parsers ignore, so
real events still flush immediately and in order while only idle gaps trigger a ping.

• New parameter keepalive_interval_seconds on add_agent_framework_fastapi_endpoint
  (default 15.0; pass None to disable). Non-positive values raise ValueError.
• The upstream generator is drained by a single dedicated task feeding an asyncio.Queue,
  keeping its entire lifecycle in one contextvars context. This is required by the agent
  run/telemetry pipeline, whose streaming cleanup resets ContextVar tokens that must be
  reset in the context that created them (racing individual pulls with asyncio.wait_for
  breaks that and was observed to raise "Token was created in a different Context"). Client
  disconnects cancel the producer so the upstream generator closes cleanly.
• No new dependency (kept StreamingResponse; did not pull in sse-starlette).

Tests

Added focused async tests in tests/ag_ui/test_endpoint.py: a keepalive is emitted during
an idle gap while real events still pass through in order; no keepalive when events flow
back-to-back; empty stream and upstream-error paths; and endpoint-level acceptance,
disable (None), and validation. Full test_endpoint.py (59) and the HTTP round-trip
suites pass; ruff and mypy are clean.

Note

The keepalive defaults on (15s) and emits only benign SSE comment lines during idle gaps.
If you'd prefer it opt-in, defaulting keepalive_interval_seconds to None is a one-line
change.

[Copilot]

Pull request overview

This PR adds an application-level SSE keepalive mechanism to the Python AG-UI FastAPI endpoint so that long idle gaps in the event stream don’t get dropped by idle-timeout proxies (e.g., Azure ingress/nginx), while preserving event order and immediate flushing for real events.

Changes:

• Introduces an SSE keepalive wrapper (_with_sse_keepalive) that injects : keepalive comment frames during idle periods.
• Adds a new keepalive_interval_seconds parameter to add_agent_framework_fastapi_endpoint (default 15.0, None disables; non-positive rejected).
• Adds unit/acceptance tests covering wrapper behavior and endpoint configuration/validation.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py	Adds keepalive wrapper and wires it into the StreamingResponse pipeline via a new endpoint parameter.
python/packages/ag-ui/tests/ag_ui/test_endpoint.py	Adds tests for keepalive wrapper behavior plus endpoint-level enable/disable/validation checks.

[Vaibhav Patel (vaibhav-patel)]

Good point — bounded the queue (maxsize=1) so the producer applies backpressure to the upstream generator; done in b528fdf.

[Vaibhav Patel (vaibhav-patel)]

Agreed — bumped the test keepalive interval to 50ms to avoid CI scheduling flakiness in b528fdf.

[Evan Mattson (moonbox3)]

Thanks for putting this together, and for clearly identifying the AG-UI SSE keepalive gap in #6611. This PR is addressing the right underlying issue: the FastAPI endpoint needs transport-level SSE comment keepalives during idle gaps so healthy long-running streams are not dropped by clients or intermediaries.

We now have #6980 open for the same fix. It implements the keepalive at the endpoint layer as well, using EventSourceResponse, preserves the old StreamingResponse path when keepalive is disabled, includes the needed dependency-bound updates, and has broader current test coverage/CI.

To avoid carrying two competing implementations for the same issue, I’m going to close this PR as superseded by #6980. I appreciate the contribution here; it helped validate the direction and the behavior we need.