Jaroslav Beneš
1c659a5640
Three features turn out to be one mechanism: a command waiting to be approved, a question the model wants answered, and "this reply is waiting for you" are all — stop the generation, put an interactive block in the bubble, wait for a POST, carry on. So there is one primitive, and the only thing using it so far is `ask_user`: a model can offer you a few answers and a box to write your own. The shell executor is not here yet. This lands first on purpose, because it is the riskiest machinery in the feature and it is worth having working before any subprocess exists to complicate it. Two things about where the pause sits. It pauses a round, not a call: a round's calls run together under a semaphore, and parking four coroutines on four separate answers inside that gather would queue them behind each other invisibly. And Stop had to be taught about it — `cancel` is read between streamed chunks and there are no chunks while paused, so the button did nothing at all until `request_stop` learned to resolve the pause itself. Also here: a risk class on every tool (read, write, execute), which is what the four permission modes will be a table over, and the systemd unit loses ProtectKernelTunables. That last one is not tidying — it bind-mounts /proc/sys read-only, which stops bubblewrap mounting /proc at all, and the obvious workaround would expose this process's environment and with it the encryption key. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
430 lines
16 KiB
Python
430 lines
16 KiB
Python
"""Pausing a reply to ask the reader something.
|
|
|
|
Driven through the real generation loop with a scripted endpoint, because the
|
|
things worth pinning here are all about the loop: that Stop still works while
|
|
nothing is streaming, that a decision lands at the right index, and that the
|
|
card clears itself.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
from sqlalchemy import select
|
|
|
|
from lembas.db.models import ROLE_ASSISTANT, Chat, Connection, Message, Model
|
|
from lembas.services import generation as generation_service
|
|
from lembas.services import interaction, settings_store
|
|
|
|
|
|
def _chat_that_can_ask(db, user_id):
|
|
connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted="")
|
|
db.add(connection)
|
|
db.commit()
|
|
db.add(Model(connection_id=connection.id, model_id="m", capabilities_json={"tools": True}))
|
|
db.commit()
|
|
chat = Chat(user_id=user_id, model_id="m", connection_id=connection.id)
|
|
db.add(chat)
|
|
db.commit()
|
|
db.add(Message(chat_id=chat.id, role="user", content="Which one?", complete=True))
|
|
db.commit()
|
|
assistant = Message(chat_id=chat.id, role=ROLE_ASSISTANT, content="", complete=False)
|
|
db.add(assistant)
|
|
db.commit()
|
|
return chat.id, assistant.id
|
|
|
|
|
|
def _ask_chunk(question: str, options: list[str] | None = None, *, index: int = 0, call_id="c1"):
|
|
import json as _json
|
|
|
|
arguments = {"question": question}
|
|
if options is not None:
|
|
arguments["options"] = options
|
|
return {
|
|
"choices": [
|
|
{
|
|
"delta": {
|
|
"tool_calls": [
|
|
{
|
|
"index": index,
|
|
"id": call_id,
|
|
"function": {
|
|
"name": "ask_user",
|
|
"arguments": _json.dumps(arguments),
|
|
},
|
|
}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
|
|
def _text_chunk(text: str) -> dict:
|
|
return {"choices": [{"delta": {"content": text}}]}
|
|
|
|
|
|
def _stub_stream(rounds, seen):
|
|
async def stream_chat(_endpoint, payload):
|
|
seen.append(payload)
|
|
for chunk in rounds[min(len(seen) - 1, len(rounds) - 1)]:
|
|
yield chunk
|
|
|
|
return stream_chat
|
|
|
|
|
|
async def _until_paused(generation, *, timeout: float = 2.0):
|
|
"""Wait for the card to go up."""
|
|
deadline = asyncio.get_running_loop().time() + timeout
|
|
while asyncio.get_running_loop().time() < deadline:
|
|
if generation.pending is not None:
|
|
return generation.pending
|
|
await asyncio.sleep(0.01)
|
|
raise AssertionError("the reply never paused")
|
|
|
|
|
|
@pytest.fixture
|
|
def scripted(db, user_id, monkeypatch):
|
|
"""A reply that asks one question, then answers with whatever it was told."""
|
|
chat_id, message_id = _chat_that_can_ask(db, user_id)
|
|
payloads: list[dict] = []
|
|
monkeypatch.setattr(
|
|
generation_service,
|
|
"stream_chat",
|
|
_stub_stream(
|
|
[[_ask_chunk("Tea or coffee?", ["Tea", "Coffee"])], [_text_chunk("Right you are.")]],
|
|
payloads,
|
|
),
|
|
)
|
|
generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
|
|
return generation, payloads, chat_id, message_id
|
|
|
|
|
|
# --- The tool is offered at all ----------------------------------------------
|
|
def test_ask_user_is_offered_to_an_ordinary_chat(db, user_id):
|
|
"""Not an agent feature. A model in a plain conversation should be able to
|
|
stop and ask which of two things you meant."""
|
|
from lembas.db.models import User
|
|
from lembas.services import tools as tools_service
|
|
|
|
chat_id, _message_id = _chat_that_can_ask(db, user_id)
|
|
chat = db.get(Chat, chat_id)
|
|
offered = tools_service.resolve_tools(db, chat, db.get(User, user_id))
|
|
assert "ask_user" in offered.by_name
|
|
|
|
|
|
def test_ask_user_is_withheld_without_the_permission(db, user_id):
|
|
from lembas.db.models import User
|
|
from lembas.services import tools as tools_service
|
|
|
|
user = db.get(User, user_id)
|
|
# Administrators are given every permission, so the baseline only bites a
|
|
# plain account.
|
|
user.role = "user"
|
|
settings_store.update(db, {"default_permissions": {"tools.ask": False}})
|
|
db.commit()
|
|
|
|
chat = db.get(Chat, _chat_that_can_ask(db, user_id)[0])
|
|
assert "ask_user" not in tools_service.resolve_tools(db, chat, user).by_name
|
|
|
|
|
|
# --- The pause ---------------------------------------------------------------
|
|
async def test_the_reply_pauses_and_the_card_describes_the_question(scripted):
|
|
generation, _payloads, _chat_id, _message_id = scripted
|
|
task = asyncio.create_task(generation_service._run(generation))
|
|
|
|
pending = await _until_paused(generation)
|
|
assert pending.kind == interaction.KIND_QUESTION
|
|
assert pending.items[0].title == "Tea or coffee?"
|
|
assert pending.options == ("Tea", "Coffee")
|
|
assert "Waiting for your answer" in generation.status
|
|
|
|
pending.resolve(interaction.ANSWER, text="Tea")
|
|
await task
|
|
|
|
|
|
async def test_the_answer_reaches_the_model_as_a_tool_result(scripted):
|
|
generation, payloads, _chat_id, _message_id = scripted
|
|
task = asyncio.create_task(generation_service._run(generation))
|
|
|
|
pending = await _until_paused(generation)
|
|
pending.resolve(interaction.ANSWER, text="Coffee, please")
|
|
await task
|
|
|
|
turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"]
|
|
assert len(turns) == 1
|
|
assert "Coffee, please" in turns[0]["content"]
|
|
assert turns[0]["tool_call_id"] == "c1"
|
|
|
|
|
|
async def test_the_card_is_cleared_once_it_is_answered(scripted):
|
|
generation, _payloads, _chat_id, _message_id = scripted
|
|
task = asyncio.create_task(generation_service._run(generation))
|
|
|
|
pending = await _until_paused(generation)
|
|
version = generation.version
|
|
pending.resolve(interaction.ANSWER, text="Tea")
|
|
await task
|
|
|
|
assert generation.pending is None
|
|
assert generation.version > version, "clearing has to bump the version or no frame is sent"
|
|
|
|
|
|
async def test_the_transcript_keeps_what_was_asked_and_answered(scripted):
|
|
generation, _payloads, _chat_id, _message_id = scripted
|
|
task = asyncio.create_task(generation_service._run(generation))
|
|
|
|
pending = await _until_paused(generation)
|
|
pending.resolve(interaction.ANSWER, text="Tea")
|
|
await task
|
|
|
|
event = generation.tool_events[0]
|
|
assert event["kind"] == "ask"
|
|
assert event["query"] == "Tea or coffee?"
|
|
assert event["text"] == "Tea"
|
|
|
|
|
|
# --- Stop, while nothing is streaming ----------------------------------------
|
|
async def test_stop_ends_a_reply_that_is_waiting_for_an_answer(db, user_id, monkeypatch):
|
|
"""`cancel` is read only between streamed chunks, and there are no chunks
|
|
while the card is up. Without the wakeup in request_stop the button does
|
|
nothing at all here."""
|
|
chat_id, message_id = _chat_that_can_ask(db, user_id)
|
|
monkeypatch.setattr(
|
|
generation_service,
|
|
"stream_chat",
|
|
_stub_stream([[_ask_chunk("Tea or coffee?")], [_text_chunk("unreachable")]], []),
|
|
)
|
|
|
|
generation = generation_service.ensure(chat_id, message_id)
|
|
await _until_paused(generation)
|
|
|
|
assert generation_service.request_stop(message_id) is True
|
|
await asyncio.wait_for(asyncio.shield(generation_service._TASKS[message_id]), timeout=2)
|
|
|
|
assert generation.stopped is True
|
|
assert generation.pending is None
|
|
|
|
|
|
# --- Timeout ------------------------------------------------------------------
|
|
async def test_an_unanswered_question_expires_and_the_reply_finishes(db, user_id, monkeypatch):
|
|
chat_id, message_id = _chat_that_can_ask(db, user_id)
|
|
payloads: list[dict] = []
|
|
monkeypatch.setattr(
|
|
generation_service,
|
|
"stream_chat",
|
|
_stub_stream([[_ask_chunk("Tea or coffee?")], [_text_chunk("Never mind.")]], payloads),
|
|
)
|
|
# The clamp floor is 60s, so the timeout is forced directly rather than
|
|
# through the settings.
|
|
monkeypatch.setattr(
|
|
generation_service.tools_service,
|
|
"context_for",
|
|
lambda *a, **k: _fast_context(*a, **k),
|
|
)
|
|
|
|
generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
|
|
await asyncio.wait_for(generation_service._run(generation), timeout=5)
|
|
|
|
turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"]
|
|
assert "did not answer" in turns[0]["content"]
|
|
assert generation.tool_events[0]["status"] == "error"
|
|
|
|
|
|
def _fast_context(db, user, chat=None, *, tools=None):
|
|
from lembas.services import tools as tools_service
|
|
|
|
context = tools_service.ToolContext(
|
|
owner_id=user.id if user else "",
|
|
tools=tools.by_name if tools is not None else None,
|
|
)
|
|
context.interaction_timeout = 0.2
|
|
return context
|
|
|
|
|
|
# --- The primitive on its own -------------------------------------------------
|
|
def test_resolving_twice_only_counts_once():
|
|
"""Two tabs, or a double click. The second answer must not win."""
|
|
|
|
async def go():
|
|
pause = interaction.build("abc", [_item()], timeout=5)
|
|
assert pause.resolve(interaction.ANSWER, text="first") is True
|
|
assert pause.resolve(interaction.ANSWER, text="second") is False
|
|
assert (await pause._future).text == "first"
|
|
|
|
asyncio.run(go())
|
|
|
|
|
|
def test_an_interruption_with_no_future_cannot_be_resolved():
|
|
pause = interaction.Interruption(id="abc", items=(_item(),))
|
|
assert pause.resolve(interaction.ANSWER) is False
|
|
|
|
|
|
def _item() -> interaction.Item:
|
|
return interaction.Item(
|
|
index=0,
|
|
kind=interaction.KIND_QUESTION,
|
|
tool_name="ask_user",
|
|
title="Tea or coffee?",
|
|
)
|
|
|
|
|
|
# --- Resolving one -------------------------------------------------------------
|
|
# `answer()` is exercised in-process rather than over the TestClient, because a
|
|
# future belongs to the loop that made it and TestClient runs the app on its
|
|
# own. In production both are the single uvicorn loop, which is the arrangement
|
|
# the in-process registry already requires.
|
|
async def test_answer_finds_the_pause_and_resolves_it(db, user_id):
|
|
chat_id, message_id = _chat_that_can_ask(db, user_id)
|
|
generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
|
|
generation.pending = interaction.build("pause-1", [_item()], timeout=30)
|
|
generation_service._RUNNING[message_id] = generation
|
|
try:
|
|
assert generation_service.answer(chat_id, "pause-1", choice="Tea", text="") is True
|
|
assert (await generation.pending._future).text == "Tea"
|
|
finally:
|
|
generation_service._RUNNING.pop(message_id, None)
|
|
|
|
|
|
async def test_answer_ignores_a_pause_in_another_chat(db, user_id):
|
|
"""The endpoint checks ownership of the chat, so the lookup is scoped to it
|
|
-- an id on its own would not be an authorisation."""
|
|
chat_id, message_id = _chat_that_can_ask(db, user_id)
|
|
generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
|
|
generation.pending = interaction.build("pause-1", [_item()], timeout=30)
|
|
generation_service._RUNNING[message_id] = generation
|
|
try:
|
|
assert generation_service.answer("another-chat", "pause-1", choice="x", text="") is False
|
|
assert not generation.pending._future.done()
|
|
finally:
|
|
generation_service._RUNNING.pop(message_id, None)
|
|
|
|
|
|
async def test_an_allow_choice_is_kept_as_a_verdict_not_as_typed_text():
|
|
""""allow" is a decision, not something somebody wrote in the box."""
|
|
pause = interaction.build("p", [_item()], timeout=5)
|
|
pause.resolve(interaction.ALLOW)
|
|
reply = await pause._future
|
|
assert reply.permitted is True
|
|
assert reply.outcome == interaction.ALLOW
|
|
|
|
|
|
def test_answering_something_that_has_gone_says_so(client, db, registered, user_id):
|
|
chat_id, _message_id = _chat_that_can_ask(db, user_id)
|
|
response = client.post(f"/api/chats/{chat_id}/interaction/nope", data={"choice": "Tea"})
|
|
assert response.status_code == 204
|
|
assert "no longer waiting" in response.headers["HX-Trigger"]
|
|
|
|
|
|
def test_another_account_cannot_answer_your_question(client, db, registered, user_id):
|
|
"""Without the ownership check, a guessed id would be answering -- and
|
|
later, approving a command in -- somebody else's conversation."""
|
|
from lembas.db.models import User
|
|
|
|
chat_id, message_id = _chat_that_can_ask(db, user_id)
|
|
generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
|
|
# No future: the request must be refused before anything tries to resolve
|
|
# it, so there is nothing here for it to reach.
|
|
pause = interaction.Interruption(id="pause-2", items=(_item(),))
|
|
generation.pending = pause
|
|
generation_service._RUNNING[message_id] = generation
|
|
try:
|
|
client.post("/auth/logout")
|
|
client.post(
|
|
"/auth/register",
|
|
data={"name": "Sam", "email": "sam@shire.test", "password": "correct horse battery"},
|
|
follow_redirects=False,
|
|
)
|
|
intruder = db.scalar(select(User).where(User.email == "sam@shire.test"))
|
|
intruder.role = "user"
|
|
intruder.active = True
|
|
db.commit()
|
|
client.post(
|
|
"/auth/login",
|
|
data={"email": "sam@shire.test", "password": "correct horse battery"},
|
|
follow_redirects=False,
|
|
)
|
|
|
|
response = client.post(
|
|
f"/api/chats/{chat_id}/interaction/pause-2", data={"choice": "Tea"}
|
|
)
|
|
assert response.status_code == 404
|
|
assert pause is generation.pending, "still waiting for the person it belongs to"
|
|
finally:
|
|
generation_service._RUNNING.pop(message_id, None)
|
|
|
|
|
|
# --- The card ------------------------------------------------------------------
|
|
def _render(pending) -> str:
|
|
from lembas.api.chats import _ask_html
|
|
|
|
return _ask_html("chat-1", pending)
|
|
|
|
|
|
def test_nothing_pending_renders_nothing():
|
|
"""The frame is sent unconditionally so the card can clear itself. An empty
|
|
string is how it does that."""
|
|
assert _render(None) == ""
|
|
|
|
|
|
def test_the_card_shows_the_question_and_its_options():
|
|
pause = interaction.Interruption(
|
|
id="p1",
|
|
items=(
|
|
interaction.Item(
|
|
index=0,
|
|
kind=interaction.KIND_QUESTION,
|
|
tool_name="ask_user",
|
|
title="Tea or coffee?",
|
|
options=("Tea", "Coffee"),
|
|
),
|
|
),
|
|
)
|
|
html = _render(pause)
|
|
assert "Tea or coffee?" in html
|
|
assert 'value="Tea"' in html and 'value="Coffee"' in html
|
|
assert 'hx-post="/api/chats/chat-1/interaction/p1"' in html
|
|
assert "The model is asking you" in html, "attributed to the model, not to LLeMbas"
|
|
|
|
|
|
def test_the_card_never_offers_a_password_field():
|
|
"""A model talked into asking for a credential must not be handed a field
|
|
that looks built for one."""
|
|
pause = interaction.Interruption(
|
|
id="p1",
|
|
items=(
|
|
interaction.Item(
|
|
index=0,
|
|
kind=interaction.KIND_QUESTION,
|
|
tool_name="ask_user",
|
|
title="Confirm your password to continue:",
|
|
),
|
|
),
|
|
)
|
|
html = _render(pause)
|
|
assert 'type="password"' not in html
|
|
assert 'type="text"' in html
|
|
|
|
|
|
def test_everything_on_the_card_is_escaped():
|
|
"""The question is model output, and the model may have been reading
|
|
somebody else's file a moment ago."""
|
|
pause = interaction.Interruption(
|
|
id="p1",
|
|
items=(
|
|
interaction.Item(
|
|
index=0,
|
|
kind=interaction.KIND_QUESTION,
|
|
tool_name="ask_user",
|
|
title="<img src=x onerror=alert(1)>",
|
|
detail="rm -rf / <script>",
|
|
options=('"><script>bad</script>',),
|
|
),
|
|
),
|
|
)
|
|
html = _render(pause)
|
|
assert "<img" not in html
|
|
assert "<script>" not in html
|
|
assert "<img" in html
|