DRU-1005 -- Write sandbox uploads in blocks the SFTP server accepts

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}
 
 

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."""
 

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,

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(

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
 ):