feat(plugins): Add flush mechanism to BigQueryAgentAnalyticsPlugin

This change introduces a `flush` method to the `BigQueryAgentAnalyticsPlugin`. This ensures that all pending log events are written to BigQuery before the agent's run completes.

Key changes:

- Added `flush()` method to `BigQueryAgentAnalyticsPlugin` to force write of pending events.

PiperOrigin-RevId: 859263853

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

@@ -699,6 +699,13 @@ class BatchProcessor:
     self._batch_processor_task: Optional[asyncio.Task] = None
     self._shutdown = False
 
+  async def flush(self) -> None:
+    """Flushes the queue by waiting for it to be empty."""
+    if self._queue.empty():
+      return
+    # Wait for all items in the queue to be processed
+    await self._queue.join()
+
   async def start(self):
     """Starts the batch writer worker task."""
     if self._batch_processor_task is None:
@@ -1516,6 +1523,11 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
       logger.warning("Content formatter failed: %s", e)
       return "[FORMATTING FAILED]", False
 
+  async def flush(self) -> None:
+    """Flushes any pending events to BigQuery."""
+    if self.batch_processor:
+      await self.batch_processor.flush()
+
   async def _lazy_setup(self, **kwargs) -> None:
     """Performs lazy initialization of BigQuery clients and resources."""
     if self._started:
@@ -1947,6 +1959,8 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
     await self._log_event(
         "INVOCATION_COMPLETED", CallbackContext(invocation_context)
     )
+    # Ensure all logs are flushed before the agent returns
+    await self.flush()
 
   async def before_agent_callback(
       self, *, agent: Any, callback_context: CallbackContext, **kwargs

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

@@ -2061,3 +2061,28 @@ class TestBigQueryAgentAnalyticsPlugin:
       assert finished_spans[0].name == "test_span"
       assert format(finished_spans[0].context.span_id, "016x") == span_id
       assert format(finished_spans[0].context.trace_id, "032x") == trace_id
+
+  @pytest.mark.asyncio
+  async def test_flush_mechanism(
+      self,
+      bq_plugin_inst,
+      mock_write_client,
+      dummy_arrow_schema,
+      invocation_context,
+  ):
+    """Verifies that flush() forces pending events to be written."""
+    # Log an event
+    bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
+    await bq_plugin_inst.before_run_callback(
+        invocation_context=invocation_context
+    )
+
+    # Call flush - this should block until the event is written
+    await bq_plugin_inst.flush()
+
+    # Verify write called
+    mock_write_client.append_rows.assert_called_once()
+    log_entry = await _get_captured_event_dict_async(
+        mock_write_client, dummy_arrow_schema
+    )
+    assert log_entry["event_type"] == "INVOCATION_STARTING"