Regenerate actually regenerates

`ensure` is keyed on message_id and idempotent on purpose -- a page load
finding an unfinished reply must attach to it rather than start a second
one, and `_follow` calls it too. But finished generations linger in the
registry for KEEP_FINISHED so a follower arriving at the last moment still
gets the final frames, and regenerate is the only caller that reuses an
existing Message row instead of creating a new one. So `ensure` handed back
the finished generation: no request was made, `_follow` replayed the old
answer, and the `done` frame re-rendered a streaming shell because the row
said incomplete. That is the reconnect loop, and the Send button stuck on
Stop. It appeared to work after five minutes only by accident, and only
sometimes: `_prune` sat below the early return, so it was unreachable for
exactly the message that needed it.

`restart()` is the explicit opposite of `ensure`, and regenerate calls it.
`_prune` moves above the lookup.

Cancelling a live predecessor makes its `finally:` run `_persist` on the
same row, which would overwrite the reply that replaced it. `_persist` now
refuses when another generation owns the message -- "someone else owns this
row now", not "this one is registered", so a direct call still writes.

Three things found next door, all in the same area and all bugs:

  - `done` was set before `_persist` committed, while `_follow`'s docstring
    claimed the opposite. `_follow` breaks out the instant it sees the flag
    and re-renders the bubble from the row, so the row has to be right
    first. Harmless today, a guaranteed loss once metrics land there.
  - Live reasoning duplicated quadratically. The frame carries the whole
    block each time, exactly as `render` and `tools` do, but the target
    swapped it `beforeend`.
  - `sse.KEEPALIVE` was defined and never yielded. A model thinking for
    ninety seconds emits nothing, and an idle connection is what a proxy
    closes.

There was no test for regenerate at all. There is now.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
2026-08-01 00:28:19 +02:00
parent 2c8c274850
commit 85f18e99b2
4 changed files with 337 additions and 6 deletions
+19 -2
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import asyncio import asyncio
import json import json
import logging import logging
import time
from collections.abc import AsyncIterator from collections.abc import AsyncIterator
from fastapi import APIRouter, Depends, Form, HTTPException, Request, status from fastapi import APIRouter, Depends, Form, HTTPException, Request, status
@@ -28,6 +29,10 @@ log = logging.getLogger(__name__)
router = APIRouter(prefix="/api/chats", tags=["chats"]) router = APIRouter(prefix="/api/chats", tags=["chats"])
# Seconds of silence before a comment frame is sent to hold the connection open.
# Well under nginx's 60s default; see services/sse.py:KEEPALIVE.
KEEPALIVE_AFTER = 15.0
def _owned_chat(db: DBSession, chat_id: str, user_id: str) -> Chat: def _owned_chat(db: DBSession, chat_id: str, user_id: str) -> Chat:
chat = db.get(Chat, chat_id) chat = db.get(Chat, chat_id)
# 404 rather than 403 for someone else's chat: whether a given id exists is # 404 rather than 403 for someone else's chat: whether a given id exists is
@@ -240,6 +245,7 @@ async def _follow(chat_id: str, message_id: str) -> AsyncIterator[str]:
generation = generation_service.ensure(chat_id, message_id) generation = generation_service.ensure(chat_id, message_id)
generation.followers += 1 generation.followers += 1
seen = -1 seen = -1
last_frame = time.monotonic()
try: try:
while True: while True:
@@ -251,9 +257,18 @@ async def _follow(chat_id: str, message_id: str) -> AsyncIterator[str]:
yield sse.event("tools", _tool_activity(generation.tool_events)) yield sse.event("tools", _tool_activity(generation.tool_events))
if generation.content: if generation.content:
yield sse.event("render", render_markdown(generation.text)) yield sse.event("render", render_markdown(generation.text))
last_frame = time.monotonic()
if generation.done: if generation.done:
break break
# A reasoning model can think for a minute or more without emitting
# anything, and an idle connection is what a proxy closes. The
# comment frame keeps it open and is ignored by the browser.
if time.monotonic() - last_frame > KEEPALIVE_AFTER:
yield sse.KEEPALIVE
last_frame = time.monotonic()
# Polling rather than per-follower wakeups: the producer already # Polling rather than per-follower wakeups: the producer already
# works in RENDER_INTERVAL steps, so a short sleep is simpler and # works in RENDER_INTERVAL steps, so a short sleep is simpler and
# cannot drop a notification. # cannot drop a notification.
@@ -261,7 +276,7 @@ async def _follow(chat_id: str, message_id: str) -> AsyncIterator[str]:
finally: finally:
generation.followers = max(0, generation.followers - 1) generation.followers = max(0, generation.followers - 1)
# The producer writes the message before marking itself done, so by here # The producer commits the message before marking itself done, so by here
# the row is authoritative and the final bubble can be rendered from it. # the row is authoritative and the final bubble can be rendered from it.
with session_scope() as db: with session_scope() as db:
message = db.get(Message, message_id) message = db.get(Message, message_id)
@@ -592,7 +607,9 @@ async def regenerate(
message.complete = False message.complete = False
message.model_id = chat.model_id message.model_id = chat.model_id
db.commit() db.commit()
generation_service.ensure(chat.id, message.id) # restart, not ensure: this is the one caller that reuses a Message row, and
# the finished generation for it is still registered.
generation_service.restart(chat.id, message.id)
return templates.TemplateResponse( return templates.TemplateResponse(
request, request,
+52 -3
View File
@@ -120,18 +120,48 @@ def ensure(chat_id: str, message_id: str) -> Generation:
Idempotent, because more than one thing can ask for it: the route that Idempotent, because more than one thing can ask for it: the route that
created the message, and any page load that finds the message unfinished. created the message, and any page load that finds the message unfinished.
`_prune` runs first, not after the lookup. Below it, a stale entry could
never expire: the early return is the only path a repeated id takes, so the
sweep was unreachable for exactly the message that needed it.
""" """
_prune()
existing = _RUNNING.get(message_id) existing = _RUNNING.get(message_id)
if existing is not None: if existing is not None:
return existing return existing
_prune()
generation = Generation(chat_id=chat_id, message_id=message_id) generation = Generation(chat_id=chat_id, message_id=message_id)
_RUNNING[message_id] = generation _RUNNING[message_id] = generation
_TASKS[message_id] = asyncio.create_task(_run(generation)) _TASKS[message_id] = asyncio.create_task(_run(generation))
return generation return generation
def restart(chat_id: str, message_id: str) -> Generation:
"""Produce this reply again, discarding any finished attempt at it.
`ensure` is idempotent on purpose, and that is load-bearing: a page load
finding an unfinished reply must attach to it rather than start a second
one, and `_follow` calls it too. Regeneration is the one caller that means
the opposite.
It is also the one caller that reuses an existing Message row -- blanked and
marked incomplete -- rather than creating a new one. The finished Generation
for that id is still in the registry, because finished ones linger
KEEP_FINISHED so a follower arriving at the last moment still gets the final
frames. `ensure` handed that one straight back: no request was made,
`_follow` replayed the previous answer, and the `done` frame re-rendered a
streaming shell because the row said incomplete. That was the reconnect loop,
and the Send button stuck on Stop.
"""
previous = _RUNNING.pop(message_id, None)
task = _TASKS.pop(message_id, None)
if previous is not None and not previous.done:
previous.cancel = True
if task is not None:
task.cancel()
return ensure(chat_id, message_id)
async def shutdown() -> None: async def shutdown() -> None:
"""Stop every running generation, keeping what each has produced.""" """Stop every running generation, keeping what each has produced."""
for task in list(_TASKS.values()): for task in list(_TASKS.values()):
@@ -295,10 +325,14 @@ async def _run(generation: Generation) -> None:
) )
title = title or chat_service.fallback_title(question) title = title or chat_service.fallback_title(question)
# Written *before* `done`, because `_follow` breaks out of its loop the
# moment it sees that flag and immediately re-renders the bubble from
# the row. The other order left a window in which the finished frame
# showed the previous turn's stored values.
_persist(generation, title, time.monotonic() - started)
generation.done = True generation.done = True
generation.finished_at = datetime.now(UTC) generation.finished_at = datetime.now(UTC)
generation.touch() generation.touch()
_persist(generation, title, time.monotonic() - started)
def _question_from(payload: dict) -> str: def _question_from(payload: dict) -> str:
@@ -319,7 +353,21 @@ def _question_from(payload: dict) -> str:
def _persist(generation: Generation, title: str, elapsed: float) -> None: def _persist(generation: Generation, title: str, elapsed: float) -> None:
"""Write the finished reply, name the chat, and set the unread flag.""" """Write the finished reply, name the chat, and set the unread flag.
A generation another one has replaced may not write. A regeneration cancels
its predecessor, whose `finally:` then runs this on the same row -- and it
would overwrite the fresh reply with the abandoned one.
The test is "someone else owns this row now", not "this one is registered":
an unregistered generation still writes, because that is a direct call
rather than a superseded one.
"""
owner = _RUNNING.get(generation.message_id)
if owner is not None and owner is not generation:
log.debug("skipping persist for superseded generation %s", generation.message_id)
return
try: try:
with session_scope() as db: with session_scope() as db:
message = db.get(Message, generation.message_id) message = db.get(Message, generation.message_id)
@@ -364,5 +412,6 @@ __all__ = [
"ensure", "ensure",
"get", "get",
"request_stop", "request_stop",
"restart",
"shutdown", "shutdown",
] ]
+4 -1
View File
@@ -98,7 +98,10 @@
<span class="reasoning__label">Thinking…</span> <span class="reasoning__label">Thinking…</span>
{{ icon("chevron-down", "icon--sm reasoning__chevron") }} {{ icon("chevron-down", "icon--sm reasoning__chevron") }}
</summary> </summary>
<div class="reasoning__body" sse-swap="reasoning" hx-swap="beforeend"></div> {# innerHTML, not beforeend: the frame carries the whole block of
thinking each time, exactly as `render` and `tools` do. Appending it
repeated everything already shown, so the panel grew quadratically. #}
<div class="reasoning__body" sse-swap="reasoning" hx-swap="innerHTML"></div>
</details> </details>
{# Tool activity as it happens. Empty until the model asks for something, {# Tool activity as it happens. Empty until the model asks for something,
+262
View File
@@ -0,0 +1,262 @@
"""Regenerating a reply, and the registry rules that make it work.
Regeneration is the one caller that reuses a Message row rather than creating a
new one, which is why it is the one caller `ensure` was wrong for.
"""
from __future__ import annotations
import asyncio
from datetime import UTC, datetime, timedelta
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import select
from lembas.db.models import Chat, Connection, Message, Model
from lembas.services import generation as generation_service
from lembas.services.crypto import encrypt
@pytest.fixture(autouse=True)
def empty_registry():
"""The registry is module state. A test that leaves an entry behind changes
what the next one sees."""
yield
generation_service._RUNNING.clear()
generation_service._TASKS.clear()
@pytest.fixture
def no_upstream(monkeypatch):
"""Replace the producer, so a test can watch the registry without a server.
Records the generations it was asked to run.
"""
started: list = []
async def _fake_run(generation):
started.append(generation)
monkeypatch.setattr(generation_service, "_run", _fake_run)
return started
def _connection(db) -> Connection:
connection = Connection(
name="Test", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")
)
db.add(connection)
db.commit()
db.add(Model(connection_id=connection.id, model_id="test-model"))
db.commit()
return connection
def _reply(db, chat_id: str, *, complete: bool = True) -> Message:
db.add(Message(chat_id=chat_id, role="user", content="what is lembas?"))
reply = Message(chat_id=chat_id, role="assistant", content="Waybread.", complete=complete)
db.add(reply)
db.commit()
return reply
def _finished(chat_id: str, message_id: str, *, text: str = "old") -> generation_service.Generation:
"""A generation in the state a just-completed reply leaves behind."""
generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
generation.content.append(text)
generation.done = True
generation.finished_at = datetime.now(UTC)
generation_service._RUNNING[message_id] = generation
return generation
# --- The bug -----------------------------------------------------------------
def test_regenerate_starts_a_fresh_generation(
client: TestClient, db, registered, make_chat, no_upstream
):
"""The bug: `ensure` handed back the finished generation still sitting in
the registry, so no request was ever made and the browser reconnected to a
stream that had nothing left to say."""
_connection(db)
chat_id = make_chat()
reply = _reply(db, chat_id)
stale = _finished(chat_id, reply.id)
response = client.post(f"/api/chats/{chat_id}/messages/{reply.id}/regenerate")
assert response.status_code == 200
fresh = generation_service.get(reply.id)
assert fresh is not stale
assert not fresh.done
assert fresh.text == ""
assert no_upstream == [fresh]
def test_regenerate_blanks_the_row_and_returns_a_streaming_shell(
client: TestClient, db, registered, make_chat, no_upstream
):
_connection(db)
chat_id = make_chat()
reply = _reply(db, chat_id)
body = client.post(f"/api/chats/{chat_id}/messages/{reply.id}/regenerate").text
assert "sse-connect=" in body
assert 'sse-swap="render"' in body
db.expire_all()
row = db.get(Message, reply.id)
assert row.complete is False
assert row.content == ""
def test_regenerating_twice_in_a_row_works(
client: TestClient, db, registered, make_chat, no_upstream
):
"""The reported symptom was that it worked once, sometimes."""
_connection(db)
chat_id = make_chat()
reply = _reply(db, chat_id)
client.post(f"/api/chats/{chat_id}/messages/{reply.id}/regenerate")
first = generation_service.get(reply.id)
first.done = True
first.finished_at = datetime.now(UTC)
client.post(f"/api/chats/{chat_id}/messages/{reply.id}/regenerate")
assert generation_service.get(reply.id) is not first
assert len(no_upstream) == 2
def test_regenerating_someone_elses_reply_is_not_found(
client: TestClient, db, registered, make_chat, no_upstream
):
from lembas.db.models import User
from lembas.security.passwords import hash_password
_connection(db)
chat_id = make_chat()
reply = _reply(db, chat_id)
someone_else = User(name="Sam", email="s@shire.test", password_hash=hash_password("x"))
db.add(someone_else)
db.commit()
db.scalar(select(Chat).where(Chat.id == chat_id)).user_id = someone_else.id
db.commit()
response = client.post(f"/api/chats/{chat_id}/messages/{reply.id}/regenerate")
assert response.status_code == 404
# --- Registry rules ----------------------------------------------------------
async def test_ensure_still_attaches_to_an_unfinished_reply(db, no_upstream):
"""Idempotence is load-bearing: a page load that finds an unfinished reply
must attach to it, not start a second one."""
first = generation_service.ensure("chat", "message")
assert generation_service.ensure("chat", "message") is first
await asyncio.sleep(0) # let the scheduled task actually start
assert len(no_upstream) == 1
async def test_prune_expires_a_finished_generation_before_the_lookup(db, no_upstream):
"""`_prune` used to sit below the early return, where it could never reach
the one entry that needed it."""
stale = _finished("chat", "message")
stale.finished_at = datetime.now(UTC) - generation_service.KEEP_FINISHED - timedelta(minutes=1)
assert generation_service.ensure("chat", "message") is not stale
def test_a_superseded_generation_does_not_write_the_row(db, registered, make_chat):
"""A cancelled predecessor's `finally:` runs _persist on the same message,
and it must not overwrite the reply that replaced it."""
chat_id = make_chat()
reply = _reply(db, chat_id)
abandoned = generation_service.Generation(chat_id=chat_id, message_id=reply.id)
abandoned.content.append("the abandoned attempt")
current = generation_service.Generation(chat_id=chat_id, message_id=reply.id)
current.content.append("the reply that replaced it")
generation_service._RUNNING[reply.id] = current
generation_service._persist(abandoned, "", 0.0)
db.expire_all()
assert db.get(Message, reply.id).content == "Waybread."
async def test_restart_cancels_a_generation_still_running(db, no_upstream):
live = generation_service.ensure("chat", "message")
generation_service.restart("chat", "message")
assert live.cancel is True
assert generation_service.get("message") is not live
# --- Ordering ----------------------------------------------------------------
async def test_the_reply_is_persisted_before_it_is_marked_done(
db, registered, make_chat, monkeypatch
):
"""`_follow` breaks out the instant it sees `done` and re-renders the bubble
from the row, so the row has to be right first."""
chat_id = make_chat()
reply = _reply(db, chat_id, complete=False)
seen: list[bool] = []
def _record(generation, title, elapsed):
seen.append(generation.done)
monkeypatch.setattr(generation_service, "_persist", _record)
generation = generation_service.Generation(chat_id=chat_id, message_id=reply.id)
generation_service._RUNNING[reply.id] = generation
# No connection row, so resolve_endpoint fails and _run goes straight to
# its finally: which is the part under test.
await generation_service._run(generation)
assert seen == [False]
assert generation.done is True
# --- The streaming shell ------------------------------------------------------
def test_live_reasoning_is_replaced_not_appended(
client: TestClient, db, registered, make_chat, no_upstream
):
"""The frame carries the whole block each time. Appending it repeated
everything already shown, so the panel grew quadratically."""
_connection(db)
chat_id = make_chat()
reply = _reply(db, chat_id)
body = client.post(f"/api/chats/{chat_id}/messages/{reply.id}/regenerate").text
line = next(line for line in body.splitlines() if 'sse-swap="reasoning"' in line)
assert 'hx-swap="innerHTML"' in line
assert "beforeend" not in line
async def test_a_silent_generation_gets_a_keepalive(db, registered, make_chat, monkeypatch):
"""A model thinking for a minute emits nothing, and an idle connection is
what a proxy closes."""
from lembas.api import chats as chats_api
chat_id = make_chat()
reply = _reply(db, chat_id, complete=False)
generation = generation_service.Generation(chat_id=chat_id, message_id=reply.id)
generation_service._RUNNING[reply.id] = generation
monkeypatch.setattr(chats_api, "KEEPALIVE_AFTER", 0.0)
frames: list[str] = []
stream = chats_api._follow(chat_id, reply.id).__aiter__()
async def _finish():
await asyncio.sleep(0.05)
generation.done = True
task = asyncio.create_task(_finish())
async for frame in stream:
frames.append(frame)
if frame.startswith("event: close"):
break
await task
assert any(frame == ": keepalive\n\n" for frame in frames)