Add cursor pagination for GET api assets (#14014)

Amp-Thread-ID: https://ampcode.com/threads/T-019e4ca5-b71a-7168-8f56-58b2325f34c3
Co-authored-by: Amp <amp@ampcode.com>

app/assets/api/routes.py

@@ -39,6 +39,7 @@ from app.assets.services import (
     update_asset_metadata,
     upload_from_temp_path,
 )
+from app.assets.services.cursor import InvalidCursorError
 from app.assets.services.path_utils import compute_paths_for_response
 from app.assets.services.tagging import list_tag_histogram
 
@@ -182,7 +183,7 @@ def _build_asset_response(result: schemas.AssetDetailResult | schemas.UploadResu
         user_metadata=result.ref.user_metadata or {},
         metadata=result.ref.system_metadata,
         job_id=result.ref.job_id,
-        prompt_id=result.ref.job_id,  # deprecated: mirrors job_id for cloud compat
+        prompt_id=result.ref.job_id,  # deprecated alias of job_id, kept for compatibility
         created_at=result.ref.created_at,
         updated_at=result.ref.updated_at,
         last_access_time=result.ref.last_access_time,
@@ -219,24 +220,37 @@ async def list_assets_route(request: web.Request) -> web.Response:
     order_candidate = (q.order or "desc").lower()
     order = order_candidate if order_candidate in {"asc", "desc"} else "desc"
 
-    result = list_assets_page(
-        owner_id=USER_MANAGER.get_request_user_id(request),
-        include_tags=q.include_tags,
-        exclude_tags=q.exclude_tags,
-        name_contains=q.name_contains,
-        metadata_filter=q.metadata_filter,
-        limit=q.limit,
-        offset=q.offset,
-        sort=sort,
-        order=order,
-    )
+    try:
+        result = list_assets_page(
+            owner_id=USER_MANAGER.get_request_user_id(request),
+            include_tags=q.include_tags,
+            exclude_tags=q.exclude_tags,
+            name_contains=q.name_contains,
+            metadata_filter=q.metadata_filter,
+            limit=q.limit,
+            offset=q.offset,
+            sort=sort,
+            order=order,
+            after=q.after,
+        )
+    except InvalidCursorError as e:
+        return _build_error_response(400, "INVALID_CURSOR", str(e))
 
     summaries = [_build_asset_response(item) for item in result.items]
 
+    # has_more semantics differ by mode:
+    #   - cursor mode: a non-empty next_cursor means there are more results.
+    #   - offset mode: derived from total - (offset + page size).
+    if q.after is not None:
+        has_more = result.next_cursor is not None
+    else:
+        has_more = (q.offset + len(summaries)) < result.total
+
     payload = schemas_out.AssetsList(
         assets=summaries,
         total=result.total,
-        has_more=(q.offset + len(summaries)) < result.total,
+        has_more=has_more,
+        next_cursor=result.next_cursor,
     )
     return web.json_response(payload.model_dump(mode="json", exclude_none=True))
 

app/assets/api/schemas_in.py

@@ -59,6 +59,11 @@ class ListAssetsQuery(BaseModel):
 
     limit: conint(ge=1, le=500) = 20
     offset: conint(ge=0) = 0
+    # Opaque keyset cursor. When supplied, `offset` is ignored. Cursor pagination
+    # is supported for sort values `created_at`, `updated_at`, `name`, `size`.
+    # Supplying `after` together with `sort=last_access_time` returns
+    # 400 INVALID_CURSOR; that sort only supports offset/limit.
+    after: str | None = None
 
     sort: Literal["name", "created_at", "updated_at", "size", "last_access_time"] = (
         "created_at"

app/assets/api/schemas_out.py

@@ -43,6 +43,8 @@ class AssetsList(BaseModel):
     assets: list[Asset]
     total: int
     has_more: bool
+    # Opaque cursor for the next page. Omitted when there are no more results.
+    next_cursor: str | None = None
 
 
 class TagUsage(BaseModel):

app/assets/database/queries/asset_reference.py

@@ -266,9 +266,18 @@ def list_references_page(
     metadata_filter: dict | None = None,
     sort: str | None = None,
     order: str | None = None,
+    after_cursor_value: object | None = None,
+    after_cursor_id: str | None = None,
 ) -> tuple[list[AssetReference], dict[str, list[str]], int]:
     """List references with pagination, filtering, and sorting.
 
+    When ``after_cursor_value``/``after_cursor_id`` are supplied the query uses
+    keyset pagination — ``offset`` is ignored and a WHERE clause selects rows
+    strictly after the given ``(sort_col, id)`` position in the active sort
+    direction. The cursor value must already be typed for the column
+    (datetime for time sorts, int for size, str for name); the caller decodes
+    the opaque cursor string and resolves to the typed value.
+
     Returns (references, tag_map, total_count).
     """
     base = (
@@ -297,9 +306,31 @@ def list_references_page(
         "size": Asset.size_bytes,
     }
     sort_col = sort_map.get(sort, AssetReference.created_at)
-    sort_exp = sort_col.desc() if order == "desc" else sort_col.asc()
+    descending = order == "desc"
 
-    base = base.order_by(sort_exp).limit(limit).offset(offset)
+    # Keyset WHERE: (sort_col, id) strictly less-than / greater-than the cursor.
+    # Equivalent to: sort_col <op> v  OR  (sort_col = v AND id <op> cursor_id).
+    if after_cursor_value is not None and after_cursor_id is not None:
+        if descending:
+            keyset = sa.or_(
+                sort_col < after_cursor_value,
+                sa.and_(sort_col == after_cursor_value, AssetReference.id < after_cursor_id),
+            )
+        else:
+            keyset = sa.or_(
+                sort_col > after_cursor_value,
+                sa.and_(sort_col == after_cursor_value, AssetReference.id > after_cursor_id),
+            )
+        base = base.where(keyset)
+
+    # Secondary ORDER BY id (matching the primary direction) gives the keyset
+    # comparison a deterministic tiebreaker on duplicate sort_col values.
+    id_exp = AssetReference.id.desc() if descending else AssetReference.id.asc()
+    sort_exp = sort_col.desc() if descending else sort_col.asc()
+
+    base = base.order_by(sort_exp, id_exp).limit(limit)
+    if after_cursor_id is None:
+        base = base.offset(offset)
 
     count_stmt = (
         select(sa.func.count())

app/assets/services/asset_management.py

@@ -1,8 +1,19 @@
 import contextlib
 import mimetypes
 import os
+from datetime import timezone
 from typing import Sequence
 
+from app.assets.services.cursor import (
+    CursorPayload,
+    InvalidCursorError,
+    decode_cursor,
+    decode_cursor_int,
+    decode_cursor_time,
+    encode_cursor,
+    encode_cursor_from_time,
+)
+
 
 from app.assets.database.models import Asset
 from app.assets.database.queries import (
@@ -242,6 +253,11 @@ def get_asset_by_hash(asset_hash: str) -> AssetData | None:
         return extract_asset_data(asset)
 
 
+# Sort fields that support cursor pagination. `last_access_time` is not
+# in this list — it falls back to offset/limit.
+_CURSOR_SORT_FIELDS = ("created_at", "updated_at", "name", "size")
+
+
 def list_assets_page(
     owner_id: str = "",
     include_tags: Sequence[str] | None = None,
@@ -252,7 +268,39 @@ def list_assets_page(
     offset: int = 0,
     sort: str = "created_at",
     order: str = "desc",
+    after: str | None = None,
 ) -> ListAssetsResult:
+    """List assets with optional cursor pagination.
+
+    When ``after`` is supplied it overrides ``offset``. The cursor's sort field
+    must match ``sort`` and be in the cursor-supported allowlist; mismatches
+    raise InvalidCursorError so the handler can map to 400 INVALID_CURSOR.
+    """
+    cursor_value: object | None = None
+    cursor_id: str | None = None
+    # Mint next_cursor on every page where the sort is cursor-supported, not
+    # only when the request itself arrived with a cursor. Otherwise a first
+    # request (no `after`) returns next_cursor=None and the client can never
+    # enter cursor mode.
+    mint_cursor = sort in _CURSOR_SORT_FIELDS
+
+    if after is not None:
+        if sort not in _CURSOR_SORT_FIELDS:
+            raise InvalidCursorError(
+                f"cursor pagination is not supported for sort={sort!r}"
+            )
+        payload = decode_cursor(after, _CURSOR_SORT_FIELDS, expected_order=order)
+        if payload.sort_field != sort:
+            raise InvalidCursorError(
+                f"cursor sort field {payload.sort_field!r} does not match request sort {sort!r}"
+            )
+        cursor_value, cursor_id = _resolve_cursor_value(payload), payload.id
+
+    # Over-fetch by one row so we can distinguish "exactly `limit` rows total
+    # remaining" from "more rows past this page" without a second query. Drop
+    # the sentinel before returning.
+    fetch_limit = limit + 1 if mint_cursor else limit
+
     with create_session() as session:
         refs, tag_map, total = list_references_page(
             session,
@@ -261,12 +309,22 @@ def list_assets_page(
             exclude_tags=exclude_tags,
             name_contains=name_contains,
             metadata_filter=metadata_filter,
-            limit=limit,
+            limit=fetch_limit,
             offset=offset,
             sort=sort,
             order=order,
+            after_cursor_value=cursor_value,
+            after_cursor_id=cursor_id,
         )
 
+        next_cursor: str | None = None
+        if mint_cursor and len(refs) > limit:
+            # There's at least one more row past this page — mint a cursor from
+            # the last row of the page (i.e. index `limit - 1`, since we
+            # over-fetched), and drop the sentinel.
+            next_cursor = _encode_next_cursor(refs[limit - 1], sort, order)
+            refs = refs[:limit]
+
         items: list[AssetSummaryData] = []
         for ref in refs:
             items.append(
@@ -277,7 +335,39 @@ def list_assets_page(
                 )
             )
 
-        return ListAssetsResult(items=items, total=total)
+        return ListAssetsResult(items=items, total=total, next_cursor=next_cursor)
+
+
+def _resolve_cursor_value(payload: CursorPayload) -> object:
+    """Map a decoded cursor payload to a column-typed Python value."""
+    if payload.sort_field in ("created_at", "updated_at"):
+        # DB stores naive UTC; strip tzinfo so the comparison binds against a
+        # `TIMESTAMP WITHOUT TIME ZONE` column without an offset shift.
+        return decode_cursor_time(payload).replace(tzinfo=None)
+    if payload.sort_field == "size":
+        return decode_cursor_int(payload)
+    return payload.value  # name, str-typed
+
+
+def _encode_next_cursor(ref, sort: str, order: str) -> str | None:
+    """Mint a cursor pointing at *ref* for the given sort dimension.
+
+    Returns None when the boundary row carries a NULL sort value (e.g. an asset
+    record whose size_bytes hasn't been backfilled). Continuing pagination
+    across a NULL boundary is undefined under keyset ordering — better to
+    truncate cleanly here than to mint a cursor that mis-positions.
+    """
+    if sort == "name":
+        return encode_cursor("name", ref.name, ref.id, order=order)
+    if sort == "size":
+        if ref.asset is None or ref.asset.size_bytes is None:
+            return None
+        return encode_cursor("size", str(ref.asset.size_bytes), ref.id, order=order)
+    # created_at / updated_at — DB datetimes are naive UTC; attach tz before encoding.
+    value = ref.created_at if sort == "created_at" else ref.updated_at
+    if value is None:
+        return None
+    return encode_cursor_from_time(sort, value.replace(tzinfo=timezone.utc), ref.id, order=order)
 
 
 def resolve_hash_to_path(

app/assets/services/cursor.py

@@ -0,0 +1,225 @@
+"""Opaque keyset-pagination cursor for /api/assets.
+
+Payload JSON uses short keys to keep the encoded length small:
+
+    {"s": <sort_field>, "v": <value>, "id": <id>, "o": <order>}
+
+The `o` key binds the cursor to the sort direction it was minted under,
+so replaying a `desc` cursor against an `asc` request fails with
+``INVALID_CURSOR`` rather than silently walking the wrong direction.
+`o` is mandatory on every payload — a cursor without it is rejected as
+malformed.
+
+Encoding is base64url with no padding. JSON serialization escapes `<`,
+`>`, `&`, U+2028, and U+2029 in encoded string values so asset names
+containing those characters produce a stable, byte-identical wire form
+across any compatible implementation of the same payload format.
+
+Time values are serialized as Unix microseconds (UTC) — microsecond
+precision is sufficient to round-trip the timestamps stored by the
+database without rounding rows in the same millisecond bucket.
+"""
+from __future__ import annotations
+
+import base64
+import json
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Iterable, Optional
+
+
+class InvalidCursorError(ValueError):
+    """Raised on a malformed, oversized, or unsupported-sort-field cursor.
+
+    Map to a 400 response with code ``INVALID_CURSOR`` at the handler.
+    """
+
+
+# Wire-format length caps. Cursors are user-controlled, so caps protect the
+# decode path from oversized allocations and downstream SQL predicates from
+# unbounded strings.
+#
+# MAX_CURSOR_VALUE_LENGTH is 512 to fit the `AssetReference.name` column max
+# (`String(512)`) — otherwise a long-named asset would mint a cursor the same
+# server then refuses on the next request.
+MAX_ENCODED_CURSOR_LENGTH = 1024
+MAX_CURSOR_VALUE_LENGTH = 512
+MAX_CURSOR_ID_LENGTH = 128
+
+
+@dataclass(frozen=True)
+class CursorPayload:
+    sort_field: str
+    value: str
+    id: str
+    order: str
+
+
+_VALID_ORDERS = ("asc", "desc")
+
+
+def encode_cursor(sort_field: str, value: str, id: str, order: str = "desc") -> str:
+    """Encode a cursor payload as a base64url (no-padding) string.
+
+    `order` binds the cursor to the sort direction it was minted under so a
+    later request with a flipped `order` query parameter is rejected with
+    ``INVALID_CURSOR`` rather than silently walking the wrong direction.
+    """
+    if order not in _VALID_ORDERS:
+        raise InvalidCursorError(f"order must be one of {_VALID_ORDERS}, got {order!r}")
+    # Symmetric input validation: the encoder must reject anything the
+    # decoder rejects, or the same server will mint cursors it then 400s on
+    # the next request.
+    if not id:
+        raise InvalidCursorError("id must be non-empty")
+    if len(id) > MAX_CURSOR_ID_LENGTH:
+        raise InvalidCursorError("id exceeds maximum length")
+    if len(value) > MAX_CURSOR_VALUE_LENGTH:
+        raise InvalidCursorError("value exceeds maximum length")
+    payload = {"s": sort_field, "v": value, "id": id, "o": order}
+    raw = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
+    # Match the default JSON escaping of HTML-significant characters and JS
+    # line/paragraph separators (U+2028 / U+2029) so an asset name carrying
+    # any of them encodes to identical bytes across runtimes. None of these
+    # characters appear in JSON structural syntax, so a global replace on the
+    # serialized output can only touch encoded values. Use explicit \uXXXX
+    # escapes for U+2028 / U+2029 so the source survives any editor / git
+    # tooling that normalizes invisible separators.
+    raw = (
+        raw.replace("<", "\\u003c")
+           .replace(">", "\\u003e")
+           .replace("&", "\\u0026")
+           .replace("\u2028", "\\u2028")
+           .replace("\u2029", "\\u2029")
+    )
+    encoded = base64.urlsafe_b64encode(raw.encode("utf-8")).rstrip(b"=").decode("ascii")
+    # Final wire-size guard: the per-field caps above are char-counted, but the
+    # wire cap applies to the base64url of the UTF-8-encoded, escape-expanded
+    # payload. A value full of multibyte or HTML-significant characters (e.g.
+    # 512 \u00d7 "\u00e9" or 512 \u00d7 "<") inflates well past MAX_ENCODED_CURSOR_LENGTH even
+    # though it passes the char-count check. Refuse to mint a cursor the decoder
+    # on the next request would reject.
+    if len(encoded) > MAX_ENCODED_CURSOR_LENGTH:
+        raise InvalidCursorError("encoded cursor exceeds maximum length")
+    return encoded
+
+
+def encode_cursor_from_time(sort_field: str, t: datetime, id: str, order: str = "desc") -> str:
+    """Encode a time-typed cursor at Unix microsecond precision.
+
+    Accepts an aware datetime (any timezone) and normalizes to UTC. Naive
+    datetimes are rejected so callers can't accidentally encode the local
+    wall-clock value of a UTC-stored timestamp.
+    """
+    if t.tzinfo is None:
+        raise ValueError("encode_cursor_from_time requires an aware datetime")
+    micros = _datetime_to_unix_micros(t.astimezone(timezone.utc))
+    return encode_cursor(sort_field, str(micros), id, order=order)
+
+
+def decode_cursor(
+    cursor: str,
+    allowed_sort_fields: Iterable[str],
+    expected_order: str | None = None,
+) -> CursorPayload:
+    """Parse an opaque cursor.
+
+    ``allowed_sort_fields`` is the endpoint's accepted sort-field list — a
+    cursor carrying a field outside this set is rejected so a cursor minted
+    for one column can't be replayed against another (e.g. a ``created_at``
+    timestamp string compared against a ``name`` column).
+
+    ``expected_order`` (``"asc"``/``"desc"``), when supplied, must match the
+    payload's ``o`` field. ``o`` is required on every payload; a cursor
+    missing it is rejected as malformed.
+
+    Passing no allowed fields rejects every cursor.
+    """
+    if len(cursor) > MAX_ENCODED_CURSOR_LENGTH:
+        raise InvalidCursorError("cursor exceeds maximum length")
+
+    try:
+        # urlsafe_b64decode requires correct padding; we strip on encode, so
+        # restore the trailing '=' pad here.
+        padding = "=" * (-len(cursor) % 4)
+        raw = base64.urlsafe_b64decode(cursor + padding)
+    except (ValueError, base64.binascii.Error) as e:
+        raise InvalidCursorError(f"encoding: {e}") from e
+
+    try:
+        decoded = json.loads(raw)
+    except (json.JSONDecodeError, UnicodeDecodeError) as e:
+        raise InvalidCursorError(f"payload: {e}") from e
+
+    if not isinstance(decoded, dict):
+        raise InvalidCursorError("payload: expected object")
+
+    sort_field = decoded.get("s")
+    value = decoded.get("v")
+    id = decoded.get("id")
+    order = decoded.get("o")
+
+    if not isinstance(sort_field, str) or not isinstance(value, str) or not isinstance(id, str):
+        raise InvalidCursorError("payload: missing or non-string s/v/id")
+
+    if id == "":
+        raise InvalidCursorError("missing id")
+    if len(id) > MAX_CURSOR_ID_LENGTH:
+        raise InvalidCursorError("id exceeds maximum length")
+    if len(value) > MAX_CURSOR_VALUE_LENGTH:
+        raise InvalidCursorError("value exceeds maximum length")
+
+    if sort_field not in allowed_sort_fields:
+        raise InvalidCursorError(f"unsupported sort field {sort_field!r}")
+
+    if not isinstance(order, str):
+        raise InvalidCursorError("missing or non-string o")
+    if order not in _VALID_ORDERS:
+        raise InvalidCursorError(f"unsupported order {order!r}")
+    if expected_order is not None and order != expected_order:
+        raise InvalidCursorError(
+            f"cursor order {order!r} does not match request order {expected_order!r}"
+        )
+
+    return CursorPayload(sort_field=sort_field, value=value, id=id, order=order)
+
+
+def decode_cursor_time(payload: Optional[CursorPayload]) -> datetime:
+    """Parse a time-typed cursor value as Unix microseconds, returning UTC."""
+    if payload is None:
+        raise InvalidCursorError("nil cursor payload")
+    try:
+        micros = int(payload.value)
+    except ValueError as e:
+        raise InvalidCursorError(f"value is not a valid timestamp: {e}") from e
+    try:
+        return _unix_micros_to_datetime(micros)
+    except (OverflowError, OSError, ValueError) as e:
+        # Crafted out-of-range microseconds (e.g. > datetime.MAX_YEAR) blow up
+        # in fromtimestamp / datetime construction. Map to 400, not 500.
+        raise InvalidCursorError(f"value is out of representable range: {e}") from e
+
+
+def decode_cursor_int(payload: Optional[CursorPayload]) -> int:
+    """Parse a cursor value as a base-10 integer."""
+    if payload is None:
+        raise InvalidCursorError("nil cursor payload")
+    try:
+        return int(payload.value)
+    except ValueError as e:
+        raise InvalidCursorError(f"value is not a valid integer: {e}") from e
+
+
+_EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc)
+
+
+def _datetime_to_unix_micros(t: datetime) -> int:
+    """Convert an aware UTC datetime to Unix microseconds (integer math)."""
+    delta = t - _EPOCH
+    return (delta.days * 86_400 + delta.seconds) * 1_000_000 + delta.microseconds
+
+
+def _unix_micros_to_datetime(micros: int) -> datetime:
+    """Convert Unix microseconds to a UTC datetime, preserving precision."""
+    seconds, micro_remainder = divmod(micros, 1_000_000)
+    return datetime.fromtimestamp(seconds, tz=timezone.utc).replace(microsecond=micro_remainder)

app/assets/services/schemas.py

@@ -71,6 +71,7 @@ class AssetSummaryData:
 class ListAssetsResult:
     items: list[AssetSummaryData]
     total: int
+    next_cursor: str | None = None
 
 
 @dataclass(frozen=True)