diff --git a/agent/api/flow.py b/agent/api/flow.py index 1373e113b..e215fbb31 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -1,8 +1,15 @@ """Direct Flow API endpoints — for manual operations outside the queue.""" from fastapi import APIRouter, HTTPException from pydantic import BaseModel -from typing import Optional +from typing import Literal, Optional + from agent.services.flow_client import get_flow_client +from agent.services.omni_flash import ( + check_omni_flash_status, + generate_omni_flash_first_frame_video, + generate_omni_flash_first_last_video, + generate_omni_flash_video, +) router = APIRouter(prefix="/flow", tags=["flow"]) @@ -23,6 +30,9 @@ class GenerateVideoRequest(BaseModel): aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT" end_image_media_id: Optional[str] = None user_paygate_tier: str = "PAYGATE_TIER_ONE" + # Backward compatible: legacy requests remain Veo unless explicitly set. + model_family: Literal["veo", "omni_flash"] = "veo" + duration_s: int = 8 class GenerateVideoRefsRequest(BaseModel): @@ -32,6 +42,20 @@ class GenerateVideoRefsRequest(BaseModel): scene_id: str aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT" user_paygate_tier: str = "PAYGATE_TIER_ONE" + # Backward compatible: existing callers keep the Veo R2V path unless they + # explicitly opt into Omni Flash. + model_family: Literal["veo", "omni_flash"] = "veo" + duration_s: int = 8 + + +class GenerateOmniFlashVideoRequest(BaseModel): + reference_media_ids: list[str] + prompt: str + project_id: str + scene_id: str = "" + duration_s: int = 8 + aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT" + user_paygate_tier: str = "PAYGATE_TIER_ONE" class UpscaleVideoRequest(BaseModel): @@ -48,7 +72,19 @@ class UploadImageRequest(BaseModel): class CheckStatusRequest(BaseModel): - operations: list[dict] + operations: list[dict] = [] + # Omni/workflow-mode callers should pass workflow descriptors instead of + # operation handles. If workflows is set, /check-status automatically uses + # authenticated Flow project polling. + workflows: Optional[list[dict]] = None + project_id: str = "" + include_encoded_video: bool = False + + +class CheckOmniStatusRequest(BaseModel): + workflows: list[dict] + project_id: str = "" + include_encoded_video: bool = False class EditImageRequest(BaseModel): @@ -95,11 +131,45 @@ async def generate_image(body: GenerateImageRequest): @router.post("/generate-video") async def generate_video(body: GenerateVideoRequest): - """Submit video generation (returns operations for polling).""" + """Submit frame-conditioned video generation using Veo or Omni Flash. + + Existing callers default to Veo. For Omni set ``model_family=omni_flash`` + and ``duration_s`` to 4/6/8/10. With only ``start_image_media_id`` the + request uses Omni First frame. When ``end_image_media_id`` is also present, + it uses Omni First+Last frames. + + Omni responses include ``flowkitPolling.workflows`` and must use workflow + media polling rather than legacy operation polling. + """ client = get_flow_client() if not client.connected: raise HTTPException(503, "Extension not connected") - result = await client.generate_video(**body.model_dump(exclude_none=True)) + + if body.model_family == "omni_flash": + try: + common = dict( + start_image_media_id=body.start_image_media_id, + prompt=body.prompt, + project_id=body.project_id, + scene_id=body.scene_id, + duration_s=body.duration_s, + aspect_ratio=body.aspect_ratio, + user_paygate_tier=body.user_paygate_tier, + ) + if body.end_image_media_id: + result = await generate_omni_flash_first_last_video( + end_image_media_id=body.end_image_media_id, + **common, + ) + else: + result = await generate_omni_flash_first_frame_video(**common) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + else: + result = await client.generate_video( + **body.model_dump(exclude={"model_family", "duration_s"}, exclude_none=True) + ) + 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"))) return result.get("data", result) @@ -107,11 +177,54 @@ async def generate_video(body: GenerateVideoRequest): @router.post("/generate-video-refs") async def generate_video_refs(body: GenerateVideoRefsRequest): - """Submit r2v video generation from reference images.""" + """Submit reference-to-video generation using Veo or Gemini Omni Flash. + + Existing requests default to ``model_family=veo``. Set + ``model_family=omni_flash`` and ``duration_s`` to 4/6/8/10 to use Omni. + Omni responses include ``flowkitPolling.workflows``; poll those workflows, + not the operation-looking handles in the raw Flow response. + """ + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + + if body.model_family == "omni_flash": + try: + result = await generate_omni_flash_video( + reference_media_ids=body.reference_media_ids, + prompt=body.prompt, + project_id=body.project_id, + scene_id=body.scene_id, + duration_s=body.duration_s, + aspect_ratio=body.aspect_ratio, + user_paygate_tier=body.user_paygate_tier, + ) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + else: + result = await client.generate_video_from_references( + **body.model_dump(exclude={"model_family", "duration_s"}) + ) + + 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"))) + return result.get("data", result) + + +@router.post("/generate-video-omni") +async def generate_video_omni(body: GenerateOmniFlashVideoRequest): + """Submit Gemini Omni Flash reference-to-video generation. + + The response includes ``flowkitPolling.workflows`` for the correct + workflow/media polling path. + """ client = get_flow_client() if not client.connected: raise HTTPException(503, "Extension not connected") - result = await client.generate_video_from_references(**body.model_dump()) + try: + result = await generate_omni_flash_video(**body.model_dump()) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc 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"))) return result.get("data", result) @@ -131,16 +244,56 @@ async def upscale_video(body: UpscaleVideoRequest): @router.post("/check-status") async def check_status(body: CheckStatusRequest): - """Check video generation status.""" + """Check Veo operation status or Omni workflow/media status. + + Veo: pass ``operations``. + Omni Flash: pass ``workflows`` from submit ``flowkitPolling.workflows``. + """ client = get_flow_client() if not client.connected: raise HTTPException(503, "Extension not connected") + + if body.workflows: + try: + return await check_omni_flash_status( + body.workflows, + include_encoded_video=body.include_encoded_video, + project_id=body.project_id, + ) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + except RuntimeError as exc: + raise HTTPException(502, str(exc)) from exc + + if not body.operations: + raise HTTPException(400, "Provide operations for Veo or workflows for Omni Flash") + result = await client.check_video_status(body.operations) if result.get("error"): raise HTTPException(502, result["error"]) + if isinstance(result.get("status"), int) and result["status"] >= 400: + raise HTTPException(result["status"], result.get("data", "Flow polling failed")) return result.get("data", result) +@router.post("/check-omni-status") +async def check_omni_status(body: CheckOmniStatusRequest): + """Poll Gemini Omni Flash jobs via workflow primary media IDs.""" + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + try: + return await check_omni_flash_status( + body.workflows, + include_encoded_video=body.include_encoded_video, + project_id=body.project_id, + ) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + except RuntimeError as exc: + raise HTTPException(502, str(exc)) from exc + + @router.post("/refresh-urls/{project_id}") async def refresh_project_urls(project_id: str): """Bulk refresh all media URLs for a project via per-media get_media calls.""" @@ -157,8 +310,8 @@ async def refresh_project_urls(project_id: str): async def get_media(media_id: str): """Get media metadata + fresh signed URL from Google Flow. - Returns the raw response which should contain a fresh fifeUrl/servingUri. - Use this to refresh expired GCS signed URLs. + Returns the raw response which may contain ``video.encodedVideo`` for + workflow-backed video generations. """ client = get_flow_client() if not client.connected: diff --git a/agent/api/models.py b/agent/api/models.py index aea93127e..4b1f35885 100644 --- a/agent/api/models.py +++ b/agent/api/models.py @@ -25,7 +25,11 @@ def _write_models(data: dict): def _reload_config(data: dict): - """Hot-reload model keys into the running config module.""" + """Hot-reload model keys into the running config module. + + Omni Flash is resolved directly from models.json at submit time, so no + config-module mirror is required for that section. + """ config.VIDEO_MODELS.clear() config.VIDEO_MODELS.update(data["video_models"]) config.UPSCALE_MODELS.clear() @@ -54,18 +58,35 @@ async def patch_models(body: dict): } } } + + Omni Flash duration keys are configurable too: + { + "omni_flash_models": { + "reference_to_video": {"10": "abra_r2v_10s"} + } + } """ current = _read_models() - # Deep merge: only update keys that are provided - for section in ("video_models", "image_models", "upscale_models"): + # Deep merge: only update keys that are provided. + for section in ( + "video_models", + "omni_flash_models", + "image_models", + "upscale_models", + ): if section not in body: continue - if section == "upscale_models" or section == "image_models": - # Flat dict — direct merge - current[section].update(body[section]) + if section in ("upscale_models", "image_models"): + # Flat dict — direct merge. + current.setdefault(section, {}).update(body[section]) + elif section == "omni_flash_models": + # Nested by generation mode -> duration -> model key. + target = current.setdefault(section, {}) + for mode, durations in body[section].items(): + target.setdefault(mode, {}).update(durations) else: - # Nested dict — merge per tier, per gen_type + # Nested dict — merge per tier, per gen_type. for tier, gen_types in body[section].items(): if tier not in current[section]: current[section][tier] = {} diff --git a/agent/models.json b/agent/models.json index 631b43a41..c44c789b5 100644 --- a/agent/models.json +++ b/agent/models.json @@ -29,6 +29,26 @@ } } }, + "omni_flash_models": { + "frame_to_video": { + "4": "abra_i2v_4s", + "6": "abra_i2v_6s", + "8": "abra_i2v_8s", + "10": "abra_i2v_10s" + }, + "start_end_frame_to_video": { + "4": "abra_i2v_4s", + "6": "abra_i2v_6s", + "8": "abra_i2v_8s", + "10": "abra_i2v_10s" + }, + "reference_to_video": { + "4": "abra_r2v_4s", + "6": "abra_r2v_6s", + "8": "abra_r2v_8s", + "10": "abra_r2v_10s" + } + }, "upscale_models": { "VIDEO_RESOLUTION_4K": "veo_3_1_upsampler_4k", "VIDEO_RESOLUTION_1080P": "veo_3_1_upsampler_1080p" diff --git a/agent/services/omni_flash.py b/agent/services/omni_flash.py new file mode 100644 index 000000000..2324b812f --- /dev/null +++ b/agent/services/omni_flash.py @@ -0,0 +1,547 @@ +"""Gemini Omni Flash video generation through the Google Flow bridge. + +Supported Omni surfaces in this module: + +* first frame -> video via ``batchAsyncGenerateVideoStartImage`` +* first + last frame -> video via ``batchAsyncGenerateVideoStartAndEndImage`` +* reference images -> video via ``batchAsyncGenerateVideoReferenceImages`` + +Omni duration-specific model keys live in ``agent/models.json``. First-frame +Flow requests have been captured with ``abra_i2v_s``. The current +First+Last rollout uses the same Omni I2V family but the StartAndEnd endpoint; +that mapping is deliberately configurable separately so it can be changed +without a code release if Google's rollout rotates the wire key. + +Important: Omni submit responses may contain operation-looking handles, but +those handles are not compatible with the legacy +``batchCheckAsyncVideoGenerationStatus`` polling endpoint. Omni jobs are +workflow-backed and are polled through Flow's authenticated project data. +""" + +from __future__ import annotations + +import json +import time +import uuid +from pathlib import Path +from urllib.parse import quote + +from agent.services.flow_client import get_flow_client +from agent.services.headers import random_headers + +_MODELS_FILE = Path(__file__).parent.parent / "models.json" + +OMNI_FLASH_VALID_DURATIONS = (4, 6, 8, 10) +OMNI_FLASH_VALID_ASPECTS = { + "VIDEO_ASPECT_RATIO_PORTRAIT", + "VIDEO_ASPECT_RATIO_LANDSCAPE", +} +OMNI_FLASH_MAX_REFERENCE_IMAGES = 7 +# Informational only. Flow pricing can be promotional/variable. +OMNI_FLASH_CREDIT_COST = {4: 15, 6: 20, 8: 25, 10: 30} + + +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, + ) + + +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, + ) + + +def _validate_duration(duration_s: int) -> None: + if duration_s not in OMNI_FLASH_VALID_DURATIONS: + raise ValueError( + f"Omni Flash duration {duration_s}s is unsupported; " + f"choose one of {list(OMNI_FLASH_VALID_DURATIONS)}" + ) + + +def _validate_aspect(aspect_ratio: str) -> None: + if aspect_ratio not in OMNI_FLASH_VALID_ASPECTS: + raise ValueError( + f"Omni Flash aspect ratio {aspect_ratio!r} is unsupported; " + "use VIDEO_ASPECT_RATIO_PORTRAIT or VIDEO_ASPECT_RATIO_LANDSCAPE" + ) + + +def _load_model_key(duration_s: int, mode: str = "reference_to_video") -> str: + """Resolve a configured Omni Flash model key for ``mode`` + duration.""" + _validate_duration(duration_s) + + with open(_MODELS_FILE, encoding="utf-8") as f: + models = json.load(f) + + key = ( + models.get("omni_flash_models", {}) + .get(mode, {}) + .get(str(duration_s)) + ) + if not key: + raise ValueError( + f"No Omni Flash model key configured for mode {mode!r}, {duration_s}s" + ) + return key + + +def _validate_reference_inputs( + reference_media_ids: list[str], + duration_s: int, + aspect_ratio: str, +) -> list[str]: + _validate_duration(duration_s) + _validate_aspect(aspect_ratio) + + refs = [mid for mid in (reference_media_ids or []) if isinstance(mid, str) and mid] + if not refs: + raise ValueError("Omni Flash requires at least one reference image") + if len(refs) > OMNI_FLASH_MAX_REFERENCE_IMAGES: + raise ValueError( + f"Omni Flash accepts at most {OMNI_FLASH_MAX_REFERENCE_IMAGES} reference images" + ) + return refs + + +def _validate_frame_inputs( + start_image_media_id: str, + end_image_media_id: str | None, + duration_s: int, + aspect_ratio: str, +) -> None: + _validate_duration(duration_s) + _validate_aspect(aspect_ratio) + if not isinstance(start_image_media_id, str) or not start_image_media_id: + raise ValueError("Omni Flash first-frame generation requires start_image_media_id") + if end_image_media_id is not None and ( + not isinstance(end_image_media_id, str) or not end_image_media_id + ): + 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 + + +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 + + +async def _submit_omni_frame_video( + *, + start_image_media_id: str, + end_image_media_id: str | None, + prompt: str, + project_id: str, + scene_id: str = "", + duration_s: int = 8, + aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT", + user_paygate_tier: str = "PAYGATE_TIER_ONE", + seed: int | None = None, +) -> dict: + """Submit Omni first-frame or First+Last generation.""" + _validate_frame_inputs( + start_image_media_id, + end_image_media_id, + duration_s, + aspect_ratio, + ) + + mode = ( + "start_end_frame_to_video" + if end_image_media_id is not None + else "frame_to_video" + ) + endpoint = ( + "generate_video_start_end" + if end_image_media_id is not None + else "generate_video" + ) + model_key = _load_model_key(duration_s, mode=mode) + client = get_flow_client() + ts = int(time.time() * 1000) + + request_item = { + "aspectRatio": aspect_ratio, + "textInput": {"structuredPrompt": {"parts": [{"text": prompt}]}}, + "videoModelKey": model_key, + "seed": seed if seed is not None else ts % 1_000_000, + "metadata": {"sceneId": scene_id} if scene_id else {}, + "startImage": {"mediaId": start_image_media_id}, + } + if end_image_media_id is not None: + request_item["endImage"] = {"mediaId": end_image_media_id} + + context = client._client_context(project_id, user_paygate_tier) + body = { + "mediaGenerationContext": {"batchId": str(uuid.uuid4())}, + "clientContext": {**context, "sessionId": f";{ts}"}, + "requests": [request_item], + "useV2ModelConfig": True, + } + + result = await client._send( + "api_request", + { + "url": client._build_url(endpoint), + "method": "POST", + "headers": random_headers(), + "body": body, + "captchaAction": "VIDEO_GENERATION", + }, + timeout=60, + ) + return _annotate_polling(result, project_id) + + +async def generate_omni_flash_first_frame_video( + start_image_media_id: str, + prompt: str, + project_id: str, + scene_id: str = "", + duration_s: int = 8, + aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT", + user_paygate_tier: str = "PAYGATE_TIER_ONE", + seed: int | None = None, +) -> dict: + """Submit Omni Flash First frame -> video.""" + return await _submit_omni_frame_video( + start_image_media_id=start_image_media_id, + end_image_media_id=None, + prompt=prompt, + project_id=project_id, + scene_id=scene_id, + duration_s=duration_s, + aspect_ratio=aspect_ratio, + user_paygate_tier=user_paygate_tier, + seed=seed, + ) + + +async def generate_omni_flash_first_last_video( + start_image_media_id: str, + end_image_media_id: str, + prompt: str, + project_id: str, + scene_id: str = "", + duration_s: int = 8, + aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT", + user_paygate_tier: str = "PAYGATE_TIER_ONE", + seed: int | None = None, +) -> dict: + """Submit Omni Flash First + Last frame -> video.""" + return await _submit_omni_frame_video( + start_image_media_id=start_image_media_id, + end_image_media_id=end_image_media_id, + prompt=prompt, + project_id=project_id, + scene_id=scene_id, + duration_s=duration_s, + aspect_ratio=aspect_ratio, + user_paygate_tier=user_paygate_tier, + seed=seed, + ) + + +async def generate_omni_flash_video( + reference_media_ids: list[str], + prompt: str, + project_id: str, + scene_id: str = "", + duration_s: int = 8, + aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT", + user_paygate_tier: str = "PAYGATE_TIER_ONE", + seed: int | None = None, +) -> dict: + """Submit an Omni Flash reference-to-video generation. + + Successful responses are annotated with ``data.flowkitPolling`` containing + the workflow names and primary media IDs required by the Omni polling path. + Do not feed Omni operation handles to ``check_video_status``. + """ + refs = _validate_reference_inputs(reference_media_ids, duration_s, aspect_ratio) + model_key = _load_model_key(duration_s, mode="reference_to_video") + client = get_flow_client() + + ts = int(time.time() * 1000) + request_item = { + "aspectRatio": aspect_ratio, + "textInput": {"structuredPrompt": {"parts": [{"text": prompt}]}}, + "videoModelKey": model_key, + "seed": seed if seed is not None else ts % 1_000_000, + "metadata": {"sceneId": scene_id} if scene_id else {}, + "referenceImages": [ + {"mediaId": mid, "imageUsageType": "IMAGE_USAGE_TYPE_ASSET"} + for mid in refs + ], + } + + context = client._client_context(project_id, user_paygate_tier) + body = { + "mediaGenerationContext": { + "batchId": str(uuid.uuid4()), + "audioFailurePreference": "BLOCK_SILENCED_VIDEOS", + }, + "clientContext": {**context, "sessionId": f";{ts}"}, + "requests": [request_item], + "useV2ModelConfig": True, + } + + url = client._build_url("generate_video_references") + result = await client._send( + "api_request", + { + "url": url, + "method": "POST", + "headers": random_headers(), + "body": body, + "captchaAction": "VIDEO_GENERATION", + }, + timeout=60, + ) + return _annotate_polling(result, project_id) + + +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. + + Flow's production UI exposes workflow status through its authenticated + ``flow.projectInitialData`` tRPC response. The old ``/v1/media`` transport + currently returns ``INVALID_ARGUMENT`` for these workflow media IDs. + """ + 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")), + "", + ) + 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/OMNI_FLASH.md b/docs/OMNI_FLASH.md new file mode 100644 index 000000000..93aa1e94c --- /dev/null +++ b/docs/OMNI_FLASH.md @@ -0,0 +1,275 @@ +# Gemini Omni Flash: integration guide for agents + +FlowKit exposes Gemini Omni Flash video generation through the authenticated Google Flow session in its persistent Chrome profile. An integrating service talks only to the FlowKit REST API; it must not call Google Flow endpoints or the extension WebSocket directly. + +## Prerequisites and base URL + +Before submitting work, both checks must pass: + +```bash +curl -fsS "$FLOWKIT_BASE_URL/health" +curl -fsS "$FLOWKIT_BASE_URL/api/flow/status" +``` + +Expected state: + +```json +{"status":"ok","extension_connected":true} +{"connected":true,"flow_key_present":true} +``` + +Use `http://127.0.0.1:8100` when the caller runs on the FlowKit host. For a remote integration, set `FLOWKIT_BASE_URL` to the protected HTTPS reverse-proxy URL and allow only the required source IPs or private network. Do not expose Chrome, VNC/noVNC, the extension WebSocket, or port 8100 publicly. + +## Supported modes + +| Mode | Inputs | Endpoint | Internal model family | +|---|---|---|---| +| First frame to video | one uploaded start image | `POST /api/flow/generate-video` | `abra_i2v_s` | +| First + Last frame to video | uploaded start and end images | `POST /api/flow/generate-video` | `abra_i2v_s` | +| References to video | 1-7 uploaded reference images | `POST /api/flow/generate-video-omni` | `abra_r2v_s` | + +Supported durations are `4`, `6`, `8`, and `10` seconds. Supported aspect ratios are: + +- `VIDEO_ASPECT_RATIO_PORTRAIT` (`9:16`) +- `VIDEO_ASPECT_RATIO_LANDSCAPE` (`16:9`) + +First + Last generation with `batchAsyncGenerateVideoStartAndEndImage` and the current `abra_i2v_*` mapping has been verified with a real Flow generation. + +## End-to-end integration flow + +An integration agent should implement this state machine: + +1. Check `/health` and `/api/flow/status`. +2. Make each source image readable on the FlowKit server. +3. Call `/api/flow/upload-image` for every source image and retain each returned `media_id`. +4. Submit exactly one Omni request and persist its complete `flowkitPolling` object. +5. Poll `/api/flow/check-omni-status` every 10-20 seconds using `project_id` and `workflows` from `flowkitPolling`. +6. On `PENDING`, continue polling. On `FAILED`, stop and report the returned error. On `COMPLETED`, immediately download every non-null `media.url`. +7. Store the downloaded video in the project's own durable storage. The returned Google URL is signed and short-lived. + +Do not send Omni workflow names to the legacy Veo `batchCheckAsyncVideoGenerationStatus` operation poller. Do not use the obsolete `/v1/media/` polling path. + +## Supplying images + +`POST /api/flow/upload-image` is not a multipart upload endpoint. Its `file_path` is an absolute path on the **FlowKit server**, not on the calling server. + +For a remote integration, first stage the file on the FlowKit host using an authenticated transfer such as SFTP/SCP, a private shared volume, or a separately secured upload service. Use a unique per-job directory, validate file size/type, and make the file readable by the FlowKit service account. Then call: + +```bash +curl -fsS -X POST "$FLOWKIT_BASE_URL/api/flow/upload-image" \ + -H 'Content-Type: application/json' \ + -d '{ + "file_path": "/var/lib/flowkit/input/JOB_ID/start.jpg", + "project_id": "FLOW_PROJECT_ID", + "file_name": "start.jpg" + }' +``` + +Response: + +```json +{"media_id":"FLOW_MEDIA_ID","raw":{}} +``` + +Use the returned `media_id` in generation requests. Never pass a caller-local path such as `/tmp/image.jpg` unless that exact file also exists on the FlowKit server. + +## Prompts + +The input images define identity, appearance, objects, and composition. The prompt should primarily describe motion, camera behavior, timing, and audio/dialogue. Avoid restating a person's detailed appearance when reference images already provide it. + +Example: + +```text +The woman turns naturally toward the camera, smiles, then raises one leg into the final pose. Subtle handheld camera movement, realistic cloth and hair motion, stable face and body proportions. +``` + +For longer clips, timed instructions are useful, for example: `0-3s: ...; 3-6s: ...; 6-8s: ...`. + +## First frame to video + +```bash +curl -fsS -X POST "$FLOWKIT_BASE_URL/api/flow/generate-video" \ + -H 'Content-Type: application/json' \ + -d '{ + "model_family": "omni_flash", + "start_image_media_id": "START_MEDIA_ID", + "prompt": "The subject looks toward the camera and smiles; subtle cinematic push-in, natural motion.", + "project_id": "FLOW_PROJECT_ID", + "scene_id": "JOB_ID", + "duration_s": 4, + "aspect_ratio": "VIDEO_ASPECT_RATIO_PORTRAIT", + "user_paygate_tier": "PAYGATE_TIER_ONE" + }' +``` + +This uses `batchAsyncGenerateVideoStartImage`. + +## First + Last frame to video + +```bash +curl -fsS -X POST "$FLOWKIT_BASE_URL/api/flow/generate-video" \ + -H 'Content-Type: application/json' \ + -d '{ + "model_family": "omni_flash", + "start_image_media_id": "START_MEDIA_ID", + "end_image_media_id": "END_MEDIA_ID", + "prompt": "The subject moves naturally from the first pose to the final pose; stable identity and smooth realistic motion.", + "project_id": "FLOW_PROJECT_ID", + "scene_id": "JOB_ID", + "duration_s": 4, + "aspect_ratio": "VIDEO_ASPECT_RATIO_PORTRAIT", + "user_paygate_tier": "PAYGATE_TIER_ONE" + }' +``` + +This uses `batchAsyncGenerateVideoStartAndEndImage` and sends both `startImage` and `endImage`. + +## References to video + +Use 1-7 media IDs. Reference images act as components/identity/style guidance; they are not treated as fixed first and last frames. + +```bash +curl -fsS -X POST "$FLOWKIT_BASE_URL/api/flow/generate-video-omni" \ + -H 'Content-Type: application/json' \ + -d '{ + "reference_media_ids": ["REFERENCE_MEDIA_ID_1", "REFERENCE_MEDIA_ID_2"], + "prompt": "Cinematic handheld shot with natural character motion and consistent referenced subjects.", + "project_id": "FLOW_PROJECT_ID", + "scene_id": "JOB_ID", + "duration_s": 4, + "aspect_ratio": "VIDEO_ASPECT_RATIO_LANDSCAPE", + "user_paygate_tier": "PAYGATE_TIER_ONE" + }' +``` + +The compatible generic endpoint is `POST /api/flow/generate-video-refs` with the same fields plus `"model_family":"omni_flash"`. + +## Submit response and polling + +Persist the entire normalized polling descriptor returned by submit: + +```json +{ + "flowkitPolling": { + "mode": "project_media", + "project_id": "FLOW_PROJECT_ID", + "workflows": [ + { + "name": "WORKFLOW_NAME", + "primary_media_id": "PRIMARY_MEDIA_ID", + "project_id": "FLOW_PROJECT_ID" + } + ] + } +} +``` + +Poll using those values without transforming workflow names into operation handles: + +```bash +curl -fsS -X POST "$FLOWKIT_BASE_URL/api/flow/check-omni-status" \ + -H 'Content-Type: application/json' \ + -d '{ + "project_id": "FLOW_PROJECT_ID", + "workflows": [ + { + "name": "WORKFLOW_NAME", + "primary_media_id": "PRIMARY_MEDIA_ID", + "project_id": "FLOW_PROJECT_ID" + } + ], + "include_encoded_video": false + }' +``` + +The generic `POST /api/flow/check-status` endpoint also accepts the same `workflows` and automatically selects Omni project polling. + +Pending response: + +```json +{ + "project_id": "FLOW_PROJECT_ID", + "done": false, + "status": "PENDING", + "workflows": [{"done":false,"status":"PENDING"}] +} +``` + +Successful response: + +```json +{ + "project_id": "FLOW_PROJECT_ID", + "done": true, + "status": "COMPLETED", + "workflows": [ + { + "done": true, + "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", + "media": { + "media_id": "PRIMARY_MEDIA_ID", + "url": "https://flow-content.google/...", + "encoded_video_available": false + } + } + ] +} +``` + +Keep `include_encoded_video` set to `false`. FlowKit resolves the signed download URL without buffering the MP4 through Chrome or the extension bridge. + +## Retry and failure policy + +- HTTP `400`: request/contract error. Do not retry unchanged input. +- HTTP `503`: Chrome extension is disconnected. Pause submission and alert or retry health checks with bounded backoff. +- HTTP `502`: Flow/bridge failure. Retry a small bounded number of times with exponential backoff; preserve the original workflow descriptor. +- Poll result `PENDING`: poll again after 10-20 seconds. Do not submit the generation again. +- Poll result `FAILED`: stop polling and surface the workflow error. +- `COMPLETED` with a null URL or `url_error`: poll again to obtain a fresh signed URL; do not regenerate the video. + +Submission is credit-consuming and is not guaranteed to be idempotent. Never blindly retry a timed-out submit unless the integration can determine that no workflow was created. + +## Model configuration + +Mappings live in `agent/models.json`: + +```json +{ + "omni_flash_models": { + "frame_to_video": { + "4": "abra_i2v_4s", + "6": "abra_i2v_6s", + "8": "abra_i2v_8s", + "10": "abra_i2v_10s" + }, + "start_end_frame_to_video": { + "4": "abra_i2v_4s", + "6": "abra_i2v_6s", + "8": "abra_i2v_8s", + "10": "abra_i2v_10s" + }, + "reference_to_video": { + "4": "abra_r2v_4s", + "6": "abra_r2v_6s", + "8": "abra_r2v_8s", + "10": "abra_r2v_10s" + } + } +} +``` + +The mappings can be changed through `PATCH /api/models` if Google rotates internal keys. Treat configured credit-cost estimates as informational only: Google can change pricing, so use `GET /api/flow/credits` and the submit response's `remainingCredits` where available. + +## Minimal agent checklist + +- Use the `/api/flow/...` paths exactly. +- Select Omni explicitly with `model_family: "omni_flash"` on shared endpoints. +- Upload inputs once and reuse returned Flow media IDs. +- Persist `project_id`, workflow `name`, and `primary_media_id` before polling. +- Poll workflows through project media status, never through legacy Veo operations. +- Download signed output URLs immediately into durable project storage. +- Do not log Google auth data, extension messages, or complete signed URLs. +- Use a stable `scene_id`/job ID for traceability. +- Start with 4 seconds for smoke tests to limit credit usage. + +Omni Flash is an unofficial integration over Google Flow's internal interfaces. Endpoints, model keys, and response shapes may change when Flow changes; integrations should fail visibly and retain job metadata for diagnosis. diff --git a/extension/background.js b/extension/background.js index fb1b1b415..e0f601a55 100644 --- a/extension/background.js +++ b/extension/background.js @@ -354,7 +354,7 @@ async function handleSolveCaptcha(msg) { async function handleTrpcRequest(msg) { const { id, params } = msg; - const { url, method = 'POST', headers = {}, body } = params; + const { url, method = 'POST', headers = {}, body, responseMode = 'json' } = params; if (!url || !url.startsWith('https://labs.google/')) { sendToAgent({ id, error: 'INVALID_TRPC_URL' }); @@ -380,7 +380,19 @@ async function handleTrpcRequest(msg) { body: body ? JSON.stringify(body) : undefined, credentials: 'include', }); - const data = await resp.json(); + let data; + if (responseMode === 'url') { + // fetch() has already followed the authenticated Flow redirect. Return + // only the final signed URL and cancel the body so large videos are not + // buffered in the extension or copied through the WebSocket bridge. + data = { + url: resp.url, + contentType: resp.headers.get('content-type'), + }; + await resp.body?.cancel(); + } else { + data = await resp.json(); + } chrome.storage.local.set({ metrics }); updateRequestLog(logId, { status: 'success' }); sendToAgent({ id, status: resp.status, data }); diff --git a/extension/manifest.json b/extension/manifest.json index dff8b6ef4..01105d548 100644 --- a/extension/manifest.json +++ b/extension/manifest.json @@ -1,12 +1,13 @@ { "manifest_version": 3, "name": "Flow Kit", - "version": "0.2.0", + "version": "0.2.1", "description": "Local agent bridge for Google Flow API — captures tokens, solves reCAPTCHA, proxies API calls", "permissions": ["storage", "alarms", "tabs", "webRequest", "scripting", "declarativeNetRequest", "sidePanel"], "host_permissions": [ "https://labs.google/*", "https://aisandbox-pa.googleapis.com/*", + "https://flow-content.google/*", "http://127.0.0.1:8100/*" ], "background": { diff --git a/tests/unit/test_omni_flash.py b/tests/unit/test_omni_flash.py new file mode 100644 index 000000000..38d5ed9fc --- /dev/null +++ b/tests/unit/test_omni_flash.py @@ -0,0 +1,464 @@ +"""Unit tests for Gemini Omni Flash Flow submissions and workflow polling.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from agent.services.omni_flash import ( + OMNI_FLASH_MAX_REFERENCE_IMAGES, + _fetch_media_url, + _fetch_project_initial_data, + _load_model_key, + check_omni_flash_status, + extract_omni_workflows, + generate_omni_flash_first_frame_video, + generate_omni_flash_first_last_video, + generate_omni_flash_video, +) + + +@pytest.mark.parametrize( + ("duration", "expected"), + [ + (4, "abra_r2v_4s"), + (6, "abra_r2v_6s"), + (8, "abra_r2v_8s"), + (10, "abra_r2v_10s"), + ], +) +def test_omni_reference_duration_model_keys(duration, expected): + assert _load_model_key(duration) == expected + + +@pytest.mark.parametrize( + ("duration", "expected"), + [ + (4, "abra_i2v_4s"), + (6, "abra_i2v_6s"), + (8, "abra_i2v_8s"), + (10, "abra_i2v_10s"), + ], +) +def test_omni_first_frame_model_keys(duration, expected): + assert _load_model_key(duration, mode="frame_to_video") == expected + + +@pytest.mark.parametrize( + ("duration", "expected"), + [ + (4, "abra_i2v_4s"), + (6, "abra_i2v_6s"), + (8, "abra_i2v_8s"), + (10, "abra_i2v_10s"), + ], +) +def test_omni_first_last_model_keys_are_independently_configured(duration, expected): + assert _load_model_key(duration, mode="start_end_frame_to_video") == expected + + +def test_invalid_duration_fails_before_submit(): + with pytest.raises(ValueError, match="duration 5s is unsupported"): + _load_model_key(5) + + +def test_extract_omni_workflows_uses_primary_media_id(): + result = { + "status": 200, + "data": { + "operations": [ + { + "operation": {"name": "operation-looking-handle"}, + "status": "MEDIA_GENERATION_STATUS_PENDING", + } + ], + "workflows": [ + { + "name": "workflow-1", + "metadata": {"primaryMediaId": "media-1"}, + } + ], + }, + } + assert extract_omni_workflows(result) == [ + {"name": "workflow-1", "primary_media_id": "media-1"} + ] + + +def _mock_submit_client(): + client = MagicMock() + client._client_context.return_value = { + "projectId": "project-1", + "tool": "PINHOLE", + "userPaygateTier": "PAYGATE_TIER_ONE", + "recaptchaContext": { + "applicationType": "RECAPTCHA_APPLICATION_TYPE_WEB", + "token": "", + }, + "sessionId": ";old", + } + client._build_url.side_effect = lambda name: f"https://example.test/{name}" + client._send = AsyncMock( + return_value={ + "status": 200, + "data": { + "operations": [ + { + "operation": {"name": "op-1"}, + "status": "MEDIA_GENERATION_STATUS_PENDING", + } + ], + "workflows": [ + { + "name": "workflow-1", + "metadata": {"primaryMediaId": "media-1"}, + } + ], + }, + } + ) + return client + + +@pytest.mark.asyncio +async def test_submit_builds_flow_omni_first_frame_request_and_poll_descriptor(): + client = _mock_submit_client() + + with patch("agent.services.omni_flash.get_flow_client", return_value=client): + result = await generate_omni_flash_first_frame_video( + start_image_media_id="start-1", + prompt="Camera slowly pushes in", + project_id="project-1", + scene_id="scene-1", + duration_s=10, + aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE", + user_paygate_tier="PAYGATE_TIER_ONE", + seed=321, + ) + + client._build_url.assert_called_once_with("generate_video") + method, params = client._send.await_args.args[:2] + assert method == "api_request" + assert params["captchaAction"] == "VIDEO_GENERATION" + + body = params["body"] + assert body["useV2ModelConfig"] is True + assert set(body["mediaGenerationContext"]) == {"batchId"} + request = body["requests"][0] + assert request["videoModelKey"] == "abra_i2v_10s" + assert request["startImage"] == {"mediaId": "start-1"} + assert "endImage" not in request + assert request["seed"] == 321 + assert result["data"]["flowkitPolling"]["mode"] == "project_media" + assert result["data"]["flowkitPolling"]["project_id"] == "project-1" + + +@pytest.mark.asyncio +async def test_submit_builds_flow_omni_first_last_request(): + client = _mock_submit_client() + + with patch("agent.services.omni_flash.get_flow_client", return_value=client): + result = await generate_omni_flash_first_last_video( + start_image_media_id="start-1", + end_image_media_id="end-1", + prompt="Move naturally from the first composition to the final composition", + project_id="project-1", + scene_id="scene-1", + duration_s=8, + aspect_ratio="VIDEO_ASPECT_RATIO_PORTRAIT", + user_paygate_tier="PAYGATE_TIER_ONE", + seed=456, + ) + + client._build_url.assert_called_once_with("generate_video_start_end") + method, params = client._send.await_args.args[:2] + assert method == "api_request" + body = params["body"] + request = body["requests"][0] + assert request["videoModelKey"] == "abra_i2v_8s" + assert request["startImage"] == {"mediaId": "start-1"} + assert request["endImage"] == {"mediaId": "end-1"} + assert request["aspectRatio"] == "VIDEO_ASPECT_RATIO_PORTRAIT" + assert request["seed"] == 456 + assert result["data"]["flowkitPolling"]["workflows"] == [ + { + "name": "workflow-1", + "primary_media_id": "media-1", + "project_id": "project-1", + } + ] + + +@pytest.mark.asyncio +async def test_first_frame_rejects_missing_start_before_submit(): + with pytest.raises(ValueError, match="requires start_image_media_id"): + await generate_omni_flash_first_frame_video( + start_image_media_id="", + prompt="test", + project_id="p", + ) + + +@pytest.mark.asyncio +async def test_first_last_rejects_missing_end_before_submit(): + with pytest.raises(ValueError, match="non-empty end_image_media_id"): + await generate_omni_flash_first_last_video( + start_image_media_id="start", + end_image_media_id="", + prompt="test", + project_id="p", + ) + + +@pytest.mark.asyncio +async def test_submit_builds_flow_omni_r2v_request_and_poll_descriptor(): + client = _mock_submit_client() + + with patch("agent.services.omni_flash.get_flow_client", return_value=client): + result = await generate_omni_flash_video( + reference_media_ids=["ref-1", "ref-2"], + prompt="Two friends talking in a cafe", + project_id="project-1", + scene_id="scene-1", + duration_s=10, + aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE", + user_paygate_tier="PAYGATE_TIER_ONE", + seed=123, + ) + + assert result["status"] == 200 + assert result["data"]["flowkitPolling"] == { + "mode": "project_media", + "project_id": "project-1", + "workflows": [ + { + "name": "workflow-1", + "primary_media_id": "media-1", + "project_id": "project-1", + } + ], + } + client._build_url.assert_called_once_with("generate_video_references") + client._send.assert_awaited_once() + + method, params = client._send.await_args.args[:2] + assert method == "api_request" + assert params["captchaAction"] == "VIDEO_GENERATION" + + body = params["body"] + assert body["useV2ModelConfig"] is True + assert body["mediaGenerationContext"]["audioFailurePreference"] == "BLOCK_SILENCED_VIDEOS" + assert body["clientContext"]["projectId"] == "project-1" + + request = body["requests"][0] + assert request["videoModelKey"] == "abra_r2v_10s" + assert request["aspectRatio"] == "VIDEO_ASPECT_RATIO_LANDSCAPE" + assert request["seed"] == 123 + assert request["metadata"] == {"sceneId": "scene-1"} + assert request["referenceImages"] == [ + {"mediaId": "ref-1", "imageUsageType": "IMAGE_USAGE_TYPE_ASSET"}, + {"mediaId": "ref-2", "imageUsageType": "IMAGE_USAGE_TYPE_ASSET"}, + ] + + +def _project_response(generation_status="MEDIA_GENERATION_STATUS_PENDING", include_media=True): + media = [] + if include_media: + media.append( + { + "name": "media-1", + "projectId": "project-1", + "workflowId": "workflow-1", + "mediaMetadata": { + "mediaStatus": {"mediaGenerationStatus": generation_status} + }, + "video": {"generatedVideo": {"model": "abra_r2v_4s"}}, + } + ) + return { + "status": 200, + "data": { + "result": { + "data": { + "json": { + "projectContents": { + "workflows": [ + { + "name": "workflow-1", + "projectId": "project-1", + "metadata": {"primaryMediaId": "media-1"}, + } + ], + "media": media, + } + } + } + } + }, + } + + +@pytest.mark.asyncio +async def test_project_poll_fetch_matches_live_flow_trpc_contract(): + client = MagicMock() + client._send = 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, + ) + + +@pytest.mark.asyncio +async def test_media_redirect_fetch_requests_url_only_mode(): + client = MagicMock() + client._send = 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, + ) + + +@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.check_video_status = AsyncMock(side_effect=AssertionError("operation poll forbidden")) + + 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", + } + ] + ) + + 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" + + +@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" + }, + }, + ] + ) + + 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", + ) + + assert result["done"] is True + assert result["status"] == "COMPLETED" + item = result["workflows"][0] + assert item["status"] == "MEDIA_GENERATION_STATUS_SUCCESSFUL" + assert item["media"]["media_id"] == "media-1" + 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"] + + +@pytest.mark.asyncio +async def test_omni_poll_missing_media_is_pending(): + client = MagicMock() + client._send = AsyncMock(return_value=_project_response(include_media=False)) + + 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", + ) + + assert result["done"] is False + assert result["workflows"][0]["status"] == "PENDING" + + +@pytest.mark.asyncio +async def test_omni_poll_requires_project_id_for_legacy_descriptors(): + with pytest.raises(ValueError, match="requires project_id"): + await check_omni_flash_status( + [{"name": "workflow-1", "primary_media_id": "media-1"}] + ) + + +@pytest.mark.asyncio +async def test_submit_accepts_seven_references(): + client = MagicMock() + client._client_context.return_value = {"projectId": "p"} + client._build_url.return_value = "https://example.test/omni" + client._send = AsyncMock(return_value={"status": 200, "data": {}}) + refs = [f"ref-{i}" for i in range(OMNI_FLASH_MAX_REFERENCE_IMAGES)] + + with patch("agent.services.omni_flash.get_flow_client", return_value=client): + await generate_omni_flash_video( + reference_media_ids=refs, + prompt="test", + project_id="p", + duration_s=4, + ) + + request = client._send.await_args.args[1]["body"]["requests"][0] + assert len(request["referenceImages"]) == 7 + + +@pytest.mark.asyncio +async def test_submit_rejects_more_than_seven_references(): + refs = [f"ref-{i}" for i in range(OMNI_FLASH_MAX_REFERENCE_IMAGES + 1)] + + with pytest.raises(ValueError, match="at most 7 reference images"): + await generate_omni_flash_video( + reference_media_ids=refs, + prompt="test", + project_id="p", + duration_s=8, + ) + + +@pytest.mark.asyncio +async def test_submit_rejects_empty_reference_set(): + with pytest.raises(ValueError, match="requires at least one reference image"): + await generate_omni_flash_video( + reference_media_ids=[], + prompt="test", + project_id="p", + duration_s=8, + )