fix(skills): bound Gemini video operation polling (#7579)

skills/public/video-generation/scripts/generate.py

@@ -82,6 +82,41 @@ def _poll_video_task(host: str, auth: str, task_id: str, max_attempts: int = 120
     raise Exception(f"MiniMax video task {task_id} timed out after {max_attempts} polls")
 
 
+def _poll_gemini_operation(api_key: str, operation_name: str, max_attempts: int = 240, interval: int = 3) -> str:
+    """Wait for a Veo long-running operation and return the generated video URI.
+
+    Bounded like ``_poll_video_task`` so a stuck or rejected operation fails with a
+    diagnosis instead of polling forever. The default budget (12 minutes) leaves
+    room for queueing, which routinely outlasts a short generation.
+    """
+    for _ in range(max_attempts):
+        response = requests.get(
+            f"https://generativelanguage.googleapis.com/v1beta/{operation_name}",
+            headers={"x-goog-api-key": api_key},
+            timeout=30,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        if not payload.get("done", False):
+            time.sleep(interval)
+            continue
+        error = payload.get("error") or {}
+        if error:
+            raise Exception(f"Gemini video operation {operation_name} failed: {error.get('code')} {error.get('message')}")
+        video_response = (payload.get("response") or {}).get("generateVideoResponse") or {}
+        samples = video_response.get("generatedSamples") or []
+        if not samples:
+            # A filtered prompt finishes successfully without any sample.
+            reasons = "; ".join(video_response.get("raiMediaFilteredReasons") or [])
+            detail = f": {reasons}" if reasons else ""
+            raise Exception(f"Gemini video operation {operation_name} returned no generated sample{detail}")
+        uri = (samples[0].get("video") or {}).get("uri")
+        if not uri:
+            raise Exception(f"Gemini video operation {operation_name} returned a sample without a video URI")
+        return uri
+    raise Exception(f"Gemini video operation {operation_name} timed out after {max_attempts} polls")
+
+
 def _retrieve_file_url(host: str, auth: str, file_id: str) -> str:
     response = requests.get(
         f"{host}/v1/files/retrieve",
@@ -163,19 +198,8 @@ def _generate_video_gemini(prompt: str, reference_images: list[str], output_file
     response.raise_for_status()
     data = response.json()
     operation_name = data["name"]
-    while True:
-        response = requests.get(
-            f"https://generativelanguage.googleapis.com/v1beta/{operation_name}",
-            headers={"x-goog-api-key": gemini_api_key},
-            timeout=30,
-        )
-        response.raise_for_status()
-        data = response.json()
-        if data.get("done", False):
-            sample = data["response"]["generateVideoResponse"]["generatedSamples"][0]
-            download(sample["video"]["uri"], output_file)
-            break
-        time.sleep(3)
+    download_url = _poll_gemini_operation(gemini_api_key, operation_name)
+    download(download_url, output_file)
     return f"The video has been generated successfully to {output_file}"
 
 

tests/skills/test_video_generation.py

@@ -191,6 +191,91 @@ def test_gemini_post_raises_on_http_error(monkeypatch, tmp_path):
         vid.generate_video(str(pf), [], str(tmp_path / "v.mp4"), "16:9")
 
 
+def test_gemini_poll_is_bounded(monkeypatch):
+    # The MiniMax path bounds its polling; the Gemini operation path must too.
+    polls = {"count": 0}
+
+    def fake_get(url, headers=None, **kw):
+        polls["count"] += 1
+        return FakeResp(json_data={"done": False})
+
+    monkeypatch.setattr(vid.requests, "get", fake_get)
+    with pytest.raises(Exception, match="timed out"):
+        vid._poll_gemini_operation("g", "operations/op", max_attempts=3, interval=0)
+    assert polls["count"] == 3
+
+
+def test_gemini_operation_error_is_reported(monkeypatch, tmp_path):
+    # A finished operation carrying an error must not surface as a KeyError on
+    # the missing "response" field.
+    monkeypatch.setenv("GEMINI_API_KEY", "g")
+
+    def fake_post(url, headers=None, json=None, **kw):
+        return FakeResp(json_data={"name": "operations/op"})
+
+    def fake_get(url, headers=None, **kw):
+        return FakeResp(json_data={"done": True, "error": {"code": 3, "message": "prompt rejected"}})
+
+    monkeypatch.setattr(vid.requests, "post", fake_post)
+    monkeypatch.setattr(vid.requests, "get", fake_get)
+    pf = tmp_path / "p.json"
+    pf.write_text("x", encoding="utf-8")
+    out = tmp_path / "v.mp4"
+    with pytest.raises(Exception, match="prompt rejected"):
+        vid.generate_video(str(pf), [], str(out), "16:9")
+    assert not out.exists()
+
+
+def test_gemini_filtered_sample_reports_reason(monkeypatch, tmp_path):
+    monkeypatch.setenv("GEMINI_API_KEY", "g")
+    reason = "The generated video was filtered by the safety system."
+
+    def fake_post(url, headers=None, json=None, **kw):
+        return FakeResp(json_data={"name": "operations/op"})
+
+    def fake_get(url, headers=None, **kw):
+        return FakeResp(
+            json_data={
+                "done": True,
+                "response": {
+                    "generateVideoResponse": {
+                        "raiMediaFilteredCount": 1,
+                        "raiMediaFilteredReasons": [reason],
+                    }
+                },
+            }
+        )
+
+    monkeypatch.setattr(vid.requests, "post", fake_post)
+    monkeypatch.setattr(vid.requests, "get", fake_get)
+    pf = tmp_path / "p.json"
+    pf.write_text("x", encoding="utf-8")
+    with pytest.raises(Exception, match="filtered by the safety system"):
+        vid.generate_video(str(pf), [], str(tmp_path / "v.mp4"), "16:9")
+
+
+def test_gemini_sample_without_uri_is_reported(monkeypatch, tmp_path):
+    monkeypatch.setenv("GEMINI_API_KEY", "g")
+
+    def fake_post(url, headers=None, json=None, **kw):
+        return FakeResp(json_data={"name": "operations/op"})
+
+    def fake_get(url, headers=None, **kw):
+        return FakeResp(
+            json_data={
+                "done": True,
+                "response": {"generateVideoResponse": {"generatedSamples": [{"video": {}}]}},
+            }
+        )
+
+    monkeypatch.setattr(vid.requests, "post", fake_post)
+    monkeypatch.setattr(vid.requests, "get", fake_get)
+    pf = tmp_path / "p.json"
+    pf.write_text("x", encoding="utf-8")
+    with pytest.raises(Exception, match="without a video URI"):
+        vid.generate_video(str(pf), [], str(tmp_path / "v.mp4"), "16:9")
+
+
 def test_gemini_forwards_aspect_ratio_to_predict_request(monkeypatch, tmp_path):
     monkeypatch.setenv("GEMINI_API_KEY", "g")
     captured = {}