"""Chat creation, messaging and the streaming reply endpoint.""" from __future__ import annotations import asyncio import json import logging import time from collections.abc import AsyncIterator from datetime import UTC, datetime from fastapi import APIRouter, Depends, Form, HTTPException, Request, status from fastapi.responses import HTMLResponse, JSONResponse, Response, StreamingResponse from sqlalchemy import func, select from sqlalchemy.orm import Session as DBSession from lembas.api.deps import Db, RequiredUser, require_permission from lembas.db.models import ( KIND_AGENT, KIND_CHAT, ROLE_ASSISTANT, ROLE_USER, Chat, Folder, Message, Model, User, ) from lembas.db.session import session_scope from lembas.security import permissions from lembas.services import audio as audio_service from lembas.services import chat as chat_service from lembas.services import compaction as compaction_service from lembas.services import files as files_service from lembas.services import generation as generation_service from lembas.services import interaction, settings_store, sse from lembas.services import metrics as metrics_service from lembas.services import prompts as prompts_service from lembas.services import tools as tools_service from lembas.services.agent import policy as agent_policy from lembas.services.agent import terminal as terminal_service from lembas.services.markdown import escape_text, render_markdown from lembas.web.templating import render, templates log = logging.getLogger(__name__) 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 # How many prompts may wait behind a reply at once. The terminal panel's Auto # send is what this exists for: a `for` loop in a shell can produce commands # faster than any model answers them, and a bound with a sentence attached is # better than four hundred rows nobody meant to write. MAX_QUEUED = 10 # How many things one chat may have switched off. There are a dozen families and # sixty skills at most, so this is not a limit anybody reaches by hand -- it is # there so a crafted POST cannot grow the column without bound. MAX_SCOPE_KEYS = 200 def _owned_chat(db: DBSession, chat_id: str, user_id: str) -> Chat: chat = db.get(Chat, chat_id) # 404 rather than 403 for someone else's chat: whether a given id exists is # not information this endpoint should hand out. if chat is None or chat.user_id != user_id: raise HTTPException(status.HTTP_404_NOT_FOUND, "That chat no longer exists.") return chat def _new_chat( db: DBSession, user: User, *, folder_id: str = "", model_id: str = "", temporary: bool = False, kind: str = KIND_CHAT, ssh_profile_id: str = "", project_dir: str = "", agent_mode: str = "", reasoning_effort: str = "", ) -> Chat: """Create a chat row, resolving which model it should use. An agent chat's connection is settled here and never again. That is the lock: the harness, the tools offered and the approval loop all differ, so a conversation whose earlier turns ran somewhere else is not one conversation. The mode is *not* part of that lock and is accepted here so it can be chosen before the first word. Without it, reaching Plan mode meant starting a chat in Manual, sending something to make the chat exist, and only then being offered the control -- by which point the model had already answered under the wrong rules. The reasoning effort is accepted for the same reason, and wins over the model's default: an explicit choice beats an inherited one. A folder's own defaults fill in anything the request left empty, and nothing it filled in. That order is the point: the folder says what this piece of work usually needs, and the screen in front of somebody says what they want this time. The folder's system prompt is deliberately not among them -- it is read at request time so that editing the folder later reaches the chats already in it. """ folder = db.get(Folder, folder_id) if folder_id else None if folder is not None and folder.user_id != user.id: folder = None if folder is not None: model_id = model_id or folder.model_id kind = kind or folder.kind if kind == KIND_AGENT: ssh_profile_id = ssh_profile_id or folder.ssh_profile_id project_dir = project_dir or folder.project_dir agent_mode = agent_mode or folder.agent_mode chosen = None if model_id: match = next( (m for m in chat_service.available_models(db, user) if m.model_id == model_id), None ) if match is not None: chosen = (match.model_id, match.connection_id) if chosen is None: chosen = chat_service.default_model(db, user) # `Model.params_json` has said "default sampling params applied to new chats # using this model" since it was added and has been applied nowhere. It is # empty on every existing row, so honouring it now changes nothing until an # administrator sets something -- and it is what makes a per-model default # reasoning effort possible without a second column meaning the same thing. defaults: dict = {} if chosen is not None: model = db.scalar( select(Model).where( Model.model_id == chosen[0], Model.connection_id == chosen[1] ) ) if model is not None: defaults = dict(model.params_json or {}) profile = _agent_target(db, user, kind, ssh_profile_id) chat = Chat( user_id=user.id, folder_id=folder_id or None, model_id=chosen[0] if chosen else "", connection_id=chosen[1] if chosen else None, temporary=temporary, kind=KIND_AGENT if profile is not None else KIND_CHAT, ssh_profile_id=profile.id if profile is not None else None, project_dir=(project_dir.strip() or profile.default_dir) if profile is not None else "", params_json=defaults, ) # Ignored rather than refused when it is not a mode, matching how every # other bad value here collapses: somebody who mistypes should get a chat # under the safest rules, not an error page holding their message hostage. # Left alone entirely on a plain chat, where it means nothing. if profile is not None and agent_mode.strip() in agent_policy.MODES: chat.agent_mode = agent_mode.strip() # After the model's defaults, so choosing one on the new-chat screen wins # over the administrator's. # # `"off"` is a sentinel, and it has to be: `reasoning_effort` arrives as # `Form("")`, so an absent field and an empty one are indistinguishable -- # the FastAPI trap this codebase has already been bitten by once. With # `value=""` on the off option, the reader would pick "off", the value would # fall out of EFFORTS, the model's default seeded above would stay, and they # would silently get "high". The picker shows what will be sent, so the two # have to agree. wanted_effort = reasoning_effort.strip().lower() if wanted_effort == "off": chat.params_json = {k: v for k, v in chat.params_json.items() if k != "reasoning_effort"} elif wanted_effort in chat_service.EFFORTS: chat.params_json = {**chat.params_json, "reasoning_effort": wanted_effort} db.add(chat) db.commit() return chat @router.post("/start", dependencies=[Depends(require_permission("chat.create"))]) async def start_chat( db: Db, user: RequiredUser, content: str = Form(""), file_ids: list[str] = Form(default=[]), folder_id: str = Form(""), model_id: str = Form(""), temporary: bool = Form(False), kind: str = Form(KIND_CHAT), ssh_profile_id: str = Form(""), project_dir: str = Form(""), agent_mode: str = Form(""), reasoning_effort: str = Form(""), ) -> Response: """Create a chat from its first message. Chats are made here rather than by a "New chat" button so that an opened- and-abandoned chat never exists: the row appears only once there is something in it. The reply then streams the same way as any other, because /chat/{id} renders the unfinished assistant message with its sse-connect. """ content = content.strip() if not content and not file_ids: return Response(status_code=status.HTTP_204_NO_CONTENT) chat = _new_chat( db, user, folder_id=folder_id, model_id=model_id, temporary=temporary, kind=kind, ssh_profile_id=ssh_profile_id, project_dir=project_dir, agent_mode=agent_mode, reasoning_effort=reasoning_effort, ) user_message = chat_service.create_message(db, chat, ROLE_USER, content) if file_ids: files_service.claim(db, ids=file_ids, user_id=user.id, message_id=user_message.id) assistant = chat_service.create_message( db, chat, ROLE_ASSISTANT, "", complete_=False, model_id=chat.model_id ) generation_service.ensure(chat.id, assistant.id) response = Response(status_code=status.HTTP_204_NO_CONTENT) response.headers["HX-Redirect"] = f"/chat/{chat.id}" return response def _agent_target(db: DBSession, user: User, kind: str, profile_id: str): """The connection an agent chat is being pointed at, or None. Every "no" collapses to None and the chat is an ordinary one: not asked for, no permission, the feature off, or a profile that is not this person's. Refusing outright would be worse -- somebody whose permission was withdrawn between opening the composer and sending would lose the message. """ from lembas.db.models import SshProfile from lembas.security import permissions if kind != KIND_AGENT or not profile_id: return None if not permissions.has(db, user, "tools.agent"): return None if not settings_store.agents(db).get("enabled"): return None profile = db.get(SshProfile, profile_id) # Ownership re-checked rather than trusted from the form: an id in a POST is # not an authorisation, and these are credentials to somebody's machine. if profile is None or profile.owner_id != user.id or not profile.enabled: return None return profile # There is deliberately no route that creates an empty chat. Starting one is # navigation to /chat (optionally ?model=...), and the row is written by # /start when the first message is actually sent. @router.get("/{chat_id}/inspect") async def inspect_chat(request: Request, db: Db, user: RequiredUser, chat_id: str) -> Response: """What this chat would send upstream right now. Owner-checked *and* admin-checked, not admin alone. `permissions.resolve` giving an admin everything is about configuration, which they can grant themselves anyway; reading someone's conversation is a different act, which is why `sharing.visible_to` has no admin branch either. An inspector that could dump any user's transcript would be that branch under another name. Rebuilt, not recorded. Recording every request would store a copy of the whole conversation against every message, which grows quadratically with chat length -- and the thing an administrator actually wants to see is what the current configuration produces. The panel says so in as many words. """ chat = _owned_chat(db, chat_id, user.id) if not user.is_admin: raise HTTPException( status.HTTP_403_FORBIDDEN, "The inspector is restricted to administrators." ) offered = tools_service.enabled_tools(db, chat, user) payload = chat_service.build_request(db, chat, tools=offered, user=user) last = db.scalar( select(Message) .where(Message.chat_id == chat.id, Message.role == ROLE_ASSISTANT) .order_by(Message.created_at.desc()) ) messages = payload.get("messages") or [] system = messages[0]["content"] if messages and messages[0].get("role") == "system" else "" return render( request, "chat/_inspector_body.html", { "chat": chat, "system": system, "request_json": _pretty(_redact(payload)), "row": last, "metrics": metrics_service.from_message(last.usage_json if last else None), "tool_names": [ (t.get("function") or {}).get("name", "") for t in offered ], "model": chat_service.model_for(db, chat), }, ) @router.get("/{chat_id}/usage") async def chat_usage(request: Request, db: Db, user: RequiredUser, chat_id: str) -> Response: """What this conversation has cost, and how full the window is. Owner-checked and nothing else: it is your own chat's totals. Unlike the inspector next door there is no admin branch, because there is no reason for one -- the numbers describe a conversation, and reading somebody's conversation is exactly what `sharing` has no admin branch for either. Summed from what each reply recorded rather than recomputed: an endpoint that reported no usage contributed an estimate at the time, and re-deriving it now with a different estimator would make the totals move under a chat that had not changed. """ chat = _owned_chat(db, chat_id, user.id) replies = list( db.scalars( select(Message) .where(Message.chat_id == chat.id, Message.role == ROLE_ASSISTANT) .order_by(Message.created_at) ) ) totals = {"prompt": 0, "completion": 0, "total": 0} estimated = False for reply in replies: usage = metrics_service.from_message(reply.usage_json) totals["prompt"] += usage.prompt_tokens totals["completion"] += usage.completion_tokens totals["total"] += usage.total_tokens estimated = estimated or usage.estimated last = replies[-1] if replies else None return render( request, "chat/_usage.html", { "chat": chat, "totals": totals, "estimated": estimated, "replies": len(replies), "metrics": metrics_service.from_message(last.usage_json if last else None), "model": chat_service.model_for(db, chat), }, ) # Roughly what a downscaled phone photo comes to as base64. The exact figure # does not matter; putting megabytes of it into the DOM does. _REDACTED_URI = "data:…base64 image omitted…" MAX_INSPECT_CHARS = 40_000 def _redact(payload: dict) -> dict: """Replace image data URIs before dumping. Nothing else is hidden -- fidelity is the whole point of the panel, and API keys never appear because `build_request` returns a body, not headers. """ messages = [] for message in payload.get("messages") or []: content = message.get("content") if isinstance(content, list): parts = [] for part in content: if isinstance(part, dict) and part.get("type") == "image_url": parts.append({"type": "image_url", "image_url": {"url": _REDACTED_URI}}) else: parts.append(part) message = {**message, "content": parts} messages.append(message) return {**payload, "messages": messages} def _pretty(payload: dict) -> str: text = json.dumps(payload, indent=2, ensure_ascii=False, default=str) if len(text) > MAX_INSPECT_CHARS: return text[:MAX_INSPECT_CHARS] + "\n… truncated" return text @router.post("/{chat_id}/compact") async def compact_chat(request: Request, db: Db, user: RequiredUser, chat_id: str) -> Response: """Summarise the earlier turns and stop sending them. No permission of its own: compaction changes only what one chat sends upstream, and gating it would mean answering "why can this user not tidy their own conversation". """ chat = _owned_chat(db, chat_id, user.id) unfinished = db.scalar( select(Message).where(Message.chat_id == chat.id, Message.complete.is_(False)) ) if unfinished is not None: # Summarising a transcript that is still being written races # build_request. Queuing it is a state machine nobody asked for. raise HTTPException( status.HTTP_409_CONFLICT, "Wait for the current reply to finish, then compact." ) template = prompts_service.resolve(db, "task.compact") if not template.strip(): raise HTTPException( status.HTTP_409_CONFLICT, "Compaction is turned off: its prompt is empty under Admin → Prompts.", ) upto = compaction_service.last_complete(db, chat) if upto is None: raise HTTPException( status.HTTP_409_CONFLICT, "There is nothing here to summarise yet." ) endpoint, model_id = chat_service.resolve_endpoint(db, chat) transcript = compaction_service.transcript(db, chat, upto=upto) previous = compaction_service.previous_summary_block(chat) summary = await chat_service.summarise_for_compaction( endpoint, model_id, transcript=transcript, previous_summary=previous, template=template, ) if not summary: raise HTTPException( status.HTTP_409_CONFLICT, "The model returned no summary, so nothing changed." ) compaction_service.apply(chat, summary=summary, upto=upto) db.commit() log.info("chat %s compacted through %s", chat.id, upto.id) return templates.TemplateResponse( request, "chat/_thread.html", {"request": request, **_thread_context(db, chat, user)} ) @router.post("/{chat_id}/index") async def reindex_chat(db: Db, user: RequiredUser, chat_id: str) -> Response: """Walk the project directory again, now. The listing is cached for five minutes and only ever built when a reply starts, so a tree that has just changed under somebody's hands -- a checkout, a build, anything done in the terminal panel rather than through `file_write` -- stays wrong until the next reply after the TTL lapses. This is the "look again" that was missing. Read-only, and therefore outside `agent/policy.py` for the reason the directory browser is: it is LLeMbas acting on a person's instruction, not a model choosing to look, and a listing that asked permission would be useless. The gate is ownership of the connection, checked here rather than trusted from the chat. """ from lembas.db.models import SshProfile from lembas.services.agent import index as index_service from lembas.services.agent import ssh as ssh_service chat = _owned_chat(db, chat_id, user.id) if chat.kind != KIND_AGENT or not chat.ssh_profile_id: raise HTTPException(status.HTTP_409_CONFLICT, "This chat has no project directory.") if not permissions.has(db, user, "tools.agent"): raise HTTPException(status.HTTP_403_FORBIDDEN, "You may not use agent connections.") profile = db.get(SshProfile, chat.ssh_profile_id) if profile is None or profile.owner_id != user.id or not profile.enabled: raise HTTPException(status.HTTP_404_NOT_FOUND, "That connection is not available.") if not profile.host_key: raise HTTPException( status.HTTP_409_CONFLICT, "This connection's host key has not been accepted yet.", ) project_dir = chat.project_dir or profile.default_dir or "" try: found = await index_service.ensure( ssh_service.SshExecutor(ssh_service.spec_from(profile), project_dir), profile.id, project_dir, refresh=True, ) except Exception as exc: # noqa: BLE001 - surfaced to the reader, not swallowed log.warning("could not index %s for chat %s: %s", project_dir, chat.id, exc) raise HTTPException( status.HTTP_502_BAD_GATEWAY, "Could not read the project directory." ) from exc listed = len(found.paths) return JSONResponse( { "ok": True, "files": found.total, "listed": listed, "truncated": found.truncated, "message": ( f"{found.total} files under {project_dir or '~'}" + (f", {listed} listed." if listed != found.total else ".") ), } ) @router.post("/{chat_id}/bases") async def attach_base( request: Request, db: Db, user: RequiredUser, chat_id: str, base_id: str = Form("") ) -> Response: """Scope this chat to a knowledge base, from the `@` menu. A base is a *reference*, not an attachment: `Chat.knowledge_bases` already narrows `knowledge_search`, and the harness already names the attached bases so the model can tell "there is nothing about this" from "I can only see this folder". Copying a folder of documents into the window instead would cost the context on every request forever to answer one question. Additive, and idempotent -- choosing the same base twice is not an error. Removing one is a checkbox in the chat's settings, where the whole set is visible at once. """ from lembas.db.models import KnowledgeBase from lembas.services.library import documents as documents_service chat = _owned_chat(db, chat_id, user.id) if not permissions.has(db, user, "library.use"): raise HTTPException(status.HTTP_403_FORBIDDEN, "You may not use the library.") base = db.scalar( documents_service.visible_bases(db, user).where(KnowledgeBase.id == base_id) ) if base is None: raise HTTPException(status.HTTP_404_NOT_FOUND, "That knowledge base is not available.") if base.id not in {b.id for b in chat.knowledge_bases}: chat.knowledge_bases = [*chat.knowledge_bases, base] db.commit() return templates.TemplateResponse( request, "chat/_base_chip.html", {"request": request, "base": base} ) @router.post("/{chat_id}/scope") async def set_scope( db: Db, user: RequiredUser, chat_id: str, kind: str = Form(""), name: str = Form(""), on: bool = Form(False), ) -> Response: """Turn one thing this chat may use on or off. **Narrowing only.** Nothing here widens anything. `resolve_tools` applies this *after* the model's capabilities, the reader's permissions and the instance configuration, so a crafted POST turning something on reaches a tool those gates have already removed -- there is a test for exactly that. On is stored by **removing** the key rather than by writing True, so absent stays the single representation of "on" and the column cannot grow a row per family per chat. Bounded, so a crafted request cannot grow it either. JSON reassignment rather than mutation: a plain dict assignment into a JSON column is not detected. """ chat = _owned_chat(db, chat_id, user.id) bucket = {"family": "families", "skill": "skills"}.get(kind.strip()) wanted = name.strip()[:64] if bucket is None or not wanted: raise HTTPException(status.HTTP_400_BAD_REQUEST, "Say what to turn on or off.") scope = dict(chat.scope_json or {}) entries = dict(scope.get(bucket) or {}) if on: entries.pop(wanted, None) else: if len(entries) >= MAX_SCOPE_KEYS: raise HTTPException(status.HTTP_409_CONFLICT, "Too many things switched off.") entries[wanted] = False if entries: scope[bucket] = entries else: scope.pop(bucket, None) chat.scope_json = scope db.commit() return Response(status_code=status.HTTP_204_NO_CONTENT) @router.post("/{chat_id}/keep") async def keep_chat(db: Db, user: RequiredUser, chat_id: str) -> Response: """Stop a temporary chat being temporary. A conversation that turns out to matter has to have a way out; without one, the sweep destroys it a day later with no recourse, and people discover that exactly once. """ chat = _owned_chat(db, chat_id, user.id) chat.temporary = False db.commit() response = Response(status_code=status.HTTP_204_NO_CONTENT) # The sidebar has to gain a row and the topbar has to lose a badge; a full # refresh is one line against a handful of out-of-band fragments. response.headers["HX-Refresh"] = "true" return response @router.get("/unread") async def unread_poll(db: Db, user: RequiredUser) -> Response: """Dots for the sidebar, and a toast for anything newly arrived. Polled rather than pushed: a browser sitting on a different chat has no open connection to the one that finished, and a second always-on channel per tab is a lot of machinery for a green dot. Returns out-of-band spans so only the dots change -- re-rendering the whole sidebar would reset the folder open/closed state on every tick. """ chats = list( db.scalars( select(Chat).where( Chat.user_id == user.id, Chat.archived.is_(False), # A temporary chat has no sidebar row, so a dot has nowhere to # land and the toast would name a chat nobody can navigate to. Chat.temporary.is_(False), ) ) ) fresh = [c for c in chats if c.unread and not c.unread_notified] for chat in fresh: chat.unread_notified = True if fresh: db.commit() markup = "".join( f'' for c in chats ) response = HTMLResponse(markup) if fresh: # HX-Trigger carries the toast; ui.js listens for it. response.headers["HX-Trigger"] = json.dumps( {"lembas:unread": {"titles": [c.title for c in fresh]}} ) return response @router.post("/{chat_id}/messages") async def post_message( request: Request, db: Db, user: RequiredUser, chat_id: str, content: str = Form(""), file_ids: list[str] = Form(default=[]), ) -> Response: """Persist the user's turn and hand back the pair of bubbles. The assistant bubble comes back empty, carrying the sse-connect attribute that opens the stream below. Splitting it this way means the POST returns immediately and the slow part is a separate, resumable connection. """ chat = _owned_chat(db, chat_id, user.id) content = content.strip() # "Here, look at this" with no words is a legitimate turn, so an empty # message is only empty when it carries nothing at all. if not content and not file_ids: return Response(status_code=status.HTTP_204_NO_CONTENT) return _send(request, db, chat, user, content, file_ids=file_ids) def _reply_in_flight(db: DBSession, chat: Chat) -> bool: """Whether this chat already has a reply being written. The row is the authority, not the registry: a restart leaves an incomplete assistant message behind with no `Generation` anywhere, and that row is what starts the reply again on the next page load. The registry is consulted too, for the sliver in which a generation is still running and its row has already been written -- `_persist` sets `complete` before `_run` sets `done`. """ unfinished = db.scalar( select(Message).where(Message.chat_id == chat.id, Message.complete.is_(False)) ) return unfinished is not None or generation_service.running_for(chat.id) is not None def _note_rewind(chat: Chat) -> None: """Record that an agent chat's transcript went back and the machine did not. Deliberately no attempt to undo anything out there. The project directory is somebody's real working tree, and deleting their work to match a rewound transcript would be far worse than the inconsistency. So the model is told instead -- see the `tool.agent_rewound` fragment -- and can look rather than assume. """ if chat.kind == KIND_AGENT: chat.rewound_at = datetime.now(UTC) def _send( request: Request, db: Db, chat: Chat, user: User, content: str, *, file_ids: list[str] | None = None, ) -> Response: """Write a turn, start the reply, and hand back the pair of bubbles. Shared by the composer and by anything else that puts words into a conversation on somebody's behalf -- carrying out a plan, for one. One path rather than two, so a second way of sending cannot drift from the first. If a reply is already being written, the turn is *queued* instead: written, shown, and not sent. Starting a second reply here is what used to happen, and it produced two generations answering the same chat from two different prefixes of it, with Stop pointing at whichever bubble came first in the document. """ if queued := _reply_in_flight(db, chat): waiting = db.scalar( select(func.count()) .select_from(Message) .where(Message.chat_id == chat.id, Message.queued.is_(True)) ) if waiting >= MAX_QUEUED: raise HTTPException( status.HTTP_409_CONFLICT, f"There are already {MAX_QUEUED} messages waiting to be sent.", ) user_message = chat_service.create_message(db, chat, ROLE_USER, content, queued=queued) if file_ids: files_service.claim(db, ids=file_ids, user_id=user.id, message_id=user_message.id) db.refresh(user_message) if queued: # One bubble and no assistant placeholder. The streaming shell is the # only thing that starts a generation, so a placeholder here would be a # second concurrent reply -- exactly what the queue exists to prevent. if (live := generation_service.running_for(chat.id)) is not None: # So a reply between two rounds of tool calls notices it, and so # anybody following sees the status change. live.touch() return templates.TemplateResponse( request, "chat/_message.html", { "request": request, "message": user_message, "chat": chat, "user": user, "models_by_id": { m.model_id: m for m in chat_service.available_models(db, user) }, **audio_service.template_flags(db, user), }, ) assistant_message = chat_service.create_message( db, chat, ROLE_ASSISTANT, "", complete_=False, model_id=chat.model_id ) generation_service.ensure(chat.id, assistant_message.id) # `user` is required by the shared message template, which renders both # roles; without it the user bubble's initial blows up. return templates.TemplateResponse( request, "chat/_turn.html", { "request": request, "user_message": user_message, "assistant_message": assistant_message, "chat": chat, "user": user, "models_by_id": { m.model_id: m for m in chat_service.available_models(db, user) }, **audio_service.template_flags(db, user), }, ) @router.get("/{chat_id}/messages/{message_id}/stream") async def stream_message( db: Db, user: RequiredUser, chat_id: str, message_id: str, ) -> Response: """Stream the assistant's reply as server-sent events. Emits `token` events carrying escaped text, then a single `done` event carrying the finished bubble rendered from Markdown, then `close`. """ chat = _owned_chat(db, chat_id, user.id) message = db.get(Message, message_id) if message is None or message.chat_id != chat.id: raise HTTPException(status.HTTP_404_NOT_FOUND, "That message no longer exists.") return StreamingResponse( _follow(chat.id, message.id), media_type="text/event-stream", headers={ "Cache-Control": "no-cache, no-transform", "Connection": "keep-alive", # nginx buffers proxied responses by default, which turns a stream # into one delivery at the end. This is the documented opt-out. "X-Accel-Buffering": "no", }, ) def _tool_activity(events: list[dict], *, live: bool = True) -> str: """Render the tool block. Whole, never a delta, like every other frame.""" return templates.get_template("chat/_tool_activity.html").render( {"tool_events": events, "live": live} ) def _ask_html(chat_id: str, pending) -> str: """The card asking the reader something, or nothing at all. Returns "" when there is nothing pending, and the frame is sent unconditionally, because this is one of the few blocks that has to be able to *clear* itself: the card must vanish the moment it is answered. `reasoning`, `tools` and `render` are the opposite -- guarded by truthiness so a frame can never blank them. """ if pending is None: return "" return templates.get_template("chat/_interaction.html").render( {"ask": pending, "chat_id": chat_id} ) async def _follow(chat_id: str, message_id: str) -> AsyncIterator[str]: """Stream a generation that is running independently of this request. This connection only *watches*. Closing it -- navigating away, opening another chat -- leaves the reply being written, and reconnecting replays the whole state immediately rather than starting over. Both `render` and `reasoning` carry the complete block each time rather than a delta, which is what makes reattaching mid-reply work at all: a follower arriving late has no earlier fragments to append to. """ generation = generation_service.ensure(chat_id, message_id) generation.followers += 1 seen = -1 last_frame = time.monotonic() try: while True: if generation.version != seen: seen = generation.version if generation.thinking: yield sse.event("reasoning", escape_text(generation.thinking)) if generation.tool_events: yield sse.event("tools", _tool_activity(generation.tool_events)) if generation.content: yield sse.event("render", render_markdown(generation.text)) yield sse.event("metrics", _metrics_html(generation)) yield sse.event("status", escape_text(generation.status)) yield sse.event("ask", _ask_html(chat_id, generation.pending)) last_frame = time.monotonic() if generation.done: 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 # works in RENDER_INTERVAL steps, so a short sleep is simpler and # cannot drop a notification. await asyncio.sleep(generation_service.RENDER_INTERVAL * 0.8) finally: generation.followers = max(0, generation.followers - 1) # 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. with session_scope() as db: message = db.get(Message, message_id) chat = db.get(Chat, chat_id) if message is None or chat is None: yield sse.event("close", "") return owner = db.get(User, chat.user_id) final_html = templates.get_template("chat/_message.html").render( { "message": message, "body_html": render_markdown(message.content), "chat": chat, # Passed even though an assistant bubble never reads it: the # template shares both roles, and a missing `user` would only # blow up on whichever branch is not being exercised here. "user": owner, "models_by_id": { m.model_id: m for m in chat_service.available_models(db, None) }, # This frame replaces the whole bubble, so it has to carry the # speaker button's conditions too -- and the owner's, not the # follower's: there is no request here to ask who is watching. **audio_service.template_flags(db, owner), # The one render that means "this reply just landed", which is # what read-aloud-automatically keys off. A page load must not # set it or reopening a chat would start talking. "just_finished": True, } ) title_html = templates.get_template("chat/_title_oob.html").render({"chat": chat}) # What the queue did while this reply was running. There is no push # channel that outlives one message's stream, and this is the last frame # that reaches the browser -- so it carries the rest out of band, the # way the chat title already does. moved_html, queue_html = _queue_frames(db, chat, owner, generation) yield sse.event("done", moved_html + final_html + queue_html + title_html) yield sse.event("close", "") def _render_bubble(db: DBSession, chat: Chat, owner: User | None, message: Message) -> str: """One finished bubble, rendered the way the `done` frame renders its own.""" return templates.get_template("chat/_message.html").render( { "message": message, "body_html": ( render_markdown(message.content) if message.role == ROLE_ASSISTANT else "" ), "chat": chat, "user": owner, "models_by_id": {m.model_id: m for m in chat_service.available_models(db, None)}, **audio_service.template_flags(db, owner), } ) def _queue_frames( db: DBSession, chat: Chat, owner: User | None, generation ) -> tuple[str, str]: """The bubbles the queue produced during this reply, as out-of-band HTML. Two pieces, because they swap differently. Anything taken *into* this reply mid-round now sorts before it, so it is rendered ahead of the finished bubble in the same `outerHTML` swap and its stale node is deleted out of band -- one frame, and the DOM ends up in the order the database is in. Anything drained *after* the reply is a new pair appended to the thread. """ moved: list[str] = [] out_of_band: list[str] = [] for injected_id in generation.injected_ids: row = db.get(Message, injected_id) if row is None: continue moved.append(_render_bubble(db, chat, owner, row)) # Removed where it was; it is about to reappear above the reply. out_of_band.append(f'') if generation.drained: fresh = list( db.scalars( select(Message) .where(Message.chat_id == chat.id, Message.complete.is_(False)) .order_by(Message.created_at) ) ) for assistant in fresh: # The user turn that was waiting has just lost its Send now and # Discard, so it is re-rendered in place. delivered = db.scalars( select(Message) .where( Message.chat_id == chat.id, Message.role == ROLE_USER, Message.created_at <= assistant.created_at, ) .order_by(Message.created_at.desc()) .limit(1) ).first() if delivered is not None: out_of_band.append( f'