Files
LLeMbas/src/lembas/services/generation.py
T
Jaroslav Beneš a58e48fce5 Say what a tool did, not where it ran
An agent event set its label to the SSH profile's name, so the transcript read
"homeserver · ls -la" -- naming the machine rather than the thing that was done.
Built-in tools set no label at all and fell back to the function name, so a
saved memory read "memory_add". The status line said "Running shell_run…" and
the approval card had its own hand-written wording. Four places, four answers,
nothing checking that any of them agreed.

services/tool_labels.py is the one table all of them read now. Bash, Read,
Write, List, Web search, Memory saved; an icon each, instead of everything
being the sparkle.

The precedence is inverted on purpose. Tool events are persisted in
Message.tool_calls_json, so every agent row already on disk carries the profile
name -- a resolver that preferred the stored value would fix nothing for any
transcript that already exists. So a name the table knows resolves from the
table, and a name it does not -- a custom HTTP tool, an MCP tool, whose labels
are per row and cannot be tabulated -- keeps its own. One rule, both cases
correct. The machine moves to `detail`, where "where this ran" belongs.

tool_label and tool_icon are Jinja globals because a message bubble is rendered
from four handlers, and a fifth thing each of them must remember to pass is a
fifth thing one of them will forget.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-03 10:58:02 +02:00

1261 lines
52 KiB
Python

"""Background reply generation.
Generation used to be driven by the SSE request: the browser opening the stream
was what produced the tokens, so navigating away cancelled the reply mid-
sentence. Here it runs as its own task instead, and the SSE endpoint merely
*follows* it. Closing the page, opening another chat, or starting a new one
leaves the answer being written; coming back attaches to it and immediately
receives everything produced so far.
The registry is in-process, which is right for the single-worker deployment
this ships with. Several workers would need the state in the database or a
broker, because the request that follows a generation would not necessarily
land in the process running it.
"""
from __future__ import annotations
import asyncio
import contextlib
import json
import logging
import time
import uuid
from dataclasses import dataclass, field, replace
from datetime import UTC, datetime, timedelta
from sqlalchemy import select
from lembas.db.models import KIND_AGENT, ROLE_ASSISTANT, ROLE_USER, Chat, Message, User
from lembas.db.session import session_scope
from lembas.services import chat as chat_service
from lembas.services import compaction as compaction_service
from lembas.services import interaction, tokens, tool_labels
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.llm.openai_client import (
LLMError,
chunk_usage,
delta_reasoning,
delta_text,
delta_tool_calls,
stream_chat,
)
from lembas.services.reasoning import REASONING, ReasoningSplitter
from lembas.services.tools import ToolOutcome
log = logging.getLogger(__name__)
# How often the partial answer is offered to followers. Markdown is re-rendered
# whole each time -- a list or a code fence is only correct once its context
# exists -- so this trades a little work for formatting that appears as the
# model writes. 100ms is below the threshold where the eye reads it as stepping.
RENDER_INTERVAL = 0.1
# Finished generations linger so a follower attaching at the last moment still
# gets the final frames, then are pruned.
KEEP_FINISHED = timedelta(minutes=5)
@dataclass
class Generation:
"""The live state of one reply being written."""
chat_id: str
message_id: str
content: list[str] = field(default_factory=list)
reasoning: list[str] = field(default_factory=list)
reasoning_ms: int = 0
# One entry per tool call made while producing this reply, in order. Shown
# live as the model works and kept on the message afterwards.
tool_events: list[dict] = field(default_factory=list)
# --- What it cost --------------------------------------------------------
# Prompt and completion are summed across tool rounds: what the reply cost.
# context_tokens is overwritten each round with that round's prompt plus
# completion, because a three-round reply pays for its prompt three times
# but only ever occupies the window once.
prompt_tokens: int = 0
completion_tokens: int = 0
context_tokens: int = 0
context_limit: int = 0
# Filled from the assembled request before the first chunk, so a follower
# has a percentage to show while the reply is still being written -- real
# usage only arrives in a single chunk at the very end.
prompt_estimate: int = 0
rounds: int = 0
# time.monotonic() at the start. A field rather than a local in `_run`
# because `_follow` is a different function that sees only this object, and
# without it there is nothing to compute a live tokens/second against.
started_at: float = 0.0
elapsed_ms: int = 0
error: str = ""
stopped: bool = False
done: bool = False
# Bumped on every change. Followers compare against it rather than being
# woken individually: with a 100ms cadence a short poll is simpler than
# future bookkeeping, and cannot drop a wakeup.
version: int = 0
# What the reply is doing when it is not producing tokens. Shown in the
# streaming bubble, because a silent multi-second pause before the first
# token is what a hang looks like.
status: str = ""
# Number of browsers currently watching. Decides whether a finished reply
# counts as unread.
followers: int = 0
finished_at: datetime | None = None
cancel: bool = False
# Set while the reply is stopped waiting for a person -- an approval, or a
# question the model asked. None at every other moment. Read by `_follow`,
# which sends the card, and by `request_stop`, which resolves it: `cancel`
# is otherwise only ever read between streamed chunks, and there are no
# chunks while this is set.
pending: interaction.Interruption | None = None
# Seconds spent waiting for a person, cumulative. Taken off the wall-clock
# budget so that thinking time is the model's and not the reader's.
waited: float = 0.0
# How much tool output this reply has handed back, against the agent budget.
# A model that fills its own context with build logs has no room left to
# answer with.
output_bytes: int = 0
# A plan proposed in Plan mode: {"title": str, "steps": [str, ...]}. Ends
# the reply and is written onto the message, so the Execute button sends
# exactly what was proposed rather than something parsed back out of prose.
plan: dict | None = None
# The queue, seen from the reply's side. `drained` says this reply's ending
# handed the next waiting prompt to a fresh one; `injected_ids` names the
# prompts taken into *this* reply between two rounds of tool calls. Both are
# read only by `_follow`, which turns them into bubbles on the `done` frame
# -- the one frame that reaches a browser after a reply is over.
drained: bool = False
injected_ids: list[str] = field(default_factory=list)
def touch(self) -> None:
self.version += 1
@property
def text(self) -> str:
return "".join(self.content)
@property
def thinking(self) -> str:
return "".join(self.reasoning)
_RUNNING: dict[str, Generation] = {}
_TASKS: dict[str, asyncio.Task] = {}
def get(message_id: str) -> Generation | None:
return _RUNNING.get(message_id)
def request_stop(message_id: str) -> bool:
"""Ask a running generation to stop. Returns whether one was found."""
generation = _RUNNING.get(message_id)
if generation is None or generation.done:
return False
generation.cancel = True
# A paused reply produces no chunks, and the chunk loop is the only place
# `cancel` is ever read -- so without this, Stop does nothing at all while
# an approval card is on screen. Resolving the pause is the wakeup; `_run`
# then takes its ordinary stopped path rather than needing a second branch.
if generation.pending is not None:
generation.pending.resolve(interaction.CANCELLED)
return True
def answer(
chat_id: str,
interaction_id: str,
*,
verdict: str = "",
answers: dict[str, str] | None = None,
) -> bool:
"""Resolve whichever running reply is parked on this interruption.
A linear scan of the registry: it holds one entry per reply in flight, and
this runs at human speed. Scoped to the chat because the caller has already
checked that this reader owns *that* chat, and an id alone would not.
"""
for generation in _RUNNING.values():
pending = generation.pending
if generation.chat_id != chat_id or pending is None or pending.id != interaction_id:
continue
outcome = verdict if verdict in _VERDICTS else interaction.ANSWER
return pending.resolve(outcome, answers=answers)
return False
def running_for(chat_id: str) -> Generation | None:
"""The reply being written in this chat, if there is one.
A linear scan for the reason `answer` gives above: one entry per reply in
flight, consulted at human speed. `_prune` first, because a finished
generation lingers `KEEP_FINISHED` so that late followers still get the
final frames -- and without the sweep those five minutes would look like a
chat that is permanently busy, and queue everything typed into it.
"""
_prune()
for generation in _RUNNING.values():
if generation.chat_id == chat_id and not generation.done:
return generation
return None
_VERDICTS = (interaction.ALLOW, interaction.ALLOW_ALWAYS, interaction.DENY)
def _prune() -> None:
cutoff = datetime.now(UTC) - KEEP_FINISHED
now = time.monotonic()
for message_id, generation in list(_RUNNING.items()):
# A paused reply is deliberately not `done` -- a page reload has to be
# able to reattach to it. Its timeout is what stops it lingering, and
# this is the belt to that pair of braces: a deadline long past means
# the timeout did not fire, and a task parked forever is worse than one
# that gives up.
pending = generation.pending
if pending is not None and now > pending.expires_at + KEEP_FINISHED.total_seconds():
log.warning("resolving a stuck interaction on message %s", message_id)
pending.resolve(interaction.EXPIRED)
if generation.done and generation.finished_at and generation.finished_at < cutoff:
_RUNNING.pop(message_id, None)
_TASKS.pop(message_id, None)
def ensure(chat_id: str, message_id: str) -> Generation:
"""Start generating this reply if it is not already under way.
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.
`_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)
if existing is not None:
return existing
generation = Generation(chat_id=chat_id, message_id=message_id)
_RUNNING[message_id] = generation
_TASKS[message_id] = asyncio.create_task(_run(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:
"""Stop every running generation, keeping what each has produced."""
for task in list(_TASKS.values()):
task.cancel()
for task in list(_TASKS.values()):
with contextlib.suppress(asyncio.CancelledError, Exception):
await task
async def _run(generation: Generation) -> None:
"""Produce one reply, then persist it. Never raises into the task.
A reply is not necessarily one request. When tools are offered and the
model asks to use one, the loop below runs it, appends the result to the
conversation and asks again -- up to tools_service.MAX_ROUNDS times, after
which the model has to answer with what it has. Text produced before a tool
call is kept, so a model that narrates what it is about to look up does not
lose that when the results come back.
"""
splitter = ReasoningSplitter()
started = time.monotonic()
generation.started_at = started
reasoning_started: float | None = None
question = ""
endpoint = model_id = None
needs_title = False
title_prompt = ""
# Bound before the try, because the finally clears the credential on it and
# a chat that has been deleted returns before it would otherwise be set.
tool_context = None
try:
# Before the request is assembled, so build_request is called once and
# what goes out is the compacted conversation -- there is no second
# assembly path. Here rather than in post_message because that route's
# whole contract is to return immediately, and a three-second
# summarisation in front of it would break exactly that.
await _maybe_compact(generation)
# Before the session opens, for the same reason compaction is: the
# listing is an SSH round trip, and holding a database session across
# one to save opening a second is the wrong trade. `build_request`
# below reads whatever this left in the cache and never fetches.
await _warm_index(generation)
with session_scope() as db:
chat = db.get(Chat, generation.chat_id)
message = db.get(Message, generation.message_id)
if chat is None or message is None:
generation.error = "That chat no longer exists."
return
endpoint, model_id = chat_service.resolve_endpoint(db, chat)
owner = db.get(User, chat.user_id)
# Read while the session is open: everything below outlives it.
# Resolved once, so that what the loop is allowed to *run* is the
# same set the endpoint was *offered* -- not whatever happens to
# exist by the time a call comes back.
toolset = tools_service.resolve_tools(db, chat, owner)
offered = toolset.schemas
payload = chat_service.build_request(
db, chat, upto=message, tools=offered, user=owner
)
question = _question_from(payload)
needs_title = not chat.title_generated
# Read here, with the rest, because titling happens after this
# session has closed and must not open another one.
title_prompt = prompts_service.resolve(db, "task.title")
tool_context = tools_service.context_for(db, owner, chat, tools=toolset)
model = chat_service.model_for(db, chat)
generation.context_limit = model.context_length if model is not None else 0
# Kept for `_inject`, which builds a user turn after this session
# has closed. A turn taken in mid-reply has to be shaped exactly as
# the same words typed a moment later would have been -- images to a
# vision model, a plain string to anything else, or the endpoint
# rejects the whole request.
vision = chat_service.model_supports(db, chat, "vision")
generation.prompt_estimate = tokens.estimate_request(payload)
limits = tool_context.agent.limits if tool_context.agent else None
budget = limits.steps if limits else tools_service.MAX_ROUNDS
for round_number in range(budget + 1):
generation.rounds = round_number + 1
# Checked between rounds, never mid-stream: cutting a reply off in
# the middle of a sentence to enforce a budget produces garbage, and
# Stop already covers the mid-stream case. Time spent waiting for a
# person is subtracted -- somebody who thinks for ten minutes about
# one command should not thereby spend the whole allowance.
if limits is not None and round_number:
spent = (time.monotonic() - started) - generation.waited
if spent > limits.wall_seconds:
_gave_up(generation, f"after {spent / 60:.0f} minutes")
break
if generation.output_bytes > limits.output_bytes:
_gave_up(generation, "with too much output to read")
break
accumulator = tools_service.ToolCallAccumulator()
# Text the model produced in *this* round, needed separately from
# generation.content when echoing the assistant turn back.
round_text: list[str] = []
async for chunk in stream_chat(endpoint, payload):
counts = chunk_usage(chunk)
if counts is not None:
generation.prompt_tokens += counts.get("prompt_tokens", 0)
generation.completion_tokens += counts.get("completion_tokens", 0)
# Overwritten, not summed: this round's prompt already
# contains every earlier round.
generation.context_tokens = counts.get("prompt_tokens", 0) + counts.get(
"completion_tokens", 0
)
generation.touch()
thought = delta_reasoning(chunk)
if thought:
if reasoning_started is None:
reasoning_started = time.monotonic()
generation.reasoning.append(thought)
generation.touch()
if offered:
fragments = delta_tool_calls(chunk)
if fragments:
accumulator.feed(fragments)
text = delta_text(chunk)
if text:
for kind, piece in splitter.feed(text):
if kind == REASONING:
if reasoning_started is None:
reasoning_started = time.monotonic()
generation.reasoning.append(piece)
else:
if reasoning_started is not None and not generation.reasoning_ms:
generation.reasoning_ms = int(
(time.monotonic() - reasoning_started) * 1000
)
generation.content.append(piece)
round_text.append(piece)
generation.touch()
if generation.cancel:
generation.stopped = True
break
# Let followers and other tasks run between chunks.
await asyncio.sleep(0)
calls = accumulator.calls
if generation.stopped or not calls:
break
if round_number == budget:
# Out of rounds with the model still asking for tools. Recorded
# rather than silently dropped: an answer that stops here needs
# to be explicable.
#
# `budget`, not `MAX_ROUNDS`. The loop is sized by the budget on
# the line above and the message below has always reported it,
# but the comparison was against the global 3 -- so an agent
# chat allowed forty steps stopped after three and said it had
# taken forty. Two numbers, one of them wrong, in code whose
# whole job is to say what happened.
generation.tool_events.append(
{
"name": calls[0]["name"],
"status": "error",
"error": (
f"Stopped after {budget} rounds of tool calls "
f"without an answer."
),
}
)
generation.touch()
break
messages = [
*payload["messages"],
tools_service.assistant_turn(calls, "".join(round_text)),
]
# Decided before anything runs, never during. A round's calls run
# together under a semaphore, and four people-shaped pauses inside
# that gather would queue behind each other invisibly -- see
# services/interaction.py.
decided, allowed = await _authorise(generation, tool_context, calls)
if generation.stopped:
break
generation.status = _tool_status(calls)
generation.touch()
try:
outcomes = await _run_calls(
tool_context, calls, decided=decided, allowed=allowed
)
finally:
generation.status = ""
generation.touch()
for call, outcome in zip(calls, outcomes, strict=True):
generation.tool_events.append(outcome.event)
generation.output_bytes += len(outcome.content)
messages.append(tools_service.tool_turn(call, outcome.content))
if outcome.event.get("plan"):
generation.plan = outcome.event["plan"]
generation.touch()
# Something typed while this reply was working. Taken in here, at a
# round boundary, rather than made to wait for the whole reply: an
# agent that has just finished one loop and is about to start
# another is exactly when "actually, do it the other way" is worth
# having.
#
# Only while there is a round left to answer in. Injecting into the
# last one would deliver the prompt into a reply that then runs out
# of budget without addressing it -- and it is marked delivered, so
# nothing would ever send it again. Below that line it waits for
# `_drain`, which always gives it a reply of its own.
if round_number + 1 < budget and (
added := _inject(generation, generation.chat_id, vision)
):
messages.append(added)
payload = {**payload, "messages": messages}
# A plan ends the turn. One more request so the model can say what
# it proposed and why -- a bubble containing only a card reads as
# though it had nothing to add -- but with the tools withdrawn, so
# "one more round" cannot become three rounds of it changing its
# mind about a plan the reader is being asked to approve.
if generation.plan is not None:
offered = []
payload.pop("tools", None)
for kind, piece in splitter.flush():
(generation.reasoning if kind == REASONING else generation.content).append(piece)
generation.touch()
except LLMError as exc:
generation.error = exc.message
log.info("generation failed for chat %s: %s", generation.chat_id, exc.message)
except asyncio.CancelledError:
# Shutdown, not a reader navigating away -- that no longer reaches here.
generation.stopped = True
raise
except Exception: # noqa: BLE001 - a task that dies silently is worse
generation.error = "Something went wrong while generating this reply."
log.exception("unexpected generation failure for chat %s", generation.chat_id)
finally:
if reasoning_started is not None and not generation.reasoning_ms:
generation.reasoning_ms = int((time.monotonic() - reasoning_started) * 1000)
generation.elapsed_ms = int((time.monotonic() - started) * 1000)
if not generation.completion_tokens:
# The endpoint reported nothing, so fall back to the estimate. Marked
# as such everywhere it is shown -- four characters to a token is
# wrong enough on code and CJK to be worth saying out loud.
generation.completion_tokens = tokens.estimate(
generation.text + generation.thinking
)
generation.prompt_tokens = generation.prompt_estimate
generation.context_tokens = generation.prompt_tokens + generation.completion_tokens
# Naming the chat is a second, short completion, so it has to happen
# here rather than in the synchronous persist step below. Best-effort:
# a chat title is never worth surfacing an error for.
title = ""
if needs_title and question:
if generation.error or endpoint is None:
title = chat_service.fallback_title(question)
else:
with contextlib.suppress(Exception):
title = await chat_service.generate_title(
endpoint,
model_id,
question,
generation.text,
template=title_prompt,
)
title = title or chat_service.fallback_title(question)
# The decrypted SSH credential dies with the reply rather than with the
# object holding it. A finished Generation lingers KEEP_FINISHED so a
# follower arriving at the last moment still gets the final frames, and
# a private key should not sit in memory for five minutes waiting on
# that.
if tool_context is not None and getattr(tool_context, "agent", None) is not None:
tool_context.agent.clear()
# 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)
# After the row is authoritative and before `done`, for the same reason
# `_persist` is: `_follow` breaks the instant it sees that flag, and the
# frame it then sends is the one that has to carry the next turn's
# bubbles. There is no push channel that outlives a single reply.
_drain(generation)
generation.done = True
generation.finished_at = datetime.now(UTC)
generation.touch()
# How long a reply will wait for a directory listing before starting without
# one. Short on purpose: the listing is a convenience and the reply is the
# thing somebody is waiting for. A walk that outruns this keeps going in the
# background and the next turn has it.
INDEX_WAIT = 6.0
async def _warm_index(generation: Generation) -> None:
"""Build this chat's project listing, or leave whatever is cached.
Never raises and never blocks for long. `harness` reads the cache
synchronously while assembling the system message, so something has to fill
it, and this is the one place in a reply's life that is both asynchronous
and already doing network work.
The first reply in a brand-new chat on a big tree may start before the walk
finishes. That is deliberate: the fragment carrying the listing vanishes
when it is empty rather than appearing as a heading with nothing under it,
and by the following turn it is there.
"""
from lembas.services import settings_store
from lembas.services.agent import index as index_service
from lembas.services.agent import session as agent_session
try:
with session_scope() as db:
if not settings_store.agents(db).get("index_enabled"):
return
chat = db.get(Chat, generation.chat_id)
if chat is None or chat.kind != KIND_AGENT:
return
owner = db.get(User, chat.user_id)
context = agent_session.resolve(db, chat, owner)
profile_id = chat.ssh_profile_id or ""
if context is None or not profile_id:
return
if index_service.cached(profile_id, context.project_dir) is not None:
return
await asyncio.wait_for(
index_service.ensure(context.executor(), profile_id, context.project_dir),
timeout=INDEX_WAIT,
)
except TimeoutError:
log.debug("index for chat %s outran its wait; carrying on", generation.chat_id)
except Exception as exc: # noqa: BLE001 - a missing listing is not a failed reply
log.info("could not warm the index for chat %s: %s", generation.chat_id, exc)
async def _maybe_compact(generation: Generation) -> None:
"""Summarise the earlier turns if the window is about to be full.
Never raises. A failed compaction logs and sends the uncompacted request,
which either works or fails upstream with a message that says what actually
happened -- refusing to answer because the summariser was unavailable would
be a worse trade.
The awaited call is deliberately outside any session, the same shape titling
uses: read everything needed, close, ask, reopen to write.
"""
try:
with session_scope() as db:
chat = db.get(Chat, generation.chat_id)
message = db.get(Message, generation.message_id)
if chat is None or message is None:
return
pending = _pending_text(db, message)
if not compaction_service.should_compact(db, chat, pending=pending):
return
template = prompts_service.resolve(db, "task.compact")
upto = compaction_service.last_complete(db, chat)
if not template.strip() or upto is None:
return
endpoint, model_id = chat_service.resolve_endpoint(db, chat)
transcript = compaction_service.transcript(db, chat, upto=upto)
previous = compaction_service.previous_summary_block(chat)
upto_id = upto.id
generation.status = "Summarising earlier messages…"
generation.touch()
summary = await chat_service.summarise_for_compaction(
endpoint,
model_id,
transcript=transcript,
previous_summary=previous,
template=template,
)
if not summary:
return
with session_scope() as db:
chat = db.get(Chat, generation.chat_id)
upto = db.get(Message, upto_id)
if chat is None or upto is None:
return
compaction_service.apply(chat, summary=summary, upto=upto)
db.commit()
log.info("chat %s compacted automatically through %s", chat.id, upto_id)
except Exception: # noqa: BLE001 - the reply matters more than the tidy-up
log.exception("automatic compaction failed for chat %s", generation.chat_id)
finally:
generation.status = ""
generation.touch()
# How many of a round's tool calls may be in flight at once. A bound rather
# than none: a model that asks for eight would otherwise open eight sockets and
# eight database sessions at the same moment.
MAX_PARALLEL_TOOLS = 4
def _gave_up(generation, why: str) -> None:
"""Stop, and leave something in the transcript saying why.
A reply that simply stopped would look like the model losing interest. The
event is the same shape the out-of-rounds branch uses, so it renders with
everything else.
"""
generation.tool_events.append(
{
"name": "budget",
"kind": "agent",
"status": "error",
"results": [],
"error": f"Stopped {why}. Ask again to carry on from here.",
}
)
generation.touch()
def _written(generation: Generation) -> int:
"""How much this reply has written so far, in tokens, reported or estimated.
Both, because neither alone is enough. `completion_tokens` is only populated
when the endpoint sends a usage block, and a good half of the ones this
talks to -- llama.cpp, Ollama and friends -- never do; the fallback estimate
is otherwise computed once, in `_run`'s `finally:`, long after the loop that
needs it. A ceiling reading only the reported figure would work on OpenAI
and silently do nothing everywhere else, which is the worst kind of limit:
one that looks configured.
Reasoning counts. It was generated and it was paid for, even though it is
deliberately never replayed as context.
"""
return max(
generation.completion_tokens,
tokens.estimate(generation.text + generation.thinking),
)
def _tool_status(calls: list[dict]) -> str:
"""What to show while tools run.
A remote tool -- an HTTP endpoint, an MCP server -- can take seconds with
nothing streaming, and a silent pause is exactly what a hang looks like.
"""
if len(calls) == 1:
return f"Running {tool_labels.label_for(calls[0]['name'])}…"
return f"Running {len(calls)} tools…"
def _arguments_of(call: dict) -> dict:
try:
args = json.loads(call["arguments"] or "{}")
except json.JSONDecodeError:
return {}
return args if isinstance(args, dict) else {}
def _describe(name: str, args: dict) -> tuple[str, str]:
"""What an approval card says about one call: a title, and the detail.
The detail is the thing being agreed to -- the command line, the path -- and
is shown verbatim and escaped. A summary that paraphrased it would be a card
approving something other than what runs.
Delegated to services/tool_labels.py, which the transcript and the status
line read too. This used to be a hand-written if-chain and was the fourth
place with its own wording for the same tool.
"""
return tool_labels.describe(name, args)
def _approvals(context, calls: list[dict]) -> list[interaction.Item]:
"""The calls in this round that a person has to allow before they run.
Only in an agent chat: `context.agent` is None everywhere else, and an
ordinary conversation behaves exactly as it did. Within one, *every* call
goes through the table, including the built-in ones -- `notes_edit` writes,
and Plan mode meaning "look but do not touch" has to mean that too.
"""
agent = getattr(context, "agent", None)
if agent is None:
return []
book = context.tools if context.tools is not None else tools_service.REGISTRY
items: list[interaction.Item] = []
for index, call in enumerate(calls):
tool = book.get(call["name"])
if tool is None or tool.risk == tools_service.RISK_ASK:
continue # unknown names are refused by run_tool; questions are their own card
args = _arguments_of(call)
command = str(args.get("command") or "") if call["name"] == "shell_run" else ""
decision = agent_policy.decide(
mode=agent.mode,
risk=tool.risk,
tool_name=call["name"],
command=command,
allow=agent.allow,
deny=agent.deny,
)
if decision.verdict == agent_policy.ALLOW:
continue
title, detail = _describe(call["name"], args)
items.append(
interaction.Item(
index=index,
key=f"a{index}",
kind=interaction.KIND_APPROVAL,
tool_name=call["name"],
title=f"{title} on {agent.label}",
detail=detail,
reason=decision.reason,
)
)
return items
def _ask_items(context, calls: list[dict]) -> list[interaction.Item]:
"""Which of this round's calls need a person, and what to show about each.
Looked up through `context.tools`, the map of what was actually offered --
the same authority `run_tool` uses. A name that is not in it is left alone
here and refused there, so an unknown tool cannot smuggle itself past by
being unclassifiable.
"""
book = context.tools if context.tools is not None else tools_service.REGISTRY
items: list[interaction.Item] = []
for index, call in enumerate(calls):
tool = book.get(call["name"])
if tool is None or tool.risk != tools_service.RISK_ASK:
continue
args = _arguments_of(call)
for asked in _questions_in(args):
options = [str(o).strip() for o in (asked.get("options") or []) if str(o).strip()]
items.append(
interaction.Item(
index=index,
key=f"q{len(items)}",
kind=interaction.KIND_QUESTION,
tool_name=call["name"],
title=str(asked.get("question") or "").strip() or "A question for you",
options=tuple(options[: interaction.MAX_OPTIONS]),
)
)
return items
def _questions_in(args: dict) -> list[dict]:
"""The questions in one `ask_user` call, however it was spelled.
The schema asks for a list of objects, and a capable model sends that. A
small one sends a bare `question` string, or a list of plain strings, or
one object where a list belonged -- all of which mean something obvious, so
they are read rather than refused. Getting this wrong costs a whole round
trip and produces a card saying "A question for you" and nothing else.
"""
raw = args.get("questions")
if raw is None:
raw = args.get("question")
if raw is None:
return []
if isinstance(raw, str | dict):
raw = [raw]
if not isinstance(raw, list):
return []
out: list[dict] = []
for entry in raw[: interaction.MAX_QUESTIONS]:
if isinstance(entry, str) and entry.strip():
# A bare string, possibly alongside a sibling `options` that was
# meant to go with it -- which only makes sense for a lone question.
out.append({"question": entry, "options": args.get("options") if len(raw) == 1 else []})
elif isinstance(entry, dict) and str(entry.get("question") or "").strip():
out.append(entry)
return out
async def _authorise(
generation, context, calls: list[dict]
) -> tuple[dict[int, ToolOutcome], set[int]]:
"""Which of this round's calls may run, and what the others answer instead.
Returns outcomes keyed by the call's index. Every index the caller does not
find here is cleared to run; every index it does find is answered without
the runner being reached at all. That is what keeps
`zip(calls, outcomes, strict=True)` aligned -- an endpoint matching on
`tool_call_id` pairs the wrong content with the right id otherwise.
Also returns the indices a person explicitly allowed, so the runners can be
told. They re-check the mode as a backstop and would otherwise refuse the
very thing that was just approved -- the mode says "ask", and asking is what
happened.
"""
questions = _ask_items(context, calls)
approvals = _approvals(context, calls)
items = [*approvals, *questions]
if not items:
return {}, set()
timeout = float(context.interaction_timeout or 900)
pause = interaction.build(uuid.uuid4().hex, items, timeout=timeout)
generation.status = interaction.summarise(pause.items)
reply = await interaction.wait_for(generation, pause, timeout=timeout)
generation.status = ""
if reply.ended:
generation.stopped = True
return {}, set()
decided: dict[int, ToolOutcome] = {}
allowed: set[int] = set()
# An approval that came back as a refusal answers its call without the
# runner being reached; one that came back allowed is simply left out, which
# is how `_run_calls` is told to go ahead.
for item in approvals:
if reply.permitted:
allowed.add(item.index)
continue
decided[item.index] = _not_allowed(item, reply)
# Questions are grouped back by call, because one `ask_user` call may have
# carried several and the endpoint expects exactly one tool turn per call.
grouped: dict[int, list[interaction.Item]] = {}
for item in questions:
grouped.setdefault(item.index, []).append(item)
for index, asked in grouped.items():
decided[index] = _answered(asked, reply)
return decided, allowed
def _not_allowed(item: interaction.Item, reply: interaction.Reply) -> ToolOutcome:
"""What the model is told when a person declined, or never answered.
Told plainly, and told to stop rather than to try again: a model that reads
"not allowed" as "not allowed *that way*" will spend the rest of the reply
looking for a way round, which is the opposite of what the refusal meant.
"""
event = {
"name": item.tool_name,
"kind": "agent",
"label": item.title,
"query": item.detail,
"results": [],
}
if reply.outcome == interaction.EXPIRED:
return ToolOutcome(
"Nobody answered, so this was not run. Stop and say what you were "
"about to do and why.",
{**event, "status": "error", "error": "Not answered."},
)
return ToolOutcome(
"They declined this. Do not try it another way — say what you were "
"going to do and ask what they would prefer.",
{**event, "status": "error", "error": "Declined."},
)
def _answered(items: list[interaction.Item], reply: interaction.Reply) -> ToolOutcome:
"""What one `ask_user` call gets back, however many questions it put."""
event = {
"name": items[0].tool_name,
"kind": "ask",
"label": "Asked you",
"query": "; ".join(item.title for item in items),
"results": [],
}
if reply.outcome == interaction.EXPIRED:
return ToolOutcome(
"They did not answer. Carry on as best you can without it, or say "
"what you still need.",
{**event, "status": "error", "error": "No answer.", "text": ""},
)
answered = [(item, reply.answer_to(item)) for item in items]
given = [(item, text) for item, text in answered if text]
if not given:
return ToolOutcome(
"They closed the question without answering.",
{**event, "status": "error", "error": "No answer.", "text": ""},
)
# Each answer is quoted next to the question it belongs to. With four
# questions on one card, a bare list of answers would leave the model
# matching them up by position and sometimes getting it wrong.
lines = [f"{item.title}\n{text}" for item, text in given]
skipped = [item for item, text in answered if not text]
if skipped:
lines.append(
"They left unanswered: " + "; ".join(item.title for item in skipped)
)
body = "\n\n".join(lines)
return ToolOutcome(f"They answered:\n\n{body}", {**event, "status": "ok", "text": body})
async def _run_calls(
context,
calls: list[dict],
*,
decided: dict[int, ToolOutcome] | None = None,
allowed: set[int] | None = None,
) -> list:
"""Run one round's calls together, results in call order.
Sequential was right when every tool was a local database read. A remote one
is latency-bound, and three two-second calls in a row are six seconds of a
reply looking hung -- while the model has already been told it may ask for
several at once.
The results are indexed rather than appended as they finish, because each
tool turn has to line up with the assistant turn's `tool_calls`: an endpoint
matching on `tool_call_id` would otherwise pair the right id with the wrong
content the moment two calls came back out of order.
Safe to run together because `run_tool` never raises, so no failure cancels
its siblings, and each runner opens its own `session_scope()` against a
database in WAL mode with a busy timeout.
"""
limit = asyncio.Semaphore(MAX_PARALLEL_TOOLS)
async def one(index: int, call: dict):
# Already answered by a person, or refused before it got here. It still
# occupies its index, because the tool turns have to line up.
if decided and index in decided:
return decided[index]
# A per-call copy for anything a person allowed, so the runner's own
# check does not undo their decision. A copy rather than a flag on the
# shared context, because a round runs its calls together and only some
# of them were approved.
ctx = context
if allowed and index in allowed and getattr(context, "agent", None) is not None:
ctx = replace(context, agent=context.agent.as_approved())
async with limit:
return await tools_service.run_tool(ctx, call["name"], call["arguments"])
return list(await asyncio.gather(*(one(i, c) for i, c in enumerate(calls))))
def _pending_text(db, message: Message) -> str:
"""The user turn this reply is answering, for the size estimate."""
previous = db.scalars(
select(Message)
.where(Message.chat_id == message.chat_id, Message.created_at < message.created_at)
.order_by(Message.created_at.desc())
.limit(1)
).first()
return previous.content if previous is not None else ""
def _question_from(payload: dict) -> str:
"""The last thing the user said, for auto-titling."""
for entry in reversed(payload.get("messages", [])):
if entry.get("role") != ROLE_USER:
continue
content = entry.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
return " ".join(
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
).strip()
return ""
def _next_waiting(db, chat_id: str) -> Message | None:
"""The oldest prompt in this chat that has not been sent."""
return db.scalars(
select(Message)
.where(
Message.chat_id == chat_id,
Message.role == ROLE_USER,
Message.queued.is_(True),
)
.order_by(Message.created_at)
.limit(1)
).first()
def _drain(generation: Generation) -> None:
"""Hand the next waiting prompt to a reply of its own, if there is one.
Exactly one, not all of them. Draining the lot would put two consecutive
user turns into the next request, which several local chat templates refuse
outright -- `build_messages` already goes to some trouble over that around
the compaction lead. "One after another" is also what was asked for: the
second waiting prompt is drained by the reply the first one starts, and so
on down the chain.
Three refusals, and none of them is a special case:
- **Superseded.** The same test `_persist` makes, for the same reason: a
regeneration cancels its predecessor and the predecessor's `finally:`
still runs. Without this, regenerating would drain the queue *and* leave
a third generation running.
- **Stopped.** Stop means stop, and the queue stays visible and
undelivered with Send now beside it. This is also what makes shutdown
safe -- cancellation sets `stopped`, so a restart never fires off a reply
with nobody watching.
- **Errored.** The endpoint has just failed. Feeding the next prompt into it
produces a second failure and spends somebody's words to do it.
"""
owner = _RUNNING.get(generation.message_id)
if owner is not None and owner is not generation:
return
if generation.stopped or generation.error:
return
try:
with session_scope() as db:
chat = db.get(Chat, generation.chat_id)
if chat is None:
return
waiting = _next_waiting(db, chat.id)
if waiting is None:
return
waiting.queued = False
assistant = chat_service.create_message(
db, chat, ROLE_ASSISTANT, "", complete_=False, model_id=chat.model_id
)
chat_id, assistant_id = chat.id, assistant.id
except Exception: # noqa: BLE001 - the reply is over either way
log.exception("could not drain the queue for chat %s", generation.chat_id)
return
# Outside the session: this starts a task, and a task is not something to
# hold a database session open across.
ensure(chat_id, assistant_id)
generation.drained = True
def _inject(generation: Generation, chat_id: str, vision: bool) -> dict | None:
"""Take the oldest waiting prompt into this reply, between two rounds.
Marked delivered and committed *before* the request goes out, so this is
at-most-once. A crash in between loses the turn, which is recoverable --
the words are still in the transcript with Send now beside them. The other
way round would ask the same question twice and let an agent act on it
twice, which is not.
Sent verbatim, in the user role, with no framing. Everything else this
codebase injects is quoted and attributed because it came out of a file, a
page or a machine; this one genuinely *is* the person at the keyboard,
authenticated by the session cookie and stored as a `Message` whose role
says so. Wrapping it would teach a model that a user turn can be a
quotation, which is the exact distinction the other two rely on. What the
model needs -- that this can happen at all -- is one sentence in the
harness, where authored wording lives.
"""
try:
with session_scope() as db:
waiting = _next_waiting(db, chat_id)
if waiting is None:
return None
waiting.queued = False
entry = chat_service.message_payload(waiting, vision=vision)
# The reply that answers it must sort *before* it, or the next
# turn's transcript reads "answer, then the question it answered"
# and a small model dutifully answers again. Moving the placeholder
# rather than the prompt keeps several interjections in the order
# they were typed.
placeholder = db.get(Message, generation.message_id)
if placeholder is not None:
placeholder.created_at = datetime.now(UTC)
generation.injected_ids.append(waiting.id)
except Exception: # noqa: BLE001 - a lost interjection is not a failed reply
log.exception("could not take a queued prompt into chat %s", chat_id)
return None
generation.status = "Taking in what you just added…"
generation.touch()
return entry
def _persist(generation: Generation, title: str, elapsed: float) -> None:
"""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:
with session_scope() as db:
message = db.get(Message, generation.message_id)
chat = db.get(Chat, generation.chat_id)
if message is None or chat is None:
return
message.content = generation.text
message.reasoning = generation.thinking
message.reasoning_ms = generation.reasoning_ms
message.tool_calls_json = generation.tool_events
message.plan_json = generation.plan or {}
message.usage_json = metrics_service.to_json(
metrics_service.from_generation(generation)
)
message.error = generation.error
message.stopped = generation.stopped
message.complete = True
if title and not chat.title_generated:
chat.title = title
chat.title_generated = True
# Nobody watching when it landed, so it is news. The chat page
# clears this when it is next opened. Not for a temporary chat:
# there is no sidebar row for the dot, and the toast would name a
# chat nobody can navigate to.
if generation.followers == 0 and not chat.temporary:
chat.unread = True
chat.unread_notified = False
db.commit()
log.debug(
"chat %s finished: %d chars, %d reasoning, %.1fs",
generation.chat_id,
len(message.content),
len(message.reasoning),
elapsed,
)
except Exception: # noqa: BLE001 - the task is ending either way
log.exception("could not persist generation for chat %s", generation.chat_id)
__all__ = [
"RENDER_INTERVAL",
"Generation",
"ROLE_ASSISTANT",
"ensure",
"get",
"request_stop",
"restart",
"shutdown",
]