[feat] OpenAI data URLs: decode in bounded chunks, test cancelled decoding
Review follow-ups for the inline video/audio reference decoder: - decode_data_url now undoes percent escapes and strips whitespace in 64 KiB chunks and stops as soon as the payload must decode to more than server.media_reference_max_bytes. An oversized request costs memory in proportion to the limit, not to its own size. Escapes are never split across chunks. - Data URL checks lower-case only the scheme instead of copying the whole URL. - Tests: the wrapped/percent-escaped payload decodes the same at chunk sizes that split escapes and line breaks; an oversized payload is refused after a few chunks; and a request cancelled before its decoding thread writes the file still leaves no upload once the thread finishes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
1fa29f201e
commit
2b13dfd01c
@@ -26,6 +26,7 @@ import httpx
|
||||
MAX_REDIRECTS = 5
|
||||
READ_TIMEOUT_S = 60.0
|
||||
DATA_URL_HEADER_MAX_CHARS = 256
|
||||
DATA_URL_CHUNK_CHARS = 1 << 16 # at least 3, so a percent escape fits in one chunk
|
||||
|
||||
|
||||
class MediaPolicyError(ValueError):
|
||||
@@ -94,14 +95,37 @@ def data_url_media_type(source: str) -> str:
|
||||
return parameters[0].strip().lower()
|
||||
|
||||
|
||||
def _base64_payload(source: str, start: int, max_bytes: int) -> str:
|
||||
"""``source[start:]`` without percent escapes or whitespace, normalized in bounded chunks.
|
||||
|
||||
Stops as soon as the payload must decode to more than ``max_bytes``, so an oversized request costs memory
|
||||
in proportion to the limit, not to its own size.
|
||||
"""
|
||||
chunks: list[str] = []
|
||||
length = 0
|
||||
while start < len(source):
|
||||
end = min(start + DATA_URL_CHUNK_CHARS, len(source))
|
||||
escape = source.find("%", max(start + 1, end - 2), end) if end < len(source) else -1
|
||||
if escape >= 0:
|
||||
end = escape # an escape that may continue past the chunk starts the next one
|
||||
# Same leniency as SGLang Diffusion: percent escapes and line-wrapped payloads decode.
|
||||
chunk = source[start:end]
|
||||
if "%" in chunk:
|
||||
chunk = unquote(chunk)
|
||||
chunk = "".join(chunk.split())
|
||||
length += len(chunk)
|
||||
# Base64 decodes to 3 bytes per 4 characters, less at most 2 bytes of padding.
|
||||
if length // 4 * 3 - 2 > max_bytes:
|
||||
raise _too_large(max_bytes)
|
||||
chunks.append(chunk)
|
||||
start = end
|
||||
return "".join(chunks)
|
||||
|
||||
|
||||
def decode_data_url(source: str, target: str, *, max_bytes: int) -> None:
|
||||
"""Decode a base64 ``data:`` URL to ``target``; nothing is written if it is invalid or over ``max_bytes``."""
|
||||
data_url_media_type(source)
|
||||
payload = source[source.find(",") + 1:]
|
||||
# Same leniency as SGLang Diffusion: percent escapes and line-wrapped payloads decode.
|
||||
if "%" in payload:
|
||||
payload = unquote(payload)
|
||||
payload = "".join(payload.split())
|
||||
payload = _base64_payload(source, source.find(",") + 1, max_bytes)
|
||||
padding = len(payload) - len(payload.rstrip("="))
|
||||
if len(payload) // 4 * 3 - padding > max_bytes:
|
||||
raise _too_large(max_bytes)
|
||||
|
||||
@@ -152,7 +152,8 @@ async def _materialize_media(source: str,
|
||||
source = source.strip()
|
||||
if not source:
|
||||
raise RequestAdaptationError(f"{kind.capitalize()} references must not be empty.")
|
||||
if source.lower().startswith(("http://", "https://")):
|
||||
# Only the scheme is lower-cased: copying a whole data: URL would double its memory.
|
||||
if source[:8].lower().startswith(("http://", "https://")):
|
||||
suffix = os.path.splitext(urlparse(source).path)[1].lower()
|
||||
if suffix not in _MEDIA_SUFFIXES[kind]:
|
||||
suffix = _DEFAULT_SUFFIXES[kind]
|
||||
@@ -164,7 +165,7 @@ async def _materialize_media(source: str,
|
||||
except Exception as error:
|
||||
raise RequestAdaptationError(f"Unable to fetch {kind} reference: {error}") from error
|
||||
return target, True
|
||||
if source.lower().startswith("data:"):
|
||||
if source[:5].lower() == "data:":
|
||||
return await _decode_data_media(source, target_base, kind, max_bytes=max_bytes), True
|
||||
return _local_media(source, policy), False
|
||||
|
||||
@@ -213,7 +214,7 @@ async def prepare_reference_media(
|
||||
local_path = track(await save_base64_image_to_path(source, target))
|
||||
except Exception as error:
|
||||
raise RequestAdaptationError(f"Unable to decode image reference: {error}") from error
|
||||
elif source.lower().startswith("data:"):
|
||||
elif source[:5].lower() == "data:":
|
||||
raise RequestAdaptationError("Image references need a data:image/... URL.")
|
||||
elif is_remote(source):
|
||||
local_path = track(await _materialize_image_url(source, target, max_bytes=max_bytes, policy=policy))
|
||||
|
||||
@@ -17,6 +17,7 @@ from fastapi.testclient import TestClient
|
||||
|
||||
from fastvideo.api.parser import parse_config
|
||||
from fastvideo.api.schema import GenerationRequest
|
||||
from fastvideo.entrypoints.openai import media_policy, request_adapter
|
||||
from fastvideo.entrypoints.openai.media_policy import MediaPolicy
|
||||
from fastvideo.entrypoints.openai.protocol import UrlAudioReference, VideoGenerationRequest
|
||||
from fastvideo.entrypoints.openai.request_adapter import (
|
||||
@@ -293,11 +294,17 @@ def test_inline_video_and_audio_references_are_decoded_in_order(tmp_path: Path)
|
||||
assert not _uploads(tmp_path)
|
||||
|
||||
|
||||
def test_inline_reference_payload_may_be_wrapped_and_percent_escaped(tmp_path: Path) -> None:
|
||||
track = b"wrapped-track" * 8
|
||||
encoded = base64.b64encode(track).decode().replace("=", "%3D")
|
||||
wrapped = "\n".join(encoded[start:start + 16] for start in range(0, len(encoded), 16))
|
||||
request = VideoGenerationRequest(prompt="waves", audio_reference={"audio_url": f"data:audio/flac;base64,{wrapped}"})
|
||||
@pytest.mark.parametrize("chunk_chars", [media_policy.DATA_URL_CHUNK_CHARS, 3, 4, 5, 7])
|
||||
def test_inline_reference_payload_may_be_wrapped_and_percent_escaped(tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
|
||||
chunk_chars: int) -> None:
|
||||
# Small chunks put chunk boundaries inside percent escapes and line breaks.
|
||||
monkeypatch.setattr(media_policy, "DATA_URL_CHUNK_CHARS", chunk_chars)
|
||||
track = bytes(range(256)) * 2
|
||||
encoded = base64.b64encode(track).decode()
|
||||
lines = [encoded[start:start + 16] for start in range(0, len(encoded), 16)]
|
||||
wrapped = "".join(line + ("%0A" if number % 2 else "\r\n") for number, line in enumerate(lines))
|
||||
escaped = wrapped.replace("+", "%2B").replace("/", "%2f").replace("=", "%3D")
|
||||
request = VideoGenerationRequest(prompt="waves", audio_reference={"audio_url": f"data:audio/flac;base64,{escaped}"})
|
||||
created = asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
|
||||
assert Path(created[0]).read_bytes() == track
|
||||
assert created[0].endswith(".flac")
|
||||
@@ -333,6 +340,24 @@ def test_inline_reference_over_the_media_limit_is_refused(tmp_path: Path) -> Non
|
||||
assert Path(created[0]).read_bytes() == b"1234"
|
||||
|
||||
|
||||
def test_oversized_inline_payload_is_refused_before_it_is_all_read(tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(media_policy, "DATA_URL_CHUNK_CHARS", 8)
|
||||
unescaped: list[int] = []
|
||||
original_unquote = media_policy.unquote
|
||||
|
||||
def counting_unquote(chunk: str) -> str:
|
||||
unescaped.append(len(chunk))
|
||||
return original_unquote(chunk)
|
||||
|
||||
monkeypatch.setattr(media_policy, "unquote", counting_unquote)
|
||||
source = "data:video/mp4;base64," + "QU%4AB\n" * 10_000
|
||||
with pytest.raises(media_policy.MediaPolicyError, match="4-byte limit"):
|
||||
media_policy.decode_data_url(source, str(tmp_path / "clip.mp4"), max_bytes=4)
|
||||
assert max(unescaped) <= 8 and sum(unescaped) < 64
|
||||
assert not (tmp_path / "clip.mp4").exists()
|
||||
|
||||
|
||||
def test_decoded_references_are_removed_when_a_later_reference_fails(tmp_path: Path) -> None:
|
||||
request = VideoGenerationRequest(
|
||||
prompt="waves",
|
||||
@@ -344,6 +369,41 @@ def test_decoded_references_are_removed_when_a_later_reference_fails(tmp_path: P
|
||||
assert not _uploads(tmp_path)
|
||||
|
||||
|
||||
def test_cancelled_decoding_removes_its_file_once_the_thread_finishes(tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
asyncio.run(_assert_cancelled_decoding_removes_its_file(tmp_path, monkeypatch))
|
||||
|
||||
|
||||
async def _assert_cancelled_decoding_removes_its_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
started, release, finished = threading.Event(), threading.Event(), threading.Event()
|
||||
written: list[bool] = []
|
||||
|
||||
def decode_after_release(source: str, target: str, *, max_bytes: int) -> None:
|
||||
started.set()
|
||||
release.wait(timeout=5)
|
||||
try:
|
||||
media_policy.decode_data_url(source, target, max_bytes=max_bytes)
|
||||
written.append(os.path.exists(target))
|
||||
finally:
|
||||
finished.set()
|
||||
|
||||
monkeypatch.setattr(request_adapter, "decode_data_url", decode_after_release)
|
||||
request = VideoGenerationRequest(prompt="waves", video_reference={"video_url": _data_url("video/mp4", b"clip")})
|
||||
preparing = asyncio.create_task(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
|
||||
assert await asyncio.to_thread(started.wait, 5)
|
||||
preparing.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await preparing
|
||||
# The thread writes its file only after the request is gone, so cleanup has to wait for the thread.
|
||||
release.set()
|
||||
assert await asyncio.to_thread(finished.wait, 5)
|
||||
assert written == [True]
|
||||
deadline = time.monotonic() + 5
|
||||
while _uploads(tmp_path) and time.monotonic() < deadline:
|
||||
await asyncio.sleep(0.01)
|
||||
assert not _uploads(tmp_path)
|
||||
|
||||
|
||||
def test_local_audio_paths_need_a_configured_root(tmp_path: Path) -> None:
|
||||
track = tmp_path / "track.wav"
|
||||
track.write_bytes(b"x")
|
||||
|
||||
Reference in New Issue
Block a user