improve polling for FLUX3 video gen (#75010)

* wait between polls

tests/tools/test_flux3_video_tool.py

@@ -44,10 +44,14 @@ class _FakeResponse:
 
 
 class _FakeClient:
-    """Captures the one request each handler makes."""
+    """Captures each request a handler makes.
+
+    A list of responses is served in order, with the last one repeating, so a
+    poll that looks twice can be given a job that finishes between looks.
+    """
 
     def __init__(self, response, sink):
-        self._response = response
+        self._responses = list(response) if isinstance(response, list) else [response]
         self._sink = sink
 
     async def __aenter__(self):
@@ -58,9 +62,10 @@ class _FakeClient:
 
     async def request(self, method, url, headers=None, json=None):
         self._sink.append({"method": method, "url": url, "headers": headers or {}, "json": json})
-        if isinstance(self._response, Exception):
-            raise self._response
-        return self._response
+        response = self._responses[min(len(self._sink) - 1, len(self._responses) - 1)]
+        if isinstance(response, Exception):
+            raise response
+        return response
 
 
 class _FakeStream:
@@ -115,6 +120,13 @@ def _run(coro):
     return asyncio.run(coro)
 
 
+def _record_sleep(sink):
+    async def _sleep(seconds):
+        sink.append(seconds)
+
+    return _sleep
+
+
 def _call(handler, args, response, headers=None):
     """Invoke a handler with the transport stubbed; returns (parsed, requests)."""
     sink = []
@@ -249,18 +261,82 @@ class TestSubmitTransport:
         assert "error" in parsed
 
 
+@pytest.fixture(autouse=True)
+def _no_real_poll_wait(monkeypatch):
+    """Keep the in-call wait out of the test clock."""
+    monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 0)
+
+
 class TestPollTransport:
-    def test_poll_gets_the_job_and_returns_guidance(self):
-        response = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
+    def test_a_terminal_status_returns_at_once_without_waiting(self):
+        response = _FakeResponse(200, {"id": "bfl_job_1", "status": "Error", "guidance": "The job is over."})
 
         parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, response)
 
         assert requests[0]["method"] == "GET"
         assert requests[0]["url"] == f"{BASE_URL}/generations/bfl_job_1"
         assert requests[0]["json"] is None
-        assert parsed["result"] == "Still going."
+        assert len(requests) == 1
+        assert parsed["result"] == "The job is over."
+
+    def test_a_running_job_is_waited_out_inside_the_call(self, monkeypatch):
+        # A model has no clock, so telling it to pause produced a burst of polls
+        # instead of a paced one. The wait lives here where it cannot be skipped.
+        monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 45.0)
+        running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
+
+        slept = []
+        with patch.object(flux3.asyncio, "sleep", new=_record_sleep(slept)):
+            parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, running)
+
+        assert len(requests) == 2, "should look again after waiting"
+        assert sum(slept) == 45.0
         assert parsed["details"]["status"] == "Generating"
 
+    def test_the_wait_is_answerable_to_a_stop(self, monkeypatch):
+        # Nothing outside the tool can end a call that has already started —
+        # the executor only checks for an interrupt between tools — so /stop
+        # has to land inside the wait rather than at the end of it.
+        monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 45.0)
+        from tools import interrupt as interrupt_module
+
+        running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
+        looks = []
+
+        def _stop_after_one_slice():
+            looks.append(True)
+            return len(looks) > 1
+
+        slept = []
+        with patch.object(interrupt_module, "is_interrupted", _stop_after_one_slice), \
+                patch.object(flux3.asyncio, "sleep", new=_record_sleep(slept)):
+            parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, running)
+
+        assert slept == [flux3._POLL_WAIT_SLICE_SECONDS], "the rest of the wait is abandoned"
+        assert len(requests) == 1, "and so is the second look"
+        assert parsed["details"]["status"] == "Generating"
+
+    def test_the_call_returns_as_soon_as_the_job_finishes(self):
+        # The point of waiting in here is that the caller gets the result on the
+        # wait it was already taking, not one round trip later.
+        running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
+        done = _FakeResponse(200, {"id": "bfl_job_1", "status": "Error", "guidance": "That job failed."})
+
+        parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, [running, done])
+
+        assert len(requests) == 2
+        assert parsed["result"] == "That job failed."
+
+    def test_a_refusal_is_returned_immediately_rather_than_waited_on(self):
+        # A 429 carries its own retry guidance; sleeping on it would only delay
+        # showing the model what to do, and spend the poll budget twice.
+        response = _FakeResponse(429, {"error": {"message": "Too many polls. Wait 30 seconds."}})
+
+        parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, response)
+
+        assert len(requests) == 1
+        assert "Too many polls" in parsed["error"]
+
     def test_ready_saves_the_clip_and_never_returns_the_signed_url(self, tmp_path):
         # The signed URL is a bearer credential for the clip and it used to be
         # re-keyed into a shell command by hand, dropping characters. Neither

tools/flux3_video_tool.py

@@ -179,6 +179,58 @@ _MEDIA_FIELDS = {
 _MAX_IMAGES = 10
 
 
+# How long one get_result call waits before its second look. The bound that
+# matters is the whole call rather than this number: model_tools' async bridge
+# abandons a tool at 300s, and one call spends two polls (each bounded
+# server-side at 45s), this wait, and on Ready the download of the clip. There
+# is room to roughly double this; the reason not to is that a finished job is
+# only noticed at the next look, so the wait is also the notice delay.
+_POLL_FOLLOW_UP_WAIT_SECONDS = 45.0
+# Taken in slices so the wait is answerable. Nothing outside a tool can end a
+# call that has already started — the executor only checks for an interrupt
+# between tools — so a tool that blocks this long has to watch the flag itself.
+_POLL_WAIT_SLICE_SECONDS = 1.0
+
+# Mirrors the gateway's BFL statuses
+_TERMINAL_POLL_STATUSES = frozenset(
+    {"Ready", "Error", "Request Moderated", "Content Moderated", "Task not found"}
+)
+
+
+def _poll_is_finished(raw: str) -> bool:
+    """True when there is nothing to wait for: done, refused, or unreadable."""
+    try:
+        payload = json.loads(raw)
+    except Exception:
+        return True
+    if not isinstance(payload, dict) or "error" in payload:
+        # A refusal carries its own guidance (a wait, a limit, a dead job).
+        # Sleeping on it would only delay showing the model what to do.
+        return True
+    details = payload.get("details")
+    status = details.get("status") if isinstance(details, dict) else None
+    return not isinstance(status, str) or status in _TERMINAL_POLL_STATUSES
+
+
+async def _wait_before_second_look() -> bool:
+    """Hold the call open between looks; False if the user interrupted.
+
+    Counted down rather than clock-driven: this paces polling, so a slice
+    that runs long changes nothing, and the loop stays testable without a
+    fake clock.
+    """
+    from tools.interrupt import is_interrupted
+
+    remaining = _POLL_FOLLOW_UP_WAIT_SECONDS
+    while remaining > 0:
+        if is_interrupted():
+            return False
+        this_slice = min(_POLL_WAIT_SLICE_SECONDS, remaining)
+        await asyncio.sleep(this_slice)
+        remaining -= this_slice
+    return True
+
+
 def _warm_nous_token() -> None:
     """Refresh the Nous token once, before any parallel upload needs it.
 
@@ -471,8 +523,25 @@ async def _handle_get_result(args: dict, **kwargs) -> str:
         return _error("BFL video generation is not available in this build.")
     from urllib.parse import quote
 
-    raw = await _call_gateway("GET", f"{endpoints['base_url']}/generations/{quote(job_id.strip(), safe='')}")
-    return await _save_if_ready(raw, (args or {}).get("save_to"))
+    url = f"{endpoints['base_url']}/generations/{quote(job_id.strip(), safe='')}"
+    save_to = (args or {}).get("save_to")
+
+    raw = await _call_gateway("GET", url)
+    if _poll_is_finished(raw):
+        return await _save_if_ready(raw, save_to)
+
+    # Still running, so absorb the wait here instead of asking the model to
+    # take it. A model has no clock: told to wait it emits "I'll wait a minute"
+    # and its next action lands immediately, so the guidance produced a burst of
+    # polls rather than a paced one. Waiting inside the call cannot be skipped,
+    # needs no shell, and works the same on every platform. One call therefore
+    # covers a couple of minutes and returns as soon as a look finds it done.
+    if not await _wait_before_second_look():
+        # Interrupted mid-wait: hand back the status we already have rather
+        # than spending a round trip the user has just asked us to stop for.
+        return await _save_if_ready(raw, save_to)
+    raw = await _call_gateway("GET", url)
+    return await _save_if_ready(raw, save_to)
 
 
 async def _handle_prompting_guide(args: dict, **kwargs) -> str:
@@ -682,8 +751,9 @@ GET_RESULT_SCHEMA = {
     "name": "bfl_flux3_get_result",
     "description": (
         "Poll a FLUX 3 video job by the job id a generate tool returned. Generation takes minutes "
-        "and a long Generating phase is normal. Every response states the job's status, how long "
-        "to wait before polling again, and what you may do next — follow it rather than guessing. "
+        "and a long Generating phase is normal. This call waits for you while the job runs, so it "
+        "may take a couple of minutes; if it returns still generating, just call it again. Do not "
+        "sleep between calls. "
         "On Ready the clip is downloaded for you and the response gives its local path; your only "
         "remaining step is to deliver that file as the response describes."
     ),
@@ -813,11 +883,15 @@ auto-downscaled and output tops out at 720p.
 Submit returns a job id immediately — the video does not exist yet. Poll
 bfl_flux3_get_result with that id; generation takes several minutes and a long
 Generating phase is normal, not a stall. Nothing reaches disk before the job is
-Ready, so checking folders mid-run tells you nothing. Every response
-states the job status, how long to wait before polling again, and what you may
-do next — follow it literally rather than guessing an interval. A job survives
-client restarts: re-poll the same id rather than resubmitting, which would only
-spend your budgets on duplicate work.
+Ready, so checking folders mid-run tells you nothing.
+
+The waiting is not yours to do. bfl_flux3_get_result takes the pause itself
+while a job is still running, so one call can occupy a couple of minutes and
+comes back the moment the job finishes. If it returns still generating, just
+call it again — no sleeping, no interval to judge, nothing to time.
+
+A job survives client restarts: re-poll the same id rather than resubmitting,
+which would only spend your budgets on duplicate work.
 
 ## Save and deliver