Skip to content

Commit 7c4ef94

Browse files
committed
fix: add log_state_changes toggle and deduplicate STATE_DELTA logging
- Add log_state_changes config flag (default False) to BigQueryLoggerConfig for explicit opt-in to STATE_DELTA logging via the existing after_tool_callback inline path - Add event ID dedup guard in Runner._exec_with_plugin to prevent the same event from triggering on_state_change_callback twice - Add tests for toggle enabled and disabled behavior
1 parent 87c46a7 commit 7c4ef94

3 files changed

Lines changed: 31 additions & 19 deletions

File tree

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,8 @@ class BigQueryLoggerConfig:
410410
# Format: "location.connection_id" (e.g. "us.my-connection")
411411
connection_id: Optional[str] = None
412412

413+
# Toggle for state change (STATE_DELTA) logging via on_state_change_callback
414+
log_state_changes: bool = False
413415
# Toggle for session metadata (e.g. gchat thread-id)
414416
log_session_metadata: bool = True
415417
# Static custom tags (e.g. {"agent_role": "sales"})
@@ -2510,7 +2512,7 @@ async def after_tool_callback(
25102512
parent_span_id_override=parent_span_id,
25112513
)
25122514

2513-
if tool_context.actions.state_delta:
2515+
if tool_context.actions.state_delta and self.config.log_state_changes:
25142516
await self._log_event(
25152517
"STATE_DELTA",
25162518
tool_context,

src/google/adk/runners.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -774,6 +774,7 @@ async def _exec_with_plugin(
774774
# transcription event.
775775
buffered_events: list[Event] = []
776776
is_transcribing: bool = False
777+
notified_state_change_event_ids: set[str] = set()
777778

778779
async with Aclosing(execute_fn(invocation_context)) as agen:
779780
async for event in agen:
@@ -845,7 +846,11 @@ async def _exec_with_plugin(
845846
yield final_event
846847

847848
# Step 3b: Notify plugins of state changes, if any.
848-
if final_event.actions.state_delta:
849+
if (
850+
final_event.actions.state_delta
851+
and final_event.id not in notified_state_change_event_ids
852+
):
853+
notified_state_change_event_ids.add(final_event.id)
849854
from .agents.callback_context import CallbackContext
850855

851856
await plugin_manager.run_on_state_change_callback(

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

Lines changed: 22 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1563,6 +1563,8 @@ async def test_after_tool_callback_logs_correctly(
15631563
async def test_after_tool_callback_state_delta_logging(
15641564
self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema
15651565
):
1566+
"""STATE_DELTA is logged inline when log_state_changes is True."""
1567+
bq_plugin_inst.config.log_state_changes = True
15661568
mock_tool = mock.create_autospec(
15671569
base_tool_lib.BaseTool, instance=True, spec_set=True
15681570
)
@@ -1608,29 +1610,32 @@ async def test_after_tool_callback_state_delta_logging(
16081610
assert state_delta_event["content"] is None
16091611

16101612
@pytest.mark.asyncio
1611-
async def test_on_state_change_callback_logs_correctly(
1612-
self,
1613-
bq_plugin_inst,
1614-
mock_write_client,
1615-
callback_context,
1616-
dummy_arrow_schema,
1613+
async def test_after_tool_callback_state_delta_disabled(
1614+
self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema
16171615
):
1618-
state_delta = {"key": "value", "new_key": 123}
1619-
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
1620-
await bq_plugin_inst.on_state_change_callback(
1621-
callback_context=callback_context, state_delta=state_delta
1616+
"""STATE_DELTA is not logged when log_state_changes is False (default)."""
1617+
mock_tool = mock.create_autospec(
1618+
base_tool_lib.BaseTool, instance=True, spec_set=True
1619+
)
1620+
type(mock_tool).name = mock.PropertyMock(return_value="StateTool")
1621+
type(mock_tool).description = mock.PropertyMock(return_value="Sets state")
1622+
1623+
tool_context.actions.state_delta["new_key"] = "new_value"
1624+
1625+
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
1626+
await bq_plugin_inst.after_tool_callback(
1627+
tool=mock_tool,
1628+
tool_args={"arg1": "val1"},
1629+
tool_context=tool_context,
1630+
result={"res": "success"},
16221631
)
16231632
await asyncio.sleep(0.01)
1633+
1634+
# Only TOOL_COMPLETED should be logged, not STATE_DELTA
16241635
log_entry = await _get_captured_event_dict_async(
16251636
mock_write_client, dummy_arrow_schema
16261637
)
1627-
_assert_common_fields(log_entry, "STATE_DELTA")
1628-
# content should be None (as raw_content was not passed)
1629-
assert log_entry["content"] is None
1630-
1631-
# state_delta should be in attributes
1632-
attributes = json.loads(log_entry["attributes"])
1633-
assert attributes["state_delta"] == state_delta
1638+
assert log_entry["event_type"] == "TOOL_COMPLETED"
16341639

16351640
@pytest.mark.asyncio
16361641
async def test_log_event_with_session_metadata(

0 commit comments

Comments
 (0)