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

@@ -120,15 +120,20 @@ async def create_file_or_directory(
                         $ref: "#/components/schemas/FileCreateResponse"
     """
     try:
-        parsed = await parse_multipart_request(
+        async with 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"),
-        )
+        ) as parsed:
+            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

@@ -5,6 +5,7 @@
 import subprocess
 import sys
 import webbrowser
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from pathlib import Path
 from shutil import which
@@ -25,6 +26,9 @@
 from marimo._utils.parse_dataclass import parse_raw
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from starlette.datastructures import UploadFile
     from starlette.requests import Request
 
     from marimo._session.session import Session
@@ -44,20 +48,31 @@ 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]
 
 
+@asynccontextmanager
 async def parse_multipart_request(
     request: Request, cls: type[S]
-) -> MultipartRequest[S]:
-    """Parse a multipart/form-data body into a msgspec.Struct + file bytes.
+) -> AsyncIterator[MultipartRequest[S]]:
+    """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).
+    yielded as `UploadFile` objects (un-read) in `files`, keyed by
+    form-field name, so callers can stream them to disk instead of buffering.
+
+    Used as an async context manager so the underlying form (and its
+    spooled temp files / fds) are closed automatically when the caller
+    is done — `UploadFile` parts remain readable for the entire body of
+    the `async with` block.
 
     Raises msgspec.ValidationError if required string fields are missing
     or invalid.
@@ -66,18 +81,16 @@ 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] = {}
+        files: dict[str, UploadFile] = {}
         for key, value in form.multi_items():
             if isinstance(value, UploadFile):
-                files[key] = await value.read()
+                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)
+        body = msgspec.convert(string_payload, cls, strict=False)
+        yield MultipartRequest(body=body, files=files)
 
 
 @runtime_checkable

marimo/_server/files/os_file_system.py

@@ -8,9 +8,10 @@
 import re
 import shutil
 import subprocess
+import tempfile
 from collections import deque
 from pathlib import Path
-from typing import Literal
+from typing import Literal, Protocol
 
 from marimo import _loggers
 from marimo._server.files.file_system import FileSystem
@@ -36,6 +37,25 @@
     "..",
 ]
 
+# 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
+
+
+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,62 @@ 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)
+
+        # `NamedTemporaryFile` gives us a guaranteed-unique sibling path so
+        # concurrent uploads racing through `_generate_unique_path` can't
+        # collide on the same `.part` file.
+        tmp = tempfile.NamedTemporaryFile(
+            dir=full_path.parent,
+            prefix=full_path.name + ".",
+            suffix=".part",
+            delete=False,
+        )
+        tmp_path = tmp.name
+        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 tmp:
+                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"
+                        )
+                    tmp.write(chunk)
+            os.replace(tmp_path, full_path)
+        except BaseException:
+            try:
+                os.unlink(tmp_path)
+            except FileNotFoundError:
+                pass
+            raise
+
+        # Use the metadata-only helper: `get_details` would re-read the
+        # file contents (and base64-encode binary), defeating the point of
+        # streaming for large uploads.
+        return self._get_file_info(str(full_path))
+
     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

@@ -23,16 +23,19 @@ class _SampleForm(msgspec.Struct):
 
 def _build_app(captured: dict[str, object]) -> TestClient:
     async def endpoint(request: Request) -> JSONResponse:
-        parsed = await parse_multipart_request(request, _SampleForm)
-        captured["body"] = parsed.body
-        captured["files"] = parsed.files
+        async with parse_multipart_request(request, _SampleForm) as parsed:
+            captured["body"] = parsed.body
+            captured["files"] = dict(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: