fix(plugins): sanitize the whole BigQuery analytics attributes tree (v1)

Before, `_enrich_attributes` ran the redacting sanitizer over `usage_metadata`,
`cache_metadata` and `session.state`, but copied `extra_attributes` and
`custom_tags` in untouched and then serialized the result. A session state
delta or a configured tag holding a key such as `api_key` or `refresh_token`
therefore reached the `attributes` column in the clear.

The assembled tree now goes through `_recursive_smart_truncate` once more,
immediately before serialization, so every value in the column has seen the
sensitive-key redaction regardless of which producer put it there.

That pass walks objects the sanitizer never saw before, which exposed a second
problem. `_recursive_smart_truncate` detects cycles by object id, and an object
whose `model_dump`, `dict` or `to_dict` returns a freshly built wrapper defeats
that, because the walk never sees the same id twice and keeps descending. It
now stops at a depth of 50 and substitutes `[MAX_DEPTH_EXCEEDED]`.

Depth is only half a bound, because it says nothing about width. An object that
hands back two fresh children on every access fills the 50 levels beneath it
with tens of millions of nodes, and one such value in a state delta held the
event loop for over a minute in testing. The walk now also carries a budget of
100,000 nodes for the whole invocation and replaces the remainder with
`[SANITIZE_BUDGET_EXCEEDED]`. A directly redacted key spends budget too, so a
wide `temp:`-scoped mapping cannot slip past the bound, and each container loop
stops at the budget rather than emitting one sentinel per remaining element.

The `mock_agent` test fixture now returns itself from `root_agent`, matching
the pattern already used elsewhere in the file, so `root_agent_name` is a real
name instead of a bare mock object.

Behaviour change: the `attributes` column can now contain `[REDACTED]` and
truncation markers where it previously carried raw values, and `is_truncated`
is set when the attributes pass alters anything. A value already truncated by
`_enrich_attributes` gains a second `...[TRUNCATED]` marker.

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

@@ -280,9 +280,23 @@ _REDACTED_URI = "[REDACTED_SENSITIVE_URI]"
 # A URI longer than this is replaced wholesale rather than parsed.
 _MAX_URI_LENGTH = 8192
 
+# Deepest nesting _recursive_smart_truncate walks before replacing a value.
+_MAX_SANITIZE_DEPTH = 50
+
+# Total nodes one sanitizer invocation may visit. Depth and per-string size
+# are bounded, but width was not: a million-element list, or an object that
+# manufactures two fresh children per access, walks tens of millions of
+# nodes inside the depth cap. The remainder is replaced with a sentinel and
+# the row is flagged truncated.
+_MAX_SANITIZE_NODES = 100_000
+
 
 def _recursive_smart_truncate(
-    obj: Any, max_len: int, seen: Optional[set[int]] = None
+    obj: Any,
+    max_len: int,
+    seen: Optional[set[int]] = None,
+    depth: int = 0,
+    budget: Optional[list[int]] = None,
 ) -> tuple[Any, bool]:
   """Recursively truncates string values within a dict or list.
 
@@ -293,12 +307,29 @@ def _recursive_smart_truncate(
       obj: The object to truncate.
       max_len: Maximum length for string values.
       seen: Set of object IDs visited in the current recursion stack.
+      depth: Current recursion depth.
+      budget: Single-element list holding the nodes left in the shared work
+        budget for this invocation.
 
   Returns:
       A tuple of (truncated_object, is_truncated).
   """
   if seen is None:
     seen = set()
+  if budget is None:
+    budget = [_MAX_SANITIZE_NODES]
+  budget[0] -= 1
+  if budget[0] < 0:
+    return "[SANITIZE_BUDGET_EXCEEDED]", True
+
+  # The id()-based cycle detection below cannot catch an object graph that
+  # manufactures a new object on every duck-typed access, which is what any
+  # object whose model_dump()/dict()/to_dict() returns a fresh wrapper does.
+  # Such a graph recurses until the interpreter's own limit. The replacement
+  # discards real data, so unlike "[CIRCULAR_REFERENCE]" it reports
+  # truncation.
+  if depth >= _MAX_SANITIZE_DEPTH:
+    return "[MAX_DEPTH_EXCEEDED]", True
 
   obj_id = id(obj)
   if obj_id in seen:
@@ -327,13 +358,26 @@ def _recursive_smart_truncate(
       # but explicit loop is fine for clarity given recursive nature.
       new_dict = {}
       for k, v in obj.items():
+        # Stop iterating once the budget is exhausted. Recursing on every
+        # remaining entry still did work proportional to the input and
+        # produced one sentinel per entry; a single remainder sentinel
+        # stands in for everything dropped.
+        if budget[0] <= 0:
+          new_dict["[SANITIZE_BUDGET_EXCEEDED]"] = "[SANITIZE_BUDGET_EXCEEDED]"
+          truncated_any = True
+          break
         if isinstance(k, str):
           k_lower = k.lower()
           if k_lower in _SENSITIVE_KEYS or k_lower.startswith("temp:"):
+            # A directly redacted entry costs budget too, otherwise a wide
+            # "temp:" mapping bypasses the bound entirely.
+            budget[0] -= 1
             new_dict[k] = "[REDACTED]"
             continue
 
-        val, trunc = _recursive_smart_truncate(v, max_len, seen)
+        val, trunc = _recursive_smart_truncate(
+            v, max_len, seen, depth + 1, budget
+        )
         if trunc:
           truncated_any = True
         new_dict[k] = val
@@ -343,7 +387,14 @@ def _recursive_smart_truncate(
       new_list = []
       # Explicit loop to handle flag propagation
       for i in obj:
-        val, trunc = _recursive_smart_truncate(i, max_len, seen)
+        # Same bound as the mapping loop.
+        if budget[0] <= 0:
+          new_list.append("[SANITIZE_BUDGET_EXCEEDED]")
+          truncated_any = True
+          break
+        val, trunc = _recursive_smart_truncate(
+            i, max_len, seen, depth + 1, budget
+        )
         if trunc:
           truncated_any = True
         new_list.append(val)
@@ -351,23 +402,31 @@ def _recursive_smart_truncate(
     elif dataclasses.is_dataclass(obj) and not isinstance(obj, type):
       # Manually iterate fields to preserve 'seen' context, avoiding dataclasses.asdict recursion
       as_dict = {f.name: getattr(obj, f.name) for f in dataclasses.fields(obj)}
-      return _recursive_smart_truncate(as_dict, max_len, seen)
+      return _recursive_smart_truncate(
+          as_dict, max_len, seen, depth + 1, budget
+      )
     elif hasattr(obj, "model_dump") and callable(obj.model_dump):
       # Pydantic v2
       try:
-        return _recursive_smart_truncate(obj.model_dump(), max_len, seen)
+        return _recursive_smart_truncate(
+            obj.model_dump(), max_len, seen, depth + 1, budget
+        )
       except Exception:
         pass
     elif hasattr(obj, "dict") and callable(obj.dict):
       # Pydantic v1
       try:
-        return _recursive_smart_truncate(obj.dict(), max_len, seen)
+        return _recursive_smart_truncate(
+            obj.dict(), max_len, seen, depth + 1, budget
+        )
       except Exception:
         pass
     elif hasattr(obj, "to_dict") and callable(obj.to_dict):
       # Common pattern for custom objects
       try:
-        return _recursive_smart_truncate(obj.to_dict(), max_len, seen)
+        return _recursive_smart_truncate(
+            obj.to_dict(), max_len, seen, depth + 1, budget
+        )
       except Exception:
         pass
     elif obj is None or isinstance(obj, (int, float, bool)):
@@ -2972,6 +3031,15 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
     latency_json = self._extract_latency(event_data)
     attributes = self._enrich_attributes(event_data, callback_context)
 
+    # Final pass over the complete assembled tree. _enrich_attributes copies
+    # extra_attributes (which carries session state deltas) and custom_tags in
+    # untouched, so this is the only point at which every value is guaranteed
+    # to have seen the sensitive-key redaction.
+    attributes, attrs_truncated = _recursive_smart_truncate(
+        attributes, self.config.max_content_length
+    )
+    is_truncated = is_truncated or attrs_truncated
+
     # Serialize attributes to JSON string
     try:
       attributes_json = json.dumps(attributes)

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

@@ -74,6 +74,9 @@ def mock_agent():
   # Mock the 'name' property
   type(mock_a).name = mock.PropertyMock(return_value="MyTestAgent")
   type(mock_a).instruction = mock.PropertyMock(return_value="Test Instruction")
+  # root_agent returns itself (no parent), so root_agent.name is a real name
+  # rather than a bare mock.
+  mock_a.root_agent = mock_a
   return mock_a
 
 
@@ -426,6 +429,105 @@ def test_recursive_smart_truncate_redaction():
   assert truncated["nested"]["normal"] == "value"
 
 
+def test_recursive_smart_truncate_bounds_self_generating_objects():
+  """An object that makes a fresh wrapper on each access stops at the cap."""
+
+  class Endless:
+    """to_dict() hands back a new object every time, so ids never repeat."""
+
+    def to_dict(self):
+      return {"next": Endless()}
+
+  truncated, is_truncated = (
+      bigquery_agent_analytics_plugin._recursive_smart_truncate(
+          {"root": Endless()}, 1000
+      )
+  )
+
+  assert is_truncated
+  flattened = json.dumps(truncated)
+  assert "[MAX_DEPTH_EXCEEDED]" in flattened
+
+
+def test_recursive_smart_truncate_bounds_branching_self_generating_objects():
+  """The depth cap alone does not bound an object that branches.
+
+  Every access hands back two fresh children, so ids never repeat and the
+  50-level cap still leaves tens of millions of nodes below it. The node
+  budget is what stops the walk.
+  """
+  max_nodes = bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES
+  # Stop manufacturing children well past the budget, so a walk that is not
+  # bounded fails this test in a second or two rather than the minute-plus
+  # it takes to exhaust the depth cap on its own.
+  safety_limit = 5 * max_nodes
+  visited = 0
+
+  class Branching:
+    """to_dict() hands back two new objects every time."""
+
+    def to_dict(self):
+      nonlocal visited
+      visited += 1
+      if visited > safety_limit:
+        return {"left": "stopped", "right": "stopped"}
+      return {"left": Branching(), "right": Branching()}
+
+  truncated, is_truncated = (
+      bigquery_agent_analytics_plugin._recursive_smart_truncate(
+          {"root": Branching()}, 1000
+      )
+  )
+
+  assert is_truncated
+  assert visited <= max_nodes
+  assert "[SANITIZE_BUDGET_EXCEEDED]" in json.dumps(truncated)
+
+
+def test_recursive_smart_truncate_elides_the_remainder_of_a_wide_value():
+  """A wide value stops at the budget and leaves one remainder sentinel."""
+  max_nodes = bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES
+  wide = list(range(max_nodes * 2))
+
+  truncated, is_truncated = (
+      bigquery_agent_analytics_plugin._recursive_smart_truncate(
+          {"wide": wide}, 1000
+      )
+  )
+
+  assert is_truncated
+  assert truncated["wide"][-1] == "[SANITIZE_BUDGET_EXCEEDED]"
+  assert truncated["wide"].count("[SANITIZE_BUDGET_EXCEEDED]") == 1
+  # Bounded output: budget entries plus the single remainder sentinel.
+  assert len(truncated["wide"]) <= max_nodes + 1
+
+
+def test_recursive_smart_truncate_charges_directly_redacted_keys():
+  """Redacted keys cost budget, so a wide temp: mapping cannot bypass it."""
+  max_nodes = bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES
+  wide_temp = {f"temp:{i}": i for i in range(max_nodes * 2)}
+
+  truncated, is_truncated = (
+      bigquery_agent_analytics_plugin._recursive_smart_truncate(wide_temp, 1000)
+  )
+
+  assert is_truncated
+  assert len(truncated) <= max_nodes + 1
+  assert "[SANITIZE_BUDGET_EXCEEDED]" in truncated
+
+
+def test_recursive_smart_truncate_keeps_ordinary_nesting():
+  """Nesting well inside the cap is copied through untouched."""
+  obj = {"a": {"b": {"c": {"d": "leaf"}}}}
+
+  truncated, is_truncated = (
+      bigquery_agent_analytics_plugin._recursive_smart_truncate(obj, 1000)
+  )
+
+  assert not is_truncated
+  assert truncated == obj
+
+
 class TestBigQueryAgentAnalyticsPlugin:
   """Tests for the BigQueryAgentAnalyticsPlugin."""
 
@@ -1754,6 +1856,44 @@ class TestBigQueryAgentAnalyticsPlugin:
     attributes = json.loads(log_entry["attributes"])
     assert attributes["custom_tags"] == custom_tags
 
+  @pytest.mark.asyncio
+  async def test_custom_tags_and_extra_attributes_are_redacted(
+      self,
+      bq_plugin_inst,
+      mock_write_client,
+      callback_context,
+      dummy_arrow_schema,
+  ):
+    """Sensitive keys are redacted wherever they sit in the attributes tree."""
+    bq_plugin_inst.config.custom_tags = {
+        "env": "prod",
+        "api_key": "sk-live-should-not-be-stored",
+    }
+
+    await bq_plugin_inst._log_event(
+        "TEST_EVENT",
+        callback_context,
+        raw_content="test content",
+        event_data=bigquery_agent_analytics_plugin.EventData(
+            extra_attributes={
+                "tool": "search",
+                "nested": {"refresh_token": "rt-should-not-be-stored"},
+            }
+        ),
+    )
+    await asyncio.sleep(0.01)
+    log_entry = await _get_captured_event_dict_async(
+        mock_write_client, dummy_arrow_schema
+    )
+
+    attributes_json = log_entry["attributes"]
+    assert "should-not-be-stored" not in attributes_json
+    attributes = json.loads(attributes_json)
+    assert attributes["custom_tags"]["api_key"] == "[REDACTED]"
+    assert attributes["custom_tags"]["env"] == "prod"
+    assert attributes["nested"]["refresh_token"] == "[REDACTED]"
+    assert attributes["tool"] == "search"
+
   @pytest.mark.asyncio
   async def test_on_model_error_callback_logs_correctly(
       self,