A directory the model knows about, and @ to name a file in it
An agent chat used to open with the model knowing the name of a machine and nothing about what was on it, so the first two rounds of every reply went on finding out. It now gets a listing: one read-only command, `git ls-files` where that works and `find` otherwise, falling back to an SFTP walk that always does. git first because a repository already carries somebody's considered list of what is not part of the project, and reproducing it by hand is how an index ends up mostly build output. The listing is budgeted rather than dumped. A tree of a thousand files is worse than no tree -- it costs the window on every request forever and buries the four names that mattered -- so directories that will not fit are shown as a count and the model is told to open one itself. Collapsing picks the deepest and largest first: by saving alone it would take `src/` before `src/web/static/vendor/`, because it contains it, and lose every name worth having. Read from a cache and never fetched. `harness.context_variables` is synchronous and sits on the request path; the walk happens in the generation setup, which is async and already doing network work, with a short wait. A chat whose first reply outruns its first walk simply has no listing that turn and the fragment disappears rather than appearing as an empty heading. Then `@`, over the same index and over the library, and `/` for commands with an Alt-based keyboard for the same jobs. A mentioned file arrives as contents, not a reference -- a small model asked to call file_read often does not bother -- and it arrives with its absolute path and the machine it came from, because a model handed `main.py` cannot tell which of four it is and cannot name it back when asked to change something. The rule that matters for `/`: a message that merely starts with a slash still sends. `//` escapes and an unrecognised command is posted as written. Swallowing somebody's message is a much worse failure than an unknown command. Two exceptions to Manual mode now, not one. Browsing and indexing are a person acting, not a model, so neither passes through policy.py -- the same argument the terminal panel rests on. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
@@ -102,6 +102,21 @@ def fresh_terminal_registry() -> Iterator[None]:
|
||||
_clear()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def fresh_project_index() -> Iterator[None]:
|
||||
"""Empty the directory-listing cache between tests, for the third time.
|
||||
|
||||
Keyed on (profile, directory) and both are recycled freely by fixtures, so
|
||||
without this a test asserting "the listing said X" can be answered by the
|
||||
previous test's walk of an entirely different tmp_path.
|
||||
"""
|
||||
from lembas.services.agent import index as index_service
|
||||
|
||||
index_service.clear()
|
||||
yield
|
||||
index_service.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db() -> Iterator[Session]:
|
||||
session = get_session_factory()()
|
||||
|
||||
@@ -0,0 +1,252 @@
|
||||
"""Listing a project directory: the ladder, the cache, and the budget.
|
||||
|
||||
The ladder is exercised against a fake executor rather than a real host,
|
||||
because what is being tested is *which rung is chosen* and what happens when
|
||||
one falls through -- and a real box either has git or does not, which would
|
||||
make the interesting cases unreachable.
|
||||
|
||||
`scan_dir` itself is tested against a real SFTP server in test_agent_browse.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from lembas.services.agent import index as index_service
|
||||
from lembas.services.agent.base import ExecError, ExecResult, RemoteEntry
|
||||
|
||||
|
||||
class _Fake:
|
||||
"""An executor that answers whatever the test says, and records the asking."""
|
||||
|
||||
def __init__(self, *, answers: dict[str, ExecResult] | None = None, tree=None):
|
||||
self.answers = answers or {}
|
||||
self.tree = tree or {}
|
||||
self.commands: list[str] = []
|
||||
self.scans: list[str] = []
|
||||
|
||||
async def run(self, request) -> ExecResult:
|
||||
self.commands.append(request.command)
|
||||
for fragment, result in self.answers.items():
|
||||
if fragment in request.command:
|
||||
return result
|
||||
return ExecResult(exit_status=1, output="")
|
||||
|
||||
async def scan_dir(self, path: str) -> list[RemoteEntry]:
|
||||
self.scans.append(path)
|
||||
if path not in self.tree:
|
||||
raise ExecError(f"There is no directory at {path}.")
|
||||
return self.tree[path]
|
||||
|
||||
|
||||
def _ok(output: str) -> ExecResult:
|
||||
return ExecResult(exit_status=0, output=output)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clean_cache():
|
||||
index_service.clear()
|
||||
yield
|
||||
index_service.clear()
|
||||
|
||||
|
||||
# --- The ladder --------------------------------------------------------------
|
||||
async def test_git_is_tried_first_and_wins_where_it_works():
|
||||
"""`--exclude-standard` is the whole reason. A repository already carries
|
||||
somebody's considered list of what is not part of the project, and
|
||||
reproducing it by hand is how an index ends up mostly build output."""
|
||||
executor = _Fake(answers={"git ls-files": _ok("README.md\nsrc/main.py\n")})
|
||||
|
||||
found = await index_service.build(executor, "/work")
|
||||
|
||||
assert found.source == "git"
|
||||
assert found.paths == ("README.md", "src/main.py")
|
||||
assert executor.commands == [executor.commands[0]]
|
||||
assert "--exclude-standard" in executor.commands[0]
|
||||
|
||||
|
||||
async def test_find_takes_over_when_the_directory_is_not_a_repository():
|
||||
executor = _Fake(answers={"find .": _ok("./README.md\n./src/main.py\n")})
|
||||
|
||||
found = await index_service.build(executor, "/work")
|
||||
|
||||
assert found.source == "find"
|
||||
assert found.paths == ("README.md", "src/main.py")
|
||||
assert len(executor.commands) == 2 # git was tried, and fell through
|
||||
|
||||
|
||||
async def test_find_prunes_the_usual_noise():
|
||||
"""Not applied to the git rung, which has already applied the repository's
|
||||
own rules -- a checked-in `vendor/` is checked in on purpose."""
|
||||
executor = _Fake(answers={"find .": _ok("./a\n")})
|
||||
|
||||
await index_service.build(executor, "/work")
|
||||
|
||||
assert "node_modules" in executor.commands[1]
|
||||
assert "'.git'" in executor.commands[1]
|
||||
|
||||
|
||||
async def test_sftp_is_the_last_resort_and_always_works():
|
||||
executor = _Fake(
|
||||
tree={
|
||||
"/work": [RemoteEntry("src", True), RemoteEntry("README.md", False)],
|
||||
"src": [RemoteEntry("main.py", False)],
|
||||
}
|
||||
)
|
||||
|
||||
found = await index_service.build(executor, "/work")
|
||||
|
||||
assert found.source == "sftp"
|
||||
assert "README.md" in found.paths
|
||||
assert "src/main.py" in found.paths
|
||||
|
||||
|
||||
async def test_a_directory_that_cannot_be_read_at_all_is_empty_not_an_error():
|
||||
"""A reply must not fail because a listing did. The fragment carrying it
|
||||
disappears instead, which is what `requires` is for."""
|
||||
executor = _Fake(tree={})
|
||||
|
||||
found = await index_service.build(executor, "/work")
|
||||
|
||||
assert found.paths == ()
|
||||
assert not found.ok
|
||||
|
||||
|
||||
async def test_control_characters_are_stripped_from_a_filename():
|
||||
"""These end up inside a system prompt. An escape sequence in a filename
|
||||
could otherwise repaint the transcript it is quoted in."""
|
||||
executor = _Fake(answers={"git ls-files": _ok("ok.py\nevil\x1b[31m.py\n")})
|
||||
|
||||
found = await index_service.build(executor, "/work")
|
||||
|
||||
assert "evil[31m.py" in found.paths
|
||||
assert "\x1b" not in "".join(found.paths)
|
||||
|
||||
|
||||
async def test_the_number_of_paths_is_capped_and_says_so():
|
||||
""""There is nothing else here" and "I stopped looking" are different
|
||||
answers, and it must not give the first for the second."""
|
||||
many = "\n".join(f"f{i}.py" for i in range(index_service.MAX_ENTRIES + 50))
|
||||
executor = _Fake(answers={"git ls-files": _ok(many)})
|
||||
|
||||
found = await index_service.build(executor, "/work")
|
||||
|
||||
assert len(found.paths) == index_service.MAX_ENTRIES
|
||||
assert found.truncated
|
||||
|
||||
|
||||
# --- The cache ---------------------------------------------------------------
|
||||
async def test_a_second_ask_does_not_walk_again():
|
||||
executor = _Fake(answers={"git ls-files": _ok("a.py\n")})
|
||||
|
||||
await index_service.ensure(executor, "profile-1", "/work")
|
||||
await index_service.ensure(executor, "profile-1", "/work")
|
||||
|
||||
assert len(executor.commands) == 1
|
||||
|
||||
|
||||
async def test_refreshing_walks_again():
|
||||
executor = _Fake(answers={"git ls-files": _ok("a.py\n")})
|
||||
|
||||
await index_service.ensure(executor, "profile-1", "/work")
|
||||
await index_service.ensure(executor, "profile-1", "/work", refresh=True)
|
||||
|
||||
assert len(executor.commands) == 2
|
||||
|
||||
|
||||
async def test_a_stale_index_is_not_returned():
|
||||
executor = _Fake(answers={"git ls-files": _ok("a.py\n")})
|
||||
await index_service.ensure(executor, "profile-1", "/work")
|
||||
|
||||
index_service._CACHE[("profile-1", "/work")] = index_service.ProjectIndex(
|
||||
paths=("a.py",), total=1, source="git", built_at=time.monotonic() - index_service.TTL - 1
|
||||
)
|
||||
|
||||
assert index_service.cached("profile-1", "/work") is None
|
||||
|
||||
|
||||
async def test_forgetting_a_connection_drops_its_listings():
|
||||
"""A connection somebody has just revoked must not leave a listing of the
|
||||
machine behind it. Called from the same three places that close its
|
||||
terminals."""
|
||||
executor = _Fake(answers={"git ls-files": _ok("a.py\n")})
|
||||
await index_service.ensure(executor, "profile-1", "/work")
|
||||
await index_service.ensure(executor, "profile-2", "/work")
|
||||
|
||||
assert index_service.forget("profile-1") == 1
|
||||
assert index_service.cached("profile-1", "/work") is None
|
||||
assert index_service.cached("profile-2", "/work") is not None
|
||||
|
||||
|
||||
def test_reading_the_cache_never_does_work():
|
||||
"""`harness` calls this synchronously while assembling the system message,
|
||||
so it must never be the thing that opens a connection."""
|
||||
assert index_service.cached("nobody", "/nowhere") is None
|
||||
|
||||
|
||||
# --- The budget --------------------------------------------------------------
|
||||
def _index(paths) -> index_service.ProjectIndex:
|
||||
return index_service.ProjectIndex(
|
||||
paths=tuple(sorted(paths)), total=len(paths), source="git", built_at=time.monotonic()
|
||||
)
|
||||
|
||||
|
||||
def test_a_small_tree_is_shown_whole():
|
||||
text = index_service.render(_index(["README.md", "src/main.py"]), 2000)
|
||||
|
||||
assert "README.md" in text
|
||||
assert "main.py" in text
|
||||
assert "files)" not in text
|
||||
|
||||
|
||||
def test_a_big_directory_becomes_a_count():
|
||||
paths = ["README.md"] + [f"vendor/f{i}.js" for i in range(400)]
|
||||
|
||||
text = index_service.render(_index(paths), 300)
|
||||
|
||||
assert "vendor/ (400 files)" in text
|
||||
assert "README.md" in text
|
||||
assert "file_list" in text
|
||||
|
||||
|
||||
def test_the_deepest_big_directory_is_collapsed_first():
|
||||
"""Collapsing by size alone takes `src/` before `src/.../vendor/` -- it is
|
||||
bigger because it *contains* it -- and loses every name worth having in
|
||||
order to fold away one directory of third-party files."""
|
||||
paths = [f"src/api/{n}.py" for n in "abcde"] + [
|
||||
f"src/web/vendor/f{i}.js" for i in range(300)
|
||||
]
|
||||
|
||||
text = index_service.render(_index(paths), 400)
|
||||
|
||||
assert "vendor/ (300 files)" in text
|
||||
assert "a.py" in text # src/api survived
|
||||
|
||||
|
||||
def test_a_flat_directory_of_thousands_is_still_bounded():
|
||||
"""The one shape collapsing cannot help with: no directory to fold them
|
||||
into, so the running budget has to bite instead."""
|
||||
text = index_service.render(_index([f"dump{i:05d}.log" for i in range(5000)]), 400)
|
||||
|
||||
assert len(text) < 1200
|
||||
assert "more files" in text
|
||||
|
||||
|
||||
def test_a_budget_of_zero_renders_nothing():
|
||||
"""Which is how "index it for the picker, but say nothing to the model" is
|
||||
expressed. The fragment vanishes rather than appearing empty."""
|
||||
assert index_service.render(_index(["a.py"]), 0) == ""
|
||||
|
||||
|
||||
def test_an_empty_index_renders_nothing():
|
||||
assert index_service.render(index_service.ProjectIndex(), 2000) == ""
|
||||
|
||||
|
||||
def test_a_truncated_index_says_it_is_a_sample():
|
||||
sample = index_service.ProjectIndex(
|
||||
paths=("a.py", "b.py"), total=9000, truncated=True, source="find", built_at=time.monotonic()
|
||||
)
|
||||
|
||||
assert "sample" in index_service.render(sample, 2000)
|
||||
@@ -257,3 +257,103 @@ def test_the_tools_array_rides_along(db, owner):
|
||||
offered = tools_service.enabled_tools(db, chat, owner)
|
||||
body = chat_service.build_request(db, chat, tools=offered, user=owner)
|
||||
assert body["tools"] == offered
|
||||
|
||||
|
||||
# --- The project listing -----------------------------------------------------
|
||||
# Injected from a cache that something else fills, because this module runs
|
||||
# synchronously on the request path and an SFTP round trip here would hold a
|
||||
# request open while somebody's machine thought about it.
|
||||
def _agent_chat(db, owner):
|
||||
from lembas.db.models import KIND_AGENT, SshProfile
|
||||
|
||||
settings_store.update(db, {"enabled": True}, key=settings_store.AGENTS)
|
||||
profile = SshProfile(
|
||||
owner_id=owner.id,
|
||||
name="Test box",
|
||||
host="127.0.0.1",
|
||||
port=22,
|
||||
username="tester",
|
||||
host_key="host key",
|
||||
host_fingerprint="SHA256:x",
|
||||
default_dir="/work",
|
||||
)
|
||||
db.add(profile)
|
||||
db.commit()
|
||||
chat = Chat(
|
||||
user_id=owner.id,
|
||||
kind=KIND_AGENT,
|
||||
ssh_profile_id=profile.id,
|
||||
project_dir="/work",
|
||||
)
|
||||
db.add(chat)
|
||||
db.commit()
|
||||
return chat, profile
|
||||
|
||||
|
||||
def _agent_tools(db):
|
||||
"""The agent tools as offered, which REGISTRY does not carry.
|
||||
|
||||
`REGISTRY` is built at import time and holds the built-ins alone; the agent
|
||||
tools are listed by `registry(db)`, unbound to any chat. That is the same
|
||||
lookup the harness does to map `shell_run` back to the `agent` family, and
|
||||
the reason it exists at all.
|
||||
"""
|
||||
book = tools_service.registry(db)
|
||||
return [book["shell_run"].schema]
|
||||
|
||||
|
||||
def _cache(profile_id, paths):
|
||||
import time
|
||||
|
||||
from lembas.services.agent import index as index_service
|
||||
|
||||
index_service._CACHE[(profile_id, "/work")] = index_service.ProjectIndex(
|
||||
paths=tuple(paths), total=len(paths), source="git", built_at=time.monotonic()
|
||||
)
|
||||
|
||||
|
||||
def test_the_project_listing_reaches_the_model(db, owner):
|
||||
chat, profile = _agent_chat(db, owner)
|
||||
_cache(profile.id, ["README.md", "src/main.py"])
|
||||
|
||||
text = harness.compose(db, owner, _agent_tools(db), chat=chat)
|
||||
|
||||
assert "Files in /work" in text
|
||||
assert "README.md" in text
|
||||
|
||||
|
||||
def test_nothing_cached_means_no_section_at_all(db, owner):
|
||||
"""Not an empty heading. `Fragment.requires` makes the whole thing vanish,
|
||||
which is what lets the first reply in a new chat outrun the first walk
|
||||
without saying anything strange."""
|
||||
chat, _profile = _agent_chat(db, owner)
|
||||
|
||||
text = harness.compose(db, owner, _agent_tools(db), chat=chat)
|
||||
|
||||
assert "Files in" not in text
|
||||
|
||||
|
||||
def test_a_budget_of_zero_keeps_the_listing_out_of_the_prompt(db, owner):
|
||||
"""The listing is still built and the file picker still uses it. This is
|
||||
the only way to say "index it, but do not spend context on it"."""
|
||||
chat, profile = _agent_chat(db, owner)
|
||||
_cache(profile.id, ["README.md"])
|
||||
settings_store.update(db, {"index_chars": 0}, key=settings_store.AGENTS)
|
||||
|
||||
assert "Files in" not in harness.compose(db, owner, _agent_tools(db), chat=chat)
|
||||
|
||||
|
||||
def test_switching_the_listing_off_keeps_it_out(db, owner):
|
||||
chat, profile = _agent_chat(db, owner)
|
||||
_cache(profile.id, ["README.md"])
|
||||
settings_store.update(db, {"index_enabled": False}, key=settings_store.AGENTS)
|
||||
|
||||
assert "Files in" not in harness.compose(db, owner, _agent_tools(db), chat=chat)
|
||||
|
||||
|
||||
def test_a_plain_chat_is_told_nothing_about_files(db, owner):
|
||||
chat = Chat(user_id=owner.id)
|
||||
db.add(chat)
|
||||
db.commit()
|
||||
|
||||
assert "Files in" not in harness.compose(db, owner, _tools("web_search"), chat=chat)
|
||||
|
||||
@@ -0,0 +1,331 @@
|
||||
"""`@` attachments: what the picker offers, and what the model is told it got.
|
||||
|
||||
The point of the second half is the one the feature was asked for: a model
|
||||
handed a file called `main.py` cannot tell which of four it is looking at, and
|
||||
cannot name it back when asked to change something. So the path and the machine
|
||||
travel with the contents.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy import select
|
||||
|
||||
from lembas.db.models import Attachment, Message, SshProfile, User
|
||||
from lembas.services import chat as chat_service
|
||||
from lembas.services import settings_store
|
||||
from lembas.services.agent import index as index_service
|
||||
from lembas.services.agent import ssh as ssh_service
|
||||
|
||||
asyncssh = pytest.importorskip("asyncssh")
|
||||
|
||||
|
||||
class _Server(asyncssh.SSHServer):
|
||||
def begin_auth(self, username: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def box(tmp_path):
|
||||
"""A real SFTP server on its own loop, with a small tree to mention from."""
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
root = tmp_path / "work"
|
||||
(root / "src").mkdir(parents=True)
|
||||
(root / "src" / "main.py").write_text("print('hello')\n")
|
||||
(root / "README.md").write_text("# Project\n")
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
thread = threading.Thread(target=loop.run_forever, daemon=True)
|
||||
thread.start()
|
||||
|
||||
async def start():
|
||||
server = await asyncssh.create_server(
|
||||
_Server,
|
||||
"127.0.0.1",
|
||||
0,
|
||||
server_host_keys=[asyncssh.generate_private_key("ssh-ed25519")],
|
||||
sftp_factory=True,
|
||||
)
|
||||
port = next(iter(server.sockets)).getsockname()[1]
|
||||
line, fingerprint = await ssh_service.capture_host_key("127.0.0.1", port)
|
||||
return server, port, line, fingerprint
|
||||
|
||||
server, port, line, fingerprint = asyncio.run_coroutine_threadsafe(start(), loop).result(10)
|
||||
try:
|
||||
yield {"port": port, "host_key": line, "fingerprint": fingerprint, "root": str(root)}
|
||||
finally:
|
||||
|
||||
async def stop():
|
||||
server.close()
|
||||
await server.wait_closed()
|
||||
|
||||
asyncio.run_coroutine_threadsafe(stop(), loop).result(10)
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
def _profile(db, box) -> SshProfile:
|
||||
settings_store.update(db, {"enabled": True}, key=settings_store.AGENTS)
|
||||
user = db.scalars(select(User)).first()
|
||||
profile = SshProfile(
|
||||
owner_id=user.id,
|
||||
name="Container",
|
||||
host="127.0.0.1",
|
||||
port=box["port"],
|
||||
username="tester",
|
||||
host_key=box["host_key"],
|
||||
host_fingerprint=box["fingerprint"],
|
||||
default_dir=box["root"],
|
||||
)
|
||||
db.add(profile)
|
||||
db.commit()
|
||||
return profile
|
||||
|
||||
|
||||
def _index(profile, box, paths=("README.md", "src/main.py")):
|
||||
index_service._CACHE[(profile.id, box["root"])] = index_service.ProjectIndex(
|
||||
paths=tuple(paths), total=len(paths), source="git", built_at=time.monotonic()
|
||||
)
|
||||
|
||||
|
||||
# --- The picker --------------------------------------------------------------
|
||||
def test_the_picker_offers_project_files(client: TestClient, db, registered, box):
|
||||
profile = _profile(db, box)
|
||||
_index(profile, box)
|
||||
|
||||
body = client.get(
|
||||
"/api/files/mention-picker",
|
||||
params={"q": "main", "profile_id": profile.id, "project_dir": box["root"]},
|
||||
).text
|
||||
|
||||
assert "src/main.py" in body
|
||||
assert "README.md" not in body # filtered by the query
|
||||
|
||||
|
||||
def test_the_picker_never_waits_on_a_machine(client: TestClient, db, registered, box):
|
||||
"""No listing built yet means no files offered, not a connection opened.
|
||||
|
||||
This is a keystroke-latency path. Building the index here would put an SSH
|
||||
round trip between a letter and the menu.
|
||||
"""
|
||||
profile = _profile(db, box)
|
||||
|
||||
body = client.get(
|
||||
"/api/files/mention-picker",
|
||||
params={"profile_id": profile.id, "project_dir": box["root"]},
|
||||
).text
|
||||
|
||||
assert "main.py" not in body
|
||||
|
||||
|
||||
def test_the_picker_refuses_somebody_elses_connection(client: TestClient, db, registered, box):
|
||||
"""An id in a query string is not an authorisation, and this lists the
|
||||
contents of somebody's machine."""
|
||||
profile = _profile(db, box)
|
||||
_index(profile, box)
|
||||
|
||||
client.post("/auth/logout", follow_redirects=False)
|
||||
client.post(
|
||||
"/auth/register",
|
||||
data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
|
||||
follow_redirects=False,
|
||||
)
|
||||
|
||||
body = client.get(
|
||||
"/api/files/mention-picker",
|
||||
params={"profile_id": profile.id, "project_dir": box["root"]},
|
||||
).text
|
||||
|
||||
assert "main.py" not in body
|
||||
|
||||
|
||||
def test_a_plain_chat_gets_a_picker_with_no_file_half(client: TestClient, db, registered):
|
||||
"""`@` works everywhere; only the project half needs a connection."""
|
||||
response = client.get("/api/files/mention-picker")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "In the project" not in response.text
|
||||
|
||||
|
||||
# --- Attaching ---------------------------------------------------------------
|
||||
def test_a_mentioned_file_arrives_with_its_contents(client: TestClient, db, registered, box):
|
||||
profile = _profile(db, box)
|
||||
|
||||
client.post(
|
||||
"/api/files/from-project",
|
||||
data={"profile_id": profile.id, "path": "src/main.py"},
|
||||
)
|
||||
|
||||
attachment = db.scalars(select(Attachment)).one()
|
||||
assert "print('hello')" in attachment.extracted_text
|
||||
assert attachment.filename == "main.py"
|
||||
|
||||
|
||||
def test_the_model_is_told_which_file_and_where(client: TestClient, db, registered, box):
|
||||
"""The whole reason the columns exist. `main.py` alone is not an answer to
|
||||
"which one", and a model cannot name a file back that it was never given
|
||||
the path of."""
|
||||
profile = _profile(db, box)
|
||||
client.post(
|
||||
"/api/files/from-project",
|
||||
data={"profile_id": profile.id, "path": "src/main.py"},
|
||||
)
|
||||
|
||||
attachment = db.scalars(select(Attachment)).one()
|
||||
message = Message(chat_id=None, role="user", content="")
|
||||
message.attachments = [attachment]
|
||||
|
||||
block = chat_service.document_context(message)
|
||||
|
||||
assert 'path="src/main.py"' in block
|
||||
assert 'from="Container"' in block
|
||||
assert 'name="main.py"' in block
|
||||
|
||||
|
||||
def test_a_quote_in_a_path_cannot_break_out_of_the_tag(client: TestClient, db, registered):
|
||||
"""These are attribute values in a tag we write. A path containing a quote
|
||||
would otherwise close it early and the rest would read as instructions."""
|
||||
from lembas.services import files as files_service
|
||||
|
||||
user = db.scalars(select(User)).first()
|
||||
attachment = files_service.store_text(
|
||||
db,
|
||||
user_id=user.id,
|
||||
chat_id=None,
|
||||
filename="x.txt",
|
||||
text="body",
|
||||
source_path='/tmp/a"><script>.txt',
|
||||
source_label='Box" evil',
|
||||
)
|
||||
message = Message(chat_id=None, role="user", content="")
|
||||
message.attachments = [attachment]
|
||||
|
||||
block = chat_service.document_context(message)
|
||||
|
||||
assert "<script>" not in block
|
||||
assert block.count("<document") == 1
|
||||
|
||||
|
||||
def test_a_mentioned_directory_attaches_its_listing(client: TestClient, db, registered, box):
|
||||
""""@ that folder" is a reasonable thing to mean, and the listing is what
|
||||
it means -- better than refusing."""
|
||||
profile = _profile(db, box)
|
||||
|
||||
client.post("/api/files/from-project", data={"profile_id": profile.id, "path": "src/"})
|
||||
|
||||
attachment = db.scalars(select(Attachment)).one()
|
||||
assert "main.py" in attachment.extracted_text
|
||||
|
||||
|
||||
def test_somebody_elses_connection_cannot_be_read_from(client: TestClient, db, registered, box):
|
||||
profile = _profile(db, box)
|
||||
|
||||
client.post("/auth/logout", follow_redirects=False)
|
||||
client.post(
|
||||
"/auth/register",
|
||||
data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
|
||||
follow_redirects=False,
|
||||
)
|
||||
settings_store.update(db, {"default_permissions": {"tools.agent": True}})
|
||||
|
||||
body = client.post(
|
||||
"/api/files/from-project",
|
||||
data={"profile_id": profile.id, "path": "src/main.py"},
|
||||
).text
|
||||
|
||||
assert "not available" in body
|
||||
assert db.scalars(select(Attachment)).all() == []
|
||||
|
||||
|
||||
def test_a_missing_file_is_an_error_chip_not_a_crash(client: TestClient, db, registered, box):
|
||||
profile = _profile(db, box)
|
||||
|
||||
response = client.post(
|
||||
"/api/files/from-project",
|
||||
data={"profile_id": profile.id, "path": "nowhere.txt"},
|
||||
)
|
||||
|
||||
# 200 with a readable chip, like every other attach failure: htmx swaps the
|
||||
# body either way, and an error somebody can read beats a console message.
|
||||
assert response.status_code == 200
|
||||
assert "no file" in response.text.lower()
|
||||
|
||||
|
||||
def test_an_ordinary_upload_carries_no_provenance(client: TestClient, db, registered):
|
||||
"""A file dragged in from a laptop has no address this instance could
|
||||
honestly report, so the tag stays as it was."""
|
||||
from lembas.services import files as files_service
|
||||
|
||||
user = db.scalars(select(User)).first()
|
||||
attachment = files_service.store(
|
||||
db, user_id=user.id, chat_id=None, payload=b"hello", filename="notes.txt"
|
||||
)
|
||||
message = Message(chat_id=None, role="user", content="")
|
||||
message.attachments = [attachment]
|
||||
|
||||
block = chat_service.document_context(message)
|
||||
|
||||
assert "path=" not in block
|
||||
assert "from=" not in block
|
||||
|
||||
|
||||
# --- /usage ------------------------------------------------------------------
|
||||
def test_usage_sums_what_every_reply_recorded(client: TestClient, db, registered, make_chat):
|
||||
"""Summed from what was stored rather than recomputed. An endpoint that
|
||||
reported nothing contributed an estimate at the time, and re-deriving it
|
||||
now would make the totals move under a chat nobody had touched."""
|
||||
from lembas.db.models import ROLE_ASSISTANT, Chat
|
||||
|
||||
chat_id = make_chat()
|
||||
chat = db.get(Chat, chat_id)
|
||||
for prompt, completion in ((100, 20), (150, 30)):
|
||||
db.add(
|
||||
Message(
|
||||
chat_id=chat.id,
|
||||
role=ROLE_ASSISTANT,
|
||||
content="hi",
|
||||
usage_json={
|
||||
"prompt_tokens": prompt,
|
||||
"completion_tokens": completion,
|
||||
"total_tokens": prompt + completion,
|
||||
},
|
||||
)
|
||||
)
|
||||
db.commit()
|
||||
|
||||
body = client.get(f"/api/chats/{chat_id}/usage").text
|
||||
|
||||
assert "300" in body # 250 sent + 50 written
|
||||
assert "2 replies" in body
|
||||
|
||||
|
||||
def test_usage_is_owner_only(client: TestClient, db, registered, make_chat):
|
||||
"""No admin branch, unlike the inspector beside it: these numbers describe
|
||||
somebody's conversation, and reading one is not a configuration act."""
|
||||
chat_id = make_chat()
|
||||
|
||||
client.post("/auth/logout", follow_redirects=False)
|
||||
client.post(
|
||||
"/auth/register",
|
||||
data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
|
||||
follow_redirects=False,
|
||||
)
|
||||
|
||||
assert client.get(f"/api/chats/{chat_id}/usage").status_code == 404
|
||||
|
||||
|
||||
def test_an_unknown_window_size_is_not_reported_as_zero(
|
||||
client: TestClient, db, registered, make_chat
|
||||
):
|
||||
"""Unknown is not zero. Nobody has said how big this model's window is, so
|
||||
there is no percentage and compaction will never fire."""
|
||||
chat_id = make_chat()
|
||||
|
||||
body = client.get(f"/api/chats/{chat_id}/usage").text
|
||||
|
||||
assert "never compact itself" in body
|
||||
Reference in New Issue
Block a user