From dd28d1f811b6394ba0f6dca4ff5bbc1ce4994f2b Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 8 Sep 2026 09:14:03 +0300 Subject: [PATCH 1/3] feat(flow): expand migrated image API --- README.md | 8 ++ agent/api/flow.py | 68 +++++++++++++++-- agent/config.py | 5 +- agent/models.json | 3 +- agent/services/flow_batch.py | 107 +++++++++++++++++++++------ agent/services/flow_client.py | 69 +++++++++++++---- docs/IMAGE_API.md | 93 +++++++++++++++++++++++ tests/unit/test_flow_batch.py | 65 +++++++++++++++- tests/unit/test_flow_client_batch.py | 46 +++++++++++- 9 files changed, 410 insertions(+), 54 deletions(-) create mode 100644 docs/IMAGE_API.md diff --git a/README.md b/README.md index e4dd67ea..a6f98f4f 100644 --- a/README.md +++ b/README.md @@ -244,6 +244,14 @@ You can also pass `flow_project_id` per project on `POST /api/projects`. | `FLOW_ALLOW_DEGRADED` | `0` | `1` lets scene chaining and r2v fall back to plain i2v instead of failing. | | `DEFAULT_PAYGATE_TIER` | `PAYGATE_TIER_TWO` | Carried for the DB and dashboard; no longer selects a model. | +### Image API + +The migrated image path supports Nano Banana Pro, Nano Banana 2 and Nano Banana +2 Lite, all five current aspect ratios, 1-4 outputs, true base-image editing and +native 2K image export. Exact future Flow image model wire ids pass through +without being silently replaced by the default model. See +[`docs/IMAGE_API.md`](docs/IMAGE_API.md). + ### What does not work on the new API yet Three capabilities have no captured payload, so they fail with diff --git a/agent/api/flow.py b/agent/api/flow.py index 99394096..4062eedd 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -1,6 +1,6 @@ """Direct Flow API endpoints — for manual operations outside the queue.""" -from fastapi import APIRouter, HTTPException -from pydantic import BaseModel +from fastapi import APIRouter, HTTPException, Response +from pydantic import BaseModel, Field from typing import Literal, Optional from agent.config import USE_BATCH_RPC, FLOW_PROJECT_ID, FLOW_ALLOW_DEGRADED @@ -20,6 +20,10 @@ class GenerateImageRequest(BaseModel): project_id: str aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT" user_paygate_tier: str = "PAYGATE_TIER_ONE" + image_model: Optional[str] = None + count: int = Field(default=1, ge=1, le=4) + seed: Optional[int] = Field(default=None, ge=1, le=1_000_000_000) + reference_media_ids: Optional[list[str]] = None character_media_ids: Optional[list[str]] = None @@ -94,6 +98,16 @@ class EditImageRequest(BaseModel): project_id: str aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT" user_paygate_tier: str = "PAYGATE_TIER_ONE" + image_model: Optional[str] = None + count: int = Field(default=1, ge=1, le=4) + seed: Optional[int] = Field(default=None, ge=1, le=1_000_000_000) + reference_media_ids: Optional[list[str]] = None + + +class UpscaleImageRequest(BaseModel): + media_id: str + project_id: str + quality: Literal["2k", "4k"] = "2k" @router.get("/status") @@ -127,11 +141,14 @@ async def get_credits(): @router.post("/generate-image") async def generate_image(body: GenerateImageRequest): - """Generate image directly (bypasses queue).""" + """Generate 1-4 images with an explicit Flow image model.""" client = get_flow_client() if not client.connected: raise HTTPException(503, "Extension not connected") - result = await client.generate_images(**body.model_dump()) + data = body.model_dump(exclude={"reference_media_ids"}) + refs = list(dict.fromkeys((body.reference_media_ids or []) + (body.character_media_ids or []))) + data["character_media_ids"] = refs or None + result = await client.generate_images(**data) 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) @@ -335,20 +352,59 @@ async def get_media(media_id: str): @router.post("/edit-image") async def edit_image(body: EditImageRequest): - """Edit an existing image using IMAGE_INPUT_TYPE_BASE_IMAGE (bypasses queue).""" + """Edit an existing image using the current Flow BASE_IMAGE wire input.""" client = get_flow_client() if not client.connected: raise HTTPException(503, "Extension not connected") result = await client.edit_image( - body.prompt, body.source_media_id, body.project_id, + body.prompt, + body.source_media_id, + body.project_id, aspect_ratio=body.aspect_ratio, user_paygate_tier=body.user_paygate_tier, + character_media_ids=body.reference_media_ids, + image_model=body.image_model, + count=body.count, + seed=body.seed, ) 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("/export-image") +@router.post("/upscale-image", include_in_schema=False) +async def export_image(body: UpscaleImageRequest): + """Download a generated Flow image at 2K (or plan-gated 4K).""" + import base64 + import binascii + + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + result = await client.upscale_image( + body.media_id, + body.project_id, + resolution=body.quality.upper(), + ) + 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"))) + data = result.get("data", result) + try: + content = base64.b64decode(data["encodedImage"], validate=True) + except (KeyError, TypeError, binascii.Error) as exc: + raise HTTPException(502, "Flow image upscale returned invalid image data") from exc + quality = body.quality.lower() + return Response( + content=content, + media_type=data.get("contentType", "image/jpeg"), + headers={ + "Content-Disposition": f'attachment; filename="flow-{body.media_id}-{quality}.jpg"', + "X-Flow-Image-Quality": quality, + }, + ) + + @router.post("/upload-image") async def upload_image(body: UploadImageRequest): """Upload a local image file to Google Flow and get a media_id.""" diff --git a/agent/config.py b/agent/config.py index 5d6edf2f..a02a51bd 100644 --- a/agent/config.py +++ b/agent/config.py @@ -61,8 +61,9 @@ VIDEO_MODELS = _MODELS["video_models"] UPSCALE_MODELS = _MODELS["upscale_models"] IMAGE_MODELS = _MODELS["image_models"] -# Nickname from image_models. The batch path accepts GEM_PIX_2 (Nano Banana Pro) -# and NARWHAL (Banana 2) and rejects everything else. +# Nickname from image_models. Known aliases live in models.json, while the +# batch path also accepts syntactically valid Flow wire model ids directly so +# newly introduced image models do not require a Flow Kit release. DEFAULT_IMAGE_MODEL = _MODELS.get("default_image_model", "NANO_BANANA_PRO") # ─── API Endpoints ─────────────────────────────────────────── diff --git a/agent/models.json b/agent/models.json index 8c690e59..7830223e 100644 --- a/agent/models.json +++ b/agent/models.json @@ -55,7 +55,8 @@ }, "image_models": { "NANO_BANANA_PRO": "GEM_PIX_2", - "NANO_BANANA_2": "NARWHAL" + "NANO_BANANA_2": "NARWHAL", + "NANO_BANANA_2_LITE": "HARBOR_SEAL" }, "default_image_model": "NANO_BANANA_PRO", "batch_video_models": { diff --git a/agent/services/flow_batch.py b/agent/services/flow_batch.py index f0111632..2549a511 100644 --- a/agent/services/flow_batch.py +++ b/agent/services/flow_batch.py @@ -40,6 +40,7 @@ RPC_PROJECT_MEDIA = "Zzl0ze" RPC_MEDIA = "as29s" RPC_UPLOAD_IMAGE = "maseQ" +RPC_UPSCALE_IMAGE = "SPrCad" CAPTCHA_IMAGE = "IMAGE_GENERATION" CAPTCHA_VIDEO = "VIDEO_GENERATION" @@ -49,15 +50,21 @@ #: happen in the page, moments before the request leaves. CAPTCHA_SLOT = "__CAPTCHA__" -#: Wire names this path accepts. Everything else is rejected outright by Flow. -#: ``GEM_PIX_2`` is Nano Banana Pro, ``NARWHAL`` is Banana 2. Flow Kit uses Pro -#: by default (see agent/models.json), which is also what the new path defaults -#: to; a caller that wants Banana 2 has to name it. -IMAGE_MODELS = {"GEM_PIX_2", "NARWHAL"} +#: Current image model wire ids observed in the Flow frontend. Unknown future +#: ids are accepted by ``resolve_image_model`` instead of being silently +#: replaced by the default, so callers can opt into a newly exposed model +#: before Flow Kit itself ships another release. +IMAGE_MODELS = {"GEM_PIX_2", "NARWHAL", "HARBOR_SEAL"} IMAGE_MODEL = "GEM_PIX_2" -#: The nicknames models.json speaks, resolved to wire names. -IMAGE_MODEL_BY_NICKNAME = {"NANO_BANANA_PRO": "GEM_PIX_2", "NANO_BANANA_2": "NARWHAL"} +#: Friendly aliases. Exact Flow wire ids work too. +IMAGE_MODEL_BY_NICKNAME = { + "NANO_BANANA_PRO": "GEM_PIX_2", + "NANO_BANANA_2": "NARWHAL", + "NANO_BANANA_2_LITE": "HARBOR_SEAL", + "NANO_BANANA_LITE": "HARBOR_SEAL", +} +IMAGE_MODEL_ID_RE = re.compile(r"^[A-Z][A-Z0-9_]{1,95}$") #: Image aspect ratios, measured by generating one of each and reading the #: JPEG header. This slot was mistaken for a variant count at first — 1 means @@ -73,8 +80,16 @@ "IMAGE_ASPECT_RATIO_SQUARE": ASPECT_SQUARE, "IMAGE_ASPECT_RATIO_PORTRAIT": ASPECT_PORTRAIT, "IMAGE_ASPECT_RATIO_LANDSCAPE": ASPECT_LANDSCAPE, + # Current spelling plus the old Flow Kit alias for compatibility. + "IMAGE_ASPECT_RATIO_PORTRAIT_THREE_FOUR": ASPECT_PORTRAIT_4_3, "IMAGE_ASPECT_RATIO_PORTRAIT_FOUR_THREE": ASPECT_PORTRAIT_4_3, "IMAGE_ASPECT_RATIO_LANDSCAPE_FOUR_THREE": ASPECT_LANDSCAPE_4_3, + # Friendly API aliases. + "1:1": ASPECT_SQUARE, + "9:16": ASPECT_PORTRAIT, + "16:9": ASPECT_LANDSCAPE, + "3:4": ASPECT_PORTRAIT_4_3, + "4:3": ASPECT_LANDSCAPE_4_3, } #: Video models this path accepts. The REST-era map was keyed by @@ -117,10 +132,11 @@ #: reframed by hand: a hair inside the edges, spanning 128/129 of the frame. FULL_FRAME_CROP = [None, 0.0038759689922481244, 1, 0.9961240310077519] -#: A reference image, as the UI sends it: the media id FIRST and a type flag -#: four slots later. Probing never found this — the id sat in the wrong -#: position, the payload was accepted, and the picture quietly ignored it. +#: Image inputs put the media id first and the input type four slots later. +#: Type 1 is a reference; type 2 is the image being edited (BASE_IMAGE). REF_TYPE_IMAGE = 1 +BASE_TYPE_IMAGE = 2 +IMAGE_UPSCALE_RESOLUTIONS = {"2K": 1, "4K": 2} class RpcError(RuntimeError): @@ -180,12 +196,18 @@ class MediaUrls: # ── model / aspect resolvers ───────────────────────────────────────────────── def resolve_image_model(key: Optional[str]) -> str: - """Nickname or wire name in, wire name out; anything unknown coerces.""" + """Nickname or wire id in, wire id out. + + Flow can add image models independently of Flow Kit releases. A + syntactically valid, previously unseen wire id therefore passes through + unchanged instead of being silently replaced by the default model. + """ if isinstance(key, str): - if key in IMAGE_MODEL_BY_NICKNAME: - return IMAGE_MODEL_BY_NICKNAME[key] - if key in IMAGE_MODELS: - return key + normalized = key.strip().upper().replace("-", "_") + if normalized in IMAGE_MODEL_BY_NICKNAME: + return IMAGE_MODEL_BY_NICKNAME[normalized] + if IMAGE_MODEL_ID_RE.fullmatch(normalized): + return normalized return IMAGE_MODEL @@ -211,8 +233,10 @@ def resolve_video_model(key: Optional[str]) -> str: def resolve_aspect(aspect: Any) -> int: - """Take either the wire value or the REST-era name.""" + """Take a current/legacy wire name, friendly ratio, or integer 1-5.""" if isinstance(aspect, int): + if aspect not in (1, 2, 3, 4, 5): + raise ValueError(f"image aspect must be 1-5, got {aspect}") return aspect try: return ASPECT_BY_NAME[aspect] @@ -313,15 +337,24 @@ def _context(project_id: str) -> list: [CAPTCHA_SLOT, 1]] +def _image_input(media_id: str, input_type: int) -> list: + return [media_id, None, None, None, input_type] + + def _reference(media_id: str) -> list: - return [media_id, None, None, None, REF_TYPE_IMAGE] + return _image_input(media_id, REF_TYPE_IMAGE) + + +def _base_image(media_id: str) -> list: + return _image_input(media_id, BASE_TYPE_IMAGE) def image_request(prompt: str, project_id: str, count: int = 1, aspect: Any = ASPECT_SQUARE, seed: Optional[int] = None, prompts: Optional[list[str]] = None, model: str = IMAGE_MODEL, - ref_media_ids: Optional[list[str]] = None) -> str: + ref_media_ids: Optional[list[str]] = None, + base_media_id: Optional[str] = None) -> str: """One request item per variant, exactly as the REST payload did it. There is no "how many" field: Flow returns one image per item in the list, @@ -329,19 +362,39 @@ def image_request(prompt: str, project_id: str, count: int = 1, the result on images already in the project — this is what keeps a character the same person from beat to beat. """ + if not isinstance(count, int) or isinstance(count, bool) or not 1 <= count <= 4: + raise ValueError("image count must be an integer from 1 to 4") ratio = resolve_aspect(aspect) + resolved_model = resolve_image_model(model) base = seed if seed is not None else random.randint(1, 10**9) items = [] - for index in range(max(1, count)): + for index in range(count): text = prompts[index] if prompts and index < len(prompts) else prompt - refs = [_reference(mid) for mid in (ref_media_ids or [])] or None - items.append([None, None, refs, base + index * 9973, ratio, model, None, - _context(project_id), [[[text]]], None, None, None, - _client_uuid(), _client_uuid()]) + image_inputs = [] + if base_media_id: + image_inputs.append(_base_image(base_media_id)) + image_inputs.extend( + _reference(mid) for mid in (ref_media_ids or []) if mid != base_media_id + ) + items.append([None, None, image_inputs or None, base + index * 9973, ratio, + resolved_model, None, _context(project_id), [[[text]]], + None, None, None, _client_uuid(), _client_uuid()]) return build_envelope(RPC_GEN_IMAGE, [None, items, 1, _context(project_id), [_client_uuid()]]) +def image_upscale_request(media_id: str, resolution: str = "2K") -> str: + """Build the current FlowService.UpsampleImage request (RPC SPrCad).""" + key = str(resolution).strip().upper() + if key.startswith("UPSAMPLE_IMAGE_RESOLUTION_"): + key = key.removeprefix("UPSAMPLE_IMAGE_RESOLUTION_") + try: + code = IMAGE_UPSCALE_RESOLUTIONS[key] + except KeyError: + raise ValueError("image upscale resolution must be 2K or 4K") from None + return build_envelope(RPC_UPSCALE_IMAGE, [media_id, code, _context(None)]) + + def video_request(prompt: str, project_id: str, source_media_id: str, crop: Optional[list] = None, aspect: Any = VIDEO_ASPECT_LANDSCAPE, @@ -429,6 +482,14 @@ def read_uploaded_media_id(payload: Any) -> str: return media_id +def read_upscaled_image(payload: Any) -> str: + """Return the base64 image body from FlowService.UpsampleImage.""" + encoded = payload[1] if isinstance(payload, list) and len(payload) > 1 else None + if not isinstance(encoded, str) or len(encoded) < 100: + raise FlowBatchError("image upscale response carried no encoded image") + return encoded + + def read_operation(payload: Any) -> Operation: """`[null, 50, [[opId, projectId, sceneId, status, …]]]`. diff --git a/agent/services/flow_client.py b/agent/services/flow_client.py index 667e2d51..fb3b78f2 100644 --- a/agent/services/flow_client.py +++ b/agent/services/flow_client.py @@ -572,7 +572,10 @@ async def generate_images(self, prompt: str, project_id: str, aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT", user_paygate_tier: str = "PAYGATE_TIER_TWO", character_media_ids: list[str] = None, - image_model: str = None) -> dict: + image_model: str = None, + count: int = 1, + seed: int | None = None, + base_media_id: str | None = None) -> dict: """Generate image(s). ``character_media_ids`` are attached as reference images, which is what @@ -587,9 +590,10 @@ async def generate_images(self, prompt: str, project_id: str, try: pid = self._batch_project_id(project_id) freq = fb.image_request( - prompt, pid, count=1, aspect=aspect_ratio, + prompt, pid, count=count, aspect=aspect_ratio, seed=seed, model=self._batch_image_model(image_model), ref_media_ids=list(character_media_ids or []) or None, + base_media_id=base_media_id, ) payload = await self._batch_payload(fb.RPC_GEN_IMAGE, freq, fb.CAPTCHA_IMAGE) except Exception as e: @@ -604,29 +608,62 @@ async def edit_image(self, prompt: str, source_media_id: str, project_id: str, aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT", user_paygate_tier: str = "PAYGATE_TIER_ONE", - character_media_ids: list[str] = None) -> dict: - """Regenerate from an existing image plus any entity references. - - The REST path had a dedicated base-image input type; the new payload's - reference slot was captured but a base-image variant of it was not, so - here the source rides in as the first reference. In practice that - conditions the result on the source rather than editing it in place — - good enough for continuation scenes, not identical to the old edit. - Capturing the real slot is the fix; see docs/CAPTURE.md. + character_media_ids: list[str] = None, + image_model: str = None, + count: int = 1, + seed: int | None = None) -> dict: + """Edit an image with the source encoded as Flow's BASE_IMAGE input. + + Additional references remain REFERENCE inputs. Sending the source as a + generic reference conditions a fresh generation; BASE_IMAGE is the wire + shape the current Flow editor uses for an actual image edit/refine. """ if not USE_BATCH_RPC: return await self._legacy_edit_image( prompt, source_media_id, project_id, aspect_ratio, user_paygate_tier, character_media_ids) - refs = [source_media_id] + [ - mid for mid in (character_media_ids or []) if mid != source_media_id - ] + refs = [mid for mid in (character_media_ids or []) if mid != source_media_id] return await self.generate_images( - prompt=prompt, project_id=project_id, aspect_ratio=aspect_ratio, - user_paygate_tier=user_paygate_tier, character_media_ids=refs, + prompt=prompt, + project_id=project_id, + aspect_ratio=aspect_ratio, + user_paygate_tier=user_paygate_tier, + character_media_ids=refs, + image_model=image_model, + count=count, + seed=seed, + base_media_id=source_media_id, ) + async def upscale_image(self, media_id: str, project_id: str, + resolution: str = "2K") -> dict: + """Return Flow's synchronous 2K/4K image upscale as base64 JPEG data.""" + if not USE_BATCH_RPC: + return {"status": 400, "error": "Image upscale requires the flow.google.com batch transport"} + try: + pid = self._batch_project_id(project_id) + freq = fb.image_upscale_request(media_id, resolution) + payload = await self._batch_payload( + fb.RPC_UPSCALE_IMAGE, + freq, + fb.CAPTCHA_IMAGE, + timeout=150, + ) + encoded = fb.read_upscaled_image(payload) + except Exception as e: + return _batch_error(e) + return { + "status": 200, + "data": { + "media_id": media_id, + "project_id": pid, + "resolution": str(resolution).upper(), + "encodedImage": encoded, + "contentType": "image/jpeg", + }, + } + async def generate_video(self, start_image_media_id: str, prompt: str, project_id: str, scene_id: str, aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT", diff --git a/docs/IMAGE_API.md b/docs/IMAGE_API.md new file mode 100644 index 00000000..e0447176 --- /dev/null +++ b/docs/IMAGE_API.md @@ -0,0 +1,93 @@ +# Migrated Image API + +Flow's September 2026 frontend exposes image generation through the +`flow.google.com` batchexecute transport. This document covers the image +capabilities currently wired by Flow Kit. + +## Generate images + +`POST /api/flow/generate-image` + +```json +{ + "prompt": "A red paper boat on a calm pond", + "project_id": "", + "image_model": "HARBOR_SEAL", + "aspect_ratio": "16:9", + "count": 2, + "seed": 12345, + "reference_media_ids": [] +} +``` + +Current Flow UI model ids observed on the migrated frontend: + +- `GEM_PIX_2` — Nano Banana Pro +- `NARWHAL` — Nano Banana 2 +- `HARBOR_SEAL` — Nano Banana 2 Lite + +Friendly aliases from `models.json` continue to work. Flow Kit also passes a +syntactically valid future wire model id through unchanged instead of silently +coercing it to the default, so a newly exposed model can be selected before the +next Flow Kit release once its id is known. + +`count` accepts 1-4, matching the Flow UI. Each output is a separate request +item with its own seed. + +## Aspect ratios + +All five current image ratios are supported, using either the friendly ratio or +wire enum: + +| Ratio | Wire name | +|---|---| +| `1:1` | `IMAGE_ASPECT_RATIO_SQUARE` | +| `9:16` | `IMAGE_ASPECT_RATIO_PORTRAIT` | +| `16:9` | `IMAGE_ASPECT_RATIO_LANDSCAPE` | +| `3:4` | `IMAGE_ASPECT_RATIO_PORTRAIT_THREE_FOUR` | +| `4:3` | `IMAGE_ASPECT_RATIO_LANDSCAPE_FOUR_THREE` | + +The older Flow Kit spelling `IMAGE_ASPECT_RATIO_PORTRAIT_FOUR_THREE` remains an +alias for compatibility. + +## Edit an image + +`POST /api/flow/edit-image` + +The source image is sent as Flow's `BASE_IMAGE` input (wire type 2). Additional +`reference_media_ids` remain reference inputs (wire type 1). This fixes the old +batch-path behavior where the source itself was only a generic reference and +therefore conditioned a fresh generation instead of performing a true edit. + +```json +{ + "prompt": "Make the paper boat blue", + "source_media_id": "", + "project_id": "", + "image_model": "GEM_PIX_2", + "aspect_ratio": "16:9", + "count": 1 +} +``` + +## Export / upscale image + +`POST /api/flow/export-image` + +```json +{ + "media_id": "", + "project_id": "", + "quality": "2k" +} +``` + +The migrated frontend uses RPC `SPrCad` (`FlowService.UpsampleImage`). The call +is synchronous and returns the encoded JPEG, which the HTTP endpoint returns as +a downloadable image. + +- `2k` — standard high-resolution download; live verified +- `4k` — same RPC with target code 2; availability is account/plan-gated + +Live verification on the current Flow frontend produced a 2752x1536 JPEG from a +1376x768 source. diff --git a/tests/unit/test_flow_batch.py b/tests/unit/test_flow_batch.py index 63b8761d..29c75389 100644 --- a/tests/unit/test_flow_batch.py +++ b/tests/unit/test_flow_batch.py @@ -85,6 +85,21 @@ def test_reference_puts_the_media_id_first_and_the_type_flag_fourth(self): item = inner(fb.image_request("a cat", self.PID, ref_media_ids=["mid-1"]))[1][0] assert item[2] == [["mid-1", None, None, None, fb.REF_TYPE_IMAGE]] + def test_base_image_and_references_use_distinct_wire_types(self): + item = inner(fb.image_request( + "make the boat blue", self.PID, + base_media_id="base-1", ref_media_ids=["ref-1", "base-1"], + ))[1][0] + assert item[2] == [ + ["base-1", None, None, None, fb.BASE_TYPE_IMAGE], + ["ref-1", None, None, None, fb.REF_TYPE_IMAGE], + ] + + @pytest.mark.parametrize("bad_count", [0, 5, -1, True]) + def test_count_outside_flow_ui_range_is_rejected(self, bad_count): + with pytest.raises(ValueError): + fb.image_request("a cat", self.PID, count=bad_count) + def test_no_references_leaves_the_slot_null_rather_than_empty(self): assert inner(fb.image_request("a cat", self.PID))[1][0][2] is None @@ -99,6 +114,25 @@ def test_the_model_is_named_in_slot_5(self): assert item[5] == "NARWHAL" +class TestImageUpscaleRequest: + @pytest.mark.parametrize("resolution,wire", [ + ("2k", 1), + ("4K", 2), + ("UPSAMPLE_IMAGE_RESOLUTION_2K", 1), + ]) + def test_resolution_maps_to_live_sprcad_wire(self, resolution, wire): + payload = inner(fb.image_upscale_request("media-1", resolution)) + assert payload[0] == "media-1" + assert payload[1] == wire + assert payload[2][1] == fb.SURFACE_ID + assert payload[2][5] is None + assert fb.CAPTCHA_SLOT in json.dumps(payload) + + def test_unknown_resolution_is_rejected(self): + with pytest.raises(ValueError): + fb.image_upscale_request("media-1", "8K") + + class TestVideoRequest: PID = "11111111-2222-3333-4444-555555555555" @@ -135,6 +169,11 @@ def test_a_repeated_url_is_not_a_second_variant(self): url = f"https://{fb.MEDIA_HOST}/image/{self.MID}?sig=x" assert len(fb.read_images([url, url])) == 1 + 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_operation_reads_the_id_and_status(self): op = fb.read_operation([None, 50, [[self.OP, "proj", "scene", "CAE"]]]) assert (op.operation_id, op.status, op.done) == (self.OP, "CAE", True) @@ -204,12 +243,32 @@ def test_an_unknown_aspect_name_raises_rather_than_defaulting(self): def test_nicknames_resolve_to_wire_names(self): assert fb.resolve_image_model("NANO_BANANA_PRO") == "GEM_PIX_2" assert fb.resolve_image_model("NANO_BANANA_2") == "NARWHAL" + assert fb.resolve_image_model("NANO_BANANA_2_LITE") == "HARBOR_SEAL" - def test_a_wire_name_passes_through(self): + def test_wire_names_and_future_model_ids_pass_through(self): assert fb.resolve_image_model("NARWHAL") == "NARWHAL" + assert fb.resolve_image_model("harbor_seal") == "HARBOR_SEAL" + assert fb.resolve_image_model("FUTURE_BANANA_3") == "FUTURE_BANANA_3" + + def test_invalid_or_blank_model_falls_back_to_default(self): + assert fb.resolve_image_model("") == fb.IMAGE_MODEL + assert fb.resolve_image_model("not a wire id!") == fb.IMAGE_MODEL + + @pytest.mark.parametrize("name,expected", [ + ("1:1", fb.ASPECT_SQUARE), + ("9:16", fb.ASPECT_PORTRAIT), + ("16:9", fb.ASPECT_LANDSCAPE), + ("3:4", fb.ASPECT_PORTRAIT_4_3), + ("4:3", fb.ASPECT_LANDSCAPE_4_3), + ("IMAGE_ASPECT_RATIO_PORTRAIT_THREE_FOUR", fb.ASPECT_PORTRAIT_4_3), + ("IMAGE_ASPECT_RATIO_PORTRAIT_FOUR_THREE", fb.ASPECT_PORTRAIT_4_3), + ]) + def test_all_current_image_aspect_names_and_friendly_aliases(self, name, expected): + assert fb.resolve_aspect(name) == expected - def test_an_unknown_image_model_coerces_to_the_default(self): - assert fb.resolve_image_model("SOMETHING_ELSE") == fb.IMAGE_MODEL + def test_invalid_integer_image_aspect_is_rejected(self): + with pytest.raises(ValueError): + fb.resolve_aspect(6) @pytest.mark.parametrize("legacy,expected", [ ("veo_3_1_i2v_s_fast_ultra_relaxed", "veo_3_1_i2v_s_fast_ultra"), diff --git a/tests/unit/test_flow_client_batch.py b/tests/unit/test_flow_client_batch.py index 733f5244..a9ae44cd 100644 --- a/tests/unit/test_flow_client_batch.py +++ b/tests/unit/test_flow_client_batch.py @@ -72,6 +72,26 @@ async def test_character_refs_ride_in_the_reference_slot(self, client): assert item[2] == [["ref-a", None, None, None, fb.REF_TYPE_IMAGE], ["ref-b", None, None, None, fb.REF_TYPE_IMAGE]] + async def test_explicit_model_and_count_reach_the_batch_wire(self, client): + second = "12345678-1234-1234-1234-1234567890ac" + client.responses[fb.RPC_GEN_IMAGE] = { + "data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL], [IMAGE_URL.replace(MEDIA, second)]]) + } + result = await client.generate_images( + "a cat", PROJECT, image_model="HARBOR_SEAL", count=2, seed=100, + ) + items = json.loads(json.loads(client.calls[0]["freq"])[0][0][1])[1] + assert len(items) == 2 + assert [item[5] for item in items] == ["HARBOR_SEAL", "HARBOR_SEAL"] + assert [item[3] for item in items] == [100, 100 + 9973] + assert len(result["data"]["media"]) == 2 + + async def test_future_wire_model_is_not_silently_replaced(self, client): + client.responses[fb.RPC_GEN_IMAGE] = {"data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL]])} + await client.generate_images("a cat", PROJECT, image_model="FUTURE_BANANA_3") + item = json.loads(json.loads(client.calls[0]["freq"])[0][0][1])[1][0] + assert item[5] == "FUTURE_BANANA_3" + async def test_a_project_less_call_falls_back_to_the_pinned_project(self, client): client.responses[fb.RPC_GEN_IMAGE] = {"data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL]])} await client.generate_images("a cat", "0") @@ -96,19 +116,39 @@ async def test_no_project_anywhere_is_a_named_failure(self, client, monkeypatch) class TestEditImage: - async def test_the_source_leads_the_reference_list(self, client): + async def test_source_is_base_image_and_extra_inputs_are_references(self, client): client.responses[fb.RPC_GEN_IMAGE] = {"data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL]])} await client.edit_image("redraw", "src-1", PROJECT, character_media_ids=["ref-a"]) item = json.loads(json.loads(client.calls[0]["freq"])[0][0][1])[1][0] - assert [ref[0] for ref in item[2]] == ["src-1", "ref-a"] + assert item[2] == [ + ["src-1", None, None, None, fb.BASE_TYPE_IMAGE], + ["ref-a", None, None, None, fb.REF_TYPE_IMAGE], + ] - async def test_the_source_is_not_repeated_when_it_is_also_a_character(self, client): + async def test_source_is_not_repeated_when_also_supplied_as_reference(self, client): client.responses[fb.RPC_GEN_IMAGE] = {"data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL]])} await client.edit_image("redraw", "src-1", PROJECT, character_media_ids=["src-1", "ref-a"]) item = json.loads(json.loads(client.calls[0]["freq"])[0][0][1])[1][0] assert [ref[0] for ref in item[2]] == ["src-1", "ref-a"] + assert [ref[4] for ref in item[2]] == [fb.BASE_TYPE_IMAGE, fb.REF_TYPE_IMAGE] + + +class TestUpscaleImage: + async def test_2k_upscale_uses_sprcad_and_returns_encoded_image(self, client): + encoded = "A" * 200 + client.responses[fb.RPC_UPSCALE_IMAGE] = { + "data": envelope(fb.RPC_UPSCALE_IMAGE, [["media-record"], encoded]) + } + result = await client.upscale_image(MEDIA, PROJECT, "2K") + assert result["data"]["encodedImage"] == encoded + assert result["data"]["resolution"] == "2K" + assert client.calls[0]["rpcid"] == fb.RPC_UPSCALE_IMAGE + assert client.calls[0]["captcha"] == fb.CAPTCHA_IMAGE + payload = json.loads(json.loads(client.calls[0]["freq"])[0][0][1]) + assert payload[0] == MEDIA + assert payload[1] == 1 class TestGenerateVideo: From 59cc4970c1b320da9a381a4f39e80f159f7d11b7 Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:29:17 +0300 Subject: [PATCH 2/3] fix(flow): match image variant submit behavior --- agent/services/flow_client.py | 40 +++++++++++++++++++++------- docs/IMAGE_API.md | 8 ++++-- extension/background.js | 9 +++++-- extension/injected.js | 20 +++++++++++--- tests/unit/test_flow_client_batch.py | 39 ++++++++++++++++++++++----- 5 files changed, 92 insertions(+), 24 deletions(-) diff --git a/agent/services/flow_client.py b/agent/services/flow_client.py index fb3b78f2..4352f1f0 100644 --- a/agent/services/flow_client.py +++ b/agent/services/flow_client.py @@ -33,6 +33,11 @@ logger = logging.getLogger(__name__) +# Captured from the current Flow image composer. x4 launches independent +# ogiZ0b requests at roughly 0.0s, 0.5s, 1.5s and 2.5s rather than bursting +# all variants at once. The generation work still overlaps after submission. +IMAGE_UI_SUBMIT_OFFSETS_S = (0.0, 0.5, 1.5, 2.5) + class FlowClient: """Sends commands to Chrome extension via WebSocket.""" @@ -588,20 +593,35 @@ async def generate_images(self, prompt: str, project_id: str, prompt, project_id, aspect_ratio, user_paygate_tier, character_media_ids) try: + if not isinstance(count, int) or isinstance(count, bool) or not 1 <= count <= 4: + raise ValueError("image count must be an integer from 1 to 4") pid = self._batch_project_id(project_id) - freq = fb.image_request( - prompt, pid, count=count, aspect=aspect_ratio, seed=seed, - model=self._batch_image_model(image_model), - ref_media_ids=list(character_media_ids or []) or None, - base_media_id=base_media_id, - ) - payload = await self._batch_payload(fb.RPC_GEN_IMAGE, freq, fb.CAPTCHA_IMAGE) + model = self._batch_image_model(image_model) + refs = list(character_media_ids or []) or None + + async def submit_one(index: int): + # Flow UI implements x2/x3/x4 as independent single-image RPCs. + # It also staggers their launch to avoid an artificial burst. + offset = IMAGE_UI_SUBMIT_OFFSETS_S[index] + if offset: + await asyncio.sleep(offset) + request_seed = seed + index * 9973 if seed is not None else None + freq = fb.image_request( + prompt, pid, count=1, aspect=aspect_ratio, seed=request_seed, + model=model, ref_media_ids=refs, base_media_id=base_media_id, + ) + payload = await self._batch_payload( + fb.RPC_GEN_IMAGE, freq, fb.CAPTCHA_IMAGE + ) + generated = fb.read_images(payload) + if not generated: + raise fb.FlowBatchError("Image generation returned no media url") + return generated[0] + + images = await asyncio.gather(*(submit_one(i) for i in range(count))) except Exception as e: return _batch_error(e) - images = fb.read_images(payload) - if not images: - return {"status": 502, "error": "Image generation returned no media url"} return {"status": 200, "data": {"media": [_as_media_record(i) for i in images]}} async def edit_image(self, prompt: str, source_media_id: str, diff --git a/docs/IMAGE_API.md b/docs/IMAGE_API.md index e0447176..e6c931a2 100644 --- a/docs/IMAGE_API.md +++ b/docs/IMAGE_API.md @@ -31,8 +31,12 @@ syntactically valid future wire model id through unchanged instead of silently coercing it to the default, so a newly exposed model can be selected before the next Flow Kit release once its id is known. -`count` accepts 1-4, matching the Flow UI. Each output is a separate request -item with its own seed. +`count` accepts 1-4, matching the Flow UI. Flow itself implements x2/x3/x4 +as independent single-image `ogiZ0b` RPCs, each with a single-use reCAPTCHA; +it also staggers x4 launches at roughly 0.0 / 0.5 / 1.5 / 2.5 seconds. FlowKit +mirrors that behavior instead of sending an unofficial multi-item burst inside +one RPC. Each variant keeps its own seed; when `seed` is supplied, later +variants use a deterministic stride. ## Aspect ratios diff --git a/extension/background.js b/extension/background.js index 98228da8..cc75d683 100644 --- a/extension/background.js +++ b/extension/background.js @@ -480,10 +480,15 @@ async function runBatchRpc(cmd) { const bl = wiz.cfb2h; if (!at) return { error: 'NO_AT_TOKEN' }; const reqid = Math.floor(Math.random() * 900000) + 100000; + // Match Flow's own WIZ metadata. GEM_PIX_2 (Nano Banana Pro) rejects + // image generation when source-path is missing even though Lite may not. + const sourcePath = location.pathname || '/'; + const hl = (document.documentElement.lang || navigator.language || 'en').split('-')[0]; const url = `/_/AiSandboxAngularFrontend/data/batchexecute?rpcids=${encodeURIComponent(rpcid)}` + - `&f.sid=${encodeURIComponent(sid || '')}&bl=${encodeURIComponent(bl || '')}` + - `&hl=en-AU&_reqid=${reqid}&rt=c`; + `&source-path=${encodeURIComponent(sourcePath)}` + + `&bl=${encodeURIComponent(bl || '')}&f.sid=${encodeURIComponent(sid || '')}` + + `&hl=${encodeURIComponent(hl)}&_reqid=${reqid}&rt=c`; const resp = await fetch(url, { method: 'POST', credentials: 'include', diff --git a/extension/injected.js b/extension/injected.js index e26b690a..a955a5f8 100644 --- a/extension/injected.js +++ b/extension/injected.js @@ -34,13 +34,27 @@ window.fetch = async function (...args) { }; -window.addEventListener('GET_CAPTCHA', async ({ detail }) => { - const { requestId, pageAction } = detail; +let captchaMintTail = Promise.resolve(); + +async function mintCaptcha(pageAction) { + const previous = captchaMintTail.catch(() => {}); + let release; + captchaMintTail = new Promise((resolve) => { release = resolve; }); + await previous; try { await waitForGrecaptcha(); - const token = await window.grecaptcha.enterprise.execute(SITE_KEY, { + return await window.grecaptcha.enterprise.execute(SITE_KEY, { action: pageAction, }); + } finally { + release(); + } +} + +window.addEventListener('GET_CAPTCHA', async ({ detail }) => { + const { requestId, pageAction } = detail; + try { + const token = await mintCaptcha(pageAction); window.dispatchEvent(new CustomEvent('CAPTCHA_RESULT', { detail: { requestId, token }, })); diff --git a/tests/unit/test_flow_client_batch.py b/tests/unit/test_flow_client_batch.py index a9ae44cd..a3235a7e 100644 --- a/tests/unit/test_flow_client_batch.py +++ b/tests/unit/test_flow_client_batch.py @@ -72,20 +72,45 @@ async def test_character_refs_ride_in_the_reference_slot(self, client): assert item[2] == [["ref-a", None, None, None, fb.REF_TYPE_IMAGE], ["ref-b", None, None, None, fb.REF_TYPE_IMAGE]] - async def test_explicit_model_and_count_reach_the_batch_wire(self, client): - second = "12345678-1234-1234-1234-1234567890ac" + async def test_explicit_model_and_count_dispatch_as_ui_style_rpcs(self, client, monkeypatch): + import agent.services.flow_client as module + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(module.asyncio, "sleep", fake_sleep) client.responses[fb.RPC_GEN_IMAGE] = { - "data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL], [IMAGE_URL.replace(MEDIA, second)]]) + "data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL]]) } result = await client.generate_images( "a cat", PROJECT, image_model="HARBOR_SEAL", count=2, seed=100, ) - items = json.loads(json.loads(client.calls[0]["freq"])[0][0][1])[1] - assert len(items) == 2 - assert [item[5] for item in items] == ["HARBOR_SEAL", "HARBOR_SEAL"] - assert [item[3] for item in items] == [100, 100 + 9973] + assert len(client.calls) == 2 + items = [json.loads(json.loads(call["freq"])[0][0][1])[1] for call in client.calls] + assert [len(group) for group in items] == [1, 1] + assert [group[0][5] for group in items] == ["HARBOR_SEAL", "HARBOR_SEAL"] + assert [group[0][3] for group in items] == [100, 100 + 9973] + assert all(call["captcha"] == fb.CAPTCHA_IMAGE for call in client.calls) + assert sleeps == [module.IMAGE_UI_SUBMIT_OFFSETS_S[1]] assert len(result["data"]["media"]) == 2 + async def test_count_four_uses_captured_ui_launch_offsets(self, client, monkeypatch): + import agent.services.flow_client as module + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(module.asyncio, "sleep", fake_sleep) + client.responses[fb.RPC_GEN_IMAGE] = { + "data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL]]) + } + result = await client.generate_images("a cat", PROJECT, count=4) + assert len(client.calls) == 4 + assert sleeps == list(module.IMAGE_UI_SUBMIT_OFFSETS_S[1:4]) + assert len(result["data"]["media"]) == 4 + async def test_future_wire_model_is_not_silently_replaced(self, client): client.responses[fb.RPC_GEN_IMAGE] = {"data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL]])} await client.generate_images("a cat", PROJECT, image_model="FUTURE_BANANA_3") From 005ad119a13315302d49ad8ec272bba2dbc4c8da Mon Sep 17 00:00:00 2001 From: Bl0ck <36800583+Bl0ck154@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:02:17 +0300 Subject: [PATCH 3/3] fix(flow): retry image failures by settled wave --- agent/services/flow_client.py | 79 +++++++++++++++++++++++++--- docs/IMAGE_API.md | 8 +++ tests/unit/test_flow_client_batch.py | 71 +++++++++++++++++++++++++ 3 files changed, 150 insertions(+), 8 deletions(-) diff --git a/agent/services/flow_client.py b/agent/services/flow_client.py index 4352f1f0..03501e8c 100644 --- a/agent/services/flow_client.py +++ b/agent/services/flow_client.py @@ -38,6 +38,14 @@ # all variants at once. The generation work still overlaps after submission. IMAGE_UI_SUBMIT_OFFSETS_S = (0.0, 0.5, 1.5, 2.5) +# RPC [8] is a transient Flow-side generation rejection seen under image load. +# A short 6s retry was still rejected in live testing, so use one bounded +# cooldown retry rather than hot-looping or multiplying duplicate generations. +# This is FlowKit resilience policy; the current UI was not observed to retry +# automatically after the same failure. +IMAGE_TRANSIENT_RETRY_DELAY_S = 34.0 +IMAGE_TRANSIENT_MAX_ATTEMPTS = 2 + class FlowClient: """Sends commands to Chrome extension via WebSocket.""" @@ -599,12 +607,9 @@ async def generate_images(self, prompt: str, project_id: str, model = self._batch_image_model(image_model) refs = list(character_media_ids or []) or None - async def submit_one(index: int): - # Flow UI implements x2/x3/x4 as independent single-image RPCs. - # It also staggers their launch to avoid an artificial burst. - offset = IMAGE_UI_SUBMIT_OFFSETS_S[index] - if offset: - await asyncio.sleep(offset) + async def submit_once(index: int, launch_offset: float = 0.0): + if launch_offset: + await asyncio.sleep(launch_offset) request_seed = seed + index * 9973 if seed is not None else None freq = fb.image_request( prompt, pid, count=1, aspect=aspect_ratio, seed=request_seed, @@ -618,11 +623,69 @@ async def submit_one(index: int): raise fb.FlowBatchError("Image generation returned no media url") return generated[0] - images = await asyncio.gather(*(submit_one(i) for i in range(count))) + async def run_wave(indices: list[int]) -> dict[int, object]: + # Flow's UI starts variants as separate single-image RPCs with a + # short cadence instead of a burst. Apply the cadence relative + # to each wave, while Google still performs the generation work + # concurrently after each request has been accepted. + tasks = [ + submit_once(index, IMAGE_UI_SUBMIT_OFFSETS_S[position]) + for position, index in enumerate(indices) + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + return dict(zip(indices, results)) + + results = await run_wave(list(range(count))) + retry_indices = [ + index for index, result in results.items() + if isinstance(result, fb.RpcError) + and result.rpcid == fb.RPC_GEN_IMAGE + and result.detail == [8] + ] + if retry_indices: + logger.warning( + "Flow image wave had transient [8] for variant(s) %s; " + "retrying after %.0fs cooldown once the first wave is fully settled", + ",".join(str(i + 1) for i in retry_indices), + IMAGE_TRANSIENT_RETRY_DELAY_S, + ) + await asyncio.sleep(IMAGE_TRANSIENT_RETRY_DELAY_S) + retried = await run_wave(retry_indices) + results.update(retried) + + images_by_index = { + index: result + for index, result in results.items() + if not isinstance(result, BaseException) + } + failures = { + index: result + for index, result in results.items() + if isinstance(result, BaseException) + } + if not images_by_index: + first_error = failures[min(failures)] if failures else fb.FlowBatchError( + "Image generation returned no media url" + ) + raise first_error + + images = [images_by_index[index] for index in sorted(images_by_index)] + except Exception as e: return _batch_error(e) - return {"status": 200, "data": {"media": [_as_media_record(i) for i in images]}} + data = { + "media": [_as_media_record(i) for i in images], + "requested_count": count, + "generated_count": len(images), + "complete": len(images) == count, + } + if failures: + data["failed_variants"] = [ + {"index": index + 1, "error": str(error)} + for index, error in sorted(failures.items()) + ] + return {"status": 200, "data": data} async def edit_image(self, prompt: str, source_media_id: str, project_id: str, diff --git a/docs/IMAGE_API.md b/docs/IMAGE_API.md index e6c931a2..64130dde 100644 --- a/docs/IMAGE_API.md +++ b/docs/IMAGE_API.md @@ -38,6 +38,14 @@ mirrors that behavior instead of sending an unofficial multi-item burst inside one RPC. Each variant keeps its own seed; when `seed` is supplied, later variants use a deterministic stride. +FlowKit treats only image RPC `[8]` as transient. It waits for the entire first +variant wave to settle, cools down for 34 seconds, then retries only failed `[8]` +variants once with fresh request UUIDs and the same seed. This is a FlowKit +resilience policy; live UI capture did not show the same automatic retry. If a +multi-image request still has a failed variant after retry, already successful +images are preserved and the response reports `complete=false`, +`generated_count`, and `failed_variants` instead of discarding the whole batch. + ## Aspect ratios All five current image ratios are supported, using either the friendly ratio or diff --git a/tests/unit/test_flow_client_batch.py b/tests/unit/test_flow_client_batch.py index a3235a7e..c38c09ed 100644 --- a/tests/unit/test_flow_client_batch.py +++ b/tests/unit/test_flow_client_batch.py @@ -111,6 +111,77 @@ async def fake_sleep(delay): assert sleeps == list(module.IMAGE_UI_SUBMIT_OFFSETS_S[1:4]) assert len(result["data"]["media"]) == 4 + async def test_rpc_error_8_retries_once_after_cooldown(self, client, monkeypatch): + import agent.services.flow_client as module + + attempts = 0 + sleeps = [] + + async def fake_payload(rpcid, freq, captcha_action=None, timeout=300): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise fb.RpcError(fb.RPC_GEN_IMAGE, [8]) + return [[IMAGE_URL]] + + async def fake_sleep(delay): + sleeps.append(delay) + + client._batch_payload = fake_payload + monkeypatch.setattr(module.asyncio, "sleep", fake_sleep) + + result = await client.generate_images("a cat", PROJECT, count=1) + assert not _is_error(result) + assert attempts == 2 + assert sleeps == [module.IMAGE_TRANSIENT_RETRY_DELAY_S] + + async def test_non_transient_rpc_error_is_not_retried(self, client, monkeypatch): + import agent.services.flow_client as module + + attempts = 0 + sleeps = [] + + async def fake_payload(rpcid, freq, captcha_action=None, timeout=300): + nonlocal attempts + attempts += 1 + raise fb.RpcError(fb.RPC_GEN_IMAGE, [5]) + + async def fake_sleep(delay): + sleeps.append(delay) + + client._batch_payload = fake_payload + monkeypatch.setattr(module.asyncio, "sleep", fake_sleep) + + result = await client.generate_images("a cat", PROJECT, count=1) + assert _is_error(result) + assert attempts == 1 + assert sleeps == [] + + async def test_partial_batch_keeps_successes_and_reports_failed_variants(self, client, monkeypatch): + import agent.services.flow_client as module + + async def fake_sleep(_delay): + return None + + async def fake_payload(rpcid, freq, captcha_action=None, timeout=300): + item = json.loads(json.loads(freq)[0][0][1])[1][0] + if item[3] == 100 + 9973: + raise fb.RpcError(fb.RPC_GEN_IMAGE, [5]) + return [[IMAGE_URL]] + + client._batch_payload = fake_payload + monkeypatch.setattr(module.asyncio, "sleep", fake_sleep) + + result = await client.generate_images("a cat", PROJECT, count=2, seed=100) + assert not _is_error(result) + data = result["data"] + assert len(data["media"]) == 1 + assert data["requested_count"] == 2 + assert data["generated_count"] == 1 + assert data["complete"] is False + assert data["failed_variants"][0]["index"] == 2 + assert "[5]" in data["failed_variants"][0]["error"] + async def test_future_wire_model_is_not_silently_replaced(self, client): client.responses[fb.RPC_GEN_IMAGE] = {"data": envelope(fb.RPC_GEN_IMAGE, [[IMAGE_URL]])} await client.generate_images("a cat", PROJECT, image_model="FUTURE_BANANA_3")