fix: Enhance BigQuery Plugin Robustness and Schema Accuracy

This update improves the `BigQueryAgentAnalyticsPlugin` in several ways:

*   Corrects the PyArrow schema generation to accurately reflect BigQuery field nullability based on the `mode` attribute.
*   Introduces a configurable `shutdown_timeout` in `BigQueryLoggerConfig` to manage how long the plugin waits for pending logs to flush during shutdown.
*   Adds more robust error handling within the `shutdown` method and background write tasks, particularly for event loop closure issues.
*   Improves internal logging to provide better diagnostics.
*   Ensures consistent use of safe content formatting.

PiperOrigin-RevId: 831225837

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

@@ -146,14 +146,14 @@ def mock_write_client():
 @pytest.fixture
 def dummy_arrow_schema():
   return pa.schema([
-      pa.field("timestamp", pa.timestamp("us", tz="UTC")),
-      pa.field("event_type", pa.string()),
-      pa.field("agent", pa.string()),
-      pa.field("session_id", pa.string()),
-      pa.field("invocation_id", pa.string()),
-      pa.field("user_id", pa.string()),
-      pa.field("content", pa.string()),
-      pa.field("error_message", pa.string()),
+      pa.field("timestamp", pa.timestamp("us", tz="UTC"), nullable=False),
+      pa.field("event_type", pa.string(), nullable=True),
+      pa.field("agent", pa.string(), nullable=True),
+      pa.field("session_id", pa.string(), nullable=True),
+      pa.field("invocation_id", pa.string(), nullable=True),
+      pa.field("user_id", pa.string(), nullable=True),
+      pa.field("content", pa.string(), nullable=True),
+      pa.field("error_message", pa.string(), nullable=True),
   ])
 
 
@@ -391,6 +391,102 @@ class TestBigQueryAgentAnalyticsPlugin:
     log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
     assert log_entry["content"] == "User Content: [FORMATTING FAILED]"
 
+  @pytest.mark.asyncio
+  async def test_max_content_length(
+      self,
+      mock_write_client,
+      invocation_context,
+      callback_context,
+      mock_auth_default,
+      mock_bq_client,
+      mock_to_arrow_schema,
+      dummy_arrow_schema,
+      mock_asyncio_to_thread,
+  ):
+    config = BigQueryLoggerConfig(max_content_length=40)
+    plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
+        PROJECT_ID, DATASET_ID, TABLE_ID, config
+    )
+    await plugin._ensure_init()
+    mock_write_client.append_rows.reset_mock()
+
+    # Test User Message Truncation
+    user_message = types.Content(
+        parts=[types.Part(text="12345678901234567890123456789012345678901")]
+    )  # 41 chars
+    await plugin.on_user_message_callback(
+        invocation_context=invocation_context, user_message=user_message
+    )
+    await asyncio.sleep(0.01)
+    mock_write_client.append_rows.assert_called_once()
+    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
+    assert (
+        log_entry["content"]
+        == "User Content: text: '1234567890123456789012345678901234567890...' "
+    )
+    mock_write_client.append_rows.reset_mock()
+
+    # Test before_model_callback full content truncation
+    llm_request = llm_request_lib.LlmRequest(
+        model="gemini-pro",
+        config=types.GenerateContentConfig(
+            system_instruction=types.Content(
+                parts=[types.Part(text="System Instruction")]
+            )
+        ),
+        contents=[
+            types.Content(role="user", parts=[types.Part(text="Prompt")])
+        ],
+    )
+    await plugin.before_model_callback(
+        callback_context=callback_context, llm_request=llm_request
+    )
+    await asyncio.sleep(0.01)
+    mock_write_client.append_rows.assert_called_once()
+    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
+    # Full content: "Model: gemini-pro | Prompt: user: text: 'Prompt' | System Prompt: System Instruction"
+    # Truncated to 40 chars + ...:
+    expected_content = "Model: gemini-pro | Prompt: user: text: ..."
+    assert log_entry["content"] == expected_content
+
+  @pytest.mark.asyncio
+  async def test_max_content_length_tool_args(
+      self,
+      mock_write_client,
+      tool_context,
+      mock_auth_default,
+      mock_bq_client,
+      mock_to_arrow_schema,
+      dummy_arrow_schema,
+      mock_asyncio_to_thread,
+  ):
+    config = BigQueryLoggerConfig(max_content_length=10)
+    plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
+        PROJECT_ID, DATASET_ID, TABLE_ID, config
+    )
+    await plugin._ensure_init()
+    mock_write_client.append_rows.reset_mock()
+
+    mock_tool = mock.create_autospec(
+        base_tool_lib.BaseTool, instance=True, spec_set=True
+    )
+    type(mock_tool).name = mock.PropertyMock(return_value="MyTool")
+    type(mock_tool).description = mock.PropertyMock(return_value="Description")
+
+    # Args length > 10
+    # {"param": "long_value"} is ~24 chars
+    await plugin.before_tool_callback(
+        tool=mock_tool,
+        tool_args={"param": "long_value"},
+        tool_context=tool_context,
+    )
+    await asyncio.sleep(0.01)
+    mock_write_client.append_rows.assert_called_once()
+    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
+    # JSON string: '{"param": "long_value"}'
+    # Truncated to 10: '{"param": ...'
+    assert 'Arguments: {"param": ...' in log_entry["content"]
+
   @pytest.mark.asyncio
   async def test_on_user_message_callback_logs_correctly(
       self,
@@ -430,7 +526,7 @@ class TestBigQueryAgentAnalyticsPlugin:
     await asyncio.sleep(0.01)
     log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
     _assert_common_fields(log_entry, "TOOL_CALL", agent="MyTestAgent")
-    assert '"name": "get_weather"' in log_entry["content"]
+    assert "call: get_weather" in log_entry["content"]
     assert log_entry["timestamp"] == datetime.datetime(
         2025, 10, 22, 10, 0, 0, tzinfo=datetime.timezone.utc
     )
@@ -456,7 +552,7 @@ class TestBigQueryAgentAnalyticsPlugin:
     await asyncio.sleep(0.01)
     log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
     _assert_common_fields(log_entry, "MODEL_RESPONSE", agent="MyTestAgent")
-    assert '"text": "Hello there!"' in log_entry["content"]
+    assert "text: 'Hello there!'" in log_entry["content"]
     assert log_entry["timestamp"] == datetime.datetime(
         2025, 10, 22, 11, 0, 0, tzinfo=datetime.timezone.utc
     )
@@ -485,7 +581,7 @@ class TestBigQueryAgentAnalyticsPlugin:
           user_message=types.Content(parts=[types.Part(text="Test")]),
       )
       await asyncio.sleep(0.01)
-      mock_log_error.assert_any_call("BQ Init Failed: Auth failed")
+      mock_log_error.assert_any_call("BQ Plugin: Init Failed:", exc_info=True)
     mock_write_client.append_rows.assert_not_called()
 
   @pytest.mark.asyncio
@@ -509,16 +605,20 @@ class TestBigQueryAgentAnalyticsPlugin:
           user_message=types.Content(parts=[types.Part(text="Test")]),
       )
       await asyncio.sleep(0.01)
-      mock_log_error.assert_called_with("BQ Write Error: Test BQ Error")
+      mock_log_error.assert_called_with("BQ Plugin: Write Error: Test BQ Error")
     mock_write_client.append_rows.assert_called_once()
 
   @pytest.mark.asyncio
   async def test_close(self, bq_plugin_inst, mock_bq_client, mock_write_client):
     await bq_plugin_inst.close()
     mock_write_client.transport.close.assert_called_once()
-    mock_bq_client.close.assert_called_once()
+    # bq_client might not be closed if it wasn't created or if close() failed,
+    # but here it should be.
+    # in the new implementation we verify attributes are reset
+    assert bq_plugin_inst._write_client is None
+    assert bq_plugin_inst._bq_client is None
+    assert bq_plugin_inst._is_shutting_down is False
 
-  # ... other tests remain the same ...
   @pytest.mark.asyncio
   async def test_before_run_callback_logs_correctly(
       self,
@@ -595,7 +695,9 @@ class TestBigQueryAgentAnalyticsPlugin:
   ):
     llm_request = llm_request_lib.LlmRequest(
         model="gemini-pro",
-        contents=[types.Content(parts=[types.Part(text="Prompt")])],
+        contents=[
+            types.Content(role="user", parts=[types.Part(text="Prompt")])
+        ],
     )
     await bq_plugin_inst.before_model_callback(
         callback_context=callback_context, llm_request=llm_request
@@ -603,7 +705,11 @@ class TestBigQueryAgentAnalyticsPlugin:
     await asyncio.sleep(0.01)
     log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
     _assert_common_fields(log_entry, "LLM_REQUEST")
-    assert log_entry["content"] == "Model: gemini-pro | System Prompt: Empty"
+    assert (
+        log_entry["content"]
+        == "Model: gemini-pro | Prompt: user: text: 'Prompt' | System Prompt:"
+        " Empty"
+    )
 
   @pytest.mark.asyncio
   async def test_after_model_callback_text_response(