Jaroslav Beneš
b6aab8de55
The four tools an agent chat has -- shell_run, file_read, file_write, file_list -- and the mode table wired into the loop that decides which of them stop for approval. Verified end to end against a real Kali container over SSH: the card shows the command, allowing it runs it there, and the file it writes is visible from outside. The mode is enforced in `_authorise`, in the generation loop, server-side, keyed on each tool's declared risk. Not in the prompt: a model is told which mode it is in so it behaves sensibly, but everything it reads -- a web page, a README, the output of the last command -- is untrusted, and a rule written only into a system message is one a poisoned file can argue with. Within an agent chat every call goes through the table, including the built-in ones, because notes_edit writes and Plan mode meaning "look but do not touch" has to mean that too. Two things this turned up. The runners re-check the mode as a backstop, and that backstop refused the very thing a person had just approved -- the mode says "ask", and asking was exactly what happened. Approval is now threaded per call, on a copy of the context, because a round runs its calls together and only some of them were allowed. And the harness said nothing at all, because `registry` maps an offered tool *name* back to a family and did not know the agent tools existed. So shell_run resolved to no family and the fragment naming the machine, the directory and the mode was never admitted. The same omission cost custom tools their guidance once already; there is a test for it now. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1007 lines
40 KiB
Python
1007 lines
40 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 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
|
|
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
|
|
|
|
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
|
|
|
|
|
|
_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)
|
|
|
|
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
|
|
|
|
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 == tools_service.MAX_ROUNDS:
|
|
# Out of rounds with the model still asking for tools. Recorded
|
|
# rather than silently dropped: an answer that stops here needs
|
|
# to be explicable.
|
|
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))
|
|
generation.touch()
|
|
|
|
payload = {**payload, "messages": messages}
|
|
|
|
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)
|
|
generation.done = True
|
|
generation.finished_at = datetime.now(UTC)
|
|
generation.touch()
|
|
|
|
|
|
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 _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 {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.
|
|
"""
|
|
if name == "shell_run":
|
|
return "Run a command", str(args.get("command") or "")
|
|
if name == "file_write":
|
|
return "Write a file", str(args.get("path") or "")
|
|
if name == "file_read":
|
|
return "Read a file", str(args.get("path") or "")
|
|
if name == "file_list":
|
|
return "List a directory", str(args.get("path") or "")
|
|
detail = ", ".join(f"{k}={v!r}" for k, v in list(args.items())[:4])
|
|
return f"Use {name}", detail[:400]
|
|
|
|
|
|
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 _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.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",
|
|
]
|