feat(BA-6159): add /upload/status endpoint and progress headers

Standard TUS HEAD only returns Upload-Offset, which is ambiguous when
chunks arrive out of order: the client cannot tell which byte ranges are
still missing. This sub-task surfaces the session metadata so resume is
precise.

  - GET /upload/status?token=… returns JSON: session, total_size,
    committed_offset, chunks_received (list of offsets),
    missing_ranges (list of {offset, length} gaps),
    progress_percent, status. Authenticated via the same upload JWT.
  - PATCH responses gain three additive headers:
    X-Backend-Ai-Chunks-Received, X-Backend-Ai-Progress-Percent,
    X-Backend-Ai-Total-Expected. Headers are exposed via CORS
    Access-Control-Expose-Headers so browser clients can read them.

Tests cover the status endpoint for empty / partial / out-of-order /
missing-session cases plus the new PATCH response headers.

Resolves BA-6159. Part of epic BA-6153 (implements BA-3974).

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

Author: jopemachine

changes/11770.feature.md

@@ -0,0 +1 @@
+Add `GET /upload/status` endpoint and `X-Backend-Ai-*` progress response headers so clients can recover after a storage proxy crash by resending only the byte ranges still missing.

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

@@ -50,7 +50,10 @@
 from ai.backend.storage.services.file_stream.zip import (
     ZipArchiveStreamReader,
 )
-from ai.backend.storage.services.upload.tus_session import TusUploadSession
+from ai.backend.storage.services.upload.tus_session import (
+    SessionState,
+    TusUploadSession,
+)
 from ai.backend.storage.types import SENTINEL, TusChunkUploadStreamReader
 from ai.backend.storage.utils import (
     CheckParamSource,
@@ -459,21 +462,90 @@ class Params(TypedDict):
                 upload_offset=state.committed_offset,
                 upload_length=total_size,
             )
+            headers.update(_progress_response_headers(state))
     return web.Response(status=HTTPStatus.NO_CONTENT, headers=headers)
 
 
+async def tus_upload_status(request: web.Request) -> web.Response:
+    """
+    GET /upload/status — JSON snapshot of an in-flight upload session.
+
+    Used by clients to recover after a Storage Proxy crash or network error:
+    they discover which byte ranges are still missing and resend only those.
+    """
+    ctx: RootContext = request.app["ctx"]
+    secret = ctx.local_config.storage_proxy.secret
+
+    class Params(TypedDict):
+        token: UploadTokenData
+        dst_dir: str
+
+    async with cast(
+        AbstractAsyncContextManager[Params],
+        check_params(
+            request,
+            t.Dict(
+                {
+                    t.Key("token"): tx.JsonWebToken(
+                        secret=secret,
+                        inner_iv=upload_token_data_iv,
+                    ),
+                    t.Key("dst_dir", default=None): t.Null | t.String,
+                },
+            ),
+            read_from=CheckParamSource.QUERY,
+        ),
+    ) as params:
+        token_data = params["token"]
+        total_size = int(token_data["size"])
+        async with ctx.get_volume(token_data["volume"]) as volume:
+            session_dir = _resolve_session_dir(volume, token_data)
+            if not session_dir.exists():
+                raise web.HTTPNotFound(
+                    body=dump_json_str(
+                        {
+                            "title": "No such upload session",
+                            "type": ("https://api.backend.ai/probs/storage/no-such-upload-session"),
+                        },
+                    ),
+                    content_type="application/problem+json",
+                )
+            session = TusUploadSession(
+                session_dir,
+                session_id=token_data["session"],
+                total_size=total_size,
+            )
+            state = await session.read_state()
+            body = {
+                "session": token_data["session"],
+                "total_size": total_size,
+                "committed_offset": state.committed_offset,
+                "chunks_received": [rec.offset for rec in state.received],
+                "missing_ranges": [
+                    {"offset": offset, "length": length}
+                    for offset, length in state.missing_ranges()
+                ],
+                "progress_percent": state.progress_percent(),
+                "status": state.status,
+            }
+    return web.json_response(body, status=HTTPStatus.OK)
+
+
 def _resolve_session_dir(volume: AbstractVolume, token_data: UploadTokenData) -> Path:
     return volume.mangle_vfpath(token_data["vfid"]) / ".upload" / token_data["session"]
 
 
 _TUS_HEADER_LIST = "Tus-Resumable, Upload-Length, Upload-Metadata, Upload-Offset, Content-Type"
+_PROGRESS_HEADER_LIST = (
+    "X-Backend-Ai-Chunks-Received, X-Backend-Ai-Progress-Percent, X-Backend-Ai-Total-Expected"
+)
 
 
 def _tus_response_headers(*, upload_offset: int, upload_length: int) -> dict[str, str]:
     return {
         "Access-Control-Allow-Origin": "*",
         "Access-Control-Allow-Headers": _TUS_HEADER_LIST,
-        "Access-Control-Expose-Headers": _TUS_HEADER_LIST,
+        "Access-Control-Expose-Headers": f"{_TUS_HEADER_LIST}, {_PROGRESS_HEADER_LIST}",
         "Access-Control-Allow-Methods": "*",
         "Cache-Control": "no-store",
         "Tus-Resumable": "1.0.0",
@@ -482,6 +554,14 @@ def _tus_response_headers(*, upload_offset: int, upload_length: int) -> dict[str
     }
 
 
+def _progress_response_headers(state: SessionState) -> dict[str, str]:
+    return {
+        "X-Backend-Ai-Chunks-Received": str(len(state.received)),
+        "X-Backend-Ai-Progress-Percent": str(state.progress_percent()),
+        "X-Backend-Ai-Total-Expected": str(state.total_size),
+    }
+
+
 async def tus_options(request: web.Request) -> web.Response:
     """
     Let clients discover the supported features of our tus.io server-side implementation.
@@ -580,4 +660,7 @@ async def init_client_app(ctx: RootContext) -> web.Application:
     r.add_route("HEAD", tus_check_session)
     r.add_route("PATCH", tus_upload_part)
 
+    status_resource = app.router.add_resource("/upload/status")
+    status_resource.add_route("GET", tus_upload_status)
+
     return app

tests/unit/storage/api/test_tus_upload.py

@@ -10,14 +10,15 @@
 from __future__ import annotations
 
 import dataclasses
+import json
 from pathlib import Path
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from aiohttp import web
 
-from ai.backend.storage.api.client import tus_upload_part
+from ai.backend.storage.api.client import tus_upload_part, tus_upload_status
 from ai.backend.storage.errors import InvalidAPIParameters, UploadOffsetMismatchError
 
 
@@ -298,3 +299,138 @@ async def test_chunk_exceeding_declared_size_raises(self, patch_env: _PatchEnv)
         chunks_dir = patch_env.session_dir / "chunks"
         if chunks_dir.exists():
             assert list(chunks_dir.glob("*.tmp")) == []
+
+
+def _body_bytes(response: web.Response) -> bytes:
+    body = response.body
+    assert isinstance(body, (bytes, bytearray))
+    return bytes(body)
+
+
+class TestProgressHeaders:
+    async def test_patch_response_includes_progress_headers(self, patch_env: _PatchEnv) -> None:
+        token_data = _token_data(session_id="test-session", total_size=2048, relpath="result.bin")
+        cp = _patch_handler_params(token_data)
+        try:
+            request = _build_request(
+                vfpath=patch_env.vfpath,
+                session_id="test-session",
+                total_size=2048,
+                body=b"A" * 1024,
+                offset_header="0",
+            )
+            response = await tus_upload_part(request)
+        finally:
+            cp.stop()
+
+        assert response.headers["X-Backend-Ai-Chunks-Received"] == "1"
+        assert response.headers["X-Backend-Ai-Total-Expected"] == "2048"
+        assert response.headers["X-Backend-Ai-Progress-Percent"] == "50.0"
+
+
+class TestUploadStatus:
+    async def test_status_for_empty_session(self, patch_env: _PatchEnv) -> None:
+        token_data = _token_data(session_id="test-session", total_size=2048, relpath="result.bin")
+        request = _build_request(
+            vfpath=patch_env.vfpath,
+            session_id="test-session",
+            total_size=2048,
+            body=None,
+            offset_header=None,
+        )
+        cp = _patch_handler_params(token_data)
+        try:
+            response = await tus_upload_status(request)
+        finally:
+            cp.stop()
+
+        payload = json.loads(_body_bytes(response))
+        assert payload["total_size"] == 2048
+        assert payload["committed_offset"] == 0
+        assert payload["chunks_received"] == []
+        assert payload["missing_ranges"] == [{"offset": 0, "length": 2048}]
+        assert payload["progress_percent"] == 0.0
+        assert payload["status"] == "pending"
+
+    async def test_status_after_partial_upload(self, patch_env: _PatchEnv) -> None:
+        token_data = _token_data(session_id="test-session", total_size=4096, relpath="result.bin")
+        cp = _patch_handler_params(token_data)
+        try:
+            patch_request = _build_request(
+                vfpath=patch_env.vfpath,
+                session_id="test-session",
+                total_size=4096,
+                body=b"A" * 1024,
+                offset_header="0",
+            )
+            await tus_upload_part(patch_request)
+
+            status_request = _build_request(
+                vfpath=patch_env.vfpath,
+                session_id="test-session",
+                total_size=4096,
+                body=None,
+                offset_header=None,
+            )
+            response = await tus_upload_status(status_request)
+        finally:
+            cp.stop()
+
+        payload = json.loads(_body_bytes(response))
+        assert payload["committed_offset"] == 1024
+        assert payload["chunks_received"] == [0]
+        assert payload["missing_ranges"] == [{"offset": 1024, "length": 3072}]
+        assert payload["progress_percent"] == 25.0
+        assert payload["status"] == "pending"
+
+    async def test_status_with_out_of_order_chunks_reports_gaps(self, patch_env: _PatchEnv) -> None:
+        token_data = _token_data(session_id="test-session", total_size=4096, relpath="result.bin")
+        cp = _patch_handler_params(token_data)
+        try:
+            # Commit the *second* chunk first, leaving [0, 1024) missing.
+            second = _build_request(
+                vfpath=patch_env.vfpath,
+                session_id="test-session",
+                total_size=4096,
+                body=b"B" * 1024,
+                offset_header="1024",
+            )
+            await tus_upload_part(second)
+
+            status_request = _build_request(
+                vfpath=patch_env.vfpath,
+                session_id="test-session",
+                total_size=4096,
+                body=None,
+                offset_header=None,
+            )
+            response = await tus_upload_status(status_request)
+        finally:
+            cp.stop()
+
+        payload = json.loads(_body_bytes(response))
+        # Contiguous prefix has not advanced because [0,1024) is missing.
+        assert payload["committed_offset"] == 0
+        assert payload["chunks_received"] == [1024]
+        assert payload["missing_ranges"] == [
+            {"offset": 0, "length": 1024},
+            {"offset": 2048, "length": 2048},
+        ]
+
+    async def test_status_for_missing_session_returns_404(self, tmp_path: Path) -> None:
+        vfpath = tmp_path / "vfpath"
+        vfpath.mkdir()
+        token_data = _token_data(session_id="missing", total_size=1024, relpath="result.bin")
+        request = _build_request(
+            vfpath=vfpath,
+            session_id="missing",
+            total_size=1024,
+            body=None,
+            offset_header=None,
+        )
+        cp = _patch_handler_params(token_data)
+        try:
+            with pytest.raises(web.HTTPNotFound):
+                await tus_upload_status(request)
+        finally:
+            cp.stop()