Houmeres/LLeMbas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Agent chats run commands, and stop to ask first

Browse Source 
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>
This commit is contained in:
Jaroslav Beneš
2026-08-02 00:08:48 +02:00
parent 6a849dc1ec
commit b6aab8de55
14 changed files with 1579 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent_tools.py
+485
View File
@@ -0,0 +1,485 @@
"""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 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"