Flush buffered events before the async tracking client's worker exits

ui/sdk/src/hamilton_sdk/adapters.py

@@ -469,6 +469,14 @@ async def ainit(self):
         self.initialized = True
         return self
 
+    async def stop(self):
+        """Flush any buffered tracking data and stop the background worker.
+
+        Call this before your event loop shuts down; nothing else flushes
+        the tracker's queue for you.
+        """
+        await self.client.stop()
+
     async def post_graph_construct(
         self, graph: h_graph.FunctionGraph, modules: list[ModuleType], config: dict[str, Any]
     ):

ui/sdk/src/hamilton_sdk/api/clients.py

@@ -585,9 +585,24 @@ def __init__(
         self.data_queue = asyncio.Queue()
         self.running = True
         self.max_batch_size = 100
+        self.worker_task: asyncio.Task | None = None
 
     async def ainit(self):
-        asyncio.create_task(self.worker())
+        self.worker_task = asyncio.create_task(self.worker())
+
+    async def stop(self):
+        """Signal the worker to flush whatever is buffered and exit.
+
+        Unlike BasicSynchronousHamiltonClient, this is not wired to atexit:
+        the worker is a task on this event loop, and by the time atexit
+        callbacks run at interpreter shutdown the loop may already be
+        closed, so callers must await this explicitly before the loop
+        that owns it stops running.
+        """
+        if self.worker_task is None or self.worker_task.done():
+            return
+        await self.data_queue.put(None)
+        await self.worker_task
 
     async def flush(self, batch):
         """Flush the batch (send it to the backend or process it)."""
@@ -624,14 +639,14 @@ async def worker(self):
             )
             try:
                 item = await asyncio.wait_for(self.data_queue.get(), timeout=self.flush_interval)
-                batch.append(item)
             except asyncio.TimeoutError:
                 # This is fine, we just keep waiting
                 pass
             else:
                 if item is None:
                     await self.flush(batch)
                     return
+                batch.append(item)
 
             # Check if batch is full or flush interval has passed
             if (

ui/sdk/tests/test_adapters.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import asyncio
 import os.path
 
 import pytest
@@ -162,6 +163,32 @@ def test_parallel_ray_sample_error():
         shutdown()
 
 
+def test_async_hamilton_tracker_stop_delegates_to_client():
+    """Regression for #1214: AsyncHamiltonTracker had no way for a caller to
+    flush its client's buffered tracking data before the event loop shuts
+    down. stop() must delegate to the underlying client's stop()."""
+
+    class _FakeAsyncClient:
+        def __init__(self, *args, **kwargs):
+            self.stop_calls = 0
+
+        async def stop(self):
+            self.stop_calls += 1
+
+    async def run():
+        tracker = adapters.AsyncHamiltonTracker(
+            project_id=1,
+            username="test-user",
+            dag_name="test-dag",
+            client_factory=_FakeAsyncClient,
+            api_key="test-key",
+        )
+        await tracker.stop()
+        assert tracker.client.stop_calls == 1
+
+    asyncio.run(run())
+
+
 if __name__ == "__main__":
     # test_adapters()
 

ui/sdk/tests/test_clients.py

@@ -81,3 +81,48 @@ async def run():
             await client.flush(batch)
 
     asyncio.run(run())
+
+
+def test_async_client_stop_flushes_buffered_item_and_exits_worker():
+    """Regression for #1214: worker()'s only exit path waits for a `None`
+    sentinel on data_queue, and nothing ever put one there for the async
+    client, so a buffered item sat unflushed until the flush_interval timer
+    happened to fire, or was lost outright if the process exited first.
+    stop() must send that sentinel and wait for the worker to flush and
+    exit, well before the flush_interval timeout."""
+
+    async def run():
+        client = _make_client()
+        session_cm = _mock_session_with_status(200)
+        with patch("aiohttp.ClientSession", return_value=session_cm):
+            await client.ainit()
+            await client.update_tasks(dag_run_id=1, attributes=[], task_updates=[])
+
+            assert not client.worker_task.done()
+
+            await client.stop()
+
+            assert client.worker_task.done()
+            assert client.data_queue.empty()
+            session_cm.__aenter__.return_value.put.assert_called_once()
+
+    asyncio.run(run())
+
+
+def test_async_client_stop_is_idempotent():
+    async def run():
+        client = _make_client()
+        with patch("aiohttp.ClientSession", return_value=_mock_session_with_status(200)):
+            await client.ainit()
+            await client.stop()
+            await client.stop()  # must not hang or raise
+
+    asyncio.run(run())
+
+
+def test_async_client_stop_before_ainit_is_a_noop():
+    async def run():
+        client = _make_client()
+        await client.stop()  # worker_task is None, must not raise
+
+    asyncio.run(run())