From 37512acee0ae02d28ee7d2b6fce72f69c6482621 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 8 Sep 2026 08:08:30 +0300 Subject: [PATCH 1/3] feat(flow): restore 1080p export on migrated API --- README.md | 7 +- agent/api/flow.py | 1 + agent/api/upscale_status.py | 87 +++++++++++++ agent/main.py | 2 + agent/services/flow_batch.py | 21 +++ agent/services/flow_client.py | 31 ++++- agent/services/upscale_polling.py | 185 +++++++++++++++++++++++++++ docs/VIDEO_EXPORTS.md | 50 ++++++++ tests/unit/test_flow_batch.py | 15 +++ tests/unit/test_flow_client_batch.py | 25 +++- 10 files changed, 410 insertions(+), 14 deletions(-) create mode 100644 agent/api/upscale_status.py create mode 100644 agent/services/upscale_polling.py create mode 100644 docs/VIDEO_EXPORTS.md diff --git a/README.md b/README.md index 8084b1de..65cced40 100644 --- a/README.md +++ b/README.md @@ -254,12 +254,13 @@ without being silently replaced by the default model. See ### What does not work on the new API yet -Three capabilities have no captured payload, so they fail with -`UNSUPPORTED_ON_BATCH_API` rather than quietly producing the wrong thing: +1080p export is ported: FlowKit mirrors the Flow UI's `p0UkFb` high-resolution +Download request and polls the resulting media through `as29s`. The remaining +capabilities below still fail loudly rather than quietly producing the wrong thing: | Capability | Status | Workaround | |---|---|---| -| 4K/1080p upscale (`/fk-pipeline` last step) | unported | none — keep the 1080p render | +| 4K export | plan-gated / not live-verified | use 1080p; Google Flow exposes Full HD as the standard high-resolution export | | Reference-to-video (r2v) | unported | `FLOW_ALLOW_DEGRADED=1` → i2v off the first reference | | Start+end-frame chaining (`/fk-gen-chain-videos`) | unported | `FLOW_ALLOW_DEGRADED=1` → i2v off the start frame | | Omni Flash text-to-video | ported | `POST /api/flow/generate-video-omni-text` (4/6/8/10s) | diff --git a/agent/api/flow.py b/agent/api/flow.py index 5aeed8fa..1f585ea8 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -78,6 +78,7 @@ class UpscaleVideoRequest(BaseModel): scene_id: str aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT" resolution: str = "VIDEO_RESOLUTION_4K" + project_id: Optional[str] = None class UploadImageRequest(BaseModel): diff --git a/agent/api/upscale_status.py b/agent/api/upscale_status.py new file mode 100644 index 00000000..6f9c0f9f --- /dev/null +++ b/agent/api/upscale_status.py @@ -0,0 +1,87 @@ +"""Explicit Full HD / 4K export endpoints for Google Flow videos.""" + +from typing import Literal + +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel + +from agent.services.flow_client import get_flow_client +from agent.services.upscale_polling import annotate_upscale_polling, check_upscale_status + +router = APIRouter(prefix="/flow", tags=["flow"]) + + +class ExportVideoRequest(BaseModel): + media_id: str + scene_id: str = "export" + quality: Literal["1080p", "4k"] = "1080p" + aspect_ratio: str = "VIDEO_ASPECT_RATIO_LANDSCAPE" + project_id: str | None = None + + +class CheckExportStatusRequest(BaseModel): + workflows: list[dict] + + +@router.post("/export-video") +async def export_video(body: ExportVideoRequest): + """Start Google's native Full HD/4K export. + + This is the same Flow upsample operation exposed by the UI, presented as an + export/download-quality choice. 1080p is the default because Omni Flash's + generated file is normally 720p and Full HD is the expected downloadable + master. + """ + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + + resolution = ( + "VIDEO_RESOLUTION_1080P" + if body.quality == "1080p" + else "VIDEO_RESOLUTION_4K" + ) + result = await client.upscale_video( + media_id=body.media_id, + scene_id=body.scene_id, + aspect_ratio=body.aspect_ratio, + resolution=resolution, + project_id=body.project_id, + ) + if result.get("error") or ( + isinstance(result.get("status"), int) and result["status"] >= 400 + ): + raise HTTPException( + result.get("status", 502), + result.get("error", result.get("data")), + ) + + annotated = annotate_upscale_polling(result) + data = annotated.get("data", annotated) + if isinstance(data, dict): + data["export"] = { + "quality": body.quality, + "resolution": resolution, + "native_flow_export": True, + "next": "/api/flow/check-export-status", + } + return data + + +@router.post("/check-export-status") +@router.post("/check-upscale-status", include_in_schema=False) +async def check_export_status(body: CheckExportStatusRequest): + """Return a signed downloadable URL when the native Flow export is ready.""" + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + try: + result = await check_upscale_status(body.workflows) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + + if result.get("status") == "COMPLETED": + result["download_ready"] = True + else: + result["download_ready"] = False + return result diff --git a/agent/main.py b/agent/main.py index 9649c807..14b30ce4 100644 --- a/agent/main.py +++ b/agent/main.py @@ -17,6 +17,7 @@ from agent.api.scenes import router as scenes_router from agent.api.requests import router as requests_router from agent.api.flow import router as flow_router +from agent.api.upscale_status import router as upscale_status_router from agent.api.reviews import router as reviews_router from agent.api.tts import router as tts_router from agent.api.materials import router as materials_router @@ -128,6 +129,7 @@ async def lifespan(app: FastAPI): app.include_router(scenes_router, prefix="/api") app.include_router(requests_router, prefix="/api") app.include_router(flow_router, prefix="/api") +app.include_router(upscale_status_router, prefix="/api") app.include_router(reviews_router, prefix="/api") app.include_router(tts_router, prefix="/api") app.include_router(materials_router, prefix="/api") diff --git a/agent/services/flow_batch.py b/agent/services/flow_batch.py index f81c1e44..2071dcdf 100644 --- a/agent/services/flow_batch.py +++ b/agent/services/flow_batch.py @@ -42,6 +42,7 @@ RPC_MEDIA = "as29s" RPC_UPLOAD_IMAGE = "maseQ" RPC_UPSCALE_IMAGE = "SPrCad" +RPC_UPSCALE = "p0UkFb" CAPTCHA_IMAGE = "IMAGE_GENERATION" CAPTCHA_VIDEO = "VIDEO_GENERATION" @@ -426,6 +427,20 @@ def text_video_request(prompt: str, project_id: str, [request], _context(project_id), [_client_uuid(), 1], +def upscale_request(media_id: str, project_id: str, + aspect: Any = VIDEO_ASPECT_LANDSCAPE, + model: str = "veo_3_1_upsampler_1080p") -> str: + """Build Flow's migrated high-resolution download request (RPC p0UkFb).""" + item = [None] * 32 + item[0] = [None, media_id] + item[2] = 1 + item[4] = [None, str(uuid.uuid4()), None, None, _client_uuid()] + item[6] = resolve_video_aspect(aspect) + item[31] = model + return build_envelope(RPC_UPSCALE, [ + [item], + _context(project_id), + [_client_uuid()], ]) @@ -526,6 +541,12 @@ def read_upscaled_image(payload: Any) -> str: if not isinstance(encoded, str) or len(encoded) < 100: raise FlowBatchError("image upscale response carried no encoded image") return encoded +def read_upscaled_media_id(payload: Any) -> str: + """Return the media id created by the p0UkFb upscale submit.""" + for text in _walk_strings(payload): + if isinstance(text, str) and text.endswith("_upsampled"): + return text + raise FlowBatchError("upscale response carried no upsampled media id") def read_operation(payload: Any) -> Operation: diff --git a/agent/services/flow_client.py b/agent/services/flow_client.py index be56e708..ec66a1ae 100644 --- a/agent/services/flow_client.py +++ b/agent/services/flow_client.py @@ -839,14 +839,33 @@ async def generate_video_from_references(self, reference_media_ids: list[str], async def upscale_video(self, media_id: str, scene_id: str, aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT", - resolution: str = "VIDEO_RESOLUTION_4K") -> dict: - """Upscale a video.""" + resolution: str = "VIDEO_RESOLUTION_4K", + project_id: str | None = None) -> dict: + """Upscale/export a video using Flow's migrated p0UkFb RPC.""" if not USE_BATCH_RPC: return await self._legacy_upscale_video(media_id, scene_id, aspect_ratio, resolution) - return {"error": _unsupported( - "video upscale", - "no upsampler rpc appears in the new frontend's captures", - )} + + model = UPSCALE_MODELS.get(resolution) + if not model: + return {"status": 400, "error": f"Unsupported upscale resolution: {resolution}"} + try: + pid = self._batch_project_id(project_id or "") + freq = fb.upscale_request(media_id, pid, aspect=aspect_ratio, model=model) + payload = await self._batch_payload( + fb.RPC_UPSCALE, freq, fb.CAPTCHA_VIDEO, timeout=120) + upscaled_media_id = fb.read_upscaled_media_id(payload) + except Exception as e: + return _batch_error(e) + + workflow = { + "name": upscaled_media_id, + "primary_media_id": upscaled_media_id, + "project_id": pid, + } + return {"status": 200, "data": { + "media": [{"name": upscaled_media_id}], + "workflows": [workflow], + }} async def check_video_status(self, operations: list[dict]) -> dict: """One poll round for each submitted operation. diff --git a/agent/services/upscale_polling.py b/agent/services/upscale_polling.py new file mode 100644 index 00000000..29d705e5 --- /dev/null +++ b/agent/services/upscale_polling.py @@ -0,0 +1,185 @@ +"""Headless polling for Flow video upscales. + +Google Flow's upsampler returns workflow descriptors whose logical +``primaryMediaId`` may not appear in ``flow.projectInitialData``. The browser UI +can still resolve completed media through ``media.getMediaUrlRedirect``. + +This module exposes a small active poller that treats a successful authenticated +media redirect as the completion signal, avoiding the legacy +``batchCheckAsyncVideoGenerationStatus`` and ``/v1/media/{id}`` paths. +""" + +from __future__ import annotations + +from urllib.parse import quote + +from agent.config import USE_BATCH_RPC +from agent.services.flow_client import get_flow_client + +_ALLOWED_MEDIA_URL_PREFIX = "https://flow-content.google/" + + +def _normalize_workflow(workflow: dict) -> dict | None: + 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_upscale_workflows(result: dict) -> list[dict]: + 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_upscale_polling(result: dict) -> dict: + workflows = extract_upscale_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": "media_redirect", + "workflows": workflows, + } + return result + + +async def _fetch_media_url(client, media_id: str) -> dict: + if USE_BATCH_RPC: + result = await client.get_media(media_id) + data = result.get("data") if isinstance(result.get("data"), dict) else {} + video = data.get("video") if isinstance(data, dict) else None + candidate = video.get("fifeUrl") if isinstance(video, dict) else None + return { + "status": result.get("status", 200), + "data": { + "url": candidate, + "contentType": "video/mp4" if candidate else None, + }, + "error": result.get("error"), + } + + url = ( + "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" + f"?name={quote(media_id, safe='')}" + ) + return await client._send( + "trpc_request", + { + "url": url, + "method": "GET", + "headers": {"content-type": "application/json"}, + "responseMode": "url", + }, + timeout=15, + ) + + +def _parse_media_redirect(response: dict) -> tuple[str | None, str | None, str | None]: + if not isinstance(response, dict): + return None, None, "Flow media redirect returned an invalid response" + status = response.get("status") + data = response.get("data") if isinstance(response.get("data"), dict) else {} + candidate = data.get("url") + content_type = data.get("contentType") + if ( + isinstance(status, int) + and status < 400 + and isinstance(candidate, str) + and candidate.startswith(_ALLOWED_MEDIA_URL_PREFIX) + ): + return candidate, content_type if isinstance(content_type, str) else None, None + error = response.get("error") + if not error and isinstance(status, int) and status >= 400: + error = f"API_{status}" + if not error: + error = "media redirect not ready" + return None, content_type if isinstance(content_type, str) else None, str(error) + + +async def check_upscale_status( + workflows: list[dict], + include_encoded_video: bool = False, +) -> dict: + """Poll native Flow Full HD/4K export workflows without buffering the MP4.""" + normalized = [] + for workflow in workflows or []: + item = _normalize_workflow(workflow) + if item: + normalized.append(item) + if not normalized: + raise ValueError( + "Export polling requires workflow descriptors with name and " + "primary_media_id (or raw Flow metadata.primaryMediaId)" + ) + + client = get_flow_client() + items = [] + for workflow in normalized: + media_id = workflow["primary_media_id"] + response = await _fetch_media_url(client, media_id) + url, content_type, diagnostic = _parse_media_redirect(response) + if url: + media = { + "media_id": media_id, + "url": url, + "encoded_video_available": False, + "resolved_via": "as29s" if USE_BATCH_RPC else "media.getMediaUrlRedirect", + } + if content_type: + media["content_type"] = content_type + if include_encoded_video: + media["encoded_video"] = None + items.append({ + "name": workflow["name"], + "primary_media_id": media_id, + "done": True, + "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", + "error": None, + "media": media, + }) + continue + + probe = {} + if isinstance(response, dict): + if isinstance(response.get("status"), int): + probe["http_status"] = response["status"] + data = response.get("data") + if isinstance(data, dict) and isinstance(data.get("url"), str): + probe["resolved_url"] = data["url"] + if diagnostic: + probe["diagnostic"] = diagnostic + item = { + "name": workflow["name"], + "primary_media_id": media_id, + "done": False, + "status": "PENDING", + "error": None, + } + if probe: + item["probe"] = probe + items.append(item) + + all_done = bool(items) and all(item["done"] for item in items) + return { + "done": all_done, + "status": "COMPLETED" if all_done else "PENDING", + "workflows": items, + } diff --git a/docs/VIDEO_EXPORTS.md b/docs/VIDEO_EXPORTS.md new file mode 100644 index 00000000..de056793 --- /dev/null +++ b/docs/VIDEO_EXPORTS.md @@ -0,0 +1,50 @@ +# Video exports + +FlowKit treats output quality as an explicit export/download choice rather than +an implementation detail. + +## Omni Flash + +Gemini Omni Flash generation normally produces a 720p source video. Google Flow +provides a native Full HD export for that result. FlowKit exposes it directly as +**Export 1080p**; internally Google names the operation an upsample, but callers +do not need to reason about that implementation detail. + +## Export Full HD (recommended) + +```bash +curl -sS -X POST http://127.0.0.1:8100/api/flow/export-video \ + -H 'Content-Type: application/json' \ + -d '{ + "media_id": "", + "scene_id": "job-1", + "quality": "1080p", + "aspect_ratio": "VIDEO_ASPECT_RATIO_LANDSCAPE" + }' +``` + +`quality` defaults to `1080p`. `4k` remains an explicit API option, but it is plan-gated by Google Flow; the currently verified account exposes 4K as disabled while 1080p is available. + +The response contains `flowkitPolling.workflows`. Poll those descriptors: + +```bash +curl -sS -X POST http://127.0.0.1:8100/api/flow/check-export-status \ + -H 'Content-Type: application/json' \ + -d '{"workflows": }' +``` + +When `download_ready` becomes `true`, use +`workflows[].media.url` immediately. It is a short-lived signed +`flow-content.google` URL. + +## Compatibility + +The older endpoints remain supported: + +- `POST /api/flow/upscale-video` +- `POST /api/flow/check-upscale-status` + +They are aliases/low-level surfaces for the same Google Flow capability. New +integrations should use `export-video` and `check-export-status`, because those +names describe the user-visible operation: selecting the downloadable output +quality. diff --git a/tests/unit/test_flow_batch.py b/tests/unit/test_flow_batch.py index db1f3457..0c70d672 100644 --- a/tests/unit/test_flow_batch.py +++ b/tests/unit/test_flow_batch.py @@ -169,6 +169,19 @@ def test_text_video_matches_the_captured_yhhmef_shape(self): assert len(request[4]) == 6 assert payload[1][5] == self.PID assert payload[2][1] == 1 + def test_upscale_matches_the_captured_1080p_slots(self): + payload = inner(fb.upscale_request( + "media-1", self.PID, + aspect="VIDEO_ASPECT_RATIO_LANDSCAPE", + model="veo_3_1_upsampler_1080p", + )) + item = payload[0][0] + assert len(item) == 32 + assert item[0] == [None, "media-1"] + assert item[2] == 1 + assert item[6] == fb.VIDEO_ASPECT_LANDSCAPE + assert item[31] == "veo_3_1_upsampler_1080p" + assert payload[1][5] == self.PID assert fb.CAPTCHA_SLOT in json.dumps(payload) @@ -198,6 +211,8 @@ def test_image_upscale_reads_synchronous_encoded_image(self): assert fb.read_upscaled_image([["media"], "A" * 200]) == "A" * 200 with pytest.raises(fb.FlowBatchError): fb.read_upscaled_image([["media"], "short"]) + def test_upscale_submit_reads_the_new_media_id(self): + assert fb.read_upscaled_media_id([[self.MID + "_upsampled"]]) == self.MID + "_upsampled" def test_operation_reads_the_id_and_status(self): op = fb.read_operation([None, 50, [[self.OP, "proj", "scene", "CAE"]]]) diff --git a/tests/unit/test_flow_client_batch.py b/tests/unit/test_flow_client_batch.py index c38c09ed..f9882ff2 100644 --- a/tests/unit/test_flow_client_batch.py +++ b/tests/unit/test_flow_client_batch.py @@ -294,11 +294,26 @@ async def test_degraded_r2v_uses_the_first_reference_as_the_start_frame(self, cl payload = json.loads(json.loads(client.calls[0]["freq"])[0][0][1]) assert payload[0][0][4][1] == "ref-a" - async def test_upscale_is_unported_and_has_no_fallback(self, client, monkeypatch): - import agent.services.flow_client as module - monkeypatch.setattr(module, "FLOW_ALLOW_DEGRADED", True) - result = await client.upscale_video(MEDIA, "scene-1") - assert "UNSUPPORTED_ON_BATCH_API" in result["error"] + + +class TestUpscaleVideo: + async def test_1080p_submit_returns_a_pollable_workflow(self, client): + upscaled = MEDIA + "_upsampled" + client.responses[fb.RPC_UPSCALE] = { + "data": envelope(fb.RPC_UPSCALE, [[[[upscaled], "", None, None, 1]]]) + } + + result = await client.upscale_video( + MEDIA, + "scene-1", + aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE", + resolution="VIDEO_RESOLUTION_1080P", + ) + workflow = result["data"]["workflows"][0] + assert workflow["primary_media_id"] == upscaled + assert workflow["project_id"] == PROJECT + assert client.calls[0]["rpcid"] == fb.RPC_UPSCALE + assert client.calls[0]["captcha"] == fb.CAPTCHA_VIDEO class TestCheckVideoStatus: From 679d0c8373cf341c12c41dbd41a1fa1f13c47ae9 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 8 Sep 2026 08:10:21 +0300 Subject: [PATCH 2/3] test(flow): cover migrated 1080p polling --- tests/unit/test_upscale_polling.py | 67 ++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 tests/unit/test_upscale_polling.py diff --git a/tests/unit/test_upscale_polling.py b/tests/unit/test_upscale_polling.py new file mode 100644 index 00000000..14826628 --- /dev/null +++ b/tests/unit/test_upscale_polling.py @@ -0,0 +1,67 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import agent.services.upscale_polling as polling +from agent.services.upscale_polling import annotate_upscale_polling, check_upscale_status + + +MEDIA = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + + +def test_annotate_upscale_polling_uses_workflow_primary_media_id(): + result = { + "status": 200, + "data": { + "workflows": [ + {"name": "wf-1", "metadata": {"primaryMediaId": MEDIA + "_upsampled"}} + ] + }, + } + annotated = annotate_upscale_polling(result) + assert annotated["data"]["flowkitPolling"] == { + "mode": "media_redirect", + "workflows": [ + {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} + ], + } + + +@pytest.mark.asyncio +async def test_batch_poll_resolves_completed_1080_media_through_as29s(monkeypatch): + monkeypatch.setattr(polling, "USE_BATCH_RPC", True) + client = MagicMock() + client.get_media = AsyncMock(return_value={ + "status": 200, + "data": { + "video": { + "fifeUrl": "https://flow-content.google/video/out?Signature=test" + } + }, + }) + with patch("agent.services.upscale_polling.get_flow_client", return_value=client): + result = await check_upscale_status([ + {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} + ]) + + assert result["done"] is True + assert result["status"] == "COMPLETED" + media = result["workflows"][0]["media"] + assert media["resolved_via"] == "as29s" + assert media["url"].startswith("https://flow-content.google/video/") + client.get_media.assert_awaited_once_with(MEDIA + "_upsampled") + + +@pytest.mark.asyncio +async def test_batch_poll_stays_pending_until_video_url_exists(monkeypatch): + monkeypatch.setattr(polling, "USE_BATCH_RPC", True) + client = MagicMock() + client.get_media = AsyncMock(return_value={"status": 200, "data": {"video": {}}}) + with patch("agent.services.upscale_polling.get_flow_client", return_value=client): + result = await check_upscale_status([ + {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} + ]) + + assert result["done"] is False + assert result["status"] == "PENDING" + assert result["workflows"][0]["status"] == "PENDING" From a7a7dd9cd4dd2b8e96545174124e7b4cb40a4513 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:04:50 +0300 Subject: [PATCH 3/3] refactor(flow): share workflow polling for video export --- agent/api/flow.py | 58 +++++- agent/api/upscale_status.py | 87 -------- agent/main.py | 2 - agent/services/flow_batch.py | 5 + agent/services/flow_client.py | 49 +++++ agent/services/flow_poll.py | 314 +++++++++++++++++++++++++++++ agent/services/omni_flash.py | 311 ++-------------------------- agent/services/upscale_polling.py | 185 ----------------- docs/VIDEO_EXPORTS.md | 14 +- tests/unit/test_flow_poll.py | 134 ++++++++++++ tests/unit/test_omni_flash.py | 79 +++----- tests/unit/test_upscale_polling.py | 67 ------ 12 files changed, 609 insertions(+), 696 deletions(-) delete mode 100644 agent/api/upscale_status.py create mode 100644 agent/services/flow_poll.py delete mode 100644 agent/services/upscale_polling.py create mode 100644 tests/unit/test_flow_poll.py delete mode 100644 tests/unit/test_upscale_polling.py diff --git a/agent/api/flow.py b/agent/api/flow.py index 1f585ea8..7fff86a7 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -5,6 +5,7 @@ from agent.config import USE_BATCH_RPC, FLOW_PROJECT_ID, FLOW_ALLOW_DEGRADED from agent.services.flow_client import get_flow_client +from agent.services.flow_poll import annotate_polling, check_workflow_status from agent.services.omni_flash import ( check_omni_flash_status, generate_omni_flash_first_frame_video, @@ -81,6 +82,14 @@ class UpscaleVideoRequest(BaseModel): project_id: Optional[str] = None +class ExportVideoRequest(BaseModel): + media_id: str + scene_id: str = "export" + quality: Literal["1080p", "4k"] = "1080p" + aspect_ratio: str = "VIDEO_ASPECT_RATIO_LANDSCAPE" + project_id: Optional[str] = None + + class UploadImageRequest(BaseModel): file_path: str # absolute path to local image file project_id: str = "" @@ -95,6 +104,7 @@ class CheckStatusRequest(BaseModel): workflows: Optional[list[dict]] = None project_id: str = "" include_encoded_video: bool = False + mode: Literal["omni", "export"] = "omni" class CheckOmniStatusRequest(BaseModel): @@ -301,12 +311,44 @@ async def upscale_video(body: UpscaleVideoRequest): return result.get("data", result) +@router.post("/export-video") +async def export_video(body: ExportVideoRequest): + """Start Google's native Full HD/4K export and return workflow polling data.""" + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + resolution = "VIDEO_RESOLUTION_1080P" if body.quality == "1080p" else "VIDEO_RESOLUTION_4K" + result = await client.upscale_video( + media_id=body.media_id, + scene_id=body.scene_id, + aspect_ratio=body.aspect_ratio, + resolution=resolution, + project_id=body.project_id, + ) + if result.get("error") or ( + isinstance(result.get("status"), int) and result["status"] >= 400 + ): + raise HTTPException(result.get("status", 502), result.get("error", result.get("data"))) + + annotated = annotate_polling(result, mode="media_redirect") + data = annotated.get("data", annotated) + if isinstance(data, dict): + data["export"] = { + "quality": body.quality, + "resolution": resolution, + "native_flow_export": True, + "next": "/api/flow/check-status", + "check_status_mode": "export", + } + return data + + @router.post("/check-status") async def check_status(body: CheckStatusRequest): - """Check Veo operation status or Omni workflow/media status. + """Check Veo operations, Omni workflows, or native export workflows. - Veo: pass ``operations``. - Omni Flash: pass ``workflows`` from submit ``flowkitPolling.workflows``. + Veo: pass ``operations``. Workflow callers pass ``workflows``; native video + export additionally sets ``mode=\"export\"``. """ client = get_flow_client() if not client.connected: @@ -314,6 +356,16 @@ async def check_status(body: CheckStatusRequest): if body.workflows: try: + if body.mode == "export": + result = await check_workflow_status( + body.workflows, + mode="media_redirect", + include_encoded_video=body.include_encoded_video, + project_id=body.project_id, + client=client, + ) + result["download_ready"] = result.get("status") == "COMPLETED" + return result return await check_omni_flash_status( body.workflows, include_encoded_video=body.include_encoded_video, diff --git a/agent/api/upscale_status.py b/agent/api/upscale_status.py deleted file mode 100644 index 6f9c0f9f..00000000 --- a/agent/api/upscale_status.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Explicit Full HD / 4K export endpoints for Google Flow videos.""" - -from typing import Literal - -from fastapi import APIRouter, HTTPException -from pydantic import BaseModel - -from agent.services.flow_client import get_flow_client -from agent.services.upscale_polling import annotate_upscale_polling, check_upscale_status - -router = APIRouter(prefix="/flow", tags=["flow"]) - - -class ExportVideoRequest(BaseModel): - media_id: str - scene_id: str = "export" - quality: Literal["1080p", "4k"] = "1080p" - aspect_ratio: str = "VIDEO_ASPECT_RATIO_LANDSCAPE" - project_id: str | None = None - - -class CheckExportStatusRequest(BaseModel): - workflows: list[dict] - - -@router.post("/export-video") -async def export_video(body: ExportVideoRequest): - """Start Google's native Full HD/4K export. - - This is the same Flow upsample operation exposed by the UI, presented as an - export/download-quality choice. 1080p is the default because Omni Flash's - generated file is normally 720p and Full HD is the expected downloadable - master. - """ - client = get_flow_client() - if not client.connected: - raise HTTPException(503, "Extension not connected") - - resolution = ( - "VIDEO_RESOLUTION_1080P" - if body.quality == "1080p" - else "VIDEO_RESOLUTION_4K" - ) - result = await client.upscale_video( - media_id=body.media_id, - scene_id=body.scene_id, - aspect_ratio=body.aspect_ratio, - resolution=resolution, - project_id=body.project_id, - ) - if result.get("error") or ( - isinstance(result.get("status"), int) and result["status"] >= 400 - ): - raise HTTPException( - result.get("status", 502), - result.get("error", result.get("data")), - ) - - annotated = annotate_upscale_polling(result) - data = annotated.get("data", annotated) - if isinstance(data, dict): - data["export"] = { - "quality": body.quality, - "resolution": resolution, - "native_flow_export": True, - "next": "/api/flow/check-export-status", - } - return data - - -@router.post("/check-export-status") -@router.post("/check-upscale-status", include_in_schema=False) -async def check_export_status(body: CheckExportStatusRequest): - """Return a signed downloadable URL when the native Flow export is ready.""" - client = get_flow_client() - if not client.connected: - raise HTTPException(503, "Extension not connected") - try: - result = await check_upscale_status(body.workflows) - except ValueError as exc: - raise HTTPException(400, str(exc)) from exc - - if result.get("status") == "COMPLETED": - result["download_ready"] = True - else: - result["download_ready"] = False - return result diff --git a/agent/main.py b/agent/main.py index 14b30ce4..9649c807 100644 --- a/agent/main.py +++ b/agent/main.py @@ -17,7 +17,6 @@ from agent.api.scenes import router as scenes_router from agent.api.requests import router as requests_router from agent.api.flow import router as flow_router -from agent.api.upscale_status import router as upscale_status_router from agent.api.reviews import router as reviews_router from agent.api.tts import router as tts_router from agent.api.materials import router as materials_router @@ -129,7 +128,6 @@ async def lifespan(app: FastAPI): app.include_router(scenes_router, prefix="/api") app.include_router(requests_router, prefix="/api") app.include_router(flow_router, prefix="/api") -app.include_router(upscale_status_router, prefix="/api") app.include_router(reviews_router, prefix="/api") app.include_router(tts_router, prefix="/api") app.include_router(materials_router, prefix="/api") diff --git a/agent/services/flow_batch.py b/agent/services/flow_batch.py index 2071dcdf..8d89cf8a 100644 --- a/agent/services/flow_batch.py +++ b/agent/services/flow_batch.py @@ -427,6 +427,9 @@ def text_video_request(prompt: str, project_id: str, [request], _context(project_id), [_client_uuid(), 1], + ]) + + def upscale_request(media_id: str, project_id: str, aspect: Any = VIDEO_ASPECT_LANDSCAPE, model: str = "veo_3_1_upsampler_1080p") -> str: @@ -541,6 +544,8 @@ def read_upscaled_image(payload: Any) -> str: if not isinstance(encoded, str) or len(encoded) < 100: raise FlowBatchError("image upscale response carried no encoded image") return encoded + + def read_upscaled_media_id(payload: Any) -> str: """Return the media id created by the p0UkFb upscale submit.""" for text in _walk_strings(payload): diff --git a/agent/services/flow_client.py b/agent/services/flow_client.py index ec66a1ae..f3104b1a 100644 --- a/agent/services/flow_client.py +++ b/agent/services/flow_client.py @@ -20,6 +20,7 @@ import time import uuid from typing import Optional +from urllib.parse import quote from agent.config import ( GOOGLE_FLOW_API, GOOGLE_API_KEY, ENDPOINTS, @@ -1035,6 +1036,54 @@ async def get_media(self, media_id: str) -> dict: data["image"] = {"fifeUrl": urls.image} return {"status": 200, "data": data} + async def resolve_media_url(self, media_id: str) -> dict: + """Resolve a signed Flow media URL without downloading the media body.""" + if USE_BATCH_RPC: + result = await self.get_media(media_id) + data = result.get("data") if isinstance(result.get("data"), dict) else {} + video = data.get("video") if isinstance(data, dict) else None + candidate = video.get("fifeUrl") if isinstance(video, dict) else None + return { + "status": result.get("status", 200), + "data": { + "url": candidate, + "contentType": "video/mp4" if candidate else None, + }, + "error": result.get("error"), + } + + url = ( + "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" + f"?name={quote(media_id, safe='')}" + ) + return await self._send( + "trpc_request", + { + "url": url, + "method": "GET", + "headers": {"content-type": "application/json"}, + "responseMode": "url", + }, + timeout=15, + ) + + async def get_project_initial_data(self, project_id: str) -> dict: + """Fetch the authenticated legacy Flow project snapshot used for polling.""" + query = quote( + json.dumps({"json": {"projectId": project_id}}, separators=(",", ":")), + safe="", + ) + url = f"https://labs.google/fx/api/trpc/flow.projectInitialData?input={query}" + return await self._send( + "trpc_request", + { + "url": url, + "method": "GET", + "headers": {"content-type": "application/json"}, + }, + timeout=15, + ) + async def upload_image(self, image_base64: str, mime_type: str = "image/jpeg", project_id: str = "", file_name: str = "image.jpg") -> dict: """Upload an image into the project so it can be used as a reference.""" diff --git a/agent/services/flow_poll.py b/agent/services/flow_poll.py new file mode 100644 index 00000000..3d3e9571 --- /dev/null +++ b/agent/services/flow_poll.py @@ -0,0 +1,314 @@ +"""Shared polling helpers for Flow workflow-backed video jobs. + +Workflow-producing surfaces (Omni and native video export) return the same +logical pair: a workflow name plus a primary media id. This module owns the +normalization, response annotation and one-pass polling so each feature does +not grow its own copy of transport/status handling. +""" +from __future__ import annotations + +from agent.config import USE_BATCH_RPC +from agent.services.flow_client import get_flow_client + +_ALLOWED_MEDIA_URL_PREFIX = "https://flow-content.google/" +_POLL_MODES = {"project_media", "batch_media", "media_redirect"} + + +def normalize_workflow(workflow: dict) -> dict | None: + """Normalize a raw Flow workflow or a FlowKit polling descriptor.""" + if not isinstance(workflow, dict): + return None + name = workflow.get("name") + primary_media_id = workflow.get("primary_media_id") + if not primary_media_id: + metadata = workflow.get("metadata") + if isinstance(metadata, dict): + primary_media_id = metadata.get("primaryMediaId") + if not isinstance(name, str) or not name: + return None + if not isinstance(primary_media_id, str) or not primary_media_id: + return None + item = {"name": name, "primary_media_id": primary_media_id} + project_id = workflow.get("project_id") or workflow.get("projectId") + if isinstance(project_id, str) and project_id: + item["project_id"] = project_id + return item + + +def extract_workflows(result: dict) -> list[dict]: + """Extract normalized workflow descriptors from a submit response.""" + if not isinstance(result, dict): + return [] + data = result.get("data") if isinstance(result.get("data"), dict) else result + workflows = data.get("workflows", []) if isinstance(data, dict) else [] + return [item for workflow in workflows if (item := normalize_workflow(workflow))] + + +def annotate_polling(result: dict, *, mode: str, project_id: str = "") -> dict: + """Attach the descriptor callers can pass back to ``/flow/check-status``.""" + if mode not in _POLL_MODES: + raise ValueError(f"Unknown Flow polling mode: {mode}") + workflows = extract_workflows(result) + if not workflows: + return result + if project_id: + for workflow in workflows: + workflow["project_id"] = project_id + data = result.get("data") if isinstance(result.get("data"), dict) else result + if isinstance(data, dict): + descriptor = {"mode": mode, "workflows": workflows} + if project_id: + descriptor["project_id"] = project_id + data["flowkitPolling"] = descriptor + return result + + +def _normalize_workflows(workflows: list[dict]) -> list[dict]: + return [item for workflow in (workflows or []) if (item := normalize_workflow(workflow))] + + +def _parse_media_redirect(response: dict) -> tuple[str | None, str | None, str | None]: + if not isinstance(response, dict): + return None, None, "Flow media redirect returned an invalid response" + status = response.get("status") + data = response.get("data") if isinstance(response.get("data"), dict) else {} + candidate = data.get("url") + content_type = data.get("contentType") + if ( + isinstance(status, int) + and status < 400 + and isinstance(candidate, str) + and candidate.startswith(_ALLOWED_MEDIA_URL_PREFIX) + ): + return candidate, content_type if isinstance(content_type, str) else None, None + error = response.get("error") + if not error and isinstance(status, int) and status >= 400: + error = f"API_{status}" + return None, content_type if isinstance(content_type, str) else None, str(error or "media redirect not ready") + + +async def _check_media_ready( + workflows: list[dict], + *, + client, + include_encoded_video: bool, + project_id: str, +) -> dict: + resolved_project_id = project_id or next( + (item.get("project_id", "") for item in workflows if item.get("project_id")), "" + ) + items = [] + for workflow in workflows: + media_id = workflow["primary_media_id"] + response = await client.resolve_media_url(media_id) + url, content_type, diagnostic = _parse_media_redirect(response) + item_project = workflow.get("project_id") or resolved_project_id + if url: + media = { + "media_id": media_id, + "url": url, + "encoded_video_available": False, + "resolved_via": "as29s" if USE_BATCH_RPC else "media.getMediaUrlRedirect", + } + if content_type: + media["content_type"] = content_type + if include_encoded_video: + media["encoded_video"] = None + item = { + "name": workflow["name"], + "primary_media_id": media_id, + "done": True, + "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", + "error": None, + "media": media, + } + else: + item = { + "name": workflow["name"], + "primary_media_id": media_id, + "done": False, + "status": "PENDING", + "error": None, + } + if diagnostic: + item["probe"] = {"diagnostic": diagnostic} + if isinstance(response, dict) and isinstance(response.get("status"), int): + item["probe"]["http_status"] = response["status"] + if item_project: + item["project_id"] = item_project + items.append(item) + + all_done = bool(items) and all(item["done"] for item in items) + result = { + "done": all_done, + "status": "COMPLETED" if all_done else "PENDING", + "workflows": items, + } + if resolved_project_id: + result["project_id"] = resolved_project_id + return result + + +async def _check_project_media( + workflows: list[dict], + *, + client, + include_encoded_video: bool, + project_id: str, +) -> dict: + resolved_project_id = project_id or next( + (item.get("project_id", "") for item in workflows if item.get("project_id")), "" + ) + if not resolved_project_id: + raise ValueError( + "Flow project polling requires project_id. Use the project_id returned " + "inside flowkitPolling or pass project_id explicitly." + ) + if any( + item.get("project_id") and item["project_id"] != resolved_project_id + for item in workflows + ): + raise ValueError("All workflows in one project poll must belong to the same project_id") + + response = await client.get_project_initial_data(resolved_project_id) + http_status = response.get("status") if isinstance(response, dict) else None + if isinstance(http_status, int) and http_status >= 400: + data = response.get("data") if isinstance(response.get("data"), dict) else {} + error = data.get("error") if isinstance(data, dict) else None + if isinstance(error, dict): + error = error.get("message") or error.get("code") + raise RuntimeError(error or response.get("error") or f"Flow project poll failed: API_{http_status}") + + envelope = response.get("data") if isinstance(response, dict) else None + result = envelope.get("result") if isinstance(envelope, dict) else None + result_data = result.get("data") if isinstance(result, dict) else None + project_json = result_data.get("json") if isinstance(result_data, dict) else None + contents = project_json.get("projectContents") if isinstance(project_json, dict) else None + if not isinstance(contents, dict): + raise RuntimeError("Flow project poll returned an unexpected response shape") + + project_workflows = contents.get("workflows") + project_media = contents.get("media") + project_workflows = project_workflows if isinstance(project_workflows, list) else [] + project_media = project_media if isinstance(project_media, list) else [] + known_workflow_names = { + item.get("name") for item in project_workflows + if isinstance(item, dict) and isinstance(item.get("name"), str) + } + media_by_id = { + item.get("name"): item for item in project_media + if isinstance(item, dict) and isinstance(item.get("name"), str) + } + media_by_workflow = { + item.get("workflowId"): item for item in project_media + if isinstance(item, dict) and isinstance(item.get("workflowId"), str) + } + + items = [] + for workflow in workflows: + name = workflow["name"] + media_id = workflow["primary_media_id"] + payload = media_by_id.get(media_id) or media_by_workflow.get(name) + if not isinstance(payload, dict): + items.append({ + "name": name, + "primary_media_id": media_id, + "project_id": resolved_project_id, + "done": False, + "status": "PENDING", + "error": None, + "workflow_present": name in known_workflow_names, + }) + continue + + metadata = payload.get("mediaMetadata") + metadata = metadata if isinstance(metadata, dict) else {} + media_status = metadata.get("mediaStatus") + media_status = media_status if isinstance(media_status, dict) else {} + generation_status = media_status.get("mediaGenerationStatus") + if isinstance(generation_status, str) and ( + generation_status.endswith("FAILED") or generation_status.endswith("CANCELLED") + ): + items.append({ + "name": name, + "primary_media_id": media_id, + "project_id": resolved_project_id, + "done": True, + "status": "FAILED", + "error": generation_status, + }) + continue + if generation_status != "MEDIA_GENERATION_STATUS_SUCCESSFUL": + items.append({ + "name": name, + "primary_media_id": media_id, + "project_id": resolved_project_id, + "done": False, + "status": "PENDING", + "error": None, + }) + continue + + url_response = await client.resolve_media_url(media_id) + url, _, url_error = _parse_media_redirect(url_response) + media = { + "media_id": media_id, + "url": url, + "encoded_video_available": False, + } + if include_encoded_video: + media["encoded_video"] = None + if url_error: + media["url_error"] = url_error + items.append({ + "name": name, + "primary_media_id": media_id, + "project_id": resolved_project_id, + "done": True, + "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", + "error": None, + "media": media, + }) + + all_done = bool(items) and all(item["done"] for item in items) + any_failed = any(item.get("status") == "FAILED" for item in items) + return { + "project_id": resolved_project_id, + "done": all_done, + "status": "FAILED" if any_failed else ("COMPLETED" if all_done else "PENDING"), + "workflows": items, + } + + +async def check_workflow_status( + workflows: list[dict], + *, + mode: str, + include_encoded_video: bool = False, + project_id: str = "", + client=None, +) -> dict: + """Perform one non-blocking poll pass for a workflow-backed Flow job.""" + if mode not in _POLL_MODES: + raise ValueError(f"Unknown Flow polling mode: {mode}") + normalized = _normalize_workflows(workflows) + if not normalized: + raise ValueError( + "Flow polling requires workflow descriptors with name and primary_media_id " + "(or raw Flow metadata.primaryMediaId)" + ) + + client = client or get_flow_client() + if mode in {"batch_media", "media_redirect"}: + return await _check_media_ready( + normalized, + client=client, + include_encoded_video=include_encoded_video, + project_id=project_id, + ) + return await _check_project_media( + normalized, + client=client, + include_encoded_video=include_encoded_video, + project_id=project_id, + ) diff --git a/agent/services/omni_flash.py b/agent/services/omni_flash.py index 796d6a31..8c53f91b 100644 --- a/agent/services/omni_flash.py +++ b/agent/services/omni_flash.py @@ -24,11 +24,10 @@ import time import uuid from pathlib import Path -from urllib.parse import quote - from agent.config import USE_BATCH_RPC from agent.services import flow_batch as fb from agent.services.flow_client import get_flow_client +from agent.services.flow_poll import annotate_polling, check_workflow_status, extract_workflows from agent.services.headers import random_headers _MODELS_FILE = Path(__file__).parent.parent / "models.json" @@ -62,39 +61,13 @@ def _batch_path_blocks_omni() -> dict | None: async def _fetch_project_initial_data(client, project_id: str) -> dict: - """Fetch the same authenticated project snapshot used by the Flow UI.""" - query = quote( - json.dumps({"json": {"projectId": project_id}}, separators=(",", ":")), - safe="", - ) - url = f"https://labs.google/fx/api/trpc/flow.projectInitialData?input={query}" - return await client._send( - "trpc_request", - { - "url": url, - "method": "GET", - "headers": {"content-type": "application/json"}, - }, - timeout=15, - ) + """Compatibility wrapper around FlowClient's public project poll surface.""" + return await client.get_project_initial_data(project_id) async def _fetch_media_url(client, media_id: str) -> dict: - """Resolve Flow's authenticated media redirect without buffering the file.""" - url = ( - "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" - f"?name={quote(media_id, safe='')}" - ) - return await client._send( - "trpc_request", - { - "url": url, - "method": "GET", - "headers": {"content-type": "application/json"}, - "responseMode": "url", - }, - timeout=15, - ) + """Compatibility wrapper around FlowClient's public URL resolver.""" + return await client.resolve_media_url(media_id) def _validate_duration(duration_s: int) -> None: @@ -166,56 +139,14 @@ def _validate_frame_inputs( raise ValueError("Omni Flash First+Last requires a non-empty end_image_media_id") -def _normalize_workflow(workflow: dict) -> dict | None: - """Normalize a raw Flow workflow or FlowKit polling descriptor.""" - if not isinstance(workflow, dict): - return None - name = workflow.get("name") - primary_media_id = workflow.get("primary_media_id") - if not primary_media_id: - metadata = workflow.get("metadata") - if isinstance(metadata, dict): - primary_media_id = metadata.get("primaryMediaId") - if not isinstance(name, str) or not name: - return None - if not isinstance(primary_media_id, str) or not primary_media_id: - return None - item = {"name": name, "primary_media_id": primary_media_id} - project_id = workflow.get("project_id") or workflow.get("projectId") - if isinstance(project_id, str) and project_id: - item["project_id"] = project_id - return item - - def extract_omni_workflows(result: dict) -> list[dict]: - """Extract ``name`` + ``primaryMediaId`` pairs from an Omni submit.""" - if not isinstance(result, dict): - return [] - data = result.get("data") if isinstance(result.get("data"), dict) else result - workflows = data.get("workflows", []) if isinstance(data, dict) else [] - normalized = [] - for workflow in workflows: - item = _normalize_workflow(workflow) - if item: - normalized.append(item) - return normalized + """Backward-compatible Omni name for the shared workflow extractor.""" + return extract_workflows(result) def _annotate_polling(result: dict, project_id: str) -> dict: - """Add an explicit FlowKit polling descriptor to a successful submit.""" - workflows = extract_omni_workflows(result) - if not workflows: - return result - data = result.get("data") if isinstance(result.get("data"), dict) else result - if isinstance(data, dict): - for workflow in workflows: - workflow["project_id"] = project_id - data["flowkitPolling"] = { - "mode": "project_media", - "project_id": project_id, - "workflows": workflows, - } - return result + """Annotate legacy Omni submits using the shared polling descriptor.""" + return annotate_polling(result, mode="project_media", project_id=project_id) async def generate_omni_flash_text_video( @@ -447,226 +378,16 @@ async def generate_omni_flash_video( return _annotate_polling(result, project_id) -async def _check_omni_batch_media( - workflows: list[dict], - include_encoded_video: bool = False, - project_id: str = "", -) -> dict: - normalized = [item for workflow in (workflows or []) if (item := _normalize_workflow(workflow))] - if not normalized: - raise ValueError("Omni polling requires workflow descriptors with name and primary_media_id") - resolved_project_id = project_id or next( - (item.get("project_id", "") for item in normalized if item.get("project_id")), "") - client = get_flow_client() - items = [] - for workflow in normalized: - media_id = workflow["primary_media_id"] - response = await client.get_media(media_id) - data = response.get("data") if isinstance(response, dict) else None - video = data.get("video") if isinstance(data, dict) else None - url = video.get("fifeUrl") if isinstance(video, dict) else None - if isinstance(url, str) and url.startswith("https://flow-content.google/video/"): - media = { - "media_id": media_id, - "url": url, - "encoded_video_available": False, - "resolved_via": "as29s", - } - if include_encoded_video: - media["encoded_video"] = None - items.append({ - "name": workflow["name"], - "primary_media_id": media_id, - "project_id": workflow.get("project_id") or resolved_project_id, - "done": True, - "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", - "error": None, - "media": media, - }) - else: - items.append({ - "name": workflow["name"], - "primary_media_id": media_id, - "project_id": workflow.get("project_id") or resolved_project_id, - "done": False, - "status": "PENDING", - "error": None, - }) - all_done = bool(items) and all(item["done"] for item in items) - return { - "project_id": resolved_project_id or None, - "done": all_done, - "status": "COMPLETED" if all_done else "PENDING", - "workflows": items, - } - - async def check_omni_flash_status( workflows: list[dict], include_encoded_video: bool = False, project_id: str = "", ) -> dict: - """Perform one non-blocking poll pass for Omni workflow-backed jobs.""" - if USE_BATCH_RPC: - return await _check_omni_batch_media(workflows, include_encoded_video, project_id) - normalized = [] - for workflow in workflows or []: - item = _normalize_workflow(workflow) - if item: - normalized.append(item) - if not normalized: - raise ValueError( - "Omni polling requires workflow descriptors with name and primary_media_id " - "(or raw Flow metadata.primaryMediaId)" - ) - - resolved_project_id = project_id or next( - (item.get("project_id", "") for item in normalized if item.get("project_id")), - "", + """Perform one non-blocking poll pass using the shared Flow poller.""" + return await check_workflow_status( + workflows, + mode="batch_media" if USE_BATCH_RPC else "project_media", + include_encoded_video=include_encoded_video, + project_id=project_id, + client=get_flow_client(), ) - if not resolved_project_id: - raise ValueError( - "Omni project polling requires project_id. Use the project_id returned " - "inside flowkitPolling or pass project_id explicitly." - ) - if any( - item.get("project_id") and item["project_id"] != resolved_project_id - for item in normalized - ): - raise ValueError( - "All Omni workflows in one poll must belong to the same project_id" - ) - - client = get_flow_client() - response = await _fetch_project_initial_data(client, resolved_project_id) - http_status = response.get("status") if isinstance(response, dict) else None - if isinstance(http_status, int) and http_status >= 400: - data = response.get("data") if isinstance(response.get("data"), dict) else {} - error = data.get("error") if isinstance(data, dict) else None - if isinstance(error, dict): - error = error.get("message") or error.get("code") - raise RuntimeError( - error - or response.get("error") - or f"Flow project poll failed: API_{http_status}" - ) - - envelope = response.get("data") if isinstance(response, dict) else None - result = envelope.get("result") if isinstance(envelope, dict) else None - result_data = result.get("data") if isinstance(result, dict) else None - project_json = result_data.get("json") if isinstance(result_data, dict) else None - contents = project_json.get("projectContents") if isinstance(project_json, dict) else None - if not isinstance(contents, dict): - raise RuntimeError("Flow project poll returned an unexpected response shape") - - project_workflows = contents.get("workflows") - project_media = contents.get("media") - project_workflows = project_workflows if isinstance(project_workflows, list) else [] - project_media = project_media if isinstance(project_media, list) else [] - known_workflow_names = { - item.get("name") - for item in project_workflows - if isinstance(item, dict) and isinstance(item.get("name"), str) - } - media_by_id = { - item.get("name"): item - for item in project_media - if isinstance(item, dict) and isinstance(item.get("name"), str) - } - media_by_workflow = { - item.get("workflowId"): item - for item in project_media - if isinstance(item, dict) and isinstance(item.get("workflowId"), str) - } - - items = [] - - for workflow in normalized: - name = workflow["name"] - media_id = workflow["primary_media_id"] - payload = media_by_id.get(media_id) or media_by_workflow.get(name) - if not isinstance(payload, dict): - items.append({ - "name": name, - "primary_media_id": media_id, - "project_id": resolved_project_id, - "done": False, - "status": "PENDING", - "error": None, - "workflow_present": name in known_workflow_names, - }) - continue - - metadata = payload.get("mediaMetadata") - metadata = metadata if isinstance(metadata, dict) else {} - media_status = metadata.get("mediaStatus") - media_status = media_status if isinstance(media_status, dict) else {} - generation_status = media_status.get("mediaGenerationStatus") - - if isinstance(generation_status, str) and ( - generation_status.endswith("FAILED") or generation_status.endswith("CANCELLED") - ): - items.append({ - "name": name, - "primary_media_id": media_id, - "project_id": resolved_project_id, - "done": True, - "status": "FAILED", - "error": generation_status, - }) - continue - - if generation_status != "MEDIA_GENERATION_STATUS_SUCCESSFUL": - items.append({ - "name": name, - "primary_media_id": media_id, - "project_id": resolved_project_id, - "done": False, - "status": "PENDING", - "error": None, - }) - continue - - url = None - url_error = None - url_response = await _fetch_media_url(client, media_id) - if isinstance(url_response, dict) and url_response.get("status", 500) < 400: - url_data = url_response.get("data") - candidate = url_data.get("url") if isinstance(url_data, dict) else None - if isinstance(candidate, str) and candidate.startswith("https://flow-content.google/"): - url = candidate - else: - url_error = "Flow media redirect returned no allowed URL" - else: - url_error = ( - url_response.get("error") - if isinstance(url_response, dict) - else "Flow media redirect failed" - ) - item = { - "name": name, - "primary_media_id": media_id, - "project_id": resolved_project_id, - "done": True, - "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", - "error": None, - "media": { - "media_id": media_id, - "url": url, - "encoded_video_available": False, - }, - } - if include_encoded_video: - item["media"]["encoded_video"] = None - if url_error: - item["media"]["url_error"] = url_error - items.append(item) - - all_done = bool(items) and all(item["done"] for item in items) - any_failed = any(item.get("status") == "FAILED" for item in items) - return { - "project_id": resolved_project_id, - "done": all_done, - "status": "FAILED" if any_failed else ("COMPLETED" if all_done else "PENDING"), - "workflows": items, - } diff --git a/agent/services/upscale_polling.py b/agent/services/upscale_polling.py deleted file mode 100644 index 29d705e5..00000000 --- a/agent/services/upscale_polling.py +++ /dev/null @@ -1,185 +0,0 @@ -"""Headless polling for Flow video upscales. - -Google Flow's upsampler returns workflow descriptors whose logical -``primaryMediaId`` may not appear in ``flow.projectInitialData``. The browser UI -can still resolve completed media through ``media.getMediaUrlRedirect``. - -This module exposes a small active poller that treats a successful authenticated -media redirect as the completion signal, avoiding the legacy -``batchCheckAsyncVideoGenerationStatus`` and ``/v1/media/{id}`` paths. -""" - -from __future__ import annotations - -from urllib.parse import quote - -from agent.config import USE_BATCH_RPC -from agent.services.flow_client import get_flow_client - -_ALLOWED_MEDIA_URL_PREFIX = "https://flow-content.google/" - - -def _normalize_workflow(workflow: dict) -> dict | None: - 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_upscale_workflows(result: dict) -> list[dict]: - 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_upscale_polling(result: dict) -> dict: - workflows = extract_upscale_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": "media_redirect", - "workflows": workflows, - } - return result - - -async def _fetch_media_url(client, media_id: str) -> dict: - if USE_BATCH_RPC: - result = await client.get_media(media_id) - data = result.get("data") if isinstance(result.get("data"), dict) else {} - video = data.get("video") if isinstance(data, dict) else None - candidate = video.get("fifeUrl") if isinstance(video, dict) else None - return { - "status": result.get("status", 200), - "data": { - "url": candidate, - "contentType": "video/mp4" if candidate else None, - }, - "error": result.get("error"), - } - - url = ( - "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" - f"?name={quote(media_id, safe='')}" - ) - return await client._send( - "trpc_request", - { - "url": url, - "method": "GET", - "headers": {"content-type": "application/json"}, - "responseMode": "url", - }, - timeout=15, - ) - - -def _parse_media_redirect(response: dict) -> tuple[str | None, str | None, str | None]: - if not isinstance(response, dict): - return None, None, "Flow media redirect returned an invalid response" - status = response.get("status") - data = response.get("data") if isinstance(response.get("data"), dict) else {} - candidate = data.get("url") - content_type = data.get("contentType") - if ( - isinstance(status, int) - and status < 400 - and isinstance(candidate, str) - and candidate.startswith(_ALLOWED_MEDIA_URL_PREFIX) - ): - return candidate, content_type if isinstance(content_type, str) else None, None - error = response.get("error") - if not error and isinstance(status, int) and status >= 400: - error = f"API_{status}" - if not error: - error = "media redirect not ready" - return None, content_type if isinstance(content_type, str) else None, str(error) - - -async def check_upscale_status( - workflows: list[dict], - include_encoded_video: bool = False, -) -> dict: - """Poll native Flow Full HD/4K export workflows without buffering the MP4.""" - normalized = [] - for workflow in workflows or []: - item = _normalize_workflow(workflow) - if item: - normalized.append(item) - if not normalized: - raise ValueError( - "Export polling requires workflow descriptors with name and " - "primary_media_id (or raw Flow metadata.primaryMediaId)" - ) - - client = get_flow_client() - items = [] - for workflow in normalized: - media_id = workflow["primary_media_id"] - response = await _fetch_media_url(client, media_id) - url, content_type, diagnostic = _parse_media_redirect(response) - if url: - media = { - "media_id": media_id, - "url": url, - "encoded_video_available": False, - "resolved_via": "as29s" if USE_BATCH_RPC else "media.getMediaUrlRedirect", - } - if content_type: - media["content_type"] = content_type - if include_encoded_video: - media["encoded_video"] = None - items.append({ - "name": workflow["name"], - "primary_media_id": media_id, - "done": True, - "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", - "error": None, - "media": media, - }) - continue - - probe = {} - if isinstance(response, dict): - if isinstance(response.get("status"), int): - probe["http_status"] = response["status"] - data = response.get("data") - if isinstance(data, dict) and isinstance(data.get("url"), str): - probe["resolved_url"] = data["url"] - if diagnostic: - probe["diagnostic"] = diagnostic - item = { - "name": workflow["name"], - "primary_media_id": media_id, - "done": False, - "status": "PENDING", - "error": None, - } - if probe: - item["probe"] = probe - items.append(item) - - all_done = bool(items) and all(item["done"] for item in items) - return { - "done": all_done, - "status": "COMPLETED" if all_done else "PENDING", - "workflows": items, - } diff --git a/docs/VIDEO_EXPORTS.md b/docs/VIDEO_EXPORTS.md index de056793..9071660a 100644 --- a/docs/VIDEO_EXPORTS.md +++ b/docs/VIDEO_EXPORTS.md @@ -28,9 +28,9 @@ curl -sS -X POST http://127.0.0.1:8100/api/flow/export-video \ The response contains `flowkitPolling.workflows`. Poll those descriptors: ```bash -curl -sS -X POST http://127.0.0.1:8100/api/flow/check-export-status \ +curl -sS -X POST http://127.0.0.1:8100/api/flow/check-status \ -H 'Content-Type: application/json' \ - -d '{"workflows": }' + -d '{"mode":"export","workflows": }' ``` When `download_ready` becomes `true`, use @@ -39,12 +39,4 @@ When `download_ready` becomes `true`, use ## Compatibility -The older endpoints remain supported: - -- `POST /api/flow/upscale-video` -- `POST /api/flow/check-upscale-status` - -They are aliases/low-level surfaces for the same Google Flow capability. New -integrations should use `export-video` and `check-export-status`, because those -names describe the user-visible operation: selecting the downloadable output -quality. +`POST /api/flow/upscale-video` remains the low-level submit surface. New integrations should use `export-video`, then poll the returned workflows through the existing `POST /api/flow/check-status` endpoint with `mode: "export"`. This keeps all Flow polling behind one API surface. diff --git a/tests/unit/test_flow_poll.py b/tests/unit/test_flow_poll.py new file mode 100644 index 00000000..ea31c9aa --- /dev/null +++ b/tests/unit/test_flow_poll.py @@ -0,0 +1,134 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import agent.services.flow_poll as polling +import agent.services.flow_client as flow_client_module +from agent.services.flow_client import FlowClient +from agent.services.flow_poll import annotate_polling, check_workflow_status + +MEDIA = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + + +def test_annotate_export_polling_uses_workflow_primary_media_id(): + result = { + "status": 200, + "data": { + "workflows": [ + {"name": "wf-1", "metadata": {"primaryMediaId": MEDIA + "_upsampled"}} + ] + }, + } + annotated = annotate_polling(result, mode="media_redirect") + assert annotated["data"]["flowkitPolling"] == { + "mode": "media_redirect", + "workflows": [ + {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} + ], + } + + +@pytest.mark.asyncio +async def test_batch_poll_resolves_completed_media_through_public_resolver(monkeypatch): + monkeypatch.setattr(polling, "USE_BATCH_RPC", True) + client = MagicMock() + client.resolve_media_url = AsyncMock(return_value={ + "status": 200, + "data": { + "url": "https://flow-content.google/video/out?Signature=test", + "contentType": "video/mp4", + }, + }) + with patch("agent.services.flow_poll.get_flow_client", return_value=client): + result = await check_workflow_status( + [{"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"}], + mode="media_redirect", + ) + + assert result["done"] is True + assert result["status"] == "COMPLETED" + assert result["workflows"][0]["media"]["resolved_via"] == "as29s" + client.resolve_media_url.assert_awaited_once_with(MEDIA + "_upsampled") + + +@pytest.mark.asyncio +async def test_batch_poll_stays_pending_until_video_url_exists(monkeypatch): + monkeypatch.setattr(polling, "USE_BATCH_RPC", True) + client = MagicMock() + client.resolve_media_url = AsyncMock(return_value={ + "status": 404, + "error": "not ready", + "data": {}, + }) + with patch("agent.services.flow_poll.get_flow_client", return_value=client): + result = await check_workflow_status( + [{"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"}], + mode="media_redirect", + ) + + assert result["done"] is False + assert result["status"] == "PENDING" + assert result["workflows"][0]["status"] == "PENDING" + + +def test_unknown_poll_mode_is_rejected(): + with pytest.raises(ValueError, match="Unknown Flow polling mode"): + annotate_polling({"data": {"workflows": []}}, mode="mystery") + + +@pytest.mark.asyncio +async def test_flow_client_project_snapshot_keeps_trpc_wire_contract(monkeypatch): + monkeypatch.setattr(flow_client_module, "USE_BATCH_RPC", False) + client = FlowClient() + client._send = AsyncMock(return_value={"status": 200, "data": {}}) + + await client.get_project_initial_data("project-1") + + client._send.assert_awaited_once_with( + "trpc_request", + { + "url": ( + "https://labs.google/fx/api/trpc/flow.projectInitialData?input=" + "%7B%22json%22%3A%7B%22projectId%22%3A%22project-1%22%7D%7D" + ), + "method": "GET", + "headers": {"content-type": "application/json"}, + }, + timeout=15, + ) + + +@pytest.mark.asyncio +async def test_flow_client_media_resolver_keeps_legacy_redirect_contract(monkeypatch): + monkeypatch.setattr(flow_client_module, "USE_BATCH_RPC", False) + client = FlowClient() + client._send = AsyncMock(return_value={"status": 200, "data": {}}) + + await client.resolve_media_url("media-1") + + client._send.assert_awaited_once_with( + "trpc_request", + { + "url": "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect?name=media-1", + "method": "GET", + "headers": {"content-type": "application/json"}, + "responseMode": "url", + }, + timeout=15, + ) + + +@pytest.mark.asyncio +async def test_flow_client_media_resolver_uses_get_media_on_batch(monkeypatch): + monkeypatch.setattr(flow_client_module, "USE_BATCH_RPC", True) + client = FlowClient() + client.get_media = AsyncMock(return_value={ + "status": 200, + "data": {"video": {"fifeUrl": "https://flow-content.google/video/media-1?sig=x"}}, + }) + + result = await client.resolve_media_url("media-1") + + assert result["data"]["url"].startswith("https://flow-content.google/video/") + assert result["data"]["contentType"] == "video/mp4" + client.get_media.assert_awaited_once_with("media-1") diff --git a/tests/unit/test_omni_flash.py b/tests/unit/test_omni_flash.py index cf3b2fb1..036ee267 100644 --- a/tests/unit/test_omni_flash.py +++ b/tests/unit/test_omni_flash.py @@ -342,55 +342,35 @@ def _project_response(generation_status="MEDIA_GENERATION_STATUS_PENDING", inclu @pytest.mark.asyncio -async def test_project_poll_fetch_matches_live_flow_trpc_contract(): +async def test_project_poll_fetch_uses_public_flow_client_surface(): client = MagicMock() - client._send = AsyncMock(return_value={"status": 200, "data": {}}) + client.get_project_initial_data = AsyncMock(return_value={"status": 200, "data": {}}) await _fetch_project_initial_data(client, "project-1") - client._send.assert_awaited_once_with( - "trpc_request", - { - "url": ( - "https://labs.google/fx/api/trpc/flow.projectInitialData?input=" - "%7B%22json%22%3A%7B%22projectId%22%3A%22project-1%22%7D%7D" - ), - "method": "GET", - "headers": {"content-type": "application/json"}, - }, - timeout=15, - ) + client.get_project_initial_data.assert_awaited_once_with("project-1") @pytest.mark.asyncio -async def test_media_redirect_fetch_requests_url_only_mode(): +async def test_media_redirect_fetch_uses_public_flow_client_surface(): client = MagicMock() - client._send = AsyncMock(return_value={"status": 200, "data": {}}) + client.resolve_media_url = AsyncMock(return_value={"status": 200, "data": {}}) await _fetch_media_url(client, "media-1") - client._send.assert_awaited_once_with( - "trpc_request", - { - "url": ( - "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" - "?name=media-1" - ), - "method": "GET", - "headers": {"content-type": "application/json"}, - "responseMode": "url", - }, - timeout=15, - ) + client.resolve_media_url.assert_awaited_once_with("media-1") @pytest.mark.asyncio async def test_batch_omni_poll_uses_as29s_media(monkeypatch): monkeypatch.setattr(omni_flash, "USE_BATCH_RPC", True) client = MagicMock() - client.get_media = AsyncMock(return_value={ + client.resolve_media_url = AsyncMock(return_value={ "status": 200, - "data": {"video": {"fifeUrl": "https://flow-content.google/video/media-1?Signature=test"}}, + "data": { + "url": "https://flow-content.google/video/media-1?Signature=test", + "contentType": "video/mp4", + }, }) with patch("agent.services.omni_flash.get_flow_client", return_value=client): result = await check_omni_flash_status([{ @@ -398,14 +378,15 @@ async def test_batch_omni_poll_uses_as29s_media(monkeypatch): }]) assert result["done"] is True assert result["workflows"][0]["media"]["resolved_via"] == "as29s" - client.get_media.assert_awaited_once_with("media-1") + client.resolve_media_url.assert_awaited_once_with("media-1") @pytest.mark.asyncio async def test_omni_poll_pending_uses_project_snapshot_not_legacy_transports(): client = MagicMock() - client._send = AsyncMock(return_value=_project_response()) - client.get_media = AsyncMock(side_effect=AssertionError("legacy get_media forbidden")) + client.get_project_initial_data = AsyncMock(return_value=_project_response()) + client.resolve_media_url = AsyncMock(side_effect=AssertionError("URL lookup should wait for success")) + client.get_media = AsyncMock(side_effect=AssertionError("direct media lookup forbidden")) client.check_video_status = AsyncMock(side_effect=AssertionError("operation poll forbidden")) with patch("agent.services.omni_flash.get_flow_client", return_value=client): @@ -419,9 +400,10 @@ async def test_omni_poll_pending_uses_project_snapshot_not_legacy_transports(): ] ) + client.get_project_initial_data.assert_awaited_once_with("project-1") + client.resolve_media_url.assert_not_awaited() client.get_media.assert_not_awaited() client.check_video_status.assert_not_awaited() - client._send.assert_awaited_once() assert result["done"] is False assert result["status"] == "PENDING" assert result["workflows"][0]["status"] == "PENDING" @@ -430,17 +412,15 @@ async def test_omni_poll_pending_uses_project_snapshot_not_legacy_transports(): @pytest.mark.asyncio async def test_omni_poll_completed_returns_signed_url_without_buffering_video(): client = MagicMock() - client._send = AsyncMock( - side_effect=[ - _project_response("MEDIA_GENERATION_STATUS_SUCCESSFUL"), - { - "status": 200, - "data": { - "url": "https://flow-content.google/video/media-1?Signature=test" - }, - }, - ] + client.get_project_initial_data = AsyncMock( + return_value=_project_response("MEDIA_GENERATION_STATUS_SUCCESSFUL") ) + client.resolve_media_url = AsyncMock(return_value={ + "status": 200, + "data": { + "url": "https://flow-content.google/video/media-1?Signature=test" + }, + }) with patch("agent.services.omni_flash.get_flow_client", return_value=client): result = await check_omni_flash_status( @@ -456,12 +436,17 @@ async def test_omni_poll_completed_returns_signed_url_without_buffering_video(): assert item["media"]["url"].startswith("https://flow-content.google/video/") assert item["media"]["encoded_video_available"] is False assert "encoded_video" not in item["media"] + client.get_project_initial_data.assert_awaited_once_with("project-1") + client.resolve_media_url.assert_awaited_once_with("media-1") @pytest.mark.asyncio async def test_omni_poll_missing_media_is_pending(): client = MagicMock() - client._send = AsyncMock(return_value=_project_response(include_media=False)) + client.get_project_initial_data = AsyncMock( + return_value=_project_response(include_media=False) + ) + client.resolve_media_url = AsyncMock(side_effect=AssertionError("missing media must stay pending")) with patch("agent.services.omni_flash.get_flow_client", return_value=client): result = await check_omni_flash_status( @@ -471,6 +456,8 @@ async def test_omni_poll_missing_media_is_pending(): assert result["done"] is False assert result["workflows"][0]["status"] == "PENDING" + client.get_project_initial_data.assert_awaited_once_with("project-1") + client.resolve_media_url.assert_not_awaited() @pytest.mark.asyncio diff --git a/tests/unit/test_upscale_polling.py b/tests/unit/test_upscale_polling.py deleted file mode 100644 index 14826628..00000000 --- a/tests/unit/test_upscale_polling.py +++ /dev/null @@ -1,67 +0,0 @@ -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -import agent.services.upscale_polling as polling -from agent.services.upscale_polling import annotate_upscale_polling, check_upscale_status - - -MEDIA = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" - - -def test_annotate_upscale_polling_uses_workflow_primary_media_id(): - result = { - "status": 200, - "data": { - "workflows": [ - {"name": "wf-1", "metadata": {"primaryMediaId": MEDIA + "_upsampled"}} - ] - }, - } - annotated = annotate_upscale_polling(result) - assert annotated["data"]["flowkitPolling"] == { - "mode": "media_redirect", - "workflows": [ - {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} - ], - } - - -@pytest.mark.asyncio -async def test_batch_poll_resolves_completed_1080_media_through_as29s(monkeypatch): - monkeypatch.setattr(polling, "USE_BATCH_RPC", True) - client = MagicMock() - client.get_media = AsyncMock(return_value={ - "status": 200, - "data": { - "video": { - "fifeUrl": "https://flow-content.google/video/out?Signature=test" - } - }, - }) - with patch("agent.services.upscale_polling.get_flow_client", return_value=client): - result = await check_upscale_status([ - {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} - ]) - - assert result["done"] is True - assert result["status"] == "COMPLETED" - media = result["workflows"][0]["media"] - assert media["resolved_via"] == "as29s" - assert media["url"].startswith("https://flow-content.google/video/") - client.get_media.assert_awaited_once_with(MEDIA + "_upsampled") - - -@pytest.mark.asyncio -async def test_batch_poll_stays_pending_until_video_url_exists(monkeypatch): - monkeypatch.setattr(polling, "USE_BATCH_RPC", True) - client = MagicMock() - client.get_media = AsyncMock(return_value={"status": 200, "data": {"video": {}}}) - with patch("agent.services.upscale_polling.get_flow_client", return_value=client): - result = await check_upscale_status([ - {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} - ]) - - assert result["done"] is False - assert result["status"] == "PENDING" - assert result["workflows"][0]["status"] == "PENDING"