adk-python/tests/unittests/plugins/test_bigquery_logging_plugin.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import datetime
import json
import logging
from unittest import mock

from google.adk.agents import base_agent
from google.adk.agents import callback_context as callback_context_lib
from google.adk.agents import invocation_context as invocation_context_lib
from google.adk.events import event as event_lib
from google.adk.models import llm_request as llm_request_lib
from google.adk.models import llm_response as llm_response_lib
from google.adk.plugins import bigquery_logging_plugin
from google.adk.plugins import plugin_manager as plugin_manager_lib
from google.adk.sessions import base_session_service as base_session_service_lib
from google.adk.sessions import session as session_lib
from google.adk.tools import base_tool as base_tool_lib
from google.adk.tools import tool_context as tool_context_lib
import google.auth
from google.auth import exceptions as auth_exceptions
import google.auth.credentials
from google.cloud import bigquery
from google.genai import types
import pyarrow as pa
import pytest

BigQueryLoggerConfig = bigquery_logging_plugin.BigQueryLoggerConfig

PROJECT_ID = "test-gcp-project"
DATASET_ID = "adk_logs"
TABLE_ID = "agent_events"
DEFAULT_STREAM_NAME = (
    f"projects/{PROJECT_ID}/datasets/{DATASET_ID}/tables/{TABLE_ID}/_default"
)

# --- Pytest Fixtures ---


@pytest.fixture
def mock_session():
  mock_s = mock.create_autospec(
      session_lib.Session, instance=True, spec_set=True
  )
  type(mock_s).id = mock.PropertyMock(return_value="session-123")
  type(mock_s).user_id = mock.PropertyMock(return_value="user-456")
  type(mock_s).app_name = mock.PropertyMock(return_value="test_app")
  type(mock_s).state = mock.PropertyMock(return_value={})
  return mock_s


@pytest.fixture
def mock_agent():
  mock_a = mock.create_autospec(
      base_agent.BaseAgent, instance=True, spec_set=True
  )
  # Mock the 'name' property
  type(mock_a).name = mock.PropertyMock(return_value="MyTestAgent")
  return mock_a


@pytest.fixture
def invocation_context(mock_agent, mock_session):
  mock_session_service = mock.create_autospec(
      base_session_service_lib.BaseSessionService, instance=True, spec_set=True
  )
  mock_plugin_manager = mock.create_autospec(
      plugin_manager_lib.PluginManager, instance=True, spec_set=True
  )
  return invocation_context_lib.InvocationContext(
      agent=mock_agent,
      session=mock_session,
      invocation_id="inv-789",
      session_service=mock_session_service,
      plugin_manager=mock_plugin_manager,
  )


@pytest.fixture
def callback_context(invocation_context):
  return callback_context_lib.CallbackContext(
      invocation_context=invocation_context
  )


@pytest.fixture
def tool_context(invocation_context):
  return tool_context_lib.ToolContext(invocation_context=invocation_context)


@pytest.fixture
def mock_auth_default():
  mock_creds = mock.create_autospec(
      google.auth.credentials.Credentials, instance=True, spec_set=True
  )
  with mock.patch.object(
      google.auth,
      "default",
      autospec=True,
      return_value=(mock_creds, PROJECT_ID),
  ) as mock_auth:
    yield mock_auth


@pytest.fixture
def mock_bq_client():
  with mock.patch.object(bigquery, "Client", autospec=True) as mock_cls:
    yield mock_cls.return_value


@pytest.fixture
def mock_write_client():
  with mock.patch.object(
      bigquery_logging_plugin, "BigQueryWriteAsyncClient", autospec=True
  ) as mock_cls:
    mock_client = mock_cls.return_value
    mock_append_rows_response = mock.MagicMock()
    # Configure the 'row_errors' attribute on the mock object.
    mock_append_rows_response.row_errors = []
    mock_append_rows_response.error = mock.MagicMock()
    mock_append_rows_response.error.code = 0  # OK status

    mock_client.append_rows.return_value = _async_gen(mock_append_rows_response)
    yield mock_client


@pytest.fixture
def dummy_arrow_schema():
  return pa.schema([
      pa.field(
          "timestamp", pa.string()
      ),  # Store as string for simplicity in test
      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()),
  ])


@pytest.fixture
def mock_to_arrow_schema(dummy_arrow_schema):
  with mock.patch.object(
      bigquery_logging_plugin,
      "to_arrow_schema",
      autospec=True,
      return_value=dummy_arrow_schema,
  ) as mock_func:
    yield mock_func


@pytest.fixture
def mock_asyncio_to_thread():
  async def fake_to_thread(func, *args, **kwargs):
    return func(*args, **kwargs)

  with mock.patch(
      "asyncio.to_thread", side_effect=fake_to_thread
  ) as mock_async:
    yield mock_async


@pytest.fixture
def bq_plugin_inst(
    mock_auth_default,
    mock_bq_client,
    mock_write_client,
    mock_to_arrow_schema,
):
  plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin(
      project_id=PROJECT_ID,
      dataset_id=DATASET_ID,
      table_id=TABLE_ID,
  )
  # Trigger lazy initialization
  plugin._ensure_initialized_sync()
  mock_write_client.append_rows.reset_mock()
  return plugin


# --- Helper Functions ---


async def _async_gen(val):
  yield val


def _get_captured_event_dict(mock_write_client, expected_schema):
  """Helper to get the event_dict passed to append_rows."""
  mock_write_client.append_rows.assert_called_once()
  call_args = mock_write_client.append_rows.call_args
  requests = call_args.kwargs["requests"]
  assert len(requests) == 1
  request = requests[0]
  assert request.write_stream == DEFAULT_STREAM_NAME

  arrow_rows = request.arrow_rows
  message = pa.ipc.read_message(arrow_rows.rows.serialized_record_batch)
  batch = pa.ipc.read_record_batch(message, schema=expected_schema)
  table = pa.Table.from_batches([batch])
  assert table.schema.equals(
      expected_schema
  ), f"Schema mismatch: Expected {expected_schema}, got {table.schema}"
  pydict = table.to_pydict()
  return {k: v[0] for k, v in pydict.items()}


def _assert_common_fields(log_entry, event_type, agent="MyTestAgent"):
  assert log_entry["event_type"] == event_type
  assert log_entry["agent"] == agent
  assert log_entry["session_id"] == "session-123"
  assert log_entry["invocation_id"] == "inv-789"
  assert log_entry["user_id"] == "user-456"
  assert "timestamp" in log_entry
  assert isinstance(log_entry["timestamp"], str)


# --- Test Class ---


class TestBigQueryAgentAnalyticsPlugin:

  @pytest.mark.asyncio
  async def test_plugin_disabled(
      self,
      mock_auth_default,
      mock_bq_client,
      mock_write_client,
      invocation_context,
  ):
    config = BigQueryLoggerConfig(enabled=False)
    plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin(
        project_id=PROJECT_ID,
        dataset_id=DATASET_ID,
        table_id=TABLE_ID,
        config=config,
    )
    plugin._ensure_initialized_sync()  # Should do nothing
    user_message = types.Content(parts=[types.Part(text="Test")])

    await plugin.on_user_message_callback(
        invocation_context=invocation_context, user_message=user_message
    )
    mock_auth_default.assert_not_called()
    mock_bq_client.assert_not_called()
    mock_write_client.append_rows.assert_not_called()

  @pytest.mark.asyncio
  async def test_event_allowlist(
      self,
      mock_write_client,
      callback_context,
      invocation_context,
      mock_auth_default,
      mock_bq_client,
      mock_to_arrow_schema,
      dummy_arrow_schema,
  ):
    config = BigQueryLoggerConfig(event_allowlist=["LLM_REQUEST"])
    plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin(
        PROJECT_ID, DATASET_ID, TABLE_ID, config
    )
    plugin._ensure_initialized_sync()
    mock_write_client.append_rows.reset_mock()

    llm_request = llm_request_lib.LlmRequest(
        model="gemini-pro",
        contents=[types.Content(parts=[types.Part(text="Prompt")])],
    )
    await plugin.before_model_callback(
        callback_context=callback_context, llm_request=llm_request
    )
    mock_write_client.append_rows.assert_called_once()
    mock_write_client.append_rows.reset_mock()

    user_message = types.Content(parts=[types.Part(text="What is up?")])
    await plugin.on_user_message_callback(
        invocation_context=invocation_context, user_message=user_message
    )
    mock_write_client.append_rows.assert_not_called()

  @pytest.mark.asyncio
  async def test_event_denylist(
      self,
      mock_write_client,
      invocation_context,
      mock_auth_default,
      mock_bq_client,
      mock_to_arrow_schema,
      dummy_arrow_schema,
  ):
    config = BigQueryLoggerConfig(event_denylist=["USER_MESSAGE_RECEIVED"])
    plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin(
        PROJECT_ID, DATASET_ID, TABLE_ID, config
    )
    plugin._ensure_initialized_sync()
    mock_write_client.append_rows.reset_mock()

    user_message = types.Content(parts=[types.Part(text="What is up?")])
    await plugin.on_user_message_callback(
        invocation_context=invocation_context, user_message=user_message
    )
    mock_write_client.append_rows.assert_not_called()

    await plugin.before_run_callback(invocation_context=invocation_context)
    mock_write_client.append_rows.assert_called_once()

  @pytest.mark.asyncio
  async def test_content_formatter(
      self,
      mock_write_client,
      invocation_context,
      mock_auth_default,
      mock_bq_client,
      mock_to_arrow_schema,
      dummy_arrow_schema,
  ):
    def redact_content(content):
      return "[REDACTED]"

    config = BigQueryLoggerConfig(content_formatter=redact_content)
    plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin(
        PROJECT_ID, DATASET_ID, TABLE_ID, config
    )
    plugin._ensure_initialized_sync()
    mock_write_client.append_rows.reset_mock()

    user_message = types.Content(parts=[types.Part(text="Secret message")])
    await plugin.on_user_message_callback(
        invocation_context=invocation_context, user_message=user_message
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "USER_MESSAGE_RECEIVED")
    assert log_entry["content"] == "[REDACTED]"

  @pytest.mark.asyncio
  async def test_content_formatter_error(
      self,
      mock_write_client,
      invocation_context,
      mock_auth_default,
      mock_bq_client,
      mock_to_arrow_schema,
      dummy_arrow_schema,
  ):
    def error_formatter(content):
      raise ValueError("Formatter failed")

    config = BigQueryLoggerConfig(content_formatter=error_formatter)
    plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin(
        PROJECT_ID, DATASET_ID, TABLE_ID, config
    )
    plugin._ensure_initialized_sync()
    mock_write_client.append_rows.reset_mock()
    user_message = types.Content(parts=[types.Part(text="Test")])

    with mock.patch.object(logging, "warning") as mock_log_warning:
      await plugin.on_user_message_callback(
          invocation_context=invocation_context, user_message=user_message
      )
      mock_log_warning.assert_called_once_with(
          "Error applying custom content formatter for event type %s: %s",
          "USER_MESSAGE_RECEIVED",
          mock.ANY,
      )

    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "USER_MESSAGE_RECEIVED")
    assert "User Content: text: 'Test'" in log_entry["content"]

  @pytest.mark.asyncio
  async def test_on_user_message_callback_logs_correctly(
      self,
      bq_plugin_inst,
      mock_write_client,
      invocation_context,
      dummy_arrow_schema,
  ):
    user_message = types.Content(parts=[types.Part(text="What is up?")])
    await bq_plugin_inst.on_user_message_callback(
        invocation_context=invocation_context, user_message=user_message
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "USER_MESSAGE_RECEIVED")
    assert log_entry["content"] == "User Content: text: 'What is up?'"

  @pytest.mark.asyncio
  async def test_on_event_callback_tool_call(
      self,
      bq_plugin_inst,
      mock_write_client,
      invocation_context,
      dummy_arrow_schema,
  ):
    tool_fc = types.FunctionCall(name="get_weather", args={"location": "Paris"})
    event = event_lib.Event(
        author="MyTestAgent",
        content=types.Content(parts=[types.Part(function_call=tool_fc)]),
        timestamp=datetime.datetime(
            2025, 10, 22, 10, 0, 0, tzinfo=datetime.timezone.utc
        ).timestamp(),
    )
    await bq_plugin_inst.on_event_callback(
        invocation_context=invocation_context, event=event
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "TOOL_CALL")
    logged_content = json.loads(log_entry["content"])
    assert logged_content[0]["function_call"]["args"] == {"location": "Paris"}
    assert logged_content[0]["function_call"]["name"] == "get_weather"
    assert log_entry["timestamp"] == "2025-10-22T10:00:00+00:00"

  @pytest.mark.asyncio
  async def test_on_event_callback_model_response(
      self,
      bq_plugin_inst,
      mock_write_client,
      invocation_context,
      dummy_arrow_schema,
  ):
    event = event_lib.Event(
        author="MyTestAgent",
        content=types.Content(parts=[types.Part(text="Hello there!")]),
        timestamp=datetime.datetime(
            2025, 10, 22, 11, 0, 0, tzinfo=datetime.timezone.utc
        ).timestamp(),
    )
    await bq_plugin_inst.on_event_callback(
        invocation_context=invocation_context, event=event
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "MODEL_RESPONSE")
    logged_content = json.loads(log_entry["content"])
    assert logged_content[0]["text"] == "Hello there!"
    assert log_entry["timestamp"] == "2025-10-22T11:00:00+00:00"

  @pytest.mark.asyncio
  async def test_bigquery_client_initialization_failure(
      self, mock_auth_default, mock_write_client, invocation_context
  ):
    mock_auth_default.side_effect = auth_exceptions.GoogleAuthError(
        "Auth failed"
    )
    plugin_with_fail = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin(
        project_id=PROJECT_ID,
        dataset_id=DATASET_ID,
        table_id=TABLE_ID,
    )
    with mock.patch.object(logging, "exception") as mock_log_exception:
      await plugin_with_fail.on_user_message_callback(
          invocation_context=invocation_context,
          user_message=types.Content(parts=[types.Part(text="Test")]),
      )
      mock_log_exception.assert_called_once_with(
          "Failed to initialize BigQuery client or table: %s", mock.ANY
      )
    mock_write_client.append_rows.assert_not_called()

  @pytest.mark.asyncio
  async def test_bigquery_insert_error_does_not_raise(
      self, bq_plugin_inst, mock_write_client, invocation_context
  ):
    mock_append_rows_response = mock.MagicMock()
    mock_append_rows_response.row_errors = [mock.MagicMock()]
    mock_append_rows_response.error = mock.MagicMock()
    mock_append_rows_response.error.code = 0
    mock_write_client.append_rows.return_value = _async_gen(
        mock_append_rows_response
    )

    with mock.patch.object(logging, "error") as mock_log_error:
      await bq_plugin_inst.on_user_message_callback(
          invocation_context=invocation_context,
          user_message=types.Content(parts=[types.Part(text="Test")]),
      )
      mock_log_error.assert_called_with(
          "Errors occurred while writing to BigQuery (Storage Write API): %s",
          mock_append_rows_response.row_errors,
      )
    mock_write_client.append_rows.assert_called_once()

  @pytest.mark.asyncio
  async def test_before_run_callback_logs_correctly(
      self,
      bq_plugin_inst,
      mock_write_client,
      invocation_context,
      dummy_arrow_schema,
  ):
    await bq_plugin_inst.before_run_callback(
        invocation_context=invocation_context
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "INVOCATION_STARTING")
    assert log_entry["content"] is None

  @pytest.mark.asyncio
  async def test_after_run_callback_logs_correctly(
      self,
      bq_plugin_inst,
      mock_write_client,
      invocation_context,
      dummy_arrow_schema,
  ):
    await bq_plugin_inst.after_run_callback(
        invocation_context=invocation_context
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "INVOCATION_COMPLETED")
    assert log_entry["content"] is None

  @pytest.mark.asyncio
  async def test_before_agent_callback_logs_correctly(
      self,
      bq_plugin_inst,
      mock_write_client,
      mock_agent,
      callback_context,
      dummy_arrow_schema,
  ):
    await bq_plugin_inst.before_agent_callback(
        agent=mock_agent, callback_context=callback_context
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "AGENT_STARTING")
    assert log_entry["content"] == "Agent Name: MyTestAgent"

  @pytest.mark.asyncio
  async def test_after_agent_callback_logs_correctly(
      self,
      bq_plugin_inst,
      mock_write_client,
      mock_agent,
      callback_context,
      dummy_arrow_schema,
  ):
    await bq_plugin_inst.after_agent_callback(
        agent=mock_agent, callback_context=callback_context
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "AGENT_COMPLETED")
    assert log_entry["content"] == "Agent Name: MyTestAgent"

  @pytest.mark.asyncio
  async def test_before_model_callback_logs_correctly(
      self,
      bq_plugin_inst,
      mock_write_client,
      callback_context,
      dummy_arrow_schema,
  ):
    llm_request = llm_request_lib.LlmRequest(
        model="gemini-pro",
        contents=[types.Content(parts=[types.Part(text="Prompt")])],
        config=types.GenerateContentConfig(
            temperature=0.5,
            top_p=0.9,
            max_output_tokens=100,
            system_instruction=types.Content(
                parts=[types.Part(text="Be helpful")]
            ),
        ),
        tools_dict={
            "my_tool": mock.create_autospec(
                base_tool_lib.BaseTool, instance=True, spec_set=True
            )
        },
    )
    await bq_plugin_inst.before_model_callback(
        callback_context=callback_context, llm_request=llm_request
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "LLM_REQUEST")
    assert "Model: gemini-pro" in log_entry["content"]
    assert "System Prompt: Be helpful" in log_entry["content"]
    assert (
        "Params: {temperature=0.5, top_p=0.9, max_output_tokens=100}"
        in log_entry["content"]
    )
    assert "Available Tools: ['my_tool']" in log_entry["content"]

  @pytest.mark.asyncio
  async def test_after_model_callback_text_response(
      self,
      bq_plugin_inst,
      mock_write_client,
      callback_context,
      dummy_arrow_schema,
  ):
    llm_response = llm_response_lib.LlmResponse(
        content=types.Content(parts=[types.Part(text="Model response")]),
        usage_metadata=types.UsageMetadata(
            prompt_token_count=10, total_token_count=15
        ),
    )
    await bq_plugin_inst.after_model_callback(
        callback_context=callback_context, llm_response=llm_response
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "LLM_RESPONSE")
    assert (
        "Tool Name: text_response, text: 'Model response'"
        in log_entry["content"]
    )
    assert "Token Usage: {prompt: 10," in log_entry["content"]
    assert log_entry["error_message"] is None

  @pytest.mark.asyncio
  async def test_after_model_callback_tool_call(
      self,
      bq_plugin_inst,
      mock_write_client,
      callback_context,
      dummy_arrow_schema,
  ):
    llm_response = llm_response_lib.LlmResponse(
        content=types.Content(
            parts=[
                types.Part(
                    function_call=types.FunctionCall(name="tool1", args={})
                )
            ]
        ),
    )
    await bq_plugin_inst.after_model_callback(
        callback_context=callback_context, llm_response=llm_response
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "LLM_RESPONSE")
    assert "Tool Name: tool1" in log_entry["content"]

  @pytest.mark.asyncio
  async def test_before_tool_callback_logs_correctly(
      self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema
  ):
    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="Does something"
    )
    await bq_plugin_inst.before_tool_callback(
        tool=mock_tool, tool_args={"param": "value"}, tool_context=tool_context
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "TOOL_STARTING")
    assert "Tool Name: MyTool" in log_entry["content"]
    assert "Arguments: {'param': 'value'}" in log_entry["content"]

  @pytest.mark.asyncio
  async def test_after_tool_callback_logs_correctly(
      self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema
  ):
    mock_tool = mock.create_autospec(
        base_tool_lib.BaseTool, instance=True, spec_set=True
    )
    type(mock_tool).name = mock.PropertyMock(return_value="MyTool")
    await bq_plugin_inst.after_tool_callback(
        tool=mock_tool,
        tool_args={},
        tool_context=tool_context,
        result={"status": "success"},
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "TOOL_COMPLETED")
    assert "Tool Name: MyTool" in log_entry["content"]
    assert "Result: {'status': 'success'}" in log_entry["content"]

  @pytest.mark.asyncio
  async def test_on_model_error_callback_logs_correctly(
      self,
      bq_plugin_inst,
      mock_write_client,
      callback_context,
      dummy_arrow_schema,
  ):
    llm_request = llm_request_lib.LlmRequest(
        model="gemini-pro",
        contents=[types.Content(parts=[types.Part(text="Prompt")])],
    )
    error = ValueError("LLM failed")
    await bq_plugin_inst.on_model_error_callback(
        callback_context=callback_context, llm_request=llm_request, error=error
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "LLM_ERROR")
    assert log_entry["content"] is None
    assert log_entry["error_message"] == "LLM failed"

  @pytest.mark.asyncio
  async def test_on_tool_error_callback_logs_correctly(
      self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema
  ):
    mock_tool = mock.create_autospec(
        base_tool_lib.BaseTool, instance=True, spec_set=True
    )
    type(mock_tool).name = mock.PropertyMock(return_value="MyTool")
    error = TimeoutError("Tool timed out")
    await bq_plugin_inst.on_tool_error_callback(
        tool=mock_tool,
        tool_args={"param": "value"},
        tool_context=tool_context,
        error=error,
    )
    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
    _assert_common_fields(log_entry, "TOOL_ERROR")
    assert "Tool Name: MyTool" in log_entry["content"]
    assert "Arguments: {'param': 'value'}" in log_entry["content"]
    assert log_entry["error_message"] == "Tool timed out"