diff --git a/README.md b/README.md index 77d56ab3..c8aecbbb 100644 --- a/README.md +++ b/README.md @@ -475,12 +475,13 @@ When `project_id` is omitted on the maintained session-project path, the upload ### What does not work on the new API yet -Three capabilities have no captured payload, so they fail with -`UNSUPPORTED_ON_BATCH_API` rather than quietly producing the wrong thing: +1080p export is ported: FlowKit mirrors the Flow UI's `p0UkFb` high-resolution +Download request and polls the resulting media through `as29s`. The remaining +capabilities below still fail loudly rather than quietly producing the wrong thing: | Capability | Status | Workaround | |---|---|---| -| 4K/1080p upscale (`/fk-pipeline` last step) | unported | none — keep the 1080p render | +| 4K export | plan-gated / not live-verified | use 1080p; Google Flow exposes Full HD as the standard high-resolution export | | Veo reference-to-video (r2v) | unported | Omni r2v (`model_family=omni_flash`), or `FLOW_ALLOW_DEGRADED=1` → i2v off the first reference | | Veo start+end-frame chaining (`/fk-gen-chain-videos`) | unported | Omni first+last (`model_family=omni_flash`), or `FLOW_ALLOW_DEGRADED=1` → i2v off the start frame | | Omni Flash text-to-video | ported | `POST /api/flow/generate-video-omni-text` (4/6/8/10s) | diff --git a/agent/api/flow.py b/agent/api/flow.py index 18f9859b..6b51a214 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -13,6 +13,7 @@ ) from agent.services.flow_client import get_flow_client from agent.services.flow_project_session import current_session_project, ensure_session_project +from agent.services.flow_poll import annotate_polling, check_workflow_status from agent.services.omni_flash import ( check_omni_flash_status, generate_omni_flash_first_frame_video, @@ -96,6 +97,15 @@ class UpscaleVideoRequest(BaseModel): scene_id: str aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT" resolution: str = "VIDEO_RESOLUTION_4K" + project_id: Optional[str] = None + + +class ExportVideoRequest(BaseModel): + media_id: str + scene_id: str = "export" + quality: Literal["1080p", "4k"] = "1080p" + aspect_ratio: str = "VIDEO_ASPECT_RATIO_LANDSCAPE" + project_id: Optional[str] = None class UploadImageRequest(BaseModel): @@ -133,6 +143,7 @@ class CheckStatusRequest(BaseModel): workflows: Optional[list[dict]] = None project_id: str = "" include_encoded_video: bool = False + mode: Literal["omni", "export"] = "omni" class CheckOmniStatusRequest(BaseModel): @@ -390,12 +401,45 @@ async def upscale_video(body: UpscaleVideoRequest): return result.get("data", result) +@router.post("/export-video") +async def export_video(body: ExportVideoRequest): + """Start Google's native Full HD/4K export and return workflow polling data.""" + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + resolution = "VIDEO_RESOLUTION_1080P" if body.quality == "1080p" else "VIDEO_RESOLUTION_4K" + project_id = await _resolve_direct_project(client, body.project_id or "") + result = await client.upscale_video( + media_id=body.media_id, + scene_id=body.scene_id, + aspect_ratio=body.aspect_ratio, + resolution=resolution, + project_id=project_id, + ) + if result.get("error") or ( + isinstance(result.get("status"), int) and result["status"] >= 400 + ): + raise HTTPException(_safe_status_code(result.get("status")), result.get("error", result.get("data"))) + + annotated = annotate_polling(result, mode="media_redirect") + data = annotated.get("data", annotated) + if isinstance(data, dict): + data["export"] = { + "quality": body.quality, + "resolution": resolution, + "native_flow_export": True, + "next": "/api/flow/check-status", + "check_status_mode": "export", + } + return data + + @router.post("/check-status") async def check_status(body: CheckStatusRequest): - """Check Veo operation status or Omni workflow/media status. + """Check Veo operations, Omni workflows, or native export workflows. - Veo: pass ``operations``. - Omni Flash: pass ``workflows`` from submit ``flowkitPolling.workflows``. + Veo: pass ``operations``. Workflow callers pass ``workflows``; native video + export additionally sets ``mode=\"export\"``. """ client = get_flow_client() if not client.connected: @@ -403,6 +447,16 @@ async def check_status(body: CheckStatusRequest): if body.workflows: try: + if body.mode == "export": + result = await check_workflow_status( + body.workflows, + mode="media_redirect", + include_encoded_video=body.include_encoded_video, + project_id=body.project_id, + client=client, + ) + result["download_ready"] = result.get("status") == "COMPLETED" + return result return await check_omni_flash_status( body.workflows, include_encoded_video=body.include_encoded_video, diff --git a/agent/services/flow_batch.py b/agent/services/flow_batch.py index 852478ec..61b5d9a5 100644 --- a/agent/services/flow_batch.py +++ b/agent/services/flow_batch.py @@ -45,6 +45,7 @@ RPC_MEDIA = "as29s" RPC_UPLOAD_IMAGE = "maseQ" RPC_UPSCALE_IMAGE = "SPrCad" +RPC_UPSCALE = "p0UkFb" CAPTCHA_IMAGE = "IMAGE_GENERATION" CAPTCHA_VIDEO = "VIDEO_GENERATION" @@ -572,6 +573,23 @@ def text_video_request(prompt: str, project_id: str, ]) +def upscale_request(media_id: str, project_id: str, + aspect: Any = VIDEO_ASPECT_LANDSCAPE, + model: str = "veo_3_1_upsampler_1080p") -> str: + '''Build Flow's native high-resolution export request (RPC p0UkFb).''' + item = [None] * 32 + item[0] = [None, media_id] + item[2] = 1 + item[4] = [None, str(uuid.uuid4()), None, None, _client_uuid()] + item[6] = resolve_video_aspect(aspect) + item[31] = model + return build_envelope(RPC_UPSCALE, [ + [item], + _context(project_id), + [_client_uuid()], + ]) + + def upload_request(image_b64: str, project_id: str, mime_type: str = "image/jpeg", file_name: str = "upload.jpg") -> str: """Put a local image into the project so it can be used as a reference. @@ -687,6 +705,14 @@ def read_upscaled_image(payload: Any) -> str: return encoded +def read_upscaled_media_id(payload: Any) -> str: + """Return the media id created by the p0UkFb upscale submit.""" + for text in _walk_strings(payload): + if isinstance(text, str) and text.endswith("_upsampled"): + return text + raise FlowBatchError("upscale response carried no upsampled media id") + + def read_operation(payload: Any) -> Operation: """`[null, 50, [[opId, projectId, sceneId, status, …]]]`. diff --git a/agent/services/flow_client.py b/agent/services/flow_client.py index 874a2f57..52a362f2 100644 --- a/agent/services/flow_client.py +++ b/agent/services/flow_client.py @@ -23,6 +23,7 @@ import time import uuid from typing import Optional +from urllib.parse import quote from agent.config import ( VIDEO_MODELS, @@ -951,12 +952,30 @@ async def generate_video_from_references(self, reference_media_ids: list[str], async def upscale_video(self, media_id: str, scene_id: str, aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT", - resolution: str = "VIDEO_RESOLUTION_4K") -> dict: - """Upscale a video.""" - return {"error": _unsupported( - "video upscale", - "no upsampler rpc appears in the new frontend's captures", - )} + resolution: str = "VIDEO_RESOLUTION_4K", + project_id: str | None = None) -> dict: + '''Upscale/export a video using Flow's migrated p0UkFb RPC.''' + model = _config.UPSCALE_MODELS.get(resolution) + if not model: + return {"status": 400, "error": f"Unsupported upscale resolution: {resolution}"} + try: + pid = self._batch_project_id(project_id or "") + freq = fb.upscale_request(media_id, pid, aspect=aspect_ratio, model=model) + payload = await self._batch_payload( + fb.RPC_UPSCALE, freq, fb.CAPTCHA_VIDEO, timeout=120) + upscaled_media_id = fb.read_upscaled_media_id(payload) + except Exception as e: + return _batch_error(e) + + workflow = { + "name": upscaled_media_id, + "primary_media_id": upscaled_media_id, + "project_id": pid, + } + return {"status": 200, "data": { + "media": [{"name": upscaled_media_id}], + "workflows": [workflow], + }} async def check_video_status(self, operations: list[dict]) -> dict: """One poll round for each submitted operation. @@ -1134,6 +1153,21 @@ async def get_media(self, media_id: str) -> dict: data["image"] = {"fifeUrl": urls.image} return {"status": 200, "data": data} + async def resolve_media_url(self, media_id: str) -> dict: + """Resolve a signed Flow video URL through the authenticated batch media API.""" + result = await self.get_media(media_id) + data = result.get("data") if isinstance(result.get("data"), dict) else {} + video = data.get("video") if isinstance(data, dict) else None + candidate = video.get("fifeUrl") if isinstance(video, dict) else None + return { + "status": result.get("status", 200), + "data": { + "url": candidate, + "contentType": "video/mp4" if candidate else None, + }, + "error": result.get("error"), + } + async def upload_image(self, image_base64: str, mime_type: str = "image/jpeg", project_id: str = "", file_name: str = "image.jpg") -> dict: """Upload an image into the project so it can be used as a reference.""" diff --git a/agent/services/flow_poll.py b/agent/services/flow_poll.py new file mode 100644 index 00000000..fb00faf2 --- /dev/null +++ b/agent/services/flow_poll.py @@ -0,0 +1,175 @@ +"""Shared polling helpers for Flow workflow-backed video jobs. + +Workflow-producing surfaces (Omni and native video export) return the same +logical pair: a workflow name plus a primary media id. This module owns the +normalization, response annotation and one-pass polling so each feature does +not grow its own copy of transport/status handling. +""" +from __future__ import annotations + +from agent.services.flow_client import get_flow_client + +_ALLOWED_MEDIA_URL_PREFIX = "https://flow-content.google/video/" +_POLL_MODES = {"batch_media", "media_redirect"} + + +def normalize_workflow(workflow: dict) -> dict | None: + """Normalize a raw Flow workflow or a FlowKit polling descriptor.""" + if not isinstance(workflow, dict): + return None + name = workflow.get("name") + primary_media_id = workflow.get("primary_media_id") + if not primary_media_id: + metadata = workflow.get("metadata") + if isinstance(metadata, dict): + primary_media_id = metadata.get("primaryMediaId") + if not isinstance(name, str) or not name: + return None + if not isinstance(primary_media_id, str) or not primary_media_id: + return None + item = {"name": name, "primary_media_id": primary_media_id} + project_id = workflow.get("project_id") or workflow.get("projectId") + if isinstance(project_id, str) and project_id: + item["project_id"] = project_id + return item + + +def extract_workflows(result: dict) -> list[dict]: + """Extract normalized workflow descriptors from a submit response.""" + if not isinstance(result, dict): + return [] + data = result.get("data") if isinstance(result.get("data"), dict) else result + workflows = data.get("workflows", []) if isinstance(data, dict) else [] + return [item for workflow in workflows if (item := normalize_workflow(workflow))] + + +def annotate_polling(result: dict, *, mode: str, project_id: str = "") -> dict: + """Attach the descriptor callers can pass back to ``/flow/check-status``.""" + if mode not in _POLL_MODES: + raise ValueError(f"Unknown Flow polling mode: {mode}") + workflows = extract_workflows(result) + if not workflows: + return result + if project_id: + for workflow in workflows: + workflow["project_id"] = project_id + data = result.get("data") if isinstance(result.get("data"), dict) else result + if isinstance(data, dict): + descriptor = {"mode": mode, "workflows": workflows} + if project_id: + descriptor["project_id"] = project_id + data["flowkitPolling"] = descriptor + return result + + +def _normalize_workflows(workflows: list[dict]) -> list[dict]: + return [item for workflow in (workflows or []) if (item := normalize_workflow(workflow))] + + +def _parse_media_redirect(response: dict) -> tuple[str | None, str | None, str | None]: + if not isinstance(response, dict): + return None, None, "Flow media redirect returned an invalid response" + status = response.get("status") + data = response.get("data") if isinstance(response.get("data"), dict) else {} + candidate = data.get("url") + content_type = data.get("contentType") + if ( + isinstance(status, int) + and status < 400 + and isinstance(candidate, str) + and candidate.startswith(_ALLOWED_MEDIA_URL_PREFIX) + ): + return candidate, content_type if isinstance(content_type, str) else None, None + error = response.get("error") + if not error and isinstance(status, int) and status >= 400: + error = f"API_{status}" + return None, content_type if isinstance(content_type, str) else None, str(error or "media redirect not ready") + + +async def _check_media_ready( + workflows: list[dict], + *, + client, + include_encoded_video: bool, + project_id: str, +) -> dict: + resolved_project_id = project_id or next( + (item.get("project_id", "") for item in workflows if item.get("project_id")), "" + ) + items = [] + for workflow in workflows: + media_id = workflow["primary_media_id"] + response = await client.resolve_media_url(media_id) + url, content_type, diagnostic = _parse_media_redirect(response) + item_project = workflow.get("project_id") or resolved_project_id + if url: + media = { + "media_id": media_id, + "url": url, + "encoded_video_available": False, + "resolved_via": "as29s", + } + if content_type: + media["content_type"] = content_type + if include_encoded_video: + media["encoded_video"] = None + item = { + "name": workflow["name"], + "primary_media_id": media_id, + "done": True, + "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", + "error": None, + "media": media, + } + else: + item = { + "name": workflow["name"], + "primary_media_id": media_id, + "done": False, + "status": "PENDING", + "error": None, + } + if diagnostic: + item["probe"] = {"diagnostic": diagnostic} + if isinstance(response, dict) and isinstance(response.get("status"), int): + item["probe"]["http_status"] = response["status"] + if item_project: + item["project_id"] = item_project + items.append(item) + + all_done = bool(items) and all(item["done"] for item in items) + result = { + "done": all_done, + "status": "COMPLETED" if all_done else "PENDING", + "workflows": items, + } + if resolved_project_id: + result["project_id"] = resolved_project_id + return result + + +async def check_workflow_status( + workflows: list[dict], + *, + mode: str, + include_encoded_video: bool = False, + project_id: str = "", + client=None, +) -> dict: + """Perform one non-blocking poll pass for a workflow-backed Flow job.""" + if mode not in _POLL_MODES: + raise ValueError(f"Unknown Flow polling mode: {mode}") + normalized = _normalize_workflows(workflows) + if not normalized: + raise ValueError( + "Flow polling requires workflow descriptors with name and primary_media_id " + "(or raw Flow metadata.primaryMediaId)" + ) + + client = client or get_flow_client() + return await _check_media_ready( + normalized, + client=client, + include_encoded_video=include_encoded_video, + project_id=project_id, + ) diff --git a/agent/services/omni_flash.py b/agent/services/omni_flash.py index 277a8736..a8dd28ae 100644 --- a/agent/services/omni_flash.py +++ b/agent/services/omni_flash.py @@ -24,6 +24,7 @@ from agent.services import flow_batch as fb from agent.services.flow_client import get_flow_client +from agent.services.flow_poll import check_workflow_status, extract_workflows _MODELS_FILE = Path(__file__).parent.parent / "models.json" @@ -134,39 +135,9 @@ def _validate_frame_inputs( raise ValueError("Omni Flash First+Last requires a non-empty end_image_media_id") -def _normalize_workflow(workflow: dict) -> dict | None: - """Normalize a raw Flow workflow or FlowKit polling descriptor.""" - if not isinstance(workflow, dict): - return None - name = workflow.get("name") - primary_media_id = workflow.get("primary_media_id") - if not primary_media_id: - metadata = workflow.get("metadata") - if isinstance(metadata, dict): - primary_media_id = metadata.get("primaryMediaId") - if not isinstance(name, str) or not name: - return None - if not isinstance(primary_media_id, str) or not primary_media_id: - return None - item = {"name": name, "primary_media_id": primary_media_id} - project_id = workflow.get("project_id") or workflow.get("projectId") - if isinstance(project_id, str) and project_id: - item["project_id"] = project_id - return item - - def extract_omni_workflows(result: dict) -> list[dict]: - """Extract ``name`` + ``primaryMediaId`` pairs from an Omni submit.""" - if not isinstance(result, dict): - return [] - data = result.get("data") if isinstance(result.get("data"), dict) else result - workflows = data.get("workflows", []) if isinstance(data, dict) else [] - normalized = [] - for workflow in workflows: - item = _normalize_workflow(workflow) - if item: - normalized.append(item) - return normalized + """Backward-compatible Omni name for the shared workflow extractor.""" + return extract_workflows(result) async def generate_omni_flash_text_video( @@ -363,64 +334,16 @@ async def generate_omni_flash_video( batch_model = f"abra_r2v_{duration_s}s" + ("_360p" if resolution == "360p" else "") return _batch_operation_result(operation, pid, batch_model, duration_s, resolution) -async def _check_omni_batch_media( - workflows: list[dict], - include_encoded_video: bool = False, - project_id: str = "", -) -> dict: - normalized = [item for workflow in (workflows or []) if (item := _normalize_workflow(workflow))] - if not normalized: - raise ValueError("Omni polling requires workflow descriptors with name and primary_media_id") - resolved_project_id = project_id or next( - (item.get("project_id", "") for item in normalized if item.get("project_id")), "") - client = get_flow_client() - items = [] - for workflow in normalized: - media_id = workflow["primary_media_id"] - response = await client.get_media(media_id) - data = response.get("data") if isinstance(response, dict) else None - video = data.get("video") if isinstance(data, dict) else None - url = video.get("fifeUrl") if isinstance(video, dict) else None - if isinstance(url, str) and url.startswith("https://flow-content.google/video/"): - media = { - "media_id": media_id, - "url": url, - "encoded_video_available": False, - "resolved_via": "as29s", - } - if include_encoded_video: - media["encoded_video"] = None - items.append({ - "name": workflow["name"], - "primary_media_id": media_id, - "project_id": workflow.get("project_id") or resolved_project_id, - "done": True, - "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", - "error": None, - "media": media, - }) - else: - items.append({ - "name": workflow["name"], - "primary_media_id": media_id, - "project_id": workflow.get("project_id") or resolved_project_id, - "done": False, - "status": "PENDING", - "error": None, - }) - all_done = bool(items) and all(item["done"] for item in items) - return { - "project_id": resolved_project_id or None, - "done": all_done, - "status": "COMPLETED" if all_done else "PENDING", - "workflows": items, - } - - async def check_omni_flash_status( workflows: list[dict], include_encoded_video: bool = False, project_id: str = "", ) -> dict: - """Perform one non-blocking poll pass for Omni workflow-backed jobs.""" - return await _check_omni_batch_media(workflows, include_encoded_video, project_id) + """Poll Omni text-to-video results using the same batch media resolver as exports.""" + return await check_workflow_status( + workflows, + mode="batch_media", + include_encoded_video=include_encoded_video, + project_id=project_id, + client=get_flow_client(), + ) diff --git a/docs/VIDEO_EXPORTS.md b/docs/VIDEO_EXPORTS.md new file mode 100644 index 00000000..19b9f586 --- /dev/null +++ b/docs/VIDEO_EXPORTS.md @@ -0,0 +1,45 @@ +# Video exports + +FlowKit treats output quality as an explicit export/download choice rather than +an implementation detail. + +## Omni Flash + +Gemini Omni Flash generation normally produces a 720p source video. Google Flow +provides a native Full HD export for that result. FlowKit exposes it directly as +**Export 1080p**; internally Google names the operation an upsample, but callers +do not need to reason about that implementation detail. + +## Export Full HD (recommended) + +```bash +curl -sS -X POST http://127.0.0.1:8100/api/flow/export-video \ + -H 'Content-Type: application/json' \ + -d '{ + "media_id": "", + "scene_id": "job-1", + "quality": "1080p", + "aspect_ratio": "VIDEO_ASPECT_RATIO_LANDSCAPE" + }' +``` + +`quality` defaults to `1080p`. `4k` remains an explicit API option, but it is plan-gated by Google Flow; the currently verified account exposes 4K as disabled while 1080p is available. + +If `project_id` is omitted, FlowKit uses the current ad-hoc Flow session project. +Pass `project_id` explicitly when exporting media from an existing Flow project. + +The response contains `flowkitPolling.workflows`. Poll those descriptors: + +```bash +curl -sS -X POST http://127.0.0.1:8100/api/flow/check-status \ + -H 'Content-Type: application/json' \ + -d '{"mode":"export","workflows": }' +``` + +When `download_ready` becomes `true`, use +`workflows[].media.url` immediately. It is a short-lived signed +`flow-content.google` URL. + +## Compatibility + +`POST /api/flow/upscale-video` remains the low-level submit surface. New integrations should use `export-video`, then poll the returned workflows through the existing `POST /api/flow/check-status` endpoint with `mode: "export"`. This keeps all Flow polling behind one API surface. diff --git a/tests/unit/test_flow_batch.py b/tests/unit/test_flow_batch.py index fb48f08e..2a39c053 100644 --- a/tests/unit/test_flow_batch.py +++ b/tests/unit/test_flow_batch.py @@ -207,6 +207,21 @@ def test_text_video_matches_the_captured_yhhmef_shape(self): assert payload[2][1] == 2 assert fb.CAPTCHA_SLOT in json.dumps(payload) + def test_upscale_matches_the_captured_1080p_slots(self): + payload = inner(fb.upscale_request( + "media-1", self.PID, + aspect="VIDEO_ASPECT_RATIO_LANDSCAPE", + model="veo_3_1_upsampler_1080p", + )) + item = payload[0][0] + assert len(item) == 32 + assert item[0] == [None, "media-1"] + assert item[2] == 1 + assert item[6] == fb.VIDEO_ASPECT_LANDSCAPE + assert item[31] == "veo_3_1_upsampler_1080p" + assert payload[1][5] == self.PID + assert fb.CAPTCHA_SLOT in json.dumps(payload) + def test_text_video_360p_matches_current_ui_slots(self): payload = inner(fb.text_video_request( "a boat", @@ -261,6 +276,8 @@ def test_image_upscale_reads_synchronous_encoded_image(self): assert fb.read_upscaled_image([["media"], "A" * 200]) == "A" * 200 with pytest.raises(fb.FlowBatchError): fb.read_upscaled_image([["media"], "short"]) + def test_upscale_submit_reads_the_new_media_id(self): + assert fb.read_upscaled_media_id([[self.MID + "_upsampled"]]) == self.MID + "_upsampled" def test_operation_reads_the_id_and_status(self): op = fb.read_operation([None, 50, [[self.OP, "proj", "scene", "CAE"]]]) diff --git a/tests/unit/test_flow_client_batch.py b/tests/unit/test_flow_client_batch.py index 688cd744..299f8b17 100644 --- a/tests/unit/test_flow_client_batch.py +++ b/tests/unit/test_flow_client_batch.py @@ -293,11 +293,26 @@ async def test_degraded_r2v_uses_the_first_reference_as_the_start_frame(self, cl payload = json.loads(json.loads(client.calls[0]["freq"])[0][0][1]) assert payload[0][0][4][1] == "ref-a" - async def test_upscale_is_unported_and_has_no_fallback(self, client, monkeypatch): - import agent.services.flow_client as module - monkeypatch.setattr(module, "FLOW_ALLOW_DEGRADED", True) - result = await client.upscale_video(MEDIA, "scene-1") - assert "UNSUPPORTED_ON_BATCH_API" in result["error"] + + +class TestUpscaleVideo: + async def test_1080p_submit_returns_a_pollable_workflow(self, client): + upscaled = MEDIA + "_upsampled" + client.responses[fb.RPC_UPSCALE] = { + "data": envelope(fb.RPC_UPSCALE, [[[[upscaled], "", None, None, 1]]]) + } + + result = await client.upscale_video( + MEDIA, + "scene-1", + aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE", + resolution="VIDEO_RESOLUTION_1080P", + ) + workflow = result["data"]["workflows"][0] + assert workflow["primary_media_id"] == upscaled + assert workflow["project_id"] == PROJECT + assert client.calls[0]["rpcid"] == fb.RPC_UPSCALE + assert client.calls[0]["captcha"] == fb.CAPTCHA_VIDEO class TestCheckVideoStatus: diff --git a/tests/unit/test_flow_export_api.py b/tests/unit/test_flow_export_api.py new file mode 100644 index 00000000..1fd80360 --- /dev/null +++ b/tests/unit/test_flow_export_api.py @@ -0,0 +1,82 @@ +"""Regression tests for the public Flow native video export API.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from agent.api import flow as api + + +@pytest.mark.asyncio +async def test_export_uses_session_project_and_returns_polling_descriptor(): + client = MagicMock(connected=True) + client.upscale_video = AsyncMock(return_value={ + "status": 200, + "data": {"workflows": [{"name": "upscale-job", "primary_media_id": "video_upscaled"}]}, + }) + with ( + patch.object(api, "get_flow_client", return_value=client), + patch.object(api, "ensure_session_project", new_callable=AsyncMock) as session, + ): + session.return_value = {"project_id": "project-123"} + result = await api.export_video(api.ExportVideoRequest(media_id="video")) + + client.upscale_video.assert_awaited_once_with( + media_id="video", scene_id="export", + aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE", + resolution="VIDEO_RESOLUTION_1080P", + project_id="project-123", + ) + assert result["flowkitPolling"]["mode"] == "media_redirect" + assert result["flowkitPolling"]["workflows"][0]["primary_media_id"] == "video_upscaled" + assert result["export"]["check_status_mode"] == "export" + + +@pytest.mark.asyncio +async def test_export_with_explicit_project_does_not_allocate_session(): + client = MagicMock(connected=True) + client.upscale_video = AsyncMock(return_value={"status": 200, "data": {"workflows": []}}) + with ( + patch.object(api, "get_flow_client", return_value=client), + patch.object(api, "ensure_session_project", new_callable=AsyncMock) as session, + ): + await api.export_video(api.ExportVideoRequest(media_id="video", project_id="existing-project")) + session.assert_not_awaited() + assert client.upscale_video.await_args.kwargs["project_id"] == "existing-project" + + +@pytest.mark.asyncio +async def test_export_polls_through_shared_batch_media_resolver(): + client = MagicMock(connected=True) + with ( + patch.object(api, "get_flow_client", return_value=client), + patch.object(api, "check_workflow_status", new_callable=AsyncMock) as polling, + ): + polling.return_value = {"status": "COMPLETED", "done": True, "workflows": []} + response = await api.check_status(api.CheckStatusRequest( + mode="export", + project_id="existing-project", + workflows=[{"name": "job", "primary_media_id": "video_upscaled"}], + )) + assert response["download_ready"] is True + assert polling.await_args.kwargs["mode"] == "media_redirect" + assert polling.await_args.kwargs["project_id"] == "existing-project" + + +@pytest.mark.asyncio +async def test_export_surfaces_provider_error(): + client = MagicMock(connected=True) + client.upscale_video = AsyncMock(return_value={"status": 403, "error": "plan required"}) + with patch.object(api, "get_flow_client", return_value=client): + with pytest.raises(HTTPException) as exc: + await api.export_video(api.ExportVideoRequest(media_id="video", project_id="project-123")) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_export_fails_cleanly_without_browser_connection(): + with patch.object(api, "get_flow_client", return_value=MagicMock(connected=False)): + with pytest.raises(HTTPException) as exc: + await api.export_video(api.ExportVideoRequest(media_id="video")) + assert exc.value.status_code == 503 diff --git a/tests/unit/test_flow_poll.py b/tests/unit/test_flow_poll.py new file mode 100644 index 00000000..bee92e7d --- /dev/null +++ b/tests/unit/test_flow_poll.py @@ -0,0 +1,88 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import agent.services.flow_poll as polling +from agent.services.flow_client import FlowClient +from agent.services.flow_poll import annotate_polling, check_workflow_status + +MEDIA = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + + +def test_annotate_export_polling_uses_workflow_primary_media_id(): + result = { + "status": 200, + "data": { + "workflows": [ + {"name": "wf-1", "metadata": {"primaryMediaId": MEDIA + "_upsampled"}} + ] + }, + } + annotated = annotate_polling(result, mode="media_redirect") + assert annotated["data"]["flowkitPolling"] == { + "mode": "media_redirect", + "workflows": [ + {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} + ], + } + + +@pytest.mark.asyncio +async def test_batch_poll_resolves_completed_media_through_public_resolver(): + client = MagicMock() + client.resolve_media_url = AsyncMock(return_value={ + "status": 200, + "data": { + "url": "https://flow-content.google/video/out?Signature=test", + "contentType": "video/mp4", + }, + }) + with patch("agent.services.flow_poll.get_flow_client", return_value=client): + result = await check_workflow_status( + [{"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"}], + mode="media_redirect", + ) + + assert result["done"] is True + assert result["status"] == "COMPLETED" + assert result["workflows"][0]["media"]["resolved_via"] == "as29s" + client.resolve_media_url.assert_awaited_once_with(MEDIA + "_upsampled") + + +@pytest.mark.asyncio +async def test_batch_poll_stays_pending_until_video_url_exists(): + client = MagicMock() + client.resolve_media_url = AsyncMock(return_value={ + "status": 404, + "error": "not ready", + "data": {}, + }) + with patch("agent.services.flow_poll.get_flow_client", return_value=client): + result = await check_workflow_status( + [{"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"}], + mode="media_redirect", + ) + + assert result["done"] is False + assert result["status"] == "PENDING" + assert result["workflows"][0]["status"] == "PENDING" + + +def test_unknown_poll_mode_is_rejected(): + with pytest.raises(ValueError, match="Unknown Flow polling mode"): + annotate_polling({"data": {"workflows": []}}, mode="mystery") + + +@pytest.mark.asyncio +async def test_flow_client_media_resolver_uses_get_media_on_batch(): + client = FlowClient() + client.get_media = AsyncMock(return_value={ + "status": 200, + "data": {"video": {"fifeUrl": "https://flow-content.google/video/media-1?sig=x"}}, + }) + + result = await client.resolve_media_url("media-1") + + assert result["data"]["url"].startswith("https://flow-content.google/video/") + assert result["data"]["contentType"] == "video/mp4" + client.get_media.assert_awaited_once_with("media-1") diff --git a/tests/unit/test_omni_flash.py b/tests/unit/test_omni_flash.py index 7bcb137e..09ceb75c 100644 --- a/tests/unit/test_omni_flash.py +++ b/tests/unit/test_omni_flash.py @@ -190,65 +190,59 @@ async def test_first_last_rejects_missing_end_before_submit(): @pytest.mark.asyncio -async def test_batch_omni_poll_uses_as29s_media(monkeypatch): +async def test_batch_omni_poll_uses_as29s_media(): client = MagicMock() - client.get_media = AsyncMock(return_value={ + client.resolve_media_url = AsyncMock(return_value={ "status": 200, "data": { - "video": { - "fifeUrl": "https://flow-content.google/video/media-1?Signature=test" - } + "url": "https://flow-content.google/video/media-1?Signature=test", + "contentType": "video/mp4", }, }) - with patch("agent.services.omni_flash.get_flow_client", return_value=client): - result = await check_omni_flash_status([ - { - "name": "workflow-1", - "primary_media_id": "media-1", - "project_id": "project-1", - } - ]) + result = await check_omni_flash_status([{ + "name": "workflow-1", "primary_media_id": "media-1", "project_id": "project-1", + }]) assert result["done"] is True assert result["status"] == "COMPLETED" assert result["workflows"][0]["media"]["resolved_via"] == "as29s" - client.get_media.assert_awaited_once_with("media-1") + client.resolve_media_url.assert_awaited_once_with("media-1") @pytest.mark.asyncio -async def test_batch_poll_hands_back_the_signed_url_and_buffers_nothing(): - """The poller must return Flow's signed url, never the bytes: buffering a - finished clip through the agent is what the as29s lookup exists to avoid.""" +async def test_batch_poll_hands_back_signed_url_without_buffering(): client = MagicMock() - client.get_media = AsyncMock(return_value={"status": 200, "data": {"video": { - "fifeUrl": "https://flow-content.google/video/media-1?Signature=test"}}}) - + client.resolve_media_url = AsyncMock(return_value={ + "status": 200, "data": { + "url": "https://flow-content.google/video/media-1?Signature=test", + "contentType": "video/mp4", + }, + }) with patch("agent.services.omni_flash.get_flow_client", return_value=client): result = await check_omni_flash_status( [{"name": "workflow-1", "primary_media_id": "media-1"}], - include_encoded_video=True, project_id="project-1") - + include_encoded_video=True, project_id="project-1", + ) media = result["workflows"][0]["media"] assert result["done"] is True assert media["url"].startswith("https://flow-content.google/video/") assert media["encoded_video_available"] is False assert media["encoded_video"] is None + client.resolve_media_url.assert_awaited_once_with("media-1") @pytest.mark.asyncio async def test_batch_poll_treats_a_media_record_without_video_as_pending(): - """A media id exists before its clip does — the record serves the poster - image first. Reading that as done is what saves a still instead of a video.""" client = MagicMock() - client.get_media = AsyncMock(return_value={"status": 200, "data": {"image": { - "fifeUrl": "https://flow-content.google/image/media-1?Signature=test"}}}) - + client.resolve_media_url = AsyncMock(return_value={ + "status": 200, "data": {"url": None, "contentType": None}, + }) with patch("agent.services.omni_flash.get_flow_client", return_value=client): result = await check_omni_flash_status( [{"name": "workflow-1", "primary_media_id": "media-1"}], - project_id="project-1") - + project_id="project-1", + ) assert result["done"] is False assert result["workflows"][0]["status"] == "PENDING"