diff --git a/README.md b/README.md index 8084b1de..65cced40 100644 --- a/README.md +++ b/README.md @@ -254,12 +254,13 @@ without being silently replaced by the default model. See ### 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 | | Reference-to-video (r2v) | unported | `FLOW_ALLOW_DEGRADED=1` → i2v off the first reference | | Start+end-frame chaining (`/fk-gen-chain-videos`) | unported | `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 5aeed8fa..7fff86a7 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -5,6 +5,7 @@ from agent.config import USE_BATCH_RPC, FLOW_PROJECT_ID, FLOW_ALLOW_DEGRADED from agent.services.flow_client import get_flow_client +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, @@ -78,6 +79,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): @@ -94,6 +104,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): @@ -300,12 +311,44 @@ 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" + result = await client.upscale_video( + media_id=body.media_id, + scene_id=body.scene_id, + aspect_ratio=body.aspect_ratio, + resolution=resolution, + project_id=body.project_id, + ) + if result.get("error") or ( + isinstance(result.get("status"), int) and result["status"] >= 400 + ): + raise HTTPException(result.get("status", 502), 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: @@ -313,6 +356,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 f81c1e44..8d89cf8a 100644 --- a/agent/services/flow_batch.py +++ b/agent/services/flow_batch.py @@ -42,6 +42,7 @@ RPC_MEDIA = "as29s" RPC_UPLOAD_IMAGE = "maseQ" RPC_UPSCALE_IMAGE = "SPrCad" +RPC_UPSCALE = "p0UkFb" CAPTCHA_IMAGE = "IMAGE_GENERATION" CAPTCHA_VIDEO = "VIDEO_GENERATION" @@ -429,6 +430,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 migrated high-resolution download 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. @@ -528,6 +546,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 be56e708..f3104b1a 100644 --- a/agent/services/flow_client.py +++ b/agent/services/flow_client.py @@ -20,6 +20,7 @@ import time import uuid from typing import Optional +from urllib.parse import quote from agent.config import ( GOOGLE_FLOW_API, GOOGLE_API_KEY, ENDPOINTS, @@ -839,14 +840,33 @@ 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.""" + resolution: str = "VIDEO_RESOLUTION_4K", + project_id: str | None = None) -> dict: + """Upscale/export a video using Flow's migrated p0UkFb RPC.""" if not USE_BATCH_RPC: return await self._legacy_upscale_video(media_id, scene_id, aspect_ratio, resolution) - return {"error": _unsupported( - "video upscale", - "no upsampler rpc appears in the new frontend's captures", - )} + + model = 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. @@ -1016,6 +1036,54 @@ 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 media URL without downloading the media body.""" + if USE_BATCH_RPC: + 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"), + } + + url = ( + "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" + f"?name={quote(media_id, safe='')}" + ) + return await self._send( + "trpc_request", + { + "url": url, + "method": "GET", + "headers": {"content-type": "application/json"}, + "responseMode": "url", + }, + timeout=15, + ) + + async def get_project_initial_data(self, project_id: str) -> dict: + """Fetch the authenticated legacy Flow project snapshot used for polling.""" + query = quote( + json.dumps({"json": {"projectId": project_id}}, separators=(",", ":")), + safe="", + ) + url = f"https://labs.google/fx/api/trpc/flow.projectInitialData?input={query}" + return await self._send( + "trpc_request", + { + "url": url, + "method": "GET", + "headers": {"content-type": "application/json"}, + }, + timeout=15, + ) + 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..3d3e9571 --- /dev/null +++ b/agent/services/flow_poll.py @@ -0,0 +1,314 @@ +"""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.config import USE_BATCH_RPC +from agent.services.flow_client import get_flow_client + +_ALLOWED_MEDIA_URL_PREFIX = "https://flow-content.google/" +_POLL_MODES = {"project_media", "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 USE_BATCH_RPC else "media.getMediaUrlRedirect", + } + 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_project_media( + 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")), "" + ) + if not resolved_project_id: + raise ValueError( + "Flow project polling requires project_id. Use the project_id returned " + "inside flowkitPolling or pass project_id explicitly." + ) + if any( + item.get("project_id") and item["project_id"] != resolved_project_id + for item in workflows + ): + raise ValueError("All workflows in one project poll must belong to the same project_id") + + response = await client.get_project_initial_data(resolved_project_id) + http_status = response.get("status") if isinstance(response, dict) else None + if isinstance(http_status, int) and http_status >= 400: + data = response.get("data") if isinstance(response.get("data"), dict) else {} + error = data.get("error") if isinstance(data, dict) else None + if isinstance(error, dict): + error = error.get("message") or error.get("code") + raise RuntimeError(error or response.get("error") or f"Flow project poll failed: API_{http_status}") + + envelope = response.get("data") if isinstance(response, dict) else None + result = envelope.get("result") if isinstance(envelope, dict) else None + result_data = result.get("data") if isinstance(result, dict) else None + project_json = result_data.get("json") if isinstance(result_data, dict) else None + contents = project_json.get("projectContents") if isinstance(project_json, dict) else None + if not isinstance(contents, dict): + raise RuntimeError("Flow project poll returned an unexpected response shape") + + project_workflows = contents.get("workflows") + project_media = contents.get("media") + project_workflows = project_workflows if isinstance(project_workflows, list) else [] + project_media = project_media if isinstance(project_media, list) else [] + known_workflow_names = { + item.get("name") for item in project_workflows + if isinstance(item, dict) and isinstance(item.get("name"), str) + } + media_by_id = { + item.get("name"): item for item in project_media + if isinstance(item, dict) and isinstance(item.get("name"), str) + } + media_by_workflow = { + item.get("workflowId"): item for item in project_media + if isinstance(item, dict) and isinstance(item.get("workflowId"), str) + } + + items = [] + for workflow in workflows: + name = workflow["name"] + media_id = workflow["primary_media_id"] + payload = media_by_id.get(media_id) or media_by_workflow.get(name) + if not isinstance(payload, dict): + items.append({ + "name": name, + "primary_media_id": media_id, + "project_id": resolved_project_id, + "done": False, + "status": "PENDING", + "error": None, + "workflow_present": name in known_workflow_names, + }) + continue + + metadata = payload.get("mediaMetadata") + metadata = metadata if isinstance(metadata, dict) else {} + media_status = metadata.get("mediaStatus") + media_status = media_status if isinstance(media_status, dict) else {} + generation_status = media_status.get("mediaGenerationStatus") + if isinstance(generation_status, str) and ( + generation_status.endswith("FAILED") or generation_status.endswith("CANCELLED") + ): + items.append({ + "name": name, + "primary_media_id": media_id, + "project_id": resolved_project_id, + "done": True, + "status": "FAILED", + "error": generation_status, + }) + continue + if generation_status != "MEDIA_GENERATION_STATUS_SUCCESSFUL": + items.append({ + "name": name, + "primary_media_id": media_id, + "project_id": resolved_project_id, + "done": False, + "status": "PENDING", + "error": None, + }) + continue + + url_response = await client.resolve_media_url(media_id) + url, _, url_error = _parse_media_redirect(url_response) + media = { + "media_id": media_id, + "url": url, + "encoded_video_available": False, + } + if include_encoded_video: + media["encoded_video"] = None + if url_error: + media["url_error"] = url_error + items.append({ + "name": name, + "primary_media_id": media_id, + "project_id": resolved_project_id, + "done": True, + "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", + "error": None, + "media": media, + }) + + all_done = bool(items) and all(item["done"] for item in items) + any_failed = any(item.get("status") == "FAILED" for item in items) + return { + "project_id": resolved_project_id, + "done": all_done, + "status": "FAILED" if any_failed else ("COMPLETED" if all_done else "PENDING"), + "workflows": items, + } + + +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() + if mode in {"batch_media", "media_redirect"}: + return await _check_media_ready( + normalized, + client=client, + include_encoded_video=include_encoded_video, + project_id=project_id, + ) + return await _check_project_media( + 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 796d6a31..8c53f91b 100644 --- a/agent/services/omni_flash.py +++ b/agent/services/omni_flash.py @@ -24,11 +24,10 @@ import time import uuid from pathlib import Path -from urllib.parse import quote - from agent.config import USE_BATCH_RPC from agent.services import flow_batch as fb from agent.services.flow_client import get_flow_client +from agent.services.flow_poll import annotate_polling, check_workflow_status, extract_workflows from agent.services.headers import random_headers _MODELS_FILE = Path(__file__).parent.parent / "models.json" @@ -62,39 +61,13 @@ def _batch_path_blocks_omni() -> dict | None: async def _fetch_project_initial_data(client, project_id: str) -> dict: - """Fetch the same authenticated project snapshot used by the Flow UI.""" - query = quote( - json.dumps({"json": {"projectId": project_id}}, separators=(",", ":")), - safe="", - ) - url = f"https://labs.google/fx/api/trpc/flow.projectInitialData?input={query}" - return await client._send( - "trpc_request", - { - "url": url, - "method": "GET", - "headers": {"content-type": "application/json"}, - }, - timeout=15, - ) + """Compatibility wrapper around FlowClient's public project poll surface.""" + return await client.get_project_initial_data(project_id) async def _fetch_media_url(client, media_id: str) -> dict: - """Resolve Flow's authenticated media redirect without buffering the file.""" - url = ( - "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" - f"?name={quote(media_id, safe='')}" - ) - return await client._send( - "trpc_request", - { - "url": url, - "method": "GET", - "headers": {"content-type": "application/json"}, - "responseMode": "url", - }, - timeout=15, - ) + """Compatibility wrapper around FlowClient's public URL resolver.""" + return await client.resolve_media_url(media_id) def _validate_duration(duration_s: int) -> None: @@ -166,56 +139,14 @@ 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) def _annotate_polling(result: dict, project_id: str) -> dict: - """Add an explicit FlowKit polling descriptor to a successful submit.""" - workflows = extract_omni_workflows(result) - if not workflows: - return result - data = result.get("data") if isinstance(result.get("data"), dict) else result - if isinstance(data, dict): - for workflow in workflows: - workflow["project_id"] = project_id - data["flowkitPolling"] = { - "mode": "project_media", - "project_id": project_id, - "workflows": workflows, - } - return result + """Annotate legacy Omni submits using the shared polling descriptor.""" + return annotate_polling(result, mode="project_media", project_id=project_id) async def generate_omni_flash_text_video( @@ -447,226 +378,16 @@ async def generate_omni_flash_video( return _annotate_polling(result, project_id) -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.""" - if USE_BATCH_RPC: - return await _check_omni_batch_media(workflows, include_encoded_video, project_id) - normalized = [] - for workflow in workflows or []: - item = _normalize_workflow(workflow) - if item: - normalized.append(item) - if not normalized: - raise ValueError( - "Omni polling requires workflow descriptors with name and primary_media_id " - "(or raw Flow metadata.primaryMediaId)" - ) - - resolved_project_id = project_id or next( - (item.get("project_id", "") for item in normalized if item.get("project_id")), - "", + """Perform one non-blocking poll pass using the shared Flow poller.""" + return await check_workflow_status( + workflows, + mode="batch_media" if USE_BATCH_RPC else "project_media", + include_encoded_video=include_encoded_video, + project_id=project_id, + client=get_flow_client(), ) - if not resolved_project_id: - raise ValueError( - "Omni project polling requires project_id. Use the project_id returned " - "inside flowkitPolling or pass project_id explicitly." - ) - if any( - item.get("project_id") and item["project_id"] != resolved_project_id - for item in normalized - ): - raise ValueError( - "All Omni workflows in one poll must belong to the same project_id" - ) - - client = get_flow_client() - response = await _fetch_project_initial_data(client, resolved_project_id) - http_status = response.get("status") if isinstance(response, dict) else None - if isinstance(http_status, int) and http_status >= 400: - data = response.get("data") if isinstance(response.get("data"), dict) else {} - error = data.get("error") if isinstance(data, dict) else None - if isinstance(error, dict): - error = error.get("message") or error.get("code") - raise RuntimeError( - error - or response.get("error") - or f"Flow project poll failed: API_{http_status}" - ) - - envelope = response.get("data") if isinstance(response, dict) else None - result = envelope.get("result") if isinstance(envelope, dict) else None - result_data = result.get("data") if isinstance(result, dict) else None - project_json = result_data.get("json") if isinstance(result_data, dict) else None - contents = project_json.get("projectContents") if isinstance(project_json, dict) else None - if not isinstance(contents, dict): - raise RuntimeError("Flow project poll returned an unexpected response shape") - - project_workflows = contents.get("workflows") - project_media = contents.get("media") - project_workflows = project_workflows if isinstance(project_workflows, list) else [] - project_media = project_media if isinstance(project_media, list) else [] - known_workflow_names = { - item.get("name") - for item in project_workflows - if isinstance(item, dict) and isinstance(item.get("name"), str) - } - media_by_id = { - item.get("name"): item - for item in project_media - if isinstance(item, dict) and isinstance(item.get("name"), str) - } - media_by_workflow = { - item.get("workflowId"): item - for item in project_media - if isinstance(item, dict) and isinstance(item.get("workflowId"), str) - } - - items = [] - - for workflow in normalized: - name = workflow["name"] - media_id = workflow["primary_media_id"] - payload = media_by_id.get(media_id) or media_by_workflow.get(name) - if not isinstance(payload, dict): - items.append({ - "name": name, - "primary_media_id": media_id, - "project_id": resolved_project_id, - "done": False, - "status": "PENDING", - "error": None, - "workflow_present": name in known_workflow_names, - }) - continue - - metadata = payload.get("mediaMetadata") - metadata = metadata if isinstance(metadata, dict) else {} - media_status = metadata.get("mediaStatus") - media_status = media_status if isinstance(media_status, dict) else {} - generation_status = media_status.get("mediaGenerationStatus") - - if isinstance(generation_status, str) and ( - generation_status.endswith("FAILED") or generation_status.endswith("CANCELLED") - ): - items.append({ - "name": name, - "primary_media_id": media_id, - "project_id": resolved_project_id, - "done": True, - "status": "FAILED", - "error": generation_status, - }) - continue - - if generation_status != "MEDIA_GENERATION_STATUS_SUCCESSFUL": - items.append({ - "name": name, - "primary_media_id": media_id, - "project_id": resolved_project_id, - "done": False, - "status": "PENDING", - "error": None, - }) - continue - - url = None - url_error = None - url_response = await _fetch_media_url(client, media_id) - if isinstance(url_response, dict) and url_response.get("status", 500) < 400: - url_data = url_response.get("data") - candidate = url_data.get("url") if isinstance(url_data, dict) else None - if isinstance(candidate, str) and candidate.startswith("https://flow-content.google/"): - url = candidate - else: - url_error = "Flow media redirect returned no allowed URL" - else: - url_error = ( - url_response.get("error") - if isinstance(url_response, dict) - else "Flow media redirect failed" - ) - item = { - "name": name, - "primary_media_id": media_id, - "project_id": resolved_project_id, - "done": True, - "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", - "error": None, - "media": { - "media_id": media_id, - "url": url, - "encoded_video_available": False, - }, - } - if include_encoded_video: - item["media"]["encoded_video"] = None - if url_error: - item["media"]["url_error"] = url_error - items.append(item) - - all_done = bool(items) and all(item["done"] for item in items) - any_failed = any(item.get("status") == "FAILED" for item in items) - return { - "project_id": resolved_project_id, - "done": all_done, - "status": "FAILED" if any_failed else ("COMPLETED" if all_done else "PENDING"), - "workflows": items, - } diff --git a/docs/VIDEO_EXPORTS.md b/docs/VIDEO_EXPORTS.md new file mode 100644 index 00000000..9071660a --- /dev/null +++ b/docs/VIDEO_EXPORTS.md @@ -0,0 +1,42 @@ +# 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. + +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 db1f3457..0c70d672 100644 --- a/tests/unit/test_flow_batch.py +++ b/tests/unit/test_flow_batch.py @@ -169,6 +169,19 @@ def test_text_video_matches_the_captured_yhhmef_shape(self): assert len(request[4]) == 6 assert payload[1][5] == self.PID assert payload[2][1] == 1 + 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) @@ -198,6 +211,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 c38c09ed..f9882ff2 100644 --- a/tests/unit/test_flow_client_batch.py +++ b/tests/unit/test_flow_client_batch.py @@ -294,11 +294,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_poll.py b/tests/unit/test_flow_poll.py new file mode 100644 index 00000000..ea31c9aa --- /dev/null +++ b/tests/unit/test_flow_poll.py @@ -0,0 +1,134 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import agent.services.flow_poll as polling +import agent.services.flow_client as flow_client_module +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(monkeypatch): + monkeypatch.setattr(polling, "USE_BATCH_RPC", True) + 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(monkeypatch): + monkeypatch.setattr(polling, "USE_BATCH_RPC", True) + 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_project_snapshot_keeps_trpc_wire_contract(monkeypatch): + monkeypatch.setattr(flow_client_module, "USE_BATCH_RPC", False) + client = FlowClient() + client._send = AsyncMock(return_value={"status": 200, "data": {}}) + + await client.get_project_initial_data("project-1") + + client._send.assert_awaited_once_with( + "trpc_request", + { + "url": ( + "https://labs.google/fx/api/trpc/flow.projectInitialData?input=" + "%7B%22json%22%3A%7B%22projectId%22%3A%22project-1%22%7D%7D" + ), + "method": "GET", + "headers": {"content-type": "application/json"}, + }, + timeout=15, + ) + + +@pytest.mark.asyncio +async def test_flow_client_media_resolver_keeps_legacy_redirect_contract(monkeypatch): + monkeypatch.setattr(flow_client_module, "USE_BATCH_RPC", False) + client = FlowClient() + client._send = AsyncMock(return_value={"status": 200, "data": {}}) + + await client.resolve_media_url("media-1") + + client._send.assert_awaited_once_with( + "trpc_request", + { + "url": "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect?name=media-1", + "method": "GET", + "headers": {"content-type": "application/json"}, + "responseMode": "url", + }, + timeout=15, + ) + + +@pytest.mark.asyncio +async def test_flow_client_media_resolver_uses_get_media_on_batch(monkeypatch): + monkeypatch.setattr(flow_client_module, "USE_BATCH_RPC", True) + 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 cf3b2fb1..036ee267 100644 --- a/tests/unit/test_omni_flash.py +++ b/tests/unit/test_omni_flash.py @@ -342,55 +342,35 @@ def _project_response(generation_status="MEDIA_GENERATION_STATUS_PENDING", inclu @pytest.mark.asyncio -async def test_project_poll_fetch_matches_live_flow_trpc_contract(): +async def test_project_poll_fetch_uses_public_flow_client_surface(): client = MagicMock() - client._send = AsyncMock(return_value={"status": 200, "data": {}}) + client.get_project_initial_data = AsyncMock(return_value={"status": 200, "data": {}}) await _fetch_project_initial_data(client, "project-1") - client._send.assert_awaited_once_with( - "trpc_request", - { - "url": ( - "https://labs.google/fx/api/trpc/flow.projectInitialData?input=" - "%7B%22json%22%3A%7B%22projectId%22%3A%22project-1%22%7D%7D" - ), - "method": "GET", - "headers": {"content-type": "application/json"}, - }, - timeout=15, - ) + client.get_project_initial_data.assert_awaited_once_with("project-1") @pytest.mark.asyncio -async def test_media_redirect_fetch_requests_url_only_mode(): +async def test_media_redirect_fetch_uses_public_flow_client_surface(): client = MagicMock() - client._send = AsyncMock(return_value={"status": 200, "data": {}}) + client.resolve_media_url = AsyncMock(return_value={"status": 200, "data": {}}) await _fetch_media_url(client, "media-1") - client._send.assert_awaited_once_with( - "trpc_request", - { - "url": ( - "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" - "?name=media-1" - ), - "method": "GET", - "headers": {"content-type": "application/json"}, - "responseMode": "url", - }, - timeout=15, - ) + client.resolve_media_url.assert_awaited_once_with("media-1") @pytest.mark.asyncio async def test_batch_omni_poll_uses_as29s_media(monkeypatch): monkeypatch.setattr(omni_flash, "USE_BATCH_RPC", True) 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"}}, + "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([{ @@ -398,14 +378,15 @@ async def test_batch_omni_poll_uses_as29s_media(monkeypatch): }]) assert result["done"] is True 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_omni_poll_pending_uses_project_snapshot_not_legacy_transports(): client = MagicMock() - client._send = AsyncMock(return_value=_project_response()) - client.get_media = AsyncMock(side_effect=AssertionError("legacy get_media forbidden")) + client.get_project_initial_data = AsyncMock(return_value=_project_response()) + client.resolve_media_url = AsyncMock(side_effect=AssertionError("URL lookup should wait for success")) + client.get_media = AsyncMock(side_effect=AssertionError("direct media lookup forbidden")) client.check_video_status = AsyncMock(side_effect=AssertionError("operation poll forbidden")) with patch("agent.services.omni_flash.get_flow_client", return_value=client): @@ -419,9 +400,10 @@ async def test_omni_poll_pending_uses_project_snapshot_not_legacy_transports(): ] ) + client.get_project_initial_data.assert_awaited_once_with("project-1") + client.resolve_media_url.assert_not_awaited() client.get_media.assert_not_awaited() client.check_video_status.assert_not_awaited() - client._send.assert_awaited_once() assert result["done"] is False assert result["status"] == "PENDING" assert result["workflows"][0]["status"] == "PENDING" @@ -430,17 +412,15 @@ async def test_omni_poll_pending_uses_project_snapshot_not_legacy_transports(): @pytest.mark.asyncio async def test_omni_poll_completed_returns_signed_url_without_buffering_video(): client = MagicMock() - client._send = AsyncMock( - side_effect=[ - _project_response("MEDIA_GENERATION_STATUS_SUCCESSFUL"), - { - "status": 200, - "data": { - "url": "https://flow-content.google/video/media-1?Signature=test" - }, - }, - ] + client.get_project_initial_data = AsyncMock( + return_value=_project_response("MEDIA_GENERATION_STATUS_SUCCESSFUL") ) + client.resolve_media_url = AsyncMock(return_value={ + "status": 200, + "data": { + "url": "https://flow-content.google/video/media-1?Signature=test" + }, + }) with patch("agent.services.omni_flash.get_flow_client", return_value=client): result = await check_omni_flash_status( @@ -456,12 +436,17 @@ async def test_omni_poll_completed_returns_signed_url_without_buffering_video(): assert item["media"]["url"].startswith("https://flow-content.google/video/") assert item["media"]["encoded_video_available"] is False assert "encoded_video" not in item["media"] + client.get_project_initial_data.assert_awaited_once_with("project-1") + client.resolve_media_url.assert_awaited_once_with("media-1") @pytest.mark.asyncio async def test_omni_poll_missing_media_is_pending(): client = MagicMock() - client._send = AsyncMock(return_value=_project_response(include_media=False)) + client.get_project_initial_data = AsyncMock( + return_value=_project_response(include_media=False) + ) + client.resolve_media_url = AsyncMock(side_effect=AssertionError("missing media must stay pending")) with patch("agent.services.omni_flash.get_flow_client", return_value=client): result = await check_omni_flash_status( @@ -471,6 +456,8 @@ async def test_omni_poll_missing_media_is_pending(): assert result["done"] is False assert result["workflows"][0]["status"] == "PENDING" + client.get_project_initial_data.assert_awaited_once_with("project-1") + client.resolve_media_url.assert_not_awaited() @pytest.mark.asyncio