wehub/google--adk-python
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
73fe961a8faf182ba47f9ede20fe7c59509dfc68
google--adk-python/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py
T
wehub-resource-sync ec2b666284
Continuous Integration / Pre-commit Linter (push) Waiting to run
Details
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Details
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Details
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Details
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Details
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Details
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Details
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Details
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Details
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Details
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Details
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Details
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Details
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Details
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Details
Copybara PR Handler / close-imported-pr (push) Waiting to run
Details
chore: import upstream snapshot with attribution
2026-07-13 13:25:13 +08:00

9701 lines
336 KiB
Python
Raw Blame History

# Copyright 2026 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 asyncio
import contextlib
import dataclasses
import json
import os
from unittest import mock
from google.adk.agents import base_agent
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.invocation_context import InvocationContext
from google.adk.events import event as event_lib
from google.adk.events import event_actions as event_actions_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_agent_analytics_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
from google.adk.utils._telemetry_context import _is_visual_builder
from google.adk.version import __version__
import google.auth
from google.auth import exceptions as auth_exceptions
import google.auth.credentials
from google.cloud import bigquery
from google.cloud import exceptions as cloud_exceptions
from google.genai import types
from opentelemetry import trace
import pyarrow as pa
import pytest
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")
type(mock_a).instruction = mock.PropertyMock(return_value="Test Instruction")
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 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 CallbackContext(invocation_context=invocation_context)
@pytest.fixture
def tool_context(invocation_context):
return tool_context_lib.ToolContext(invocation_context=invocation_context)
class FakeCredentials(google.auth.credentials.Credentials):
def __init__(self):
pass
def refresh(self, request):
pass
@pytest.fixture
def mock_auth_default():
mock_creds = FakeCredentials()
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_agent_analytics_plugin, "BigQueryWriteAsyncClient", autospec=True
) as mock_cls:
mock_client = mock_cls.return_value
mock_client.transport = mock.AsyncMock()
async def fake_append_rows(requests, **kwargs):
# This function is now async, so `await client.append_rows` works.
mock_append_rows_response = mock.MagicMock()
mock_append_rows_response.row_errors = []
mock_append_rows_response.error = mock.MagicMock()
mock_append_rows_response.error.code = 0 # OK status
# This a gen is what's returned *after* the await.
return _async_gen(mock_append_rows_response)
mock_client.append_rows.side_effect = fake_append_rows
yield mock_client
@pytest.fixture
def dummy_arrow_schema():
return pa.schema([
pa.field("timestamp", pa.timestamp("us", tz="UTC"), nullable=False),
pa.field("root_agent_name", pa.string(), nullable=True),
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("trace_id", pa.string(), nullable=True),
pa.field("span_id", pa.string(), nullable=True),
pa.field("parent_span_id", pa.string(), nullable=True),
pa.field(
"content", pa.string(), nullable=True
), # JSON stored as string in Arrow
pa.field(
"content_parts",
pa.list_(
pa.struct([
pa.field("mime_type", pa.string(), nullable=True),
pa.field("uri", pa.string(), nullable=True),
pa.field(
"object_ref",
pa.struct([
pa.field("uri", pa.string(), nullable=True),
pa.field("authorizer", pa.string(), nullable=True),
pa.field("version", pa.string(), nullable=True),
pa.field(
"details",
pa.string(),
nullable=True,
metadata={
b"ARROW:extension:name": (
b"google:sqlType:json"
)
},
),
]),
nullable=True,
),
pa.field("text", pa.string(), nullable=True),
pa.field("part_index", pa.int64(), nullable=True),
pa.field("part_attributes", pa.string(), nullable=True),
pa.field("storage_mode", pa.string(), nullable=True),
])
),
nullable=True,
),
pa.field("attributes", pa.string(), nullable=True),
pa.field("latency_ms", pa.string(), nullable=True),
pa.field("status", pa.string(), nullable=True),
pa.field("error_message", pa.string(), nullable=True),
pa.field("is_truncated", pa.bool_(), nullable=True),
])
@pytest.fixture
def mock_to_arrow_schema(dummy_arrow_schema):
with mock.patch.object(
bigquery_agent_analytics_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 mock_storage_client():
with mock.patch("google.cloud.storage.Client") as mock_client:
yield mock_client
@pytest.fixture
async def bq_plugin_inst(
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
await plugin._ensure_started() # Ensure clients are initialized
mock_write_client.append_rows.reset_mock()
yield plugin
await plugin.shutdown()
@contextlib.asynccontextmanager
async def managed_plugin(*args, **kwargs):
"""Async context manager to ensure plugin shutdown."""
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
*args, **kwargs
)
try:
yield plugin
finally:
await plugin.shutdown()
# --- Helper Functions ---
async def _async_gen(val):
yield val
async def _get_captured_event_dict_async(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_iter = call_args.args[0]
requests = []
if hasattr(requests_iter, "__aiter__"):
async for req in requests_iter:
requests.append(req)
else:
requests = list(requests_iter)
assert len(requests) == 1
request = requests[0]
assert request.write_stream == DEFAULT_STREAM_NAME
assert request.trace_id.startswith("google-adk-bq-logger")
assert request.trace_id.endswith(f"/{__version__}")
# Parse the Arrow batch back to a dict for verification
try:
reader = pa.ipc.open_stream(request.arrow_rows.rows.serialized_record_batch)
table = reader.read_all()
except Exception:
# Fallback: try reading as a single batch
buf = pa.py_buffer(request.arrow_rows.rows.serialized_record_batch)
batch = pa.ipc.read_record_batch(buf, 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()}
async def _get_captured_rows_async(mock_write_client, expected_schema):
"""Helper to get all rows passed to append_rows."""
all_rows = []
for call in mock_write_client.append_rows.call_args_list:
requests_iter = call.args[0]
requests = []
if hasattr(requests_iter, "__aiter__"):
async for req in requests_iter:
requests.append(req)
else:
requests = list(requests_iter)
for request in requests:
# Parse the Arrow batch back to a dict for verification
try:
reader = pa.ipc.open_stream(
request.arrow_rows.rows.serialized_record_batch
)
table = reader.read_all()
except Exception:
# Fallback: try reading as a single batch
buf = pa.py_buffer(request.arrow_rows.rows.serialized_record_batch)
batch = pa.ipc.read_record_batch(buf, expected_schema)
table = pa.Table.from_batches([batch])
pydict = table.to_pylist()
all_rows.extend(pydict)
return all_rows
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"
def test_recursive_smart_truncate():
"""Test recursive smart truncate."""
obj = {
"a": "long string" * 10,
"b": ["short", "long string" * 10],
"c": {"d": "long string" * 10},
}
max_len = 10
truncated, is_truncated = (
bigquery_agent_analytics_plugin._recursive_smart_truncate(obj, max_len)
)
assert is_truncated
assert truncated["a"] == "long strin...[TRUNCATED]"
assert truncated["b"][0] == "short"
assert truncated["b"][1] == "long strin...[TRUNCATED]"
assert truncated["c"]["d"] == "long strin...[TRUNCATED]"
def test_recursive_smart_truncate_with_dataclasses():
"""Test recursive smart truncate with dataclasses."""
@dataclasses.dataclass
class LocalMissedKPI:
kpi: str
value: float
@dataclasses.dataclass
class LocalIncident:
id: str
kpi_missed: list[LocalMissedKPI]
status: str
incident = LocalIncident(
id="inc-123",
kpi_missed=[LocalMissedKPI(kpi="latency", value=99.9)],
status="active",
)
content = {"result": incident}
max_len = 1000
truncated, is_truncated = (
bigquery_agent_analytics_plugin._recursive_smart_truncate(
content, max_len
)
)
assert not is_truncated
assert isinstance(truncated["result"], dict)
assert truncated["result"]["id"] == "inc-123"
assert isinstance(truncated["result"]["kpi_missed"][0], dict)
assert truncated["result"]["kpi_missed"][0]["kpi"] == "latency"
def test_recursive_smart_truncate_redaction():
"""Test that sensitive keys and temp: state keys are redacted."""
obj = {
"client_secret": "super-secret-123",
"access_token": "ya29.blah",
"refresh_token": "1//0g",
"id_token": "eyJhb",
"api_key": "AIza",
"password": "my-password",
"safe_key": "safe-value",
"temp:auth_state": "some-auth-state",
"nested": {
"CLIENT_SECRET": "nested-secret",
"normal": "value",
},
}
max_len = 1000
truncated, is_truncated = (
bigquery_agent_analytics_plugin._recursive_smart_truncate(obj, max_len)
)
assert not is_truncated
assert truncated["client_secret"] == "[REDACTED]"
assert truncated["access_token"] == "[REDACTED]"
assert truncated["refresh_token"] == "[REDACTED]"
assert truncated["id_token"] == "[REDACTED]"
assert truncated["api_key"] == "[REDACTED]"
assert truncated["password"] == "[REDACTED]"
assert truncated["safe_key"] == "safe-value"
assert truncated["temp:auth_state"] == "[REDACTED]"
assert truncated["nested"]["CLIENT_SECRET"] == "[REDACTED]"
assert truncated["nested"]["normal"] == "value"
class TestBigQueryAgentAnalyticsPlugin:
"""Tests for the BigQueryAgentAnalyticsPlugin."""
@pytest.mark.asyncio
async def test_plugin_disabled(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
invocation_context,
):
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(enabled=False)
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=config,
) as plugin:
# user_message = types.Content(parts=[types.Part(text="Test")])
await plugin.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Test")]),
)
mock_auth_default.assert_not_called()
mock_bq_client.assert_not_called()
@pytest.mark.asyncio
async def test_enriched_metadata_logging(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
dummy_arrow_schema,
callback_context,
):
# Setup
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
async with managed_plugin(PROJECT_ID, DATASET_ID, config=config) as plugin:
# Mock root agent
mock_root = mock.create_autospec(
base_agent.BaseAgent, instance=True, spec_set=True
)
type(mock_root).name = mock.PropertyMock(return_value="RootAgent")
callback_context._invocation_context.agent.root_agent = mock_root
# 1. Test root_agent_name and model extraction from request
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Hi")])],
)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
# 2. Test model_version and usage_metadata extraction from response
usage = types.GenerateContentResponseUsageMetadata(
prompt_token_count=10, candidates_token_count=20, total_token_count=30
)
llm_response = llm_response_lib.LlmResponse(
content=types.Content(parts=[types.Part(text="Hello")]),
usage_metadata=usage,
model_version="v1.2.3",
)
await plugin.after_model_callback(
callback_context=callback_context, llm_response=llm_response
)
# Verify captured rows from mock client
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
assert len(rows) == 2
# Check LLM_REQUEST row
# Sort by event_type to ensure consistent indexing
rows.sort(key=lambda x: x["event_type"])
request_row = rows[0] # LLM_REQUEST
response_row = rows[1] # LLM_RESPONSE
assert request_row["event_type"] == "LLM_REQUEST"
attr_req = json.loads(request_row["attributes"])
assert attr_req["root_agent_name"] == "RootAgent"
assert attr_req["model"] == "gemini-pro"
# Check LLM_RESPONSE row
assert response_row["event_type"] == "LLM_RESPONSE"
attr_res = json.loads(response_row["attributes"])
assert attr_res["root_agent_name"] == "RootAgent"
assert attr_res["model_version"] == "v1.2.3"
usage_meta = attr_res["usage_metadata"]
assert "prompt_token_count" in usage_meta
assert usage_meta["prompt_token_count"] == 10
mock_write_client.append_rows.assert_called()
@pytest.mark.asyncio
async def test_concurrent_span_management(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
callback_context,
):
# Setup
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
PROJECT_ID, DATASET_ID, config=config
)
# Initialize trace in main context
bigquery_agent_analytics_plugin.TraceManager.init_trace(callback_context)
async def branch_1():
s_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, span_name="span-1"
)
await asyncio.sleep(0.02)
current_s_id = (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
)
assert s_id == current_s_id
bigquery_agent_analytics_plugin.TraceManager.pop_span()
return s_id
async def branch_2():
s_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, span_name="span-2"
)
await asyncio.sleep(0.02)
current_s_id = (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
)
assert s_id == current_s_id
bigquery_agent_analytics_plugin.TraceManager.pop_span()
return s_id
# Run concurrently
results = await asyncio.gather(branch_1(), branch_2())
# If they shared the same list/dict, they would interfere.
assert results[0] is not None
assert results[1] is not None
assert results[0] != results[1]
@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,
mock_asyncio_to_thread,
):
_ = mock_auth_default
_ = mock_bq_client
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
event_allowlist=["LLM_REQUEST"]
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Prompt")])],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01) # Allow background task to run
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?")])
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await plugin.on_user_message_callback(
invocation_context=invocation_context, user_message=user_message
)
await asyncio.sleep(0.01) # Allow background task to run
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,
mock_asyncio_to_thread,
):
_ = mock_auth_default
_ = mock_bq_client
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
event_denylist=["USER_MESSAGE_RECEIVED"]
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
user_message = types.Content(parts=[types.Part(text="What is up?")])
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
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_not_called()
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await plugin.before_run_callback(invocation_context=invocation_context)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
@pytest.mark.asyncio
async def test_append_rows_sets_regional_routing_header(
self,
mock_write_client,
callback_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""Regression test for cross-region writes (issue #262).
The Storage Write API streaming AppendRows RPC does not
auto-populate the request-routing header, so writes to a dataset
outside the US multiregion (e.g. northamerica-northeast1) fail with
a "session not found" / stream-not-found error unless the header is
set explicitly. Assert the header is passed to append_rows so the
request reaches the region that owns the write stream.
"""
_ = mock_auth_default
_ = mock_bq_client
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
async with managed_plugin(
PROJECT_ID,
DATASET_ID,
table_id=TABLE_ID,
config=config,
location="northamerica-northeast1",
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Prompt")])],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01) # Allow background task to run
mock_write_client.append_rows.assert_called_once()
metadata = mock_write_client.append_rows.call_args.kwargs.get("metadata")
assert metadata is not None, "append_rows must receive routing metadata"
assert (
"x-goog-request-params",
f"write_stream={DEFAULT_STREAM_NAME}",
) in tuple(metadata)
@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,
mock_asyncio_to_thread,
):
"""Test content formatter."""
_ = mock_auth_default
_ = mock_bq_client
def redact_content(content, event_type):
return "[REDACTED]"
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
content_formatter=redact_content
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
user_message = types.Content(parts=[types.Part(text="Secret message")])
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
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 = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
# If the formatter returns a string, it's stored directly.
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,
mock_asyncio_to_thread,
):
"""Test content formatter error handling."""
_ = mock_auth_default
_ = mock_bq_client
def error_formatter(content, event_type):
raise ValueError("Formatter failed")
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
content_formatter=error_formatter
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
user_message = types.Content(parts=[types.Part(text="Secret message")])
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
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 = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
# If formatter fails, it logs a warning and continues with original content.
assert log_entry["content"] == '{"text_summary": "Secret message"}'
@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,
):
_ = mock_auth_default
_ = mock_bq_client
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
max_content_length=40
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
# Test User Message Truncation
user_message = types.Content(
parts=[types.Part(text="12345678901234567890123456789012345678901")]
) # 41 chars
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
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 = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
assert (
log_entry["content"]
== '{"text_summary":'
' "1234567890123456789012345678901234567890...[TRUNCATED]"}'
)
assert log_entry["is_truncated"]
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")])
],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
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 = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
# Full content: {"prompt": "text: 'Prompt'",
# "system_prompt": "text: 'System Instruction'"}
# In our new logic, we don't truncate the whole JSON string if it's valid JSON.
# Instead, we should have truncated the values within the dict, but currently we don't.
# For now, update test to reflect current behavior (valid JSON, no truncation of the whole string).
assert log_entry["content"].startswith(
'{"prompt": [{"role": "user", "content": "Prompt"}]'
)
assert log_entry["is_truncated"] is False
@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,
):
_ = mock_auth_default
_ = mock_bq_client
_ = mock_to_arrow_schema
_ = mock_asyncio_to_thread
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
max_content_length=80
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
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 > 80
# {"param": "A" * 100} is > 100 chars.
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await plugin.before_tool_callback(
tool=mock_tool,
tool_args={"param": "A" * 100},
tool_context=tool_context,
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "TOOL_STARTING")
# Now we do truncate nested values, and is_truncated flag is True
assert log_entry["is_truncated"]
content_dict = json.loads(log_entry["content"])
assert content_dict["tool"] == "MyTool"
assert content_dict["args"]["param"].endswith("...[TRUNCATED]")
@pytest.mark.asyncio
async def test_max_content_length_tool_args_no_truncation(
self,
mock_write_client,
tool_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
max_content_length=-1
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
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 > 80
# {"param": "A" * 100} is > 100 chars.
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await plugin.before_tool_callback(
tool=mock_tool,
tool_args={"param": "A" * 100},
tool_context=tool_context,
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "TOOL_STARTING")
# No truncation
assert not log_entry["is_truncated"]
content_dict = json.loads(log_entry["content"])
assert content_dict["tool"] == "MyTool"
assert content_dict["args"]["param"] == "A" * 100
@pytest.mark.asyncio
async def test_max_content_length_tool_result(
self,
mock_write_client,
tool_context,
mock_auth_default,
mock_bq_client,
mock_asyncio_to_thread,
mock_to_arrow_schema,
dummy_arrow_schema,
):
"""Test max content length for tool result."""
_ = mock_auth_default
_ = mock_bq_client
_ = mock_to_arrow_schema
_ = mock_asyncio_to_thread
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
max_content_length=80
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
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")
# Result length > 80
# {"res": "A" * 100} is > 100 chars.
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await plugin.after_tool_callback(
tool=mock_tool,
tool_args={},
tool_context=tool_context,
result={"res": "A" * 100},
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "TOOL_COMPLETED")
# Now we do truncate nested values, and is_truncated flag is True
assert log_entry["is_truncated"]
content_dict = json.loads(log_entry["content"])
assert content_dict["tool"] == "MyTool"
assert content_dict["result"]["res"].endswith("...[TRUNCATED]")
@pytest.mark.asyncio
async def test_max_content_length_tool_result_no_truncation(
self,
mock_write_client,
tool_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""Test max content length for tool result with no truncation."""
_ = mock_auth_default
_ = mock_bq_client
_ = mock_to_arrow_schema
_ = mock_asyncio_to_thread
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
max_content_length=-1
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
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")
# Result length > 80
# {"res": "A" * 100} is > 100 chars.
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await plugin.after_tool_callback(
tool=mock_tool,
tool_args={},
tool_context=tool_context,
result={"res": "A" * 100},
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "TOOL_COMPLETED")
# No truncation
assert not log_entry["is_truncated"]
content_dict = json.loads(log_entry["content"])
assert content_dict["tool"] == "MyTool"
assert content_dict["result"]["res"] == "A" * 100
@pytest.mark.asyncio
async def test_max_content_length_tool_error(
self,
mock_write_client,
tool_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
max_content_length=80
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
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")
# Args length > 80
# {"arg": "A" * 100} is > 100 chars.
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await plugin.on_tool_error_callback(
tool=mock_tool,
tool_args={"arg": "A" * 100},
tool_context=tool_context,
error=ValueError("Oops"),
)
await asyncio.sleep(0.01)
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["content"].startswith(
'{"tool": "MyTool", "args": {"arg": "AAAAA'
)
# Check for truncation in the nested value
content_dict = json.loads(log_entry["content"])
assert content_dict["args"]["arg"].endswith("...[TRUNCATED]")
assert log_entry["is_truncated"]
assert log_entry["error_message"] == "Oops"
@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?")])
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context, user_message=user_message
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "USER_MESSAGE_RECEIVED")
assert log_entry["content"] == '{"text_summary": "What is up?"}'
@pytest.mark.asyncio
async def test_offloading_with_connection_id(
self,
mock_write_client,
invocation_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
mock_storage_client,
):
_ = mock_auth_default
_ = mock_bq_client
_ = mock_to_arrow_schema
_ = mock_asyncio_to_thread
# Mock GCS bucket
mock_bucket = mock.Mock()
mock_blob = mock.Mock()
mock_bucket.blob.return_value = mock_blob
mock_bucket.name = "my-bucket"
mock_storage_client.return_value.bucket.return_value = mock_bucket
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
gcs_bucket_name="my-bucket",
connection_id="us.my-connection",
max_content_length=20, # Small limit to force offloading
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started(
storage_client=mock_storage_client.return_value
)
mock_write_client.append_rows.reset_mock()
# Create mixed content: one small inline, one large offloaded
small_text = "Small inline text"
large_text = "A" * 100
user_message = types.Content(
parts=[types.Part(text=small_text), types.Part(text=large_text)]
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
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 = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
# Verify content parts
assert len(log_entry["content_parts"]) == 2
# Part 0: Inline
part0 = log_entry["content_parts"][0]
assert part0["storage_mode"] == "INLINE"
assert part0["text"] == small_text
assert part0["object_ref"] is None
# Part 1: Offloaded
part1 = log_entry["content_parts"][1]
assert part1["storage_mode"] == "GCS_REFERENCE"
assert part1["uri"].startswith("gs://my-bucket/")
assert part1["object_ref"]["uri"] == part1["uri"]
assert part1["object_ref"]["authorizer"] == "us.my-connection"
assert json.loads(part1["object_ref"]["details"]) == {
"gcs_metadata": {"content_type": "text/plain"}
}
# Removed on_event_callback tests as they are no longer applicable in V2
@pytest.mark.asyncio
async def test_bigquery_client_initialization_failure(
self,
mock_auth_default,
mock_write_client,
invocation_context,
mock_asyncio_to_thread,
):
_ = mock_asyncio_to_thread
mock_auth_default.side_effect = auth_exceptions.GoogleAuthError(
"Auth failed"
)
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
) as plugin_with_fail:
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.logger"
) as mock_logger:
bigquery_agent_analytics_plugin.TraceManager.push_span(
invocation_context
)
await plugin_with_fail.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Test")]),
)
await asyncio.sleep(0.01)
mock_logger.error.assert_called_with(
"Failed to initialize BigQuery Plugin: %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
):
_ = bq_plugin_inst
async def fake_append_rows_with_error(requests, **kwargs):
mock_append_rows_response = mock.MagicMock()
mock_append_rows_response.row_errors = [] # No row errors
mock_append_rows_response.error = mock.MagicMock()
mock_append_rows_response.error.code = 3 # INVALID_ARGUMENT
mock_append_rows_response.error.message = "Test BQ Error"
return _async_gen(mock_append_rows_response)
mock_write_client.append_rows.side_effect = fake_append_rows_with_error
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.logger"
) as mock_logger:
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Test")]),
)
await asyncio.sleep(0.01)
# The logger is called multiple times, check that one of them is the error message
# Or just check that it was called with the expected message at some point
mock_logger.error.assert_any_call(
"Non-retryable BigQuery error: %s", "Test BQ Error"
)
mock_write_client.append_rows.assert_called_once()
@pytest.mark.asyncio
async def test_bigquery_insert_retryable_error(
self, bq_plugin_inst, mock_write_client, invocation_context
):
"""Test that retryable BigQuery errors are logged and retried."""
async def fake_append_rows_with_retryable_error(requests, **kwargs):
mock_append_rows_response = mock.MagicMock()
mock_append_rows_response.row_errors = [] # No row errors
mock_append_rows_response.error = mock.MagicMock()
mock_append_rows_response.error.code = 10 # ABORTED (retryable)
mock_append_rows_response.error.message = "Test BQ Retryable Error"
return _async_gen(mock_append_rows_response)
mock_write_client.append_rows.side_effect = (
fake_append_rows_with_retryable_error
)
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.logger"
) as mock_logger:
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Test")]),
)
await asyncio.sleep(0.01)
mock_logger.warning.assert_any_call(
"BigQuery Write API returned error code %s: %s",
10,
"Test BQ Retryable Error",
)
# Should be called at least once. Retries are hard to test due to async backoff.
assert mock_write_client.append_rows.call_count >= 1
@pytest.mark.asyncio
async def test_schema_mismatch_error_handling(
self, bq_plugin_inst, mock_write_client, invocation_context
):
async def fake_append_rows_with_schema_error(requests, **kwargs):
mock_resp = mock.MagicMock()
mock_resp.row_errors = []
mock_resp.error = mock.MagicMock()
mock_resp.error.code = 3
mock_resp.error.message = (
"Schema mismatch: Field 'new_field' not found in table."
)
return _async_gen(mock_resp)
mock_write_client.append_rows.side_effect = (
fake_append_rows_with_schema_error
)
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.logger"
) as mock_logger:
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Test")]),
)
await asyncio.sleep(0.01)
mock_logger.error.assert_called_with(
"BigQuery Schema Mismatch: %s. This usually means the"
" table schema does not match the expected schema.",
"Schema mismatch: Field 'new_field' not found in table.",
)
@pytest.mark.asyncio
async def test_close(self, bq_plugin_inst, mock_bq_client, mock_write_client):
"""Test plugin shutdown."""
await bq_plugin_inst.shutdown()
# shutdown calls transport.close() on all clients
assert mock_write_client.transport.close.call_count >= 1
# Verify loop states are cleared
assert not bq_plugin_inst._loop_state_by_loop
@pytest.mark.asyncio
async def test_before_run_callback_logs_correctly(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""Test before_run_callback logs correctly."""
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
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,
):
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.after_run_callback(
invocation_context=invocation_context
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
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,
):
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "AGENT_STARTING")
assert log_entry["content"] == "Test Instruction"
@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,
):
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.after_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "AGENT_COMPLETED")
assert log_entry["content"] is None
# Latency should be an int >= 0 now that we instrument it
assert log_entry["latency_ms"] is not None
latency_dict = json.loads(log_entry["latency_ms"])
assert latency_dict["total_ms"] >= 0
@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(role="user", parts=[types.Part(text="Prompt")])
],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "LLM_REQUEST")
assert "Prompt" in log_entry["content"]
@pytest.mark.asyncio
async def test_before_model_callback_with_params_and_tools(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
config=types.GenerateContentConfig(
temperature=0.5,
top_p=0.9,
system_instruction=types.Content(parts=[types.Part(text="Sys")]),
),
contents=[types.Content(role="user", parts=[types.Part(text="User")])],
)
# Manually set tools_dict as it is excluded from init
llm_request.tools_dict = {"tool1": "func1", "tool2": "func2"}
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "LLM_REQUEST")
# Verify content is JSON and has correct fields
assert "content" in log_entry
content_dict = json.loads(log_entry["content"])
assert content_dict["prompt"] == [{"role": "user", "content": "User"}]
assert content_dict["system_prompt"] == "Sys"
# Verify attributes
assert "attributes" in log_entry
attributes = json.loads(log_entry["attributes"])
assert attributes["llm_config"]["temperature"] == 0.5
assert attributes["llm_config"]["top_p"] == 0.9
assert attributes["llm_config"]["top_p"] == 0.9
# Tools without a name/description/declaration fall back to just the key.
assert attributes["tools"] == [{"name": "tool1"}, {"name": "tool2"}]
@pytest.mark.asyncio
async def test_before_model_callback_logs_tool_declarations(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""LLM_REQUEST tools carry name, description, and parameter schema."""
class _FakeTool(base_tool_lib.BaseTool):
def __init__(self, name, description, declaration):
super().__init__(name=name, description=description)
self._declaration = declaration
def _get_declaration(self):
return self._declaration
execute_sql = _FakeTool(
name="execute_sql",
description="Run a SQL query against BigQuery.",
declaration=types.FunctionDeclaration(
name="execute_sql",
description="Run a SQL query against BigQuery.",
parameters=types.Schema(
type=types.Type.OBJECT,
properties={
"query": types.Schema(
type=types.Type.STRING,
description="The SQL query to run.",
)
},
required=["query"],
),
),
)
# A tool without a declaration still contributes name + description.
list_datasets = _FakeTool(
name="list_dataset_ids",
description="List available datasets.",
declaration=None,
)
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(role="user", parts=[types.Part(text="hi")])],
)
llm_request.tools_dict = {
"execute_sql": execute_sql,
"list_dataset_ids": list_datasets,
}
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "LLM_REQUEST")
attributes = json.loads(log_entry["attributes"])
tools_by_name = {t["name"]: t for t in attributes["tools"]}
assert tools_by_name["execute_sql"]["description"] == (
"Run a SQL query against BigQuery."
)
params = tools_by_name["execute_sql"]["parameters"]
assert params["type"] == "OBJECT"
assert params["properties"]["query"]["type"] == "STRING"
assert params["required"] == ["query"]
assert tools_by_name["list_dataset_ids"]["description"] == (
"List available datasets."
)
assert "parameters" not in tools_by_name["list_dataset_ids"]
def test_extract_tool_declarations_declaration_error_is_isolated(self):
"""A tool whose _get_declaration raises still yields name + description."""
class _RaisingTool(base_tool_lib.BaseTool):
def _get_declaration(self):
raise ValueError("boom")
class _OkTool(base_tool_lib.BaseTool):
def _get_declaration(self):
return None
result = bigquery_agent_analytics_plugin._extract_tool_declarations({
"raiser": _RaisingTool(name="raiser", description="Raises."),
"ok": _OkTool(name="ok", description="Fine."),
})
by_name = {t["name"]: t for t in result}
# The raising tool is not dropped; other tools are unaffected.
assert by_name["raiser"] == {"name": "raiser", "description": "Raises."}
assert by_name["ok"] == {"name": "ok", "description": "Fine."}
def test_extract_tool_declarations_parameters_serialization_error(self):
"""A parameters object that fails to serialize is dropped, not fatal."""
class _BadParams:
def model_dump(self, *args, **kwargs):
raise ValueError("cannot serialize")
class _BadDecl:
description = None
parameters = _BadParams()
class _BadParamTool(base_tool_lib.BaseTool):
def _get_declaration(self):
return _BadDecl()
result = bigquery_agent_analytics_plugin._extract_tool_declarations(
{"bad_params": _BadParamTool(name="bad_params", description="Bad.")}
)
# Name + description survive; the unserializable parameters key is omitted.
assert result == [{"name": "bad_params", "description": "Bad."}]
def test_extract_tool_declarations_uses_parameters_json_schema(self):
"""Declarations exposing parameters_json_schema log that raw schema."""
json_schema = {
"type": "object",
"properties": {"path": {"type": "string"}},
"required": ["path"],
}
class _JsonSchemaTool(base_tool_lib.BaseTool):
def _get_declaration(self):
return types.FunctionDeclaration(
name="read_file",
description="Read a file.",
parameters_json_schema=json_schema,
)
result = bigquery_agent_analytics_plugin._extract_tool_declarations(
{"read_file": _JsonSchemaTool(name="read_file", description="Read.")}
)
# parameters_json_schema is logged verbatim (preferred over `parameters`).
assert result == [{
"name": "read_file",
"description": "Read.",
"parameters": json_schema,
}]
@pytest.mark.asyncio
async def test_before_model_callback_with_full_config(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Test that all config fields, including falsy values and labels, are logged."""
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
config=types.GenerateContentConfig(
temperature=0.0,
top_p=0.1,
top_k=5.0,
candidate_count=5,
max_output_tokens=65000,
stop_sequences=["STOP"],
presence_penalty=0.1,
frequency_penalty=0.5,
seed=42,
response_logprobs=True,
logprobs=3,
labels={"llm.agent.name": "test_agent"},
),
contents=[types.Content(role="user", parts=[types.Part(text="User")])],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "LLM_REQUEST")
# Verify attributes
assert "attributes" in log_entry
attributes = json.loads(log_entry["attributes"])
llm_config = attributes.get("llm_config", {})
expected_llm_config = {
"temperature": 0.0,
"top_p": 0.1,
"top_k": 5.0,
"candidate_count": 5,
"max_output_tokens": 65000,
"stop_sequences": ["STOP"],
"presence_penalty": 0.1,
"frequency_penalty": 0.5,
"seed": 42,
"response_logprobs": True,
"logprobs": 3,
}
assert llm_config == expected_llm_config
assert attributes.get("labels") == {"llm.agent.name": "test_agent"}
@pytest.mark.asyncio
async def test_before_model_callback_multipart_separator(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[
types.Content(
role="user",
parts=[types.Part(text="Part1"), types.Part(text="Part2")],
)
],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
content_dict = json.loads(log_entry["content"])
# Verify the separator is " | "
assert content_dict["prompt"][0]["content"] == "Part1 | Part2"
@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
),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.after_model_callback(
callback_context=callback_context,
llm_response=llm_response,
# latency_ms is now calculated internally via TraceManager
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "LLM_RESPONSE")
content_dict = json.loads(log_entry["content"])
assert content_dict["response"] == "text: 'Model response'"
assert content_dict["usage"]["prompt"] == 10
assert content_dict["usage"]["total"] == 15
assert log_entry["error_message"] is None
latency_dict = json.loads(log_entry["latency_ms"])
# Latency comes from time.time(), so we can't assert exact 100ms
# But it should be present
assert latency_dict["total_ms"] >= 0
# tfft is passed via kwargs if present, or we can mock it.
# In this test we didn't pass it in kwargs in the updated call above, so it might be missing unless we add it back to kwargs.
# The original test passed it as kwarg.
@pytest.mark.asyncio
async def test_after_model_callback_tool_call(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
tool_fc = types.FunctionCall(name="get_weather", args={"location": "Paris"})
llm_response = llm_response_lib.LlmResponse(
content=types.Content(parts=[types.Part(function_call=tool_fc)]),
usage_metadata=types.UsageMetadata(
prompt_token_count=10, total_token_count=15
),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.after_model_callback(
callback_context=callback_context,
llm_response=llm_response,
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "LLM_RESPONSE")
content_dict = json.loads(log_entry["content"])
assert content_dict["response"] == "call: get_weather"
assert content_dict["usage"]["prompt"] == 10
assert content_dict["usage"]["total"] == 15
assert log_entry["error_message"] is None
@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="Description")
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await bq_plugin_inst.before_tool_callback(
tool=mock_tool, tool_args={"param": "value"}, tool_context=tool_context
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "TOOL_STARTING")
content_dict = json.loads(log_entry["content"])
assert content_dict["tool"] == "MyTool"
assert content_dict["args"] == {"param": "value"}
@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")
type(mock_tool).description = mock.PropertyMock(return_value="Description")
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await bq_plugin_inst.after_tool_callback(
tool=mock_tool,
tool_args={"arg1": "val1"},
tool_context=tool_context,
result={"res": "success"},
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "TOOL_COMPLETED")
content_dict = json.loads(log_entry["content"])
assert content_dict["tool"] == "MyTool"
assert content_dict["result"] == {"res": "success"}
@pytest.mark.asyncio
async def test_after_tool_callback_no_state_delta_logging(
self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema
):
"""State deltas are now logged via on_event_callback, not after_tool."""
mock_tool = mock.create_autospec(
base_tool_lib.BaseTool, instance=True, spec_set=True
)
type(mock_tool).name = mock.PropertyMock(return_value="StateTool")
type(mock_tool).description = mock.PropertyMock(return_value="Sets state")
# Simulate a tool modifying the state
tool_context.actions.state_delta["new_key"] = "new_value"
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await bq_plugin_inst.after_tool_callback(
tool=mock_tool,
tool_args={"arg1": "val1"},
tool_context=tool_context,
result={"res": "success"},
)
await asyncio.sleep(0.01)
# Only TOOL_COMPLETED should be logged; STATE_DELTA is handled
# by on_event_callback now.
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
assert len(rows) == 1
assert rows[0]["event_type"] == "TOOL_COMPLETED"
@pytest.mark.asyncio
async def test_on_event_callback_logs_state_delta(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""on_event_callback logs STATE_DELTA for events with state changes."""
state_delta = {"key": "value", "new_key": 123}
event = event_lib.Event(
author="test_agent",
actions=event_actions_lib.EventActions(state_delta=state_delta),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
result = await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
# Must return None to not modify the event
assert result is None
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "STATE_DELTA")
assert log_entry["content"] is None
attributes = json.loads(log_entry["attributes"])
assert attributes["state_delta"] == state_delta
@pytest.mark.asyncio
async def test_on_event_callback_ignores_empty_state_delta(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""on_event_callback should not log when state_delta is empty."""
event = event_lib.Event(
author="test_agent",
actions=event_actions_lib.EventActions(state_delta={}),
)
result = await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
assert result is None
# No events should have been logged
mock_write_client.append_rows.assert_not_called()
@pytest.mark.asyncio
async def test_log_event_with_session_metadata(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Test that session metadata is logged when enabled."""
# Setup session state with user metadata
session = callback_context._invocation_context.session
type(session).state = mock.PropertyMock(
return_value={"thread_id": "gchat-123", "customer_id": "cust-42"}
)
# Ensure config enabled (default is True)
bq_plugin_inst.config.log_session_metadata = True
await bq_plugin_inst._log_event(
"TEST_EVENT",
callback_context,
raw_content="test content",
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
attributes = json.loads(log_entry["attributes"])
meta = attributes["session_metadata"]
assert meta["session_id"] == session.id
assert meta["app_name"] == session.app_name
assert meta["user_id"] == session.user_id
assert meta["state"] == {
"thread_id": "gchat-123",
"customer_id": "cust-42",
}
@pytest.mark.asyncio
async def test_log_event_with_custom_tags(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Test that custom tags are logged."""
custom_tags = {"agent_role": "sales", "env": "prod"}
bq_plugin_inst.config.custom_tags = custom_tags
await bq_plugin_inst._log_event(
"TEST_EVENT",
callback_context,
raw_content="test content",
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
attributes = json.loads(log_entry["attributes"])
assert attributes["custom_tags"] == custom_tags
def test_resolve_agent_label_prefers_running_agent(self, callback_context):
"""agent present → agent.name, regardless of any source event."""
event = event_lib.Event(author="WorkflowNodeA")
label = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._resolve_agent_label(
callback_context, event
)
assert label == "MyTestAgent"
def test_resolve_agent_label_falls_back_to_event_author(
self, callback_context
):
"""No agent + source Event → Event.author (the emitting node)."""
callback_context._invocation_context.agent = None
event = event_lib.Event(author="WorkflowNodeA")
label = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._resolve_agent_label(
callback_context, event
)
assert label == "WorkflowNodeA"
def test_resolve_agent_label_null_for_callback_only_row(
self, callback_context
):
"""No agent and no source Event → None (SQL NULL)."""
callback_context._invocation_context.agent = None
label = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._resolve_agent_label(
callback_context, None
)
assert label is None
@pytest.mark.asyncio
async def test_log_event_survives_none_agent_with_event_author(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Regression for #6063: None agent falls back to source event author."""
# Workflow-driven invocations leave ``InvocationContext.agent`` as None.
# Reading ``callback_context.agent_name`` then raised ``AttributeError``,
# which ``@_safe_callback`` swallowed, silently dropping the BigQuery row.
# The row must now be written with the source Event's author as the label.
callback_context._invocation_context.agent = None
event = event_lib.Event(author="WorkflowNodeA")
await bq_plugin_inst._log_event(
"TEST_EVENT",
callback_context,
raw_content="test content",
event_data=bigquery_agent_analytics_plugin.EventData(
source_event=event
),
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
assert log_entry["event_type"] == "TEST_EVENT"
assert log_entry["agent"] == "WorkflowNodeA"
@pytest.mark.asyncio
async def test_log_event_survives_none_agent_without_source_event(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Regression for #6063: callback-only row with no agent writes null."""
callback_context._invocation_context.agent = None
await bq_plugin_inst._log_event(
"TEST_EVENT",
callback_context,
raw_content="test content",
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
assert log_entry["event_type"] == "TEST_EVENT"
assert log_entry["agent"] is None
@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")
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.on_model_error_callback(
callback_context=callback_context, llm_request=llm_request, error=error
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
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"
assert log_entry["status"] == "ERROR"
@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")
type(mock_tool).description = mock.PropertyMock(return_value="Description")
error = TimeoutError("Tool timed out")
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
await bq_plugin_inst.on_tool_error_callback(
tool=mock_tool,
tool_args={"param": "value"},
tool_context=tool_context,
error=error,
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "TOOL_ERROR")
content_dict = json.loads(log_entry["content"])
assert content_dict["tool"] == "MyTool"
assert content_dict["args"] == {"param": "value"}
assert log_entry["error_message"] == "Tool timed out"
assert log_entry["status"] == "ERROR"
@pytest.mark.asyncio
async def test_on_agent_error_callback_logs_correctly(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
mock_agent,
dummy_arrow_schema,
):
"""on_agent_error_callback emits AGENT_ERROR with traceback."""
error = RuntimeError("Agent crashed")
try:
raise error
except RuntimeError:
pass # populate __traceback__
pushed_span_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "agent"
)
await bq_plugin_inst.on_agent_error_callback(
agent=mock_agent,
callback_context=callback_context,
error=error,
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
log_entry = next(r for r in rows if r["event_type"] == "AGENT_ERROR")
assert log_entry["error_message"] == "Agent crashed"
assert log_entry["status"] == "ERROR"
# The agent span BQAA pushed is popped and attributed to the error row.
assert log_entry["span_id"] == pushed_span_id
content = json.loads(log_entry["content"])
assert "error_traceback" in content
assert "RuntimeError: Agent crashed" in content["error_traceback"]
@pytest.mark.asyncio
async def test_on_agent_error_does_not_pop_foreign_invocation_span(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
mock_agent,
dummy_arrow_schema,
):
"""on_agent_error must not pop a span BQAA did not push for this agent.
Simulates another plugin's before_agent_callback raising before BQAA's
own before_agent_callback ran: the stack holds only the invocation root.
The guarded pop must leave the invocation span in place so the
subsequent INVOCATION_ERROR keeps correct span/latency data.
"""
trace_manager = bigquery_agent_analytics_plugin.TraceManager
inv_span_id = trace_manager.push_span(callback_context, "invocation")
error = RuntimeError("other plugin's before_agent failed")
try:
raise error
except RuntimeError:
pass
await bq_plugin_inst.on_agent_error_callback(
agent=mock_agent,
callback_context=callback_context,
error=error,
)
await asyncio.sleep(0.05)
# The invocation root was NOT consumed by the agent-error pop.
assert trace_manager.get_current_span_id() == inv_span_id
# The AGENT_ERROR row is still emitted.
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
log_entry = next(r for r in rows if r["event_type"] == "AGENT_ERROR")
assert log_entry["error_message"] == "other plugin's before_agent failed"
@pytest.mark.asyncio
async def test_on_run_error_callback_logs_correctly(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""on_run_error_callback emits INVOCATION_ERROR with traceback."""
error = ValueError("Invocation failed")
try:
raise error
except ValueError:
pass
bigquery_agent_analytics_plugin.TraceManager.push_span(
invocation_context, "invocation"
)
await bq_plugin_inst.on_run_error_callback(
invocation_context=invocation_context,
error=error,
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
log_entry = next(r for r in rows if r["event_type"] == "INVOCATION_ERROR")
assert log_entry["error_message"] == "Invocation failed"
assert log_entry["status"] == "ERROR"
content = json.loads(log_entry["content"])
assert "error_traceback" in content
assert "ValueError: Invocation failed" in content["error_traceback"]
@pytest.mark.asyncio
async def test_on_run_error_callback_cleanup_runs_on_log_failure(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""on_run_error_callback cleans up even when _log_event raises."""
# Push spans and set context vars to simulate active invocation
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
bigquery_agent_analytics_plugin._active_invocation_id_ctx.set("test-inv")
bigquery_agent_analytics_plugin._root_agent_name_ctx.set("test-agent")
# Make _log_event raise
with mock.patch.object(
bq_plugin_inst, "_log_event", side_effect=RuntimeError("boom")
):
# @_safe_callback swallows the exception
await bq_plugin_inst.on_run_error_callback(
invocation_context=invocation_context,
error=ValueError("app error"),
)
# finally block must have cleaned up
assert (
bigquery_agent_analytics_plugin._active_invocation_id_ctx.get(None)
is None
)
assert (
bigquery_agent_analytics_plugin._root_agent_name_ctx.get(None) is None
)
@pytest.mark.asyncio
async def test_traceback_not_truncated_with_negative_max_len(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
invocation_context,
mock_agent,
dummy_arrow_schema,
):
"""Traceback is not truncated when max_content_length is -1."""
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
max_content_length=-1,
create_views=False,
)
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=config,
) as plugin:
await plugin._ensure_started()
error = RuntimeError("x" * 2000)
try:
raise error
except RuntimeError:
pass
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await plugin.on_agent_error_callback(
agent=mock_agent,
callback_context=bigquery_agent_analytics_plugin.CallbackContext(
invocation_context
),
error=error,
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
log_entry = next(r for r in rows if r["event_type"] == "AGENT_ERROR")
content = json.loads(log_entry["content"])
# Should NOT be truncated
assert "[truncated]" not in content["error_traceback"]
assert "x" * 2000 in content["error_traceback"]
@pytest.mark.asyncio
async def test_table_creation_options(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
mock_bq_client.get_table.side_effect = cloud_exceptions.NotFound(
"Not found"
)
await plugin._ensure_started()
# Verify create_table was called with correct table options
mock_bq_client.create_table.assert_called_once()
call_args = mock_bq_client.create_table.call_args
table_arg = call_args[0][0]
assert isinstance(table_arg, bigquery.Table)
assert table_arg.time_partitioning.type_ == "DAY"
assert table_arg.time_partitioning.field == "timestamp"
assert table_arg.clustering_fields == ["event_type", "agent", "user_id"]
# Verify schema descriptions are present (spot check)
timestamp_field = next(
f for f in table_arg.schema if f.name == "timestamp"
)
assert (
timestamp_field.description
== "The UTC timestamp when the event occurred. Used for ordering"
" events"
" within a session."
)
@pytest.mark.asyncio
async def test_init_in_thread_pool(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
invocation_context,
):
"""Verifies that the plugin can be initialized from a thread pool."""
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
def _run_in_thread(p):
# In a real thread pool, there might not be an event loop.
# However, since we are calling an async method (_ensure_started),
# we must run it in an event loop. The issue was that _lazy_setup
# called get_event_loop() which fails in threads without a loop.
# Here we simulate the condition by running in a thread and creating a new loop if needed,
# but the key is that the plugin's internal calls should use the correct loop.
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
# _ensure_started is called by managed_plugin, but we need to ensure
# that if it were called in a thread, it would work.
# For this test, we just ensure the plugin is accessible and started.
loop.run_until_complete(p._ensure_started())
finally:
loop.close()
# Run in a separate thread to simulate ThreadPoolExecutor-0_0
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(_run_in_thread, plugin)
future.result() # Should not raise "no current event loop"
assert plugin._started
# Verify loop states are populated
assert plugin._loop_state_by_loop
@pytest.mark.asyncio
async def test_multimodal_offloading(
self,
mock_write_client,
callback_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_storage_client,
):
# Setup
bucket_name = "test-bucket"
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
gcs_bucket_name=bucket_name
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started(
storage_client=mock_storage_client.return_value
)
# Mock GCS bucket and blob
mock_bucket = mock_storage_client.return_value.bucket.return_value
mock_bucket.name = bucket_name
mock_blob = mock_bucket.blob.return_value
# Create content with large text that should be offloaded
large_text = "A" * (32 * 1024 + 1)
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text=large_text)])],
)
# Execute
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
# Use flush instead of sleep for robustness
await plugin.flush()
# Verify GCS upload
mock_blob.upload_from_string.assert_called_once()
args, kwargs = mock_blob.upload_from_string.call_args
assert args[0] == large_text
assert kwargs["content_type"] == "text/plain"
# Verify BQ write
mock_write_client.append_rows.assert_called_once()
event_dict = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
content_parts = event_dict["content_parts"]
assert len(content_parts) == 1
assert content_parts[0]["storage_mode"] == "GCS_REFERENCE"
assert content_parts[0]["uri"].startswith(f"gs://{bucket_name}/")
@pytest.mark.asyncio
async def test_quota_project_id_used_in_client(
self,
mock_bq_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
mock_creds = mock.create_autospec(
google.auth.credentials.Credentials, instance=True, spec_set=True
)
mock_creds.quota_project_id = "quota-project"
with mock.patch.object(
google.auth,
"default",
autospec=True,
return_value=(mock_creds, PROJECT_ID),
) as mock_auth_default:
with mock.patch.object(
bigquery_agent_analytics_plugin,
"BigQueryWriteAsyncClient",
autospec=True,
) as mock_bq_write_cls:
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
) as plugin:
await plugin._ensure_started()
mock_auth_default.assert_called_once()
mock_bq_write_cls.assert_called_once()
_, kwargs = mock_bq_write_cls.call_args
assert kwargs["client_options"].quota_project_id == "quota-project"
@pytest.mark.asyncio
async def test_no_quota_project_when_creds_lack_it(
self,
mock_bq_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""Verify no quota_project_id is set when credentials don't provide one.
This is critical for Workload Identity Federation flows where setting
quota_project_id on the client breaks auth token refresh (issue #4370).
"""
mock_creds = mock.create_autospec(
google.auth.credentials.Credentials, instance=True, spec_set=True
)
mock_creds.quota_project_id = None
with mock.patch.object(
google.auth,
"default",
autospec=True,
return_value=(mock_creds, PROJECT_ID),
):
with mock.patch.object(
bigquery_agent_analytics_plugin,
"BigQueryWriteAsyncClient",
autospec=True,
) as mock_bq_write_cls:
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
) as plugin:
await plugin._ensure_started()
mock_bq_write_cls.assert_called_once()
_, kwargs = mock_bq_write_cls.call_args
assert kwargs["client_options"] is None
@pytest.mark.asyncio
async def test_custom_credentials_used(
self,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""Verify custom credentials are used and default auth is not called."""
mock_custom_creds = mock.create_autospec(
google.auth.credentials.Credentials, instance=True, spec_set=True
)
mock_custom_creds.quota_project_id = "custom-quota-project"
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
gcs_bucket_name="test-bucket",
create_views=False,
)
with mock.patch.object(
google.auth,
"default",
autospec=True,
) as mock_auth_default:
with mock.patch.object(
bigquery_agent_analytics_plugin,
"BigQueryWriteAsyncClient",
autospec=True,
) as mock_bq_write_cls:
with mock.patch(
"google.cloud.bigquery.Client", autospec=True
) as mock_bq_cls:
with mock.patch(
"google.cloud.storage.Client", autospec=True
) as mock_storage_cls:
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
credentials=mock_custom_creds,
config=config,
) as plugin:
await plugin._ensure_started()
mock_auth_default.assert_not_called()
mock_bq_write_cls.assert_called_once()
_, kwargs = mock_bq_write_cls.call_args
assert kwargs["credentials"] == mock_custom_creds
mock_bq_cls.assert_called_once()
_, kwargs = mock_bq_cls.call_args
assert kwargs["credentials"] == mock_custom_creds
mock_storage_cls.assert_called_once()
_, kwargs = mock_storage_cls.call_args
assert kwargs["credentials"] == mock_custom_creds
@pytest.mark.asyncio
async def test_pickle_safety(self, mock_auth_default, mock_bq_client):
"""Test that the plugin can be pickled safely."""
import pickle
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(enabled=True)
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
)
# Test pickling before start
pickled = pickle.dumps(plugin)
unpickled = pickle.loads(pickled)
assert unpickled.project_id == PROJECT_ID
assert unpickled._setup_lock is None
assert unpickled._executor is None
# Start the plugin
await plugin._ensure_started()
assert plugin._executor is not None
try:
# Test pickling after start
pickled_started = pickle.dumps(plugin)
unpickled_started = pickle.loads(pickled_started)
assert unpickled_started.project_id == PROJECT_ID
# Runtime objects should be None after unpickling
assert unpickled_started._setup_lock is None
assert unpickled_started._executor is None
assert not unpickled_started._loop_state_by_loop
finally:
await plugin.shutdown()
@pytest.mark.asyncio
async def test_span_hierarchy_llm_call(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Verifies that LLM events have correct Span ID hierarchy."""
# 1. Start Agent Span
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
_, _ = (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_and_parent()
)
agent_span_id = (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
)
# 2. Start LLM Span (Implicitly handled if we push it?
# Actually before_model_callback assumes a span is pushed for the LLM call if we want one?
# No, usually the Runner/Agent pushes a span BEFORE calling before_model_callback?
# Let's verify usage in agent.py or plugin.
# Plugin does NOT push spans automatically for LLM. It relies on TraceManager being managed externally
# OR it uses current span.
# Wait, the Runner pushes spans.
# 3. LLM Request
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Prompt")])],
)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
# Capture the actual LLM Span ID (pushed by before_model_callback)
llm_span_id = (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
)
# Now that we push a new span for LLM calls, it should differ from agent_span_id
assert llm_span_id != agent_span_id
log_entry_req = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
assert log_entry_req["event_type"] == "LLM_REQUEST"
assert log_entry_req["span_id"] == llm_span_id
# The parent of the LLM span should be the Agent span
assert log_entry_req["parent_span_id"] == agent_span_id
mock_write_client.append_rows.reset_mock()
# 4. LLM Response
# In the actual flow, after_model_callback pops the span.
# But explicitly via TraceManager.pop_span()?
# No, after_model_callback calls TraceManager.pop_span().
# So we should validly call it.
llm_response = llm_response_lib.LlmResponse(
content=types.Content(parts=[types.Part(text="Response")]),
)
await bq_plugin_inst.after_model_callback(
callback_context=callback_context, llm_response=llm_response
)
await asyncio.sleep(0.01)
log_entry_resp = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
assert log_entry_resp["event_type"] == "LLM_RESPONSE"
assert log_entry_resp["span_id"] == llm_span_id
# The parent of the LLM span should be the Agent span
assert log_entry_resp["parent_span_id"] == agent_span_id
# Verify LLM Span was popped and we are back to Agent Span
assert (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
== agent_span_id
)
# Clean up Agent Span
bigquery_agent_analytics_plugin.TraceManager.pop_span()
assert (
not bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
)
@pytest.mark.asyncio
async def test_custom_object_serialization(
self,
mock_write_client,
tool_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""Verifies that custom objects (Dataclasses) are serialized to dicts."""
_ = mock_auth_default
_ = mock_bq_client
@dataclasses.dataclass
class LocalMissedKPI:
kpi: str
value: float
@dataclasses.dataclass
class LocalIncident:
id: str
kpi_missed: list[LocalMissedKPI]
status: str
incident = LocalIncident(
id="inc-123",
kpi_missed=[LocalMissedKPI(kpi="latency", value=99.9)],
status="active",
)
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
content = {"result": incident}
# Verify full flow
await plugin._log_event(
"TOOL_PARTIAL",
tool_context,
raw_content=content,
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
# Content should be valid JSON string
content_json = json.loads(log_entry["content"])
assert content_json["result"]["id"] == "inc-123"
assert content_json["result"]["kpi_missed"][0]["kpi"] == "latency"
@pytest.mark.asyncio
async def test_push_pop_does_not_call_tracer_start_span(
self,
callback_context,
):
"""Regression guard for the duplicate-Cloud-Trace bug (issue #94).
The plugin must NOT call ``tracer.start_span(...)`` from
``push_span`` / ``pop_span``. Any owned OTel span goes through
the globally configured exporter (e.g. Cloud Trace via Agent
Engine telemetry) and surfaces as a duplicate span next to the
framework's real one. The plugin's internal stack is sufficient
for ``span_id`` / ``parent_span_id`` / ``trace_id`` resolution
without creating an exportable span.
"""
mock_tracer = mock.Mock()
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.tracer",
mock_tracer,
):
span_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "test_span"
)
assert isinstance(span_id, str) and len(span_id) == 16
trace_id = bigquery_agent_analytics_plugin.TraceManager.get_trace_id(
callback_context
)
assert isinstance(trace_id, str) and len(trace_id) == 32
popped_span_id, _duration_ms = (
bigquery_agent_analytics_plugin.TraceManager.pop_span()
)
assert popped_span_id == span_id
mock_tracer.start_span.assert_not_called()
@pytest.mark.asyncio
async def test_push_pop_does_not_export_spans_through_real_provider(
self, callback_context
):
"""End-to-end regression guard against #94 with a real OTel
provider + in-memory exporter.
Wires an ``InMemorySpanExporter`` to a real ``TracerProvider``,
drives a push/pop cycle through ``TraceManager``, and asserts
that **zero** spans were exported. Pre-fix behavior was to
export one span per push/pop pair — visible to Cloud Trace as
duplicate spans alongside the framework's real ones.
"""
# pylint: disable=g-import-not-at-top
from opentelemetry.sdk import trace as trace_sdk
from opentelemetry.sdk.trace import export as trace_export
from opentelemetry.sdk.trace.export import in_memory_span_exporter
# pylint: enable=g-import-not-at-top
provider = trace_sdk.TracerProvider()
exporter = in_memory_span_exporter.InMemorySpanExporter()
provider.add_span_processor(trace_export.SimpleSpanProcessor(exporter))
real_tracer = provider.get_tracer("test_tracer")
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.tracer",
real_tracer,
):
span_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "test_span"
)
assert exporter.get_finished_spans() == ()
trace_id = bigquery_agent_analytics_plugin.TraceManager.get_trace_id(
callback_context
)
assert trace_id is not None and len(trace_id) == 32
popped_span_id, _ = (
bigquery_agent_analytics_plugin.TraceManager.pop_span()
)
assert popped_span_id == span_id
assert exporter.get_finished_spans() == (), (
"Plugin must not export OTel spans; any owned span would"
" surface as a duplicate in Cloud Trace alongside the"
" framework's real spans (issue #94)."
)
provider.shutdown()
@pytest.mark.asyncio
async def test_push_span_inherits_ambient_trace_id(self, callback_context):
"""When the host has an ambient OTel span (e.g.
Agent Engine's Runner span), the plugin's ``trace_id`` MUST inherit from it
so BigQuery rows correlate with the host's Cloud Trace entries via a shared
``trace_id``.
"""
# pylint: disable=g-import-not-at-top
from opentelemetry import trace as otel_trace
from opentelemetry.sdk import trace as trace_sdk
# pylint: enable=g-import-not-at-top
provider = trace_sdk.TracerProvider()
host_tracer = provider.get_tracer("host_tracer")
# Clear any state on the plugin's contextvar stack.
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
with host_tracer.start_as_current_span("ambient-host-span") as host_span:
expected_trace_id = format(host_span.get_span_context().trace_id, "032x")
# Plugin pushes its first internal span inside the ambient span.
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "bqaa-span"
)
plugin_trace_id = (
bigquery_agent_analytics_plugin.TraceManager.get_trace_id(
callback_context
)
)
assert plugin_trace_id == expected_trace_id, (
"Plugin must inherit ambient trace_id so BigQuery rows join"
" to Cloud Trace via the same trace_id"
)
# Nested plugin push also stays under the ambient trace_id.
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "bqaa-nested"
)
assert (
bigquery_agent_analytics_plugin.TraceManager.get_trace_id(
callback_context
)
== expected_trace_id
)
bigquery_agent_analytics_plugin.TraceManager.clear_stack()
provider.shutdown()
del otel_trace # unused; imported for symmetry with provider setup
@pytest.mark.asyncio
async def test_llm_request_response_share_span_id_contract(
self, callback_context
):
"""Lifecycle contract: ``LLM_REQUEST`` and ``LLM_RESPONSE`` for the
same model call share one ``span_id`` and one ``trace_id``.
Models the structural pattern the real callbacks use:
* ``before_model_callback`` calls ``push_span(...)`` and writes
``LLM_REQUEST`` with the returned ``span_id``.
* ``after_model_callback`` calls ``get_current_span_id()`` /
``pop_span()`` and writes ``LLM_RESPONSE`` with the same
``span_id``.
A future change must not split this pair onto two different
``span_id``s — that would break the documented BigQuery query
shape and the BQAA join contract.
"""
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
TM = bigquery_agent_analytics_plugin.TraceManager
# before_model_callback path.
pushed_span_id = TM.push_span(callback_context, "llm_request")
request_trace_id = TM.get_trace_id(callback_context)
# after_model_callback (final chunk) path.
response_top_of_stack = TM.get_current_span_id()
popped_span_id, _duration_ms = TM.pop_span()
response_trace_id = TM.get_trace_id(callback_context)
assert response_top_of_stack == pushed_span_id
assert popped_span_id == pushed_span_id
# trace_id resolved on the response side may have to fall back
# past the now-empty stack — but if it does resolve, it must
# match what the request observed. An empty-stack fallback to
# invocation_id is acceptable here; what we are guarding against
# is the *pair* drifting onto two structurally different ids.
if response_trace_id is not None and len(response_trace_id) == 32:
assert response_trace_id == request_trace_id
@pytest.mark.asyncio
async def test_tool_starting_completed_share_span_id_contract(
self, callback_context
):
"""Lifecycle contract: ``TOOL_STARTING`` and ``TOOL_COMPLETED`` for
the same tool call share one ``span_id``.
Same shape as the LLM pair above — push on before, pop on after,
same id on both sides.
"""
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
TM = bigquery_agent_analytics_plugin.TraceManager
# before_tool_callback path.
pushed_span_id = TM.push_span(callback_context, "tool")
starting_trace_id = TM.get_trace_id(callback_context)
# after_tool_callback path.
popped_span_id, _duration_ms = TM.pop_span()
assert popped_span_id == pushed_span_id
assert isinstance(starting_trace_id, str) and len(starting_trace_id) == 32
@pytest.mark.asyncio
async def test_streaming_llm_response_shares_span_id_until_final_contract(
self, callback_context
):
"""Streaming-response contract.
On a streaming LLM call, ``after_model_callback`` is fired once
per partial chunk *plus* once for the final chunk. Partial fires
do NOT pop the span (see ``after_model_callback:3354-3363``) —
they only read ``get_current_span_id()`` and record first-token
timing. Only the final fire calls ``pop_span()``.
All resulting ``LLM_RESPONSE`` rows therefore share one
``span_id`` (the same as the paired ``LLM_REQUEST``). A future
change must not "dedupe" the partial rows by switching to a fresh
span id per chunk — those rows are real and intentional.
"""
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
TM = bigquery_agent_analytics_plugin.TraceManager
pushed_span_id = TM.push_span(callback_context, "llm_request")
# Simulate three partial chunks: each callback observes the same
# span_id at top of stack and does NOT pop.
for _ in range(3):
assert TM.get_current_span_id() == pushed_span_id
# Final chunk: pop_span returns the same id and a populated
# latency.
popped_span_id, duration_ms = TM.pop_span()
assert popped_span_id == pushed_span_id
assert duration_ms is not None and duration_ms >= 0
# Stack must be empty after the final chunk.
assert TM.get_current_span_id() is None
@pytest.mark.asyncio
async def test_keyword_identifiers_emission_default(
self,
mock_auth_default,
mock_bq_client,
callback_context,
):
"""Verify the default keyword flow for User-Agent and Trace-ID."""
keyword = "google-adk-bq-logger"
mock_write_client = mock.AsyncMock()
# 1. Verify User-Agent contains default keyword.
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.BigQueryWriteAsyncClient",
autospec=True,
) as mock_write_cls:
mock_write_cls.return_value = mock_write_client
async with managed_plugin(PROJECT_ID, DATASET_ID) as plugin:
await plugin._ensure_started()
_, kwargs = mock_write_cls.call_args
client_info = kwargs.get("client_info")
assert f"{keyword}/{__version__}" in client_info.user_agent
# 2. Verify Trace ID contains default keyword.
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.BigQueryWriteAsyncClient",
autospec=True,
) as mock_write_cls:
mock_write_cls.return_value = mock_write_client
async with managed_plugin(PROJECT_ID, DATASET_ID) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Hi")])],
)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await plugin.flush()
call_args = mock_write_client.append_rows.call_args
requests_iter = call_args.args[0]
requests = []
async for req in requests_iter:
requests.append(req)
assert requests[0].trace_id.startswith(keyword)
assert requests[0].trace_id.endswith(f"/{__version__}")
@pytest.mark.asyncio
async def test_visual_builder_identifiers_flow(
self,
mock_auth_default,
mock_bq_client,
callback_context,
dummy_arrow_schema,
):
"""Verify visual-builder keyword flow via contextvars."""
keyword = "google-adk-visual-builder"
mock_write_client = mock.AsyncMock()
# Simulate setting the internal flag via contextvars
token = _is_visual_builder.set(True)
try:
# 1. Verify Client User-Agent
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.BigQueryWriteAsyncClient",
autospec=True,
) as mock_write_cls:
mock_write_cls.return_value = mock_write_client
async with managed_plugin(PROJECT_ID, DATASET_ID) as plugin:
await plugin._ensure_started()
_, kwargs = mock_write_cls.call_args
client_info = kwargs.get("client_info")
assert keyword in client_info.user_agent
# 2. Verify Request Trace ID
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.BigQueryWriteAsyncClient",
autospec=True,
) as mock_write_cls:
mock_write_cls.return_value = mock_write_client
async with managed_plugin(PROJECT_ID, DATASET_ID) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Hi")])],
)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await plugin.flush()
call_args = mock_write_client.append_rows.call_args
requests_iter = call_args.args[0]
requests = []
async for req in requests_iter:
requests.append(req)
assert requests[0].trace_id.startswith(
"google-adk-bq-logger-visual-builder"
)
assert requests[0].trace_id.endswith(f"/{__version__}")
finally:
_is_visual_builder.reset(token)
@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"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"gen_config_kwargs, expected_llm_config",
[
(
{
"temperature": 0.0,
"top_k": 5.0,
"top_p": 0.1,
"candidate_count": 5,
"max_output_tokens": 65000,
"presence_penalty": 0.1,
"frequency_penalty": 0.5,
"response_logprobs": True,
"logprobs": 3,
"seed": 42,
"labels": {"llm.agent.name": "test_agent"},
},
{
"temperature": 0.0,
"top_k": 5.0,
"top_p": 0.1,
"candidate_count": 5,
"max_output_tokens": 65000,
"presence_penalty": 0.1,
"frequency_penalty": 0.5,
"response_logprobs": True,
"logprobs": 3,
"seed": 42,
},
),
],
)
async def test_generation_config_logging(
self,
bq_plugin_inst,
mock_write_client,
dummy_arrow_schema,
callback_context,
gen_config_kwargs,
expected_llm_config,
):
"""Verifies that all fields in GenerateContentConfig are logged correctly."""
gen_config = types.GenerateContentConfig(**gen_config_kwargs)
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Prompt")])],
config=gen_config,
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
# Flush
await bq_plugin_inst.flush()
# Verify
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
assert log_entry["event_type"] == "LLM_REQUEST"
attributes = json.loads(log_entry["attributes"])
llm_config = attributes.get("llm_config", {})
assert llm_config == expected_llm_config
if "labels" in gen_config_kwargs:
assert attributes.get("labels") == gen_config_kwargs["labels"]
class TestSafeCallbackDecorator:
"""Tests that _safe_callback prevents plugin errors from propagating."""
@pytest.mark.asyncio
async def test_callback_exception_does_not_propagate(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""A callback that throws should return None, not crash."""
# Force _log_event to raise
with mock.patch.object(
bq_plugin_inst,
"_log_event",
side_effect=RuntimeError("BQ network timeout"),
):
# Should NOT raise
result = await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Test")]),
)
assert result is None
@pytest.mark.asyncio
async def test_callback_exception_is_logged(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""The swallowed exception should be logged with exc_info."""
with mock.patch.object(
bq_plugin_inst,
"_log_event",
side_effect=RuntimeError("BQ write failed"),
):
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.logger"
) as mock_logger:
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context,
)
mock_logger.exception.assert_called_once_with(
"BigQuery analytics plugin error in %s; skipping.",
"before_run_callback",
)
@pytest.mark.asyncio
async def test_subsequent_callbacks_work_after_failure(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""After one callback fails, the next one should still work."""
call_count = 0
original_log_event = bq_plugin_inst._log_event
async def fail_once(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
raise RuntimeError("Transient error")
return await original_log_event(*args, **kwargs)
with mock.patch.object(bq_plugin_inst, "_log_event", side_effect=fail_once):
# First call fails silently
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Fail")]),
)
# Second call succeeds
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context,
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
@pytest.mark.asyncio
async def test_on_event_callback_exception_returns_none(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""on_event_callback should return None on error, not crash."""
event = event_lib.Event(
author="test_agent",
actions=event_actions_lib.EventActions(state_delta={"key": "value"}),
)
with mock.patch.object(
bq_plugin_inst,
"_log_event",
side_effect=Exception("serialize error"),
):
result = await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
assert result is None
@pytest.mark.asyncio
async def test_tool_callback_exception_does_not_propagate(
self,
bq_plugin_inst,
mock_write_client,
tool_context,
):
"""Tool callbacks should not crash even if plugin errors."""
mock_tool = mock.create_autospec(
base_tool_lib.BaseTool, instance=True, spec_set=True
)
type(mock_tool).name = mock.PropertyMock(return_value="MyTool")
with mock.patch.object(
bq_plugin_inst,
"_log_event",
side_effect=RuntimeError("BQ down"),
):
# before_tool_callback
result = await bq_plugin_inst.before_tool_callback(
tool=mock_tool,
tool_args={"p": "v"},
tool_context=tool_context,
)
assert result is None
# after_tool_callback
result = await bq_plugin_inst.after_tool_callback(
tool=mock_tool,
tool_args={"p": "v"},
tool_context=tool_context,
result={"r": "ok"},
)
assert result is None
# on_tool_error_callback
result = await bq_plugin_inst.on_tool_error_callback(
tool=mock_tool,
tool_args={"p": "v"},
tool_context=tool_context,
error=ValueError("tool broke"),
)
assert result is None
@pytest.mark.asyncio
async def test_model_callback_exception_does_not_propagate(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
):
"""Model callbacks should not crash even if plugin errors."""
with mock.patch.object(
bq_plugin_inst,
"_log_event",
side_effect=RuntimeError("BQ down"),
):
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Hi")])],
)
result = await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
assert result is None
llm_response = llm_response_lib.LlmResponse(
content=types.Content(parts=[types.Part(text="Hi")]),
)
result = await bq_plugin_inst.after_model_callback(
callback_context=callback_context, llm_response=llm_response
)
assert result is None
result = await bq_plugin_inst.on_model_error_callback(
callback_context=callback_context,
llm_request=llm_request_lib.LlmRequest(model="gemini-pro"),
error=ValueError("llm error"),
)
assert result is None
class TestParserReuse:
"""Tests that HybridContentParser is reused, not recreated per event."""
@pytest.mark.asyncio
async def test_parser_instance_is_reused(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""The same parser instance should be reused across _log_event calls."""
parser_after_init = bq_plugin_inst.parser
assert parser_after_init is not None
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Hello")]),
)
await asyncio.sleep(0.01)
# Parser should be the same instance, not a new one
assert bq_plugin_inst.parser is parser_after_init
@pytest.mark.asyncio
async def test_parser_trace_id_updated_per_call(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""trace_id and span_id on the parser should update per _log_event."""
parser = bq_plugin_inst.parser
original_trace_id = parser.trace_id
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Test")]),
)
await asyncio.sleep(0.01)
# After logging, trace_id/span_id should have been updated
# (they're derived from TraceManager, not the initial empty strings)
assert parser.span_id != ""
@pytest.mark.asyncio
async def test_parser_not_recreated_with_constructor(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""HybridContentParser constructor should not be called in
_log_event."""
with mock.patch.object(
bigquery_agent_analytics_plugin,
"HybridContentParser",
wraps=bigquery_agent_analytics_plugin.HybridContentParser,
) as mock_parser_cls:
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(parts=[types.Part(text="Test")]),
)
await asyncio.sleep(0.01)
# Constructor should NOT have been called during _log_event
mock_parser_cls.assert_not_called()
class TestPropertyAccessors:
"""Tests that properties work correctly after __getattribute__ removal."""
@pytest.mark.asyncio
async def testbatch_processorerty_returns_processor(self, bq_plugin_inst):
"""batch_processor property should return the processor for the
current loop."""
bp = bq_plugin_inst.batch_processor
assert bp is not None
assert isinstance(bp, bigquery_agent_analytics_plugin.BatchProcessor)
@pytest.mark.asyncio
async def test_write_client_property_returns_client(self, bq_plugin_inst):
"""write_client property should return the client for the current
loop."""
wc = bq_plugin_inst.write_client
assert wc is not None
@pytest.mark.asyncio
async def test_write_stream_property_returns_stream(self, bq_plugin_inst):
"""write_stream property should return the stream name."""
ws = bq_plugin_inst.write_stream
assert ws is not None
assert ws == DEFAULT_STREAM_NAME
@pytest.mark.asyncio
async def test_properties_return_none_when_no_loop_state(self):
"""Properties should return None when no state exists for the
current loop."""
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
assert plugin.batch_processor is None
assert plugin.write_client is None
assert plugin.write_stream is None
@pytest.mark.asyncio
async def test_regular_attributes_still_accessible(self, bq_plugin_inst):
"""Regular instance attributes should still be accessible."""
assert bq_plugin_inst.project_id == PROJECT_ID
assert bq_plugin_inst.dataset_id == DATASET_ID
assert bq_plugin_inst.table_id == TABLE_ID
assert bq_plugin_inst.config is not None
assert bq_plugin_inst._started is True
def test_properties_without_running_loop(self):
"""Properties should return None when no event loop is running."""
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
# No running loop → should return None, not crash
assert plugin.batch_processor is None
assert plugin.write_client is None
assert plugin.write_stream is None
class TestUnifiedSpanRecords:
"""Tests for the unified _SpanRecord-based TraceManager."""
@pytest.mark.asyncio
async def test_push_pop_keeps_stacks_in_sync(self, callback_context):
"""Push and pop should always leave the records stack consistent."""
TM = bigquery_agent_analytics_plugin.TraceManager
TM.init_trace(callback_context)
span_id_1 = TM.push_span(callback_context, "span-1")
span_id_2 = TM.push_span(callback_context, "span-2")
# Both should be on the stack
assert TM.get_current_span_id() == span_id_2
current, parent = TM.get_current_span_and_parent()
assert current == span_id_2
assert parent == span_id_1
# Pop span-2
popped_id, duration = TM.pop_span()
assert popped_id == span_id_2
assert duration is not None
assert TM.get_current_span_id() == span_id_1
# Pop span-1
popped_id, _ = TM.pop_span()
assert popped_id == span_id_1
assert TM.get_current_span_id() is None
@pytest.mark.asyncio
async def test_pop_empty_stack_returns_none(self, callback_context):
"""Popping an empty stack should return (None, None)."""
TM = bigquery_agent_analytics_plugin.TraceManager
TM.init_trace(callback_context)
span_id, duration = TM.pop_span()
assert span_id is None
assert duration is None
@pytest.mark.asyncio
async def test_first_token_time_stored_in_record(self, callback_context):
"""first_token_time should be stored on the span record."""
TM = bigquery_agent_analytics_plugin.TraceManager
TM.init_trace(callback_context)
span_id = TM.push_span(callback_context, "llm-span")
# No first token yet
assert TM.get_first_token_time(span_id) is None
# Record first token
assert TM.record_first_token(span_id) is True
ftt = TM.get_first_token_time(span_id)
assert ftt is not None
# Second call should return False (already recorded)
assert TM.record_first_token(span_id) is False
# Clean up
TM.pop_span()
@pytest.mark.asyncio
async def test_start_time_accessible_by_span_id(self, callback_context):
"""get_start_time should find the span by ID in the records."""
TM = bigquery_agent_analytics_plugin.TraceManager
TM.init_trace(callback_context)
span_id = TM.push_span(callback_context, "timed-span")
start = TM.get_start_time(span_id)
assert start is not None
assert start > 0
TM.pop_span()
@pytest.mark.asyncio
async def test_attach_current_span_does_not_own(self, callback_context):
"""attach_current_span should not end the span on pop."""
TM = bigquery_agent_analytics_plugin.TraceManager
TM.init_trace(callback_context)
mock_span = mock.Mock()
mock_ctx = mock.Mock()
mock_ctx.is_valid = False
mock_span.get_span_context.return_value = mock_ctx
with mock.patch(
"opentelemetry.trace.get_current_span", return_value=mock_span
):
span_id = TM.attach_current_span(callback_context)
assert span_id is not None
TM.pop_span()
# Should NOT have called span.end() since we don't own it
mock_span.end.assert_not_called()
@pytest.mark.asyncio
async def test_concurrent_tasks_have_isolated_stacks(self, callback_context):
"""Concurrent async tasks should have isolated span stacks."""
TM = bigquery_agent_analytics_plugin.TraceManager
TM.init_trace(callback_context)
async def task_a():
s = TM.push_span(callback_context, "task-a")
await asyncio.sleep(0.02)
assert TM.get_current_span_id() == s
TM.pop_span()
return s
async def task_b():
s = TM.push_span(callback_context, "task-b")
await asyncio.sleep(0.02)
assert TM.get_current_span_id() == s
TM.pop_span()
return s
results = await asyncio.gather(task_a(), task_b())
assert results[0] != results[1]
@pytest.mark.asyncio
async def test_pop_cleans_up_record_completely(self, callback_context):
"""After pop, the record should be fully removed from the stack."""
TM = bigquery_agent_analytics_plugin.TraceManager
TM.init_trace(callback_context)
span_id = TM.push_span(callback_context, "temp-span")
# Record is on the stack
assert TM.get_current_span_id() == span_id
assert TM.get_start_time(span_id) is not None
TM.pop_span()
# Record is gone
assert TM.get_current_span_id() is None
assert TM.get_start_time(span_id) is None
assert TM.get_first_token_time(span_id) is None
class TestLoopStateValidation:
"""Tests for loop state validation and stale loop cleanup."""
def _make_plugin(self):
"""Creates a plugin instance without starting it."""
return bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
def _make_loop_state(self):
"""Creates a mock _LoopState with batch_processor and write_client."""
state = mock.MagicMock()
state.batch_processor = mock.MagicMock(
spec=bigquery_agent_analytics_plugin.BatchProcessor
)
state.batch_processor.flush = mock.AsyncMock()
state.write_client = mock.MagicMock()
return state
def test_cleanup_stale_loop_states_removes_closed_loops(self):
"""Closed loops should be removed from _loop_state_by_loop."""
plugin = self._make_plugin()
closed_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop)
closed_loop.is_closed.return_value = True
plugin._loop_state_by_loop[closed_loop] = self._make_loop_state()
plugin._cleanup_stale_loop_states()
assert closed_loop not in plugin._loop_state_by_loop
def test_cleanup_stale_loop_states_keeps_open_loops(self):
"""Open loops should not be removed from _loop_state_by_loop."""
plugin = self._make_plugin()
open_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop)
open_loop.is_closed.return_value = False
plugin._loop_state_by_loop[open_loop] = self._make_loop_state()
plugin._cleanup_stale_loop_states()
assert open_loop in plugin._loop_state_by_loop
def test_cleanup_removes_only_closed_loops(self):
"""Only closed loops should be removed; open ones stay."""
plugin = self._make_plugin()
open_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop)
open_loop.is_closed.return_value = False
closed_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop)
closed_loop.is_closed.return_value = True
plugin._loop_state_by_loop[open_loop] = self._make_loop_state()
plugin._loop_state_by_loop[closed_loop] = self._make_loop_state()
plugin._cleanup_stale_loop_states()
assert open_loop in plugin._loop_state_by_loop
assert closed_loop not in plugin._loop_state_by_loop
@pytest.mark.asyncio
async def testbatch_processor_returns_processor_for_open_loop(
self,
):
"""batch_processor returns processor for the current loop."""
plugin = self._make_plugin()
loop = asyncio.get_running_loop()
state = self._make_loop_state()
plugin._loop_state_by_loop[loop] = state
assert plugin.batch_processor is state.batch_processor
# Clean up
del plugin._loop_state_by_loop[loop]
@pytest.mark.asyncio
async def testbatch_processor_cleans_closed_loop_entry(self):
"""Accessing batch_processor cleans up closed loop entries."""
plugin = self._make_plugin()
closed_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop)
closed_loop.is_closed.return_value = True
plugin._loop_state_by_loop[closed_loop] = self._make_loop_state()
# Accessing the prop should clean up the closed loop entry
_ = plugin.batch_processor
assert closed_loop not in plugin._loop_state_by_loop
@pytest.mark.asyncio
async def test_flush_cleans_stale_states(self):
"""flush() should clean up stale loop states before flushing."""
plugin = self._make_plugin()
closed_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop)
closed_loop.is_closed.return_value = True
plugin._loop_state_by_loop[closed_loop] = self._make_loop_state()
await plugin.flush()
assert closed_loop not in plugin._loop_state_by_loop
class TestAtexitCleanup:
"""Tests for the simplified _atexit_cleanup static method."""
def _make_batch_processor(self, queue_items=0):
bp = mock.MagicMock()
bp._shutdown = False
q = asyncio.Queue()
for i in range(queue_items):
q.put_nowait({"event": i})
bp._queue = q
return bp
def test_skips_none_processor(self):
"""Should return immediately when batch_processor is None."""
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._atexit_cleanup(
None
)
def test_skips_already_shutdown(self):
"""Should return immediately when batch_processor._shutdown is True."""
bp = self._make_batch_processor()
bp._shutdown = True
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._atexit_cleanup(
bp
)
def test_skips_reference_error(self):
"""Should handle ReferenceError from weakref'd processor."""
bp = mock.MagicMock()
type(bp)._shutdown = mock.PropertyMock(side_effect=ReferenceError)
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._atexit_cleanup(
bp
)
def test_empty_queue_no_warning(self):
"""Should not warn when queue is empty."""
bp = self._make_batch_processor(queue_items=0)
with mock.patch.object(
bigquery_agent_analytics_plugin.logger, "warning"
) as mock_warn:
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._atexit_cleanup(
bp
)
mock_warn.assert_not_called()
def test_remaining_items_logs_warning(self):
"""Should drain queue and log warning with count of lost items."""
bp = self._make_batch_processor(queue_items=3)
with mock.patch.object(
bigquery_agent_analytics_plugin.logger, "warning"
) as mock_warn:
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._atexit_cleanup(
bp
)
mock_warn.assert_called_once()
# Verify the warning mentions the count
call_args = mock_warn.call_args
assert "3" in str(call_args)
def test_queue_is_drained(self):
"""Should drain all items from the queue."""
bp = self._make_batch_processor(queue_items=5)
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._atexit_cleanup(
bp
)
assert bp._queue.empty()
class TestDuplicateLabels:
"""Tests that labels in before_model_callback are set exactly once."""
@pytest.mark.asyncio
async def test_labels_set_when_present(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Labels should appear in attributes when config has them."""
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
config=types.GenerateContentConfig(
labels={"env": "test"},
),
contents=[types.Content(role="user", parts=[types.Part(text="hi")])],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
attributes = json.loads(log_entry["attributes"])
assert attributes["labels"] == {"env": "test"}
@pytest.mark.asyncio
async def test_labels_absent_when_none(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Labels should not appear in attributes when config.labels is None."""
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
config=types.GenerateContentConfig(
temperature=0.5,
),
contents=[types.Content(role="user", parts=[types.Part(text="hi")])],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
attributes = json.loads(log_entry["attributes"])
assert "labels" not in attributes
@pytest.mark.asyncio
async def test_no_config_no_labels(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Labels should not appear when llm_request has no config."""
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(role="user", parts=[types.Part(text="hi")])],
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
attributes = json.loads(log_entry["attributes"])
assert "labels" not in attributes
class TestResolveIds:
"""Tests for the _resolve_ids static helper."""
def _resolve(self, ed, callback_context):
return bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._resolve_ids(
ed, callback_context
)
def test_uses_trace_manager_defaults(self, callback_context):
"""Should use TraceManager values when no overrides and no ambient."""
ed = bigquery_agent_analytics_plugin.EventData(
extra_attributes={"some_key": "value"}
)
with (
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_current_span_and_parent",
return_value=("span-1", "parent-1"),
),
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_trace_id",
return_value="trace-1",
),
):
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
assert trace_id == "trace-1"
assert span_id == "span-1"
assert parent_id == "parent-1"
def test_span_id_override(self, callback_context):
"""Should use span_id_override from EventData."""
ed = bigquery_agent_analytics_plugin.EventData(
span_id_override="custom-span"
)
with (
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_current_span_and_parent",
return_value=("span-1", "parent-1"),
),
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_trace_id",
return_value="trace-1",
),
):
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
assert span_id == "custom-span"
assert parent_id == "parent-1"
def test_parent_span_id_override(self, callback_context):
"""Should use parent_span_id_override from EventData."""
ed = bigquery_agent_analytics_plugin.EventData(
parent_span_id_override="custom-parent"
)
with (
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_current_span_and_parent",
return_value=("span-1", "parent-1"),
),
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_trace_id",
return_value="trace-1",
),
):
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
assert span_id == "span-1"
assert parent_id == "custom-parent"
def test_none_override_keeps_default(self, callback_context):
"""None overrides should keep the TraceManager defaults."""
ed = bigquery_agent_analytics_plugin.EventData(
span_id_override=None, parent_span_id_override=None
)
with (
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_current_span_and_parent",
return_value=("span-1", "parent-1"),
),
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_trace_id",
return_value="trace-1",
),
):
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
assert span_id == "span-1"
assert parent_id == "parent-1"
def test_ambient_provides_trace_id_only_when_stack_present(
self, callback_context
):
"""Plugin stack owns span_id/parent; ambient only provides trace_id."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
ed = bigquery_agent_analytics_plugin.EventData()
with real_tracer.start_as_current_span("invocation") as parent_span:
with real_tracer.start_as_current_span("agent") as agent_span:
ambient_ctx = agent_span.get_span_context()
expected_trace = format(ambient_ctx.trace_id, "032x")
# Plugin stack has spans — these should win for span/parent.
with (
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_current_span_and_parent",
return_value=("plugin-span", "plugin-parent"),
),
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_trace_id",
return_value="plugin-trace",
),
):
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
# trace_id comes from ambient OTel.
assert trace_id == expected_trace
# span_id and parent_span_id come from plugin stack.
assert span_id == "plugin-span"
assert parent_id == "plugin-parent"
provider.shutdown()
def test_ambient_fallback_when_no_plugin_stack(self, callback_context):
"""Ambient OTel provides span_id/parent when plugin stack is empty."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
ed = bigquery_agent_analytics_plugin.EventData()
with real_tracer.start_as_current_span("invocation") as parent_span:
with real_tracer.start_as_current_span("agent") as agent_span:
ambient_ctx = agent_span.get_span_context()
expected_trace = format(ambient_ctx.trace_id, "032x")
expected_span = format(ambient_ctx.span_id, "016x")
expected_parent = format(parent_span.get_span_context().span_id, "016x")
# Plugin stack returns None — ambient is the fallback.
with (
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_current_span_and_parent",
return_value=(None, None),
),
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_trace_id",
return_value=None,
),
):
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
assert trace_id == expected_trace
assert span_id == expected_span
assert parent_id == expected_parent
provider.shutdown()
def test_override_beats_ambient(self, callback_context):
"""EventData overrides take priority over ambient OTel span."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
ed = bigquery_agent_analytics_plugin.EventData(
trace_id_override="forced-trace",
span_id_override="forced-span",
parent_span_id_override="forced-parent",
)
with real_tracer.start_as_current_span("invocation"):
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
assert trace_id == "forced-trace"
assert span_id == "forced-span"
assert parent_id == "forced-parent"
provider.shutdown()
def test_plugin_stack_wins_over_ambient_root_span(self, callback_context):
"""Plugin stack span is used even when ambient root span exists."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
# Seed the plugin stack with a span.
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "plugin-child"
)
# Capture the plugin span_id that was pushed.
plugin_span_id, _ = (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_and_parent()
)
ed = bigquery_agent_analytics_plugin.EventData()
# Single root ambient span — no parent.
with real_tracer.start_as_current_span("root_invocation") as root:
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
ambient_trace = format(root.get_span_context().trace_id, "032x")
# trace_id comes from ambient.
assert trace_id == ambient_trace
# span_id comes from plugin stack, not ambient.
assert span_id == plugin_span_id
# parent is None — only one span in plugin stack.
assert parent_id is None
# Cleanup
bigquery_agent_analytics_plugin.TraceManager.pop_span()
provider.shutdown()
def test_ambient_root_fallback_no_self_parent(self, callback_context):
"""Ambient root span fallback must not produce self-parent."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
ed = bigquery_agent_analytics_plugin.EventData()
# Plugin stack empty — ambient provides the fallback.
with (
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_current_span_and_parent",
return_value=(None, None),
),
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_trace_id",
return_value=None,
),
):
with real_tracer.start_as_current_span("root") as root:
trace_id, span_id, parent_id = self._resolve(ed, callback_context)
root_span_id = format(root.get_span_context().span_id, "016x")
assert span_id == root_span_id
assert parent_id is None
provider.shutdown()
def test_plugin_stack_pairs_starting_completed(self, callback_context):
"""STARTING/COMPLETED pairing uses plugin stack, not ambient.
Post-pop callbacks now always pass explicit overrides from the
plugin stack. The plugin stack span_id is used for both events
regardless of ambient OTel state.
"""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with real_tracer.start_as_current_span("invoke_agent"):
# Simulate STARTING: plugin stack provides span_id.
with (
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_current_span_and_parent",
return_value=("plugin-agent", "plugin-inv"),
),
mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_trace_id",
return_value="plugin-trace",
),
):
ed_starting = bigquery_agent_analytics_plugin.EventData()
_, span_starting, _ = self._resolve(ed_starting, callback_context)
# Simulate COMPLETED: explicit override from popped span.
ed_completed = bigquery_agent_analytics_plugin.EventData(
span_id_override="plugin-agent",
parent_span_id_override="plugin-inv",
latency_ms=42,
)
_, span_completed, _ = self._resolve(ed_completed, callback_context)
assert span_starting == "plugin-agent"
assert span_completed == "plugin-agent"
assert span_starting == span_completed
provider.shutdown()
class TestExtractLatency:
"""Tests for the _extract_latency static helper."""
def test_no_latency_returns_none(self):
"""Should return None when no latency fields present."""
ed = bigquery_agent_analytics_plugin.EventData(
extra_attributes={"other": "val"}
)
result = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._extract_latency(
ed
)
assert result is None
def test_total_latency_only(self):
"""Should extract latency_ms into total_ms."""
ed = bigquery_agent_analytics_plugin.EventData(latency_ms=42.5)
result = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._extract_latency(
ed
)
assert result == {"total_ms": 42.5}
def test_tfft_only(self):
"""Should extract time_to_first_token_ms."""
ed = bigquery_agent_analytics_plugin.EventData(time_to_first_token_ms=10.0)
result = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._extract_latency(
ed
)
assert result == {"time_to_first_token_ms": 10.0}
def test_both_latencies(self):
"""Should extract both latency fields."""
ed = bigquery_agent_analytics_plugin.EventData(
latency_ms=100, time_to_first_token_ms=20
)
result = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin._extract_latency(
ed
)
assert result == {"total_ms": 100, "time_to_first_token_ms": 20}
class TestEnrichAttributes:
"""Tests for the _enrich_attributes helper."""
def _make_plugin(self):
with (
mock.patch(
"google.auth.default",
return_value=(mock.Mock(), PROJECT_ID),
),
mock.patch(
"google.cloud.bigquery.Client",
),
):
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
)
plugin.config.max_content_length = 10000
plugin.config.log_session_metadata = False
plugin.config.custom_tags = None
return plugin
def _make_callback_context(self):
ctx = mock.MagicMock()
session = mock.MagicMock()
session.id = "sess-001"
session.app_name = "test-app"
session.user_id = "user-001"
session.state = {"env": "test"}
ctx._invocation_context.session = session
return ctx
def test_adds_root_agent_name(self):
"""Should always add root_agent_name."""
plugin = self._make_plugin()
ed = bigquery_agent_analytics_plugin.EventData()
with mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_root_agent_name",
return_value="my-agent",
):
attrs = plugin._enrich_attributes(ed, self._make_callback_context())
assert attrs["root_agent_name"] == "my-agent"
def test_includes_model(self):
"""Should include model from EventData."""
plugin = self._make_plugin()
ed = bigquery_agent_analytics_plugin.EventData(model="gemini-pro")
with mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_root_agent_name",
return_value="agent",
):
attrs = plugin._enrich_attributes(ed, self._make_callback_context())
assert attrs["model"] == "gemini-pro"
def test_session_metadata_when_enabled(self):
"""Should add session_metadata when log_session_metadata is True."""
plugin = self._make_plugin()
plugin.config.log_session_metadata = True
ctx = self._make_callback_context()
ed = bigquery_agent_analytics_plugin.EventData()
with mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_root_agent_name",
return_value="agent",
):
attrs = plugin._enrich_attributes(ed, ctx)
meta = attrs["session_metadata"]
assert meta["session_id"] == "sess-001"
assert meta["app_name"] == "test-app"
assert meta["user_id"] == "user-001"
assert meta["state"] == {"env": "test"}
def test_session_metadata_when_disabled(self):
"""Should not add session_metadata when log_session_metadata is False."""
plugin = self._make_plugin()
plugin.config.log_session_metadata = False
ed = bigquery_agent_analytics_plugin.EventData()
with mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_root_agent_name",
return_value="agent",
):
attrs = plugin._enrich_attributes(ed, self._make_callback_context())
assert "session_metadata" not in attrs
def test_custom_tags_added(self):
"""Should add custom_tags when configured."""
plugin = self._make_plugin()
plugin.config.custom_tags = {"team": "infra"}
ed = bigquery_agent_analytics_plugin.EventData()
with mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_root_agent_name",
return_value="agent",
):
attrs = plugin._enrich_attributes(ed, self._make_callback_context())
assert attrs["custom_tags"] == {"team": "infra"}
def test_usage_metadata_truncated(self):
"""Should smart-truncate usage_metadata."""
plugin = self._make_plugin()
ed = bigquery_agent_analytics_plugin.EventData(
usage_metadata={"input_tokens": 100, "output_tokens": 50}
)
with mock.patch.object(
bigquery_agent_analytics_plugin.TraceManager,
"get_root_agent_name",
return_value="agent",
):
attrs = plugin._enrich_attributes(ed, self._make_callback_context())
assert attrs["usage_metadata"] == {
"input_tokens": 100,
"output_tokens": 50,
}
class TestMultiSubagentToolLogging:
"""Tests that tool events from different subagents are attributed correctly.
Covers:
- Tool calls from different subagents have the correct `agent` field
- Multi-turn (different invocation_ids, same session) logs correctly
- Full callback sequence across multiple subagents in one turn
- Span hierarchy is maintained per-subagent
"""
@staticmethod
def _make_invocation_context(agent_name, session, invocation_id="inv-001"):
"""Create an InvocationContext with a specific agent name."""
mock_a = mock.create_autospec(
base_agent.BaseAgent, instance=True, spec_set=True
)
type(mock_a).name = mock.PropertyMock(return_value=agent_name)
type(mock_a).instruction = mock.PropertyMock(
return_value=f"{agent_name} instruction"
)
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 InvocationContext(
agent=mock_a,
session=session,
invocation_id=invocation_id,
session_service=mock_session_service,
plugin_manager=mock_plugin_manager,
)
@staticmethod
def _make_session(session_id="session-multi", user_id="user-multi"):
mock_s = mock.create_autospec(
session_lib.Session, instance=True, spec_set=True
)
type(mock_s).id = mock.PropertyMock(return_value=session_id)
type(mock_s).user_id = mock.PropertyMock(return_value=user_id)
type(mock_s).app_name = mock.PropertyMock(return_value="test_app")
type(mock_s).state = mock.PropertyMock(return_value={})
return mock_s
@staticmethod
def _make_tool(name):
mock_tool = mock.create_autospec(
base_tool_lib.BaseTool, instance=True, spec_set=True
)
type(mock_tool).name = mock.PropertyMock(return_value=name)
type(mock_tool).description = mock.PropertyMock(
return_value=f"{name} description"
)
return mock_tool
@pytest.mark.asyncio
async def test_tool_calls_attributed_to_correct_subagent(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""Tool events from different subagents carry the correct agent name."""
session = self._make_session()
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
# --- Subagent A: schema_explorer calls list_datasets ---
inv_ctx_a = self._make_invocation_context("schema_explorer", session)
ctx_a = tool_context_lib.ToolContext(invocation_context=inv_ctx_a)
tool_a = self._make_tool("list_dataset_ids")
bigquery_agent_analytics_plugin.TraceManager.push_span(ctx_a, "tool")
await plugin.before_tool_callback(
tool=tool_a,
tool_args={"project_id": "my-project"},
tool_context=ctx_a,
)
await asyncio.sleep(0.01)
# --- Subagent B: image_describer calls describe_this_image ---
inv_ctx_b = self._make_invocation_context("image_describer", session)
ctx_b = tool_context_lib.ToolContext(invocation_context=inv_ctx_b)
tool_b = self._make_tool("describe_this_image")
bigquery_agent_analytics_plugin.TraceManager.push_span(ctx_b, "tool")
await plugin.before_tool_callback(
tool=tool_b,
tool_args={"image_uri": "gs://bucket/image.jpg"},
tool_context=ctx_b,
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
assert len(rows) == 2
# First row: schema_explorer's tool
assert rows[0]["event_type"] == "TOOL_STARTING"
assert rows[0]["agent"] == "schema_explorer"
content_a = json.loads(rows[0]["content"])
assert content_a["tool"] == "list_dataset_ids"
assert content_a["args"] == {"project_id": "my-project"}
# Second row: image_describer's tool
assert rows[1]["event_type"] == "TOOL_STARTING"
assert rows[1]["agent"] == "image_describer"
content_b = json.loads(rows[1]["content"])
assert content_b["tool"] == "describe_this_image"
assert content_b["args"] == {"image_uri": "gs://bucket/image.jpg"}
@pytest.mark.asyncio
async def test_multi_turn_tool_calls_different_invocations(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""Multi-turn: same session, different invocation IDs, tools logged."""
session = self._make_session()
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
# --- Turn 1: schema_explorer calls list_dataset_ids ---
inv_ctx_1 = self._make_invocation_context(
"schema_explorer", session, invocation_id="inv-turn1"
)
ctx_1 = tool_context_lib.ToolContext(invocation_context=inv_ctx_1)
tool_1 = self._make_tool("list_dataset_ids")
bigquery_agent_analytics_plugin.TraceManager.push_span(ctx_1, "tool")
await plugin.before_tool_callback(
tool=tool_1,
tool_args={"project_id": "proj"},
tool_context=ctx_1,
)
await asyncio.sleep(0.01)
await plugin.after_tool_callback(
tool=tool_1,
tool_args={"project_id": "proj"},
tool_context=ctx_1,
result={"datasets": ["ds1", "ds2"]},
)
await asyncio.sleep(0.01)
# --- Turn 2: query_analyst calls execute_sql ---
inv_ctx_2 = self._make_invocation_context(
"query_analyst", session, invocation_id="inv-turn2"
)
ctx_2 = tool_context_lib.ToolContext(invocation_context=inv_ctx_2)
tool_2 = self._make_tool("execute_sql")
bigquery_agent_analytics_plugin.TraceManager.push_span(ctx_2, "tool")
await plugin.before_tool_callback(
tool=tool_2,
tool_args={"sql": "SELECT * FROM t"},
tool_context=ctx_2,
)
await asyncio.sleep(0.01)
await plugin.after_tool_callback(
tool=tool_2,
tool_args={"sql": "SELECT * FROM t"},
tool_context=ctx_2,
result={"rows": [{"col": "val"}]},
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
assert len(rows) == 4
# Turn 1: TOOL_STARTING + TOOL_COMPLETED for schema_explorer
assert rows[0]["event_type"] == "TOOL_STARTING"
assert rows[0]["agent"] == "schema_explorer"
assert rows[0]["invocation_id"] == "inv-turn1"
assert rows[0]["session_id"] == "session-multi"
assert rows[1]["event_type"] == "TOOL_COMPLETED"
assert rows[1]["agent"] == "schema_explorer"
assert rows[1]["invocation_id"] == "inv-turn1"
content_1 = json.loads(rows[1]["content"])
assert content_1["tool"] == "list_dataset_ids"
assert content_1["result"] == {"datasets": ["ds1", "ds2"]}
# Turn 2: TOOL_STARTING + TOOL_COMPLETED for query_analyst
assert rows[2]["event_type"] == "TOOL_STARTING"
assert rows[2]["agent"] == "query_analyst"
assert rows[2]["invocation_id"] == "inv-turn2"
assert rows[3]["event_type"] == "TOOL_COMPLETED"
assert rows[3]["agent"] == "query_analyst"
assert rows[3]["invocation_id"] == "inv-turn2"
content_2 = json.loads(rows[3]["content"])
assert content_2["tool"] == "execute_sql"
assert content_2["result"] == {"rows": [{"col": "val"}]}
@pytest.mark.asyncio
async def test_full_subagent_callback_sequence(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""Full lifecycle: agent_start → LLM → tool → tool_done → LLM → agent_done.
Simulates a subagent that makes an LLM call, then a tool call,
then another LLM call, and completes.
"""
session = self._make_session()
inv_ctx = self._make_invocation_context("schema_explorer", session)
cb_ctx = CallbackContext(invocation_context=inv_ctx)
tool_ctx = tool_context_lib.ToolContext(invocation_context=inv_ctx)
mock_agent = inv_ctx.agent
tool = self._make_tool("get_table_info")
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
# 1. AGENT_STARTING
await plugin.before_agent_callback(
agent=mock_agent, callback_context=cb_ctx
)
await asyncio.sleep(0.01)
# 2. LLM_REQUEST (agent decides to call a tool)
llm_req = llm_request_lib.LlmRequest(
model="gemini-2.5-flash",
contents=[
types.Content(parts=[types.Part(text="What tables exist?")])
],
)
await plugin.before_model_callback(
callback_context=cb_ctx, llm_request=llm_req
)
await asyncio.sleep(0.01)
# 3. LLM_RESPONSE (function call)
llm_resp = llm_response_lib.LlmResponse(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="get_table_info",
args={"table": "events"},
)
)
]
)
)
await plugin.after_model_callback(
callback_context=cb_ctx, llm_response=llm_resp
)
await asyncio.sleep(0.01)
# 4. TOOL_STARTING
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_ctx, "tool")
await plugin.before_tool_callback(
tool=tool,
tool_args={"table": "events"},
tool_context=tool_ctx,
)
await asyncio.sleep(0.01)
# 5. TOOL_COMPLETED
await plugin.after_tool_callback(
tool=tool,
tool_args={"table": "events"},
tool_context=tool_ctx,
result={"schema": [{"name": "id", "type": "INT64"}]},
)
await asyncio.sleep(0.01)
# 6. AGENT_COMPLETED
await plugin.after_agent_callback(
agent=mock_agent, callback_context=cb_ctx
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
assert len(rows) == 6
expected_sequence = [
"AGENT_STARTING",
"LLM_REQUEST",
"LLM_RESPONSE",
"TOOL_STARTING",
"TOOL_COMPLETED",
"AGENT_COMPLETED",
]
for i, expected_type in enumerate(expected_sequence):
assert (
rows[i]["event_type"] == expected_type
), f"Row {i}: expected {expected_type}, got {rows[i]['event_type']}"
assert rows[i]["agent"] == "schema_explorer"
assert rows[i]["session_id"] == "session-multi"
# TOOL rows have correct content
tool_start = json.loads(rows[3]["content"])
assert tool_start["tool"] == "get_table_info"
assert tool_start["args"] == {"table": "events"}
tool_done = json.loads(rows[4]["content"])
assert tool_done["tool"] == "get_table_info"
assert tool_done["result"] == {"schema": [{"name": "id", "type": "INT64"}]}
# AGENT_COMPLETED and TOOL_COMPLETED should have latency
assert rows[4]["latency_ms"] is not None # TOOL_COMPLETED
assert rows[5]["latency_ms"] is not None # AGENT_COMPLETED
@pytest.mark.asyncio
async def test_tool_error_attributed_to_subagent(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""TOOL_ERROR events carry the correct subagent name."""
session = self._make_session()
inv_ctx = self._make_invocation_context("query_analyst", session)
tool_ctx = tool_context_lib.ToolContext(invocation_context=inv_ctx)
tool = self._make_tool("execute_sql")
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_ctx, "tool")
await plugin.on_tool_error_callback(
tool=tool,
tool_args={"sql": "SELECT * FROM bad_table"},
tool_context=tool_ctx,
error=RuntimeError("Table not found"),
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
assert len(rows) == 1
assert rows[0]["event_type"] == "TOOL_ERROR"
assert rows[0]["agent"] == "query_analyst"
assert rows[0]["error_message"] == "Table not found"
content = json.loads(rows[0]["content"])
assert content["tool"] == "execute_sql"
assert content["args"] == {"sql": "SELECT * FROM bad_table"}
@pytest.mark.asyncio
async def test_multi_subagent_interleaved_tool_calls(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""Two subagents call tools in same invocation — agent field is correct.
Simulates orchestrator delegating to schema_explorer first, then
image_describer, all within the same invocation.
"""
session = self._make_session()
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
# Subagent 1: schema_explorer — full tool cycle
inv_ctx_1 = self._make_invocation_context(
"schema_explorer", session, invocation_id="inv-shared"
)
ctx_1 = tool_context_lib.ToolContext(invocation_context=inv_ctx_1)
tool_1 = self._make_tool("list_table_ids")
bigquery_agent_analytics_plugin.TraceManager.push_span(ctx_1, "tool")
await plugin.before_tool_callback(
tool=tool_1,
tool_args={"dataset": "analytics"},
tool_context=ctx_1,
)
await asyncio.sleep(0.01)
await plugin.after_tool_callback(
tool=tool_1,
tool_args={"dataset": "analytics"},
tool_context=ctx_1,
result={"tables": ["events", "metrics"]},
)
await asyncio.sleep(0.01)
# Subagent 2: image_describer — full tool cycle
inv_ctx_2 = self._make_invocation_context(
"image_describer", session, invocation_id="inv-shared"
)
ctx_2 = tool_context_lib.ToolContext(invocation_context=inv_ctx_2)
tool_2 = self._make_tool("describe_this_image")
bigquery_agent_analytics_plugin.TraceManager.push_span(ctx_2, "tool")
await plugin.before_tool_callback(
tool=tool_2,
tool_args={"image_uri": "https://example.com/img.jpg"},
tool_context=ctx_2,
)
await asyncio.sleep(0.01)
await plugin.after_tool_callback(
tool=tool_2,
tool_args={"image_uri": "https://example.com/img.jpg"},
tool_context=ctx_2,
result={"description": "A photo of scones"},
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
assert len(rows) == 4
# schema_explorer tool events
assert rows[0]["agent"] == "schema_explorer"
assert rows[0]["event_type"] == "TOOL_STARTING"
assert rows[0]["invocation_id"] == "inv-shared"
assert json.loads(rows[0]["content"])["tool"] == "list_table_ids"
assert rows[1]["agent"] == "schema_explorer"
assert rows[1]["event_type"] == "TOOL_COMPLETED"
assert json.loads(rows[1]["content"])["result"]["tables"] == [
"events",
"metrics",
]
# image_describer tool events
assert rows[2]["agent"] == "image_describer"
assert rows[2]["event_type"] == "TOOL_STARTING"
assert rows[2]["invocation_id"] == "inv-shared"
assert json.loads(rows[2]["content"])["tool"] == "describe_this_image"
assert rows[3]["agent"] == "image_describer"
assert rows[3]["event_type"] == "TOOL_COMPLETED"
assert (
json.loads(rows[3]["content"])["result"]["description"]
== "A photo of scones"
)
# All share the same session and invocation
for row in rows:
assert row["session_id"] == "session-multi"
assert row["invocation_id"] == "inv-shared"
@pytest.mark.asyncio
async def test_multi_turn_multi_subagent_full_sequence(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""Multi-turn + multi-subagent: two turns, each with different subagents.
Turn 1: user asks about data → orchestrator → schema_explorer (tool)
Turn 2: user asks about image → orchestrator → image_describer (tool)
Verifies invocation_id changes, agent name changes, session stays same.
"""
session = self._make_session()
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
# ===== Turn 1: schema_explorer =====
inv_ctx_t1_orch = self._make_invocation_context(
"orchestrator", session, invocation_id="inv-t1"
)
cb_ctx_t1_orch = CallbackContext(invocation_context=inv_ctx_t1_orch)
# Orchestrator agent_starting
await plugin.before_agent_callback(
agent=inv_ctx_t1_orch.agent,
callback_context=cb_ctx_t1_orch,
)
await asyncio.sleep(0.01)
# Orchestrator delegates to schema_explorer
inv_ctx_t1_sub = self._make_invocation_context(
"schema_explorer", session, invocation_id="inv-t1"
)
cb_ctx_t1_sub = CallbackContext(invocation_context=inv_ctx_t1_sub)
tool_ctx_t1 = tool_context_lib.ToolContext(
invocation_context=inv_ctx_t1_sub
)
await plugin.before_agent_callback(
agent=inv_ctx_t1_sub.agent,
callback_context=cb_ctx_t1_sub,
)
await asyncio.sleep(0.01)
# schema_explorer calls tool
tool_1 = self._make_tool("list_dataset_ids")
bigquery_agent_analytics_plugin.TraceManager.push_span(
tool_ctx_t1, "tool"
)
await plugin.before_tool_callback(
tool=tool_1,
tool_args={"project_id": "proj"},
tool_context=tool_ctx_t1,
)
await asyncio.sleep(0.01)
await plugin.after_tool_callback(
tool=tool_1,
tool_args={"project_id": "proj"},
tool_context=tool_ctx_t1,
result={"datasets": ["ds1"]},
)
await asyncio.sleep(0.01)
# schema_explorer done
await plugin.after_agent_callback(
agent=inv_ctx_t1_sub.agent,
callback_context=cb_ctx_t1_sub,
)
await asyncio.sleep(0.01)
# Orchestrator done
await plugin.after_agent_callback(
agent=inv_ctx_t1_orch.agent,
callback_context=cb_ctx_t1_orch,
)
await asyncio.sleep(0.01)
# ===== Turn 2: image_describer =====
inv_ctx_t2_orch = self._make_invocation_context(
"orchestrator", session, invocation_id="inv-t2"
)
cb_ctx_t2_orch = CallbackContext(invocation_context=inv_ctx_t2_orch)
await plugin.before_agent_callback(
agent=inv_ctx_t2_orch.agent,
callback_context=cb_ctx_t2_orch,
)
await asyncio.sleep(0.01)
# Orchestrator delegates to image_describer
inv_ctx_t2_sub = self._make_invocation_context(
"image_describer", session, invocation_id="inv-t2"
)
cb_ctx_t2_sub = CallbackContext(invocation_context=inv_ctx_t2_sub)
tool_ctx_t2 = tool_context_lib.ToolContext(
invocation_context=inv_ctx_t2_sub
)
await plugin.before_agent_callback(
agent=inv_ctx_t2_sub.agent,
callback_context=cb_ctx_t2_sub,
)
await asyncio.sleep(0.01)
# image_describer calls tool
tool_2 = self._make_tool("describe_this_image")
bigquery_agent_analytics_plugin.TraceManager.push_span(
tool_ctx_t2, "tool"
)
await plugin.before_tool_callback(
tool=tool_2,
tool_args={"image_uri": "gs://b/img.jpg"},
tool_context=tool_ctx_t2,
)
await asyncio.sleep(0.01)
await plugin.after_tool_callback(
tool=tool_2,
tool_args={"image_uri": "gs://b/img.jpg"},
tool_context=tool_ctx_t2,
result={"desc": "Scones on a table"},
)
await asyncio.sleep(0.01)
# image_describer done
await plugin.after_agent_callback(
agent=inv_ctx_t2_sub.agent,
callback_context=cb_ctx_t2_sub,
)
await asyncio.sleep(0.01)
# Orchestrator done
await plugin.after_agent_callback(
agent=inv_ctx_t2_orch.agent,
callback_context=cb_ctx_t2_orch,
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
# Turn 1: 6 rows (orch_start, sub_start, tool_start, tool_done,
# sub_done, orch_done)
# Turn 2: 6 rows (same pattern)
assert len(rows) == 12
# --- Turn 1 validation ---
t1_rows = [r for r in rows if r["invocation_id"] == "inv-t1"]
assert len(t1_rows) == 6
assert t1_rows[0]["event_type"] == "AGENT_STARTING"
assert t1_rows[0]["agent"] == "orchestrator"
assert t1_rows[1]["event_type"] == "AGENT_STARTING"
assert t1_rows[1]["agent"] == "schema_explorer"
assert t1_rows[2]["event_type"] == "TOOL_STARTING"
assert t1_rows[2]["agent"] == "schema_explorer"
assert json.loads(t1_rows[2]["content"])["tool"] == "list_dataset_ids"
assert t1_rows[3]["event_type"] == "TOOL_COMPLETED"
assert t1_rows[3]["agent"] == "schema_explorer"
assert t1_rows[4]["event_type"] == "AGENT_COMPLETED"
assert t1_rows[4]["agent"] == "schema_explorer"
assert t1_rows[5]["event_type"] == "AGENT_COMPLETED"
assert t1_rows[5]["agent"] == "orchestrator"
# --- Turn 2 validation ---
t2_rows = [r for r in rows if r["invocation_id"] == "inv-t2"]
assert len(t2_rows) == 6
assert t2_rows[0]["event_type"] == "AGENT_STARTING"
assert t2_rows[0]["agent"] == "orchestrator"
assert t2_rows[1]["event_type"] == "AGENT_STARTING"
assert t2_rows[1]["agent"] == "image_describer"
assert t2_rows[2]["event_type"] == "TOOL_STARTING"
assert t2_rows[2]["agent"] == "image_describer"
assert json.loads(t2_rows[2]["content"])["tool"] == "describe_this_image"
assert t2_rows[3]["event_type"] == "TOOL_COMPLETED"
assert t2_rows[3]["agent"] == "image_describer"
assert t2_rows[4]["event_type"] == "AGENT_COMPLETED"
assert t2_rows[4]["agent"] == "image_describer"
assert t2_rows[5]["event_type"] == "AGENT_COMPLETED"
assert t2_rows[5]["agent"] == "orchestrator"
# All rows share the same session
for row in rows:
assert row["session_id"] == "session-multi"
class TestSchemaAutoUpgrade:
"""Tests for _ensure_schema_exists with auto_schema_upgrade."""
def _make_plugin(self, auto_schema_upgrade=False):
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
auto_schema_upgrade=auto_schema_upgrade,
)
with mock.patch("google.cloud.bigquery.Client"):
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=config,
)
plugin.client = mock.MagicMock()
plugin.full_table_id = f"{PROJECT_ID}.{DATASET_ID}.{TABLE_ID}"
plugin._schema = bigquery_agent_analytics_plugin._get_events_schema()
return plugin
def test_create_table_sets_version_label(self):
"""New tables get the schema version label."""
plugin = self._make_plugin()
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
plugin._ensure_schema_exists()
plugin.client.create_table.assert_called_once()
tbl = plugin.client.create_table.call_args[0][0]
assert (
tbl.labels[bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY]
== bigquery_agent_analytics_plugin._SCHEMA_VERSION
)
def test_no_upgrade_when_disabled(self):
"""Auto-upgrade disabled: existing table is not modified."""
plugin = self._make_plugin(auto_schema_upgrade=False)
existing = mock.MagicMock()
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
]
existing.labels = {}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
plugin.client.update_table.assert_not_called()
def test_upgrade_adds_missing_columns(self):
"""Auto-upgrade adds columns missing from existing table."""
plugin = self._make_plugin(auto_schema_upgrade=True)
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
]
existing.labels = {"other": "label"}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
plugin.client.update_table.assert_called_once()
updated_table = plugin.client.update_table.call_args[0][0]
updated_names = {f.name for f in updated_table.schema}
assert "event_type" in updated_names
assert "agent" in updated_names
assert "content" in updated_names
assert (
updated_table.labels[
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY
]
== bigquery_agent_analytics_plugin._SCHEMA_VERSION
)
def test_skip_upgrade_when_version_matches(self):
"""No update when stored version matches current."""
plugin = self._make_plugin(auto_schema_upgrade=True)
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = plugin._schema
existing.labels = {
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY: (
bigquery_agent_analytics_plugin._SCHEMA_VERSION
),
}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
plugin.client.update_table.assert_not_called()
def test_upgrade_error_is_logged_not_raised(self):
"""Schema upgrade errors are logged, not propagated."""
plugin = self._make_plugin(auto_schema_upgrade=True)
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
]
existing.labels = {}
plugin.client.get_table.return_value = existing
plugin.client.update_table.side_effect = Exception("boom")
# Should not raise
plugin._ensure_schema_exists()
def test_upgrade_preserves_existing_columns(self):
"""Existing columns are never dropped or altered during upgrade."""
plugin = self._make_plugin(auto_schema_upgrade=True)
# Simulate a table with a subset of canonical columns plus a
# user-added custom column that is NOT in the canonical schema.
custom_field = bigquery.SchemaField("my_custom_col", "STRING")
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
bigquery.SchemaField("event_type", "STRING"),
custom_field,
]
existing.labels = {}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
updated_table = plugin.client.update_table.call_args[0][0]
updated_names = [f.name for f in updated_table.schema]
# Original columns are still present and in original order.
assert updated_names[0] == "timestamp"
assert updated_names[1] == "event_type"
assert updated_names[2] == "my_custom_col"
# New canonical columns were appended after existing ones.
assert "agent" in updated_names
assert "content" in updated_names
def test_upgrade_from_no_label_treats_as_outdated(self):
"""A table with no version label is treated as needing upgrade."""
plugin = self._make_plugin(auto_schema_upgrade=True)
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = list(plugin._schema) # All columns present
existing.labels = {} # No version label
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
# update_table should be called to stamp the version label even
# though no new columns were needed.
plugin.client.update_table.assert_called_once()
updated_table = plugin.client.update_table.call_args[0][0]
assert (
updated_table.labels[
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY
]
== bigquery_agent_analytics_plugin._SCHEMA_VERSION
)
def test_upgrade_from_older_version_label(self):
"""A table with an older version label triggers upgrade."""
plugin = self._make_plugin(auto_schema_upgrade=True)
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
bigquery.SchemaField("event_type", "STRING"),
]
# Simulate a table stamped with an older version.
existing.labels = {
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY: "0",
}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
plugin.client.update_table.assert_called_once()
updated_table = plugin.client.update_table.call_args[0][0]
# Version label should be updated to current.
assert (
updated_table.labels[
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY
]
== bigquery_agent_analytics_plugin._SCHEMA_VERSION
)
# Missing columns should have been added.
updated_names = {f.name for f in updated_table.schema}
assert "agent" in updated_names
assert "content" in updated_names
def test_upgrade_is_idempotent(self):
"""Calling _ensure_schema_exists twice doesn't double-update."""
plugin = self._make_plugin(auto_schema_upgrade=True)
# First call: table exists with old schema.
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
]
existing.labels = {}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
assert plugin.client.update_table.call_count == 1
# Second call: table now has current version label.
existing.labels = {
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY: (
bigquery_agent_analytics_plugin._SCHEMA_VERSION
),
}
plugin.client.update_table.reset_mock()
plugin._ensure_schema_exists()
plugin.client.update_table.assert_not_called()
def test_update_table_receives_schema_and_labels_fields(self):
"""update_table is called with update_fields=['schema', 'labels']."""
plugin = self._make_plugin(auto_schema_upgrade=True)
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
]
existing.labels = {}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
call_args = plugin.client.update_table.call_args
update_fields = call_args[0][1]
assert "schema" in update_fields
assert "labels" in update_fields
def test_auto_schema_upgrade_defaults_to_true(self):
"""Default config has auto_schema_upgrade enabled."""
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
assert config.auto_schema_upgrade is True
def test_create_table_conflict_is_ignored(self):
"""Race condition (Conflict) during create_table is silently handled."""
plugin = self._make_plugin()
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
plugin.client.create_table.side_effect = cloud_exceptions.Conflict(
"already exists"
)
# Should not raise.
plugin._ensure_schema_exists()
class TestToolProvenance:
"""Tests for _get_tool_origin helper."""
def test_function_tool_returns_local(self):
from google.adk.tools.function_tool import FunctionTool
def dummy():
pass
tool = FunctionTool(dummy)
result = bigquery_agent_analytics_plugin._get_tool_origin(tool)
assert result == "LOCAL"
def test_agent_tool_returns_sub_agent(self):
from google.adk.tools.agent_tool import AgentTool
agent = mock.MagicMock()
agent.name = "sub"
tool = AgentTool.__new__(AgentTool)
tool.agent = agent
tool._name = "sub"
result = bigquery_agent_analytics_plugin._get_tool_origin(tool)
assert result == "SUB_AGENT"
def test_transfer_tool_returns_transfer_agent(self):
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
tool = TransferToAgentTool(agent_names=["other"])
result = bigquery_agent_analytics_plugin._get_tool_origin(tool)
assert result == "TRANSFER_AGENT"
def test_transfer_tool_without_args_returns_transfer_agent(self):
"""TransferToAgentTool without tool_args falls back to TRANSFER_AGENT."""
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
tool = TransferToAgentTool(agent_names=["remote_a2a"])
result = bigquery_agent_analytics_plugin._get_tool_origin(
tool, tool_args=None, tool_context=None
)
assert result == "TRANSFER_AGENT"
def test_transfer_to_remote_a2a_sub_agent_returns_transfer_a2a(self):
"""Transfer to a RemoteA2aAgent sub-agent is classified TRANSFER_A2A."""
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
try:
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
except ImportError:
pytest.skip("A2A agent not available")
remote_agent = mock.MagicMock(spec=RemoteA2aAgent)
remote_agent.name = "remote_a2a"
current_agent = mock.MagicMock()
current_agent.name = "root"
current_agent.sub_agents = [remote_agent]
current_agent.parent_agent = None
inv_ctx = mock.MagicMock()
inv_ctx.agent = current_agent
tool_context = mock.MagicMock()
tool_context._invocation_context = inv_ctx
tool = TransferToAgentTool(agent_names=["remote_a2a"])
result = bigquery_agent_analytics_plugin._get_tool_origin(
tool,
tool_args={"agent_name": "remote_a2a"},
tool_context=tool_context,
)
assert result == "TRANSFER_A2A"
def test_transfer_to_local_sub_agent_returns_transfer_agent(self):
"""Transfer to a local sub-agent is still classified TRANSFER_AGENT."""
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
local_agent = mock.MagicMock()
local_agent.name = "local_sub"
current_agent = mock.MagicMock()
current_agent.name = "root"
current_agent.sub_agents = [local_agent]
current_agent.parent_agent = None
inv_ctx = mock.MagicMock()
inv_ctx.agent = current_agent
tool_context = mock.MagicMock()
tool_context._invocation_context = inv_ctx
tool = TransferToAgentTool(agent_names=["local_sub"])
result = bigquery_agent_analytics_plugin._get_tool_origin(
tool,
tool_args={"agent_name": "local_sub"},
tool_context=tool_context,
)
assert result == "TRANSFER_AGENT"
def test_transfer_to_a2a_peer_returns_transfer_a2a(self):
"""Transfer to a RemoteA2aAgent peer is classified TRANSFER_A2A."""
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
try:
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
except ImportError:
pytest.skip("A2A agent not available")
remote_peer = mock.MagicMock(spec=RemoteA2aAgent)
remote_peer.name = "remote_peer"
current_agent = mock.MagicMock()
current_agent.name = "child"
current_agent.sub_agents = []
parent_agent = mock.MagicMock()
parent_agent.name = "parent"
parent_agent.sub_agents = [current_agent, remote_peer]
current_agent.parent_agent = parent_agent
inv_ctx = mock.MagicMock()
inv_ctx.agent = current_agent
tool_context = mock.MagicMock()
tool_context._invocation_context = inv_ctx
tool = TransferToAgentTool(
agent_names=["remote_peer"],
)
result = bigquery_agent_analytics_plugin._get_tool_origin(
tool,
tool_args={"agent_name": "remote_peer"},
tool_context=tool_context,
)
assert result == "TRANSFER_A2A"
def test_transfer_mixed_targets_classifies_per_call(self):
"""A single TransferToAgentTool with mixed targets classifies per call."""
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
try:
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
except ImportError:
pytest.skip("A2A agent not available")
remote_agent = mock.MagicMock(spec=RemoteA2aAgent)
remote_agent.name = "remote_a2a"
local_agent = mock.MagicMock()
local_agent.name = "local_sub"
current_agent = mock.MagicMock()
current_agent.name = "root"
current_agent.sub_agents = [remote_agent, local_agent]
current_agent.parent_agent = None
inv_ctx = mock.MagicMock()
inv_ctx.agent = current_agent
tool_context = mock.MagicMock()
tool_context._invocation_context = inv_ctx
tool = TransferToAgentTool(
agent_names=["remote_a2a", "local_sub"],
)
# Transfer to remote target → TRANSFER_A2A
result = bigquery_agent_analytics_plugin._get_tool_origin(
tool,
tool_args={"agent_name": "remote_a2a"},
tool_context=tool_context,
)
assert result == "TRANSFER_A2A"
# Transfer to local target → TRANSFER_AGENT
result = bigquery_agent_analytics_plugin._get_tool_origin(
tool,
tool_args={"agent_name": "local_sub"},
tool_context=tool_context,
)
assert result == "TRANSFER_AGENT"
@pytest.mark.asyncio
async def test_tool_error_callback_classifies_a2a_transfer(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
"""on_tool_error_callback produces TRANSFER_A2A for RemoteA2aAgent."""
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
try:
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
except ImportError:
pytest.skip("A2A agent not available")
remote_agent = mock.MagicMock(spec=RemoteA2aAgent)
remote_agent.name = "remote_a2a"
mock_agent = mock.MagicMock(spec=base_agent.BaseAgent)
mock_agent.name = "root"
mock_agent.instruction = ""
mock_agent.sub_agents = [remote_agent]
mock_agent.parent_agent = None
mock_s = mock.create_autospec(
session_lib.Session, instance=True, spec_set=True
)
type(mock_s).id = mock.PropertyMock(return_value="sess-1")
type(mock_s).user_id = mock.PropertyMock(return_value="user-1")
type(mock_s).app_name = mock.PropertyMock(return_value="test_app")
type(mock_s).state = mock.PropertyMock(return_value={})
inv_ctx = InvocationContext(
agent=mock_agent,
session=mock_s,
invocation_id="inv-err",
session_service=mock.create_autospec(
base_session_service_lib.BaseSessionService,
instance=True,
spec_set=True,
),
plugin_manager=mock.create_autospec(
plugin_manager_lib.PluginManager,
instance=True,
spec_set=True,
),
)
tool_ctx = tool_context_lib.ToolContext(invocation_context=inv_ctx)
tool = TransferToAgentTool(agent_names=["remote_a2a"])
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID
) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_ctx, "tool")
await plugin.on_tool_error_callback(
tool=tool,
tool_args={"agent_name": "remote_a2a"},
tool_context=tool_ctx,
error=RuntimeError("connection refused"),
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
assert len(rows) == 1
assert rows[0]["event_type"] == "TOOL_ERROR"
content = json.loads(rows[0]["content"])
assert content["tool_origin"] == "TRANSFER_A2A"
def test_mcp_tool_returns_mcp(self):
try:
from google.adk.tools.mcp_tool.mcp_tool import McpTool
except ImportError:
pytest.skip("MCP not installed")
tool = McpTool.__new__(McpTool)
result = bigquery_agent_analytics_plugin._get_tool_origin(tool)
assert result == "MCP"
def test_a2a_agent_tool_returns_a2a(self):
from google.adk.tools.agent_tool import AgentTool
try:
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
except ImportError:
pytest.skip("A2A agent not available")
remote_agent = mock.MagicMock(spec=RemoteA2aAgent)
remote_agent.name = "remote"
remote_agent.description = "remote a2a agent"
tool = AgentTool.__new__(AgentTool)
tool.agent = remote_agent
tool._name = "remote"
result = bigquery_agent_analytics_plugin._get_tool_origin(tool)
assert result == "A2A"
def test_unknown_tool_returns_unknown(self):
tool = mock.MagicMock(spec=base_tool_lib.BaseTool)
tool.name = "mystery"
result = bigquery_agent_analytics_plugin._get_tool_origin(tool)
assert result == "UNKNOWN"
class TestHITLTracing:
"""Tests for HITL-specific event emission via on_event_callback.
HITL events (``adk_request_credential``, ``adk_request_confirmation``,
``adk_request_input``) are synthetic function calls injected by the
framework — they never pass through ``before_tool_callback`` /
``after_tool_callback``. Detection therefore lives in
``on_event_callback``, which inspects the event stream for these
function calls and their corresponding function responses.
"""
def _make_fc_event(self, fc_name, args=None):
"""Build a mock Event containing a function call."""
event = mock.MagicMock(spec=event_lib.Event)
fc = types.FunctionCall(name=fc_name, args=args or {})
part = types.Part(function_call=fc)
event.content = types.Content(role="model", parts=[part])
event.actions = event_actions_lib.EventActions()
return event
def _make_fr_event(self, fr_name, response=None):
"""Build a mock Event containing a function response."""
event = mock.MagicMock(spec=event_lib.Event)
fr = types.FunctionResponse(name=fr_name, response=response or {})
part = types.Part(function_response=fr)
event.content = types.Content(role="user", parts=[part])
event.actions = event_actions_lib.EventActions()
return event
@pytest.mark.asyncio
async def test_hitl_confirmation_emits_additional_event(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
event = self._make_fc_event("adk_request_confirmation", {"confirm": True})
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
event_types = [r["event_type"] for r in rows]
assert "HITL_CONFIRMATION_REQUEST" in event_types
@pytest.mark.asyncio
async def test_hitl_credential_emits_additional_event(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
event = self._make_fc_event("adk_request_credential", {"auth": "oauth2"})
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
event_types = [r["event_type"] for r in rows]
assert "HITL_CREDENTIAL_REQUEST" in event_types
@pytest.mark.asyncio
async def test_hitl_completion_emits_additional_event(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
event = self._make_fr_event("adk_request_confirmation", {"confirmed": True})
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
event_types = [r["event_type"] for r in rows]
assert "HITL_CONFIRMATION_REQUEST_COMPLETED" in event_types
@pytest.mark.asyncio
async def test_regular_tool_no_hitl_event(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
event = self._make_fc_event("regular_tool", {"x": 1})
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
# No HITL events should be emitted for non-HITL function calls.
# on_event_callback only logs STATE_DELTA and HITL events; a regular
# function call produces neither.
assert mock_write_client.append_rows.call_count == 0
# ==============================================================================
# TEST CLASS: Span Hierarchy Isolation (Issue #4561)
# ==============================================================================
class TestSpanHierarchyIsolation:
"""Regression tests for span hierarchy isolation.
``push_span()`` must NOT attach its span to the ambient OTel context.
If it does, any subsequent ``tracer.start_as_current_span()`` in the
framework (e.g. ``call_llm``, ``execute_tool``) will be incorrectly
re-parented under the plugin's span.
"""
def test_push_span_does_not_change_ambient_context(self, callback_context):
"""push_span must not mutate the current OTel span."""
span_before = trace.get_current_span()
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "test_span"
)
span_after = trace.get_current_span()
assert span_after is span_before
# Cleanup
bigquery_agent_analytics_plugin.TraceManager.pop_span()
def test_attach_current_span_does_not_change_ambient_context(
self, callback_context
):
"""attach_current_span must not mutate the current OTel span."""
span_before = trace.get_current_span()
bigquery_agent_analytics_plugin.TraceManager.attach_current_span(
callback_context
)
span_after = trace.get_current_span()
assert span_after is span_before
# Cleanup
bigquery_agent_analytics_plugin.TraceManager.pop_span()
def test_pop_span_does_not_change_ambient_context(self, callback_context):
"""pop_span must not mutate the current OTel span."""
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "test_span"
)
span_before = trace.get_current_span()
bigquery_agent_analytics_plugin.TraceManager.pop_span()
span_after = trace.get_current_span()
assert span_after is span_before
def test_push_span_with_real_tracer_does_not_reparent(self, callback_context):
"""With a real OTel tracer, plugin spans must not become parents
of subsequently created framework spans.
"""
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
exporter = InMemorySpanExporter()
provider = TracerProvider()
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
provider.add_span_processor(SimpleSpanProcessor(exporter))
framework_tracer = provider.get_tracer("test-framework")
# Simulate: plugin pushes a span BEFORE the framework span
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "llm_request"
)
# Framework creates its own span via start_as_current_span
with framework_tracer.start_as_current_span("call_llm") as fw_span:
fw_context = fw_span.get_span_context()
# Pop the plugin span
bigquery_agent_analytics_plugin.TraceManager.pop_span()
provider.shutdown()
# Verify the framework span was NOT re-parented under the
# plugin's llm_request span
finished = exporter.get_finished_spans()
call_llm_spans = [s for s in finished if s.name == "call_llm"]
assert len(call_llm_spans) == 1
fw_finished = call_llm_spans[0]
# The framework span's parent should NOT be the plugin's
# llm_request span. With the fix, the plugin never
# attaches to the ambient context, so ``call_llm`` will
# have whatever parent existed before (None in this test).
assert fw_finished.parent is None
def test_multiple_push_pop_cycles_leave_context_clean(self, callback_context):
"""Multiple push/pop cycles must not leak context changes."""
original_span = trace.get_current_span()
for _ in range(5):
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "cycle_span"
)
bigquery_agent_analytics_plugin.TraceManager.pop_span()
assert trace.get_current_span() is original_span
# ==============================================================================
# TEST CLASS: End-to-End HITL Tracing via Runner
# ==============================================================================
def _hitl_my_action(
tool_context: tool_context_lib.ToolContext,
) -> dict[str, str]:
"""Tool function used by HITL end-to-end tests."""
return {"result": f"confirmed={tool_context.tool_confirmation.confirmed}"}
class TestHITLTracingEndToEnd:
"""End-to-end tests that run the full Runner + Plugin pipeline with
``FunctionTool(require_confirmation=True)`` and verify that HITL events
are logged alongside normal TOOL_* events in the BQ analytics plugin.
"""
@pytest.fixture
def _mock_bq_infra(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""Bundle all BQ mocking fixtures."""
yield mock_write_client
@pytest.mark.asyncio
async def test_confirmation_flow_emits_hitl_events(
self,
_mock_bq_infra,
dummy_arrow_schema,
):
"""Full Runner pipeline: tool with require_confirmation emits
HITL_CONFIRMATION_REQUEST and HITL_CONFIRMATION_REQUEST_COMPLETED.
"""
from google.adk.flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
from google.adk.tools.function_tool import FunctionTool
from google.genai.types import FunctionCall
from google.genai.types import FunctionResponse
from google.genai.types import Part
from .. import testing_utils
mock_write_client = _mock_bq_infra
tool = FunctionTool(func=_hitl_my_action, require_confirmation=True)
# -- Mock LLM: first response calls the tool, second is final text --
llm_responses = [
testing_utils.LlmResponse(
content=testing_utils.ModelContent(
parts=[
Part(function_call=FunctionCall(name=tool.name, args={}))
]
)
),
testing_utils.LlmResponse(
content=testing_utils.ModelContent(
parts=[Part(text="Done, action confirmed.")]
)
),
]
mock_model = testing_utils.MockModel(responses=llm_responses)
# -- Build the plugin --
bq_plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
await bq_plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
# -- Build agent + runner WITH the plugin --
from google.adk.agents.llm_agent import LlmAgent
agent = LlmAgent(name="hitl_agent", model=mock_model, tools=[tool])
runner = testing_utils.InMemoryRunner(root_agent=agent, plugins=[bq_plugin])
try:
# -- Turn 1: user query → LLM calls tool → HITL pause --
events_turn1 = await runner.run_async(
testing_utils.UserContent("run my_action")
)
# Find the adk_request_confirmation function call
confirmation_fc_id = None
for ev in events_turn1:
if ev.content and ev.content.parts:
for part in ev.content.parts:
if (
hasattr(part, "function_call")
and part.function_call
and part.function_call.name
== REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
):
confirmation_fc_id = part.function_call.id
break
if confirmation_fc_id:
break
assert (
confirmation_fc_id is not None
), "Expected adk_request_confirmation function call in turn 1"
# -- Turn 2: user sends confirmation → tool re-executes --
user_confirmation = testing_utils.UserContent(
Part(
function_response=FunctionResponse(
id=confirmation_fc_id,
name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
response={"confirmed": True},
)
)
)
events_turn2 = await runner.run_async(user_confirmation)
# -- Deterministically wait for the async BQ writer to drain --
await bq_plugin.flush()
# -- Collect all BQ rows --
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
event_types = [r["event_type"] for r in rows]
# -- Verify standard events are present --
assert "TOOL_STARTING" in event_types
assert "TOOL_COMPLETED" in event_types
# -- Verify HITL-specific events are present --
assert (
"HITL_CONFIRMATION_REQUEST" in event_types
), f"Expected HITL_CONFIRMATION_REQUEST in {event_types}"
assert (
"HITL_CONFIRMATION_REQUEST_COMPLETED" in event_types
), f"Expected HITL_CONFIRMATION_REQUEST_COMPLETED in {event_types}"
# -- Verify HITL events have correct tool name in content --
hitl_rows = [r for r in rows if r["event_type"].startswith("HITL_")]
for row in hitl_rows:
content = json.loads(row["content"]) if row["content"] else {}
assert content.get("tool") == "adk_request_confirmation", (
"HITL event should reference 'adk_request_confirmation',"
f" got {content.get('tool')}"
)
finally:
await bq_plugin.shutdown()
@pytest.mark.asyncio
async def test_regular_tool_does_not_emit_hitl_events(
self,
_mock_bq_infra,
dummy_arrow_schema,
):
"""A tool WITHOUT require_confirmation should not produce HITL events."""
from google.adk.tools.function_tool import FunctionTool
from google.genai.types import FunctionCall
from google.genai.types import Part
from .. import testing_utils
mock_write_client = _mock_bq_infra
def regular_tool() -> str:
return "done"
tool = FunctionTool(func=regular_tool)
llm_responses = [
testing_utils.LlmResponse(
content=testing_utils.ModelContent(
parts=[
Part(function_call=FunctionCall(name=tool.name, args={}))
]
)
),
testing_utils.LlmResponse(
content=testing_utils.ModelContent(parts=[Part(text="All done.")])
),
]
mock_model = testing_utils.MockModel(responses=llm_responses)
bq_plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
await bq_plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
from google.adk.agents.llm_agent import LlmAgent
agent = LlmAgent(name="regular_agent", model=mock_model, tools=[tool])
runner = testing_utils.InMemoryRunner(root_agent=agent, plugins=[bq_plugin])
try:
await runner.run_async(testing_utils.UserContent("run regular_tool"))
await bq_plugin.flush()
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
event_types = [r["event_type"] for r in rows]
# Standard tool events should be present
assert "TOOL_STARTING" in event_types
assert "TOOL_COMPLETED" in event_types
# No HITL events
hitl_events = [et for et in event_types if et.startswith("HITL_")]
assert (
hitl_events == []
), f"Expected no HITL events for regular tool, got {hitl_events}"
finally:
await bq_plugin.shutdown()
# ==============================================================================
# Fork-Safety Tests
# ==============================================================================
class TestForkSafety:
"""Tests for fork-safety via PID tracking."""
def _make_plugin(self):
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=config,
)
return plugin
@pytest.mark.asyncio
async def test_pid_change_triggers_reinit(
self, mock_auth_default, mock_bq_client, mock_write_client
):
"""Simulating a fork by changing _init_pid forces re-init."""
plugin = self._make_plugin()
await plugin._ensure_started()
assert plugin._started is True
# Simulate a fork: set _init_pid to a stale value
plugin._init_pid = -1
assert plugin._started is True # still True before check
# _ensure_started should detect PID mismatch and reset
await plugin._ensure_started()
# After reset + re-init, _init_pid should match current
assert plugin._init_pid == os.getpid()
assert plugin._started is True
await plugin.shutdown()
@pytest.mark.asyncio
async def test_pid_unchanged_skips_reset(
self, mock_auth_default, mock_bq_client, mock_write_client
):
"""Same PID should not trigger a reset."""
plugin = self._make_plugin()
await plugin._ensure_started()
# Save references to verify they are not recreated
original_client = plugin.client
original_parser = plugin.parser
await plugin._ensure_started()
assert plugin.client is original_client
assert plugin.parser is original_parser
await plugin.shutdown()
def test_reset_runtime_state_clears_fields(self):
"""_reset_runtime_state clears all runtime fields."""
plugin = self._make_plugin()
# Fake some runtime state
plugin._started = True
plugin._is_shutting_down = True
plugin.client = mock.MagicMock()
plugin._loop_state_by_loop = {"fake": "state"}
plugin._write_stream_name = "some/stream"
plugin._executor = mock.MagicMock()
plugin.offloader = mock.MagicMock()
plugin.parser = mock.MagicMock()
plugin._setup_lock = mock.MagicMock()
# Keep pure-data fields
plugin._schema = ["kept"]
plugin.arrow_schema = "kept_arrow"
plugin._reset_runtime_state()
assert plugin._started is False
assert plugin._is_shutting_down is False
assert plugin.client is None
assert plugin._loop_state_by_loop == {}
assert plugin._write_stream_name is None
assert plugin._executor is None
assert plugin.offloader is None
assert plugin.parser is None
assert plugin._setup_lock is None
# Pure-data fields are preserved
assert plugin._schema == ["kept"]
assert plugin.arrow_schema == "kept_arrow"
assert plugin._init_pid == os.getpid()
def test_getstate_resets_pid(self):
"""Pickle state should have _init_pid = 0 to force re-init."""
plugin = self._make_plugin()
state = plugin.__getstate__()
assert state["_init_pid"] == 0
assert state["_started"] is False
@pytest.mark.asyncio
async def test_unpickle_legacy_state_missing_init_pid(
self, mock_auth_default, mock_bq_client, mock_write_client
):
"""Unpickling state from older code without _init_pid should not crash."""
plugin = self._make_plugin()
state = plugin.__getstate__()
# Simulate legacy pickle state that lacks _init_pid entirely
del state["_init_pid"]
new_plugin = (
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin.__new__(
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin
)
)
new_plugin.__setstate__(state)
# _init_pid should be backfilled to 0, triggering re-init
assert new_plugin._init_pid == 0
# _ensure_started should not raise AttributeError
await new_plugin._ensure_started()
assert new_plugin._started is True
await new_plugin.shutdown()
class TestForkGrpcSafety:
"""Tests for gRPC fork safety enhancements."""
def _make_plugin(self):
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
return bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=config,
)
def test_grpc_fork_env_var_set(self):
"""GRPC_ENABLE_FORK_SUPPORT should be '1' after import."""
assert os.environ.get("GRPC_ENABLE_FORK_SUPPORT") == "1"
def test_register_at_fork_resets_all_instances(self):
"""_after_fork_in_child resets all living plugin instances."""
p1 = self._make_plugin()
p2 = self._make_plugin()
p1._started = True
p2._started = True
p1._init_pid = -1
p2._init_pid = -1
bigquery_agent_analytics_plugin._after_fork_in_child()
assert p1._started is False
assert p2._started is False
assert p1._init_pid == os.getpid()
assert p2._init_pid == os.getpid()
def test_dead_plugin_removed_from_live_set(self):
"""WeakSet should not hold dead plugin references."""
p = self._make_plugin()
assert p in bigquery_agent_analytics_plugin._LIVE_PLUGINS
pid = id(p)
del p
# After deletion, the WeakSet should no longer contain it.
for alive in bigquery_agent_analytics_plugin._LIVE_PLUGINS:
assert id(alive) != pid
def test_reset_closes_inherited_sync_transports(self):
"""_reset_runtime_state closes inherited sync gRPC channels."""
plugin = self._make_plugin()
mock_channel = mock.MagicMock()
mock_channel.close.return_value = None # sync close
mock_transport = mock.MagicMock()
mock_transport._grpc_channel = mock_channel
mock_wc = mock.MagicMock()
mock_wc.transport = mock_transport
mock_loop_state = mock.MagicMock()
mock_loop_state.write_client = mock_wc
plugin._loop_state_by_loop = {mock.MagicMock(): mock_loop_state}
plugin._init_pid = -1
plugin._reset_runtime_state()
mock_channel.close.assert_called_once()
def test_reset_discards_async_channel_close_coroutine(self):
"""Async channel close() returns a coroutine; must not warn."""
import warnings
plugin = self._make_plugin()
async def _async_close():
pass
mock_channel = mock.MagicMock()
mock_channel.close.return_value = _async_close()
mock_transport = mock.MagicMock()
mock_transport._grpc_channel = mock_channel
mock_wc = mock.MagicMock()
mock_wc.transport = mock_transport
mock_loop_state = mock.MagicMock()
mock_loop_state.write_client = mock_wc
plugin._loop_state_by_loop = {mock.MagicMock(): mock_loop_state}
plugin._init_pid = -1
with warnings.catch_warnings():
warnings.simplefilter("error", RuntimeWarning)
# Must not raise RuntimeWarning for unawaited coroutine
plugin._reset_runtime_state()
mock_channel.close.assert_called_once()
def test_transport_close_exception_swallowed(self):
"""close() raising should not prevent reset from completing."""
plugin = self._make_plugin()
mock_channel = mock.MagicMock()
mock_channel.close.side_effect = RuntimeError("broken channel")
mock_transport = mock.MagicMock()
mock_transport._grpc_channel = mock_channel
mock_wc = mock.MagicMock()
mock_wc.transport = mock_transport
mock_loop_state = mock.MagicMock()
mock_loop_state.write_client = mock_wc
plugin._loop_state_by_loop = {mock.MagicMock(): mock_loop_state}
plugin._init_pid = -1
# Should not raise
plugin._reset_runtime_state()
assert plugin._started is False
assert plugin._loop_state_by_loop == {}
def test_reset_logs_fork_warning(self):
"""_reset_runtime_state logs a warning with 'Fork detected'."""
plugin = self._make_plugin()
plugin._init_pid = -1
with mock.patch.object(
bigquery_agent_analytics_plugin.logger, "warning"
) as mock_warn:
plugin._reset_runtime_state()
mock_warn.assert_called_once()
assert "Fork detected" in mock_warn.call_args[0][0]
# ==============================================================================
# Analytics Views Tests
# ==============================================================================
class TestAnalyticsViews:
"""Tests for auto-created per-event-type BigQuery views."""
def _make_plugin(self, create_views=True, view_prefix="v", table_id=TABLE_ID):
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
create_views=create_views,
view_prefix=view_prefix,
)
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=table_id,
config=config,
)
plugin.client = mock.MagicMock()
plugin.full_table_id = f"{PROJECT_ID}.{DATASET_ID}.{table_id}"
plugin._schema = bigquery_agent_analytics_plugin._get_events_schema()
return plugin
def test_views_created_on_new_table(self):
"""NotFound path creates all views."""
plugin = self._make_plugin(create_views=True)
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
mock_query_job = mock.MagicMock()
plugin.client.query.return_value = mock_query_job
plugin._ensure_schema_exists()
expected_count = len(bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS)
assert plugin.client.query.call_count == expected_count
def test_views_created_for_existing_table(self):
"""Existing table path also creates views."""
plugin = self._make_plugin(create_views=True)
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = plugin._schema
existing.labels = {
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY: (
bigquery_agent_analytics_plugin._SCHEMA_VERSION
),
}
plugin.client.get_table.return_value = existing
mock_query_job = mock.MagicMock()
plugin.client.query.return_value = mock_query_job
plugin._ensure_schema_exists()
expected_count = len(bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS)
assert plugin.client.query.call_count == expected_count
def test_views_not_created_when_disabled(self):
"""create_views=False skips view creation."""
plugin = self._make_plugin(create_views=False)
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
plugin._ensure_schema_exists()
plugin.client.query.assert_not_called()
def test_view_creation_error_logged_not_raised(self):
"""Errors during view creation don't crash the plugin."""
plugin = self._make_plugin(create_views=True)
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
plugin.client.query.side_effect = Exception("BQ error")
# Should not raise
plugin._ensure_schema_exists()
# Verify it tried to create views (and failed gracefully)
assert plugin.client.query.call_count > 0
def test_view_sql_contains_correct_event_filter(self):
"""Each SQL has correct WHERE clause and view name."""
plugin = self._make_plugin(create_views=True)
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
mock_query_job = mock.MagicMock()
plugin.client.query.return_value = mock_query_job
plugin._ensure_schema_exists()
calls = plugin.client.query.call_args_list
for call in calls:
sql = call[0][0]
# Each SQL should have CREATE OR REPLACE VIEW
assert "CREATE OR REPLACE VIEW" in sql
# Each SQL should filter by event_type
assert "WHERE" in sql
assert "event_type = " in sql
# View name should start with v_
assert ".v_" in sql
# Verify specific views exist
all_sql = " ".join(c[0][0] for c in calls)
for event_type in bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS:
view_name = "v_" + event_type.lower()
assert view_name in all_sql, f"View {view_name} not found in SQL"
def test_error_views_contain_traceback_column(self):
"""AGENT_ERROR and INVOCATION_ERROR views include error_traceback."""
plugin = self._make_plugin(create_views=True)
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
mock_query_job = mock.MagicMock()
plugin.client.query.return_value = mock_query_job
plugin._ensure_schema_exists()
calls = plugin.client.query.call_args_list
all_sqls = {c[0][0] for c in calls}
agent_error_sqls = [s for s in all_sqls if "v_agent_error" in s]
assert len(agent_error_sqls) == 1
assert "error_traceback" in agent_error_sqls[0]
assert "total_ms" in agent_error_sqls[0]
inv_error_sqls = [s for s in all_sqls if "v_invocation_error" in s]
assert len(inv_error_sqls) == 1
assert "error_traceback" in inv_error_sqls[0]
def test_llm_response_view_exposes_token_usage_columns(self):
"""LLM_RESPONSE view surfaces cached/thinking/tool-use token columns.
These are read from the full ``usage_metadata`` proto that is already
logged to ``attributes.usage_metadata``, so they are sourced from
``attributes`` rather than the ``content.usage`` summary.
"""
plugin = self._make_plugin(create_views=True)
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
plugin.client.query.return_value = mock.MagicMock()
plugin._ensure_schema_exists()
all_sql = " ".join(c[0][0] for c in plugin.client.query.call_args_list)
assert "usage_cached_tokens" in all_sql
assert "usage_thinking_tokens" in all_sql
assert "usage_tool_use_tokens" in all_sql
assert "$.usage_metadata.thoughts_token_count" in all_sql
assert "$.usage_metadata.tool_use_prompt_token_count" in all_sql
def test_config_create_views_default_true(self):
"""Config create_views defaults to True."""
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
assert config.create_views is True
assert config.view_prefix == "v"
@pytest.mark.asyncio
async def test_create_analytics_views_ensures_started(
self, mock_auth_default, mock_bq_client, mock_write_client
):
"""Public create_analytics_views() initializes plugin first."""
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
assert plugin._started is False
await plugin.create_analytics_views()
# Plugin should be started after the call
assert plugin._started is True
# Views should have been created (query called)
expected_count = len(bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS)
# _ensure_schema_exists also creates views, so total calls
# = schema-creation views + explicit views
assert mock_bq_client.query.call_count >= expected_count
await plugin.shutdown()
def test_views_not_created_after_table_creation_failure(self):
"""View creation is skipped when create_table raises a non-Conflict error."""
plugin = self._make_plugin(create_views=True)
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
plugin.client.create_table.side_effect = RuntimeError("BQ down")
plugin._ensure_schema_exists()
# Views should NOT be attempted since table creation failed
plugin.client.query.assert_not_called()
@pytest.mark.asyncio
async def test_create_analytics_views_raises_on_startup_failure(
self, mock_auth_default, mock_write_client
):
"""create_analytics_views() raises if plugin init fails."""
# Make the BQ Client constructor raise so _lazy_setup fails
# before _started is set to True.
with mock.patch.object(
bigquery, "Client", side_effect=Exception("client boom")
):
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
with pytest.raises(
RuntimeError, match="Plugin initialization failed"
) as exc_info:
await plugin.create_analytics_views()
# Root cause should be chained for debuggability
assert exc_info.value.__cause__ is not None
assert "client boom" in str(exc_info.value.__cause__)
def test_custom_view_prefix(self):
"""Custom view_prefix namespaces view names."""
plugin = self._make_plugin(view_prefix="v_staging")
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
mock_query_job = mock.MagicMock()
plugin.client.query.return_value = mock_query_job
plugin._ensure_schema_exists()
calls = plugin.client.query.call_args_list
all_sql = " ".join(c[0][0] for c in calls)
# All views should use the custom prefix
for event_type in bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS:
expected_name = "v_staging_" + event_type.lower()
assert expected_name in all_sql, f"View {expected_name} not found in SQL"
# Default prefix should NOT appear
assert ".v_llm_request" not in all_sql
def test_default_view_prefix_preserves_names(self):
"""Default view_prefix='v' produces the same names as before."""
plugin = self._make_plugin() # default view_prefix="v"
plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found")
mock_query_job = mock.MagicMock()
plugin.client.query.return_value = mock_query_job
plugin._ensure_schema_exists()
calls = plugin.client.query.call_args_list
all_sql = " ".join(c[0][0] for c in calls)
for event_type in bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS:
view_name = "v_" + event_type.lower()
assert view_name in all_sql
def test_distinct_tables_and_prefixes_no_collision(self):
"""Two plugins targeting different tables produce disjoint views."""
plugin_a = self._make_plugin(
table_id="agent_events_prod", view_prefix="v_prod"
)
plugin_b = self._make_plugin(
table_id="agent_events_staging", view_prefix="v_staging"
)
for plugin in (plugin_a, plugin_b):
plugin.client.get_table.side_effect = cloud_exceptions.NotFound(
"not found"
)
mock_query_job = mock.MagicMock()
plugin.client.query.return_value = mock_query_job
plugin._ensure_schema_exists()
sql_a = " ".join(c[0][0] for c in plugin_a.client.query.call_args_list)
sql_b = " ".join(c[0][0] for c in plugin_b.client.query.call_args_list)
# View names use their own prefix
assert "v_prod_llm_request" in sql_a
assert "v_staging_llm_request" in sql_b
# No cross-contamination
assert "v_staging_" not in sql_a
assert "v_prod_" not in sql_b
# FROM clauses point at the correct table
assert "agent_events_prod" in sql_a
assert "agent_events_staging" not in sql_a
assert "agent_events_staging" in sql_b
assert "agent_events_prod" not in sql_b
def test_empty_view_prefix_raises(self):
"""Empty view_prefix is rejected at init."""
with pytest.raises(ValueError, match="view_prefix"):
self._make_plugin(view_prefix="")
# ==============================================================================
# Trace-ID Continuity Tests (Issue #4645)
# ==============================================================================
class TestTraceIdContinuity:
"""Tests for trace_id continuity across all events in an invocation.
When there is no ambient OTel span (e.g. Agent Engine, custom runners),
early events (USER_MESSAGE_RECEIVED, INVOCATION_STARTING) used to fall
back to ``invocation_id`` while AGENT_STARTING got a new OTel hex
trace_id from ``push_span()``. The ``ensure_invocation_span()`` fix
guarantees a root span is always on the stack before any events fire.
"""
@pytest.mark.asyncio
async def test_trace_id_continuity_no_ambient_span(self, callback_context):
"""All events share one trace_id when no ambient OTel span exists.
Simulates the #4645 scenario: OTel IS configured (real TracerProvider)
but the Runner's ambient span is NOT present (e.g. Agent Engine,
custom runners).
"""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
# Wire a real TracerProvider with an in-memory exporter so we can
# also assert the plugin path does NOT export anything through it.
# (push_span no longer creates OTel spans — see _SpanRecord; the
# exporter is here as a regression guard, not a span source.)
exporter = InMemorySpanExporter()
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
real_tracer = provider.get_tracer("test-plugin")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
# Reset the span records contextvar for a clean invocation.
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
# No ambient OTel span — we do NOT start_as_current_span.
ambient = trace.get_current_span()
assert not ambient.get_span_context().is_valid
# ensure_invocation_span should push a new span.
TM.ensure_invocation_span(callback_context)
trace_id_early = TM.get_trace_id(callback_context)
assert trace_id_early is not None
# Should NOT fall back to invocation_id — it should be
# a 32-char hex OTel trace_id.
assert trace_id_early != callback_context.invocation_id
assert len(trace_id_early) == 32
# Simulate agent callback: push_span("agent")
TM.push_span(callback_context, "agent")
trace_id_agent = TM.get_trace_id(callback_context)
# Both trace_ids must be identical.
assert trace_id_early == trace_id_agent
# Cleanup
TM.pop_span() # agent
TM.pop_span() # invocation
provider.shutdown()
@pytest.mark.asyncio
async def test_invocation_completed_trace_continuity_no_ambient(
self, callback_context
):
"""INVOCATION_COMPLETED must share trace_id with earlier events.
Reproduces the completion-event fracture: after_run_callback pops
the invocation span, then _log_event would resolve trace_id via
the fallback to invocation_id. The trace_id_override ensures the
completion event keeps the same trace_id.
"""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
exporter = InMemorySpanExporter()
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
real_tracer = provider.get_tracer("test-plugin")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
# Reset for a clean invocation; no ambient span.
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
assert not trace.get_current_span().get_span_context().is_valid
# --- Simulate the full callback lifecycle ---
# 1. before_run / on_user_message: ensure invocation span
TM.ensure_invocation_span(callback_context)
trace_id_start = TM.get_trace_id(callback_context)
# 2. before_agent: push agent span
TM.push_span(callback_context, "agent")
assert TM.get_trace_id(callback_context) == trace_id_start
# 3. after_agent: pop agent span
TM.pop_span()
# 4. after_run: capture trace_id THEN pop invocation span
trace_id_before_pop = TM.get_trace_id(callback_context)
assert trace_id_before_pop == trace_id_start
TM.pop_span()
# After popping, get_trace_id falls back to invocation_id
trace_id_after_pop = TM.get_trace_id(callback_context)
assert trace_id_after_pop == callback_context.invocation_id
# The trace_id_override preserves continuity
assert trace_id_before_pop == trace_id_start
assert trace_id_before_pop != trace_id_after_pop
provider.shutdown()
@pytest.mark.asyncio
async def test_callbacks_emit_same_trace_id_no_ambient(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
callback_context,
mock_agent,
dummy_arrow_schema,
):
"""Full callback path: all emitted rows share one trace_id.
Exercises the real before_run → before_agent → after_agent →
after_run callback chain via the plugin instance, then checks
every emitted BQ row has the same trace_id.
"""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
exporter = InMemorySpanExporter()
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
real_tracer = provider.get_tracer("test-plugin")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
# Reset span records for a clean invocation.
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
# No ambient span — simulates Agent Engine / custom runner.
assert not trace.get_current_span().get_span_context().is_valid
# Run the full callback lifecycle.
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
await bq_plugin_inst.before_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await bq_plugin_inst.after_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await bq_plugin_inst.after_run_callback(
invocation_context=invocation_context
)
await asyncio.sleep(0.01)
# Collect all emitted rows.
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
event_types = [r["event_type"] for r in rows]
assert "INVOCATION_STARTING" in event_types
assert "INVOCATION_COMPLETED" in event_types
# Every row must share the same trace_id.
trace_ids = {r["trace_id"] for r in rows}
assert len(trace_ids) == 1, (
"Expected 1 unique trace_id across all events, got"
f" {len(trace_ids)}: {trace_ids}"
)
# Should be a 32-char hex OTel trace, not the invocation_id.
sole_trace_id = trace_ids.pop()
assert sole_trace_id != invocation_context.invocation_id
assert len(sole_trace_id) == 32
provider.shutdown()
@pytest.mark.asyncio
async def test_trace_id_continuity_with_ambient_span(self, callback_context):
"""All events share one trace_id when an ambient OTel span exists."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
# Set up a real OTel tracer.
exporter = InMemorySpanExporter()
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
# Reset the span records contextvar.
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
with real_tracer.start_as_current_span("runner_invocation"):
ambient = trace.get_current_span()
assert ambient.get_span_context().is_valid
ambient_trace_id = format(ambient.get_span_context().trace_id, "032x")
# ensure_invocation_span should attach the ambient span.
TM.ensure_invocation_span(callback_context)
trace_id_early = TM.get_trace_id(callback_context)
assert trace_id_early == ambient_trace_id
# Simulate agent callback: push_span("agent")
TM.push_span(callback_context, "agent")
trace_id_agent = TM.get_trace_id(callback_context)
assert trace_id_agent == ambient_trace_id
# Cleanup
TM.pop_span() # agent
TM.pop_span() # invocation (attached, not owned)
provider.shutdown()
@pytest.mark.asyncio
async def test_invocation_root_span_isolated_across_turns(
self, callback_context
):
"""Each invocation gets its own root span; turns don't leak."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
exporter = InMemorySpanExporter()
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
# --- Turn 1 ---
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
TM.ensure_invocation_span(callback_context)
trace_id_turn1 = TM.get_trace_id(callback_context)
TM.push_span(callback_context, "agent")
assert TM.get_trace_id(callback_context) == trace_id_turn1
TM.pop_span() # agent
TM.pop_span() # invocation
# After popping, the stack should be empty.
records = bigquery_agent_analytics_plugin._span_records_ctx.get()
assert not records
# --- Turn 2 ---
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
TM.ensure_invocation_span(callback_context)
trace_id_turn2 = TM.get_trace_id(callback_context)
TM.push_span(callback_context, "agent")
assert TM.get_trace_id(callback_context) == trace_id_turn2
TM.pop_span() # agent
TM.pop_span() # invocation
# The two turns must have DIFFERENT trace_ids (different
# root spans).
assert trace_id_turn1 != trace_id_turn2
provider.shutdown()
class TestSpanIdConsistency:
"""Tests that STARTING/COMPLETED event pairs share span IDs.
Span-ID resolution contract:
- When OTel is active: BQ rows use the same trace/span/parent IDs as
Cloud Trace (ambient framework spans). STARTING and COMPLETED events
in the same lifecycle share the same span_id.
- When OTel is not active: BQ rows use the plugin's internal span
stack. STARTING gets the current top-of-stack; COMPLETED gets the
popped span.
"""
@pytest.mark.asyncio
async def test_starting_completed_same_span_with_ambient(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
callback_context,
mock_agent,
dummy_arrow_schema,
):
"""With ambient OTel, STARTING and COMPLETED get the same span_id."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
# Simulate the framework's ambient spans.
with real_tracer.start_as_current_span("invocation"):
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
with real_tracer.start_as_current_span("invoke_agent"):
await bq_plugin_inst.before_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await bq_plugin_inst.after_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await bq_plugin_inst.after_run_callback(
invocation_context=invocation_context
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
agent_starting = [r for r in rows if r["event_type"] == "AGENT_STARTING"]
agent_completed = [
r for r in rows if r["event_type"] == "AGENT_COMPLETED"
]
assert len(agent_starting) == 1
assert len(agent_completed) == 1
# Both events must share the same span_id (the plugin-internal
# agent span pushed by before_agent_callback and popped by
# after_agent_callback). The lifecycle-pair invariant holds
# regardless of whether the id comes from a plugin-minted hex
# string or an ambient OTel span.
assert agent_starting[0]["span_id"] == agent_completed[0]["span_id"]
assert (
agent_starting[0]["parent_span_id"]
== agent_completed[0]["parent_span_id"]
)
provider.shutdown()
@pytest.mark.asyncio
async def test_starting_completed_use_plugin_span_without_ambient(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
callback_context,
mock_agent,
dummy_arrow_schema,
):
"""Without ambient OTel, COMPLETED gets the popped plugin span."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
# No ambient OTel span.
assert not trace.get_current_span().get_span_context().is_valid
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
await bq_plugin_inst.before_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await bq_plugin_inst.after_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await bq_plugin_inst.after_run_callback(
invocation_context=invocation_context
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
agent_starting = [r for r in rows if r["event_type"] == "AGENT_STARTING"]
agent_completed = [
r for r in rows if r["event_type"] == "AGENT_COMPLETED"
]
assert len(agent_starting) == 1
assert len(agent_completed) == 1
# AGENT_STARTING gets the top-of-stack span; AGENT_COMPLETED
# gets the popped span via override — they should match.
assert agent_starting[0]["span_id"] == agent_completed[0]["span_id"]
provider.shutdown()
@pytest.mark.asyncio
async def test_tool_error_captures_span_id(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
callback_context,
dummy_arrow_schema,
):
"""on_tool_error_callback uses the popped span_id (bonus fix)."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
mock_tool = mock.create_autospec(base_tool_lib.BaseTool, instance=True)
type(mock_tool).name = mock.PropertyMock(return_value="my_tool")
tool_ctx = tool_context_lib.ToolContext(
invocation_context=invocation_context
)
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
# No ambient OTel — plugin span stack provides IDs.
assert not trace.get_current_span().get_span_context().is_valid
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
# Push tool span via before_tool_callback
await bq_plugin_inst.before_tool_callback(
tool=mock_tool,
tool_args={"a": 1},
tool_context=tool_ctx,
)
# Error callback should pop the tool span and use its ID
await bq_plugin_inst.on_tool_error_callback(
tool=mock_tool,
tool_args={"a": 1},
tool_context=tool_ctx,
error=RuntimeError("boom"),
)
await bq_plugin_inst.after_run_callback(
invocation_context=invocation_context
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
tool_starting = [r for r in rows if r["event_type"] == "TOOL_STARTING"]
tool_error = [r for r in rows if r["event_type"] == "TOOL_ERROR"]
assert len(tool_starting) == 1
assert len(tool_error) == 1
# The TOOL_ERROR event must have the same span_id as
# TOOL_STARTING (both correspond to the same tool span).
assert tool_starting[0]["span_id"] == tool_error[0]["span_id"]
assert tool_error[0]["span_id"] is not None
provider.shutdown()
class TestStackLeakSafety:
"""Tests for stack leak safety (P2).
Ensures the plugin's internal span stack doesn't leak records
across invocations when after_run_callback is skipped.
"""
def test_ensure_invocation_span_clears_stale_records(self, callback_context):
"""Pre-populated stack from a different invocation is cleared."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
# Simulate stale records from incomplete previous invocation.
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
# Mark the stale records as belonging to a different invocation.
bigquery_agent_analytics_plugin._active_invocation_id_ctx.set(
"old-inv-stale"
)
TM.push_span(callback_context, "stale-invocation")
TM.push_span(callback_context, "stale-agent")
stale_records = bigquery_agent_analytics_plugin._span_records_ctx.get()
assert len(stale_records) == 2
# ensure_invocation_span with the *current* invocation_id should
# detect the mismatch, clear stale records, and re-init.
TM.ensure_invocation_span(callback_context)
records = bigquery_agent_analytics_plugin._span_records_ctx.get()
# Should have exactly 1 fresh entry (the new invocation span).
assert len(records) == 1
# The fresh span should NOT be one of the stale ones.
assert records[0].span_id != stale_records[0].span_id
assert records[0].span_id != stale_records[1].span_id
provider.shutdown()
def test_clear_stack_does_not_export_spans(self, callback_context):
"""``clear_stack()`` clears the internal records but does NOT
export any OTel spans (issue #94 regression guard).
Pre-fix, ``clear_stack()`` called ``record.span.end()`` for every
owned record, which delivered the now-finished span to whatever
exporter the host had wired — duplicating it next to the
framework's real span in Cloud Trace. Post-fix the plugin owns
no OTel span at all; ``clear_stack()`` only resets the contextvar.
"""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
provider = SdkProvider()
exporter = InMemorySpanExporter()
provider.add_span_processor(SimpleSpanProcessor(exporter))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
TM.push_span(callback_context, "span-a")
TM.push_span(callback_context, "span-b")
records = list(bigquery_agent_analytics_plugin._span_records_ctx.get())
assert all(r.owns_span for r in records)
# No exported spans yet (the plugin never creates any).
assert exporter.get_finished_spans() == ()
TM.clear_stack()
# Stack must be empty after clear.
result = bigquery_agent_analytics_plugin._span_records_ctx.get()
assert result == []
# Still no exported spans — the regression guard for #94.
assert exporter.get_finished_spans() == (), (
"clear_stack() must not export OTel spans; any owned span"
" would surface as a duplicate in Cloud Trace alongside the"
" framework's real spans (issue #94)."
)
provider.shutdown()
@pytest.mark.asyncio
async def test_after_run_callback_clears_remaining_stack(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
callback_context,
mock_agent,
dummy_arrow_schema,
):
"""after_run_callback clears any leftover stack entries."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
# No ambient span.
assert not trace.get_current_span().get_span_context().is_valid
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
# Push an agent span but DON'T pop it (simulate missing
# after_agent_callback due to exception).
await bq_plugin_inst.before_agent_callback(
agent=mock_agent, callback_context=callback_context
)
# Stack now has [invocation, agent].
# after_run_callback should pop invocation + clear remaining.
await bq_plugin_inst.after_run_callback(
invocation_context=invocation_context
)
# Stack must be empty.
records = bigquery_agent_analytics_plugin._span_records_ctx.get()
assert records == []
provider.shutdown()
@pytest.mark.asyncio
async def test_next_invocation_clean_after_incomplete_previous(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
callback_context,
mock_agent,
dummy_arrow_schema,
mock_session,
):
"""Next invocation starts clean even if previous was incomplete."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
bigquery_agent_analytics_plugin._active_invocation_id_ctx.set(None)
# --- Incomplete invocation 1: no after_run_callback ---
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
await bq_plugin_inst.before_agent_callback(
agent=mock_agent, callback_context=callback_context
)
# Skip after_agent and after_run — simulates exception.
stale = bigquery_agent_analytics_plugin._span_records_ctx.get()
assert len(stale) >= 2 # invocation + agent
# --- Invocation 2 with a different invocation_id ---
mock_write_client.append_rows.reset_mock()
inv_ctx_2 = InvocationContext(
agent=mock_agent,
session=mock_session,
invocation_id="inv-NEW-002",
session_service=invocation_context.session_service,
plugin_manager=invocation_context.plugin_manager,
)
await bq_plugin_inst.before_run_callback(invocation_context=inv_ctx_2)
records = bigquery_agent_analytics_plugin._span_records_ctx.get()
# Should have exactly 1 fresh invocation span.
assert len(records) == 1
# Cleanup
await bq_plugin_inst.after_run_callback(invocation_context=inv_ctx_2)
provider.shutdown()
def test_ensure_invocation_span_idempotent_same_invocation(
self, callback_context
):
"""Calling ensure_invocation_span twice in the same invocation is a no-op."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
TM = bigquery_agent_analytics_plugin.TraceManager
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
bigquery_agent_analytics_plugin._active_invocation_id_ctx.set(None)
# First call: creates invocation span.
TM.ensure_invocation_span(callback_context)
records_after_first = list(
bigquery_agent_analytics_plugin._span_records_ctx.get()
)
assert len(records_after_first) == 1
first_span_id = records_after_first[0].span_id
# Second call (same invocation): must be a no-op.
TM.ensure_invocation_span(callback_context)
records_after_second = (
bigquery_agent_analytics_plugin._span_records_ctx.get()
)
assert len(records_after_second) == 1
assert records_after_second[0].span_id == first_span_id
# Cleanup
TM.pop_span()
provider.shutdown()
@pytest.mark.asyncio
async def test_user_message_then_before_run_same_trace_no_ambient(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
callback_context,
mock_agent,
dummy_arrow_schema,
):
"""Regression: on_user_message → before_run must share one trace_id.
Without the invocation-ID guard, the second ensure_invocation_span()
call would clear the stack and create a new root span with a
different trace_id, fracturing USER_MESSAGE_RECEIVED from
INVOCATION_STARTING.
"""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
bigquery_agent_analytics_plugin._active_invocation_id_ctx.set(None)
# No ambient span.
assert not trace.get_current_span().get_span_context().is_valid
user_msg = types.Content(parts=[types.Part(text="hello")], role="user")
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=user_msg,
)
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
await bq_plugin_inst.before_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await bq_plugin_inst.after_agent_callback(
agent=mock_agent, callback_context=callback_context
)
await bq_plugin_inst.after_run_callback(
invocation_context=invocation_context
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
event_types = [r["event_type"] for r in rows]
assert "USER_MESSAGE_RECEIVED" in event_types
assert "INVOCATION_STARTING" in event_types
# Every row must share the same trace_id.
trace_ids = {r["trace_id"] for r in rows}
assert len(trace_ids) == 1, (
"Expected 1 unique trace_id across all events, got"
f" {len(trace_ids)}: {trace_ids}"
)
provider.shutdown()
class TestRootAgentNameAcrossInvocations:
"""Regression: root_agent_name must refresh across invocations."""
@pytest.mark.asyncio
async def test_root_agent_name_updates_between_invocations(
self,
bq_plugin_inst,
mock_write_client,
mock_session,
dummy_arrow_schema,
):
"""Two invocations with different root agents must log correct names.
Previously init_trace() only set _root_agent_name_ctx when it was
None, so the second invocation would inherit the first's root agent.
"""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
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,
)
def _make_inv_ctx(agent_name, inv_id):
agent = mock.create_autospec(
base_agent.BaseAgent, instance=True, spec_set=True
)
type(agent).name = mock.PropertyMock(return_value=agent_name)
type(agent).instruction = mock.PropertyMock(return_value="")
# root_agent returns itself (no parent).
agent.root_agent = agent
return InvocationContext(
agent=agent,
session=mock_session,
invocation_id=inv_id,
session_service=mock_session_service,
plugin_manager=mock_plugin_manager,
)
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
# --- Invocation 1: root agent = "RootA" ---
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
bigquery_agent_analytics_plugin._active_invocation_id_ctx.set(None)
bigquery_agent_analytics_plugin._root_agent_name_ctx.set(None)
inv1 = _make_inv_ctx("RootA", "inv-001")
cb1 = CallbackContext(inv1)
await bq_plugin_inst.before_run_callback(invocation_context=inv1)
await bq_plugin_inst.before_agent_callback(
agent=inv1.agent, callback_context=cb1
)
await bq_plugin_inst.after_agent_callback(
agent=inv1.agent, callback_context=cb1
)
await bq_plugin_inst.after_run_callback(invocation_context=inv1)
await asyncio.sleep(0.01)
rows_inv1 = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
# --- Invocation 2: root agent = "RootB" ---
mock_write_client.append_rows.reset_mock()
inv2 = _make_inv_ctx("RootB", "inv-002")
cb2 = CallbackContext(inv2)
await bq_plugin_inst.before_run_callback(invocation_context=inv2)
await bq_plugin_inst.before_agent_callback(
agent=inv2.agent, callback_context=cb2
)
await bq_plugin_inst.after_agent_callback(
agent=inv2.agent, callback_context=cb2
)
await bq_plugin_inst.after_run_callback(invocation_context=inv2)
await asyncio.sleep(0.01)
rows_inv2 = await _get_captured_rows_async(
mock_write_client, dummy_arrow_schema
)
# Parse root_agent_name from the attributes JSON column.
def _get_root_names(rows):
names = set()
for r in rows:
attrs = r.get("attributes")
if attrs:
parsed = json.loads(attrs) if isinstance(attrs, str) else attrs
if "root_agent_name" in parsed:
names.add(parsed["root_agent_name"])
return names
names_inv1 = _get_root_names(rows_inv1)
names_inv2 = _get_root_names(rows_inv2)
# Invocation 1 should only have "RootA".
assert names_inv1 == {"RootA"}, f"Expected {{'RootA'}}, got {names_inv1}"
# Invocation 2 must have "RootB", NOT stale "RootA".
assert names_inv2 == {"RootB"}, f"Expected {{'RootB'}}, got {names_inv2}"
provider.shutdown()
class TestAfterRunCleanupExceptionSafety:
"""after_run_callback cleanup must execute even if _log_event fails."""
@pytest.mark.asyncio
async def test_cleanup_runs_when_log_event_raises(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
callback_context,
mock_agent,
):
"""Stale state is cleared even when _log_event raises."""
from opentelemetry.sdk.trace import TracerProvider as SdkProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
provider = SdkProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
real_tracer = provider.get_tracer("test")
with mock.patch.object(
bigquery_agent_analytics_plugin, "tracer", real_tracer
):
bigquery_agent_analytics_plugin._span_records_ctx.set(None)
bigquery_agent_analytics_plugin._active_invocation_id_ctx.set(None)
bigquery_agent_analytics_plugin._root_agent_name_ctx.set(None)
# Run a normal before_run to initialise state.
await bq_plugin_inst.before_run_callback(
invocation_context=invocation_context
)
await bq_plugin_inst.before_agent_callback(
agent=mock_agent, callback_context=callback_context
)
# Verify state is populated.
assert bigquery_agent_analytics_plugin._span_records_ctx.get()
assert (
bigquery_agent_analytics_plugin._active_invocation_id_ctx.get()
is not None
)
# Make _log_event raise inside after_run_callback.
with mock.patch.object(
bq_plugin_inst,
"_log_event",
side_effect=RuntimeError("boom"),
):
# _safe_callback swallows the exception, but cleanup in
# the finally block must still execute.
await bq_plugin_inst.after_run_callback(
invocation_context=invocation_context
)
# All invocation state must be cleaned up despite the error.
records = bigquery_agent_analytics_plugin._span_records_ctx.get()
assert records == [] or records is None
assert (
bigquery_agent_analytics_plugin._active_invocation_id_ctx.get()
is None
)
assert bigquery_agent_analytics_plugin._root_agent_name_ctx.get() is None
provider.shutdown()
class TestStringSystemPromptTruncation:
"""Tests that a string system prompt is truncated in parse()."""
@pytest.mark.asyncio
async def test_long_string_system_prompt_is_truncated(self):
"""A string system_instruction exceeding max_content_length is truncated."""
parser = bigquery_agent_analytics_plugin.HybridContentParser(
offloader=None,
trace_id="test-trace",
span_id="test-span",
max_length=50,
)
long_prompt = "A" * 200
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Hi")])],
config=types.GenerateContentConfig(
system_instruction=long_prompt,
),
)
payload, _, is_truncated = await parser.parse(llm_request)
assert is_truncated
assert len(payload["system_prompt"]) < 200
assert "TRUNCATED" in payload["system_prompt"]
class TestSessionStateTruncation:
"""Tests that session state is truncated in _enrich_attributes."""
@pytest.mark.asyncio
async def test_oversized_session_state_is_truncated(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
mock_session,
invocation_context,
):
"""Session state with large values is truncated."""
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
max_content_length=30,
)
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=config,
)
await plugin._ensure_started()
# Set a large session state value.
large_value = "X" * 200
type(mock_session).state = mock.PropertyMock(
return_value={"big_key": large_value}
)
callback_ctx = CallbackContext(invocation_context=invocation_context)
event_data = bigquery_agent_analytics_plugin.EventData()
attrs = plugin._enrich_attributes(event_data, callback_ctx)
state = attrs["session_metadata"]["state"]
assert len(state["big_key"]) < 200
assert "TRUNCATED" in state["big_key"]
await plugin.shutdown()
class TestSchemaUpgradeNestedFields:
"""Tests for nested RECORD field detection in schema upgrade."""
def _make_plugin(self):
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
auto_schema_upgrade=True,
)
with mock.patch("google.cloud.bigquery.Client"):
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=config,
)
plugin.client = mock.MagicMock()
plugin.full_table_id = f"{PROJECT_ID}.{DATASET_ID}.{TABLE_ID}"
return plugin
def test_nested_field_detected(self):
"""A new sub-field in a RECORD triggers an upgrade."""
plugin = self._make_plugin()
existing_record = bigquery.SchemaField(
"metadata",
"RECORD",
fields=[
bigquery.SchemaField("key", "STRING"),
],
)
desired_record = bigquery.SchemaField(
"metadata",
"RECORD",
fields=[
bigquery.SchemaField("key", "STRING"),
bigquery.SchemaField("value", "STRING"),
],
)
plugin._schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
desired_record,
]
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
existing_record,
]
existing.labels = {}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
plugin.client.update_table.assert_called_once()
updated_table = plugin.client.update_table.call_args[0][0]
# Find the metadata field and check it has both sub-fields.
metadata_field = next(
f for f in updated_table.schema if f.name == "metadata"
)
sub_names = {sf.name for sf in metadata_field.fields}
assert "key" in sub_names
assert "value" in sub_names
def test_version_label_not_stamped_on_failure(self):
"""A failed update_table does not persist the version label."""
plugin = self._make_plugin()
plugin._schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
bigquery.SchemaField("new_col", "STRING"),
]
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
]
existing.labels = {}
plugin.client.get_table.return_value = existing
plugin.client.update_table.side_effect = Exception("network error")
# Should not raise.
plugin._ensure_schema_exists()
# The label is set on the table object before update_table is
# called, but since update_table failed the label was never
# persisted remotely. On the next run the stored_version will
# still be None (from the real BQ table) so the upgrade retries.
# We verify that update_table was actually attempted.
plugin.client.update_table.assert_called_once()
def test_nested_upgrade_preserves_policy_tags(self):
"""RECORD field metadata (e.g. policy_tags) is preserved on upgrade."""
from google.cloud.bigquery import schema as bq_schema
plugin = self._make_plugin()
existing_record = bigquery.SchemaField(
"metadata",
"RECORD",
policy_tags=bq_schema.PolicyTagList(
names=["projects/p/locations/us/taxonomies/t/policyTags/pt"]
),
fields=[
bigquery.SchemaField("key", "STRING"),
],
)
desired_record = bigquery.SchemaField(
"metadata",
"RECORD",
fields=[
bigquery.SchemaField("key", "STRING"),
bigquery.SchemaField("value", "STRING"),
],
)
plugin._schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
desired_record,
]
existing = mock.MagicMock(spec=bigquery.Table)
existing.schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP"),
existing_record,
]
existing.labels = {}
plugin.client.get_table.return_value = existing
plugin._ensure_schema_exists()
plugin.client.update_table.assert_called_once()
updated_table = plugin.client.update_table.call_args[0][0]
metadata_field = next(
f for f in updated_table.schema if f.name == "metadata"
)
# Sub-fields were merged.
sub_names = {sf.name for sf in metadata_field.fields}
assert "key" in sub_names
assert "value" in sub_names
# policy_tags preserved from the existing field.
assert metadata_field.policy_tags is not None
assert (
"projects/p/locations/us/taxonomies/t/policyTags/pt"
in metadata_field.policy_tags.names
)
class TestMultiLoopShutdownDrainsOtherLoops:
"""Tests that shutdown() drains batch processors on other loops."""
@pytest.mark.asyncio
async def test_other_loop_batch_processor_drained(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""Shutdown drains batch_processor.shutdown on non-current loops."""
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
await plugin._ensure_started()
# Create a mock "other" loop with a mock batch processor.
other_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop)
other_loop.is_closed.return_value = False
mock_other_bp = mock.AsyncMock()
mock_other_write_client = mock.MagicMock()
mock_other_write_client.transport = mock.AsyncMock()
other_state = bigquery_agent_analytics_plugin._LoopState(
write_client=mock_other_write_client,
batch_processor=mock_other_bp,
)
plugin._loop_state_by_loop[other_loop] = other_state
# Patch run_coroutine_threadsafe to verify it's called for
# the other loop's batch_processor. Close the coroutine arg
# to avoid "coroutine was never awaited" RuntimeWarning.
mock_future = mock.MagicMock()
mock_future.result.return_value = None
def _fake_run_coroutine_threadsafe(coro, loop):
coro.close()
return mock_future
with mock.patch.object(
asyncio,
"run_coroutine_threadsafe",
side_effect=_fake_run_coroutine_threadsafe,
) as mock_rcts:
await plugin.shutdown()
# Verify run_coroutine_threadsafe was called with
# the other loop.
mock_rcts.assert_called()
call_args = mock_rcts.call_args
assert call_args[0][1] is other_loop
class TestCacheMetadataLogging:
"""Tests for logging cache_metadata from LlmResponse."""
@pytest.mark.asyncio
async def test_cache_metadata_logged_when_present(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Verifies cache_metadata is logged into BigQuery attributes when present."""
llm_response = llm_response_lib.LlmResponse(
content=types.Content(parts=[types.Part(text="Cache test")]),
cache_metadata={"fingerprint": "abc-123", "contents_count": 2},
)
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.after_model_callback(
callback_context=callback_context,
llm_response=llm_response,
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
log_entry = next(r for r in rows if r["event_type"] == "LLM_RESPONSE")
attributes = json.loads(log_entry["attributes"])
assert "cache_metadata" in attributes
assert attributes["cache_metadata"]["fingerprint"] == "abc-123"
assert attributes["cache_metadata"]["contents_count"] == 2
@pytest.mark.asyncio
async def test_missing_cache_metadata_does_not_crash(
self,
bq_plugin_inst,
mock_write_client,
callback_context,
dummy_arrow_schema,
):
"""Verifies missing cache_metadata gracefully defaults using getattr."""
class LegacyLlmResponse:
def __init__(self):
self.content = types.Content(parts=[types.Part(text="Mock text")])
self.usage_metadata = None
self.model_version = "v1"
self.partial = False
# Deliberately omitting cache_metadata
mock_response = LegacyLlmResponse()
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
await bq_plugin_inst.after_model_callback(
callback_context=callback_context,
llm_response=mock_response,
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
log_entry = next(r for r in rows if r["event_type"] == "LLM_RESPONSE")
attributes = json.loads(log_entry["attributes"])
assert "cache_metadata" not in attributes
# ==============================================================
# TEST CLASS: A2A_INTERACTION event logging via on_event_callback
# ==============================================================
class TestA2AInteractionLogging:
"""Tests for A2A interaction event emission via on_event_callback.
When a RemoteA2aAgent processes a response, it attaches A2A
metadata (``a2a:task_id``, ``a2a:context_id``, ``a2a:request``,
``a2a:response``) to the event's ``custom_metadata``. The
plugin's ``on_event_callback`` should detect events carrying
``a2a:request`` or ``a2a:response`` and log an
``A2A_INTERACTION`` event so the remote agent's response and
cross-reference IDs are visible in BigQuery.
"""
@pytest.mark.asyncio
async def test_a2a_interaction_logged_for_response_metadata(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""Event with a2a:response in custom_metadata emits A2A_INTERACTION."""
a2a_meta = {
"a2a:task_id": "task-abc",
"a2a:context_id": "ctx-123",
"a2a:response": {"status": "completed", "text": "result"},
}
event = event_lib.Event(
author="remote_agent",
custom_metadata=a2a_meta,
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
result = await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
assert result is None
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
event_types = [r["event_type"] for r in rows]
assert "A2A_INTERACTION" in event_types
a2a_row = next(r for r in rows if r["event_type"] == "A2A_INTERACTION")
attributes = json.loads(a2a_row["attributes"])
assert "a2a_metadata" in attributes
assert attributes["a2a_metadata"]["a2a:task_id"] == "task-abc"
assert attributes["a2a_metadata"]["a2a:context_id"] == "ctx-123"
# Content should contain the a2a:response payload.
content = json.loads(a2a_row["content"])
assert content["status"] == "completed"
@pytest.mark.asyncio
async def test_a2a_interaction_logged_for_request_metadata(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""Event with a2a:request (no a2a:response) emits A2A_INTERACTION."""
a2a_meta = {
"a2a:task_id": "task-xyz",
"a2a:request": {"message": "hello"},
}
event = event_lib.Event(
author="remote_agent",
custom_metadata=a2a_meta,
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
result = await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
assert result is None
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
event_types = [r["event_type"] for r in rows]
assert "A2A_INTERACTION" in event_types
a2a_row = next(r for r in rows if r["event_type"] == "A2A_INTERACTION")
attributes = json.loads(a2a_row["attributes"])
assert attributes["a2a_metadata"]["a2a:request"] == {"message": "hello"}
# No a2a:response → content should be None.
assert a2a_row["content"] is None
@pytest.mark.asyncio
async def test_no_a2a_interaction_for_irrelevant_metadata(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Events with only a2a:task_id (no request/response) are skipped."""
a2a_meta = {
"a2a:task_id": "task-only",
"a2a:context_id": "ctx-only",
}
event = event_lib.Event(
author="remote_agent",
custom_metadata=a2a_meta,
)
result = await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
assert result is None
await asyncio.sleep(0.05)
# No events logged — a2a:task_id alone is not a meaningful
# interaction payload.
assert mock_write_client.append_rows.call_count == 0
@pytest.mark.asyncio
async def test_no_a2a_interaction_for_no_metadata(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Events without custom_metadata produce no A2A_INTERACTION."""
event = event_lib.Event(author="regular_agent")
result = await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
assert result is None
await asyncio.sleep(0.05)
assert mock_write_client.append_rows.call_count == 0
# ================================================================
# TEST CLASS: Dataset location handling (Issue #5476)
# ================================================================
class TestDatasetLocationHandling:
"""Tests that BQ client is created without a default location.
When location is omitted from bigquery.Client(), client.query()
sends no location field in the API request, letting BigQuery
infer location from the referenced dataset. This prevents
silent view-creation failures for non-US datasets.
"""
@pytest.mark.asyncio
async def test_client_created_without_location(
self,
mock_auth_default,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""bigquery.Client is created without a location parameter."""
with mock.patch.object(bigquery, "Client", autospec=True) as mock_bq_cls:
mock_bq_cls.return_value.get_table.side_effect = (
cloud_exceptions.NotFound("table")
)
mock_bq_cls.return_value.create_table.return_value = None
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
location="europe-west1",
config=bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
create_views=False,
),
) as plugin:
await plugin._ensure_started()
mock_bq_cls.assert_called_once()
_, kwargs = mock_bq_cls.call_args
assert "location" not in kwargs
@pytest.mark.asyncio
async def test_view_query_omits_location(
self,
mock_auth_default,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""View creation DDL queries do not pass an explicit location."""
with mock.patch.object(bigquery, "Client", autospec=True) as mock_bq_cls:
mock_client = mock_bq_cls.return_value
mock_client.get_table.return_value = mock.MagicMock()
mock_client.query.return_value.result.return_value = None
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
create_views=True,
),
) as plugin:
await plugin._ensure_started()
assert mock_client.query.call_count > 0
for call in mock_client.query.call_args_list:
_, kwargs = call
# No explicit location — BQ infers from dataset
assert "location" not in kwargs
@pytest.mark.asyncio
async def test_view_error_still_logged(
self,
mock_auth_default,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""View creation errors are logged but not raised."""
with mock.patch.object(bigquery, "Client", autospec=True) as mock_bq_cls:
mock_client = mock_bq_cls.return_value
mock_client.get_table.return_value = mock.MagicMock()
mock_client.query.return_value.result.side_effect = Exception(
"view error"
)
# Should not raise
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
create_views=True,
),
) as plugin:
await plugin._ensure_started()
assert plugin._started
# ================================================================
# TEST CLASS: Fork detection after pickle (Issue #86 / PR #5528)
# ================================================================
class TestForkDetectionAfterPickle:
"""Tests that unpickled plugins do not false-positive fork detection."""
@pytest.mark.asyncio
async def test_no_reset_after_unpickle(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""Unpickled plugin does not trigger _reset_runtime_state and
records os.getpid() after startup.
"""
import pickle
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
create_views=False,
)
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
)
pickled = pickle.dumps(plugin)
unpickled = pickle.loads(pickled)
assert unpickled._init_pid == 0
with mock.patch.object(unpickled, "_reset_runtime_state") as mock_reset:
await unpickled._ensure_started()
mock_reset.assert_not_called()
assert unpickled._started
assert unpickled._init_pid == os.getpid()
await unpickled.shutdown()
@pytest.mark.asyncio
async def test_reset_on_real_fork(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
"""Plugin detects real fork when _init_pid is a real non-zero PID."""
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
create_views=False,
)
async with managed_plugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
config=config,
) as plugin:
await plugin._ensure_started()
plugin._init_pid = max(os.getpid() - 1, 1)
plugin._started = True
with mock.patch.object(
plugin, "_reset_runtime_state", wraps=plugin._reset_runtime_state
) as mock_reset:
await plugin._ensure_started()
mock_reset.assert_called_once()
# ================================================================
# TEST CLASS: GCS offload unit mismatch fix (Issue #5561)
# ================================================================
class TestOffloadUnitSeparation:
"""Tests that byte-based inline limit and character-based truncation
limit are evaluated independently for the GCS offload decision.
"""
@pytest.mark.asyncio
async def test_multibyte_text_offloaded_by_byte_limit(self):
"""Multi-byte text exceeding inline_text_limit bytes is offloaded."""
mock_offloader = mock.AsyncMock()
mock_offloader.upload_content.return_value = "gs://bucket/offloaded.txt"
parser = bigquery_agent_analytics_plugin.HybridContentParser(
offloader=mock_offloader,
trace_id="t",
span_id="s",
max_length=-1,
)
text = "\U0001f600" * 10000
assert len(text) == 10000
assert len(text.encode("utf-8")) > 32 * 1024
content = types.Content(parts=[types.Part(text=text)])
_, parts, _ = await parser._parse_content_object(content)
mock_offloader.upload_content.assert_called_once()
assert parts[0]["storage_mode"] == "GCS_REFERENCE"
@pytest.mark.asyncio
async def test_ascii_under_both_limits_stays_inline(self):
"""ASCII text under both byte and character limits stays inline."""
mock_offloader = mock.AsyncMock()
parser = bigquery_agent_analytics_plugin.HybridContentParser(
offloader=mock_offloader,
trace_id="t",
span_id="s",
max_length=50000,
)
text = "A" * 1000
content = types.Content(parts=[types.Part(text=text)])
_, parts, _ = await parser._parse_content_object(content)
mock_offloader.upload_content.assert_not_called()
assert parts[0]["storage_mode"] == "INLINE"
assert parts[0]["text"] == text
@pytest.mark.asyncio
async def test_text_exceeding_char_limit_offloaded(self):
"""ASCII text exceeding max_length characters is offloaded."""
mock_offloader = mock.AsyncMock()
mock_offloader.upload_content.return_value = "gs://bucket/big.txt"
parser = bigquery_agent_analytics_plugin.HybridContentParser(
offloader=mock_offloader,
trace_id="t",
span_id="s",
max_length=100,
)
text = "X" * 200
assert len(text.encode("utf-8")) < 32 * 1024
assert len(text) > 100
content = types.Content(parts=[types.Part(text=text)])
_, parts, _ = await parser._parse_content_object(content)
mock_offloader.upload_content.assert_called_once()
assert parts[0]["storage_mode"] == "GCS_REFERENCE"
@pytest.mark.asyncio
async def test_multibyte_under_char_and_byte_limits_stays_inline(self):
"""Regression test: 3K emoji (12K bytes) with max_length=10000
should stay inline — under both real limits.
"""
mock_offloader = mock.AsyncMock()
parser = bigquery_agent_analytics_plugin.HybridContentParser(
offloader=mock_offloader,
trace_id="t",
span_id="s",
max_length=10000,
)
text = "\U0001f600" * 3000
assert len(text) < 10000
assert len(text.encode("utf-8")) > 10000
assert len(text.encode("utf-8")) < 32 * 1024
content = types.Content(parts=[types.Part(text=text)])
_, parts, _ = await parser._parse_content_object(content)
mock_offloader.upload_content.assert_not_called()
assert parts[0]["storage_mode"] == "INLINE"
@pytest.mark.asyncio
async def test_no_offloader_falls_back_to_truncate(self):
"""Without offloader, text exceeding char limit is truncated inline."""
parser = bigquery_agent_analytics_plugin.HybridContentParser(
offloader=None,
trace_id="t",
span_id="s",
max_length=50,
)
text = "Z" * 200
content = types.Content(parts=[types.Part(text=text)])
_, parts, is_truncated = await parser._parse_content_object(content)
assert is_truncated
assert parts[0]["storage_mode"] == "INLINE"
assert "TRUNCATED" in parts[0]["text"]
# ================================================================
# TEST CLASS: AGENT_RESPONSE logging (Issue #87)
# ================================================================
class TestAgentResponseLogging:
"""Tests that final agent response events are captured correctly."""
@pytest.mark.asyncio
async def test_logs_final_text_response(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""Final text response is logged as AGENT_RESPONSE with
source_event_author from event.author.
"""
event = event_lib.Event(
author="sub_agent",
content=types.Content(parts=[types.Part(text="Here is your answer.")]),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
agent_resp_rows = [r for r in rows if r["event_type"] == "AGENT_RESPONSE"]
assert len(agent_resp_rows) == 1
row = agent_resp_rows[0]
content = json.loads(row["content"])
assert "Here is your answer" in content["response"]
attributes = json.loads(row["attributes"])
# source_event_author must come from event.author
assert attributes["source_event_author"] == "sub_agent"
@pytest.mark.asyncio
async def test_skips_function_call_events(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Events with function calls are not logged as AGENT_RESPONSE."""
fc = types.FunctionCall(name="my_tool", args={"x": 1})
event = event_lib.Event(
author="agent",
content=types.Content(parts=[types.Part(function_call=fc)]),
)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
assert mock_write_client.append_rows.call_count == 0
@pytest.mark.asyncio
async def test_skips_function_response_events(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Events with function responses are not logged as AGENT_RESPONSE."""
fr = types.FunctionResponse(name="my_tool", response={"result": "ok"})
event = event_lib.Event(
author="agent",
content=types.Content(parts=[types.Part(function_response=fr)]),
actions=event_actions_lib.EventActions(skip_summarization=True),
)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
assert mock_write_client.append_rows.call_count == 0
@pytest.mark.asyncio
async def test_skips_partial_events(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Partial streaming events are not logged as AGENT_RESPONSE."""
event = event_lib.Event(
author="agent",
content=types.Content(parts=[types.Part(text="partial chunk")]),
partial=True,
)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
assert mock_write_client.append_rows.call_count == 0
@pytest.mark.asyncio
async def test_skips_long_running_tool_events(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""Long-running tool events are not logged as AGENT_RESPONSE.
They DO emit TOOL_PAUSED — here via the unmatched-id fallback, since
the function_call part has no id matching the long_running_tool_id.
"""
fc = types.FunctionCall(name="long_tool", args={})
event = event_lib.Event(
author="agent",
content=types.Content(parts=[types.Part(function_call=fc)]),
long_running_tool_ids={"call-1"},
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
types_emitted = [r["event_type"] for r in rows]
assert "AGENT_RESPONSE" not in types_emitted
# The pause is still observable via the fallback TOOL_PAUSED row.
assert types_emitted == ["TOOL_PAUSED"]
@pytest.mark.asyncio
async def test_skips_thought_only_events(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Thought-only final events are not logged as AGENT_RESPONSE."""
event = event_lib.Event(
author="agent",
content=types.Content(
parts=[types.Part(text="internal reasoning...", thought=True)]
),
)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
assert mock_write_client.append_rows.call_count == 0
@pytest.mark.asyncio
async def test_mixed_thought_and_visible_logs_only_visible(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""Mixed thought + visible text logs only the visible portion."""
event = event_lib.Event(
author="agent",
content=types.Content(
parts=[
types.Part(text="thinking step 1...", thought=True),
types.Part(text="Here is the answer."),
]
),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
agent_resp_rows = [r for r in rows if r["event_type"] == "AGENT_RESPONSE"]
assert len(agent_resp_rows) == 1
content = json.loads(agent_resp_rows[0]["content"])
assert "Here is the answer" in content["response"]
assert "thinking step" not in content["response"]
@pytest.mark.asyncio
async def test_skips_empty_part_events(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Events with only empty Part() do not log AGENT_RESPONSE."""
event = event_lib.Event(
author="agent",
content=types.Content(parts=[types.Part()]),
)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
assert mock_write_client.append_rows.call_count == 0
@pytest.mark.asyncio
async def test_skips_empty_text_events(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Events with Part(text='') do not log AGENT_RESPONSE."""
event = event_lib.Event(
author="agent",
content=types.Content(parts=[types.Part(text="")]),
)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
assert mock_write_client.append_rows.call_count == 0
@pytest.mark.asyncio
async def test_skips_executable_code_only_events(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
):
"""Events with only executable_code parts do not log AGENT_RESPONSE."""
event = event_lib.Event(
author="agent",
content=types.Content(
parts=[
types.Part(
executable_code=types.ExecutableCode(
code="print('hi')", language="PYTHON"
)
)
]
),
)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.05)
assert mock_write_client.append_rows.call_count == 0
class TestDropStats:
"""Tests that dropped events are counted and exposed via get_drop_stats."""
def _make_processor(
self, arrow_schema, *, queue_max_size=10, retry_config=None
):
"""Builds a BatchProcessor with a mock write client (writer not started)."""
return bigquery_agent_analytics_plugin.BatchProcessor(
write_client=mock.MagicMock(),
arrow_schema=arrow_schema,
write_stream=DEFAULT_STREAM_NAME,
batch_size=1,
flush_interval=1.0,
retry_config=(
retry_config or bigquery_agent_analytics_plugin.RetryConfig()
),
queue_max_size=queue_max_size,
shutdown_timeout=10.0,
)
def _stub_arrow_prep(self, bp):
"""Stubs Arrow serialization so write tests need no real row schema."""
fake_batch = mock.MagicMock()
fake_batch.serialize.return_value.to_pybytes.return_value = b"batch"
bp._prepare_arrow_batch = mock.MagicMock(return_value=fake_batch)
@pytest.mark.asyncio
async def test_queue_full_drops_are_counted(self, dummy_arrow_schema):
# Writer is not started, so a size-1 queue fills after one append and the
# next two appends overflow and are dropped.
bp = self._make_processor(dummy_arrow_schema, queue_max_size=1)
await bp.append({"event": 0})
await bp.append({"event": 1})
await bp.append({"event": 2})
assert bp.get_drop_stats()["queue_full"] == 2
assert bp.dropped_event_count == 2
@pytest.mark.asyncio
async def test_retry_exhaustion_drops_are_counted(self, dummy_arrow_schema):
# max_retries=0 with zero delay drops on the first failure without sleeping.
retry_config = bigquery_agent_analytics_plugin.RetryConfig(
max_retries=0, initial_delay=0.0, multiplier=1.0, max_delay=0.0
)
bp = self._make_processor(dummy_arrow_schema, retry_config=retry_config)
self._stub_arrow_prep(bp)
async def fake_append_rows(requests, **kwargs):
del requests, kwargs
resp = mock.MagicMock()
resp.row_errors = []
resp.error = mock.MagicMock()
resp.error.code = bigquery_agent_analytics_plugin._GRPC_UNAVAILABLE
resp.error.message = "unavailable"
return _async_gen(resp)
bp.write_client.append_rows.side_effect = fake_append_rows
await bp._write_rows_with_retry([{"a": 1}, {"a": 2}])
assert bp.get_drop_stats()["retry_exhausted"] == 2
assert bp.dropped_event_count == 2
@pytest.mark.asyncio
async def test_non_retryable_drops_are_counted(self, dummy_arrow_schema):
bp = self._make_processor(dummy_arrow_schema)
self._stub_arrow_prep(bp)
async def fake_append_rows(requests, **kwargs):
del requests, kwargs
resp = mock.MagicMock()
resp.row_errors = []
resp.error = mock.MagicMock()
resp.error.code = 3 # INVALID_ARGUMENT, non-retryable.
resp.error.message = "bad request"
return _async_gen(resp)
bp.write_client.append_rows.side_effect = fake_append_rows
await bp._write_rows_with_retry([{"a": 1}])
assert bp.get_drop_stats()["non_retryable"] == 1
assert bp.dropped_event_count == 1
def test_plugin_get_drop_stats_aggregates_across_loops(
self, dummy_arrow_schema
):
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID, dataset_id=DATASET_ID, table_id=TABLE_ID
)
bp1 = self._make_processor(dummy_arrow_schema)
bp2 = self._make_processor(dummy_arrow_schema)
bp1._dropped["queue_full"] = 3
bp1._dropped["retry_exhausted"] = 1
bp2._dropped["queue_full"] = 4
loop1 = mock.MagicMock(spec=asyncio.AbstractEventLoop)
loop2 = mock.MagicMock(spec=asyncio.AbstractEventLoop)
plugin._loop_state_by_loop[loop1] = (
bigquery_agent_analytics_plugin._LoopState(mock.MagicMock(), bp1)
)
plugin._loop_state_by_loop[loop2] = (
bigquery_agent_analytics_plugin._LoopState(mock.MagicMock(), bp2)
)
stats = plugin.get_drop_stats()
assert stats["queue_full"] == 7
assert stats["retry_exhausted"] == 1
def test_plugin_get_drop_stats_empty_without_processor(self):
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID, dataset_id=DATASET_ID, table_id=TABLE_ID
)
assert plugin.get_drop_stats() == {}
# -----------------------------------------------------------------------------
# ADK 2.0 minimum producer cut
#
# Coverage matrix:
# A1 / A2 attributes.adk.{schema_version, app_name} on every row
# A3 attributes.adk.source_event_id on Event-originating rows
# C1 attributes.adk.node {path, run_id, parent_run_id}
# C2 attributes.adk.branch
# C3 attributes.adk.scope {id, kind}
# C4 AGENT_TRANSFER emit
# C5 EVENT_COMPACTION emit (preserves fractional float epoch)
# C6 AGENT_STATE_CHECKPOINT emit (both shapes) + id-stabilization
# C7 TOOL_PAUSED with pause_kind / function_call_id
# HITL non-routing to TOOL_COMPLETED
# user-message TOOL_COMPLETED with pause_kind='tool'
# C8 attributes.adk.{route, render_ui_widgets, rewind_before_invocation_id}
# D1 on_state_change_callback removed
# -----------------------------------------------------------------------------
def test_derive_scope_unscoped():
"""C3: None isolation_scope → scope = null."""
assert bigquery_agent_analytics_plugin._derive_scope(None) is None
def test_derive_scope_node_run_bare():
"""C3: bare 'name@run_id' classifies as node_run (not function_call)."""
scope = bigquery_agent_analytics_plugin._derive_scope("loopA@42")
assert scope == {"id": "loopA@42", "kind": "node_run"}
def test_derive_scope_node_run_path():
"""C3: 'parent/name@run_id' classifies as node_run."""
scope = bigquery_agent_analytics_plugin._derive_scope("wf/A@1/B@2")
assert scope == {"id": "wf/A@1/B@2", "kind": "node_run"}
def test_derive_scope_function_call_provider_id():
"""C3: model-provided FC IDs (call_*, toolu_*) classify as function_call."""
for fc_id in ("call_abc123", "toolu_xyz", "adk-fc-1"):
scope = bigquery_agent_analytics_plugin._derive_scope(fc_id)
assert scope == {"id": fc_id, "kind": "function_call"}, fc_id
def test_derive_scope_empty_string_unknown():
"""C3: empty/non-string anomalies classify as unknown."""
scope = bigquery_agent_analytics_plugin._derive_scope("")
assert scope == {"id": "", "kind": "unknown"}
def test_d1_on_state_change_callback_removed():
"""D1: the deprecated stub is gone from the public surface."""
assert not hasattr(
bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin,
"on_state_change_callback",
)
class TestAdkEnvelope:
"""A1 / A2 / A3 / C1 / C2 / C3 / C8 envelope shape on emitted rows."""
@pytest.mark.asyncio
async def test_envelope_on_non_event_row(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""USER_MESSAGE_RECEIVED has no source Event → A1/A2 only, A3/C1/C2/C3 null."""
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(role="user", parts=[types.Part(text="hi")]),
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "USER_MESSAGE_RECEIVED")
attributes = json.loads(log_entry["attributes"])
adk = attributes["adk"]
# A1: schema_version always present.
assert adk["schema_version"] == (
bigquery_agent_analytics_plugin._ADK_ENVELOPE_SCHEMA_VERSION
)
# A2: app_name always present (from session).
assert adk["app_name"] == "test_app"
# A3 / C1 / C2 / C3 absent on rows without an originating Event.
assert "source_event_id" not in adk
assert "node" not in adk
assert "branch" not in adk
assert "scope" not in adk
@pytest.mark.asyncio
async def test_envelope_on_event_row(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""STATE_DELTA from on_event_callback carries the full envelope."""
state_delta = {"k": "v"}
event = event_lib.Event(
author="agent_a",
branch="branch-x",
actions=event_actions_lib.EventActions(state_delta=state_delta),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
_assert_common_fields(log_entry, "STATE_DELTA")
attributes = json.loads(log_entry["attributes"])
adk = attributes["adk"]
assert adk["schema_version"] == (
bigquery_agent_analytics_plugin._ADK_ENVELOPE_SCHEMA_VERSION
)
assert adk["app_name"] == "test_app"
# A3: real Event.id (model_post_init auto-assigns a UUID).
assert adk["source_event_id"] == event.id
assert len(event.id) == 36 # sanity
# C2: branch passthrough.
assert adk["branch"] == "branch-x"
# C1: node defaults to path="" with run_id="" and parent_run_id=null
# (no synthesis). run_id / parent_run_id are NodeInfo @property values
# parsed from path.
assert adk["node"]["path"] == ""
assert adk["node"]["run_id"] == ""
assert adk["node"]["parent_run_id"] is None
@pytest.mark.asyncio
async def test_envelope_node_with_parent_run_id(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""C1: run_id / parent_run_id are derived from NodeInfo for a nested path.
For path "wf/A@1/B@2": run_id is the leaf node's run_id ("2") and
parent_run_id is the parent node's run_id ("1").
"""
event = event_lib.Event(
author="agent_b",
actions=event_actions_lib.EventActions(state_delta={"k": "v"}),
)
event.node_info.path = "wf/A@1/B@2"
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
adk = json.loads(log_entry["attributes"])["adk"]
assert adk["node"]["path"] == "wf/A@1/B@2"
assert adk["node"]["run_id"] == "2"
assert adk["node"]["parent_run_id"] == "1"
class TestC4AgentTransfer:
@pytest.mark.asyncio
async def test_agent_transfer_emits_from_to_payload(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
event = event_lib.Event(
author="root_agent",
actions=event_actions_lib.EventActions(
transfer_to_agent="specialist_agent"
),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
transfers = [r for r in rows if r["event_type"] == "AGENT_TRANSFER"]
assert len(transfers) == 1
content = json.loads(transfers[0]["content"])
assert content == {
"from_agent": "root_agent",
"to_agent": "specialist_agent",
}
class TestC5EventCompaction:
@pytest.mark.asyncio
async def test_event_compaction_preserves_float_precision(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""C5: fractional float-epoch seconds must survive the producer."""
compaction = event_actions_lib.EventCompaction(
start_timestamp=1700000000.125,
end_timestamp=1700000003.875,
compacted_content=types.Content(
role="model", parts=[types.Part(text="summary")]
),
)
event = event_lib.Event(
author="agent",
actions=event_actions_lib.EventActions(compaction=compaction),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
compactions = [r for r in rows if r["event_type"] == "EVENT_COMPACTION"]
assert len(compactions) == 1
content = json.loads(compactions[0]["content"])
assert content["start_timestamp"] == 1700000000.125
assert content["end_timestamp"] == 1700000003.875
assert content["start_timestamp"] != int(content["start_timestamp"])
class TestC6AgentStateCheckpoint:
@pytest.mark.asyncio
async def test_checkpoint_state_only(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""{agent_state: {...}, end_of_agent: None} emits a CHECKPOINT row."""
event = event_lib.Event(
author="agent",
actions=event_actions_lib.EventActions(
agent_state={"step": 3, "ctx": "abc"}
),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
cps = [r for r in rows if r["event_type"] == "AGENT_STATE_CHECKPOINT"]
assert len(cps) == 1
content = json.loads(cps[0]["content"])
assert content["agent_state"] == {"step": 3, "ctx": "abc"}
assert content["end_of_agent"] is False
@pytest.mark.asyncio
async def test_checkpoint_end_of_agent_only(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""{agent_state: None, end_of_agent: True} is a valid CHECKPOINT shape."""
event = event_lib.Event(
author="agent",
actions=event_actions_lib.EventActions(end_of_agent=True),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
cps = [r for r in rows if r["event_type"] == "AGENT_STATE_CHECKPOINT"]
assert len(cps) == 1
content = json.loads(cps[0]["content"])
assert content["agent_state"] is None
assert content["end_of_agent"] is True
@pytest.mark.asyncio
async def test_checkpoint_carries_real_source_event_id(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""v3 regression guard: Event.model_post_init auto-assigns id, so a
checkpoint Event constructed without explicit id still surfaces a real
36-char UUID in attributes.adk.source_event_id."""
event = event_lib.Event(
author="agent",
actions=event_actions_lib.EventActions(end_of_agent=True),
)
assert event.id and len(event.id) == 36
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.01)
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
cps = [r for r in rows if r["event_type"] == "AGENT_STATE_CHECKPOINT"]
assert len(cps) == 1
adk = json.loads(cps[0]["attributes"])["adk"]
assert adk["source_event_id"] == event.id
class TestC7ToolPauseAndComplete:
@pytest.mark.asyncio
async def test_tool_paused_non_hitl_pause_kind_tool(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
fc = types.FunctionCall(
id="call-1", name="long_running_search", args={"q": "x"}
)
event = event_lib.Event(
author="agent",
content=types.Content(
role="model", parts=[types.Part(function_call=fc)]
),
long_running_tool_ids={"call-1"},
actions=event_actions_lib.EventActions(),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await bq_plugin_inst.flush()
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
pauses = [r for r in rows if r["event_type"] == "TOOL_PAUSED"]
assert len(pauses) == 1
# C7 pair keys live UNDER ``attributes.adk`` so the consumer SQL on
# ``JSON_VALUE(attributes, '$.adk.function_call_id')`` resolves.
adk = json.loads(pauses[0]["attributes"])["adk"]
assert adk["pause_kind"] == "tool"
assert adk["function_call_id"] == "call-1"
@pytest.mark.asyncio
async def test_tool_paused_hitl_pause_kind(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""C7: HITL long-running call → pause_kind derived from NAME, not id."""
fc = types.FunctionCall(
id="call-hitl-1", name="adk_request_confirmation", args={}
)
event = event_lib.Event(
author="agent",
content=types.Content(
role="model", parts=[types.Part(function_call=fc)]
),
long_running_tool_ids={"call-hitl-1"},
actions=event_actions_lib.EventActions(),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await bq_plugin_inst.flush()
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
pauses = [r for r in rows if r["event_type"] == "TOOL_PAUSED"]
assert len(pauses) == 1
adk = json.loads(pauses[0]["attributes"])["adk"]
assert adk["pause_kind"] == "hitl_confirmation"
assert adk["function_call_id"] == "call-hitl-1"
@pytest.mark.asyncio
async def test_user_message_function_response_emits_tool_completed(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""C7: non-HITL function_response in a user message → TOOL_COMPLETED
with pause_kind='tool' (this is the long-running resume path)."""
fr = types.FunctionResponse(
id="call-1", name="long_running_search", response={"hits": 7}
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(
role="user", parts=[types.Part(function_response=fr)]
),
)
await bq_plugin_inst.flush()
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
completed = [r for r in rows if r["event_type"] == "TOOL_COMPLETED"]
assert len(completed) == 1
adk = json.loads(completed[0]["attributes"])["adk"]
assert adk["pause_kind"] == "tool"
assert adk["function_call_id"] == "call-1"
@pytest.mark.asyncio
async def test_hitl_user_message_does_not_emit_tool_completed(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""C7 HITL non-routing: an adk_request_confirmation function_response in
a user message emits ONLY HITL_CONFIRMATION_REQUEST_COMPLETED, never
TOOL_COMPLETED."""
fr = types.FunctionResponse(
id="call-hitl-1",
name="adk_request_confirmation",
response={"approved": True},
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_user_message_callback(
invocation_context=invocation_context,
user_message=types.Content(
role="user", parts=[types.Part(function_response=fr)]
),
)
await bq_plugin_inst.flush()
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
types_emitted = {r["event_type"] for r in rows}
assert "HITL_CONFIRMATION_REQUEST_COMPLETED" in types_emitted
assert "TOOL_COMPLETED" not in types_emitted
class TestC8ActionAttributes:
@pytest.mark.asyncio
async def test_route_and_rewind_flat_under_attributes_adk(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""C8: route / rewind_before_invocation_id mirror under
attributes.adk.* (flat-with-prefix, NOT nested under .actions.)."""
event = event_lib.Event(
author="agent",
actions=event_actions_lib.EventActions(
state_delta={"k": "v"}, # to ensure an emit happens
route="branch_b",
rewind_before_invocation_id="inv-earlier",
),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await asyncio.sleep(0.01)
log_entry = await _get_captured_event_dict_async(
mock_write_client, dummy_arrow_schema
)
adk = json.loads(log_entry["attributes"])["adk"]
# Flat-with-prefix mirror under attributes.adk.*.
assert adk["route"] == "branch_b"
assert adk["rewind_before_invocation_id"] == "inv-earlier"
# Not nested under .actions.
assert "actions" not in adk
class TestViewDefsRegistration:
"""The plugin's own per-event-type view defs cover the new types."""
def test_new_event_types_registered_in_view_defs(self):
defs = bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS
for event_type in (
"AGENT_TRANSFER",
"EVENT_COMPACTION",
"AGENT_STATE_CHECKPOINT",
"TOOL_PAUSED",
):
assert event_type in defs, f"{event_type} missing from _EVENT_VIEW_DEFS"
assert isinstance(defs[event_type], list)
def test_tool_paused_view_extracts_pair_keys(self):
cols = "\n".join(
bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS["TOOL_PAUSED"]
)
assert "$.adk.pause_kind" in cols
assert "$.adk.function_call_id" in cols
def test_compaction_view_preserves_float_and_widens(self):
cols = "\n".join(
bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS["EVENT_COMPACTION"]
)
# Float passthrough for diagnostics + TIMESTAMP_MICROS widening
# (TIMESTAMP_SECONDS would truncate fractional windows).
assert "AS FLOAT64) AS start_seconds" in cols
assert "TIMESTAMP_MICROS" in cols
assert "TIMESTAMP_SECONDS" not in cols
def test_tool_completed_view_exposes_pair_keys(self):
"""v_tool_completed can do the pause/completion join end-to-end."""
cols = "\n".join(
bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS["TOOL_COMPLETED"]
)
assert "$.adk.pause_kind" in cols
assert "$.adk.function_call_id" in cols
def test_checkpoint_view_exposes_agent_state_type(self):
"""v_agent_state_checkpoint discriminates explicit JSON null from
object-valued agent_state via JSON_TYPE(JSON_QUERY(...))."""
cols = "\n".join(
bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS[
"AGENT_STATE_CHECKPOINT"
]
)
assert "JSON_TYPE(JSON_QUERY(content," in cols
assert "AS agent_state_type" in cols
class TestUnmatchedLongRunningIdFallback:
@pytest.mark.asyncio
async def test_unmatched_long_running_id_emits_tool_paused(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
caplog,
):
"""A long_running_tool_id with no matching function_call part still
emits a pairable TOOL_PAUSED row with pause_kind='tool' + warning."""
event = event_lib.Event(
author="agent",
content=types.Content(
role="model", parts=[types.Part(text="thinking...")]
),
long_running_tool_ids={"orphan-pause-1"},
actions=event_actions_lib.EventActions(),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
with caplog.at_level("WARNING"):
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await bq_plugin_inst.flush()
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
pauses = [r for r in rows if r["event_type"] == "TOOL_PAUSED"]
assert len(pauses) == 1
adk = json.loads(pauses[0]["attributes"])["adk"]
assert adk["pause_kind"] == "tool"
assert adk["function_call_id"] == "orphan-pause-1"
assert any(
"no matching function_call part" in rec.message
for rec in caplog.records
)
@pytest.mark.asyncio
async def test_matched_id_not_double_emitted_by_fallback(
self,
bq_plugin_inst,
mock_write_client,
invocation_context,
dummy_arrow_schema,
):
"""An id with a matching part emits exactly one TOOL_PAUSED row."""
fc = types.FunctionCall(id="call-1", name="long_search", args={})
event = event_lib.Event(
author="agent",
content=types.Content(
role="model", parts=[types.Part(function_call=fc)]
),
long_running_tool_ids={"call-1"},
actions=event_actions_lib.EventActions(),
)
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
await bq_plugin_inst.on_event_callback(
invocation_context=invocation_context, event=event
)
await bq_plugin_inst.flush()
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
pauses = [r for r in rows if r["event_type"] == "TOOL_PAUSED"]
assert len(pauses) == 1
# ==============================================================================
# Observability controls (otel correlation, custom_metadata allowlist,
# column projection)
# ==============================================================================
class _FakeMetaEvent:
"""Minimal stand-in for an Event carrying custom_metadata."""
def __init__(self, custom_metadata=None):
self.custom_metadata = custom_metadata
def _make_offline_plugin(config):
"""Constructs a plugin without starting the BQ/network path."""
return bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
PROJECT_ID, DATASET_ID, config=config
)
# --- custom_metadata allowlist ---
def test_parse_custom_metadata_allowlist_exact_and_prefix():
exact, prefixes = (
bigquery_agent_analytics_plugin._parse_custom_metadata_allowlist(
["citation_metadata", "a2a:*", "tool:*"]
)
)
assert exact == frozenset({"citation_metadata"})
assert prefixes == ("a2a:", "tool:")
def test_parse_custom_metadata_allowlist_none():
exact, prefixes = (
bigquery_agent_analytics_plugin._parse_custom_metadata_allowlist(None)
)
assert exact == frozenset()
assert prefixes == ()
def test_custom_metadata_allowed_exact_and_prefix():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
custom_metadata_allowlist=["citation_metadata", "trace:*"]
)
)
assert plugin._custom_metadata_allowed("citation_metadata")
assert plugin._custom_metadata_allowed("trace:foo")
# a plain key is never treated as a prefix
assert not plugin._custom_metadata_allowed("citation")
assert not plugin._custom_metadata_allowed("other")
assert not plugin._custom_metadata_allowed(123)
def test_capture_custom_metadata_namespace_and_allowlist():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
custom_metadata_allowlist=["citation_metadata"]
)
)
event_data = bigquery_agent_analytics_plugin.EventData(
source_event=_FakeMetaEvent(
{"citation_metadata": {"c1": "sql1"}, "other": "drop"}
)
)
attrs: dict = {}
truncated = plugin._capture_custom_metadata(event_data, attrs)
assert truncated is False
assert attrs["custom_metadata"] == {"citation_metadata": {"c1": "sql1"}}
assert "other" not in attrs["custom_metadata"]
def test_capture_custom_metadata_redaction_does_not_set_flag():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
custom_metadata_allowlist=["secrets"]
)
)
event_data = bigquery_agent_analytics_plugin.EventData(
source_event=_FakeMetaEvent({"secrets": {"api_key": "abc", "ok": "v"}})
)
attrs: dict = {}
truncated = plugin._capture_custom_metadata(event_data, attrs)
# redaction returns [REDACTED] without flipping is_truncated
assert truncated is False
assert attrs["custom_metadata"]["secrets"]["api_key"] == "[REDACTED]"
assert attrs["custom_metadata"]["secrets"]["ok"] == "v"
def test_capture_custom_metadata_truncation_sets_flag():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
custom_metadata_allowlist=["big"], max_content_length=5
)
)
event_data = bigquery_agent_analytics_plugin.EventData(
source_event=_FakeMetaEvent({"big": "x" * 100})
)
attrs: dict = {}
truncated = plugin._capture_custom_metadata(event_data, attrs)
assert truncated is True
assert attrs["custom_metadata"]["big"].endswith("...[TRUNCATED]")
def test_capture_custom_metadata_non_allowlisted_absent():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
custom_metadata_allowlist=["citation_metadata"]
)
)
event_data = bigquery_agent_analytics_plugin.EventData(
source_event=_FakeMetaEvent({"unrelated": "v"})
)
attrs: dict = {}
assert plugin._capture_custom_metadata(event_data, attrs) is False
assert attrs == {}
def test_capture_custom_metadata_no_source_event():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
custom_metadata_allowlist=["x"]
)
)
attrs: dict = {}
assert (
plugin._capture_custom_metadata(
bigquery_agent_analytics_plugin.EventData(), attrs
)
is False
)
assert attrs == {}
def test_default_config_has_no_custom_metadata_capture():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
)
assert plugin._custom_metadata_exact == frozenset()
assert plugin._custom_metadata_prefixes == ()
# --- payload column projection ---
def test_validate_payload_column_denylist_accepts_payload_columns():
denied = bigquery_agent_analytics_plugin._validate_payload_column_denylist(
["content", "attributes"]
)
assert denied == frozenset({"content", "attributes"})
@pytest.mark.parametrize(
"bad",
["span_id", "trace_id", "timestamp", "event_type", "is_truncated", "nope"],
)
def test_validate_payload_column_denylist_rejects_protected_or_unknown(bad):
with pytest.raises(ValueError):
bigquery_agent_analytics_plugin._validate_payload_column_denylist([bad])
def test_plugin_construction_rejects_protected_denylist():
with pytest.raises(ValueError):
_make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
payload_column_denylist=["span_id"]
)
)
def test_project_schema_removes_denied_keeps_protected():
full = bigquery_agent_analytics_plugin._get_events_schema()
full_names = {f.name for f in full}
projected = bigquery_agent_analytics_plugin._project_schema(
full, frozenset({"content", "attributes"})
)
names = {f.name for f in projected}
assert "content" not in names and "attributes" not in names
for col in (
"timestamp",
"event_type",
"span_id",
"parent_span_id",
"is_truncated",
"latency_ms",
):
assert col in names
assert names == full_names - {"content", "attributes"}
def test_project_schema_to_arrow_consistency():
# schema-first: the Arrow schema derived from the projected BQ schema
# omits the denied column too.
projected = bigquery_agent_analytics_plugin._project_schema(
bigquery_agent_analytics_plugin._get_events_schema(),
frozenset({"content"}),
)
arrow = bigquery_agent_analytics_plugin.to_arrow_schema(projected)
assert "content" not in arrow.names
assert "span_id" in arrow.names
def test_project_view_columns_drops_denied_refs():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
payload_column_denylist=["attributes"]
)
)
exprs = [
"JSON_VALUE(attributes, '$.model') AS model",
"content AS request_content",
"CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms",
]
kept = plugin._project_view_columns(exprs)
assert "JSON_VALUE(attributes, '$.model') AS model" not in kept
assert "content AS request_content" in kept
assert any("latency_ms" in e for e in kept)
def test_project_view_columns_drops_content_and_latency_refs():
# view degradation is not attributes-only: content and latency_ms too.
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
payload_column_denylist=["content", "latency_ms"]
)
)
exprs = [
"JSON_QUERY(content, '$.response') AS response",
"CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms",
"JSON_VALUE(attributes, '$.model') AS model",
]
kept = plugin._project_view_columns(exprs)
assert kept == ["JSON_VALUE(attributes, '$.model') AS model"]
def test_project_view_columns_noop_without_denylist():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
)
exprs = ["JSON_VALUE(attributes, '$.model') AS model"]
assert plugin._project_view_columns(exprs) == exprs
# --- otel correlation ---
def test_enrich_attributes_captures_valid_ambient_otel_span(callback_context):
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
enable_otel_correlation=True
)
)
ctx = trace.SpanContext(
trace_id=0x1234567890ABCDEF1234567890ABCDEF,
span_id=0xFEEDFACECAFEBEEF,
is_remote=False,
trace_flags=trace.TraceFlags(trace.TraceFlags.SAMPLED),
)
fake_span = mock.Mock()
fake_span.get_span_context.return_value = ctx
with (
mock.patch.object(plugin, "_build_adk_envelope", return_value={}),
mock.patch.object(
bigquery_agent_analytics_plugin.trace,
"get_current_span",
return_value=fake_span,
),
):
attrs = plugin._enrich_attributes(
bigquery_agent_analytics_plugin.EventData(), callback_context
)
assert attrs["otel"]["span_id"] == format(0xFEEDFACECAFEBEEF, "016x")
assert attrs["otel"]["trace_id"] == format(
0x1234567890ABCDEF1234567890ABCDEF, "032x"
)
def test_enrich_attributes_no_otel_when_correlation_disabled(callback_context):
# enable_otel_correlation defaults to False: even with a valid ambient span,
# no attributes.otel is emitted (the feature is opt-in / off by default).
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
)
ctx = trace.SpanContext(
trace_id=0x1234567890ABCDEF1234567890ABCDEF,
span_id=0xFEEDFACECAFEBEEF,
is_remote=False,
trace_flags=trace.TraceFlags(trace.TraceFlags.SAMPLED),
)
fake_span = mock.Mock()
fake_span.get_span_context.return_value = ctx
with (
mock.patch.object(plugin, "_build_adk_envelope", return_value={}),
mock.patch.object(
bigquery_agent_analytics_plugin.trace,
"get_current_span",
return_value=fake_span,
),
):
attrs = plugin._enrich_attributes(
bigquery_agent_analytics_plugin.EventData(), callback_context
)
assert "otel" not in attrs
def test_enrich_attributes_no_otel_when_span_invalid(callback_context):
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
enable_otel_correlation=True
)
)
fake_span = mock.Mock()
fake_span.get_span_context.return_value = trace.INVALID_SPAN_CONTEXT
with (
mock.patch.object(plugin, "_build_adk_envelope", return_value={}),
mock.patch.object(
bigquery_agent_analytics_plugin.trace,
"get_current_span",
return_value=fake_span,
),
):
attrs = plugin._enrich_attributes(
bigquery_agent_analytics_plugin.EventData(), callback_context
)
assert "otel" not in attrs
class _FakeTable:
"""Minimal stand-in for a bigquery.Table for schema-upgrade tests."""
def __init__(self, schema, labels):
self.schema = schema
self.labels = labels
def test_schema_upgrade_adds_columns_when_denylist_relaxed():
# Table was created under a restrictive projection (missing content +
# attributes) but its version label is current. Relaxing the denylist must
# still add the now-desired columns instead of early-returning on the label.
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
)
full = bigquery_agent_analytics_plugin._get_events_schema()
plugin._schema = full # desired = full schema (denylist relaxed)
plugin.full_table_id = "p.d.t"
plugin.client = mock.Mock()
projected = [f for f in full if f.name not in ("content", "attributes")]
existing = _FakeTable(
schema=list(projected),
labels={
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY: (
bigquery_agent_analytics_plugin._SCHEMA_VERSION
)
},
)
plugin._maybe_upgrade_schema(existing)
plugin.client.update_table.assert_called_once()
names = {f.name for f in existing.schema}
assert "content" in names and "attributes" in names
def test_schema_upgrade_noop_when_current_and_complete():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
)
full = bigquery_agent_analytics_plugin._get_events_schema()
plugin._schema = full
plugin.full_table_id = "p.d.t"
plugin.client = mock.Mock()
existing = _FakeTable(
schema=list(full),
labels={
bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY: (
bigquery_agent_analytics_plugin._SCHEMA_VERSION
)
},
)
plugin._maybe_upgrade_schema(existing)
plugin.client.update_table.assert_not_called()
def test_attributes_denylist_with_custom_metadata_rejected():
with pytest.raises(ValueError):
_make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
payload_column_denylist=["attributes"],
custom_metadata_allowlist=["citation_metadata"],
)
)
def test_attributes_denylist_without_custom_metadata_ok():
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
payload_column_denylist=["attributes"]
)
)
assert "attributes" in plugin._denied_columns
def test_enrich_attributes_skips_otel_when_attributes_denied(callback_context):
plugin = _make_offline_plugin(
bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
enable_otel_correlation=True,
payload_column_denylist=["attributes"],
)
)
ctx = trace.SpanContext(
trace_id=0x1234567890ABCDEF1234567890ABCDEF,
span_id=0xFEEDFACECAFEBEEF,
is_remote=False,
trace_flags=trace.TraceFlags(trace.TraceFlags.SAMPLED),
)
fake_span = mock.Mock()
fake_span.get_span_context.return_value = ctx
with (
mock.patch.object(plugin, "_build_adk_envelope", return_value={}),
mock.patch.object(
bigquery_agent_analytics_plugin.trace,
"get_current_span",
return_value=fake_span,
),
):
attrs = plugin._enrich_attributes(
bigquery_agent_analytics_plugin.EventData(), callback_context
)
assert "otel" not in attrs
@pytest.mark.asyncio
async def test_content_parts_denied_disables_gcs_offload(
mock_write_client,
callback_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_storage_client,
):
# denying content_parts (which holds the offload object reference)
# must disable GCS offload, otherwise the payload is uploaded with no
# retained reference (leak + cost).
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
gcs_bucket_name="test-bucket",
payload_column_denylist=["content_parts"],
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started(
storage_client=mock_storage_client.return_value
)
assert plugin.offloader is None
mock_blob = (
mock_storage_client.return_value.bucket.return_value.blob.return_value
)
large_text = "A" * (32 * 1024 + 1)
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text=large_text)])],
)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await plugin.flush()
mock_blob.upload_from_string.assert_not_called()
@pytest.mark.asyncio
async def test_both_payload_columns_denied_skips_parse_and_offload(
mock_write_client,
callback_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_storage_client,
):
# with both content and content_parts denied, parsing is skipped
# entirely -- no inline summary, no parts, and no GCS upload work.
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
gcs_bucket_name="test-bucket",
payload_column_denylist=["content", "content_parts"],
)
async with managed_plugin(
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
) as plugin:
await plugin._ensure_started(
storage_client=mock_storage_client.return_value
)
assert plugin.offloader is None
mock_blob = (
mock_storage_client.return_value.bucket.return_value.blob.return_value
)
large_text = "A" * (32 * 1024 + 1)
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text=large_text)])],
)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await plugin.flush()
mock_blob.upload_from_string.assert_not_called()
Reference in New Issue View Git Blame Copy Permalink