"""`@` 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 # --- What `@` offers where there is no machine at all ------------------------ def _library(db): """A note, a skill and a base belonging to the registered reader.""" from lembas.db.models import KnowledgeBase, Note, Skill owner = db.scalars(select(User)).first() note = Note(owner_id=owner.id, title="Mallorn notes", body="Golden leaves.") skill = Skill( owner_id=owner.id, name="bake-lembas", description="How to bake it", body="Steps." ) base = KnowledgeBase(owner_id=owner.id, name="Contracts") db.add_all([note, skill, base]) db.commit() return note, skill, base def test_notes_and_skills_are_offered_in_a_plain_chat(client: TestClient, db, registered): """A chat with no SSH connection has no project files, which is exactly why the rest of the library has to be reachable there.""" _library(db) body = client.get("/api/files/mention-picker", params={"q": "mallorn"}).text assert "Mallorn notes" in body body = client.get("/api/files/mention-picker", params={"q": "lembas"}).text assert "bake-lembas" in body def test_a_knowledge_base_is_only_offered_inside_a_chat( client: TestClient, db, registered, make_chat ): """There is nothing to attach it to before a chat exists -- the same reason project files are absent on the new-chat screen.""" _library(db) chat_id = make_chat() assert "Contracts" not in client.get("/api/files/mention-picker").text assert "Contracts" in client.get( "/api/files/mention-picker", params={"chat_id": chat_id} ).text def test_a_url_is_offered_as_a_page_to_read(client: TestClient, db, registered): body = client.get( "/api/files/mention-picker", params={"q": "https://tolkien.test/mallorn"} ).text assert "Fetch this page" in body assert "https://tolkien.test/mallorn" in body def test_a_note_arrives_with_its_text_and_its_name(client: TestClient, db, registered, make_chat): note, _skill, _base = _library(db) chat_id = make_chat() response = client.post( "/api/files/from-note", data={"note_id": note.id, "chat_id": chat_id} ) assert response.status_code == 200 attachment = db.scalars(select(Attachment)).one() assert "Golden leaves." in attachment.extracted_text # Provenance, for the reason a project file carries it: a model handed four # documents cannot name one back when asked to work on it. assert attachment.source_label == "Note" assert attachment.source_path == "Mallorn notes" def test_a_skill_can_be_handed_over_directly(client: TestClient, db, registered, make_chat): """The index is in the harness and `skill_get` fetches on demand -- but only if the model decides to. `@` is the reader saying "use this one".""" _note, skill, _base = _library(db) chat_id = make_chat() client.post("/api/files/from-skill", data={"skill_id": skill.id, "chat_id": chat_id}) attachment = db.scalars(select(Attachment)).one() assert attachment.source_label == "Skill" assert "Steps." in attachment.extracted_text def test_a_base_is_attached_as_a_reference_not_a_copy( client: TestClient, db, registered, make_chat ): """Scoping, not copying. A folder of contracts in the window would cost the context on every request forever to answer one question.""" from lembas.db.models import Chat _note, _skill, base = _library(db) chat_id = make_chat() response = client.post(f"/api/chats/{chat_id}/bases", data={"base_id": base.id}) assert response.status_code == 200 db.expire_all() assert [b.id for b in db.get(Chat, chat_id).knowledge_bases] == [base.id] # Nothing was copied into the message. assert db.scalars(select(Attachment)).all() == [] def test_attaching_the_same_base_twice_is_not_an_error( client: TestClient, db, registered, make_chat ): from lembas.db.models import Chat _note, _skill, base = _library(db) chat_id = make_chat() client.post(f"/api/chats/{chat_id}/bases", data={"base_id": base.id}) client.post(f"/api/chats/{chat_id}/bases", data={"base_id": base.id}) db.expire_all() assert len(db.get(Chat, chat_id).knowledge_bases) == 1 def test_a_library_document_now_carries_its_provenance(client: TestClient, db, registered): """This was the one attach path that dropped it, while a project file beside it carried path and machine.""" from lembas.services.fetch import Fetched from lembas.services.library import documents as documents_service owner = db.scalars(select(User)).first() base = documents_service.default_base(db, owner) document = documents_service.store_page( db, owner=owner, base=base, page=Fetched(url="https://tolkien.test/c", title="The Contract", text="Terms."), ) client.post("/api/files/from-knowledge", data={"document_id": document.id}) attachment = db.scalars(select(Attachment)).one() assert attachment.source_path == "The Contract" assert attachment.source_label == base.name # --- 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">