feat: add heartbeats to prevent long-thinking from closing

This commit is contained in:
h
2026-05-22 23:59:00 +02:00
parent 33eece331a
commit 5fb5b6a3fb
@@ -24,12 +24,14 @@ shape.
from __future__ import annotations
import asyncio
import contextlib
import json
import logging
import os
import tempfile
import time
from collections.abc import AsyncIterator
from pathlib import Path
from typing import TYPE_CHECKING, Any
@@ -86,6 +88,13 @@ _STREAM_FLUSH_DEBOUNCE = 0.4
# disk round-trip).
_SSE_FLUSH_DEBOUNCE = 0.1
# Interval between SSE comment-frames sent when the backend is silent
# (e.g. claude is mid-thinking on a large context). The Obsidian plugin
# and any intermediate proxies will hold the connection open as long as
# bytes keep flowing; a comment-frame is the cheapest legal SSE keepalive.
# Set well under typical proxy/client idle timeouts (60s).
_SSE_HEARTBEAT_INTERVAL = 15.0
class MarkdownFrontend(Frontend):
"""FastAPI app behind ``POST /chat`` driven by Obsidian-vault files."""
@@ -620,7 +629,14 @@ class MarkdownFrontend(Frontend):
return _reattach_frontmatter(parsed.metadata, new_body)
try:
async for ev in events:
async for ev in _events_with_heartbeat(events):
if ev is None:
# Backend is quiet (claude mid-thinking, MCP slow,
# whatever). SSE comment-frame keeps the TCP socket
# warm so the plugin / uvicorn / any reverse proxy
# doesn't time the request out before we finish.
yield b": keepalive\n\n"
continue
acc.feed(ev)
now = time.monotonic()
if (
@@ -906,6 +922,36 @@ class MarkdownFrontend(Frontend):
# ---- module-level utilities ----------------------------------------------
async def _events_with_heartbeat(
events: AsyncIterator[Any], interval: float = _SSE_HEARTBEAT_INTERVAL
) -> AsyncIterator[Any]:
"""Wrap an async event stream with idle-time heartbeat markers.
Yields ``None`` every ``interval`` seconds during silence; real
events pass through unchanged. When the wrapped iterator is
exhausted, this generator returns. Cancellation propagates: if the
outer scope is cancelled we cancel the pending ``__anext__`` task
instead of leaving it dangling.
"""
src = events.__aiter__()
while True:
next_task: asyncio.Task[Any] = asyncio.ensure_future(src.__anext__())
try:
done, _pending = await asyncio.wait({next_task}, timeout=interval)
except BaseException:
next_task.cancel()
with contextlib.suppress(BaseException):
await next_task
raise
if not done:
yield None
continue
try:
yield next_task.result()
except StopAsyncIteration:
return
def _sse_pack(event: str, data: dict[str, Any]) -> bytes:
r"""Format one Server-Sent Event frame.