fix(plugins): sanitize BigQuery analytics rows and build GCS paths per call (v1) - #6808
fix(plugins): sanitize BigQuery analytics rows and build GCS paths per call (v1)#6808GWeale wants to merge 4 commits into
Conversation
…er 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.
…ics (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.
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.
…(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.
| # 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( |
There was a problem hiding this comment.
This final pass is new here, and it routes two more values into _recursive_smart_truncate than reached it before: custom_tags, and extra_attributes, which carries dict(event.actions.state_delta) from :3238 -- a value a tool writes. On the base branch only dict(session.state) reached this code.
That matters because of an unguarded line further up. At :401 the list/tuple branch reconstructs the original type:
return type(obj)(new_list), truncated_any2919bf5b -- the commit this change cites -- guards that same line at :739, and the guard is still on main today (origin/main:1422):
if type(obj) is tuple or type(obj) is list:
return type(obj)(new_list), truncated_any
# Tuple/list subclasses (e.g. namedtuples) may require positional
# constructor fields; reconstructing raised TypeError and the safe
# callback dropped the whole row. JSON does not preserve the subclass
# identity anyway -- emit a plain list.
return new_list, truncated_anyA namedtuple's __new__ needs positional fields, so type(p)(new_list) raises TypeError. There is no except between :401 and _log_event, so @_safe_callback (:113) catches it and the whole BigQuery row is dropped. A tool doing tool_context.state["pos"] = Point(1, 2) would silently lose its STATE_DELTA row.
:401 is not itself in this diff -- this line is what makes it reachable for tool-written values. Should we bring the upstream guard across as well? What do you think?
| return _recursive_smart_truncate( | ||
| as_dict, max_len, seen, depth + 1, budget | ||
| ) | ||
| elif hasattr(obj, "model_dump") and callable(obj.model_dump): |
There was a problem hiding this comment.
2919bf5b added the depth cap together with a progress requirement on these three branches (:753 onwards; still there on main at origin/main:1436-1487), and the comment there explains why they belong together:
dumped = obj.model_dump()
if isinstance(dumped, (collections.abc.Mapping, list)):
return _recursive_smart_truncate(dumped, max_len, seen, depth + 1, budget)"Mock-like objects answer every duck-typed probe with another Mock-like object, and recursing on those churns to the depth cap (falsely flagging truncation) instead of settling at the stringify fallback."
Without it, an object that hands back a fresh wrapper per access walks 50 levels, loses the str(obj) repr it used to get at :438, and sets is_truncated. The mock_agent fixture change at test_bigquery_agent_analytics_plugin.py:79 looks like exactly that symptom: root_agent_name was a bare Mock and :3114 now walks it. Could we port the progress check alongside the cap?
| @@ -315,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(): | |||
There was a problem hiding this comment.
This loop is only entered when obj is a dict (:355). 2919bf5b:706 uses collections.abc.Mapping for that branch instead, and main still does (origin/main:1310), with the note that "stringifying them in the fallback branch would bypass key redaction".
A MappingProxyType or UserDict is not a dict subclass and has no model_dump/dict/to_dict, so it falls through to str(obj) at :438 and lands as "mappingproxy({'api_key': '...'})" with truncated=False. This module already uses MappingProxyType itself at :89 and :465, and :3114 now routes tool-written values through here.
:355 is not in this diff, but the redaction loop it guards is. Given the comment at :3110 says every value is guaranteed to have seen the redaction, should the branch widen to Mapping?
| truncated_any = True | ||
| break | ||
| if isinstance(k, str): | ||
| k_lower = k.lower() |
There was a problem hiding this comment.
9adf0113:1326, the commit cited for the URI change, normalizes the key before the membership test: k_lower = k.lower().replace("-", "_") (still there at origin/main:1334). Without it api-key, access-token and refresh-token are not redacted, and header-shaped keys are a common thing for a tool to echo into state. Any reason not to bring that one line across too?
| # 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: |
There was a problem hiding this comment.
This drops the query and fragment and keeps the path verbatim, which is a different algorithm from 9adf0113 -- the commit this change cites -- and from main today (origin/main:3083-3146, byte-identical to the cited commit). Upstream walks the path segment by segment through _canonicalize_common_ascii_escapes and _is_sensitive_text_key, redacting a sensitive segment and the one after it, and keeps the query with sensitive values redacted.
Dropping the query outright is stricter, but leaving the path unexamined is looser: a token in the path, say https://api.example.com/v1/files/ya29.a0Af.../download, is stored in the clear on v1 and redacted on main. Those three helpers do not exist on v1, but _SENSITIVE_KEYS (:264) does, so a per-segment k.lower().replace("-", "_") in _SENSITIVE_KEYS check with the same redact-the-following-segment rule would close most of the gap without porting them. The description explains this is a re-implementation, but not why the path was left unexamined -- could we either follow upstream here, or say what the v1 constraint is?
Non-blocking, same function: _MAX_URI_LENGTH = 8192 at :281 replaces an over-length URI with [REDACTED_SENSITIVE_URI] outright, where upstream gates at _MAX_JSON_INSPECT_CHARS (4,000,000, origin/main:668) and then shortens with self._truncate -- which v1 has, at :1440 -- so a long URI survives in truncated form. The description does call the 8192 bound out, so this is only a question about the number and the lossy failure mode.
Four changes to
BigQueryAgentAnalyticsPluginon thev1branch, one commit each, re-implemented against v1's code rather than cherry-picked.Content formatter failure —
7c96e6b8, from2919bf5bcontent_formatterraises holds[FORMATTER_FAILED]incontent.External URIs —
1ba04e1d, from9adf0113content_parts[].urikeeps scheme, netloc and path only.urlsplitrejects becomes[REDACTED_SENSITIVE_URI].Attributes tree —
ef4b10f1, from2919bf5battributesholds[REDACTED]for_SENSITIVE_KEYSandtemp:-scoped keys fromextra_attributesandcustom_tags.[MAX_DEPTH_EXCEEDED]) and 100,000 nodes per invocation ([SANITIZE_BUDGET_EXCEEDED]).GCS offload paths —
4c34f25f, from2919bf5b<date>/<trace_id>/<span_id>_<parse_uid>_c<content_ordinal>_p<part_index><ext>; existing objects are untouched.is_truncatedis set when the attributes or URI pass alters a value. There is no opt-out, so dashboards readingcontent,content_parts,attributesoris_truncatedmay need updating. The plugin is opt-in and is not re-exported fromgoogle.adk.plugins.The File Content Compliance check fails for a pre-existing reason; this PR adds no endpoint URL.