From 02bfa603e97418717412a9cee6aa558ba9c882af Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:48:28 +0300 Subject: [PATCH 01/18] feat(flow): add Omni Flash model keys --- agent/models.json | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/agent/models.json b/agent/models.json index 631b43a41..2e79bbb04 100644 --- a/agent/models.json +++ b/agent/models.json @@ -29,6 +29,14 @@ } } }, + "omni_flash_models": { + "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" From 3d4b4acb63b314297580b1fd6ddc603d54e47432 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:48:59 +0300 Subject: [PATCH 02/18] feat(flow): add Omni Flash bridge client --- agent/services/omni_flash.py | 135 +++++++++++++++++++++++++++++++++++ 1 file changed, 135 insertions(+) create mode 100644 agent/services/omni_flash.py diff --git a/agent/services/omni_flash.py b/agent/services/omni_flash.py new file mode 100644 index 000000000..2be61b87a --- /dev/null +++ b/agent/services/omni_flash.py @@ -0,0 +1,135 @@ +"""Gemini Omni Flash video generation through the Google Flow bridge. + +Omni Flash is a separate video family from Veo. Google Flow currently +submits reference-conditioned Omni jobs through the same +``batchAsyncGenerateVideoReferenceImages`` route used by R2V, but selects an +``abra_r2v_s`` model key. Keep the model keys in ``models.json`` so +they can be updated without changing this client. +""" + +from __future__ import annotations + +import json +import time +import uuid +from pathlib import Path + +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} + + +def _load_model_key(duration_s: int) -> str: + """Resolve the configured Omni Flash R2V key for a duration.""" + 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)}" + ) + + with open(_MODELS_FILE, encoding="utf-8") as f: + models = json.load(f) + + key = ( + models.get("omni_flash_models", {}) + .get("reference_to_video", {}) + .get(str(duration_s)) + ) + if not key: + raise ValueError(f"No Omni Flash model key configured for {duration_s}s") + return key + + +def _validate_inputs(reference_media_ids: list[str], duration_s: int, aspect_ratio: str) -> list[str]: + 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" + ) + + # Preserve caller ordering while filtering accidental empty values. + 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" + ) + + # Validate duration before a request reaches the browser/credits path. + 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)}" + ) + return refs + + +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. + + The return value intentionally matches FlowClient's existing video-submit + envelope, so callers can reuse ``/flow/check-status`` and the normal Flow + operation polling path. + """ + refs = _validate_inputs(reference_media_ids, duration_s, aspect_ratio) + model_key = _load_model_key(duration_s) + 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()), + # Match current Google Flow Omni submissions: a silent-audio output + # is treated as a failed generation rather than a degraded success. + "audioFailurePreference": "BLOCK_SILENCED_VIDEOS", + }, + "clientContext": {**context, "sessionId": f";{ts}"}, + "requests": [request_item], + "useV2ModelConfig": True, + } + + url = client._build_url("generate_video_references") + return await client._send( + "api_request", + { + "url": url, + "method": "POST", + "headers": random_headers(), + "body": body, + "captchaAction": "VIDEO_GENERATION", + }, + timeout=60, + ) From 4a32523bc3340d5157e55f97bde3359bea34f7bf Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:50:07 +0300 Subject: [PATCH 03/18] feat(api): expose Omni Flash video generation --- agent/api/flow.py | 61 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 57 insertions(+), 4 deletions(-) diff --git a/agent/api/flow.py b/agent/api/flow.py index 1373e113b..3ebf165bb 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -1,8 +1,10 @@ """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 generate_omni_flash_video router = APIRouter(prefix="/flow", tags=["flow"]) @@ -32,6 +34,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): @@ -95,7 +111,7 @@ async def generate_image(body: GenerateImageRequest): @router.post("/generate-video") async def generate_video(body: GenerateVideoRequest): - """Submit video generation (returns operations for polling).""" + """Submit Veo image-to-video generation (returns operations for polling).""" client = get_flow_client() if not client.connected: raise HTTPException(503, "Extension not connected") @@ -107,11 +123,48 @@ 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. + """ client = get_flow_client() if not client.connected: raise HTTPException(503, "Extension not connected") - result = await client.generate_video_from_references(**body.model_dump()) + + 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 R2V generation (returns operations for polling).""" + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + 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) From 9266a629ea94a59e9474dd346ef272688df5c346 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:50:23 +0300 Subject: [PATCH 04/18] feat(models): make Omni Flash keys configurable --- agent/api/models.py | 35 ++++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) 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] = {} From a6bd259aff3f16510a311bf5ecb698ab4945060c Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:51:18 +0300 Subject: [PATCH 05/18] test(flow): cover Omni Flash submissions --- tests/unit/test_omni_flash.py | 140 ++++++++++++++++++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 tests/unit/test_omni_flash.py diff --git a/tests/unit/test_omni_flash.py b/tests/unit/test_omni_flash.py new file mode 100644 index 000000000..53431f9ba --- /dev/null +++ b/tests/unit/test_omni_flash.py @@ -0,0 +1,140 @@ +"""Unit tests for Gemini Omni Flash Flow submissions.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from agent.services.omni_flash import ( + OMNI_FLASH_MAX_REFERENCE_IMAGES, + _load_model_key, + 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_duration_model_keys(duration, expected): + assert _load_model_key(duration) == expected + + +def test_invalid_duration_fails_before_submit(): + with pytest.raises(ValueError, match="duration 5s is unsupported"): + _load_model_key(5) + + +@pytest.mark.asyncio +async def test_submit_builds_flow_omni_r2v_request(): + 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.return_value = ( + "https://aisandbox-pa.googleapis.com/" + "v1/video:batchAsyncGenerateVideoReferenceImages?key=test" + ) + client._send = AsyncMock( + return_value={ + "status": 200, + "data": { + "operations": [ + { + "operation": {"name": "op-1"}, + "status": "MEDIA_GENERATION_STATUS_PENDING", + } + ] + }, + } + ) + + 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 + 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"}, + ] + + +@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_first_last_style_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, + ) From e69d1507162e43797be758a0c891e493cc3d1577 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:52:24 +0300 Subject: [PATCH 06/18] docs(flow): document Omni Flash REST usage --- docs/OMNI_FLASH.md | 68 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 docs/OMNI_FLASH.md diff --git a/docs/OMNI_FLASH.md b/docs/OMNI_FLASH.md new file mode 100644 index 000000000..16260414f --- /dev/null +++ b/docs/OMNI_FLASH.md @@ -0,0 +1,68 @@ +# Gemini Omni Flash + +FlowKit can submit Gemini Omni Flash reference-to-video jobs through the same authenticated Chrome extension bridge used by the existing Google Flow integrations. + +## Supported in this integration + +- Reference-to-video with 1–7 reference images +- Portrait (`9:16`) and landscape (`16:9`) +- 4, 6, 8, and 10 second generations +- Existing `/flow/check-status` polling +- Existing Google Flow account / paygate tier handling + +First + last frame interpolation is intentionally not enabled for Omni Flash here. + +## REST API + +### Dedicated endpoint + +```bash +curl -X POST http://127.0.0.1:8100/flow/generate-video-omni \ + -H 'Content-Type: application/json' \ + -d '{ + "reference_media_ids": [""], + "prompt": "Cinematic handheld shot, natural motion and dialogue", + "project_id": "", + "duration_s": 10, + "aspect_ratio": "VIDEO_ASPECT_RATIO_LANDSCAPE", + "user_paygate_tier": "PAYGATE_TIER_ONE" + }' +``` + +### Backward-compatible reference endpoint + +Existing Veo callers do not need to change. To opt into Omni Flash, add `model_family` and `duration_s`: + +```json +{ + "reference_media_ids": [""], + "prompt": "Natural cinematic motion", + "project_id": "", + "scene_id": "scene-1", + "model_family": "omni_flash", + "duration_s": 8, + "aspect_ratio": "VIDEO_ASPECT_RATIO_PORTRAIT", + "user_paygate_tier": "PAYGATE_TIER_ONE" +} +``` + +Submit this body to `POST /flow/generate-video-refs`. + +## Model keys + +The duration-to-model mapping lives in `agent/models.json` under `omni_flash_models.reference_to_video`: + +```json +{ + "4": "abra_r2v_4s", + "6": "abra_r2v_6s", + "8": "abra_r2v_8s", + "10": "abra_r2v_10s" +} +``` + +The values can be updated through `PATCH /api/models` if Google rotates Flow's internal model keys. + +## Notes + +Omni Flash is an unofficial Google Flow integration and the underlying internal model keys/endpoints can change as Flow evolves. The request shape mirrors the Omni Flash implementation used by `crisng95/flowboard`: `batchAsyncGenerateVideoReferenceImages`, asset-typed `referenceImages`, V2 model config, and `BLOCK_SILENCED_VIDEOS` audio failure preference. From 4132481784d2af80a960d9ad89e680b95bed8155 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:30:03 +0300 Subject: [PATCH 07/18] fix(omni): poll workflow media instead of legacy operations --- agent/services/omni_flash.py | 196 ++++++++++++++++++++++++++++++++--- 1 file changed, 183 insertions(+), 13 deletions(-) diff --git a/agent/services/omni_flash.py b/agent/services/omni_flash.py index 2be61b87a..d71bd651c 100644 --- a/agent/services/omni_flash.py +++ b/agent/services/omni_flash.py @@ -1,14 +1,19 @@ """Gemini Omni Flash video generation through the Google Flow bridge. -Omni Flash is a separate video family from Veo. Google Flow currently -submits reference-conditioned Omni jobs through the same -``batchAsyncGenerateVideoReferenceImages`` route used by R2V, but selects an -``abra_r2v_s`` model key. Keep the model keys in ``models.json`` so -they can be updated without changing this client. +Omni Flash is a separate video family from Veo. Google Flow currently +submits reference-conditioned Omni jobs through +``batchAsyncGenerateVideoReferenceImages`` using ``abra_r2v_s`` +model keys. + +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 must be polled through ``/v1/media/``. """ from __future__ import annotations +import base64 import json import time import uuid @@ -57,7 +62,6 @@ def _validate_inputs(reference_media_ids: list[str], duration_s: int, aspect_rat "use VIDEO_ASPECT_RATIO_PORTRAIT or VIDEO_ASPECT_RATIO_LANDSCAPE" ) - # Preserve caller ordering while filtering accidental empty values. 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") @@ -66,7 +70,6 @@ def _validate_inputs(reference_media_ids: list[str], duration_s: int, aspect_rat f"Omni Flash accepts at most {OMNI_FLASH_MAX_REFERENCE_IMAGES} reference images" ) - # Validate duration before a request reaches the browser/credits path. if duration_s not in OMNI_FLASH_VALID_DURATIONS: raise ValueError( f"Omni Flash duration {duration_s}s is unsupported; " @@ -75,6 +78,62 @@ def _validate_inputs(reference_media_ids: list[str], duration_s: int, aspect_rat return refs +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 + return {"name": name, "primary_media_id": primary_media_id} + + +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) -> 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): + data["flowkitPolling"] = { + "mode": "workflow_media", + "workflows": workflows, + } + return result + + +def _is_complete_mp4(encoded: str) -> bool: + """Flow can expose a small placeholder payload before the final MP4.""" + if not isinstance(encoded, str) or not encoded: + return False + try: + binary = base64.b64decode(encoded, validate=False) + except Exception: + return False + return len(binary) >= 12 and binary[4:8] == b"ftyp" + + async def generate_omni_flash_video( reference_media_ids: list[str], prompt: str, @@ -87,9 +146,9 @@ async def generate_omni_flash_video( ) -> dict: """Submit an Omni Flash reference-to-video generation. - The return value intentionally matches FlowClient's existing video-submit - envelope, so callers can reuse ``/flow/check-status`` and the normal Flow - operation polling path. + 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_inputs(reference_media_ids, duration_s, aspect_ratio) model_key = _load_model_key(duration_s) @@ -112,8 +171,6 @@ async def generate_omni_flash_video( body = { "mediaGenerationContext": { "batchId": str(uuid.uuid4()), - # Match current Google Flow Omni submissions: a silent-audio output - # is treated as a failed generation rather than a degraded success. "audioFailurePreference": "BLOCK_SILENCED_VIDEOS", }, "clientContext": {**context, "sessionId": f";{ts}"}, @@ -122,7 +179,7 @@ async def generate_omni_flash_video( } url = client._build_url("generate_video_references") - return await client._send( + result = await client._send( "api_request", { "url": url, @@ -133,3 +190,116 @@ async def generate_omni_flash_video( }, timeout=60, ) + return _annotate_polling(result) + + +async def check_omni_flash_status( + workflows: list[dict], + include_encoded_video: bool = False, +) -> dict: + """Perform one non-blocking poll pass for Omni workflow-backed jobs. + + Each workflow is polled via ``GET /v1/media/``. A completed + result is only reported once ``video.encodedVideo`` decodes to an MP4 + (``ftyp`` box present), matching Flow's current workflow behavior. + """ + 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)" + ) + + client = get_flow_client() + items = [] + + for workflow in normalized: + name = workflow["name"] + media_id = workflow["primary_media_id"] + response = await client.get_media(media_id) + + http_status = response.get("status") if isinstance(response, dict) else None + if isinstance(http_status, int) and http_status >= 400: + if http_status == 404: + items.append({ + "name": name, + "primary_media_id": media_id, + "done": False, + "status": "PENDING", + "error": None, + }) + continue + data = response.get("data") if isinstance(response.get("data"), dict) else {} + error = data.get("error") if isinstance(data, dict) else None + error_message = None + if isinstance(error, dict): + error_message = error.get("message") or error.get("status") + items.append({ + "name": name, + "primary_media_id": media_id, + "done": True, + "status": "FAILED", + "error": error_message or response.get("error") or f"API_{http_status}", + }) + continue + + payload = response.get("data") if isinstance(response, dict) and isinstance(response.get("data"), dict) else response + payload = payload if isinstance(payload, dict) else {} + video = payload.get("video") if isinstance(payload.get("video"), dict) else {} + encoded = video.get("encodedVideo") if isinstance(video, dict) else None + + media_status = payload.get("mediaMetadata") if isinstance(payload.get("mediaMetadata"), dict) else {} + media_status = media_status.get("mediaStatus") if isinstance(media_status.get("mediaStatus"), dict) else {} + generation_status = media_status.get("mediaGenerationStatus") if isinstance(media_status, dict) else None + + if isinstance(generation_status, str) and generation_status.endswith("FAILED"): + items.append({ + "name": name, + "primary_media_id": media_id, + "done": True, + "status": "FAILED", + "error": generation_status, + }) + continue + + if not _is_complete_mp4(encoded): + items.append({ + "name": name, + "primary_media_id": media_id, + "done": False, + "status": "PENDING", + "error": None, + }) + continue + + url = None + if isinstance(video, dict): + url = video.get("fifeUrl") or video.get("servingUri") + url = url or payload.get("fifeUrl") or payload.get("servingUri") + item = { + "name": name, + "primary_media_id": media_id, + "done": True, + "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", + "error": None, + "media": { + "media_id": media_id, + "url": url if isinstance(url, str) else None, + "encoded_video_available": True, + }, + } + if include_encoded_video: + item["media"]["encoded_video"] = encoded + 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 { + "done": all_done, + "status": "FAILED" if any_failed else ("COMPLETED" if all_done else "PENDING"), + "workflows": items, + } From 59bf46d6cfc8d5c573ddb53cfd392ae1f6762110 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:30:45 +0300 Subject: [PATCH 08/18] fix(api): add workflow polling for Omni Flash --- agent/api/flow.py | 65 ++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 59 insertions(+), 6 deletions(-) diff --git a/agent/api/flow.py b/agent/api/flow.py index 3ebf165bb..af501690a 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -4,7 +4,10 @@ from typing import Literal, Optional from agent.services.flow_client import get_flow_client -from agent.services.omni_flash import generate_omni_flash_video +from agent.services.omni_flash import ( + check_omni_flash_status, + generate_omni_flash_video, +) router = APIRouter(prefix="/flow", tags=["flow"]) @@ -64,7 +67,17 @@ 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 + # /v1/media/ polling. + workflows: Optional[list[dict]] = None + include_encoded_video: bool = False + + +class CheckOmniStatusRequest(BaseModel): + workflows: list[dict] + include_encoded_video: bool = False class EditImageRequest(BaseModel): @@ -127,6 +140,8 @@ async def generate_video_refs(body: GenerateVideoRefsRequest): 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: @@ -157,7 +172,11 @@ async def generate_video_refs(body: GenerateVideoRefsRequest): @router.post("/generate-video-omni") async def generate_video_omni(body: GenerateOmniFlashVideoRequest): - """Submit Gemini Omni Flash R2V generation (returns operations for polling).""" + """Submit Gemini Omni Flash R2V 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") @@ -184,16 +203,50 @@ 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, + ) + except ValueError as exc: + raise HTTPException(400, 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, + ) + except ValueError as exc: + raise HTTPException(400, 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.""" @@ -210,8 +263,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: From 1885470930e323d16dd59f0de63f7f18a6832e69 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:31:09 +0300 Subject: [PATCH 09/18] docs(omni): document workflow media polling --- docs/OMNI_FLASH.md | 64 +++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 60 insertions(+), 4 deletions(-) diff --git a/docs/OMNI_FLASH.md b/docs/OMNI_FLASH.md index 16260414f..bf3f6c5ce 100644 --- a/docs/OMNI_FLASH.md +++ b/docs/OMNI_FLASH.md @@ -7,14 +7,36 @@ FlowKit can submit Gemini Omni Flash reference-to-video jobs through the same au - Reference-to-video with 1–7 reference images - Portrait (`9:16`) and landscape (`16:9`) - 4, 6, 8, and 10 second generations -- Existing `/flow/check-status` polling +- Workflow/media polling via each submit response's `primaryMediaId` - Existing Google Flow account / paygate tier handling First + last frame interpolation is intentionally not enabled for Omni Flash here. +## Important: Omni does not use legacy Veo polling + +Do **not** take the operation-looking handle returned by an Omni submit and send it to the legacy `batchCheckAsyncVideoGenerationStatus` path. Google Flow can reject that with `INVALID_ARGUMENT`. + +Omni Flash submits are workflow-backed. The submit response includes `workflows[].metadata.primaryMediaId`; FlowKit normalizes those into: + +```json +{ + "flowkitPolling": { + "mode": "workflow_media", + "workflows": [ + { + "name": "", + "primary_media_id": "" + } + ] + } +} +``` + +Poll those workflow descriptors through `POST /flow/check-omni-status`, or pass them as `workflows` to the generic `POST /flow/check-status` endpoint. + ## REST API -### Dedicated endpoint +### Submit Omni Flash ```bash curl -X POST http://127.0.0.1:8100/flow/generate-video-omni \ @@ -29,6 +51,40 @@ curl -X POST http://127.0.0.1:8100/flow/generate-video-omni \ }' ``` +Keep `flowkitPolling.workflows` from the response. + +### Poll Omni status + +```bash +curl -X POST http://127.0.0.1:8100/flow/check-omni-status \ + -H 'Content-Type: application/json' \ + -d '{ + "workflows": [ + { + "name": "", + "primary_media_id": "" + } + ] + }' +``` + +Equivalent generic form: + +```bash +curl -X POST http://127.0.0.1:8100/flow/check-status \ + -H 'Content-Type: application/json' \ + -d '{ + "workflows": [ + { + "name": "", + "primary_media_id": "" + } + ] + }' +``` + +While rendering, FlowKit returns `status: "PENDING"`. Once `/v1/media/` exposes a complete MP4 payload, it returns `MEDIA_GENERATION_STATUS_SUCCESSFUL` and the media ID. Set `include_encoded_video: true` in the poll request only if you actually need the base64 MP4 inline; by default FlowKit avoids returning that large payload. + ### Backward-compatible reference endpoint Existing Veo callers do not need to change. To opt into Omni Flash, add `model_family` and `duration_s`: @@ -46,7 +102,7 @@ Existing Veo callers do not need to change. To opt into Omni Flash, add `model_f } ``` -Submit this body to `POST /flow/generate-video-refs`. +Submit this body to `POST /flow/generate-video-refs`. Omni responses from this endpoint also include `flowkitPolling.workflows`. ## Model keys @@ -65,4 +121,4 @@ The values can be updated through `PATCH /api/models` if Google rotates Flow's i ## Notes -Omni Flash is an unofficial Google Flow integration and the underlying internal model keys/endpoints can change as Flow evolves. The request shape mirrors the Omni Flash implementation used by `crisng95/flowboard`: `batchAsyncGenerateVideoReferenceImages`, asset-typed `referenceImages`, V2 model config, and `BLOCK_SILENCED_VIDEOS` audio failure preference. +Omni Flash is an unofficial Google Flow integration and the underlying internal model keys/endpoints can change as Flow evolves. The submit and polling shapes mirror the workflow-backed Omni implementation used by `crisng95/flowboard`: `batchAsyncGenerateVideoReferenceImages` for submit, then `/v1/media/` for completion polling. From ff55ba0b0279058449f68508fb74fc199a5f0139 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:31:38 +0300 Subject: [PATCH 10/18] test(omni): cover workflow media polling --- tests/unit/test_omni_flash.py | 133 +++++++++++++++++++++++++++++++++- 1 file changed, 130 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_omni_flash.py b/tests/unit/test_omni_flash.py index 53431f9ba..2ad31d4ca 100644 --- a/tests/unit/test_omni_flash.py +++ b/tests/unit/test_omni_flash.py @@ -1,5 +1,6 @@ -"""Unit tests for Gemini Omni Flash Flow submissions.""" +"""Unit tests for Gemini Omni Flash Flow submissions and workflow polling.""" +import base64 from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -7,6 +8,8 @@ from agent.services.omni_flash import ( OMNI_FLASH_MAX_REFERENCE_IMAGES, _load_model_key, + check_omni_flash_status, + extract_omni_workflows, generate_omni_flash_video, ) @@ -29,8 +32,31 @@ def test_invalid_duration_fails_before_submit(): _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"} + ] + + @pytest.mark.asyncio -async def test_submit_builds_flow_omni_r2v_request(): +async def test_submit_builds_flow_omni_r2v_request_and_poll_descriptor(): client = MagicMock() client._client_context.return_value = { "projectId": "project-1", @@ -55,7 +81,13 @@ async def test_submit_builds_flow_omni_r2v_request(): "operation": {"name": "op-1"}, "status": "MEDIA_GENERATION_STATUS_PENDING", } - ] + ], + "workflows": [ + { + "name": "workflow-1", + "metadata": {"primaryMediaId": "media-1"}, + } + ], }, } ) @@ -73,6 +105,12 @@ async def test_submit_builds_flow_omni_r2v_request(): ) assert result["status"] == 200 + assert result["data"]["flowkitPolling"] == { + "mode": "workflow_media", + "workflows": [ + {"name": "workflow-1", "primary_media_id": "media-1"} + ], + } client._build_url.assert_called_once_with("generate_video_references") client._send.assert_awaited_once() @@ -96,6 +134,95 @@ async def test_submit_builds_flow_omni_r2v_request(): ] +@pytest.mark.asyncio +async def test_omni_poll_pending_uses_media_endpoint_not_legacy_operation_poll(): + client = MagicMock() + client.get_media = AsyncMock( + return_value={ + "status": 200, + "data": { + "name": "media-1", + "video": {"encodedVideo": ""}, + }, + } + ) + client.check_video_status = AsyncMock(side_effect=AssertionError("legacy poll must not be used")) + + 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"}] + ) + + client.get_media.assert_awaited_once_with("media-1") + client.check_video_status.assert_not_awaited() + assert result["done"] is False + assert result["status"] == "PENDING" + assert result["workflows"][0]["status"] == "PENDING" + + +@pytest.mark.asyncio +async def test_omni_poll_completed_detects_mp4_without_returning_base64_by_default(): + mp4 = b"\x00\x00\x00\x18ftypisom" + b"\x00" * 32 + encoded = base64.b64encode(mp4).decode() + client = MagicMock() + client.get_media = AsyncMock( + return_value={ + "status": 200, + "data": { + "name": "media-1", + "fifeUrl": "https://example.test/video.mp4", + "video": {"encodedVideo": encoded}, + }, + } + ) + + 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"}] + ) + + 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"] == "https://example.test/video.mp4" + assert item["media"]["encoded_video_available"] is True + assert "encoded_video" not in item["media"] + + +@pytest.mark.asyncio +async def test_omni_poll_can_return_encoded_video_when_requested(): + mp4 = b"\x00\x00\x00\x18ftypisom" + b"\x00" * 32 + encoded = base64.b64encode(mp4).decode() + client = MagicMock() + client.get_media = AsyncMock( + return_value={"status": 200, "data": {"video": {"encodedVideo": encoded}}} + ) + + with patch("agent.services.omni_flash.get_flow_client", return_value=client): + result = await check_omni_flash_status( + [{"name": "workflow-1", "primary_media_id": "media-1"}], + include_encoded_video=True, + ) + + assert result["workflows"][0]["media"]["encoded_video"] == encoded + + +@pytest.mark.asyncio +async def test_omni_poll_404_is_treated_as_not_ready(): + client = MagicMock() + client.get_media = AsyncMock(return_value={"status": 404, "data": {}}) + + 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"}] + ) + + assert result["done"] is False + assert result["workflows"][0]["status"] == "PENDING" + + @pytest.mark.asyncio async def test_submit_accepts_seven_references(): client = MagicMock() From bf69051273a547920a06c715d23749ba45404b36 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:39:33 +0300 Subject: [PATCH 11/18] feat(flow): add Omni Flash frame model keys --- agent/models.json | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/agent/models.json b/agent/models.json index 2e79bbb04..c44c789b5 100644 --- a/agent/models.json +++ b/agent/models.json @@ -30,6 +30,18 @@ } }, "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", From 637b370ed818d289a4e0c7e9a68359e34c859c47 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:40:26 +0300 Subject: [PATCH 12/18] feat(flow): add Omni Flash first and last frame generation --- agent/services/omni_flash.py | 192 +++++++++++++++++++++++++++++++---- 1 file changed, 170 insertions(+), 22 deletions(-) diff --git a/agent/services/omni_flash.py b/agent/services/omni_flash.py index d71bd651c..6dc27c21e 100644 --- a/agent/services/omni_flash.py +++ b/agent/services/omni_flash.py @@ -1,9 +1,16 @@ """Gemini Omni Flash video generation through the Google Flow bridge. -Omni Flash is a separate video family from Veo. Google Flow currently -submits reference-conditioned Omni jobs through -``batchAsyncGenerateVideoReferenceImages`` using ``abra_r2v_s`` -model keys. +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 @@ -34,33 +41,48 @@ OMNI_FLASH_CREDIT_COST = {4: 15, 6: 20, 8: 25, 10: 30} -def _load_model_key(duration_s: int) -> str: - """Resolve the configured Omni Flash R2V key for a duration.""" +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("reference_to_video", {}) + .get(mode, {}) .get(str(duration_s)) ) if not key: - raise ValueError(f"No Omni Flash model key configured for {duration_s}s") + raise ValueError( + f"No Omni Flash model key configured for mode {mode!r}, {duration_s}s" + ) return key -def _validate_inputs(reference_media_ids: list[str], duration_s: int, aspect_ratio: str) -> list[str]: - 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 _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: @@ -69,15 +91,25 @@ def _validate_inputs(reference_media_ids: list[str], duration_s: int, aspect_rat raise ValueError( f"Omni Flash accepts at most {OMNI_FLASH_MAX_REFERENCE_IMAGES} reference images" ) - - 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)}" - ) 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): @@ -134,6 +166,122 @@ def _is_complete_mp4(encoded: str) -> bool: return len(binary) >= 12 and binary[4:8] == b"ftyp" +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) + + +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, @@ -150,8 +298,8 @@ async def generate_omni_flash_video( 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_inputs(reference_media_ids, duration_s, aspect_ratio) - model_key = _load_model_key(duration_s) + 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) From fd7834b10c5a095e0c9c955c46df016764ba8027 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:41:21 +0300 Subject: [PATCH 13/18] feat(api): expose Omni Flash first and last frames --- agent/api/flow.py | 45 ++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/agent/api/flow.py b/agent/api/flow.py index af501690a..d0b14df8f 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -6,6 +6,8 @@ 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, ) @@ -28,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): @@ -124,11 +129,45 @@ async def generate_image(body: GenerateImageRequest): @router.post("/generate-video") async def generate_video(body: GenerateVideoRequest): - """Submit Veo image-to-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) @@ -172,7 +211,7 @@ async def generate_video_refs(body: GenerateVideoRefsRequest): @router.post("/generate-video-omni") async def generate_video_omni(body: GenerateOmniFlashVideoRequest): - """Submit Gemini Omni Flash R2V generation. + """Submit Gemini Omni Flash reference-to-video generation. The response includes ``flowkitPolling.workflows`` for the correct workflow/media polling path. From 45f08e32049344dde45ba272c4724745b6568304 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:42:10 +0300 Subject: [PATCH 14/18] test(flow): cover Omni first and last frame requests --- tests/unit/test_omni_flash.py | 131 +++++++++++++++++++++++++++++++--- 1 file changed, 123 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_omni_flash.py b/tests/unit/test_omni_flash.py index 2ad31d4ca..858ec47e2 100644 --- a/tests/unit/test_omni_flash.py +++ b/tests/unit/test_omni_flash.py @@ -10,6 +10,8 @@ _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, ) @@ -23,10 +25,36 @@ (10, "abra_r2v_10s"), ], ) -def test_omni_duration_model_keys(duration, expected): +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) @@ -55,8 +83,7 @@ def test_extract_omni_workflows_uses_primary_media_id(): ] -@pytest.mark.asyncio -async def test_submit_builds_flow_omni_r2v_request_and_poll_descriptor(): +def _mock_submit_client(): client = MagicMock() client._client_context.return_value = { "projectId": "project-1", @@ -68,10 +95,7 @@ async def test_submit_builds_flow_omni_r2v_request_and_poll_descriptor(): }, "sessionId": ";old", } - client._build_url.return_value = ( - "https://aisandbox-pa.googleapis.com/" - "v1/video:batchAsyncGenerateVideoReferenceImages?key=test" - ) + client._build_url.side_effect = lambda name: f"https://example.test/{name}" client._send = AsyncMock( return_value={ "status": 200, @@ -91,6 +115,97 @@ async def test_submit_builds_flow_omni_r2v_request_and_poll_descriptor(): }, } ) + 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"] == "workflow_media" + + +@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"} + ] + + +@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( @@ -257,7 +372,7 @@ async def test_submit_rejects_more_than_seven_references(): @pytest.mark.asyncio -async def test_submit_rejects_first_last_style_empty_reference_set(): +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=[], From 28863b17b8d0ab0170e9d35910362b5bf1177a34 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:42:47 +0300 Subject: [PATCH 15/18] docs(flow): document Omni first and last frame modes --- docs/OMNI_FLASH.md | 137 +++++++++++++++++++++++++++------------------ 1 file changed, 82 insertions(+), 55 deletions(-) diff --git a/docs/OMNI_FLASH.md b/docs/OMNI_FLASH.md index bf3f6c5ce..fc39380a4 100644 --- a/docs/OMNI_FLASH.md +++ b/docs/OMNI_FLASH.md @@ -1,42 +1,67 @@ # Gemini Omni Flash -FlowKit can submit Gemini Omni Flash reference-to-video jobs through the same authenticated Chrome extension bridge used by the existing Google Flow integrations. +FlowKit can submit Gemini Omni Flash video jobs through the same authenticated Chrome extension bridge used by the existing Google Flow integrations. ## Supported in this integration +- First frame → video +- First + Last frame → video - Reference-to-video with 1–7 reference images - Portrait (`9:16`) and landscape (`16:9`) - 4, 6, 8, and 10 second generations - Workflow/media polling via each submit response's `primaryMediaId` - Existing Google Flow account / paygate tier handling -First + last frame interpolation is intentionally not enabled for Omni Flash here. +## First frame / First + Last -## Important: Omni does not use legacy Veo polling +The existing `POST /flow/generate-video` endpoint remains backward compatible: requests use Veo unless `model_family` is explicitly set to `omni_flash`. -Do **not** take the operation-looking handle returned by an Omni submit and send it to the legacy `batchCheckAsyncVideoGenerationStatus` path. Google Flow can reject that with `INVALID_ARGUMENT`. +### First frame -Omni Flash submits are workflow-backed. The submit response includes `workflows[].metadata.primaryMediaId`; FlowKit normalizes those into: +```bash +curl -X POST http://127.0.0.1:8100/flow/generate-video \ + -H 'Content-Type: application/json' \ + -d '{ + "start_image_media_id": "", + "prompt": "Natural cinematic motion, keep the subject consistent", + "project_id": "", + "scene_id": "scene-1", + "model_family": "omni_flash", + "duration_s": 8, + "aspect_ratio": "VIDEO_ASPECT_RATIO_PORTRAIT", + "user_paygate_tier": "PAYGATE_TIER_ONE" + }' +``` -```json -{ - "flowkitPolling": { - "mode": "workflow_media", - "workflows": [ - { - "name": "", - "primary_media_id": "" - } - ] - } -} +This submits through `batchAsyncGenerateVideoStartImage` with the configured `abra_i2v_s` key. + +### First + Last frame + +Add `end_image_media_id` to the same request: + +```bash +curl -X POST http://127.0.0.1:8100/flow/generate-video \ + -H 'Content-Type: application/json' \ + -d '{ + "start_image_media_id": "", + "end_image_media_id": "", + "prompt": "Move naturally from the first composition to the final composition", + "project_id": "", + "scene_id": "scene-1", + "model_family": "omni_flash", + "duration_s": 8, + "aspect_ratio": "VIDEO_ASPECT_RATIO_PORTRAIT", + "user_paygate_tier": "PAYGATE_TIER_ONE" + }' ``` -Poll those workflow descriptors through `POST /flow/check-omni-status`, or pass them as `workflows` to the generic `POST /flow/check-status` endpoint. +This submits through `batchAsyncGenerateVideoStartAndEndImage` and includes both `startImage` and `endImage` in the Flow request. -## REST API +The First+Last model mapping is stored independently in `models.json`. The current default uses the same `abra_i2v_s` family as Omni First-frame generation. Keeping it as a separate mapping allows a one-line model-config update if Google changes the key during rollout. -### Submit Omni Flash +## Reference-to-video + +### Dedicated endpoint ```bash curl -X POST http://127.0.0.1:8100/flow/generate-video-omni \ @@ -51,27 +76,32 @@ curl -X POST http://127.0.0.1:8100/flow/generate-video-omni \ }' ``` -Keep `flowkitPolling.workflows` from the response. +Or use `POST /flow/generate-video-refs` with `model_family: "omni_flash"` and `duration_s`. -### Poll Omni status +## Important: Omni does not use legacy Veo polling -```bash -curl -X POST http://127.0.0.1:8100/flow/check-omni-status \ - -H 'Content-Type: application/json' \ - -d '{ +Do **not** take the operation-looking handle returned by an Omni submit and send it to the legacy `batchCheckAsyncVideoGenerationStatus` path. Google Flow can reject that with `INVALID_ARGUMENT`. + +Omni Flash submits are workflow-backed. The submit response includes `workflows[].metadata.primaryMediaId`; FlowKit normalizes those into: + +```json +{ + "flowkitPolling": { + "mode": "workflow_media", "workflows": [ { "name": "", "primary_media_id": "" } ] - }' + } +} ``` -Equivalent generic form: +Poll those workflow descriptors through `POST /flow/check-omni-status`, or pass them as `workflows` to the generic `POST /flow/check-status` endpoint. ```bash -curl -X POST http://127.0.0.1:8100/flow/check-status \ +curl -X POST http://127.0.0.1:8100/flow/check-omni-status \ -H 'Content-Type: application/json' \ -d '{ "workflows": [ @@ -83,37 +113,34 @@ curl -X POST http://127.0.0.1:8100/flow/check-status \ }' ``` -While rendering, FlowKit returns `status: "PENDING"`. Once `/v1/media/` exposes a complete MP4 payload, it returns `MEDIA_GENERATION_STATUS_SUCCESSFUL` and the media ID. Set `include_encoded_video: true` in the poll request only if you actually need the base64 MP4 inline; by default FlowKit avoids returning that large payload. - -### Backward-compatible reference endpoint - -Existing Veo callers do not need to change. To opt into Omni Flash, add `model_family` and `duration_s`: - -```json -{ - "reference_media_ids": [""], - "prompt": "Natural cinematic motion", - "project_id": "", - "scene_id": "scene-1", - "model_family": "omni_flash", - "duration_s": 8, - "aspect_ratio": "VIDEO_ASPECT_RATIO_PORTRAIT", - "user_paygate_tier": "PAYGATE_TIER_ONE" -} -``` - -Submit this body to `POST /flow/generate-video-refs`. Omni responses from this endpoint also include `flowkitPolling.workflows`. +While rendering, FlowKit returns `status: "PENDING"`. Once `/v1/media/` exposes a complete MP4 payload, it returns `MEDIA_GENERATION_STATUS_SUCCESSFUL` and the media ID. Set `include_encoded_video: true` only if you actually need the base64 MP4 inline; by default FlowKit avoids returning that large payload. ## Model keys -The duration-to-model mapping lives in `agent/models.json` under `omni_flash_models.reference_to_video`: +The duration mappings live in `agent/models.json`: ```json { - "4": "abra_r2v_4s", - "6": "abra_r2v_6s", - "8": "abra_r2v_8s", - "10": "abra_r2v_10s" + "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" + } + } } ``` @@ -121,4 +148,4 @@ The values can be updated through `PATCH /api/models` if Google rotates Flow's i ## Notes -Omni Flash is an unofficial Google Flow integration and the underlying internal model keys/endpoints can change as Flow evolves. The submit and polling shapes mirror the workflow-backed Omni implementation used by `crisng95/flowboard`: `batchAsyncGenerateVideoReferenceImages` for submit, then `/v1/media/` for completion polling. +Omni Flash is an unofficial Google Flow integration and its internal model keys/endpoints can change as Flow evolves. First-frame generation uses the wire-proven `abra_i2v_s` family and `batchAsyncGenerateVideoStartImage`. First+Last uses the same configurable Omni I2V family with `batchAsyncGenerateVideoStartAndEndImage`; this matches the current Flow UI capability while keeping the key isolated so it can be changed without another code release. Reference generation uses `abra_r2v_s` with `batchAsyncGenerateVideoReferenceImages`. From 6b2a62bd45cbdc48642edca70460fb4ca95558ac Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:51:50 +0300 Subject: [PATCH 16/18] fix(omni): mirror Flowboard workflow media fetch --- agent/services/omni_flash.py | 43 ++++++++++++++++++++++++++++++++---- 1 file changed, 39 insertions(+), 4 deletions(-) diff --git a/agent/services/omni_flash.py b/agent/services/omni_flash.py index 6dc27c21e..a757a42f0 100644 --- a/agent/services/omni_flash.py +++ b/agent/services/omni_flash.py @@ -40,6 +40,40 @@ # Informational only. Flow pricing can be promotional/variable. OMNI_FLASH_CREDIT_COST = {4: 15, 6: 20, 8: 25, 10: 30} +# Exact workflow-media fetch contract mirrored from crisng95/flowboard. +# Do not use FlowClient.get_media() for Omni/workflow polling: that legacy +# helper appends Google's public API key to the query string, while Flowboard's +# working workflow-media request intentionally does not. +_FLOWBOARD_MEDIA_HEADERS = { + "content-type": "text/plain;charset=UTF-8", + "accept": "*/*", + "origin": "https://labs.google", + "referer": "https://labs.google/", +} + + +async def _fetch_workflow_media(client, media_id: str) -> dict: + """Fetch workflow-backed media using Flowboard's exact wire contract. + + The browser extension supplies the authenticated Flow Bearer context. + The URL query is deliberately limited to ``clientContext.tool=PINHOLE``: + no ``key=`` parameter, no project/paygate fields, and no captcha action. + """ + url = ( + "https://aisandbox-pa.googleapis.com" + f"/v1/media/{media_id}?clientContext.tool=PINHOLE" + ) + return await client._send( + "api_request", + { + "url": url, + "method": "GET", + "headers": dict(_FLOWBOARD_MEDIA_HEADERS), + "body": None, + }, + timeout=15, + ) + def _validate_duration(duration_s: int) -> None: if duration_s not in OMNI_FLASH_VALID_DURATIONS: @@ -347,9 +381,10 @@ async def check_omni_flash_status( ) -> dict: """Perform one non-blocking poll pass for Omni workflow-backed jobs. - Each workflow is polled via ``GET /v1/media/``. A completed - result is only reported once ``video.encodedVideo`` decodes to an MP4 - (``ftyp`` box present), matching Flow's current workflow behavior. + Each workflow is polled using Flowboard's exact unauthenticated-query media + URL (browser Bearer auth is supplied by the extension). A completed result + is only reported once ``video.encodedVideo`` decodes to an MP4 (``ftyp`` + box present), matching Flow's workflow behavior. """ normalized = [] for workflow in workflows or []: @@ -368,7 +403,7 @@ async def check_omni_flash_status( for workflow in normalized: name = workflow["name"] media_id = workflow["primary_media_id"] - response = await client.get_media(media_id) + response = await _fetch_workflow_media(client, media_id) http_status = response.get("status") if isinstance(response, dict) else None if isinstance(http_status, int) and http_status >= 400: From 5ca0543d37536c2bfcd073e8bd99df9dadd1b78f Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:52:36 +0300 Subject: [PATCH 17/18] test(omni): assert exact Flowboard media polling wire contract --- tests/unit/test_omni_flash.py | 56 ++++++++++++++++++++++++++++++----- 1 file changed, 49 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_omni_flash.py b/tests/unit/test_omni_flash.py index 858ec47e2..af2f672c7 100644 --- a/tests/unit/test_omni_flash.py +++ b/tests/unit/test_omni_flash.py @@ -7,6 +7,7 @@ from agent.services.omni_flash import ( OMNI_FLASH_MAX_REFERENCE_IMAGES, + _fetch_workflow_media, _load_model_key, check_omni_flash_status, extract_omni_workflows, @@ -250,9 +251,39 @@ async def test_submit_builds_flow_omni_r2v_request_and_poll_descriptor(): @pytest.mark.asyncio -async def test_omni_poll_pending_uses_media_endpoint_not_legacy_operation_poll(): +async def test_workflow_media_fetch_matches_flowboard_wire_contract(): client = MagicMock() - client.get_media = AsyncMock( + client._send = AsyncMock(return_value={"status": 200, "data": {}}) + + await _fetch_workflow_media(client, "primary-vid-1") + + client._send.assert_awaited_once_with( + "api_request", + { + "url": ( + "https://aisandbox-pa.googleapis.com" + "/v1/media/primary-vid-1?clientContext.tool=PINHOLE" + ), + "method": "GET", + "headers": { + "content-type": "text/plain;charset=UTF-8", + "accept": "*/*", + "origin": "https://labs.google", + "referer": "https://labs.google/", + }, + "body": None, + }, + timeout=15, + ) + params = client._send.await_args.args[1] + assert "key=" not in params["url"] + assert "captchaAction" not in params + + +@pytest.mark.asyncio +async def test_omni_poll_pending_uses_flowboard_media_transport_not_legacy_get_media(): + client = MagicMock() + client._send = AsyncMock( return_value={ "status": 200, "data": { @@ -261,15 +292,26 @@ async def test_omni_poll_pending_uses_media_endpoint_not_legacy_operation_poll() }, } ) - client.check_video_status = AsyncMock(side_effect=AssertionError("legacy poll must not be used")) + client.get_media = AsyncMock( + side_effect=AssertionError("legacy get_media must not be used for workflow polling") + ) + client.check_video_status = AsyncMock( + side_effect=AssertionError("legacy operation poll must not be used") + ) 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"}] ) - client.get_media.assert_awaited_once_with("media-1") + client.get_media.assert_not_awaited() client.check_video_status.assert_not_awaited() + client._send.assert_awaited_once() + params = client._send.await_args.args[1] + assert params["url"].endswith( + "/v1/media/media-1?clientContext.tool=PINHOLE" + ) + assert "key=" not in params["url"] assert result["done"] is False assert result["status"] == "PENDING" assert result["workflows"][0]["status"] == "PENDING" @@ -280,7 +322,7 @@ async def test_omni_poll_completed_detects_mp4_without_returning_base64_by_defau mp4 = b"\x00\x00\x00\x18ftypisom" + b"\x00" * 32 encoded = base64.b64encode(mp4).decode() client = MagicMock() - client.get_media = AsyncMock( + client._send = AsyncMock( return_value={ "status": 200, "data": { @@ -311,7 +353,7 @@ async def test_omni_poll_can_return_encoded_video_when_requested(): mp4 = b"\x00\x00\x00\x18ftypisom" + b"\x00" * 32 encoded = base64.b64encode(mp4).decode() client = MagicMock() - client.get_media = AsyncMock( + client._send = AsyncMock( return_value={"status": 200, "data": {"video": {"encodedVideo": encoded}}} ) @@ -327,7 +369,7 @@ async def test_omni_poll_can_return_encoded_video_when_requested(): @pytest.mark.asyncio async def test_omni_poll_404_is_treated_as_not_ready(): client = MagicMock() - client.get_media = AsyncMock(return_value={"status": 404, "data": {}}) + client._send = AsyncMock(return_value={"status": 404, "data": {}}) with patch("agent.services.omni_flash.get_flow_client", return_value=client): result = await check_omni_flash_status( From be2bc963a9627e493eed07da0163fa7825b3fa39 Mon Sep 17 00:00:00 2001 From: Bl0ck154 Date: Tue, 11 Aug 2026 22:47:50 +0300 Subject: [PATCH 18/18] fix(omni): poll completed media through Flow project --- agent/api/flow.py | 10 +- agent/services/omni_flash.py | 219 +++++++++++++++++++----------- docs/OMNI_FLASH.md | 242 +++++++++++++++++++++++++--------- extension/background.js | 16 ++- extension/manifest.json | 3 +- tests/unit/test_omni_flash.py | 184 ++++++++++++++++---------- 6 files changed, 459 insertions(+), 215 deletions(-) diff --git a/agent/api/flow.py b/agent/api/flow.py index d0b14df8f..e215fbb31 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -75,13 +75,15 @@ class CheckStatusRequest(BaseModel): operations: list[dict] = [] # Omni/workflow-mode callers should pass workflow descriptors instead of # operation handles. If workflows is set, /check-status automatically uses - # /v1/media/ polling. + # 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 @@ -256,9 +258,12 @@ async def check_status(body: CheckStatusRequest): 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") @@ -281,9 +286,12 @@ async def check_omni_status(body: CheckOmniStatusRequest): 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}") diff --git a/agent/services/omni_flash.py b/agent/services/omni_flash.py index a757a42f0..2324b812f 100644 --- a/agent/services/omni_flash.py +++ b/agent/services/omni_flash.py @@ -15,16 +15,16 @@ 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 must be polled through ``/v1/media/``. +workflow-backed and are polled through Flow's authenticated project data. """ from __future__ import annotations -import base64 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 @@ -40,36 +40,38 @@ # Informational only. Flow pricing can be promotional/variable. OMNI_FLASH_CREDIT_COST = {4: 15, 6: 20, 8: 25, 10: 30} -# Exact workflow-media fetch contract mirrored from crisng95/flowboard. -# Do not use FlowClient.get_media() for Omni/workflow polling: that legacy -# helper appends Google's public API key to the query string, while Flowboard's -# working workflow-media request intentionally does not. -_FLOWBOARD_MEDIA_HEADERS = { - "content-type": "text/plain;charset=UTF-8", - "accept": "*/*", - "origin": "https://labs.google", - "referer": "https://labs.google/", -} +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_workflow_media(client, media_id: str) -> dict: - """Fetch workflow-backed media using Flowboard's exact wire contract. - The browser extension supplies the authenticated Flow Bearer context. - The URL query is deliberately limited to ``clientContext.tool=PINHOLE``: - no ``key=`` parameter, no project/paygate fields, and no captcha action. - """ +async def _fetch_media_url(client, media_id: str) -> dict: + """Resolve Flow's authenticated media redirect without buffering the file.""" url = ( - "https://aisandbox-pa.googleapis.com" - f"/v1/media/{media_id}?clientContext.tool=PINHOLE" + "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" + f"?name={quote(media_id, safe='')}" ) return await client._send( - "api_request", + "trpc_request", { "url": url, "method": "GET", - "headers": dict(_FLOWBOARD_MEDIA_HEADERS), - "body": None, + "headers": {"content-type": "application/json"}, + "responseMode": "url", }, timeout=15, ) @@ -158,7 +160,11 @@ def _normalize_workflow(workflow: dict) -> dict | None: return None if not isinstance(primary_media_id, str) or not primary_media_id: return None - return {"name": name, "primary_media_id": primary_media_id} + 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]: @@ -175,31 +181,23 @@ def extract_omni_workflows(result: dict) -> list[dict]: return normalized -def _annotate_polling(result: dict) -> dict: +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": "workflow_media", + "mode": "project_media", + "project_id": project_id, "workflows": workflows, } return result -def _is_complete_mp4(encoded: str) -> bool: - """Flow can expose a small placeholder payload before the final MP4.""" - if not isinstance(encoded, str) or not encoded: - return False - try: - binary = base64.b64decode(encoded, validate=False) - except Exception: - return False - return len(binary) >= 12 and binary[4:8] == b"ftyp" - - async def _submit_omni_frame_video( *, start_image_media_id: str, @@ -264,7 +262,7 @@ async def _submit_omni_frame_video( }, timeout=60, ) - return _annotate_polling(result) + return _annotate_polling(result, project_id) async def generate_omni_flash_first_frame_video( @@ -372,19 +370,19 @@ async def generate_omni_flash_video( }, timeout=60, ) - return _annotate_polling(result) + 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. - Each workflow is polled using Flowboard's exact unauthenticated-query media - URL (browser Bearer auth is supplied by the extension). A completed result - is only reported once ``video.encodedVideo`` decodes to an MP4 (``ftyp`` - box present), matching Flow's workflow behavior. + 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 []: @@ -397,62 +395,107 @@ async def check_omni_flash_status( "(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"] - response = await _fetch_workflow_media(client, media_id) - - http_status = response.get("status") if isinstance(response, dict) else None - if isinstance(http_status, int) and http_status >= 400: - if http_status == 404: - items.append({ - "name": name, - "primary_media_id": media_id, - "done": False, - "status": "PENDING", - "error": None, - }) - continue - data = response.get("data") if isinstance(response.get("data"), dict) else {} - error = data.get("error") if isinstance(data, dict) else None - error_message = None - if isinstance(error, dict): - error_message = error.get("message") or error.get("status") + 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, - "done": True, - "status": "FAILED", - "error": error_message or response.get("error") or f"API_{http_status}", + "project_id": resolved_project_id, + "done": False, + "status": "PENDING", + "error": None, + "workflow_present": name in known_workflow_names, }) continue - payload = response.get("data") if isinstance(response, dict) and isinstance(response.get("data"), dict) else response - payload = payload if isinstance(payload, dict) else {} - video = payload.get("video") if isinstance(payload.get("video"), dict) else {} - encoded = video.get("encodedVideo") if isinstance(video, dict) else None - - media_status = payload.get("mediaMetadata") if isinstance(payload.get("mediaMetadata"), dict) else {} - media_status = media_status.get("mediaStatus") if isinstance(media_status.get("mediaStatus"), dict) else {} - generation_status = media_status.get("mediaGenerationStatus") if isinstance(media_status, dict) else None + 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"): + 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 not _is_complete_mp4(encoded): + 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, @@ -460,28 +503,44 @@ async def check_omni_flash_status( continue url = None - if isinstance(video, dict): - url = video.get("fifeUrl") or video.get("servingUri") - url = url or payload.get("fifeUrl") or payload.get("servingUri") + 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 if isinstance(url, str) else None, - "encoded_video_available": True, + "url": url, + "encoded_video_available": False, }, } if include_encoded_video: - item["media"]["encoded_video"] = encoded + 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 index fc39380a4..93aa1e94c 100644 --- a/docs/OMNI_FLASH.md +++ b/docs/OMNI_FLASH.md @@ -1,123 +1,237 @@ -# Gemini Omni Flash +# Gemini Omni Flash: integration guide for agents -FlowKit can submit Gemini Omni Flash video jobs through the same authenticated Chrome extension bridge used by the existing Google Flow integrations. +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. -## Supported in this integration +## Prerequisites and base URL -- First frame → video -- First + Last frame → video -- Reference-to-video with 1–7 reference images -- Portrait (`9:16`) and landscape (`16:9`) -- 4, 6, 8, and 10 second generations -- Workflow/media polling via each submit response's `primaryMediaId` -- Existing Google Flow account / paygate tier handling +Before submitting work, both checks must pass: -## First frame / First + Last +```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":{}} +``` -The existing `POST /flow/generate-video` endpoint remains backward compatible: requests use Veo unless `model_family` is explicitly set to `omni_flash`. +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. -### First frame +## 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 -X POST http://127.0.0.1:8100/flow/generate-video \ +curl -fsS -X POST "$FLOWKIT_BASE_URL/api/flow/generate-video" \ -H 'Content-Type: application/json' \ -d '{ - "start_image_media_id": "", - "prompt": "Natural cinematic motion, keep the subject consistent", - "project_id": "", - "scene_id": "scene-1", "model_family": "omni_flash", - "duration_s": 8, + "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 submits through `batchAsyncGenerateVideoStartImage` with the configured `abra_i2v_s` key. +This uses `batchAsyncGenerateVideoStartImage`. -### First + Last frame - -Add `end_image_media_id` to the same request: +## First + Last frame to video ```bash -curl -X POST http://127.0.0.1:8100/flow/generate-video \ +curl -fsS -X POST "$FLOWKIT_BASE_URL/api/flow/generate-video" \ -H 'Content-Type: application/json' \ -d '{ - "start_image_media_id": "", - "end_image_media_id": "", - "prompt": "Move naturally from the first composition to the final composition", - "project_id": "", - "scene_id": "scene-1", "model_family": "omni_flash", - "duration_s": 8, + "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 submits through `batchAsyncGenerateVideoStartAndEndImage` and includes both `startImage` and `endImage` in the Flow request. - -The First+Last model mapping is stored independently in `models.json`. The current default uses the same `abra_i2v_s` family as Omni First-frame generation. Keeping it as a separate mapping allows a one-line model-config update if Google changes the key during rollout. +This uses `batchAsyncGenerateVideoStartAndEndImage` and sends both `startImage` and `endImage`. -## Reference-to-video +## References to video -### Dedicated endpoint +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 -X POST http://127.0.0.1:8100/flow/generate-video-omni \ +curl -fsS -X POST "$FLOWKIT_BASE_URL/api/flow/generate-video-omni" \ -H 'Content-Type: application/json' \ -d '{ - "reference_media_ids": [""], - "prompt": "Cinematic handheld shot, natural motion and dialogue", - "project_id": "", - "duration_s": 10, + "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" }' ``` -Or use `POST /flow/generate-video-refs` with `model_family: "omni_flash"` and `duration_s`. +The compatible generic endpoint is `POST /api/flow/generate-video-refs` with the same fields plus `"model_family":"omni_flash"`. -## Important: Omni does not use legacy Veo polling +## Submit response and polling -Do **not** take the operation-looking handle returned by an Omni submit and send it to the legacy `batchCheckAsyncVideoGenerationStatus` path. Google Flow can reject that with `INVALID_ARGUMENT`. - -Omni Flash submits are workflow-backed. The submit response includes `workflows[].metadata.primaryMediaId`; FlowKit normalizes those into: +Persist the entire normalized polling descriptor returned by submit: ```json { "flowkitPolling": { - "mode": "workflow_media", + "mode": "project_media", + "project_id": "FLOW_PROJECT_ID", "workflows": [ { - "name": "", - "primary_media_id": "" + "name": "WORKFLOW_NAME", + "primary_media_id": "PRIMARY_MEDIA_ID", + "project_id": "FLOW_PROJECT_ID" } ] } } ``` -Poll those workflow descriptors through `POST /flow/check-omni-status`, or pass them as `workflows` to the generic `POST /flow/check-status` endpoint. +Poll using those values without transforming workflow names into operation handles: ```bash -curl -X POST http://127.0.0.1:8100/flow/check-omni-status \ +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": "", - "primary_media_id": "" + "name": "WORKFLOW_NAME", + "primary_media_id": "PRIMARY_MEDIA_ID", + "project_id": "FLOW_PROJECT_ID" } - ] + ], + "include_encoded_video": false }' ``` -While rendering, FlowKit returns `status: "PENDING"`. Once `/v1/media/` exposes a complete MP4 payload, it returns `MEDIA_GENERATION_STATUS_SUCCESSFUL` and the media ID. Set `include_encoded_video: true` only if you actually need the base64 MP4 inline; by default FlowKit avoids returning that large payload. +The generic `POST /api/flow/check-status` endpoint also accepts the same `workflows` and automatically selects Omni project polling. -## Model keys +Pending response: + +```json +{ + "project_id": "FLOW_PROJECT_ID", + "done": false, + "status": "PENDING", + "workflows": [{"done":false,"status":"PENDING"}] +} +``` -The duration mappings live in `agent/models.json`: +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 { @@ -144,8 +258,18 @@ The duration mappings live in `agent/models.json`: } ``` -The values can be updated through `PATCH /api/models` if Google rotates Flow's internal model keys. +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 -## Notes +- 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 Google Flow integration and its internal model keys/endpoints can change as Flow evolves. First-frame generation uses the wire-proven `abra_i2v_s` family and `batchAsyncGenerateVideoStartImage`. First+Last uses the same configurable Omni I2V family with `batchAsyncGenerateVideoStartAndEndImage`; this matches the current Flow UI capability while keeping the key isolated so it can be changed without another code release. Reference generation uses `abra_r2v_s` with `batchAsyncGenerateVideoReferenceImages`. +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 index af2f672c7..38d5ed9fc 100644 --- a/tests/unit/test_omni_flash.py +++ b/tests/unit/test_omni_flash.py @@ -1,13 +1,13 @@ """Unit tests for Gemini Omni Flash Flow submissions and workflow polling.""" -import base64 from unittest.mock import AsyncMock, MagicMock, patch import pytest from agent.services.omni_flash import ( OMNI_FLASH_MAX_REFERENCE_IMAGES, - _fetch_workflow_media, + _fetch_media_url, + _fetch_project_initial_data, _load_model_key, check_omni_flash_status, extract_omni_workflows, @@ -148,7 +148,8 @@ async def test_submit_builds_flow_omni_first_frame_request_and_poll_descriptor() assert request["startImage"] == {"mediaId": "start-1"} assert "endImage" not in request assert request["seed"] == 321 - assert result["data"]["flowkitPolling"]["mode"] == "workflow_media" + assert result["data"]["flowkitPolling"]["mode"] == "project_media" + assert result["data"]["flowkitPolling"]["project_id"] == "project-1" @pytest.mark.asyncio @@ -179,7 +180,11 @@ async def test_submit_builds_flow_omni_first_last_request(): assert request["aspectRatio"] == "VIDEO_ASPECT_RATIO_PORTRAIT" assert request["seed"] == 456 assert result["data"]["flowkitPolling"]["workflows"] == [ - {"name": "workflow-1", "primary_media_id": "media-1"} + { + "name": "workflow-1", + "primary_media_id": "media-1", + "project_id": "project-1", + } ] @@ -222,9 +227,14 @@ async def test_submit_builds_flow_omni_r2v_request_and_poll_descriptor(): assert result["status"] == 200 assert result["data"]["flowkitPolling"] == { - "mode": "workflow_media", + "mode": "project_media", + "project_id": "project-1", "workflows": [ - {"name": "workflow-1", "primary_media_id": "media-1"} + { + "name": "workflow-1", + "primary_media_id": "media-1", + "project_id": "project-1", + } ], } client._build_url.assert_called_once_with("generate_video_references") @@ -250,92 +260,131 @@ async def test_submit_builds_flow_omni_r2v_request_and_poll_descriptor(): ] +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_workflow_media_fetch_matches_flowboard_wire_contract(): +async def test_project_poll_fetch_matches_live_flow_trpc_contract(): client = MagicMock() client._send = AsyncMock(return_value={"status": 200, "data": {}}) - await _fetch_workflow_media(client, "primary-vid-1") + await _fetch_project_initial_data(client, "project-1") client._send.assert_awaited_once_with( - "api_request", + "trpc_request", { "url": ( - "https://aisandbox-pa.googleapis.com" - "/v1/media/primary-vid-1?clientContext.tool=PINHOLE" + "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": "text/plain;charset=UTF-8", - "accept": "*/*", - "origin": "https://labs.google", - "referer": "https://labs.google/", - }, - "body": None, + "headers": {"content-type": "application/json"}, }, timeout=15, ) - params = client._send.await_args.args[1] - assert "key=" not in params["url"] - assert "captchaAction" not in params @pytest.mark.asyncio -async def test_omni_poll_pending_uses_flowboard_media_transport_not_legacy_get_media(): +async def test_media_redirect_fetch_requests_url_only_mode(): client = MagicMock() - client._send = AsyncMock( - return_value={ - "status": 200, - "data": { - "name": "media-1", - "video": {"encodedVideo": ""}, - }, - } - ) - client.get_media = AsyncMock( - side_effect=AssertionError("legacy get_media must not be used for workflow polling") - ) - client.check_video_status = AsyncMock( - side_effect=AssertionError("legacy operation poll must not be used") + 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"}] + [ + { + "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() - params = client._send.await_args.args[1] - assert params["url"].endswith( - "/v1/media/media-1?clientContext.tool=PINHOLE" - ) - assert "key=" not in params["url"] assert result["done"] is False assert result["status"] == "PENDING" assert result["workflows"][0]["status"] == "PENDING" @pytest.mark.asyncio -async def test_omni_poll_completed_detects_mp4_without_returning_base64_by_default(): - mp4 = b"\x00\x00\x00\x18ftypisom" + b"\x00" * 32 - encoded = base64.b64encode(mp4).decode() +async def test_omni_poll_completed_returns_signed_url_without_buffering_video(): client = MagicMock() client._send = AsyncMock( - return_value={ - "status": 200, - "data": { - "name": "media-1", - "fifeUrl": "https://example.test/video.mp4", - "video": {"encodedVideo": encoded}, + 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"}] + [{"name": "workflow-1", "primary_media_id": "media-1"}], + project_id="project-1", ) assert result["done"] is True @@ -343,42 +392,33 @@ async def test_omni_poll_completed_detects_mp4_without_returning_base64_by_defau item = result["workflows"][0] assert item["status"] == "MEDIA_GENERATION_STATUS_SUCCESSFUL" assert item["media"]["media_id"] == "media-1" - assert item["media"]["url"] == "https://example.test/video.mp4" - assert item["media"]["encoded_video_available"] is True + 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_can_return_encoded_video_when_requested(): - mp4 = b"\x00\x00\x00\x18ftypisom" + b"\x00" * 32 - encoded = base64.b64encode(mp4).decode() +async def test_omni_poll_missing_media_is_pending(): client = MagicMock() - client._send = AsyncMock( - return_value={"status": 200, "data": {"video": {"encodedVideo": encoded}}} - ) + 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"}], - include_encoded_video=True, + project_id="project-1", ) - assert result["workflows"][0]["media"]["encoded_video"] == encoded + assert result["done"] is False + assert result["workflows"][0]["status"] == "PENDING" @pytest.mark.asyncio -async def test_omni_poll_404_is_treated_as_not_ready(): - client = MagicMock() - client._send = AsyncMock(return_value={"status": 404, "data": {}}) - - with patch("agent.services.omni_flash.get_flow_client", return_value=client): - result = await check_omni_flash_status( +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"}] ) - assert result["done"] is False - assert result["workflows"][0]["status"] == "PENDING" - @pytest.mark.asyncio async def test_submit_accepts_seven_references():