fix(framework): Avoid race condition in agent publishing

framework/py/flwr/supercore/task_process/agent/session.py

@@ -20,7 +20,7 @@
 import json
 import time
 from collections.abc import Sequence
-from queue import Empty, Full, Queue
+from queue import Empty, Queue
 from threading import Lock, Thread
 from typing import Literal, cast
 
@@ -106,11 +106,8 @@ def close(self, timeout: float | None = None) -> None:
             self._raise_worker_error()
             return
 
-        try:
-            self._queue.put_nowait(_EVENT_PUBLISH_STOP)
-        except Full:
-            pass  # The worker will still stop due to the `_closed` flag.
         self._closed = True
+        self._queue.put(_EVENT_PUBLISH_STOP)
         self._worker.join(timeout)
         if self._worker.is_alive():
             raise TimeoutError("Timed out waiting for Agent event publisher to stop.")
@@ -127,7 +124,7 @@ def _flush(self, batch: list[TaskEvent]) -> None:
 
     def _run(self) -> None:
         """Upload queued events in small batches."""
-        while not self._closed:
+        while True:
             item = self._queue.get()
             if item is _EVENT_PUBLISH_STOP:
                 return

framework/py/flwr/supercore/task_process/agent/session_test.py

@@ -69,6 +69,31 @@ def test_emit_event_pushes_task_event() -> None:
     )
 
 
+def test_close_drains_events_before_worker_stops() -> None:
+    """Close should publish queued events before stopping the worker."""
+    stub = Mock()
+    with patch("flwr.supercore.task_process.agent.session.Thread") as thread_cls:
+        thread_cls.return_value.is_alive.return_value = False
+        events = RuntimeAgentEvents(stub)
+        worker_target = thread_cls.call_args.kwargs["target"]
+        thread_cls.return_value.join.side_effect = lambda _timeout: worker_target()
+
+        event: JSONObject = {
+            "type": "response.output_text.delta",
+            "delta": "Hello",
+        }
+        events.emit(event)
+        events.close()
+
+    expected_event = TaskEvent(
+        event="response.output_text.delta",
+        data='{"type":"response.output_text.delta","delta":"Hello"}',
+    )
+    stub.PushTaskEvents.assert_called_once_with(
+        PushTaskEventsRequest(events=[expected_event])
+    )
+
+
 def test_emit_event_requires_type() -> None:
     """Emit should reject events without a valid type."""
     stub = Mock()