javadandCursor
edbaa2acec
Connector calls, semantic rules, and DuckDB memory keep long runs stable, and stream retry stops chats from stalling. Co-authored-by: Cursor <cursoragent@cursor.com>
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
"""Keep idle SSE streams alive through reverse-proxy timeouts."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from typing import AsyncIterator
|
||
|
||
# Do not send Connection. It is hop-by-hop; proxies that honor it reuse the
|
||
# socket while the body is still streaming and the next chat drops both streams.
|
||
SSE_HEADERS = {
|
||
"Cache-Control": "no-cache, no-transform",
|
||
"X-Accel-Buffering": "no",
|
||
}
|
||
|
||
# Comment frames are ignored by EventSource; a data ping resets proxies that
|
||
# only watch application frames. 15s stays under typical 30–60s idle timeouts.
|
||
SSE_PING = ': ping\n\ndata: {"type":"ping"}\n\n'
|
||
DEFAULT_KEEPALIVE_SECONDS = 15.0
|
||
|
||
|
||
async def with_sse_keepalive(
|
||
stream: AsyncIterator[str],
|
||
*,
|
||
interval: float = DEFAULT_KEEPALIVE_SECONDS,
|
||
) -> AsyncIterator[str]:
|
||
iterator = stream.__aiter__()
|
||
pending: asyncio.Task[str] | None = asyncio.create_task(anext(iterator))
|
||
try:
|
||
while True:
|
||
assert pending is not None
|
||
done, _ = await asyncio.wait({pending}, timeout=interval)
|
||
if not done:
|
||
yield SSE_PING
|
||
continue
|
||
try:
|
||
chunk = pending.result()
|
||
except StopAsyncIteration:
|
||
return
|
||
pending = None
|
||
yield chunk
|
||
pending = asyncio.create_task(anext(iterator))
|
||
finally:
|
||
if pending is not None and not pending.done():
|
||
pending.cancel()
|
||
try:
|
||
await pending
|
||
except (asyncio.CancelledError, StopAsyncIteration):
|
||
pass
|