diff --git a/backend/druks/chat/routes.py b/backend/druks/chat/routes.py index 6da461f3..b69ea5ad 100644 --- a/backend/druks/chat/routes.py +++ b/backend/druks/chat/routes.py @@ -11,6 +11,7 @@ from druks.accounts.dependencies import current_account, current_session_account from druks.accounts.models import Account from druks.api.dependencies import SessionDep +from druks.sandbox.exceptions import SandboxUploadError from .channels.whatsapp.exceptions import NumberNotOnWhatsAppError, WahaError from .enums import MessageState @@ -239,7 +240,10 @@ async def read_file( raise HTTPException(404, str(error)) from error except (SlackApiError, RequestFailed, WahaError) as error: raise HTTPException(502, str(error)) from error - path = await upload_file(session, conversation, file_id, name, content) + try: + path = await upload_file(session, conversation, file_id, name, content) + except SandboxUploadError as error: + raise HTTPException(502, str(error)) from error return {"path": path, "name": name, "content_type": content_type} diff --git a/backend/druks/sandbox/exceptions.py b/backend/druks/sandbox/exceptions.py index 44c2a702..328438ae 100644 --- a/backend/druks/sandbox/exceptions.py +++ b/backend/druks/sandbox/exceptions.py @@ -9,6 +9,10 @@ class SandboxDownloadError(SandboxError): pass +class SandboxUploadError(SandboxError): + pass + + class SetupScriptError(SandboxError): """A declared setup path cannot resolve to package bytes.""" diff --git a/backend/druks/sandbox/host.py b/backend/druks/sandbox/host.py index 7e9e795d..74a50cef 100644 --- a/backend/druks/sandbox/host.py +++ b/backend/druks/sandbox/host.py @@ -43,7 +43,7 @@ Hook, McpServer, ) -from .exceptions import SandboxDownloadError, SandboxError +from .exceptions import SandboxDownloadError, SandboxError, SandboxUploadError from .layout import get_helper_script_path, get_work_root if TYPE_CHECKING: @@ -64,6 +64,10 @@ _CONNECT_TIMEOUT_SECONDS = 30.0 _KEEPALIVE_INTERVAL_SECONDS = 15.0 +# OpenSSH's sftp-server drops the session on a packet over 256 KiB and advertises a +# max-write-length of 255 KiB. asyncssh sizes writes from the advertised limits, and +# a gateway in front of sftp-server can advertise more than sftp-server takes. +_SFTP_WRITE_BYTES = 255 * 1024 class Host: @@ -107,14 +111,20 @@ async def upload_file( ) -> None: """SFTP-put ``local`` to ``remote``, mkdir-p its parent, chmod to ``mode``. The default mode matches credential files (operator - wants tight perms by default).""" + wants tight perms by default). A failed connect or transfer raises + ``SandboxUploadError``.""" parent = remote.rsplit("/", 1)[0] if "/" in remote.lstrip("/") else "" - if parent: - await self.exec(["mkdir", "-p", parent], timeout=10.0) - conn = await self._ensure_conn() - async with conn.start_sftp_client() as sftp: - await sftp.put(str(local), remote) - await self.exec(["chmod", f"{mode:o}", remote], timeout=10.0) + try: + if parent: + await self.exec(["mkdir", "-p", parent], timeout=10.0) + conn = await self._ensure_conn() + async with conn.start_sftp_client() as sftp: + await sftp.put(str(local), remote, block_size=_SFTP_WRITE_BYTES) + await self.exec(["chmod", f"{mode:o}", remote], timeout=10.0) + except (asyncssh.Error, OSError) as error: + raise SandboxUploadError( + f"failed to upload {remote}: {type(error).__name__}: {error}" + ) from error async def write_secret( self, diff --git a/backend/tests/test_sandbox_host.py b/backend/tests/test_sandbox_host.py index d3d26353..0bebe74e 100644 --- a/backend/tests/test_sandbox_host.py +++ b/backend/tests/test_sandbox_host.py @@ -6,9 +6,10 @@ from typing import Any, Self from unittest.mock import AsyncMock +import asyncssh import pytest from drukbox_sdk import SandboxHost as SandboxHostRecord -from druks.sandbox.exceptions import SandboxError +from druks.sandbox.exceptions import SandboxError, SandboxUploadError from druks.sandbox.host import Host, _stream_local_tar_into @@ -32,11 +33,11 @@ def returncode(self) -> int | None: class _FakeSFTP: def __init__(self) -> None: - self.puts: list[tuple[str, str]] = [] + self.puts: list[tuple[str, str, int]] = [] self.gets: list[tuple[str, str]] = [] - async def put(self, local: str, remote: str) -> None: - self.puts.append((local, remote)) + async def put(self, local: str, remote: str, *, block_size: int) -> None: + self.puts.append((local, remote, block_size)) async def get(self, remote: str, local: str) -> None: self.gets.append((remote, local)) @@ -245,7 +246,25 @@ async def test_upload_pushes_via_sftp( local.write_text('{"token": "..."}') await sandbox.upload_file(local=local, remote="/creds/anthropic.json") - assert fake_conn.sftp.puts == [(str(local), "/creds/anthropic.json")] + assert fake_conn.sftp.puts == [(str(local), "/creds/anthropic.json", 255 * 1024)] + + +async def test_upload_raises_a_sandbox_upload_error_when_the_connect_or_transfer_fails( + fake_record: SandboxHostRecord, + patched_asyncssh: tuple[AsyncMock, _FakeConnection], + tmp_path: Path, +): + connect_mock, fake_conn = patched_asyncssh + local = tmp_path / "plan.pdf" + local.write_bytes(b"%PDF") + + fake_conn.sftp.put = AsyncMock(side_effect=asyncssh.SFTPConnectionLost("Connection lost")) + with pytest.raises(SandboxUploadError, match="/work/plan.pdf: SFTPConnectionLost"): + await Host(record=fake_record).upload_file(local=local, remote="/work/plan.pdf") + + connect_mock.side_effect = ConnectionRefusedError("refused") + with pytest.raises(SandboxUploadError, match="/work/plan.pdf: ConnectionRefusedError"): + await Host(record=fake_record).upload_file(local=local, remote="/work/plan.pdf") async def test_download_fetches_via_sftp_and_creates_parent_dir( diff --git a/backend/tests/test_slack.py b/backend/tests/test_slack.py index 68d67107..9eb866e5 100644 --- a/backend/tests/test_slack.py +++ b/backend/tests/test_slack.py @@ -24,6 +24,7 @@ from druks.core.webhooks.slack import SlackEvents, SlackInteractivity from druks.models import Base from druks.redis import close_client, get_client +from druks.sandbox.exceptions import SandboxUploadError from druks.sandbox.layout import get_work_root from druks.secrets.datastructures import Audience from druks.secrets.models import VaultSecret @@ -604,6 +605,37 @@ async def upload_file(*, local, remote, mode=0o600): assert (elsewhere.status_code, outside.status_code) == (404, 409) +async def test_chat_read_file_answers_502_when_the_upload_to_the_sandbox_fails( + card, druks_db, tmp_path, monkeypatch +): + """A failed upload answers which file failed and why, not a bare 500.""" + configure_app_for_test(settings=make_settings(tmp_path), authenticated=False) + api = FastAPI(dependencies=[Depends(request_session)]) + api.add_api_route("/thread/files/{file_id}", read_file, methods=["GET"]) + + async def api_call(client, api_method, *, params=None, json=None, **request): + shared = {**shared_file(params["file"]), "channels": ["C1"], "groups": []} + return Answer(ok=True, file=shared) + + error = SandboxUploadError("failed to upload plan.pdf: SFTPConnectionLost: Connection lost") + upload_file = AsyncMock(side_effect=error) + host = SimpleNamespace(ssh_username="agent", upload_file=upload_file, aclose=AsyncMock()) + monkeypatch.setattr(SlackClient, "api_call", api_call) + monkeypatch.setattr(SlackClient, "download", AsyncMock(return_value=b"%PDF")) + monkeypatch.setattr("druks.chat.service.get_running_sandbox", AsyncMock(return_value=host)) + ana = await Account.get_or_create(druks_db, "ana@example.com") + conversation = await thread_conversation(druks_db, card, ana) + _, token = await PersonalAccessToken.create(druks_db, account_id=ana.id, name="chat") + await druks_db.commit() + headers = {"Authorization": f"Bearer {token}", CONVERSATION_HEADER: conversation.id} + + async with asgi_client(api) as client: + response = await client.get("/thread/files/F1", headers=headers) + + assert response.status_code == 502 + assert response.json()["detail"] == str(error) + + async def test_a_person_who_requires_a_tag_is_answered_only_when_they_tag( card, druks_db, slack, delivery, tmp_path ):