From 7c96e6b8318c579e223c16d48b125a791c78b394 Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 23:08:05 +0000 Subject: [PATCH 1/4] fix(plugins): fail closed when the BigQuery analytics content formatter raises (v1) Before, if the configured `content_formatter` raised, the plugin logged a warning and wrote the original, unformatted payload to BigQuery. The formatter is the operator's redaction boundary, so a failing formatter published exactly the content it was installed to remove. Now a formatter that raises causes the sentinel string `[FORMATTER_FAILED]` to be written in place of the payload. The warning it logs is a constant message plus the event type; the exception text is deliberately not interpolated, because an exception message or traceback can itself quote the protected content. Behaviour change: rows whose formatter failed now carry `[FORMATTER_FAILED]` in the `content` column instead of the original payload. --- .../plugins/bigquery_agent_analytics_plugin.py | 17 +++++++++++++++-- .../test_bigquery_agent_analytics_plugin.py | 7 ++++--- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 3a09fc942f8..13484511464 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -268,6 +268,10 @@ def _get_tool_origin( "password", }) +# Written to the content column in place of the payload when the configured +# content_formatter raises. +_FORMATTER_FAILED_SENTINEL = "[FORMATTER_FAILED]" + def _recursive_smart_truncate( obj: Any, max_len: int, seen: Optional[set[int]] = None @@ -2886,8 +2890,17 @@ async def _log_event( if self.config.content_formatter: try: raw_content = self.config.content_formatter(raw_content, event_type) - except Exception as e: - logger.warning("Content formatter failed: %s", e) + except Exception: + # Fail closed. The formatter is the operator's redaction boundary, so + # a failure must not fall back to the unformatted payload. The message + # is constant on purpose: an exception's own text or traceback can + # quote the content the formatter exists to remove. + logger.warning( + "Content formatter failed for event %s; writing a sentinel" + " instead of the original content.", + event_type, + ) + raw_content = _FORMATTER_FAILED_SENTINEL trace_id, span_id, parent_span_id = self._resolve_ids( event_data, callback_context diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 5719adf2b46..4b715400b4f 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -685,7 +685,7 @@ async def test_content_formatter_error( dummy_arrow_schema, mock_asyncio_to_thread, ): - """Test content formatter error handling.""" + """A formatter that raises must not let the original content through.""" _ = mock_auth_default _ = mock_bq_client @@ -710,8 +710,9 @@ def error_formatter(content, event_type): 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"}' + # The formatter is the redaction boundary, so its failure fails closed: + # a sentinel is written and the payload never reaches the row. + assert log_entry["content"] == "[FORMATTER_FAILED]" @pytest.mark.asyncio async def test_max_content_length( From 1ba04e1d81287a2512714370da128f7af62f064e Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 23:10:16 +0000 Subject: [PATCH 2/4] fix(plugins): strip credentials from external URIs in BigQuery analytics (v1) Before, a `file_data` part's URI was written into the `content_parts` column exactly as supplied. A signed GCS or HTTP URL carries its signature in the query string and a URI may carry `user:password@` userinfo, so the row stored a working bearer credential that anyone with read access to the table could replay. Now the URI goes through a sanitizer before it is recorded. A URI with userinfo, one longer than 8192 characters, or one that does not parse is replaced with `[REDACTED_SENSITIVE_URI]`. Otherwise the query string and fragment are dropped and the scheme, host and path are kept, which is what identifies the object. A URI with none of those components is stored unchanged. Behaviour change: URIs in `content_parts` are no longer resolvable by copy-paste when they were signed, and a row whose URI was altered now has `is_truncated` set. --- .../bigquery_agent_analytics_plugin.py | 53 +++++++++- .../test_bigquery_agent_analytics_plugin.py | 97 +++++++++++++++++++ 2 files changed, 149 insertions(+), 1 deletion(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 13484511464..b17afd91480 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -43,6 +43,8 @@ from typing import Callable from typing import Optional from typing import TYPE_CHECKING +from urllib.parse import urlsplit +from urllib.parse import urlunsplit import uuid import weakref @@ -272,6 +274,12 @@ def _get_tool_origin( # content_formatter raises. _FORMATTER_FAILED_SENTINEL = "[FORMATTER_FAILED]" +# Written in place of an external URI that cannot be stored safely. +_REDACTED_URI = "[REDACTED_SENSITIVE_URI]" + +# A URI longer than this is replaced wholesale rather than parsed. +_MAX_URI_LENGTH = 8192 + def _recursive_smart_truncate( obj: Any, max_len: int, seen: Optional[set[int]] = None @@ -1378,6 +1386,44 @@ def _truncate(self, text: str) -> tuple[str, bool]: ) return text, False + def _sanitize_external_uri( + self, uri: Optional[str] + ) -> tuple[Optional[str], bool]: + """Removes credential-bearing components from a caller-supplied URI. + + A signed GCS or HTTP URL carries its signature in the query string, and any + URI may carry ``user:password@`` userinfo. Both are bearer credentials, so + neither belongs in a stored analytics row. The scheme, host and path are + kept, which is what identifies the object. + + Args: + uri: The URI to sanitize, or None. + + Returns: + A tuple of (sanitized_uri, content_removed). + """ + if uri is None: + return None, False + if not isinstance(uri, str): + return _REDACTED_URI, True + if len(uri) > _MAX_URI_LENGTH: + return _REDACTED_URI, True + try: + parsed = urlsplit(uri) + has_userinfo = parsed.username is not None or parsed.password is not None + except ValueError: + return _REDACTED_URI, True + if has_userinfo: + # Userinfo is a credential-bearing surface by definition; do not try to + # keep the username while guessing whether it is sensitive. + return _REDACTED_URI, True + if not parsed.query and not parsed.fragment: + return uri, False + return ( + urlunsplit((parsed.scheme, parsed.netloc, parsed.path, "", "")), + True, + ) + async def _parse_content_object( self, content: types.Content | types.Part ) -> tuple[str, list[dict[str, Any]], bool]: @@ -1401,7 +1447,12 @@ async def _parse_content_object( # CASE A: It is already a URI (e.g. from user input) if hasattr(part, "file_data") and part.file_data: part_data["storage_mode"] = "EXTERNAL_URI" - part_data["uri"] = part.file_data.file_uri + safe_uri, uri_content_removed = self._sanitize_external_uri( + part.file_data.file_uri + ) + part_data["uri"] = safe_uri + if uri_content_removed: + is_truncated = True part_data["mime_type"] = part.file_data.mime_type # CASE B: It is Binary/Inline Data (Image/Blob) diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 4b715400b4f..dd911dd0281 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -7596,6 +7596,103 @@ async def test_no_offloader_falls_back_to_truncate(self): assert "TRUNCATED" in parts[0]["text"] +# ================================================================ +# TEST CLASS: external URI sanitization +# ================================================================ +class TestExternalUriSanitization: + """Tests that credentials are removed from stored external URIs.""" + + def _parser(self): + return bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=-1, + ) + + @pytest.mark.asyncio + async def test_signed_gcs_uri_loses_its_signature(self): + """A signed GCS URL is stored without its query string.""" + signed = ( + "https://storage.googleapis.com/bucket/report.pdf" + "?X-Goog-Algorithm=GOOG4-RSA-SHA256" + "&X-Goog-Signature=deadbeefcafe" + ) + content = types.Content( + parts=[ + types.Part( + file_data=types.FileData( + file_uri=signed, mime_type="application/pdf" + ) + ) + ] + ) + + _, parts, is_truncated = await self._parser()._parse_content_object(content) + + assert parts[0]["storage_mode"] == "EXTERNAL_URI" + assert parts[0]["uri"] == "https://storage.googleapis.com/bucket/report.pdf" + assert "X-Goog-Signature" not in parts[0]["uri"] + assert "deadbeefcafe" not in parts[0]["uri"] + assert is_truncated + + @pytest.mark.asyncio + async def test_uri_with_userinfo_is_replaced(self): + """A URI carrying user:password@ is replaced wholesale.""" + content = types.Content( + parts=[ + types.Part( + file_data=types.FileData( + file_uri="https://alice:hunter2@example.com/a.png", + mime_type="image/png", + ) + ) + ] + ) + + _, parts, _ = await self._parser()._parse_content_object(content) + + assert parts[0]["uri"] == "[REDACTED_SENSITIVE_URI]" + + @pytest.mark.asyncio + async def test_plain_uri_is_unchanged(self): + """A URI with no query, fragment or userinfo is stored verbatim.""" + content = types.Content( + parts=[ + types.Part( + file_data=types.FileData( + file_uri="gs://bucket/photo.png", mime_type="image/png" + ) + ) + ] + ) + + _, parts, is_truncated = await self._parser()._parse_content_object(content) + + assert parts[0]["uri"] == "gs://bucket/photo.png" + assert not is_truncated + + @pytest.mark.asyncio + async def test_missing_uri_stays_none(self): + """file_data with no file_uri still records None, not a sentinel.""" + content = types.Content( + parts=[types.Part(file_data=types.FileData(mime_type="image/png"))] + ) + + _, parts, _ = await self._parser()._parse_content_object(content) + + assert parts[0]["uri"] is None + + def test_overlong_uri_is_replaced(self): + """A URI beyond the inspection bound is not parsed at all.""" + long_uri = "https://example.com/" + "a" * 9000 + + safe_uri, removed = self._parser()._sanitize_external_uri(long_uri) + + assert safe_uri == "[REDACTED_SENSITIVE_URI]" + assert removed + + # ================================================================ # TEST CLASS: AGENT_RESPONSE logging (Issue #87) # ================================================================ From ef4b10f12779b94cd5fc78166c19fbe8252ca766 Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 23:45:04 +0000 Subject: [PATCH 3/4] 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. --- .../bigquery_agent_analytics_plugin.py | 82 +++++++++- .../test_bigquery_agent_analytics_plugin.py | 140 ++++++++++++++++++ 2 files changed, 215 insertions(+), 7 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index b17afd91480..6a0dee03031 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -280,9 +280,23 @@ def _get_tool_origin( # 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 @@ async def _log_event( 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) diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index dd911dd0281..aaaf51eaea1 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/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 @@ async def test_log_event_with_custom_tags( 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, From 4c34f25ff9a9b47b3ce5ede48ea9ec2c19f530bb Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 23:58:57 +0000 Subject: [PATCH 4/4] fix(plugins): build BigQuery analytics GCS paths from call-local ids (v1) Before, `_log_event` assigned the event's trace and span ids onto the single shared `HybridContentParser` instance and then awaited the parse. The offload path was built from those instance fields after the await, so a second event arriving in the meantime replaced them and the first event's media was written under the second event's prefix. Object names were also built from the part index alone, so two messages in one request, or two events offloading at the same moment, produced the same name and one overwrote the other. `parse` and `_parse_content_object` now take the trace and span ids as keyword arguments, defaulting to the instance fields so existing callers are unaffected. Each `parse` call generates a unique id that goes into the object name, along with the index of the message within the request. `_log_event` passes the ids instead of assigning them, so the shared parser is no longer mutated. Behaviour change: offloaded GCS object names gain a unique component and a message index, so they are no longer predictable from the trace id, span id and part index. --- .../bigquery_agent_analytics_plugin.py | 104 ++++++++++++++--- .../test_bigquery_agent_analytics_plugin.py | 105 ++++++++++++++++-- 2 files changed, 184 insertions(+), 25 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 6a0dee03031..97eb5845ab2 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -1484,9 +1484,36 @@ def _sanitize_external_uri( ) async def _parse_content_object( - self, content: types.Content | types.Part + self, + content: types.Content | types.Part, + *, + trace_id: Optional[str] = None, + span_id: Optional[str] = None, + parse_uid: str = "", + content_ordinal: int = 0, ) -> tuple[str, list[dict[str, Any]], bool]: - """Parses a Content or Part object into summary text and content parts.""" + """Parses a Content or Part object into summary text and content parts. + + Args: + content: The Content or Part to parse. + trace_id: Trace id of the calling event. GCS object paths are built + from this argument rather than the instance field, because the + parser is shared across concurrent events and an await inside this + method can resume under another event's identity. Falls back to + the instance field. + span_id: Span id of the calling event, with the same rationale. + parse_uid: Unique per parse() call. Disambiguates object names across + concurrent events. Generated here when not supplied. + content_ordinal: Index of this Content within the calling request. The + part index restarts at zero for each Content, so without this two + messages in one request collide on the same object name. + + Returns: + A tuple of (summary_text, content_parts, is_truncated). + """ + trace_id = trace_id if trace_id is not None else self.trace_id + span_id = span_id if span_id is not None else self.span_id + parse_uid = parse_uid or uuid.uuid4().hex content_parts = [] is_truncated = False summary_text = [] @@ -1518,7 +1545,10 @@ async def _parse_content_object( elif hasattr(part, "inline_data") and part.inline_data: if self.offloader: ext = mimetypes.guess_extension(part.inline_data.mime_type) or ".bin" - path = f"{datetime.now().date()}/{self.trace_id}/{self.span_id}_p{idx}{ext}" + path = ( + f"{datetime.now().date()}/{trace_id}/{span_id}_{parse_uid}" + f"_c{content_ordinal}_p{idx}{ext}" + ) try: uri = await self.offloader.upload_content( part.inline_data.data, part.inline_data.mime_type, path @@ -1557,7 +1587,10 @@ async def _parse_content_object( if self.offloader and (exceeds_inline_byte_limit or exceeds_char_limit): # Text is too big, treat as file - path = f"{datetime.now().date()}/{self.trace_id}/{self.span_id}_p{idx}.txt" + path = ( + f"{datetime.now().date()}/{trace_id}/{span_id}_{parse_uid}" + f"_c{content_ordinal}_p{idx}.txt" + ) try: uri = await self.offloader.upload_content( part.text, "text/plain", path @@ -1605,8 +1638,32 @@ async def _parse_content_object( return summary_str, content_parts, is_truncated - async def parse(self, content: Any) -> tuple[Any, list[dict[str, Any]], bool]: - """Parses content into JSON payload and content parts, potentially offloading to GCS.""" + async def parse( + self, + content: Any, + *, + trace_id: Optional[str] = None, + span_id: Optional[str] = None, + ) -> tuple[Any, list[dict[str, Any]], bool]: + """Parses content into JSON payload and content parts, potentially offloading to GCS. + + Args: + content: The content to parse. + trace_id: Trace id of the calling event, used to build GCS object + paths. Pass it per call: the parser instance is shared across + concurrent events, so a path built from the mutable instance field + can pick up another event's identity across an await. Falls back + to the instance field. + span_id: Span id of the calling event, with the same rationale. + + Returns: + A tuple of (json_payload, content_parts, is_truncated). + """ + trace_id = trace_id if trace_id is not None else self.trace_id + span_id = span_id if span_id is not None else self.span_id + # Unique per parse() call, so two events offloading at the same time + # cannot produce the same object name. + parse_uid = uuid.uuid4().hex json_payload = {} content_parts = [] is_truncated = False @@ -1622,9 +1679,15 @@ def process_text(t: str) -> tuple[str, bool]: if isinstance(content.contents, list) else [content.contents] ) - for c in contents: + for content_idx, c in enumerate(contents): role = getattr(c, "role", "unknown") - summary, parts, trunc = await self._parse_content_object(c) + summary, parts, trunc = await self._parse_content_object( + c, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + content_ordinal=content_idx, + ) if trunc: is_truncated = True content_parts.extend(parts) @@ -1642,14 +1705,25 @@ def process_text(t: str) -> tuple[str, bool]: is_truncated = True json_payload["system_prompt"] = truncated_si else: - summary, parts, trunc = await self._parse_content_object(si) + summary, parts, trunc = await self._parse_content_object( + si, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + content_ordinal=len(contents), + ) if trunc: is_truncated = True content_parts.extend(parts) json_payload["system_prompt"] = summary elif isinstance(content, (types.Content, types.Part)): - summary, parts, trunc = await self._parse_content_object(content) + summary, parts, trunc = await self._parse_content_object( + content, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + ) return {"text_summary": summary}, parts, trunc elif isinstance(content, (dict, list)): @@ -3020,11 +3094,13 @@ async def _log_event( logger.warning("Parser not initialized; skipping event %s.", event_type) return - # Update parser's trace/span IDs for GCS pathing (reuse instance) - self.parser.trace_id = trace_id or "no_trace" - self.parser.span_id = span_id or "no_span" + # Pass the ids per call rather than assigning them to the shared parser: + # two events in flight at once would otherwise overwrite each other's + # identity between the assignment and the offload that follows an await. content_json, content_parts, parser_truncated = await self.parser.parse( - raw_content + raw_content, + trace_id=trace_id or "no_trace", + span_id=span_id or "no_span", ) is_truncated = is_truncated or parser_truncated diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index aaaf51eaea1..4bef3c341e2 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -2919,27 +2919,34 @@ async def test_parser_instance_is_reused( assert bq_plugin_inst.parser is parser_after_init @pytest.mark.asyncio - async def test_parser_trace_id_updated_per_call( + async def test_parser_ids_are_not_mutated_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.""" + """_log_event passes the ids per call instead of writing them on the + + shared parser. + """ parser = bq_plugin_inst.parser original_trace_id = parser.trace_id + original_span_id = parser.span_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) + with mock.patch.object(parser, "parse", wraps=parser.parse) as mock_parse: + 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 != "" + assert parser.trace_id == original_trace_id + assert parser.span_id == original_span_id + _, kwargs = mock_parse.call_args + assert kwargs["span_id"] != "" + assert kwargs["span_id"] != original_span_id @pytest.mark.asyncio async def test_parser_not_recreated_with_constructor( @@ -7833,6 +7840,82 @@ def test_overlong_uri_is_replaced(self): assert removed +# ================================================================ +# TEST CLASS: GCS offload path identity +# ================================================================ +class TestOffloadPathIdentity: + """Tests that offload paths come from the call, not shared parser state.""" + + @pytest.mark.asyncio + async def test_concurrent_parses_keep_their_own_identity(self): + """Two parses in flight at once do not write under each other's prefix.""" + paths = [] + first_upload_started = asyncio.Event() + + async def upload_content(data, mime_type, path): + paths.append(path) + if len(paths) == 1: + # Hold the first upload open until the second has begun, so both + # parses are suspended inside _parse_content_object at once. + first_upload_started.set() + await asyncio.sleep(0.05) + return f"gs://bucket/{path}" + + offloader = mock.MagicMock() + offloader.upload_content = upload_content + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=offloader, + trace_id="", + span_id="", + max_length=10, + ) + content = types.Content(parts=[types.Part(text="X" * 200)]) + + async def parse_as(trace_id, span_id): + return await parser.parse(content, trace_id=trace_id, span_id=span_id) + + task_a = asyncio.create_task(parse_as("trace-a", "span-a")) + await asyncio.wait_for(first_upload_started.wait(), timeout=5) + task_b = asyncio.create_task(parse_as("trace-b", "span-b")) + await asyncio.gather(task_a, task_b) + + assert len(paths) == 2 + assert sum("trace-a/span-a" in p for p in paths) == 1 + assert sum("trace-b/span-b" in p for p in paths) == 1 + assert parser.trace_id == "" + assert parser.span_id == "" + + @pytest.mark.asyncio + async def test_same_part_index_in_two_messages_does_not_collide(self): + """Two messages in one request get distinct object names.""" + paths = [] + + async def upload_content(data, mime_type, path): + paths.append(path) + return f"gs://bucket/{path}" + + offloader = mock.MagicMock() + offloader.upload_content = upload_content + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=offloader, + trace_id="t", + span_id="s", + max_length=10, + ) + llm_request = llm_request_lib.LlmRequest( + model="gemini-pro", + contents=[ + types.Content(parts=[types.Part(text="A" * 200)]), + types.Content(parts=[types.Part(text="B" * 200)]), + ], + ) + + await parser.parse(llm_request, trace_id="t1", span_id="s1") + + assert len(paths) == 2 + assert paths[0] != paths[1] + + # ================================================================ # TEST CLASS: AGENT_RESPONSE logging (Issue #87) # ================================================================