feat(BA-3744): Add multiple download handler and pydantic jwt model #8424

kwonkwonn · 2026-01-29T09:37:08Z

implementation of BA-1040, multiple file zip downloader.

Checklist: (if applicable)

Milestone metadata specifying the target backport version
Mention to the original issue
Installer updates including:
- Fixtures for db schema changes
- New mandatory config options
Update of end-to-end CLI integration tests in ai.backend.test
API server-client counterparts (e.g., manager API -> client SDK)
Test case(s) to:
- Demonstrate the difference of before/after
- Demonstrate the flow of abstract/conceptual models with a concrete implementation
Documentation
- Contents in the docs directory
- docstrings in public interfaces and type annotations

Copilot

Pull request overview

This PR implements a new endpoint for downloading multiple files/directories as a single ZIP archive, addressing issue #7766 (BA-3744). The implementation introduces a streaming ZIP archive handler that leverages zipstream to avoid creating temporary files.

Changes:

Adds new ZipArchiveStreamReader class and ArchiveDownloadTokenData Pydantic model for JWT validation
Implements /folder/file/download-archive manager endpoint to generate download tokens for multiple files
Implements /download-archive client endpoint to stream requested files as a ZIP archive

Reviewed changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 16 comments.

File	Description
src/ai/backend/storage/services/file_stream/zip.py	New module containing ZIP streaming infrastructure with `ZipArchiveStreamReader` and `ArchiveDownloadTokenData` Pydantic model
src/ai/backend/storage/api/manager.py	Adds `create_archive_download_session` endpoint and route for generating JWT tokens for multi-file downloads
src/ai/backend/storage/api/client.py	Adds `download_archive` handler for streaming multiple files as ZIP using token-based authentication

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

src/ai/backend/storage/services/file_stream/zip.py

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

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

Copilot · 2026-01-29T09:46:17Z

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

+    try:
+        payload = jwt.decode(raw_token, secret, algorithms=["HS256"])
+        token_data = ArchiveDownloadTokenData.model_validate(payload)
+    except jwt.ExpiredSignatureError as e:
+        raise web.HTTPUnauthorized(reason="Token expired") from e
+    except jwt.InvalidTokenError as e:
+        raise web.HTTPUnauthorized(reason="Invalid token") from e
+    except ValidationError as e:
+        raise InvalidAPIParameters(extra_msg=str(e)) from e


Manual JWT validation is inconsistent with existing patterns in this file. The existing download and upload handlers use tx.JsonWebToken validator in the check_params context manager for automatic JWT decoding and validation (see lines 145-148, 309-312). This approach should be used here for consistency and better error handling.

tx.JsonWebToken validator in the check_params uses Trafaret based validation.
Our ultimate goal is to migrate Trafaret into pydantic inside the program, Better implement it inline the function and replace with pydantic based validator later.

But at least not in this pr.

src/ai/backend/storage/services/file_stream/zip.py

Copilot · 2026-01-29T09:46:19Z

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

+        sanitized: list[Path] = []
+        for relpath in token_data.files:
+            file_path = volume.sanitize_vfpath(token_data.vfolder_id, relpath)
+            # Ensure the file_path is within the vfolder_root boundary
+            if not file_path.exists():
+                raise web.HTTPNotFound(reason=f"File not found: {relpath}")


The implementation doesn't handle duplicate paths in the input list. If the same file or directory appears multiple times in token_data.files, it will be added to the zip archive multiple times with the same arcname, which may cause issues or confusion. Consider deduplicating the file paths before processing, or at least documenting this behavior.

Suggested change

sanitized: list[Path] = []

for relpath in token_data.files:

file_path = volume.sanitize_vfpath(token_data.vfolder_id, relpath)

# Ensure the file_path is within the vfolder_root boundary

if not file_path.exists():

raise web.HTTPNotFound(reason=f"File not found: {relpath}")

sanitized: list[Path] = []

seen_paths: set[Path] = set()

for relpath in token_data.files:

file_path = volume.sanitize_vfpath(token_data.vfolder_id, relpath)

# Ensure the file_path is within the vfolder_root boundary

if not file_path.exists():

raise web.HTTPNotFound(reason=f"File not found: {relpath}")

if file_path in seen_paths:

continue

seen_paths.add(file_path)

I will add comment on this one, but we need to decide which part should get charged of detecting duplicated files in the list

Copilot · 2026-01-29T09:46:19Z

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

+        raise web.HTTPUnauthorized(reason="Invalid token") from e
+    except ValidationError as e:
+        raise InvalidAPIParameters(extra_msg=str(e)) from e
+    # TODO: We need to implement pydantic jwt decorder to avoid this manual validation. include extract secret key from ctx.


The TODO comment indicates that the current implementation is incomplete and needs to be refactored. Since the codebase already has a tx.JsonWebToken validator that handles JWT decoding with secret extraction from context, this manual approach should be replaced with the existing pattern before merging.

It depends on trafaret based validation

Please write it based on Pydantic this time.

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

src/ai/backend/storage/services/file_stream/zip.py

kwonkwonn · 2026-01-30T08:53:24Z

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

+        vfid: VFolderID
+        files: list[str]
+
+    async with cast(


I was thinking of wrapping api handlers with pydantic, but it seems better to have it in the future PR

kwonkwonn · 2026-01-30T08:54:33Z

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

+    raw_token = request.query.get("token")
+    if not raw_token:
+        raise InvalidAPIParameters(extra_msg="Missing 'token' query parameter")
+    try:


Couldn't find pydantic wrapper to automatize validation and error handling

Will be done in following issue as migrating storage is needed

You can use the api_handler decorator, which was also used in manager.

kwonkwonn · 2026-01-30T08:58:18Z

We can add three tests.

jwt generator: unit test with request->response validation
download_archive: mock streamReader and file structure, how it detect files, how it response (unit test?)
StreamReader: needs component test and running server at least. manual test for now

kwonkwonn · 2026-01-30T08:59:34Z

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

+                raise InvalidAPIParameters(extra_msg=f"Path escapes vfolder boundary: {relpath}")
+            if not file_path.exists():
+                raise web.HTTPNotFound(reason=f"File not found: {relpath}")
+


sanitize_vfpath call mangle function which translate vfid into absolute path, not needed to use entire function.

HyeockJinKim · 2026-02-03T03:57:40Z

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

+    raw_token = request.query.get("token")
+    if not raw_token:
+        raise InvalidAPIParameters(extra_msg="Missing 'token' query parameter")
+    try:


You can use the api_handler decorator, which was also used in manager.

HyeockJinKim · 2026-02-03T03:58:12Z

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

+        raise web.HTTPUnauthorized(reason="Invalid token") from e
+    except ValidationError as e:
+        raise InvalidAPIParameters(extra_msg=str(e)) from e
+    # TODO: We need to implement pydantic jwt decorder to avoid this manual validation. include extract secret key from ctx.


Please write it based on Pydantic this time.

HyeockJinKim · 2026-02-03T03:59:20Z

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

+async def create_archive_download_session(request: web.Request) -> web.Response:
+    class Params(TypedDict):
+        volume: str
+        vfid: VFolderID
+        files: list[str]
+
+    async with cast(
+        AbstractAsyncContextManager[Params],
+        check_params(
+            request,
+            t.Dict(
+                {
+                    t.Key("volume"): t.String(),
+                    t.Key("vfid"): tx.VFolderID(),
+                    t.Key("files"): t.List(t.String, min_length=1),
+                },
+            ),
+        ),
+    ) as params:


There's no reason to keep this implementation. Please define it based on pydantic.

HyeockJinKim · 2026-02-03T03:59:55Z

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

+        token_payload = ArchiveDownloadTokenData(
+            operation="download",
+            volume=params["volume"],
+            vfolder_id=params["vfid"],
+            files=params["files"],
+        )


Defining pydantic and first validating with trafaret is highly unusual.

HyeockJinKim · 2026-02-03T04:00:40Z

src/ai/backend/storage/services/file_stream/zip.py

+class ArchiveDownloadTokenData(BaseModel):
+    """Pydantic model for validating the JWT payload of archive download tokens."""
+
+    operation: Literal["download"]
+    volume: str
+    vfolder_id: VFolderIDField
+    files: list[str] = Field(min_length=1)
+


Since the request must also be made in the manager, the values to be passed as DTOs should be defined in common/dto/storage.
And please define values like operation as an Enum rather than a Literal.

HyeockJinKim · 2026-02-03T04:01:49Z

src/ai/backend/storage/services/file_stream/zip.py

+async def stream_archive_response(
+    request: web.Request,
+    reader: StreamReader,
+    filename: str,
+) -> web.StreamResponse:
+    """Stream a StreamReader's output as an HTTP attachment response.
+
+    This helper is format-agnostic: it works with any StreamReader (ZIP, tar, etc.).
+    """
+    ascii_filename = filename.encode("ascii", errors="ignore").decode("ascii").replace('"', r"\"")
+    encoded_filename = urllib.parse.quote(filename, encoding="utf-8")
+    response = web.StreamResponse(
+        headers={
+            hdrs.CONTENT_TYPE: reader.content_type() or "application/octet-stream",
+            hdrs.CONTENT_DISPOSITION: " ".join([
+                f'attachment;filename="{ascii_filename}";',  # RFC-2616 sec2.2
+                f"filename*=UTF-8''{encoded_filename}",  # RFC-5987
+            ]),
+        },
+    )
+    await response.prepare(request)
+    async for chunk in reader.read():
+        await response.write(chunk)
+    return response


Please ensure API requests are not intertwined here. The service layer should only return a StreamReader, and we want it to use only the values returned from the API layer.

Copilot AI review requested due to automatic review settings January 29, 2026 09:37

github-actions bot assigned kwonkwonn Jan 29, 2026

github-actions bot added size:L 100~500 LoC comp:storage-proxy Related to Storage proxy component labels Jan 29, 2026

Copilot started reviewing on behalf of kwonkwonn January 29, 2026 09:37 View session

Copilot AI reviewed Jan 29, 2026

View reviewed changes

lablup deleted a comment from Copilot AI Jan 29, 2026

feat: Add multiple download handler and pydantic jwt model

0f00acc

kwonkwonn force-pushed the feat/multiple-file-stream-zipper branch from 2660427 to 0f00acc Compare January 30, 2026 08:28

kwonkwonn added 2 commits January 30, 2026 17:33

refactor: Add fragment and fix lint error

cb50d71

refactor: enforce api's to use pydantic types

6a18967

kwonkwonn commented Jan 30, 2026

View reviewed changes

kwonkwonn requested a review from fregataa January 30, 2026 08:55

kwonkwonn commented Jan 30, 2026

View reviewed changes

kwonkwonn requested a review from HyeockJinKim January 30, 2026 09:00

HyeockJinKim requested changes Feb 3, 2026

View reviewed changes

feat(BA-3744): Add multiple download handler and pydantic jwt model #8424

Are you sure you want to change the base?

feat(BA-3744): Add multiple download handler and pydantic jwt model #8424

Conversation

kwonkwonn commented Jan 29, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

kwonkwonn Jan 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

kwonkwonn Jan 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kwonkwonn commented Jan 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kwonkwonn Jan 30, 2026 •

edited

Loading

kwonkwonn Jan 30, 2026 •

edited

Loading