refactor(BA-6156): stream PATCH body via common StreamReader adapter

Wrap the raw TUS PATCH body in TusChunkUploadStreamReader (in storage/types.py,
mirroring MultipartFileUploadStreamReader) so it conforms to the common
StreamReader abstraction, and drain it through TusUploadSession.write_temp_chunk
instead of a one-off reader protocol + free function. This keeps the reader a
pure byte source and concentrates all chunk file I/O in the session, consistent
with the existing storage upload/download readers.

Also update the completed-upload assertion to match the race-safe cleanup: the
session directory and its info.json marker are intentionally kept while chunk
data is reclaimed.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: jopemachine

src/ai/backend/storage/api/client.py

@@ -50,11 +50,8 @@
 from ai.backend.storage.services.file_stream.zip import (
     ZipArchiveStreamReader,
 )
-from ai.backend.storage.services.upload.tus_session import (
-    TusUploadSession,
-    stream_chunk_to_temp,
-)
-from ai.backend.storage.types import SENTINEL
+from ai.backend.storage.services.upload.tus_session import TusUploadSession
+from ai.backend.storage.types import SENTINEL, TusChunkUploadStreamReader
 from ai.backend.storage.utils import (
     CheckParamSource,
     build_attachment_headers,
@@ -426,11 +423,13 @@ class Params(TypedDict):
             )
             await session.ensure_initialized()
 
-            temp_chunk = session.open_temp_chunk(client_offset)
+            upload_stream = TusChunkUploadStreamReader(
+                request.content, request.content_type, DEFAULT_CHUNK_SIZE
+            )
+            temp_chunk, length, sha256 = await session.write_temp_chunk(
+                client_offset, upload_stream
+            )
             try:
-                length, sha256 = await stream_chunk_to_temp(
-                    request.content, temp_chunk.path, DEFAULT_CHUNK_SIZE
-                )
                 if client_offset + length > total_size:
                     raise UploadOffsetMismatchError(
                         f"Chunk at offset {client_offset} with length {length} "

src/ai/backend/storage/types.py

@@ -10,6 +10,7 @@
 import attrs
 import trafaret as t
 from aiohttp import BodyPartReader, MultipartReader, web
+from aiohttp import StreamReader as AiohttpStreamReader
 
 from ai.backend.common import validators as tx
 from ai.backend.common.types import QuotaConfig, StreamReader, VFolderID
@@ -148,3 +149,33 @@ async def read(self) -> AsyncIterator[bytes]:
     @override
     def content_type(self) -> str | None:
         return self._content_type
+
+
+class TusChunkUploadStreamReader(StreamReader):
+    """
+    Adapts a raw TUS PATCH request body (``application/offset+octet-stream``)
+    into the common :class:`StreamReader` by reading it in fixed-size chunks,
+    so the upload writer holds only one chunk in memory at a time.
+    """
+
+    def __init__(
+        self,
+        content: AiohttpStreamReader,
+        content_type: str | None,
+        chunk_size: int,
+    ) -> None:
+        self._content = content
+        self._content_type = content_type
+        self._chunk_size = chunk_size
+
+    @override
+    async def read(self) -> AsyncIterator[bytes]:
+        while True:
+            chunk = await self._content.read(self._chunk_size)
+            if not chunk:
+                break
+            yield chunk
+
+    @override
+    def content_type(self) -> str | None:
+        return self._content_type

tests/unit/storage/api/test_tus_upload.py

@@ -212,8 +212,12 @@ async def test_single_chunk_upload_completes_and_assembles(self, patch_env: _Pat
         assert response.headers["Upload-Offset"] == str(len(payload))
         final_path = patch_env.vfpath / "result.bin"
         assert final_path.read_bytes() == payload
-        # Session directory must be cleaned up after assembly.
-        assert not patch_env.session_dir.exists()
+        # After assembly the session reclaims chunk data but keeps the small
+        # completed marker (cleanup() no longer rmtree's the dir, so a late
+        # duplicate PATCH observes status=="completed" instead of racing a
+        # directory teardown).
+        assert patch_env.session_dir.exists()
+        assert list((patch_env.session_dir / "chunks").glob("*.dat")) == []
 
     async def test_two_chunks_assemble_in_order(self, patch_env: _PatchEnv) -> None:
         token_data = _token_data(session_id="test-session", total_size=2048, relpath="result.bin")