Refactor AI Assist to background-task + buffer architecture

The assist flow previously tied the entire LLM generation to a single
POST request with no keepalives, causing NS_ERROR_NET_PARTIAL_TRANSFER
in Firefox when Hypercorn closed the connection during gaps between
chunks. This refactor decouples generation into a background task with
a buffer and a separate SSE stream — the same pattern used by chat.

- generation_buffer.py: Widen _buffers to support string keys, add
  create/get/remove_assist_buffer() using "assist:{user_id}" keys,
  fix cleanup log format for string keys
- generation_task.py: Add run_assist_generation() — lightweight
  background task with no DB persistence or title generation
- notes.py: Replace single POST SSE route with POST /api/notes/assist
  (returns 202) + GET /api/notes/assist/stream (SSE with 15s keepalives
  and Last-Event-ID reconnection); 409 if already running
- useAssist.ts: Switch from apiStreamPost to apiPost + apiSSEStream
  two-step pattern with named event mapping and stream handle cleanup

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-02-13 00:27:21 -05:00
parent 3ec49b7f24
commit a89d25f5d6
5 changed files with 201 additions and 54 deletions
@@ -70,7 +70,7 @@ class GenerationBuffer:
# Module-level singleton registry
_buffers: dict[int, GenerationBuffer] = {}
_buffers: dict[int | str, GenerationBuffer] = {}
_cleanup_task: asyncio.Task | None = None
_GRACE_PERIOD = 60.0 # seconds to keep completed buffers
@@ -92,20 +92,38 @@ def remove_buffer(conv_id: int) -> None:
_buffers.pop(conv_id, None)
def create_assist_buffer(user_id: int) -> GenerationBuffer:
key = f"assist:{user_id}"
existing = _buffers.get(key)
if existing and existing.state == GenerationState.RUNNING:
raise RuntimeError(f"Assist generation already running for user {user_id}")
buf = GenerationBuffer(conversation_id=0, assistant_message_id=0)
_buffers[key] = buf
return buf
def get_assist_buffer(user_id: int) -> GenerationBuffer | None:
return _buffers.get(f"assist:{user_id}")
def remove_assist_buffer(user_id: int) -> None:
_buffers.pop(f"assist:{user_id}", None)
async def _cleanup_loop() -> None:
"""Remove completed/errored buffers after grace period."""
while True:
await asyncio.sleep(15)
now = time.monotonic()
to_remove = [
cid for cid, buf in _buffers.items()
key for key, buf in _buffers.items()
if buf.state != GenerationState.RUNNING
and buf.finished_at is not None
and now - buf.finished_at > _GRACE_PERIOD
]
for cid in to_remove:
_buffers.pop(cid, None)
logger.debug("Cleaned up generation buffer for conversation %d", cid)
for key in to_remove:
_buffers.pop(key, None)
logger.debug("Cleaned up generation buffer for key %s", key)
def start_cleanup_loop() -> None:
@@ -134,3 +134,25 @@ async def run_generation(
buf.state = GenerationState.ERRORED
buf.finished_at = time.monotonic()
buf.append_event("error", {"error": str(e)})
async def run_assist_generation(
buf: GenerationBuffer,
messages: list[dict],
model: str,
) -> None:
"""Stream LLM response for assist into buffer. No DB persistence."""
try:
async for chunk in stream_chat(messages, model, options={"num_predict": 4096}):
buf.content_so_far += chunk
buf.append_event("chunk", {"chunk": chunk})
buf.state = GenerationState.COMPLETED
buf.finished_at = time.monotonic()
buf.append_event("done", {"done": True, "full_text": buf.content_so_far})
except Exception as e:
logger.exception("Error in assist generation task")
buf.state = GenerationState.ERRORED
buf.finished_at = time.monotonic()
buf.append_event("error", {"error": str(e)})