refactor(BA-6156): type the upload session id as TusSessionId

UploadTokenData.session is now typed as the common TusSessionId NewType so the
session id flows type-safely from the token through the handlers into the
Valkey-backed engine.

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

Author: jopemachine

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

@@ -42,18 +42,20 @@
 from ai.backend.common.metrics.http import build_api_metric_middleware
 from ai.backend.common.middlewares.exception import general_exception_middleware
 from ai.backend.common.typed_validators import PydanticJWTValidator
-from ai.backend.common.types import BinarySize, VFolderID
+from ai.backend.common.types import BinarySize, TusSessionId, VFolderID
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.storage import __version__
 from ai.backend.storage.dto.context import StorageRootCtx
-from ai.backend.storage.errors import InvalidAPIParameters, UploadOffsetMismatchError
+from ai.backend.storage.errors import (
+    InvalidAPIParameters,
+    UploadChunkExceedsTotalSizeError,
+    UploadOffsetMismatchError,
+)
 from ai.backend.storage.services.file_stream.zip import (
     ZipArchiveStreamReader,
 )
-from ai.backend.storage.services.upload.tus_session import (
-    TusUploadSession,
-    TusUploadSessionArgs,
-)
+from ai.backend.storage.services.upload.tus_session import TusUploadSession
+from ai.backend.storage.services.upload.types import TusUploadSessionArgs
 from ai.backend.storage.types import SENTINEL, TusChunkUploadStreamReader
 from ai.backend.storage.utils import (
     CheckParamSource,
@@ -99,7 +101,7 @@ class UploadTokenData(TypedDict):
     volume: str
     vfid: VFolderID
     relpath: str
-    session: str
+    session: TusSessionId
     size: int
 
 
@@ -349,10 +351,11 @@ class Params(TypedDict):
                     session_id=token_data["session"],
                     total_size=int(token_data["size"]),
                     valkey_client=ctx.valkey_tus_client,
+                    lock_factory=ctx.tus_lock_factory,
                 )
             )
             state = await session.read_state()
-            headers = _tus_response_headers(
+            headers = _prepare_tus_session_headers(
                 upload_offset=state.committed_offset,
                 upload_length=int(token_data["size"]),
             )
@@ -428,31 +431,30 @@ class Params(TypedDict):
                     session_id=token_data["session"],
                     total_size=total_size,
                     valkey_client=ctx.valkey_tus_client,
+                    lock_factory=ctx.tus_lock_factory,
                 )
             )
             await session.ensure_initialized()
 
             upload_stream = TusChunkUploadStreamReader(
                 request.content, request.content_type, DEFAULT_CHUNK_SIZE
             )
-            temp_chunk, length, sha256 = await session.write_temp_chunk(
-                client_offset, upload_stream
-            )
+            written = await session.write_temp_chunk(client_offset, upload_stream)
             try:
-                if client_offset + length > total_size:
-                    raise UploadOffsetMismatchError(
-                        f"Chunk at offset {client_offset} with length {length} "
+                if client_offset + written.length > total_size:
+                    raise UploadChunkExceedsTotalSizeError(
+                        f"Chunk at offset {client_offset} with length {written.length} "
                         f"exceeds declared size {total_size}"
                     )
                 acceptance = await session.commit_chunk(
                     offset=client_offset,
-                    chunk_path=temp_chunk.path,
-                    length=length,
-                    sha256=sha256,
+                    chunk_path=written.path,
+                    length=written.length,
+                    sha256=written.sha256,
                 )
             except BaseException:
-                if temp_chunk.path.exists():
-                    await asyncio.to_thread(temp_chunk.path.unlink)
+                if written.path.exists():
+                    await asyncio.to_thread(written.path.unlink)
                 raise
 
             state = acceptance.state
@@ -464,7 +466,7 @@ class Params(TypedDict):
                 await session.assemble(target_path)
                 await session.cleanup()
 
-            headers = _tus_response_headers(
+            headers = _prepare_tus_session_headers(
                 upload_offset=state.committed_offset,
                 upload_length=total_size,
             )
@@ -478,7 +480,7 @@ def _resolve_tus_upload_session_dir(volume: AbstractVolume, token_data: UploadTo
 _TUS_HEADER_LIST = "Tus-Resumable, Upload-Length, Upload-Metadata, Upload-Offset, Content-Type"
 
 
-def _tus_response_headers(*, upload_offset: int, upload_length: int) -> dict[str, str]:
+def _prepare_tus_session_headers(*, upload_offset: int, upload_length: int) -> dict[str, str]:
     return {
         "Access-Control-Allow-Origin": "*",
         "Access-Control-Allow-Headers": _TUS_HEADER_LIST,

src/ai/backend/storage/errors/__init__.py

@@ -64,6 +64,7 @@
 )
 from .upload import (
     ChunkConflictError,
+    UploadChunkExceedsTotalSizeError,
     UploadSessionCorruptedError,
 )
 from .vfolder import (
@@ -96,6 +97,7 @@
     "UploadOffsetMismatchError",
     # upload
     "ChunkConflictError",
+    "UploadChunkExceedsTotalSizeError",
     "UploadSessionCorruptedError",
     # vfolder
     "VFolderNotFoundError",

src/ai/backend/storage/errors/upload.py

@@ -32,6 +32,24 @@ def error_code(self) -> ErrorCode:
         )
 
 
+class UploadChunkExceedsTotalSizeError(BackendAIError, web.HTTPConflict):
+    """
+    Raised when a PATCH chunk's offset+length would write past the declared
+    ``Upload-Length`` (409 Conflict). The Upload-Offset header itself is in
+    range; the chunk's body simply overruns the remaining slot.
+    """
+
+    error_type = "https://api.backend.ai/probs/storage/upload-chunk-exceeds-total-size"
+    error_title = "Upload Chunk Exceeds Total Size"
+
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.STORAGE_PROXY,
+            operation=ErrorOperation.UPDATE,
+            error_detail=ErrorDetail.CONFLICT,
+        )
+
+
 class UploadSessionCorruptedError(BackendAIError, web.HTTPInternalServerError):
     """
     Raised when the on-disk upload session metadata cannot be parsed or is