Drain run-history writes on shutdown

.changeset/mean-tips-dig.md

@@ -0,0 +1,5 @@
+---
+"gradio": patch
+---
+
+fix:Drain run-history writes on shutdown

gradio/route_utils.py

@@ -1157,6 +1157,34 @@ async def _delete_state_handler(app: App):
         await _cancel_background_task(task)
 
 
+HISTORY_DRAIN_TIMEOUT = 5.0
+
+
+@asynccontextmanager
+async def _drain_history_tasks(app: App, timeout: float = HISTORY_DRAIN_TIMEOUT):
+    """On shutdown, give in-flight run-history writes a bounded chance to finish.
+
+    A record is filed as a detached task so a prediction never waits on the Hub,
+    and nothing else ever awaits those tasks. Without this, a shutdown cancels
+    whatever is still in `record_run`, and a record that has not yet reached the
+    Hub call is lost outright.
+    """
+    try:
+        yield
+    finally:
+        # `state` itself is absent on the stand-in apps the lifespan tests use.
+        state = getattr(app, "state", None)
+        tasks = {
+            task
+            for task in getattr(state, "history_tasks", None) or ()
+            if not task.done()
+        }
+        if tasks:
+            _, pending = await asyncio.wait(tasks, timeout=timeout)
+            for task in pending:
+                task.cancel()
+
+
 def create_lifespan_handler(
     user_lifespan: Callable[[App], AbstractAsyncContextManager] | None,
     frequency: int | None = 1,
@@ -1173,6 +1201,9 @@ async def _handler(app: App):
                 await stack.enter_async_context(_lifespan_handler(app, frequency, age))
             if user_lifespan is not None:
                 state = await stack.enter_async_context(user_lifespan(app))
+            # Entered last so it unwinds first: the records are drained while
+            # the rest of the app is still up.
+            await stack.enter_async_context(_drain_history_tasks(app))
             yield state
 
     return _handler

test/test_history.py

@@ -458,6 +458,43 @@ def test_a_hub_failure_does_not_fail_the_prediction(self, recording_app):
         assert r.json()["data"] == ["hello world"]
 
 
+def test_a_record_still_in_flight_survives_shutdown(monkeypatch):
+    """Records are filed as detached tasks, so a shutdown that does not wait for
+    them loses whatever has not reached the Hub yet. The delay goes in the
+    prelude rather than in the Hub call on purpose: `anyio.to_thread.run_sync`
+    cannot interrupt its worker thread, so a write that has already started
+    finishes either way and would not tell the two behaviours apart."""
+    monkeypatch.delenv("GRADIO_HISTORY_BUCKET", raising=False)
+    hub = FakeHub()
+    original = history_mod.externalize_assets
+
+    async def slow_externalize(*args, **kwargs):
+        await asyncio.sleep(0.3)
+        return await original(*args, **kwargs)
+
+    def greet(name):
+        return f"hello {name}"
+
+    io = gr.Interface(greet, "text", "text", api_name="greet")
+    app, _, _ = io.launch(prevent_thread_lock=True)
+    with (
+        patch("gradio.history.resolve_token", return_value="tok"),
+        patch("gradio.history.HfApi", return_value=hub),
+        patch("gradio.history.externalize_assets", slow_externalize),
+        TestClient(app) as client,
+    ):
+        r = client.post(
+            "/gradio_api/run/greet",
+            json={"data": ["world"]},
+            headers={"X-Gradio-History-Bucket": "alice/hist"},
+        )
+        assert r.status_code == 200, r.text
+    # Leaving the block runs the lifespan shutdown, which is where the drain is.
+    assert len(hub.files) == 1, "the record was dropped by the shutdown"
+    io.close()
+    close_all()
+
+
 # ------------------------------------------------- end-to-end: workflow canvas