Files
LLeMbas/tests/test_agent_tools.py
T
Jaroslav Beneš 02e60d6c6c Plan mode proposes, and you decide whether to carry it out
`plan_submit` records an ordered set of steps and ends the turn. Offered in
Plan mode and nowhere else: it stops the reply, and a model in Auto mode
proposing a plan instead of doing the work would be obeying the wrong
instinct at the worst moment.

The plan is stored on the message rather than parsed back out of the prose,
so the button sends exactly what was proposed. It gets one more request to
say what it proposed and why -- a bubble containing only a card reads as
though the model 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 somebody is being asked to approve.

Carrying it out switches to Edit, never Auto. The plan was written under a
mode where every command stopped for approval, and a button that also
removed the asking is not the button anybody pressed. It goes back quoted
and attributed, not stated: a plan whose text came out of a file the model
read must not arrive in the most trusted role in the transcript wearing the
reader's authority.

Also closes the rewind gap. Editing or regenerating a turn rewinds the
transcript and not the machine, so `rewound_at` is stamped and the harness
says so. Nothing tries to undo anything out there -- the project directory
is somebody's real working tree, and deleting their work to match a rewound
transcript would be far worse than the inconsistency.

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

654 lines
23 KiB
Python

"""The agent tools: when they are offered, and what stops them running.
Everything here goes through a real SSH server on 127.0.0.1, so the gate and the
approval loop are exercised against something that genuinely executes rather
than a stub that always agrees.
"""
from __future__ import annotations
import asyncio
import json as _json
import pytest
from lembas.db.models import (
KIND_AGENT,
ROLE_ASSISTANT,
Chat,
Connection,
Message,
Model,
SshProfile,
User,
)
from lembas.services import generation as generation_service
from lembas.services import interaction, settings_store
from lembas.services import tools as tools_service
from lembas.services.agent import policy, session
from lembas.services.agent import ssh as ssh_service
asyncssh = pytest.importorskip("asyncssh")
# --- A machine to act on --------------------------------------------------------
class _Server(asyncssh.SSHServer):
def begin_auth(self, username: str) -> bool:
return False
async def _handler(process):
command = process.command or ""
process.stdout.write(f"ran: {command}\n")
process.exit(0)
@pytest.fixture
async def machine(tmp_path):
"""A real sshd, and a profile pointing at it with its key already pinned."""
project = tmp_path / "project"
project.mkdir()
server = await asyncssh.create_server(
_Server,
"127.0.0.1",
0,
server_host_keys=[asyncssh.generate_private_key("ssh-ed25519")],
process_factory=_handler,
sftp_factory=True,
)
port = next(iter(server.sockets)).getsockname()[1]
line, fingerprint = await ssh_service.capture_host_key("127.0.0.1", port)
try:
yield {"port": port, "host_key": line, "fingerprint": fingerprint, "dir": str(project)}
finally:
server.close()
await server.wait_closed()
def _setup(db, user_id, machine, *, mode=policy.MODE_MANUAL, enabled=True, kind=KIND_AGENT):
"""An agent chat pointed at the machine, with the feature switched on."""
settings_store.update(db, {"enabled": enabled}, key=settings_store.AGENTS)
profile = SshProfile(
owner_id=user_id,
name="Test box",
host="127.0.0.1",
port=machine["port"],
username="tester",
host_key=machine["host_key"],
host_fingerprint=machine["fingerprint"],
default_dir=machine["dir"],
)
db.add(profile)
db.commit()
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,
kind=kind,
ssh_profile_id=profile.id,
project_dir=machine["dir"],
agent_mode=mode,
)
db.add(chat)
db.commit()
return chat, profile
def _offered(db, chat, user) -> set[str]:
return set(tools_service.resolve_tools(db, chat, user).by_name)
# --- The gate --------------------------------------------------------------------
async def test_the_agent_tools_are_offered_to_an_agent_chat(db, user_id, machine):
chat, _profile = _setup(db, user_id, machine)
names = _offered(db, chat, db.get(User, user_id))
assert {"shell_run", "file_read", "file_write", "file_list"} <= names
async def test_an_ordinary_chat_gets_none_of_them(db, user_id, machine):
"""A plain conversation can never shell out, whatever it is asked."""
chat, _profile = _setup(db, user_id, machine, kind="chat")
assert "shell_run" not in _offered(db, chat, db.get(User, user_id))
async def test_nothing_is_offered_while_the_feature_is_off(db, user_id, machine):
chat, _profile = _setup(db, user_id, machine, enabled=False)
assert "shell_run" not in _offered(db, chat, db.get(User, user_id))
async def test_nothing_is_offered_without_the_permission(db, user_id, machine):
chat, _profile = _setup(db, user_id, machine)
user = db.get(User, user_id)
user.role = "user" # administrators pass everything
settings_store.update(db, {"default_permissions": {"tools.agent": False}})
db.commit()
assert "shell_run" not in _offered(db, chat, user)
async def test_nothing_is_offered_without_the_model_capability(db, user_id, machine):
chat, _profile = _setup(db, user_id, machine)
model = db.scalar(tools_service.select(Model))
model.capabilities_json = {"tools": True, "tool_agent": False}
db.commit()
assert "shell_run" not in _offered(db, chat, db.get(User, user_id))
async def test_a_disabled_connection_takes_the_tools_away(db, user_id, machine):
"""Offering a tool that fails on its first call is worse than not offering
it, so every "no" collapses to an empty list."""
chat, profile = _setup(db, user_id, machine)
profile.enabled = False
db.commit()
assert "shell_run" not in _offered(db, chat, db.get(User, user_id))
async def test_somebody_elses_connection_is_not_reachable(db, user_id, machine):
from lembas.security.passwords import hash_password
chat, profile = _setup(db, user_id, machine)
intruder = User(name="Sam", email="s@shire.test", password_hash=hash_password("x"))
intruder.role = "user"
db.add(intruder)
db.commit()
assert session.resolve(db, chat, intruder) is None
# --- Running something -----------------------------------------------------------
async def test_a_command_runs_on_the_machine(db, user_id, machine):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_AUTO)
user = db.get(User, user_id)
resolved = tools_service.resolve_tools(db, chat, user)
context = tools_service.context_for(db, user, chat, tools=resolved)
outcome = await tools_service.run_tool(
context, "shell_run", '{"command": "echo hello"}'
)
assert outcome.event["status"] == "ok"
assert "echo hello" in outcome.content
async def test_files_are_written_and_read_back(db, user_id, machine, tmp_path):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_AUTO)
user = db.get(User, user_id)
resolved = tools_service.resolve_tools(db, chat, user)
context = tools_service.context_for(db, user, chat, tools=resolved)
await tools_service.run_tool(
context, "file_write", '{"path": "note.txt", "content": "a mallorn tree"}'
)
assert (tmp_path / "project" / "note.txt").read_text() == "a mallorn tree"
read = await tools_service.run_tool(context, "file_read", '{"path": "note.txt"}')
assert "mallorn" in read.content
listed = await tools_service.run_tool(context, "file_list", "{}")
assert "note.txt" in listed.content
# --- The runner backstop ----------------------------------------------------------
async def test_a_runner_refuses_what_the_mode_forbids(db, user_id, machine):
"""`_authorise` is the real gate and runs first. This is the belt to that
brace: a call arriving by some other path -- a retry, a re-run button --
must not walk past it."""
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_PLAN)
user = db.get(User, user_id)
resolved = tools_service.resolve_tools(db, chat, user)
context = tools_service.context_for(db, user, chat, tools=resolved)
outcome = await tools_service.run_tool(context, "shell_run", '{"command": "rm -rf /"}')
assert outcome.event["status"] == "error"
assert "not allowed" in outcome.content.lower()
async def test_plan_mode_still_reads(db, user_id, machine, tmp_path):
(tmp_path / "project" / "readme.txt").write_text("contents")
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_PLAN)
user = db.get(User, user_id)
resolved = tools_service.resolve_tools(db, chat, user)
context = tools_service.context_for(db, user, chat, tools=resolved)
outcome = await tools_service.run_tool(context, "file_read", '{"path": "readme.txt"}')
assert outcome.event["status"] == "ok"
assert "contents" in outcome.content
async def test_edit_mode_writes_but_will_not_run(db, user_id, machine):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_EDIT)
user = db.get(User, user_id)
resolved = tools_service.resolve_tools(db, chat, user)
context = tools_service.context_for(db, user, chat, tools=resolved)
wrote = await tools_service.run_tool(
context, "file_write", '{"path": "x.txt", "content": "hi"}'
)
assert wrote.event["status"] == "ok"
ran = await tools_service.run_tool(context, "shell_run", '{"command": "ls"}')
assert ran.event["status"] == "error"
# --- The approval card in the loop -------------------------------------------------
def _chunk(name: str, arguments: str, *, call_id="c1", index=0):
return {
"choices": [
{
"delta": {
"tool_calls": [
{
"index": index,
"id": call_id,
"function": {"name": name, "arguments": arguments},
}
]
}
}
]
}
def _text(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 = 3.0):
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 for approval")
def _pending_reply(db, chat):
db.add(Message(chat_id=chat.id, role="user", content="do it", complete=True))
db.commit()
assistant = Message(chat_id=chat.id, role=ROLE_ASSISTANT, content="", complete=False)
db.add(assistant)
db.commit()
return assistant.id
async def test_a_command_waits_for_approval_and_the_card_shows_it(
db, user_id, machine, monkeypatch
):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_MANUAL)
message_id = _pending_reply(db, chat)
monkeypatch.setattr(
generation_service,
"stream_chat",
_stub_stream(
[[_chunk("shell_run", '{"command": "rm -rf /tmp/x"}')], [_text("Done.")]], []
),
)
generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
task = asyncio.create_task(generation_service._run(generation))
pending = await _until_paused(generation)
assert pending.kind == interaction.KIND_APPROVAL
item = pending.items[0]
# The exact command, verbatim. A card that paraphrased it would be
# approving something other than what runs.
assert item.detail == "rm -rf /tmp/x"
assert "Test box" in item.title
assert "Manual" in item.reason
pending.resolve(interaction.DENY)
await task
async def test_denying_reaches_the_model_as_words_and_runs_nothing(
db, user_id, machine, monkeypatch, tmp_path
):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_MANUAL)
message_id = _pending_reply(db, chat)
payloads: list[dict] = []
monkeypatch.setattr(
generation_service,
"stream_chat",
_stub_stream(
[
[_chunk("file_write", '{"path": "never.txt", "content": "nope"}')],
[_text("Understood.")],
],
payloads,
),
)
generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
task = asyncio.create_task(generation_service._run(generation))
pending = await _until_paused(generation)
pending.resolve(interaction.DENY)
await task
turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"]
assert "declined" in turns[0]["content"].lower()
# And it says not to look for another way round, because a model reads
# "not allowed" as "not allowed like that" otherwise.
assert "another way" in turns[0]["content"]
assert not (tmp_path / "project" / "never.txt").exists()
async def test_allowing_runs_it(db, user_id, machine, monkeypatch, tmp_path):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_MANUAL)
message_id = _pending_reply(db, chat)
payloads: list[dict] = []
monkeypatch.setattr(
generation_service,
"stream_chat",
_stub_stream(
[
[_chunk("file_write", '{"path": "yes.txt", "content": "written"}')],
[_text("Done.")],
],
payloads,
),
)
generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
task = asyncio.create_task(generation_service._run(generation))
pending = await _until_paused(generation)
pending.resolve(interaction.ALLOW)
await task
assert (tmp_path / "project" / "yes.txt").read_text() == "written"
async def test_auto_mode_never_pauses(db, user_id, machine, monkeypatch, tmp_path):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_AUTO)
message_id = _pending_reply(db, chat)
monkeypatch.setattr(
generation_service,
"stream_chat",
_stub_stream(
[
[_chunk("file_write", '{"path": "auto.txt", "content": "no asking"}')],
[_text("Done.")],
],
[],
),
)
generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
await asyncio.wait_for(generation_service._run(generation), timeout=10)
assert generation.pending is None
assert (tmp_path / "project" / "auto.txt").read_text() == "no asking"
async def test_the_credential_is_cleared_when_the_reply_ends(db, user_id, machine, monkeypatch):
"""A finished Generation lingers five minutes so late followers get the
final frames. A private key should not linger with it."""
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_AUTO)
message_id = _pending_reply(db, chat)
captured = {}
original = tools_service.context_for
def capture(*args, **kwargs):
context = original(*args, **kwargs)
captured["context"] = context
return context
monkeypatch.setattr(tools_service, "context_for", capture)
monkeypatch.setattr(
generation_service, "stream_chat", _stub_stream([[_text("Nothing to do.")]], [])
)
generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
await generation_service._run(generation)
assert captured["context"].agent is not None
assert captured["context"].agent.spec == {}, "the decrypted credential is dropped"
# --- The harness has to be able to name the machine ---------------------------
def test_the_registry_maps_the_agent_tools_to_their_family(db):
"""`harness._families` maps an offered tool *name* back to a family to
decide which fragments apply, and it has no chat to resolve against. Without
the agent tools listed here, `shell_run` resolves to no family and an agent
chat is told nothing about the machine it is working on -- the same omission
that cost custom tools their guidance once already."""
book = tools_service.registry(db)
for name in ("shell_run", "file_read", "file_write", "file_list"):
assert name in book, name
assert book[name].family == "agent"
async def test_the_harness_says_where_and_under_what_rules(db, user_id, machine):
from lembas.services import harness
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_MANUAL)
user = db.get(User, user_id)
offered = tools_service.resolve_tools(db, chat, user).schemas
text = harness.compose(db, user, offered, chat)
assert "Test box" in text, "which machine"
assert machine["dir"] in text, "which directory"
assert "Manual" in text, "what the mode permits"
# The single most likely cause of "the agent seems stupid": `cd build`
# followed by `make` fails silently otherwise.
assert "fresh shell" in text
async def test_an_ordinary_chat_is_told_none_of_it(db, user_id, machine):
from lembas.services import harness
chat, _profile = _setup(db, user_id, machine, kind="chat")
user = db.get(User, user_id)
offered = tools_service.resolve_tools(db, chat, user).schemas
text = harness.compose(db, user, offered, chat)
assert "Test box" not in text
assert "fresh shell" not in text
async def test_an_agent_chat_is_told_its_real_round_budget(db, user_id, machine):
"""MAX_ROUNDS is three. An agent chat gets forty, and telling it three would
be a false fact about its own budget on every turn."""
from lembas.services import harness
settings_store.update(db, {"max_steps": 25}, key=settings_store.AGENTS)
chat, _profile = _setup(db, user_id, machine)
user = db.get(User, user_id)
values = harness.context_variables(
db, user, tools_service.resolve_tools(db, chat, user).schemas, chat
)
assert values["max_rounds"] == "25"
# --- Plan mode's artifact ---------------------------------------------------------
def test_plan_submit_is_offered_only_in_plan_mode(db, user_id, machine):
"""It ends the reply. A model in Auto mode that proposed a plan instead of
doing the work would be obeying the wrong instinct at the wrong moment."""
user = db.get(User, user_id)
planning, _p = _setup(db, user_id, machine, mode=policy.MODE_PLAN)
assert "plan_submit" in _offered(db, planning, user)
planning.agent_mode = policy.MODE_AUTO
db.commit()
assert "plan_submit" not in _offered(db, planning, user)
async def test_a_plan_ends_the_reply_and_lands_on_the_message(
db, user_id, machine, monkeypatch
):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_PLAN)
message_id = _pending_reply(db, chat)
plan = {"title": "Tidy the logs", "steps": ["Read the log", "Rotate it", "Restart"]}
rounds = [
[_chunk("plan_submit", _json.dumps(plan))],
# The model gets one wordless round to explain itself. If it tries to
# act in it -- as this one does -- there are no tools to act with.
[_chunk("shell_run", '{"command": "rm -rf /"}'), _text("Here is what I would do.")],
]
seen: list[dict] = []
monkeypatch.setattr(generation_service, "stream_chat", _stub_stream(rounds, seen))
generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
await asyncio.wait_for(generation_service._run(generation), timeout=10)
assert generation.plan == plan
assert len(seen) == 2, "one round to plan, one to say what it proposed"
assert "tools" not in seen[1], "the second round is offered nothing to act with"
assert len(generation.tool_events) == 1, "the shell call had no tool to reach"
db.expire_all()
message = db.get(Message, message_id)
assert message.plan_json == plan
async def test_a_plan_with_no_steps_is_sent_back(db, user_id, machine, monkeypatch):
chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_PLAN)
message_id = _pending_reply(db, chat)
monkeypatch.setattr(
generation_service,
"stream_chat",
_stub_stream(
[[_chunk("plan_submit", '{"title": "Nothing", "steps": []}')], [_text("Sorry.")]],
[],
),
)
generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
await asyncio.wait_for(generation_service._run(generation), timeout=10)
assert generation.plan is None
assert generation.tool_events[0]["status"] == "error"
# --- Carrying a plan out ------------------------------------------------------------
def test_executing_a_plan_switches_to_edit_and_quotes_it(client, db, registered, machine):
"""Edit, never Auto. The plan was written under a mode where every command
stopped for approval, and a button that also removed the asking is not the
button anybody pressed."""
from sqlalchemy import select as _select
user = db.scalar(_select(User))
chat, _profile = _setup(db, user.id, machine, mode=policy.MODE_PLAN)
message = Message(
chat_id=chat.id,
role=ROLE_ASSISTANT,
content="Here is what I would do.",
complete=True,
plan_json={"title": "Tidy the logs", "steps": ["Read the log", "Rotate it"]},
)
db.add(message)
db.commit()
response = client.post(f"/api/chats/{chat.id}/messages/{message.id}/execute-plan")
assert response.status_code == 200
db.refresh(chat)
assert chat.agent_mode == policy.MODE_EDIT
sent = db.scalars(
_select(Message).where(Message.chat_id == chat.id, Message.role == "user")
).all()[-1]
assert "Tidy the logs" in sent.content
assert "Read the log" in sent.content
# Quoted rather than stated. A plan whose text came out of a file the model
# read must not arrive wearing the reader's authority.
assert sent.content.lstrip().startswith("Carry out the plan you proposed above")
assert "> " in sent.content
def test_a_message_with_no_plan_cannot_be_executed(client, db, registered, machine):
from sqlalchemy import select as _select
user = db.scalar(_select(User))
chat, _profile = _setup(db, user.id, machine)
message = Message(chat_id=chat.id, role=ROLE_ASSISTANT, content="hi", complete=True)
db.add(message)
db.commit()
assert (
client.post(f"/api/chats/{chat.id}/messages/{message.id}/execute-plan").status_code
== 404
)
# --- Rewind ---------------------------------------------------------------------------
def test_editing_a_turn_records_that_the_machine_did_not_rewind(
client, db, registered, machine
):
"""The transcript goes back; the project directory does not. Deleting
somebody's real working tree to match would be far worse than the
inconsistency, so the model is told instead."""
from sqlalchemy import select as _select
user = db.scalar(_select(User))
chat, _profile = _setup(db, user.id, machine)
first = Message(chat_id=chat.id, role="user", content="do a thing", complete=True)
db.add(first)
db.commit()
assert chat.rewound_at is None
client.post(
f"/api/chats/{chat.id}/messages/{first.id}/edit", data={"content": "do another thing"}
)
db.refresh(chat)
assert chat.rewound_at is not None
def test_an_ordinary_chat_records_no_rewind(client, db, registered, machine):
from sqlalchemy import select as _select
user = db.scalar(_select(User))
chat, _profile = _setup(db, user.id, machine, kind="chat")
first = Message(chat_id=chat.id, role="user", content="hello", complete=True)
db.add(first)
db.commit()
client.post(f"/api/chats/{chat.id}/messages/{first.id}/edit", data={"content": "hi"})
db.refresh(chat)
assert chat.rewound_at is None
async def test_the_harness_warns_after_a_rewind(db, user_id, machine):
from datetime import UTC, datetime
from lembas.services import harness
chat, _profile = _setup(db, user_id, machine)
user = db.get(User, user_id)
offered = tools_service.resolve_tools(db, chat, user).schemas
assert "was rewound" not in harness.compose(db, user, offered, chat)
chat.rewound_at = datetime.now(UTC)
db.commit()
text = harness.compose(db, user, offered, chat)
assert "was rewound" in text
assert "still there" in text