feat(files): stream uploads to disk instead of buffering

Previously /api/files/create read the entire UploadFile into memory
before writing to disk, peaking at ~100 MB for a 100 MB upload. Now we
drain the upload in 1 MiB chunks straight to a `.part` temp file and
atomically rename on success, so peak memory stays bounded regardless
of file size and a failed upload never leaves a half-written file at
the final path.

- `OSFileSystem.stream_create_file` does the chunked write with atomic
  rename and cleanup on failure
- `parse_multipart_request` now hands callers un-read `UploadFile`
  handles (instead of pre-read bytes), so streaming is possible without
  giving up the small-payload `.read()` path
- Directories and the default-template notebook still go through the
  in-memory `create_file_or_directory` path; only real file content
  streams

Author: mscolnick

marimo/_server/api/endpoints/file_explorer.py

@@ -123,12 +123,17 @@ async def create_file_or_directory(
         parsed = await parse_multipart_request(
             request, FileCreateMultipartRequest
         )
-        info = file_system.create_file_or_directory(
-            parsed.body.path,
-            parsed.body.type,
-            parsed.body.name,
-            parsed.files.get("file"),
-        )
+        upload = parsed.files.get("file")
+        # Stream when there's actual file content; the in-memory create
+        # path still handles directories and the default-template notebook.
+        if upload is not None and parsed.body.type in ("file", "notebook"):
+            info = await file_system.stream_create_file(
+                parsed.body.path, parsed.body.name, upload
+            )
+        else:
+            info = file_system.create_file_or_directory(
+                parsed.body.path, parsed.body.type, parsed.body.name, None
+            )
         return FileCreateResponse(success=True, info=info)
     except Exception as e:
         LOGGER.error(f"Error creating file or directory: {e}")

marimo/_server/api/utils.py

@@ -25,6 +25,7 @@
 from marimo._utils.parse_dataclass import parse_raw
 
 if TYPE_CHECKING:
+    from starlette.datastructures import UploadFile
     from starlette.requests import Request
 
     from marimo._session.session import Session
@@ -44,20 +45,25 @@ async def parse_request(
 
 @dataclass
 class MultipartRequest(Generic[S]):
-    """Result of parsing a multipart/form-data request body."""
+    """Result of parsing a multipart/form-data request body.
+
+    `files` carries the raw `UploadFile` objects (not yet read) so callers
+    can stream chunks via `.read(size)`. For small payloads, callers can
+    simply `await upload.read()` to materialize the whole body.
+    """
 
     body: S
-    files: dict[str, bytes]
+    files: dict[str, UploadFile]
 
 
 async def parse_multipart_request(
     request: Request, cls: type[S]
 ) -> MultipartRequest[S]:
-    """Parse a multipart/form-data body into a msgspec.Struct + file bytes.
+    """Parse a multipart/form-data body into a msgspec.Struct + uploads.
 
     String form fields are validated against `cls`. File upload parts are
-    read fully into memory and returned in `files`, keyed by form-field
-    name (callers look them up explicitly rather than via the struct).
+    returned as `UploadFile` objects (un-read) in `files`, keyed by
+    form-field name, so callers can stream them to disk instead of buffering.
 
     Raises msgspec.ValidationError if required string fields are missing
     or invalid.
@@ -66,16 +72,17 @@ async def parse_multipart_request(
     # without starlette (e.g. pyodide).
     from starlette.datastructures import UploadFile
 
-    # Use as an async context manager so any spooled temp files backing
-    # UploadFile parts are closed after parsing.
-    async with request.form() as form:
-        string_payload: dict[str, Any] = {}
-        files: dict[str, bytes] = {}
-        for key, value in form.multi_items():
-            if isinstance(value, UploadFile):
-                files[key] = await value.read()
-            elif isinstance(value, str):
-                string_payload[key] = value
+    # Do NOT use `async with request.form()` here: callers stream the
+    # UploadFile parts after this function returns, and the context
+    # manager would close their spooled temp files on exit.
+    form = await request.form()
+    string_payload: dict[str, Any] = {}
+    files: dict[str, UploadFile] = {}
+    for key, value in form.multi_items():
+        if isinstance(value, UploadFile):
+            files[key] = value
+        elif isinstance(value, str):
+            string_payload[key] = value
     body = msgspec.convert(string_payload, cls, strict=False)
     return MultipartRequest(body=body, files=files)
 

marimo/_server/files/os_file_system.py

@@ -10,7 +10,7 @@
 import subprocess
 from collections import deque
 from pathlib import Path
-from typing import Literal
+from typing import Literal, Protocol, runtime_checkable
 
 from marimo import _loggers
 from marimo._server.files.file_system import FileSystem
@@ -36,6 +36,26 @@
     "..",
 ]
 
+# 1 MiB. Large enough to amortize syscall overhead, small enough to keep
+# peak memory bounded when streaming.
+_STREAM_CHUNK_SIZE = 1024 * 1024
+
+# Hard cap on streamed uploads. Streaming removes the implicit OOM ceiling
+# that buffered uploads had, so without a cap an authenticated client could
+# exhaust disk. 1 GiB covers normal notebook-data use cases with margin.
+MAX_UPLOAD_BYTES = 1024 * 1024 * 1024
+
+
+@runtime_checkable
+class AsyncByteSource(Protocol):
+    """Anything that can be drained chunk-by-chunk into a file.
+
+    Starlette's `UploadFile` satisfies this; so does any object exposing
+    an async `read(size)` returning bytes.
+    """
+
+    async def read(self, size: int = -1, /) -> bytes: ...
+
 
 class OSFileSystem(FileSystem):
     def get_root(self) -> str:
@@ -133,33 +153,32 @@ def open_file(self, path: str, encoding: str | None = None) -> str | bytes:
         except UnicodeDecodeError:
             return file_path.read_bytes()
 
-    def create_file_or_directory(
-        self,
-        path: str,
-        file_type: Literal["file", "directory", "notebook"],
-        name: str,
-        contents: bytes | None,
-    ) -> FileInfo:
+    @staticmethod
+    def _validate_create_name(name: str) -> None:
+        """Reject names that are empty, reserved, or traverse out of the
+        parent. Centralized so HTTP, WASM, and streaming paths all share it.
+        """
         if name in DISALLOWED_NAMES:
             raise ValueError(
                 f"Cannot create file or directory with name {name}"
             )
         if name.strip() == "":
             raise ValueError("Cannot create file or directory with empty name")
-        # Names that traverse out of `path` or escape via separators are
-        # rejected. Validation belongs here (not in the endpoint) so every
-        # caller of OSFileSystem — HTTP, WASM bridge, scripts — is covered.
-        if (
-            "/" in name
-            or "\\" in name
-            or "\x00" in name
-            or name in (".", "..")
-        ):
+        if "/" in name or "\\" in name or "\x00" in name:
             raise ValueError(
                 f"Invalid name {name!r}: must not contain path separators "
                 "or refer to a parent directory"
             )
 
+    def create_file_or_directory(
+        self,
+        path: str,
+        file_type: Literal["file", "directory", "notebook"],
+        name: str,
+        contents: bytes | None,
+    ) -> FileInfo:
+        self._validate_create_name(name)
+
         full_path = Path(path) / name
         full_path = _generate_unique_path(full_path)
 
@@ -192,6 +211,49 @@ def create_file_or_directory(
             ),
         ).file
 
+    async def stream_create_file(
+        self,
+        path: str,
+        name: str,
+        source: AsyncByteSource,
+    ) -> FileInfo:
+        """Stream-write an uploaded file to disk, chunk by chunk.
+
+        Avoids loading the full payload into memory (the HTTP multipart
+        path can otherwise buffer 100 MB at once). Writes to a ``.part``
+        temp file and atomically renames on success so a failed upload
+        doesn't leave a half-written file at the final path.
+        """
+        self._validate_create_name(name)
+
+        full_path = Path(path) / name
+        full_path = _generate_unique_path(full_path)
+        full_path.parent.mkdir(parents=True, exist_ok=True)
+
+        tmp_path = full_path.with_name(full_path.name + ".part")
+        try:
+            # Sync writes are bounded to ~1 MiB per chunk, with an `await`
+            # in between; event loop blockage is brief and an async file
+            # library would only add a dependency for marginal gain.
+            written = 0
+            with open(tmp_path, "wb") as out:  # noqa: ASYNC230
+                while chunk := await source.read(_STREAM_CHUNK_SIZE):
+                    written += len(chunk)
+                    if written > MAX_UPLOAD_BYTES:
+                        raise ValueError(
+                            f"Upload exceeds maximum size of "
+                            f"{MAX_UPLOAD_BYTES} bytes"
+                        )
+                    out.write(chunk)
+            tmp_path.replace(full_path)
+        except BaseException:
+            tmp_path.unlink(missing_ok=True)
+            raise
+
+        # Read details fresh from disk; we deliberately don't pass `contents`
+        # since the file may be too large to round-trip through memory.
+        return self.get_details(str(full_path)).file
+
     def delete_file_or_directory(self, path: str) -> bool:
         if os.path.isdir(path):
             safe_rmtree(path)

marimo/_server/models/files.py

@@ -58,7 +58,13 @@ class FileCreateRequest(msgspec.Struct, rename="camel"):
 
 
 class FileCreateMultipartRequest(msgspec.Struct, rename="camel"):
-    """multipart/form-data body for POST /api/files/create."""
+    """multipart/form-data body for POST /api/files/create.
+
+    Schema-only: this struct exists to describe the multipart shape in
+    OpenAPI. At runtime, the endpoint reads the string fields from
+    ``MultipartRequest.body`` and the uploaded bytes from
+    ``MultipartRequest.files["file"]`` — ``body.file`` is never populated.
+    """
 
     # The path where to create the file or directory
     path: str
@@ -68,6 +74,9 @@ class FileCreateMultipartRequest(msgspec.Struct, rename="camel"):
     name: str
     # The raw file bytes (optional). When omitted, an empty file is created
     # (or, for 'notebook' type, a default notebook template).
+    # NOTE: this field is OpenAPI-only — see class docstring. The
+    # ``format: binary`` annotation makes the generated spec emit a proper
+    # file-upload schema rather than a base64 string.
     file: Annotated[
         str | None, msgspec.Meta(extra_json_schema={"format": "binary"})
     ] = None

packages/openapi/api.yaml

@@ -1665,7 +1665,11 @@ components:
       title: FileCopyResponse
       type: object
     FileCreateMultipartRequest:
-      description: multipart/form-data body for POST /api/files/create.
+      description: "multipart/form-data body for POST /api/files/create.\n\n    Schema-only:\
+        \ this struct exists to describe the multipart shape in\n    OpenAPI. At runtime,\
+        \ the endpoint reads the string fields from\n    ``MultipartRequest.body``\
+        \ and the uploaded bytes from\n    ``MultipartRequest.files[\"file\"]`` \u2014\
+        \ ``body.file`` is never populated."
       properties:
         file:
           anyOf:

packages/openapi/src/api.ts

@@ -4359,6 +4359,11 @@ export interface components {
     /**
      * FileCreateMultipartRequest
      * @description multipart/form-data body for POST /api/files/create.
+     *
+     *         Schema-only: this struct exists to describe the multipart shape in
+     *         OpenAPI. At runtime, the endpoint reads the string fields from
+     *         ``MultipartRequest.body`` and the uploaded bytes from
+     *         ``MultipartRequest.files["file"]`` — ``body.file`` is never populated.
      */
     FileCreateMultipartRequest: {
       /**

tests/_server/api/test_api_utils.py

@@ -26,13 +26,16 @@ async def endpoint(request: Request) -> JSONResponse:
         parsed = await parse_multipart_request(request, _SampleForm)
         captured["body"] = parsed.body
         captured["files"] = parsed.files
+        upload = parsed.files.get("upload")
+        if upload is not None:
+            captured["upload_bytes"] = await upload.read()
         return JSONResponse({"ok": True})
 
     app = Starlette(routes=[Route("/test", endpoint, methods=["POST"])])
     return TestClient(app)
 
 
-def test_parse_multipart_request_strings_and_file_bytes() -> None:
+def test_parse_multipart_request_strings_and_file_upload() -> None:
     captured: dict[str, object] = {}
     client = _build_app(captured)
     response = client.post(
@@ -45,7 +48,11 @@ def test_parse_multipart_request_strings_and_file_bytes() -> None:
     assert isinstance(body, _SampleForm)
     assert body.name == "marimo"
     assert body.count == 42
-    assert captured["files"] == {"upload": b"\x00\x01\x02\xff"}
+    files = captured["files"]
+    assert isinstance(files, dict)
+    assert set(files.keys()) == {"upload"}
+    # File handle is returned un-read; bytes loaded inside the endpoint above.
+    assert captured["upload_bytes"] == b"\x00\x01\x02\xff"
 
 
 def test_parse_multipart_request_omitted_file_yields_empty_dict() -> None: