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,7 +42,7 @@
 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
@@ -99,7 +99,7 @@ class UploadTokenData(TypedDict):
     volume: str
     vfid: VFolderID
     relpath: str
-    session: str
+    session: TusSessionId
     size: int
 
 
@@ -349,10 +349,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,6 +429,7 @@ 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()
@@ -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,

tests/unit/storage/api/test_tus_upload.py

@@ -24,10 +24,12 @@
 from ai.backend.common import config
 from ai.backend.common.clients.valkey_client.valkey_tus import ValkeyTusClient
 from ai.backend.common.defs import REDIS_STREAM_LOCK, REDIS_TUS_DB
+from ai.backend.common.lock import DistributedLockFactory
 from ai.backend.common.typed_validators import HostPortPair as HostPortPairModel
 from ai.backend.common.types import RedisConnectionInfo, ValkeyTarget
 from ai.backend.storage.api.client import tus_upload_part
 from ai.backend.storage.errors import InvalidAPIParameters, UploadOffsetMismatchError
+from ai.backend.storage.services.upload.tus_session import create_tus_lock_factory
 from ai.backend.testutils.bootstrap import redis_container  # noqa: F401
 
 
@@ -36,25 +38,35 @@ async def valkey_tus_client(
     redis_container: tuple[str, HostPortPairModel],  # noqa: F811
 ) -> AsyncIterator[ValkeyTusClient]:
     hostport_pair = redis_container[1]
-    lock_redis = RedisConnectionInfo(
-        Redis.from_url(f"redis://{hostport_pair.address}/{REDIS_STREAM_LOCK}"),
-        sentinel=None,
-        name="test.tus.api.lock",
-        service_name=None,
-        redis_helper_config=config.redis_helper_default_config,
-    )
     client = await ValkeyTusClient.create(
         ValkeyTarget(addr=hostport_pair.address),
         db_id=REDIS_TUS_DB,
         human_readable_name="test.tus.api",
-        lock_redis=lock_redis,
     )
     try:
         yield client
     finally:
         await client.close()
 
 
+@pytest.fixture
+async def tus_lock_factory(
+    redis_container: tuple[str, HostPortPairModel],  # noqa: F811
+) -> AsyncIterator[DistributedLockFactory]:
+    hostport_pair = redis_container[1]
+    lock_redis = RedisConnectionInfo(
+        Redis.from_url(f"redis://{hostport_pair.address}/{REDIS_STREAM_LOCK}"),
+        sentinel=None,
+        name="test.tus.api.lock",
+        service_name=None,
+        redis_helper_config=config.redis_helper_default_config,
+    )
+    try:
+        yield create_tus_lock_factory(lock_redis)
+    finally:
+        await lock_redis.close()
+
+
 @dataclasses.dataclass(slots=True)
 class _PatchEnv:
     vfpath: Path
@@ -70,6 +82,7 @@ def _build_request(
     body: bytes | None,
     offset_header: str | None,
     valkey_client: ValkeyTusClient,
+    lock_factory: DistributedLockFactory,
 ) -> MagicMock:
     volume = MagicMock()
     volume.mangle_vfpath.return_value = vfpath
@@ -79,6 +92,7 @@ def _build_request(
     ctx.get_volume.return_value.__aenter__ = AsyncMock(return_value=volume)
     ctx.get_volume.return_value.__aexit__ = AsyncMock(return_value=None)
     ctx.valkey_tus_client = valkey_client
+    ctx.tus_lock_factory = lock_factory
 
     request = MagicMock(spec=web.Request)
     request.app = {"ctx": ctx}
@@ -144,7 +158,10 @@ def _patch_handler_params(token_data: dict[str, Any]) -> Any:
 
 class TestUploadOffsetHeaderValidation:
     async def test_missing_offset_header_raises(
-        self, patch_env: _PatchEnv, valkey_tus_client: ValkeyTusClient
+        self,
+        patch_env: _PatchEnv,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         request = _build_request(
             vfpath=patch_env.vfpath,
@@ -153,6 +170,7 @@ async def test_missing_offset_header_raises(
             body=None,
             offset_header=None,
             valkey_client=valkey_tus_client,
+            lock_factory=tus_lock_factory,
         )
         token_data = _token_data(session_id=patch_env.session_id, total_size=1024, relpath="f.bin")
         cp = _patch_handler_params(token_data)
@@ -163,7 +181,10 @@ async def test_missing_offset_header_raises(
             cp.stop()
 
     async def test_non_integer_offset_header_raises(
-        self, patch_env: _PatchEnv, valkey_tus_client: ValkeyTusClient
+        self,
+        patch_env: _PatchEnv,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         request = _build_request(
             vfpath=patch_env.vfpath,
@@ -172,6 +193,7 @@ async def test_non_integer_offset_header_raises(
             body=None,
             offset_header="not-a-number",
             valkey_client=valkey_tus_client,
+            lock_factory=tus_lock_factory,
         )
         token_data = _token_data(session_id=patch_env.session_id, total_size=1024, relpath="f.bin")
         cp = _patch_handler_params(token_data)
@@ -182,7 +204,10 @@ async def test_non_integer_offset_header_raises(
             cp.stop()
 
     async def test_negative_offset_raises_conflict(
-        self, patch_env: _PatchEnv, valkey_tus_client: ValkeyTusClient
+        self,
+        patch_env: _PatchEnv,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         request = _build_request(
             vfpath=patch_env.vfpath,
@@ -191,6 +216,7 @@ async def test_negative_offset_raises_conflict(
             body=None,
             offset_header="-1",
             valkey_client=valkey_tus_client,
+            lock_factory=tus_lock_factory,
         )
         token_data = _token_data(session_id=patch_env.session_id, total_size=1024, relpath="f.bin")
         cp = _patch_handler_params(token_data)
@@ -201,7 +227,10 @@ async def test_negative_offset_raises_conflict(
             cp.stop()
 
     async def test_offset_above_total_size_raises_conflict(
-        self, patch_env: _PatchEnv, valkey_tus_client: ValkeyTusClient
+        self,
+        patch_env: _PatchEnv,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         request = _build_request(
             vfpath=patch_env.vfpath,
@@ -210,6 +239,7 @@ async def test_offset_above_total_size_raises_conflict(
             body=None,
             offset_header="2048",
             valkey_client=valkey_tus_client,
+            lock_factory=tus_lock_factory,
         )
         token_data = _token_data(session_id=patch_env.session_id, total_size=1024, relpath="f.bin")
         cp = _patch_handler_params(token_data)
@@ -222,7 +252,10 @@ async def test_offset_above_total_size_raises_conflict(
 
 class TestSessionNotFound:
     async def test_missing_session_dir_raises_not_found(
-        self, tmp_path: Path, valkey_tus_client: ValkeyTusClient
+        self,
+        tmp_path: Path,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         vfpath = tmp_path / "vfpath"
         vfpath.mkdir()
@@ -235,6 +268,7 @@ async def test_missing_session_dir_raises_not_found(
             body=b"",
             offset_header="0",
             valkey_client=valkey_tus_client,
+            lock_factory=tus_lock_factory,
         )
         token_data = _token_data(session_id="missing-session", total_size=1024, relpath="f.bin")
         cp = _patch_handler_params(token_data)
@@ -247,7 +281,10 @@ async def test_missing_session_dir_raises_not_found(
 
 class TestHappyPath:
     async def test_single_chunk_upload_completes_and_assembles(
-        self, patch_env: _PatchEnv, valkey_tus_client: ValkeyTusClient
+        self,
+        patch_env: _PatchEnv,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         payload = b"hello world" * 100
         request = _build_request(
@@ -257,6 +294,7 @@ async def test_single_chunk_upload_completes_and_assembles(
             body=payload,
             offset_header="0",
             valkey_client=valkey_tus_client,
+            lock_factory=tus_lock_factory,
         )
         token_data = _token_data(
             session_id=patch_env.session_id,
@@ -277,7 +315,10 @@ async def test_single_chunk_upload_completes_and_assembles(
         assert list((patch_env.session_dir / "chunks").glob("*.dat")) == []
 
     async def test_two_chunks_assemble_in_order(
-        self, patch_env: _PatchEnv, valkey_tus_client: ValkeyTusClient
+        self,
+        patch_env: _PatchEnv,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         token_data = _token_data(
             session_id=patch_env.session_id, total_size=2048, relpath="result.bin"
@@ -292,6 +333,7 @@ async def test_two_chunks_assemble_in_order(
                 body=b"A" * 1024,
                 offset_header="0",
                 valkey_client=valkey_tus_client,
+                lock_factory=tus_lock_factory,
             )
             await tus_upload_part(first)
 
@@ -302,6 +344,7 @@ async def test_two_chunks_assemble_in_order(
                 body=b"B" * 1024,
                 offset_header="1024",
                 valkey_client=valkey_tus_client,
+                lock_factory=tus_lock_factory,
             )
             response = await tus_upload_part(second)
         finally:
@@ -312,7 +355,10 @@ async def test_two_chunks_assemble_in_order(
         assert final_path.read_bytes() == b"A" * 1024 + b"B" * 1024
 
     async def test_duplicate_chunk_replay_is_idempotent(
-        self, patch_env: _PatchEnv, valkey_tus_client: ValkeyTusClient
+        self,
+        patch_env: _PatchEnv,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         token_data = _token_data(
             session_id=patch_env.session_id, total_size=2048, relpath="result.bin"
@@ -327,6 +373,7 @@ async def test_duplicate_chunk_replay_is_idempotent(
                 body=payload,
                 offset_header="0",
                 valkey_client=valkey_tus_client,
+                lock_factory=tus_lock_factory,
             )
             await tus_upload_part(first)
 
@@ -338,6 +385,7 @@ async def test_duplicate_chunk_replay_is_idempotent(
                 body=payload,
                 offset_header="0",
                 valkey_client=valkey_tus_client,
+                lock_factory=tus_lock_factory,
             )
             response = await tus_upload_part(replay)
         finally:
@@ -346,7 +394,10 @@ async def test_duplicate_chunk_replay_is_idempotent(
         assert response.headers["Upload-Offset"] == "1024"
 
     async def test_chunk_exceeding_declared_size_raises(
-        self, patch_env: _PatchEnv, valkey_tus_client: ValkeyTusClient
+        self,
+        patch_env: _PatchEnv,
+        valkey_tus_client: ValkeyTusClient,
+        tus_lock_factory: DistributedLockFactory,
     ) -> None:
         token_data = _token_data(
             session_id=patch_env.session_id, total_size=10, relpath="result.bin"
@@ -358,6 +409,7 @@ async def test_chunk_exceeding_declared_size_raises(
             body=b"too-much-data",  # 13 bytes > 10
             offset_header="0",
             valkey_client=valkey_tus_client,
+            lock_factory=tus_lock_factory,
         )
         cp = _patch_handler_params(token_data)
         try: