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
This commit is contained in:
Google Team Member
2026-01-21 14:27:03 -08:00
committed by Copybara-Service
parent 3e3566bd0e
commit 9579bea05d
2 changed files with 39 additions and 0 deletions
@@ -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"