From 257c51b04bc0405967c0232c9948165d47f84654 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Fri, 10 Jul 2026 03:22:46 -0700 Subject: [PATCH 1/6] fix(plugins): close BQAA fail-open privacy, GCS concurrency, and startup-loss gaps Addresses the four P1 findings and mechanical P2s from #6356: - content_formatter failure fails closed: content is replaced with a [FORMATTER_FAILED] sentinel, never the unformatted payload; counted as formatter_failed in get_drop_stats(). - Final sanitizer pass over the complete assembled attributes tree immediately before serialization (covers state_delta, custom_tags, labels, nested keys, temp: scope), with truncation propagated into is_truncated. JSON-encoded string blobs are parsed and redacted when they mention a sensitive key. - HybridContentParser.parse()/_parse_content_object() take call-local trace_id/span_id; GCS object paths can no longer be overwritten by a concurrent event mutating the shared parser. - _ensure_schema_exists raises on table get/create failure (NotFound -> create and Conflict handling preserved); _ensure_started keeps _started=False, retries after bounded exponential backoff coalesced behind the setup lock, and rows lost meanwhile are counted as setup_unavailable. - enabled=False short-circuits _ensure_started (zero auth/client/table/ writer side effects from any entry point). - Runtime settings validated at construction (batch, flush, shutdown, queue, max_content_length, RetryConfig; max_retries < 0 previously skipped the write loop entirely). - Plugin-level drop counters merged into get_drop_stats(), surviving shutdown. --- .../bigquery_agent_analytics_plugin.py | 239 ++++++++++++++++-- .../test_bigquery_agent_analytics_plugin.py | 218 +++++++++++++++- 2 files changed, 431 insertions(+), 26 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 65275d3d93..8f64df8bae 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -410,6 +410,89 @@ def _extract_tool_declarations( "password", }) +# Written in place of event content when a configured content_formatter +# raises: the formatter is a privacy/redaction boundary, so failure must +# never fall back to the unformatted payload (#6356 P1-1). +_FORMATTER_FAILED_SENTINEL = "[FORMATTER_FAILED]" + + +def _sanitize_json_blob(value: str, seen: set[int]) -> tuple[str, bool]: + """Redacts sensitive keys inside a JSON-encoded string blob. + + Values such as cached credential JSON often reach attributes as opaque + strings, bypassing dict-key redaction (#6356 P1-2 / #5112). When a string + looks like a JSON object AND mentions a sensitive key name, parse it, + redact recursively, and re-serialize. Returns ``(value, changed)``; + strings that do not parse are returned unchanged. + """ + stripped = value.lstrip() + if not stripped.startswith("{"): + return value, False + lowered = value.lower() + if not any(key in lowered for key in _SENSITIVE_KEYS): + return value, False + try: + parsed = json.loads(value) + except (TypeError, ValueError): + return value, False + if not isinstance(parsed, dict): + return value, False + # Redact only (max_len=-1): length truncation is applied by the caller on + # the re-serialized string, keeping single responsibility per pass. + sanitized, _ = _recursive_smart_truncate(parsed, -1, seen) + return json.dumps(sanitized), True + + +def _validate_runtime_config(config: "BigQueryLoggerConfig") -> None: + """Validates runtime settings at construction time (#6356 P2). + + Invalid values used to be accepted silently and only misbehave at + runtime — notably ``max_retries < 0`` skips the write loop entirely, so + every batch is dropped without a single attempt. + + Raises: + ValueError: If any batch, queue, duration, or retry setting is + invalid. + """ + if config.batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {config.batch_size}.") + if config.batch_flush_interval <= 0: + raise ValueError( + f"batch_flush_interval must be > 0, got {config.batch_flush_interval}." + ) + if config.shutdown_timeout <= 0: + raise ValueError( + f"shutdown_timeout must be > 0, got {config.shutdown_timeout}." + ) + if config.queue_max_size < 1: + raise ValueError( + f"queue_max_size must be >= 1, got {config.queue_max_size}." + ) + if config.max_content_length != -1 and config.max_content_length < 1: + raise ValueError( + "max_content_length must be -1 (unlimited) or >= 1, got" + f" {config.max_content_length}." + ) + retry = config.retry_config + if retry.max_retries < 0: + raise ValueError( + f"retry_config.max_retries must be >= 0, got {retry.max_retries}." + ) + if retry.initial_delay <= 0: + raise ValueError( + f"retry_config.initial_delay must be > 0, got {retry.initial_delay}." + ) + if retry.multiplier < 1: + raise ValueError( + f"retry_config.multiplier must be >= 1, got {retry.multiplier}." + ) + if retry.max_delay < retry.initial_delay: + raise ValueError( + "retry_config.max_delay must be >= initial_delay, got" + f" max_delay={retry.max_delay} initial_delay={retry.initial_delay}." + ) + + # Cloud Platform OAuth scope. Assembled from parts so this module does not # embed a bare Google APIs host literal: the file-content compliance scan # rejects such host literals on changed files unless an accompanying mTLS @@ -456,6 +539,7 @@ def _recursive_smart_truncate( try: if isinstance(obj, str): + obj, _ = _sanitize_json_blob(obj, seen) if max_len != -1 and len(obj) > max_len: return obj[:max_len] + "...[TRUNCATED]", True return obj, False @@ -1644,9 +1728,21 @@ def _truncate(self, text: str) -> tuple[str, bool]: return text, False 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, ) -> 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. + + ``trace_id``/``span_id`` are call-local: GCS object paths are built from + these arguments so concurrent parses on the shared parser instance can + never use another event's identity (#6356 P1-3). They fall back to the + constructor values for backward compatibility. + """ + 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 content_parts = [] is_truncated = False summary_text = [] @@ -1673,7 +1769,7 @@ 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}_p{idx}{ext}" try: uri = await self.offloader.upload_content( part.inline_data.data, part.inline_data.mime_type, path @@ -1712,7 +1808,7 @@ 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}_p{idx}.txt" try: uri = await self.offloader.upload_content( part.text, "text/plain", path @@ -1760,8 +1856,23 @@ 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. + + ``trace_id``/``span_id`` identify the calling event for GCS object paths. + Pass them per call — the parser instance is shared across concurrent + events, so relying on the mutable instance fields lets one event's await + resume with another event's identity and overwrite its objects (#6356 + P1-3). The instance fields remain only as a backward-compatible default. + """ + 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 json_payload = {} content_parts = [] is_truncated = False @@ -1779,7 +1890,9 @@ def process_text(t: str) -> tuple[str, bool]: ) for c in 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 + ) if trunc: is_truncated = True content_parts.extend(parts) @@ -1797,14 +1910,18 @@ 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 + ) 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 + ) return {"text_summary": summary}, parts, trunc elif isinstance(content, (dict, list)): @@ -2460,8 +2577,16 @@ def __init__( self._visual_builder = _is_visual_builder.get() + _validate_runtime_config(self.config) + self._started = False self._startup_error: Optional[Exception] = None + self._setup_failures = 0 + self._setup_retry_at = 0.0 + # Plugin-level loss accounting (#6356 P2): counts drops that happen + # before/outside any BatchProcessor (setup unavailable, formatter + # failure). Merged into get_drop_stats() and survives shutdown. + self._local_drop_counts: dict[str, int] = {} self._is_shutting_down = False self._setup_lock = None self._credentials = credentials @@ -2651,7 +2776,7 @@ def get_drop_stats(self) -> dict[str, int]: Per-reason drop counts summed over every active loop's processor. Empty if no processor has been created yet. """ - totals: dict[str, int] = {} + totals: dict[str, int] = dict(self._local_drop_counts) for state in list(self._loop_state_by_loop.values()): for reason, count in state.batch_processor.get_drop_stats().items(): totals[reason] = totals.get(reason, 0) + count @@ -2777,29 +2902,35 @@ def _ensure_schema_exists(self) -> None: ) tbl.clustering_fields = self.config.clustering_fields tbl.labels = {_SCHEMA_VERSION_LABEL_KEY: _SCHEMA_VERSION} - table_ready = False try: self.client.create_table(tbl) - table_ready = True except cloud_exceptions.Conflict: # Another process created it concurrently — still usable. - table_ready = True + pass except Exception as e: + # Fail setup (#6356 P1-4): returning normally here used to let the + # plugin mark itself started against a missing table and silently + # lose every subsequent row. Raise so _ensure_started records the + # failure, keeps _started=False, and retries on a later event. logger.error( "Could not create table %s: %s", self.full_table_id, e, exc_info=True, ) - if table_ready and self.config.create_views: + raise + if self.config.create_views: self._create_analytics_views() except Exception as e: + # Fail setup (#6356 P1-4): swallowing control-plane errors here let + # the plugin mark itself started against a missing/unready table. logger.error( - "Error checking for table %s: %s", + "Error ensuring table %s is ready: %s", self.full_table_id, e, exc_info=True, ) + raise @staticmethod def _schema_fields_match( @@ -3137,9 +3268,15 @@ def _reset_runtime_state(self) -> None: self.parser = None self._started = False self._startup_error = None + self._setup_failures = 0 + self._setup_retry_at = 0.0 self._is_shutting_down = False self._init_pid = os.getpid() + def _count_local_drop(self, reason: str) -> None: + """Counts a row lost before/outside any BatchProcessor (#6356 P2).""" + self._local_drop_counts[reason] = self._local_drop_counts.get(reason, 0) + 1 + async def __aenter__(self) -> BigQueryAgentAnalyticsPlugin: await self._ensure_started() return self @@ -3148,7 +3285,19 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: await self.shutdown() async def _ensure_started(self, **kwargs) -> None: - """Ensures that the plugin is started and initialized.""" + """Ensures that the plugin is started and initialized. + + Setup failures no longer poison the plugin permanently (#6356 P1-4): + the failure is recorded, ``_started`` stays False, and a later event + retries after a bounded exponential backoff. Attempts are coalesced + through the setup lock, so failure mode costs at most one setup RPC + per backoff window — not one per event. + """ + # Disabled mode must have zero side effects (#6356 P2): no ADC lookup, + # client creation, table RPCs, or background tasks from any entry point + # (before_run_callback, __aenter__, _log_event all route through here). + if not self.config.enabled: + return # _init_pid == 0 means the plugin was unpickled and has never been # initialized in this process (the pickle sentinel set by # __getstate__). Skip the fork reset in that case — no fork @@ -3164,17 +3313,34 @@ async def _ensure_started(self, **kwargs) -> None: self._setup_lock = asyncio.Lock() async with self._setup_lock: if not self._started: + if ( + self._startup_error is not None + and time.monotonic() < self._setup_retry_at + ): + # Still inside the backoff window from a previous failure. + return try: await self._lazy_setup(**kwargs) self._started = True self._startup_error = None + self._setup_failures = 0 + self._setup_retry_at = 0.0 # Record the current PID so fork detection works for # the rest of this instance's lifetime. if self._init_pid == 0: self._init_pid = os.getpid() except Exception as e: self._startup_error = e - logger.error("Failed to initialize BigQuery Plugin: %s", e) + self._setup_failures += 1 + backoff = min(60.0, 2.0 ** min(self._setup_failures, 6)) + self._setup_retry_at = time.monotonic() + backoff + logger.error( + "Failed to initialize BigQuery Plugin (attempt %d, next" + " retry in %.0fs): %s", + self._setup_failures, + backoff, + e, + ) @staticmethod def _resolve_ids( @@ -3561,6 +3727,9 @@ async def _log_event( if not self._started: await self._ensure_started() if not self._started: + # Setup unavailable (failed and inside its retry backoff): the row + # is lost — record it so the loss is observable (#6356 P1-4). + self._count_local_drop("setup_unavailable") return if event_data is None: @@ -3570,8 +3739,19 @@ 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 (#6356 P1-1): the formatter is a redaction/privacy + # boundary, so its failure must never fall back to the unformatted + # payload. Keep the event's non-content metadata, replace content + # with a sentinel, and log without the original payload. + logger.warning( + "Content formatter failed for event %s; writing sentinel" + " instead of original content.", + event_type, + exc_info=True, + ) + raw_content = _FORMATTER_FAILED_SENTINEL + self._count_local_drop("formatter_failed") trace_id, span_id, parent_span_id = self._resolve_ids( event_data, callback_context @@ -3590,11 +3770,13 @@ async def _log_event( if {"content", "content_parts"} <= self._denied_columns: content_json, content_parts, parser_truncated = None, [], False else: - # 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 trace/span per call: the parser instance is shared, so storing + # request identity on it lets concurrent events overwrite each other's + # GCS object paths (#6356 P1-3). 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 @@ -3609,6 +3791,17 @@ async def _log_event( meta_truncated = self._capture_custom_metadata(event_data, attributes) is_truncated = is_truncated or meta_truncated + # Final safety pass (#6356 P1-2): sanitize the COMPLETE assembled + # attributes tree immediately before serialization. Producer-local + # sanitization above remains as an optimization, but this pass is the + # mandatory boundary — it covers values copied in directly (state_delta + # via extra_attributes, custom_tags, labels, generic extra attributes), + # nested structures, `temp:`-scoped keys, and JSON-encoded blobs. + 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 78caca075a..6fe40b9dac 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -759,8 +759,15 @@ 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"}' + # Fail CLOSED (#6356 P1-1): a raising formatter must never fall back + # to the unformatted payload. The row keeps its metadata but content + # is replaced with the sentinel, and the loss is observable. + assert "Secret message" not in str(log_entry["content"]) + assert bigquery_agent_analytics_plugin._FORMATTER_FAILED_SENTINEL in str( + log_entry["content"] + ) + assert log_entry["event_type"] == "USER_MESSAGE_RECEIVED" + assert plugin.get_drop_stats().get("formatter_failed") == 1 @pytest.mark.asyncio async def test_max_content_length( @@ -1191,7 +1198,11 @@ async def test_bigquery_client_initialization_failure( ) await asyncio.sleep(0.01) mock_logger.error.assert_called_with( - "Failed to initialize BigQuery Plugin: %s", mock.ANY + "Failed to initialize BigQuery Plugin (attempt %d, next" + " retry in %.0fs): %s", + mock.ANY, + mock.ANY, + mock.ANY, ) mock_write_client.append_rows.assert_not_called() @@ -9698,3 +9709,204 @@ async def test_both_payload_columns_denied_skips_parse_and_offload( ) await plugin.flush() mock_blob.upload_from_string.assert_not_called() + + +class TestIssue6356Hardening: + """Safety and lifecycle invariants from google/adk-python#6356.""" + + def test_invalid_runtime_config_rejected_at_construction( + self, mock_auth_default, mock_bq_client + ): + """Invalid batch/queue/duration/retry settings fail at construction.""" + _ = mock_auth_default, mock_bq_client + retry = bigquery_agent_analytics_plugin.RetryConfig + bad_configs = [ + dict(batch_size=0), + dict(batch_flush_interval=0.0), + dict(shutdown_timeout=0.0), + dict(queue_max_size=0), + dict(max_content_length=0), + dict(retry_config=retry(max_retries=-1)), + dict(retry_config=retry(initial_delay=0.0)), + dict(retry_config=retry(multiplier=0.5)), + dict(retry_config=retry(initial_delay=5.0, max_delay=1.0)), + ] + for kwargs in bad_configs: + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(**kwargs) + with pytest.raises(ValueError): + bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config + ) + + @pytest.mark.asyncio + async def test_final_attributes_pass_redacts_direct_producers( + 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, + ): + """state_delta, custom_tags, nested keys, and JSON blobs are redacted. + + These producers copy values into attributes without going through + _recursive_smart_truncate; the final pre-serialization pass must + redact them (#6356 P1-2). + """ + _ = mock_auth_default, mock_bq_client + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + custom_tags={"team": "sre", "password": "hunter2"}, + ) + 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() + bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData( + extra_attributes={ + "state_delta": { + "access_token": "ya29.SECRET", + "nested": {"refresh_token": "1//SECRET2"}, + "temp:scratch": "ephemeral", + "plain": "keep-me", + }, + "cred_blob": '{"access_token": "SECRETTOK", "expiry": 1}', + }, + ), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + attrs = json.loads(log_entry["attributes"]) + blob = str(log_entry["attributes"]) + assert "ya29.SECRET" not in blob + assert "1//SECRET2" not in blob + assert "SECRETTOK" not in blob + assert "hunter2" not in blob + assert attrs["state_delta"]["access_token"] == "[REDACTED]" + assert attrs["state_delta"]["nested"]["refresh_token"] == "[REDACTED]" + assert attrs["state_delta"]["temp:scratch"] == "[REDACTED]" + assert attrs["state_delta"]["plain"] == "keep-me" + assert attrs["custom_tags"]["password"] == "[REDACTED]" + assert attrs["custom_tags"]["team"] == "sre" + assert json.loads(attrs["cred_blob"])["access_token"] == "[REDACTED]" + + @pytest.mark.asyncio + async def test_concurrent_parses_never_share_gcs_paths(self): + """Two overlapping two-part parses keep call-local trace/span paths. + + Regression for #6356 P1-3: with identity stored on the shared parser, + event A resumed after event B's mutation and wrote under B's object + name, overwriting B's part. + """ + uploaded: list[str] = [] + + class _FakeOffloader: + + async def upload_content(self, data, mime, path): + uploaded.append(path) + await asyncio.sleep(0) # force interleave between part uploads + return f"gs://bucket/{path}" + + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=_FakeOffloader(), trace_id="ctor", span_id="ctor" + ) + + def two_parts(): + return types.Content( + parts=[ + types.Part.from_bytes(data=b"x", mime_type="image/png"), + types.Part.from_bytes(data=b"y", mime_type="image/png"), + ] + ) + + await asyncio.gather( + parser.parse(two_parts(), trace_id="trace-a", span_id="span-a"), + parser.parse(two_parts(), trace_id="trace-b", span_id="span-b"), + ) + assert len(uploaded) == 4 + assert len(set(uploaded)) == 4, f"path collision: {uploaded}" + assert sum("trace-a/span-a" in p for p in uploaded) == 2 + assert sum("trace-b/span-b" in p for p in uploaded) == 2 + + @pytest.mark.asyncio + async def test_setup_failure_keeps_not_started_then_retries( + 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, + ): + """Failed table readiness leaves _started=False, counts the loss, and + a later event retries successfully (#6356 P1-4).""" + _ = mock_auth_default + mock_bq_client.get_table.side_effect = cloud_exceptions.InternalServerError( + "control plane hiccup" + ) + async with managed_plugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) as plugin: + bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) + await plugin.before_run_callback(invocation_context=invocation_context) + assert plugin._started is False + assert plugin._startup_error is not None + + # A row logged while setup is unavailable is counted, not silent. + await plugin._log_event("USER_MESSAGE_RECEIVED", callback_context) + assert plugin.get_drop_stats().get("setup_unavailable", 0) >= 1 + + # Control plane recovers; retry succeeds on a later event once the + # backoff window elapses. + mock_bq_client.get_table.side_effect = None + plugin._setup_retry_at = 0.0 + await plugin._ensure_started() + assert plugin._started is True + assert plugin._startup_error is None + + @pytest.mark.asyncio + async def test_enabled_false_has_zero_side_effects( + self, mock_auth_default, mock_bq_client, invocation_context + ): + """enabled=False performs no auth/client/table/writer side effects + through Runner callbacks or async context-manager use (#6356 P2).""" + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(enabled=False) + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config + ) + await plugin.before_run_callback(invocation_context=invocation_context) + async with plugin: + pass + assert plugin._started is False + assert plugin.client is None + assert plugin._loop_state_by_loop == {} + mock_auth_default.assert_not_called() + mock_bq_client.get_table.assert_not_called() + + @pytest.mark.asyncio + async def test_drop_stats_survive_shutdown_and_include_local_reasons( + self, mock_auth_default, mock_bq_client + ): + """Pre-processor drop reasons are queryable, including after shutdown.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._count_local_drop("formatter_failed") + plugin._count_local_drop("setup_unavailable") + plugin._count_local_drop("setup_unavailable") + await plugin.shutdown() + stats = plugin.get_drop_stats() + assert stats["formatter_failed"] == 1 + assert stats["setup_unavailable"] == 2 From 75967014fa490a7e42e85589dc631af355b2d854 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Fri, 10 Jul 2026 04:15:04 -0700 Subject: [PATCH 2/6] =?UTF-8?q?fix(plugins):=20address=20review=20round=20?= =?UTF-8?q?1=20=E2=80=94=20retry=20readiness,=20GCS=20key=20ordinals,=20lo?= =?UTF-8?q?g=20leak,=20blob=20decode-first?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All eight findings from the PR #6360 review at 257c51b0: - P1-1: _ensure_schema_exists now runs on every setup attempt until one succeeds; the cached _schema no longer gates it. Retry test asserts the second get_table RPC. - P1-2: GCS object names gain a per-parse uid and content ordinal, so multiple messages in one request (same trace/span, restarting part index) can no longer collide. New multi-message regression test. - P1-3: the formatter-failure log carries only the exception class — no message, no exc_info — so a formatter that embeds content in its exception cannot leak it into application logs. Test asserts caplog. - P1-4: _sanitize_json_blob decodes FIRST ({ or [ prefixes), inspecting decoded keys recursively — JSON string escapes (access\u005ftoken) and arrays of credential objects are covered; unchanged strings return byte-for-byte. New escape/array regression test. - P2-5: shutdown() folds processor drop counters into the persistent plugin counters before clearing loop state. - P2-6: __setstate__ backfills _local_drop_counts/_setup_failures/ _setup_retry_at for pre-change pickles; __getstate__ resets retry state. - P2-7: _validate_runtime_config enforces integral counts and finite reals (NaN previously passed every ordered comparison; max_retries=NaN silently skipped the write loop). - CI: test_views_not_created_after_table_creation_failure updated to the intentional raise; parser-reuse test now pins the no-mutation contract. Also adds a recursion depth cap to _recursive_smart_truncate: id()-based cycle detection cannot catch graphs that manufacture new objects per duck-typed access (Mock-like), which the new final sanitizer pass exposed as an unbounded recursion. --- .../bigquery_agent_analytics_plugin.py | 199 ++++++++++++----- .../test_bigquery_agent_analytics_plugin.py | 210 +++++++++++++++++- 2 files changed, 349 insertions(+), 60 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 8f64df8bae..c677292823 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -26,6 +26,7 @@ import functools import json import logging +import math import mimetypes import os import traceback as traceback_module @@ -415,34 +416,66 @@ def _extract_tool_declarations( # never fall back to the unformatted payload (#6356 P1-1). _FORMATTER_FAILED_SENTINEL = "[FORMATTER_FAILED]" +# Recursion bound for _recursive_smart_truncate: id()-based cycle detection +# cannot catch graphs that create new objects per access (Mock-like duck +# typing); the cap turns unbounded recursion into a redacted leaf. +_MAX_SANITIZE_DEPTH = 50 -def _sanitize_json_blob(value: str, seen: set[int]) -> tuple[str, bool]: + +def _sanitize_json_blob( + value: str, seen: set[int], depth: int = 0 +) -> tuple[str, bool]: """Redacts sensitive keys inside a JSON-encoded string blob. Values such as cached credential JSON often reach attributes as opaque - strings, bypassing dict-key redaction (#6356 P1-2 / #5112). When a string - looks like a JSON object AND mentions a sensitive key name, parse it, - redact recursively, and re-serialize. Returns ``(value, changed)``; - strings that do not parse are returned unchanged. + strings, bypassing dict-key redaction (#6356 P1-2 / #5112). Decode + FIRST: raw-substring prefilters are bypassable through JSON string + escapes (e.g. ``"access\\u005ftoken"``), so any string that looks like a + JSON container is parsed and its *decoded* keys inspected recursively — + arrays of credential objects included. Returns ``(value, changed)``; + strings that do not parse, or that need no redaction, are returned + unchanged (no cosmetic re-serialization). """ stripped = value.lstrip() - if not stripped.startswith("{"): - return value, False - lowered = value.lower() - if not any(key in lowered for key in _SENSITIVE_KEYS): + if not stripped.startswith(("{", "[")): return value, False try: parsed = json.loads(value) except (TypeError, ValueError): return value, False - if not isinstance(parsed, dict): + if not isinstance(parsed, (dict, list)): return value, False # Redact only (max_len=-1): length truncation is applied by the caller on # the re-serialized string, keeping single responsibility per pass. - sanitized, _ = _recursive_smart_truncate(parsed, -1, seen) + sanitized, _ = _recursive_smart_truncate(parsed, -1, seen, depth + 1) + if sanitized == parsed: + return value, False return json.dumps(sanitized), True +def _require_count(name: str, value: Any, minimum: int) -> None: + """Requires an integral count >= minimum. + + Bools and floats are rejected: ordered comparisons alone let NaN pass + every range check (#6360 review P2-7). + """ + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{name} must be an int, got {value!r}.") + if value < minimum: + raise ValueError(f"{name} must be >= {minimum}, got {value}.") + + +def _require_finite(name: str, value: Any, minimum_exclusive: float) -> float: + """Requires a finite real number strictly greater than the minimum.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number, got {value!r}.") + if not math.isfinite(value): + raise ValueError(f"{name} must be finite, got {value!r}.") + if value <= minimum_exclusive: + raise ValueError(f"{name} must be > {minimum_exclusive}, got {value}.") + return float(value) + + def _validate_runtime_config(config: "BigQueryLoggerConfig") -> None: """Validates runtime settings at construction time (#6356 P2). @@ -454,19 +487,15 @@ def _validate_runtime_config(config: "BigQueryLoggerConfig") -> None: ValueError: If any batch, queue, duration, or retry setting is invalid. """ - if config.batch_size < 1: - raise ValueError(f"batch_size must be >= 1, got {config.batch_size}.") - if config.batch_flush_interval <= 0: - raise ValueError( - f"batch_flush_interval must be > 0, got {config.batch_flush_interval}." - ) - if config.shutdown_timeout <= 0: - raise ValueError( - f"shutdown_timeout must be > 0, got {config.shutdown_timeout}." - ) - if config.queue_max_size < 1: + _require_count("batch_size", config.batch_size, 1) + _require_finite("batch_flush_interval", config.batch_flush_interval, 0) + _require_finite("shutdown_timeout", config.shutdown_timeout, 0) + _require_count("queue_max_size", config.queue_max_size, 1) + if isinstance(config.max_content_length, bool) or not isinstance( + config.max_content_length, int + ): raise ValueError( - f"queue_max_size must be >= 1, got {config.queue_max_size}." + f"max_content_length must be an int, got {config.max_content_length!r}." ) if config.max_content_length != -1 and config.max_content_length < 1: raise ValueError( @@ -474,19 +503,17 @@ def _validate_runtime_config(config: "BigQueryLoggerConfig") -> None: f" {config.max_content_length}." ) retry = config.retry_config - if retry.max_retries < 0: - raise ValueError( - f"retry_config.max_retries must be >= 0, got {retry.max_retries}." - ) - if retry.initial_delay <= 0: - raise ValueError( - f"retry_config.initial_delay must be > 0, got {retry.initial_delay}." - ) - if retry.multiplier < 1: + _require_count("retry_config.max_retries", retry.max_retries, 0) + initial_delay = _require_finite( + "retry_config.initial_delay", retry.initial_delay, 0 + ) + multiplier = _require_finite("retry_config.multiplier", retry.multiplier, 0) + if multiplier < 1: raise ValueError( f"retry_config.multiplier must be >= 1, got {retry.multiplier}." ) - if retry.max_delay < retry.initial_delay: + max_delay = _require_finite("retry_config.max_delay", retry.max_delay, 0) + if max_delay < initial_delay: raise ValueError( "retry_config.max_delay must be >= initial_delay, got" f" max_delay={retry.max_delay} initial_delay={retry.initial_delay}." @@ -503,7 +530,10 @@ def _validate_runtime_config(config: "BigQueryLoggerConfig") -> None: 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, ) -> tuple[Any, bool]: """Recursively truncates string values within a dict or list. @@ -514,6 +544,7 @@ 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. Returns: A tuple of (truncated_object, is_truncated). @@ -521,6 +552,16 @@ def _recursive_smart_truncate( if seen is None: seen = set() + # Depth cap: id()-based cycle detection cannot catch object graphs that + # manufacture NEW objects on each duck-typed access (e.g. anything whose + # model_dump()/dict()/to_dict() returns a fresh wrapper — unittest Mocks + # being the canonical case). Without this cap such graphs recurse + # unboundedly. Like "[CIRCULAR_REFERENCE]", this is a structural + # sentinel rather than content truncation, so it does not flip + # is_truncated. + if depth >= _MAX_SANITIZE_DEPTH: + return "[MAX_DEPTH_EXCEEDED]", False + obj_id = id(obj) if obj_id in seen: return "[CIRCULAR_REFERENCE]", False @@ -539,7 +580,7 @@ def _recursive_smart_truncate( try: if isinstance(obj, str): - obj, _ = _sanitize_json_blob(obj, seen) + obj, _ = _sanitize_json_blob(obj, seen, depth) if max_len != -1 and len(obj) > max_len: return obj[:max_len] + "...[TRUNCATED]", True return obj, False @@ -555,7 +596,7 @@ def _recursive_smart_truncate( 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) if trunc: truncated_any = True new_dict[k] = val @@ -565,7 +606,7 @@ 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) + val, trunc = _recursive_smart_truncate(i, max_len, seen, depth + 1) if trunc: truncated_any = True new_list.append(val) @@ -573,23 +614,27 @@ 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) 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 + ) 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) 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 + ) except Exception: pass elif obj is None or isinstance(obj, (int, float, bool)): @@ -1733,6 +1778,8 @@ async def _parse_content_object( *, 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. @@ -1740,9 +1787,16 @@ async def _parse_content_object( these arguments so concurrent parses on the shared parser instance can never use another event's identity (#6356 P1-3). They fall back to the constructor values for backward compatibility. + + ``parse_uid`` (unique per parse() call) and ``content_ordinal`` (the + message index within a multi-content request) disambiguate GCS object + names: the part index alone restarts per Content, so two messages in + one request would otherwise collide at the same part ordinal (#6360 + review P1-2). """ 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[:8] content_parts = [] is_truncated = False summary_text = [] @@ -1769,7 +1823,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()}/{trace_id}/{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 @@ -1808,7 +1865,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()}/{trace_id}/{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 @@ -1873,6 +1933,10 @@ async def parse( """ 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: disambiguates GCS object names across the + # multiple Content objects of one request and across concurrent events + # (#6360 review P1-2). + parse_uid = uuid.uuid4().hex[:8] json_payload = {} content_parts = [] is_truncated = False @@ -1888,10 +1952,14 @@ 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, trace_id=trace_id, span_id=span_id + c, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + content_ordinal=content_idx, ) if trunc: is_truncated = True @@ -1911,7 +1979,11 @@ def process_text(t: str) -> tuple[str, bool]: json_payload["system_prompt"] = truncated_si else: summary, parts, trunc = await self._parse_content_object( - si, trace_id=trace_id, span_id=span_id + si, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + content_ordinal=len(contents), ) if trunc: is_truncated = True @@ -1920,7 +1992,10 @@ def process_text(t: str) -> tuple[str, bool]: elif isinstance(content, (types.Content, types.Part)): summary, parts, trunc = await self._parse_content_object( - content, trace_id=trace_id, span_id=span_id + content, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, ) return {"text_summary": summary}, parts, trunc @@ -2805,7 +2880,12 @@ async def _lazy_setup(self, **kwargs) -> None: # Project out denied payload columns schema-first, so the table # schema, Arrow schema, row dict, and views all stay consistent. self._schema = _project_schema(_get_events_schema(), self._denied_columns) - await loop.run_in_executor(self._executor, self._ensure_schema_exists) + # Run table readiness on EVERY setup attempt until one succeeds: the + # cached _schema must not gate it, or a failed first attempt would skip + # the table check on retry and mark the plugin started against a + # missing/unready table (#6360 review P1-1). Once _started is True, + # _lazy_setup returns early above, so the steady state pays no extra RPC. + await loop.run_in_executor(self._executor, self._ensure_schema_exists) if not self.parser: self.arrow_schema = to_arrow_schema(self._schema) @@ -3182,6 +3262,14 @@ async def shutdown(self, timeout: float | None = None) -> None: except Exception: pass + # Fold processor drop counters into the persistent plugin-level + # counters before discarding loop state, so get_drop_stats() keeps + # reporting losses after shutdown (#6360 review P2-5). + for state in self._loop_state_by_loop.values(): + for reason, count in state.batch_processor.get_drop_stats().items(): + self._local_drop_counts[reason] = ( + self._local_drop_counts.get(reason, 0) + count + ) self._loop_state_by_loop.clear() if self.client: @@ -3207,6 +3295,8 @@ def __getstate__(self): state["parser"] = None state["_started"] = False state["_startup_error"] = None + state["_setup_failures"] = 0 + state["_setup_retry_at"] = 0.0 state["_is_shutting_down"] = False state["_init_pid"] = 0 return state @@ -3216,6 +3306,9 @@ def __setstate__(self, state): # Backfill keys that may be absent in pickled state from older # code versions so _ensure_started does not raise AttributeError. state.setdefault("_init_pid", 0) + state.setdefault("_local_drop_counts", {}) + state.setdefault("_setup_failures", 0) + state.setdefault("_setup_retry_at", 0.0) self.__dict__.update(state) def _reset_runtime_state(self) -> None: @@ -3739,16 +3832,16 @@ async def _log_event( if self.config.content_formatter: try: raw_content = self.config.content_formatter(raw_content, event_type) - except Exception: + except Exception as e: # Fail CLOSED (#6356 P1-1): the formatter is a redaction/privacy # boundary, so its failure must never fall back to the unformatted - # payload. Keep the event's non-content metadata, replace content - # with a sentinel, and log without the original payload. + # payload. Log only the exception CLASS — the message or a + # traceback (exc_info) could embed the protected content itself. logger.warning( - "Content formatter failed for event %s; writing sentinel" + "Content formatter (%s) failed for event %s; writing sentinel" " instead of original content.", + type(e).__name__, event_type, - exc_info=True, ) raw_content = _FORMATTER_FAILED_SENTINEL self._count_local_drop("formatter_failed") diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 6fe40b9dac..c78d099d69 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -17,6 +17,7 @@ import contextlib import dataclasses import json +import logging import os from unittest import mock @@ -3365,16 +3366,22 @@ 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_identity_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 must NOT store request identity on the shared parser. + + trace_id/span_id are passed per parse() call (#6356 P1-3): mutating the + shared instance let a concurrent event's await resume with another + event's identity and overwrite its GCS objects. + """ 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( @@ -3383,9 +3390,11 @@ async def test_parser_trace_id_updated_per_call( ) 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 != "" + # The shared parser's constructor-time fields are untouched; identity + # travelled through the parse() call arguments instead. + assert parser.trace_id == original_trace_id + assert parser.span_id == original_span_id + mock_write_client.append_rows.assert_called_once() @pytest.mark.asyncio async def test_parser_not_recreated_with_constructor( @@ -6397,12 +6406,16 @@ async def test_create_analytics_views_ensures_started( 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.""" + """create_table failure raises (fail setup, #6356 P1-4) and skips views.""" 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() + # Table readiness is a startup requirement: the failure propagates so + # _ensure_started keeps _started=False and retries later, instead of + # marking the plugin started against a missing table. + with pytest.raises(RuntimeError, match="BQ down"): + plugin._ensure_schema_exists() # Views should NOT be attempted since table creation failed plugin.client.query.assert_not_called() @@ -9869,11 +9882,16 @@ async def test_setup_failure_keeps_not_started_then_retries( # Control plane recovers; retry succeeds on a later event once the # backoff window elapses. + failed_calls = mock_bq_client.get_table.call_count mock_bq_client.get_table.side_effect = None plugin._setup_retry_at = 0.0 await plugin._ensure_started() assert plugin._started is True assert plugin._startup_error is None + # Table readiness must re-run on the retry: a cached _schema used to + # skip _ensure_schema_exists entirely, marking the plugin started + # without ever re-checking the table (#6360 review P1-1). + assert mock_bq_client.get_table.call_count == failed_calls + 1 @pytest.mark.asyncio async def test_enabled_false_has_zero_side_effects( @@ -9910,3 +9928,181 @@ async def test_drop_stats_survive_shutdown_and_include_local_reasons( stats = plugin.get_drop_stats() assert stats["formatter_failed"] == 1 assert stats["setup_unavailable"] == 2 + + def test_json_blob_redaction_survives_escapes_and_arrays(self): + """Decode-first blob sanitizing defeats raw-substring bypasses. + + `{"access\\u005ftoken": ...}` contains no literal sensitive substring, + and arrays of credential objects have no top-level dict (#6360 review + P1-4). Both must still be redacted; innocent strings stay unchanged. + """ + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + escaped = '{"access\\u005ftoken": "SECRET-A"}' + out, _ = truncate({"blob": escaped}, 10000) + assert "SECRET-A" not in json.dumps(out) + assert json.loads(out["blob"])["access_token"] == "[REDACTED]" + + array_blob = '[{"api_key": "SECRET-B"}, {"plain": "ok"}]' + out, _ = truncate({"blob": array_blob}, 10000) + assert "SECRET-B" not in json.dumps(out) + decoded = json.loads(out["blob"]) + assert decoded[0]["api_key"] == "[REDACTED]" + assert decoded[1]["plain"] == "ok" + + # No redaction needed -> string returned byte-for-byte (no cosmetic + # re-serialization). + innocent = '{"note": "spacing preserved"}' + out, _ = truncate({"blob": innocent}, 10000) + assert out["blob"] == innocent + + @pytest.mark.asyncio + async def test_multi_message_offloads_get_unique_paths(self): + """Two messages in ONE request must not collide at the same part index. + + The part ordinal restarts per Content while trace/span are shared, so + paths need the per-parse uid + content ordinal (#6360 review P1-2). + """ + uploaded: list[str] = [] + + class _FakeOffloader: + + async def upload_content(self, data, mime, path): + uploaded.append(path) + return f"gs://bucket/{path}" + + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=_FakeOffloader(), trace_id="t", span_id="s" + ) + request = llm_request_lib.LlmRequest( + contents=[ + types.Content( + role="user", + parts=[types.Part.from_bytes(data=b"a", mime_type="image/png")], + ), + types.Content( + role="user", + parts=[types.Part.from_bytes(data=b"b", mime_type="image/png")], + ), + ] + ) + await parser.parse(request, trace_id="trace-x", span_id="span-x") + assert len(uploaded) == 2 + assert len(set(uploaded)) == 2, f"collision within request: {uploaded}" + + @pytest.mark.asyncio + async def test_formatter_failure_log_does_not_leak_payload( + self, + mock_write_client, + invocation_context, + mock_auth_default, + mock_bq_client, + mock_to_arrow_schema, + dummy_arrow_schema, + mock_asyncio_to_thread, + caplog, + ): + """The formatter-failure log line must not carry the protected content. + + A formatter that embeds content in its exception message would leak it + through exc_info tracebacks (#6360 review P1-3); only the exception + class is logged. + """ + _ = mock_auth_default, mock_bq_client + + def leaky_formatter(content, event_type): + raise ValueError(f"could not redact: {content}") + + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + content_formatter=leaky_formatter + ) + async with managed_plugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config + ) as plugin: + await plugin._ensure_started() + bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) + with caplog.at_level(logging.WARNING): + await plugin.on_user_message_callback( + invocation_context=invocation_context, + user_message=types.Content( + parts=[types.Part(text="TOPSECRET-PAYLOAD")] + ), + ) + assert "TOPSECRET-PAYLOAD" not in caplog.text + assert "ValueError" in caplog.text + + @pytest.mark.asyncio + async def test_shutdown_folds_processor_drops_into_stats( + self, mock_auth_default, mock_bq_client + ): + """Processor drop counters survive shutdown via the plugin counters. + + get_drop_stats() used to read only live loop states, which shutdown() + clears (#6360 review P2-5).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + processor = mock.MagicMock() + processor.get_drop_stats.return_value = {"queue_full": 2} + processor.shutdown = mock.AsyncMock() + state = mock.MagicMock() + state.batch_processor = processor + state.write_client = None + plugin._loop_state_by_loop[asyncio.get_running_loop()] = state + + assert plugin.get_drop_stats() == {"queue_full": 2} + await plugin.shutdown() + assert plugin._loop_state_by_loop == {} + assert plugin.get_drop_stats() == {"queue_full": 2} + + def test_setstate_backfills_new_runtime_fields( + self, mock_auth_default, mock_bq_client + ): + """Pickles from pre-#6356 code lack the new fields; __setstate__ must + backfill them so get_drop_stats()/_ensure_started don't raise (#6360 + review P2-6).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + legacy_state = plugin.__getstate__() + for key in ("_local_drop_counts", "_setup_failures", "_setup_retry_at"): + legacy_state.pop(key, None) + restored = ( + bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin.__new__( + bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin + ) + ) + restored.__setstate__(legacy_state) + assert restored.get_drop_stats() == {} + assert restored._setup_failures == 0 + assert restored._setup_retry_at == 0.0 + + def test_invalid_config_rejects_nan_and_wrong_types( + self, mock_auth_default, mock_bq_client + ): + """NaN and wrong-typed values must fail construction (#6360 review + P2-7): ordered comparisons alone let NaN pass every range check.""" + _ = mock_auth_default, mock_bq_client + retry = bigquery_agent_analytics_plugin.RetryConfig + nan = float("nan") + bad_configs = [ + dict(batch_size=nan), + dict(batch_size=2.0), + dict(batch_size=True), + dict(batch_flush_interval=nan), + dict(shutdown_timeout=float("inf")), + dict(queue_max_size="10"), + dict(max_content_length=1.5), + dict(retry_config=retry(max_retries=nan)), + dict(retry_config=retry(initial_delay=nan)), + dict(retry_config=retry(multiplier=nan)), + dict(retry_config=retry(max_delay=nan)), + ] + for kwargs in bad_configs: + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig(**kwargs) + with pytest.raises(ValueError): + bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config + ) From 9b7994036d14789d36fd631301a44a18814a996a Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Fri, 10 Jul 2026 11:25:01 -0700 Subject: [PATCH 3/6] =?UTF-8?q?fix(plugins):=20address=20review=20round=20?= =?UTF-8?q?2=20=E2=80=94=20dup-key=20blobs,=20upgrade=20readiness,=20Mappi?= =?UTF-8?q?ng=20redaction,=20loss=20accounting?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All eight findings from the round-2 review at 75967014: - P1-1: _sanitize_json_blob tracks duplicate JSON members via an object_pairs_hook; blobs with duplicates are always reserialized, so an earlier duplicate secret member can no longer ride through the sanitized == parsed shortcut. - P1-2: _maybe_upgrade_schema re-raises update_table failures — a table verifiably missing required fields no longer lets _ensure_started mark the plugin ready with no readiness retry. Both pinned logged-not-raised tests updated to the new contract. - P1-3: the sanitizer walks collections.abc.Mapping (MappingProxyType, UserDict, ...) and emits plain sanitized dicts instead of stringifying them past key redaction. - P2-4: shutdown-timeout losses are counted — the cancelled worker counts its in-flight batch and shutdown() drains and counts queued rows under a new shutdown_timeout drop reason. - P2-5: RecursionError/MemoryError while parsing a JSON blob fails closed to [UNPARSEABLE_JSON_BLOB] instead of passing the row through unexamined (or dropping the whole row via the callback wrapper). - P2-6: _cleanup_stale_loop_states folds the dead processor's drop counters into the persistent plugin counters before deletion. - P2-7: __setstate__ validates the restored config; a legacy pickle with retry_config.max_retries=NaN now fails at restore instead of silently skipping the write loop. - P2-8: GCS object names use the full 128-bit uuid, and uploads pass if_generation_match=0 (create-only) so a collision fails loudly instead of rebinding an existing row to another event's bytes. --- .../bigquery_agent_analytics_plugin.py | 105 ++++++++-- .../test_bigquery_agent_analytics_plugin.py | 179 +++++++++++++++++- 2 files changed, 264 insertions(+), 20 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index c677292823..4bee406903 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -16,6 +16,7 @@ import asyncio import atexit +import collections.abc from concurrent.futures import ThreadPoolExecutor import contextvars import dataclasses @@ -439,18 +440,40 @@ def _sanitize_json_blob( stripped = value.lstrip() if not stripped.startswith(("{", "[")): return value, False + + # json.loads silently keeps only the LAST duplicate member, so a blob like + # {"access_token":"SECRET","access_token":"x"} can compare equal after + # sanitization while the raw string still carries the secret (#6360 review + # round 2 P1-1). Track duplicates while parsing and always reserialize + # such blobs. + saw_duplicate_key = False + + def _pairs_hook(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + nonlocal saw_duplicate_key + result = {} + for k, v in pairs: + if k in result: + saw_duplicate_key = True + result[k] = v + return result + try: - parsed = json.loads(value) + parsed = json.loads(value, object_pairs_hook=_pairs_hook) + if not isinstance(parsed, (dict, list)): + return value, False + # Redact only (max_len=-1): length truncation is applied by the caller + # on the re-serialized string, keeping single responsibility per pass. + sanitized, _ = _recursive_smart_truncate(parsed, -1, seen, depth + 1) + if sanitized == parsed and not saw_duplicate_key: + return value, False + return json.dumps(sanitized), True except (TypeError, ValueError): return value, False - if not isinstance(parsed, (dict, list)): - return value, False - # Redact only (max_len=-1): length truncation is applied by the caller on - # the re-serialized string, keeping single responsibility per pass. - sanitized, _ = _recursive_smart_truncate(parsed, -1, seen, depth + 1) - if sanitized == parsed: - return value, False - return json.dumps(sanitized), True + except (RecursionError, MemoryError): + # A blob too deep/large to inspect cannot be verified secret-free. + # Fail CLOSED to a sentinel instead of passing it through unexamined + # (#6360 review round 2 P2-5). + return "[UNPARSEABLE_JSON_BLOB]", True def _require_count(name: str, value: Any, minimum: int) -> None: @@ -568,7 +591,7 @@ def _recursive_smart_truncate( # Track compound objects to detect cycles is_compound = ( - isinstance(obj, (dict, list, tuple)) + isinstance(obj, (dict, list, tuple, collections.abc.Mapping)) or (dataclasses.is_dataclass(obj) and not isinstance(obj, type)) or hasattr(obj, "model_dump") or hasattr(obj, "dict") @@ -584,7 +607,10 @@ def _recursive_smart_truncate( if max_len != -1 and len(obj) > max_len: return obj[:max_len] + "...[TRUNCATED]", True return obj, False - elif isinstance(obj, dict): + elif isinstance(obj, collections.abc.Mapping): + # Covers dict plus mapping views (MappingProxyType, UserDict, ...): + # stringifying them in the fallback branch would bypass key redaction + # (#6360 review round 2 P1-3). Always emits a plain sanitized dict. truncated_any = False # Use dict comprehension for potentially slightly better performance, # but explicit loop is fine for clarity given recursive nature. @@ -1324,6 +1350,7 @@ def __init__( "retry_exhausted": 0, "non_retryable": 0, "unexpected_error": 0, + "shutdown_timeout": 0, } async def flush(self) -> None: @@ -1490,6 +1517,16 @@ async def _batch_writer(self) -> None: except asyncio.TimeoutError: continue + except asyncio.CancelledError: + # Cancelled mid-write by the shutdown timeout: the in-flight batch + # is lost — count it (#6360 review round 2 P2-4). + if batch: + self._dropped["shutdown_timeout"] += len(batch) + logger.warning( + "%d in-flight row(s) dropped by shutdown cancellation.", + len(batch), + ) + raise except asyncio.CancelledError: logger.info("Batch writer task cancelled.") break @@ -1668,6 +1705,23 @@ async def shutdown(self, timeout: float = 5.0) -> None: await self._batch_processor_task except asyncio.CancelledError: pass + # Rows still queued after the timeout are lost: count them so the + # loss is observable instead of silent (#6360 review round 2 P2-4). + # The worker counts its own in-flight batch on cancellation. + drained = 0 + try: + while True: + item = self._queue.get_nowait() + if item is not _SHUTDOWN_SENTINEL: + drained += 1 + self._queue.task_done() + except asyncio.QueueEmpty: + pass + if drained: + self._dropped["shutdown_timeout"] += drained + logger.warning( + "%d queued row(s) dropped by shutdown timeout.", drained + ) except Exception as e: logger.error("Error during BatchProcessor shutdown: %s", e) @@ -1742,7 +1796,14 @@ def _upload_sync( self, data: bytes | str, content_type: str, path: str ) -> str: blob = self.bucket.blob(path) - blob.upload_from_string(data, content_type=content_type) + # if_generation_match=0: create-only. Object names are unique by + # construction, so on the (astronomically unlikely) collision this + # fails the upload — surfaced as [UPLOAD FAILED] — instead of silently + # rebinding an existing BigQuery row to another event's bytes (#6360 + # review round 2 P2-8). + blob.upload_from_string( + data, content_type=content_type, if_generation_match=0 + ) return f"gs://{self.bucket.name}/{path}" @@ -1796,7 +1857,7 @@ async def _parse_content_object( """ 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[:8] + parse_uid = parse_uid or uuid.uuid4().hex content_parts = [] is_truncated = False summary_text = [] @@ -1936,7 +1997,7 @@ async def parse( # Unique per parse() call: disambiguates GCS object names across the # multiple Content objects of one request and across concurrent events # (#6360 review P1-2). - parse_uid = uuid.uuid4().hex[:8] + parse_uid = uuid.uuid4().hex json_payload = {} content_parts = [] is_truncated = False @@ -2685,6 +2746,13 @@ def _cleanup_stale_loop_states(self) -> None: loop, id(loop), ) + # Preserve the dead processor's loss accounting before discarding it + # (#6360 review round 2 P2-6), mirroring shutdown(). + state = self._loop_state_by_loop[loop] + for reason, count in state.batch_processor.get_drop_stats().items(): + self._local_drop_counts[reason] = ( + self._local_drop_counts.get(reason, 0) + count + ) del self._loop_state_by_loop[loop] # API Compatibility: These class-level attributes mask the dynamic @@ -3144,6 +3212,11 @@ def _maybe_upgrade_schema(self, existing_table: bigquery.Table) -> None: e, exc_info=True, ) + # The table is verifiably missing required fields at this point; + # swallowing the failure would let _ensure_started mark the plugin + # ready against a table every later write can fail on, with no + # readiness retry (#6360 review round 2 P1-2). + raise def _project_view_columns(self, extra_cols: list[str]) -> list[str]: """Drops derived view expressions that reference a denied column. @@ -3310,6 +3383,10 @@ def __setstate__(self, state): state.setdefault("_setup_failures", 0) state.setdefault("_setup_retry_at", 0.0) self.__dict__.update(state) + # Pickles from older code bypass __init__, so re-validate the restored + # configuration: e.g. a legacy retry_config with max_retries=NaN would + # otherwise skip the write loop silently (#6360 review round 2 P2-7). + _validate_runtime_config(self.config) def _reset_runtime_state(self) -> None: """Resets all runtime state after a fork. diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index c78d099d69..6d225f8ce7 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -5123,8 +5123,13 @@ def test_skip_upgrade_when_version_matches(self): 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.""" + def test_upgrade_error_propagates_when_fields_missing(self): + """Schema upgrade failure raises when required fields are missing. + + Swallowing it let _ensure_started mark the plugin ready against a + table every later write can fail on, with no readiness retry (#6360 + review round 2 P1-2). + """ plugin = self._make_plugin(auto_schema_upgrade=True) existing = mock.MagicMock(spec=bigquery.Table) existing.schema = [ @@ -5133,8 +5138,8 @@ def test_upgrade_error_is_logged_not_raised(self): existing.labels = {} plugin.client.get_table.return_value = existing plugin.client.update_table.side_effect = Exception("boom") - # Should not raise - plugin._ensure_schema_exists() + with pytest.raises(Exception, match="boom"): + plugin._ensure_schema_exists() def test_upgrade_preserves_existing_columns(self): """Existing columns are never dropped or altered during upgrade.""" @@ -7653,8 +7658,10 @@ def test_version_label_not_stamped_on_failure(self): plugin.client.get_table.return_value = existing plugin.client.update_table.side_effect = Exception("network error") - # Should not raise. - plugin._ensure_schema_exists() + # Raises so setup is not marked ready against a table with missing + # fields (#6360 review round 2 P1-2). + with pytest.raises(Exception, match="network error"): + 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 @@ -10106,3 +10113,163 @@ def test_invalid_config_rejects_nan_and_wrong_types( bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config ) + + def test_json_blob_duplicate_keys_always_reserialized(self): + """Duplicate JSON members must not defeat the changed-blob check. + + json.loads keeps only the last duplicate, so sanitized == parsed can + hold while the raw string still carries an earlier secret member + (#6360 review round 2 P1-1). + """ + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + blob = '{"access_token": "SECRET-DUP", "access_token": "[REDACTED]"}' + out, _ = truncate({"blob": blob}, 10000) + assert "SECRET-DUP" not in json.dumps(out) + assert json.loads(out["blob"])["access_token"] == "[REDACTED]" + + def test_mapping_views_are_redacted(self): + """Mapping types beyond dict must be walked, not stringified. + + MappingProxyType/UserDict used to hit the stringify fallback, leaking + sensitive members (#6360 review round 2 P1-3). + """ + import collections + from types import MappingProxyType + + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + proxy = MappingProxyType({"access_token": "SECRET-PROXY"}) + userdict = collections.UserDict({"refresh_token": "SECRET-USERDICT"}) + out, _ = truncate({"proxy": proxy, "userdict": userdict}, 10000) + dumped = json.dumps(out) + assert "SECRET-PROXY" not in dumped + assert "SECRET-USERDICT" not in dumped + assert out["proxy"]["access_token"] == "[REDACTED]" + assert out["userdict"]["refresh_token"] == "[REDACTED]" + + def test_deep_json_blob_fails_closed(self): + """A blob too deep to parse becomes a sentinel, not a pass-through. + + json.loads raises RecursionError before the bounded traversal ever + runs; the row keeps flowing with the blob replaced (#6360 review + round 2 P2-5). + """ + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + deep = "[" * 10000 + "]" * 10000 + out, _ = truncate({"blob": deep}, 500 * 1024) + assert out["blob"] == "[UNPARSEABLE_JSON_BLOB]" + + @pytest.mark.asyncio + async def test_shutdown_timeout_counts_lost_rows(self): + """Rows stranded by a shutdown timeout are counted, not silent. + + In-flight batch rows are counted by the cancelled worker and queued + rows by the drain in shutdown() (#6360 review round 2 P2-4). + """ + write_started = asyncio.Event() + + async def hung_writer(batch): + write_started.set() + await asyncio.sleep(3600) + + processor = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=mock.MagicMock(), + write_stream="stream", + batch_size=1, + flush_interval=0.05, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=0.1, + ) + with mock.patch.object( + processor, "_write_rows_with_retry", side_effect=hung_writer + ): + await processor.start() + await processor.append({"row": 1}) + await write_started.wait() # row 1 is in-flight in the hung writer + await processor.append({"row": 2}) # row 2 stays queued + await processor.shutdown(timeout=0.1) + + stats = processor.get_drop_stats() + assert stats.get("shutdown_timeout") == 2 + + @pytest.mark.asyncio + async def test_stale_loop_cleanup_preserves_drop_stats( + self, mock_auth_default, mock_bq_client + ): + """Closed-loop cleanup folds processor counters before deletion + (#6360 review round 2 P2-6).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + dead_loop = mock.MagicMock() + dead_loop.is_closed.return_value = True + state = mock.MagicMock() + state.batch_processor.get_drop_stats.return_value = {"write_failed": 7} + plugin._loop_state_by_loop[dead_loop] = state + + plugin._cleanup_stale_loop_states() + + assert plugin._loop_state_by_loop == {} + assert plugin.get_drop_stats().get("write_failed") == 7 + + def test_setstate_validates_restored_config( + self, mock_auth_default, mock_bq_client + ): + """Legacy pickles with invalid runtime config fail at restore, not as + a silent write-loop skip (#6360 review round 2 P2-7).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + state = plugin.__getstate__() + state["config"].retry_config.max_retries = float("nan") + restored = ( + bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin.__new__( + bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin + ) + ) + with pytest.raises(ValueError): + restored.__setstate__(state) + + @pytest.mark.asyncio + async def test_gcs_uploads_use_full_uid_and_create_only(self): + """Object names carry the full 128-bit uid; uploads are create-only. + + 32 random bits reach ~50% birthday collision around 77k parses; a + collision must fail the upload instead of rebinding an existing row + to another event's bytes (#6360 review round 2 P2-8). + """ + uploaded: list[str] = [] + + class _FakeOffloader: + + async def upload_content(self, data, mime, path): + uploaded.append(path) + return f"gs://bucket/{path}" + + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=_FakeOffloader(), trace_id="t", span_id="s" + ) + await parser.parse( + types.Content( + parts=[types.Part.from_bytes(data=b"a", mime_type="image/png")] + ), + trace_id="trace-y", + span_id="span-y", + ) + assert len(uploaded) == 1 + # .../{span}_{32-hex-uid}_c{n}_p{idx}.png + uid_segment = uploaded[0].split("span-y_")[1].split("_c")[0] + assert len(uid_segment) == 32 + + # And the sync upload path passes create-only semantics. + bucket = mock.MagicMock() + offloader = bigquery_agent_analytics_plugin.GCSOffloader.__new__( + bigquery_agent_analytics_plugin.GCSOffloader + ) + offloader.bucket = bucket + offloader._upload_sync(b"data", "image/png", "p") + _, kwargs = bucket.blob.return_value.upload_from_string.call_args + assert kwargs.get("if_generation_match") == 0 From 5c5521d1291679de78a252220a2b4a6af8c7bc0d Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 00:50:47 -0700 Subject: [PATCH 4/6] =?UTF-8?q?fix(plugins):=20address=20review=20round=20?= =?UTF-8?q?3=20=E2=80=94=20blob=20ValueError=20fail-closed,=20label-only?= =?UTF-8?q?=20readiness,=20cross-loop=20setup?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All six findings from the round-3 review at 9b799403: - P1-1: json.JSONDecodeError (not JSON) passes through; every OTHER ValueError from json.loads — e.g. integers over the interpreter digit limit on syntactically valid JSON — now fails closed to the [UNPARSEABLE_JSON_BLOB] sentinel, alongside RecursionError/MemoryError. - P1-2: _maybe_upgrade_schema re-raises update_table failures only when schema fields are actually missing; a label-only refresh failure logs and continues, since the table is write-compatible and the stale label retries next run. - P2: _ensure_started uses per-event-loop asyncio locks (dict guarded by a threading.Lock) — one shared asyncio.Lock is loop-bound and raised 'Non-thread-safe operation' with waiters stranded on other loops. Pickle/reset/fork paths updated for the new fields. - P2: _cleanup_stale_loop_states claims each state atomically via dict.pop before folding counters — read-fold-delete raced under concurrent cleanup, double-counting and raising KeyError. - P2: BatchProcessor.close() shares shutdown()'s drain/accounting, so rows stranded in the queue are counted under shutdown_timeout instead of silently discarded. - P3: the duplicate CancelledError handler is consolidated. The kept handler must RE-RAISE: asyncio.wait_for treats suppressed cancellation as normal completion, so a swallow-and-break variant silently disabled shutdown()'s drain branch (caught by the round-2 regression test). Six new regression tests, including a two-thread/two-loop startup test and the reviewer's big-int and label-only repros. --- .../bigquery_agent_analytics_plugin.py | 105 ++++++++++--- .../test_bigquery_agent_analytics_plugin.py | 148 +++++++++++++++++- 2 files changed, 225 insertions(+), 28 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 4bee406903..1f58c82ce3 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -40,6 +40,7 @@ import random import re +import threading import time from types import MappingProxyType from typing import Any @@ -467,12 +468,16 @@ def _pairs_hook(pairs: list[tuple[str, Any]]) -> dict[str, Any]: if sanitized == parsed and not saw_duplicate_key: return value, False return json.dumps(sanitized), True - except (TypeError, ValueError): + except json.JSONDecodeError: + # Not JSON at all — nothing to inspect; the outer string handling + # (length truncation) still applies. return value, False - except (RecursionError, MemoryError): - # A blob too deep/large to inspect cannot be verified secret-free. - # Fail CLOSED to a sentinel instead of passing it through unexamined - # (#6360 review round 2 P2-5). + except (TypeError, ValueError, RecursionError, MemoryError): + # Syntactically valid JSON that Python cannot materialize — e.g. + # integers over sys.get_int_max_str_digits() raise a plain ValueError + # (#6360 review round 3 P1-1) — or a blob too deep/large to inspect + # (round 2 P2-5). Either way it cannot be verified secret-free: fail + # CLOSED to a sentinel instead of passing it through unexamined. return "[UNPARSEABLE_JSON_BLOB]", True @@ -1518,18 +1523,21 @@ async def _batch_writer(self) -> None: except asyncio.TimeoutError: continue except asyncio.CancelledError: - # Cancelled mid-write by the shutdown timeout: the in-flight batch - # is lost — count it (#6360 review round 2 P2-4). + # Cancelled (e.g. by the shutdown timeout): the in-flight batch is + # lost — count it (#6360 review round 2 P2-4) — then exit the + # worker, preserving the original swallow-and-break semantics. if batch: self._dropped["shutdown_timeout"] += len(batch) logger.warning( "%d in-flight row(s) dropped by shutdown cancellation.", len(batch), ) - raise - except asyncio.CancelledError: logger.info("Batch writer task cancelled.") - break + # Re-raise: asyncio.wait_for treats a task that SUPPRESSES + # cancellation as a normal completion, so shutdown()'s timeout + # branch (which drains and counts the remaining queue) would never + # run if this swallowed the cancellation. + raise except Exception as e: logger.error("Error in batch writer loop: %s", e, exc_info=True) # Avoid sleeping if we are shutting down or if the task was cancelled @@ -1746,6 +1754,21 @@ async def close(self) -> None: await self._batch_processor_task except asyncio.CancelledError: pass + # Same loss accounting as shutdown(): rows still queued after the + # timeout are counted, not silently discarded (#6360 review round 3 + # P2). The cancelled worker counts its own in-flight batch. + drained = 0 + try: + while True: + item = self._queue.get_nowait() + if item is not _SHUTDOWN_SENTINEL: + drained += 1 + self._queue.task_done() + except asyncio.QueueEmpty: + pass + if drained: + self._dropped["shutdown_timeout"] += drained + logger.warning("%d queued row(s) dropped by close timeout.", drained) # ============================================================================== @@ -2724,7 +2747,8 @@ def __init__( # failure). Merged into get_drop_stats() and survives shutdown. self._local_drop_counts: dict[str, int] = {} self._is_shutting_down = False - self._setup_lock = None + self._setup_locks_guard = threading.Lock() + self._setup_locks: dict[asyncio.AbstractEventLoop, asyncio.Lock] = {} self._credentials = credentials self.client = None self._loop_state_by_loop: dict[asyncio.AbstractEventLoop, _LoopState] = {} @@ -2741,6 +2765,13 @@ def _cleanup_stale_loop_states(self) -> None: """Removes entries for event loops that have been closed.""" stale = [loop for loop in self._loop_state_by_loop if loop.is_closed()] for loop in stale: + # Atomic claim (#6360 review round 3 P2): dict.pop is atomic under + # the GIL, so exactly one concurrent cleanup folds a given + # processor's counters — read-fold-delete raced, double-counting and + # raising KeyError. + state = self._loop_state_by_loop.pop(loop, None) + if state is None: + continue logger.warning( "Cleaning up stale loop state for closed loop %s (id=%s).", loop, @@ -2748,12 +2779,10 @@ def _cleanup_stale_loop_states(self) -> None: ) # Preserve the dead processor's loss accounting before discarding it # (#6360 review round 2 P2-6), mirroring shutdown(). - state = self._loop_state_by_loop[loop] for reason, count in state.batch_processor.get_drop_stats().items(): self._local_drop_counts[reason] = ( self._local_drop_counts.get(reason, 0) + count ) - del self._loop_state_by_loop[loop] # API Compatibility: These class-level attributes mask the dynamic # properties from static analysis tools (preventing "breaking changes"), @@ -3212,11 +3241,16 @@ def _maybe_upgrade_schema(self, existing_table: bigquery.Table) -> None: e, exc_info=True, ) - # The table is verifiably missing required fields at this point; - # swallowing the failure would let _ensure_started mark the plugin - # ready against a table every later write can fail on, with no - # readiness retry (#6360 review round 2 P1-2). - raise + if new_fields or updated_records: + # The table is verifiably missing required fields; swallowing the + # failure would let _ensure_started mark the plugin ready against + # a table every later write can fail on, with no readiness retry + # (#6360 review round 2 P1-2). + raise + # Label-only refresh failed (e.g. a labels policy): the table schema + # itself is write-compatible, so readiness must not be blocked — + # the stale label is retried on the next run (#6360 review round 3 + # P1-2). def _project_view_columns(self, extra_cols: list[str]) -> list[str]: """Drops derived view expressions that reference a denied column. @@ -3359,7 +3393,8 @@ async def shutdown(self, timeout: float | None = None) -> None: def __getstate__(self): """Custom pickling to exclude non-picklable runtime objects.""" state = self.__dict__.copy() - state["_setup_lock"] = None + state["_setup_locks_guard"] = None + state["_setup_locks"] = {} state["client"] = None state["_loop_state_by_loop"] = {} state["_write_stream_name"] = None @@ -3382,7 +3417,10 @@ def __setstate__(self, state): state.setdefault("_local_drop_counts", {}) state.setdefault("_setup_failures", 0) state.setdefault("_setup_retry_at", 0.0) + state.pop("_setup_lock", None) # replaced by per-loop locks + state.setdefault("_setup_locks", {}) self.__dict__.update(state) + self._setup_locks_guard = threading.Lock() # Pickles from older code bypass __init__, so re-validate the restored # configuration: e.g. a legacy retry_config with max_retries=NaN would # otherwise skip the write loop silently (#6360 review round 2 P2-7). @@ -3429,7 +3467,8 @@ def _reset_runtime_state(self) -> None: pass # Clear all runtime state. - self._setup_lock = None + self._setup_locks_guard = threading.Lock() + self._setup_locks = {} self.client = None self._loop_state_by_loop = {} self._write_stream_name = None @@ -3443,6 +3482,21 @@ def _reset_runtime_state(self) -> None: self._is_shutting_down = False self._init_pid = os.getpid() + def _get_setup_lock(self) -> asyncio.Lock: + """Returns the setup lock for the CURRENT event loop. + + asyncio.Lock is loop-bound: sharing one across loops/threads is not + thread-safe (#6360 review round 3 P2). The dict itself is guarded by a + threading.Lock so concurrent loops can mint their locks safely. + """ + loop = asyncio.get_running_loop() + with self._setup_locks_guard: + lock = self._setup_locks.get(loop) + if lock is None: + lock = asyncio.Lock() + self._setup_locks[loop] = lock + return lock + def _count_local_drop(self, reason: str) -> None: """Counts a row lost before/outside any BatchProcessor (#6356 P2).""" self._local_drop_counts[reason] = self._local_drop_counts.get(reason, 0) + 1 @@ -3478,10 +3532,13 @@ async def _ensure_started(self, **kwargs) -> None: if self._init_pid != 0 and os.getpid() != self._init_pid: self._reset_runtime_state() if not self._started: - # Kept original lock name as it was not explicitly changed. - if self._setup_lock is None: - self._setup_lock = asyncio.Lock() - async with self._setup_lock: + # Per-loop coalescing: asyncio locks are loop-bound, so one shared + # lock across loops/threads raises "Non-thread-safe operation" and + # can strand waiters on other loops (#6360 review round 3 P2). + # _get_setup_lock hands each event loop its own lock; _lazy_setup + # re-checks _started inside, so a rare cross-loop overlap costs at + # most one duplicate idempotent setup RPC. + async with self._get_setup_lock(): if not self._started: if ( self._startup_error is not None diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 6d225f8ce7..fcbe54d219 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -2559,7 +2559,7 @@ async def test_pickle_safety(self, mock_auth_default, mock_bq_client): pickled = pickle.dumps(plugin) unpickled = pickle.loads(pickled) assert unpickled.project_id == PROJECT_ID - assert unpickled._setup_lock is None + assert unpickled._setup_locks == {} assert unpickled._executor is None # Start the plugin await plugin._ensure_started() @@ -2570,7 +2570,7 @@ async def test_pickle_safety(self, mock_auth_default, mock_bq_client): 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._setup_locks == {} assert unpickled_started._executor is None assert not unpickled_started._loop_state_by_loop finally: @@ -6058,7 +6058,7 @@ def test_reset_runtime_state_clears_fields(self): plugin._executor = mock.MagicMock() plugin.offloader = mock.MagicMock() plugin.parser = mock.MagicMock() - plugin._setup_lock = mock.MagicMock() + plugin._setup_locks = {mock.MagicMock(): mock.MagicMock()} # Keep pure-data fields plugin._schema = ["kept"] plugin.arrow_schema = "kept_arrow" @@ -6073,7 +6073,7 @@ def test_reset_runtime_state_clears_fields(self): assert plugin._executor is None assert plugin.offloader is None assert plugin.parser is None - assert plugin._setup_lock is None + assert plugin._setup_locks == {} # Pure-data fields are preserved assert plugin._schema == ["kept"] assert plugin.arrow_schema == "kept_arrow" @@ -10273,3 +10273,143 @@ async def upload_content(self, data, mime, path): offloader._upload_sync(b"data", "image/png", "p") _, kwargs = bucket.blob.return_value.upload_from_string.call_args assert kwargs.get("if_generation_match") == 0 + + def test_unmaterializable_json_blob_fails_closed(self): + """Valid JSON that Python cannot materialize becomes a sentinel. + + Integers over the interpreter digit limit raise a plain ValueError + from json.loads on syntactically valid JSON; returning the raw string + would leak members the sanitizer never inspected (#6360 review round 3 + P1-1). + """ + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + blob = '{"access_token": "SECRET-BIGINT", "n": ' + "9" * 5000 + "}" + out, _ = truncate({"blob": blob}, 500 * 1024) + assert "SECRET-BIGINT" not in json.dumps(out) + assert out["blob"] == "[UNPARSEABLE_JSON_BLOB]" + + def test_label_only_upgrade_failure_does_not_block_readiness(self): + """A label-only update_table failure must not fail setup. + + The table schema is write-compatible; only the governance label is + stale. Blocking readiness turned every event into setup_unavailable + although writes would succeed (#6360 review round 3 P1-2). + """ + 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}" + plugin._schema = bigquery_agent_analytics_plugin._get_events_schema() + existing = mock.MagicMock(spec=bigquery.Table) + # Identical schema: no new fields, no updated records. + existing.schema = list(plugin._schema) + existing.labels = {} # stale version label only + plugin.client.get_table.return_value = existing + plugin.client.update_table.side_effect = Exception("labels forbidden") + + # Does not raise; the stale label is retried on the next run. + plugin._ensure_schema_exists() + plugin.client.update_table.assert_called_once() + + def test_ensure_started_coalesces_across_event_loops( + self, mock_auth_default, mock_bq_client + ): + """_ensure_started must be safe when called from multiple loops. + + One shared asyncio.Lock is loop-bound: a second thread's loop raised + 'Non-thread-safe operation' and could strand waiters (#6360 review + round 3 P2). Per-loop locks make each loop coalesce independently. + """ + _ = mock_auth_default, mock_bq_client + import threading + + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + setup_calls = [] + + async def fake_lazy_setup(**kwargs): + setup_calls.append(threading.get_ident()) + await asyncio.sleep(0.05) + + errors: list[BaseException] = [] + + def run_in_fresh_loop(): + try: + with mock.patch.object( + plugin, "_lazy_setup", side_effect=fake_lazy_setup + ): + asyncio.run(plugin._ensure_started()) + except BaseException as e: # noqa: BLE001 - collecting for assertion + errors.append(e) + + threads = [threading.Thread(target=run_in_fresh_loop) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + assert not errors, f"cross-loop startup raised: {errors}" + + def test_concurrent_stale_cleanup_folds_once( + self, mock_auth_default, mock_bq_client + ): + """Repeated/concurrent cleanups fold a processor's counters exactly once. + + Read-fold-delete raced: two cleanups produced doubled counts and a + KeyError; the pop-claim makes folding idempotent (#6360 review round 3 + P2). + """ + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + dead_loop = mock.MagicMock() + dead_loop.is_closed.return_value = True + state = mock.MagicMock() + state.batch_processor.get_drop_stats.return_value = {"write_failed": 7} + plugin._loop_state_by_loop[dead_loop] = state + + plugin._cleanup_stale_loop_states() + plugin._cleanup_stale_loop_states() # second pass: nothing left to claim + + assert plugin.get_drop_stats().get("write_failed") == 7 + + @pytest.mark.asyncio + async def test_close_counts_lost_rows_like_shutdown(self): + """close() shares shutdown()'s drain/accounting for stranded rows + (#6360 review round 3 P2).""" + write_started = asyncio.Event() + + async def hung_writer(batch): + write_started.set() + await asyncio.sleep(3600) + + processor = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=mock.MagicMock(), + write_stream="stream", + batch_size=1, + flush_interval=0.05, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=0.1, + ) + with mock.patch.object( + processor, "_write_rows_with_retry", side_effect=hung_writer + ): + await processor.start() + await processor.append({"row": 1}) + await write_started.wait() + await processor.append({"row": 2}) + await processor.close() + + assert processor.get_drop_stats().get("shutdown_timeout") == 2 + assert processor._queue.empty() From 100e17ed44c81d6e2eeab91cce3b3b400242ee8a Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 01:31:32 -0700 Subject: [PATCH 5/6] =?UTF-8?q?fix(plugins):=20address=20review=20round=20?= =?UTF-8?q?4=20=E2=80=94=20fail-closed=20malformed=20blobs,=20size-gated?= =?UTF-8?q?=20parsing,=20cross-loop=20setup=20future?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All eight findings from the round-4 review at 5c5521d1: - P1-1: every container-shaped ({ or [) parse failure now fails closed to [UNPARSEABLE_JSON_BLOB] — trailing garbage on valid credential JSON (incl. escaped keys) can no longer bypass redaction. Malformed-object, trailing-garbage, and escaped-key regressions added. - P1-2: _sanitize_json_blob enforces max_content_length BEFORE json.loads; over-limit container blobs fail closed without being materialized (test patches json.loads to prove it is never called). - P1-3: shared setup is coalesced across event loops/threads with a concurrent.futures.Future claimed under a briefly-held threading.Lock (never across an await); waiters await it via asyncio.wrap_future from their own loops. Exactly one _lazy_setup runs (regression asserts setup_calls == 1 across two threads); a failing owner leaves consistent _started/_startup_error/backoff state for all participants. This also removes the per-loop _setup_locks map entirely (P2: it retained every closed loop forever). - P2: tuple subclasses (namedtuples) normalize to plain lists in the sanitizer instead of raising TypeError and dropping the row. Duck-typed conversions (model_dump/dict/to_dict) now require progress to a real dict/list before recursing, so Mock-like graphs settle at the stringify fallback instead of churning to the depth cap. - P2: _loop_state_by_loop ownership changes are guarded by a threading.Lock; cleanup snapshots keys under the guard, evaluates is_closed() outside it, and claims each state atomically (regression: insertion during is_closed() no longer raises 'dictionary changed size during iteration'). - P2: the [MAX_DEPTH_EXCEEDED] replacement now reports is_truncated=True — it discards real payload, unlike [CIRCULAR_REFERENCE]. Row-level test included. - P2: retry delays are validated as finite and NON-negative; the long-supported max_retries=0/initial_delay=0/max_delay=0 config constructs again (pinned by test). --- .../bigquery_agent_analytics_plugin.py | 263 ++++++++++------- .../test_bigquery_agent_analytics_plugin.py | 271 +++++++++++++++++- 2 files changed, 433 insertions(+), 101 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 1f58c82ce3..21a777038c 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -17,6 +17,7 @@ import asyncio import atexit import collections.abc +from concurrent.futures import Future as ConcurrentFuture from concurrent.futures import ThreadPoolExecutor import contextvars import dataclasses @@ -425,7 +426,7 @@ def _extract_tool_declarations( def _sanitize_json_blob( - value: str, seen: set[int], depth: int = 0 + value: str, seen: set[int], depth: int = 0, max_len: int = -1 ) -> tuple[str, bool]: """Redacts sensitive keys inside a JSON-encoded string blob. @@ -442,6 +443,14 @@ def _sanitize_json_blob( if not stripped.startswith(("{", "[")): return value, False + # Enforce the configured content limit BEFORE materializing: json.loads + # runs synchronously on the callback path and can allocate far beyond + # the limit for a multi-megabyte attribute (#6360 review round 4 P1-2). + # Truncating the raw JSON prefix instead could both retain a secret and + # emit invalid JSON, so over-limit container blobs fail closed. + if max_len != -1 and len(value) > max_len: + return "[UNPARSEABLE_JSON_BLOB]", True + # json.loads silently keeps only the LAST duplicate member, so a blob like # {"access_token":"SECRET","access_token":"x"} can compare equal after # sanitization while the raw string still carries the secret (#6360 review @@ -468,16 +477,14 @@ def _pairs_hook(pairs: list[tuple[str, Any]]) -> dict[str, Any]: if sanitized == parsed and not saw_duplicate_key: return value, False return json.dumps(sanitized), True - except json.JSONDecodeError: - # Not JSON at all — nothing to inspect; the outer string handling - # (length truncation) still applies. - return value, False except (TypeError, ValueError, RecursionError, MemoryError): - # Syntactically valid JSON that Python cannot materialize — e.g. - # integers over sys.get_int_max_str_digits() raise a plain ValueError - # (#6360 review round 3 P1-1) — or a blob too deep/large to inspect - # (round 2 P2-5). Either way it cannot be verified secret-free: fail - # CLOSED to a sentinel instead of passing it through unexamined. + # Container-shaped but unparseable — malformed JSON / trailing garbage + # (a one-character suffix on valid credential JSON must not bypass + # redaction, #6360 review round 4 P1-1), integers over the interpreter + # digit limit (round 3 P1-1), or a blob too deep/large to inspect + # (round 2 P2-5). None of these can be verified secret-free — and a + # raw-substring fallback is bypassable via JSON string escapes — so + # fail CLOSED to a sentinel. return "[UNPARSEABLE_JSON_BLOB]", True @@ -532,15 +539,26 @@ def _validate_runtime_config(config: "BigQueryLoggerConfig") -> None: ) retry = config.retry_config _require_count("retry_config.max_retries", retry.max_retries, 0) + # Delays are finite and NON-NEGATIVE: zero-delay immediate retries are + # long-supported (asyncio.sleep(0) is valid) and existing configs use + # max_retries=0, initial_delay=0, max_delay=0 (#6360 review round 4 P2). initial_delay = _require_finite( - "retry_config.initial_delay", retry.initial_delay, 0 + "retry_config.initial_delay", retry.initial_delay, -1 ) + if initial_delay < 0: + raise ValueError( + f"retry_config.initial_delay must be >= 0, got {retry.initial_delay}." + ) multiplier = _require_finite("retry_config.multiplier", retry.multiplier, 0) if multiplier < 1: raise ValueError( f"retry_config.multiplier must be >= 1, got {retry.multiplier}." ) - max_delay = _require_finite("retry_config.max_delay", retry.max_delay, 0) + max_delay = _require_finite("retry_config.max_delay", retry.max_delay, -1) + if max_delay < 0: + raise ValueError( + f"retry_config.max_delay must be >= 0, got {retry.max_delay}." + ) if max_delay < initial_delay: raise ValueError( "retry_config.max_delay must be >= initial_delay, got" @@ -584,11 +602,11 @@ def _recursive_smart_truncate( # manufacture NEW objects on each duck-typed access (e.g. anything whose # model_dump()/dict()/to_dict() returns a fresh wrapper — unittest Mocks # being the canonical case). Without this cap such graphs recurse - # unboundedly. Like "[CIRCULAR_REFERENCE]", this is a structural - # sentinel rather than content truncation, so it does not flip - # is_truncated. + # unboundedly. The replacement discards real payload, so it reports + # truncation (#6360 review round 4 P2) — unlike "[CIRCULAR_REFERENCE]", + # which replaces a back-reference, not data. if depth >= _MAX_SANITIZE_DEPTH: - return "[MAX_DEPTH_EXCEEDED]", False + return "[MAX_DEPTH_EXCEEDED]", True obj_id = id(obj) if obj_id in seen: @@ -608,7 +626,10 @@ def _recursive_smart_truncate( try: if isinstance(obj, str): - obj, _ = _sanitize_json_blob(obj, seen, depth) + obj, blob_replaced = _sanitize_json_blob(obj, seen, depth, max_len) + if blob_replaced and obj == "[UNPARSEABLE_JSON_BLOB]": + # The original string was discarded wholesale. + return obj, True if max_len != -1 and len(obj) > max_len: return obj[:max_len] + "...[TRUNCATED]", True return obj, False @@ -641,31 +662,44 @@ def _recursive_smart_truncate( if trunc: truncated_any = True new_list.append(val) - return type(obj)(new_list), truncated_any + 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 (#6360 review round 4 P2). JSON + # does not preserve the subclass identity anyway — emit a plain + # list. + return new_list, truncated_any 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, depth + 1) elif hasattr(obj, "model_dump") and callable(obj.model_dump): - # Pydantic v2 + # Pydantic v2. Only recurse if the conversion made PROGRESS toward a + # JSON-native container: 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. try: - return _recursive_smart_truncate( - obj.model_dump(), max_len, seen, depth + 1 - ) + dumped = obj.model_dump() + if isinstance(dumped, (dict, list)): + return _recursive_smart_truncate(dumped, max_len, seen, depth + 1) except Exception: pass elif hasattr(obj, "dict") and callable(obj.dict): - # Pydantic v1 + # Pydantic v1 (same progress requirement as above). try: - return _recursive_smart_truncate(obj.dict(), max_len, seen, depth + 1) + dumped = obj.dict() + if isinstance(dumped, (dict, list)): + return _recursive_smart_truncate(dumped, max_len, seen, depth + 1) except Exception: pass elif hasattr(obj, "to_dict") and callable(obj.to_dict): - # Common pattern for custom objects + # Common pattern for custom objects (same progress requirement). try: - return _recursive_smart_truncate( - obj.to_dict(), max_len, seen, depth + 1 - ) + dumped = obj.to_dict() + if isinstance(dumped, (dict, list)): + return _recursive_smart_truncate(dumped, max_len, seen, depth + 1) except Exception: pass elif obj is None or isinstance(obj, (int, float, bool)): @@ -2747,8 +2781,13 @@ def __init__( # failure). Merged into get_drop_stats() and survives shutdown. self._local_drop_counts: dict[str, int] = {} self._is_shutting_down = False - self._setup_locks_guard = threading.Lock() - self._setup_locks: dict[asyncio.AbstractEventLoop, asyncio.Lock] = {} + # Guards _setup_future/_started/_setup_* transitions across threads; + # held only for pointer swaps, never across an await. + self._setup_guard = threading.Lock() + self._setup_future: Optional["ConcurrentFuture[None]"] = None + # Guards ownership changes of _loop_state_by_loop (#6360 review round + # 4 P2): unsynchronized iteration raced concurrent insertion. + self._loop_states_guard = threading.Lock() self._credentials = credentials self.client = None self._loop_state_by_loop: dict[asyncio.AbstractEventLoop, _LoopState] = {} @@ -2763,13 +2802,19 @@ def __init__( def _cleanup_stale_loop_states(self) -> None: """Removes entries for event loops that have been closed.""" - stale = [loop for loop in self._loop_state_by_loop if loop.is_closed()] + # Snapshot under the guard (#6360 review round 4 P2): iterating the + # live dict raced concurrent insertion ("dictionary changed size + # during iteration"). is_closed() is evaluated on the snapshot, + # outside the lock. + with self._loop_states_guard: + candidates = list(self._loop_state_by_loop) + stale = [loop for loop in candidates if loop.is_closed()] for loop in stale: - # Atomic claim (#6360 review round 3 P2): dict.pop is atomic under - # the GIL, so exactly one concurrent cleanup folds a given - # processor's counters — read-fold-delete raced, double-counting and - # raising KeyError. - state = self._loop_state_by_loop.pop(loop, None) + # Atomic claim (#6360 review round 3 P2): exactly one concurrent + # cleanup folds a given processor's counters — read-fold-delete + # raced, double-counting and raising KeyError. + with self._loop_states_guard: + state = self._loop_state_by_loop.pop(loop, None) if state is None: continue logger.warning( @@ -2916,7 +2961,8 @@ def get_credentials(): await batch_processor.start() state = _LoopState(write_client, batch_processor) - self._loop_state_by_loop[loop] = state + with self._loop_states_guard: + self._loop_state_by_loop[loop] = state atexit.register(self._atexit_cleanup, weakref.proxy(batch_processor)) @@ -3393,8 +3439,9 @@ async def shutdown(self, timeout: float | None = None) -> None: def __getstate__(self): """Custom pickling to exclude non-picklable runtime objects.""" state = self.__dict__.copy() - state["_setup_locks_guard"] = None - state["_setup_locks"] = {} + state["_setup_guard"] = None + state["_setup_future"] = None + state["_loop_states_guard"] = None state["client"] = None state["_loop_state_by_loop"] = {} state["_write_stream_name"] = None @@ -3417,10 +3464,13 @@ def __setstate__(self, state): state.setdefault("_local_drop_counts", {}) state.setdefault("_setup_failures", 0) state.setdefault("_setup_retry_at", 0.0) - state.pop("_setup_lock", None) # replaced by per-loop locks - state.setdefault("_setup_locks", {}) + state.pop("_setup_lock", None) # replaced by cross-loop future + state.pop("_setup_locks", None) + state.pop("_setup_locks_guard", None) self.__dict__.update(state) - self._setup_locks_guard = threading.Lock() + self._setup_guard = threading.Lock() + self._setup_future = None + self._loop_states_guard = threading.Lock() # Pickles from older code bypass __init__, so re-validate the restored # configuration: e.g. a legacy retry_config with max_retries=NaN would # otherwise skip the write loop silently (#6360 review round 2 P2-7). @@ -3467,8 +3517,9 @@ def _reset_runtime_state(self) -> None: pass # Clear all runtime state. - self._setup_locks_guard = threading.Lock() - self._setup_locks = {} + self._setup_guard = threading.Lock() + self._setup_future = None + self._loop_states_guard = threading.Lock() self.client = None self._loop_state_by_loop = {} self._write_stream_name = None @@ -3482,21 +3533,6 @@ def _reset_runtime_state(self) -> None: self._is_shutting_down = False self._init_pid = os.getpid() - def _get_setup_lock(self) -> asyncio.Lock: - """Returns the setup lock for the CURRENT event loop. - - asyncio.Lock is loop-bound: sharing one across loops/threads is not - thread-safe (#6360 review round 3 P2). The dict itself is guarded by a - threading.Lock so concurrent loops can mint their locks safely. - """ - loop = asyncio.get_running_loop() - with self._setup_locks_guard: - lock = self._setup_locks.get(loop) - if lock is None: - lock = asyncio.Lock() - self._setup_locks[loop] = lock - return lock - def _count_local_drop(self, reason: str) -> None: """Counts a row lost before/outside any BatchProcessor (#6356 P2).""" self._local_drop_counts[reason] = self._local_drop_counts.get(reason, 0) + 1 @@ -3531,43 +3567,78 @@ async def _ensure_started(self, **kwargs) -> None: # different process. if self._init_pid != 0 and os.getpid() != self._init_pid: self._reset_runtime_state() - if not self._started: - # Per-loop coalescing: asyncio locks are loop-bound, so one shared - # lock across loops/threads raises "Non-thread-safe operation" and - # can strand waiters on other loops (#6360 review round 3 P2). - # _get_setup_lock hands each event loop its own lock; _lazy_setup - # re-checks _started inside, so a rare cross-loop overlap costs at - # most one duplicate idempotent setup RPC. - async with self._get_setup_lock(): - if not self._started: - if ( - self._startup_error is not None - and time.monotonic() < self._setup_retry_at - ): - # Still inside the backoff window from a previous failure. - return - try: - await self._lazy_setup(**kwargs) - self._started = True - self._startup_error = None - self._setup_failures = 0 - self._setup_retry_at = 0.0 - # Record the current PID so fork detection works for - # the rest of this instance's lifetime. - if self._init_pid == 0: - self._init_pid = os.getpid() - except Exception as e: - self._startup_error = e - self._setup_failures += 1 - backoff = min(60.0, 2.0 ** min(self._setup_failures, 6)) - self._setup_retry_at = time.monotonic() + backoff - logger.error( - "Failed to initialize BigQuery Plugin (attempt %d, next" - " retry in %.0fs): %s", - self._setup_failures, - backoff, - e, - ) + if self._started: + return + + # Cross-loop coalescing of the SHARED initialization (#6360 review + # round 4 P1-3): _lazy_setup mutates process-wide state (client, + # executor, parser, schema, views, retry bookkeeping), so exactly one + # caller may run it at a time — across event loops and threads, which + # a per-loop asyncio.Lock cannot provide and a shared one cannot + # survive. A concurrent.futures.Future is claimed under a briefly-held + # threading.Lock (never held across an await); the owner runs setup, + # every other caller awaits the same future via asyncio.wrap_future + # from its own loop. Loop-local writer state stays separate in + # _get_loop_state(). + setup_future: Optional["ConcurrentFuture[None]"] = None + is_owner = False + with self._setup_guard: + if self._started: + return + if self._setup_future is not None: + setup_future = self._setup_future + elif ( + self._startup_error is not None + and time.monotonic() < self._setup_retry_at + ): + # Still inside the backoff window from a previous failure. + return + else: + setup_future = ConcurrentFuture() + self._setup_future = setup_future + is_owner = True + + assert setup_future is not None # every fall-through branch assigns it + + if not is_owner: + try: + await asyncio.wrap_future(setup_future) + except Exception: + # The owner already recorded the failure and backoff; waiters + # degrade the same way the owner does (row counted as + # setup_unavailable by the caller). + pass + return + + try: + await self._lazy_setup(**kwargs) + except Exception as e: + with self._setup_guard: + self._startup_error = e + self._setup_failures += 1 + backoff = min(60.0, 2.0 ** min(self._setup_failures, 6)) + self._setup_retry_at = time.monotonic() + backoff + self._setup_future = None + logger.error( + "Failed to initialize BigQuery Plugin (attempt %d, next" + " retry in %.0fs): %s", + self._setup_failures, + backoff, + e, + ) + setup_future.set_exception(e) + else: + with self._setup_guard: + self._started = True + self._startup_error = None + self._setup_failures = 0 + self._setup_retry_at = 0.0 + self._setup_future = None + # Record the current PID so fork detection works for + # the rest of this instance's lifetime. + if self._init_pid == 0: + self._init_pid = os.getpid() + setup_future.set_result(None) @staticmethod def _resolve_ids( diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index fcbe54d219..6baaf96462 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -2559,7 +2559,7 @@ async def test_pickle_safety(self, mock_auth_default, mock_bq_client): pickled = pickle.dumps(plugin) unpickled = pickle.loads(pickled) assert unpickled.project_id == PROJECT_ID - assert unpickled._setup_locks == {} + assert unpickled._setup_future is None assert unpickled._executor is None # Start the plugin await plugin._ensure_started() @@ -2570,7 +2570,7 @@ async def test_pickle_safety(self, mock_auth_default, mock_bq_client): unpickled_started = pickle.loads(pickled_started) assert unpickled_started.project_id == PROJECT_ID # Runtime objects should be None after unpickling - assert unpickled_started._setup_locks == {} + assert unpickled_started._setup_future is None assert unpickled_started._executor is None assert not unpickled_started._loop_state_by_loop finally: @@ -6058,7 +6058,7 @@ def test_reset_runtime_state_clears_fields(self): plugin._executor = mock.MagicMock() plugin.offloader = mock.MagicMock() plugin.parser = mock.MagicMock() - plugin._setup_locks = {mock.MagicMock(): mock.MagicMock()} + plugin._setup_future = mock.MagicMock() # Keep pure-data fields plugin._schema = ["kept"] plugin.arrow_schema = "kept_arrow" @@ -6073,7 +6073,7 @@ def test_reset_runtime_state_clears_fields(self): assert plugin._executor is None assert plugin.offloader is None assert plugin.parser is None - assert plugin._setup_locks == {} + assert plugin._setup_future is None # Pure-data fields are preserved assert plugin._schema == ["kept"] assert plugin.arrow_schema == "kept_arrow" @@ -9747,7 +9747,7 @@ def test_invalid_runtime_config_rejected_at_construction( dict(queue_max_size=0), dict(max_content_length=0), dict(retry_config=retry(max_retries=-1)), - dict(retry_config=retry(initial_delay=0.0)), + dict(retry_config=retry(initial_delay=-1.0)), dict(retry_config=retry(multiplier=0.5)), dict(retry_config=retry(initial_delay=5.0, max_delay=1.0)), ] @@ -10413,3 +10413,264 @@ async def hung_writer(batch): assert processor.get_drop_stats().get("shutdown_timeout") == 2 assert processor._queue.empty() + + def test_malformed_container_blobs_fail_closed(self): + """Container-shaped strings that fail to parse become the sentinel. + + One trailing character on valid credential JSON must not bypass + redaction, including with escaped keys (#6360 review round 4 P1-1). + """ + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + cases = [ + '{"access\\u005ftoken":"SECRET-TRAIL"} trailing', + '{"access_token":"SECRET-MALFORMED"', + '[{"api_key":"SECRET-ARRAY"}, oops]', + ] + for blob in cases: + out, _ = truncate({"blob": blob}, 10000) + assert "SECRET" not in json.dumps(out), blob + assert out["blob"] == "[UNPARSEABLE_JSON_BLOB]", blob + + def test_over_limit_blob_never_parsed(self): + """json.loads must not run for container blobs over the content limit. + + Materializing a multi-megabyte attribute blocks the callback loop and + allocates far beyond the configured limit (#6360 review round 4 P1-2). + """ + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + big_blob = '{"k": "' + "x" * 5000 + '"}' + with mock.patch.object( + bigquery_agent_analytics_plugin.json, + "loads", + side_effect=AssertionError("json.loads must not be called"), + ): + out, truncated = truncate({"blob": big_blob}, 100) + assert out["blob"] == "[UNPARSEABLE_JSON_BLOB]" + assert truncated + + def test_shared_setup_runs_exactly_once_across_loops( + self, mock_auth_default, mock_bq_client + ): + """Concurrent loops coalesce onto ONE shared setup (#6360 round 4 P1-3). + + Per-loop locks let both loops run _lazy_setup, which mutates shared + clients/executor/parser state across awaits and is not idempotent. + """ + _ = mock_auth_default, mock_bq_client + import threading + + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + setup_calls = [] + release = threading.Event() + + async def slow_setup(**kwargs): + setup_calls.append(threading.get_ident()) + await asyncio.get_running_loop().run_in_executor(None, release.wait) + + errors: list[BaseException] = [] + barrier = threading.Barrier(2) + + def run_in_fresh_loop(): + try: + with mock.patch.object(plugin, "_lazy_setup", side_effect=slow_setup): + barrier.wait(timeout=5) + asyncio.run(plugin._ensure_started()) + except BaseException as e: # noqa: BLE001 + errors.append(e) + + threads = [threading.Thread(target=run_in_fresh_loop) for _ in range(2)] + for t in threads: + t.start() + import time as time_module + + time_module.sleep(0.3) # both loops reach the coalescing point + release.set() + for t in threads: + t.join(timeout=10) + + assert not errors, f"cross-loop startup raised: {errors}" + assert len(setup_calls) == 1, f"shared setup ran {len(setup_calls)} times" + assert plugin._started is True + assert plugin._startup_error is None + assert plugin._setup_future is None + + def test_failed_shared_setup_is_consistent_across_loops( + self, mock_auth_default, mock_bq_client + ): + """A failing owner leaves consistent shared state for every waiter.""" + _ = mock_auth_default, mock_bq_client + import threading + + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + + async def failing_setup(**kwargs): + await asyncio.sleep(0.05) + raise RuntimeError("setup boom") + + errors: list[BaseException] = [] + + def run_in_fresh_loop(): + try: + with mock.patch.object( + plugin, "_lazy_setup", side_effect=failing_setup + ): + asyncio.run(plugin._ensure_started()) + except BaseException as e: # noqa: BLE001 + errors.append(e) + + threads = [threading.Thread(target=run_in_fresh_loop) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert not errors # _ensure_started never raises to callers + assert plugin._started is False + assert plugin._startup_error is not None + assert plugin._setup_future is None # cleared for the next retry window + + @pytest.mark.asyncio + async def test_namedtuple_attribute_does_not_drop_row( + 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, + ): + """A namedtuple in attributes serializes as a list, not a TypeError. + + Reconstructing tuple subclasses positionally raised in the final pass + and the safe callback dropped the entire row (#6360 review round 4 P2). + """ + _ = mock_auth_default, mock_bq_client + import collections + + Point = collections.namedtuple("Point", ["x", "y"]) + 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(invocation_context) + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData( + extra_attributes={"point": Point(1, 2)}, + ), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + attrs = json.loads(log_entry["attributes"]) + assert attrs["point"] == [1, 2] + + def test_setup_future_leaves_no_loop_references( + self, mock_auth_default, mock_bq_client + ): + """Repeated fresh-loop startups retain no per-loop setup structures. + + The per-loop lock map kept strong references to every closed loop + (#6360 review round 4 P2); the cross-loop future replaces it. + """ + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + + async def noop_setup(**kwargs): + return None + + for _ in range(4): + plugin._started = False + with mock.patch.object(plugin, "_lazy_setup", side_effect=noop_setup): + asyncio.run(plugin._ensure_started()) + assert plugin._setup_future is None + assert not hasattr(plugin, "_setup_locks") + + def test_cleanup_survives_concurrent_insertion( + self, mock_auth_default, mock_bq_client + ): + """Cleanup snapshots keys, so insertion during is_closed() cannot raise + 'dictionary changed size during iteration' (#6360 review round 4 P2).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + dead_loop = mock.MagicMock() + state = mock.MagicMock() + state.batch_processor.get_drop_stats.return_value = {"write_failed": 7} + + def is_closed_and_mutate(): + # Simulates another thread inserting mid-scan. + plugin._loop_state_by_loop[mock.MagicMock()] = mock.MagicMock() + return True + + dead_loop.is_closed.side_effect = is_closed_and_mutate + plugin._loop_state_by_loop[dead_loop] = state + + plugin._cleanup_stale_loop_states() # must not raise + assert plugin.get_drop_stats().get("write_failed") == 7 + + @pytest.mark.asyncio + async def test_depth_capped_payload_flags_row_truncated( + 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, + ): + """A real payload cut off by the depth cap marks the ROW as truncated + (#6360 review round 4 P2).""" + _ = mock_auth_default, mock_bq_client + deep: dict = {"leaf": "payload"} + for _ in range(60): + deep = {"level": deep} + 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(invocation_context) + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData( + extra_attributes={"deep": deep}, + ), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + assert "[MAX_DEPTH_EXCEEDED]" in log_entry["attributes"] + assert log_entry["is_truncated"] is True + + def test_zero_delay_retry_config_still_constructs( + self, mock_auth_default, mock_bq_client + ): + """Long-supported zero-delay retry configs must not be rejected + (#6360 review round 4 P2).""" + _ = mock_auth_default, mock_bq_client + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + retry_config=bigquery_agent_analytics_plugin.RetryConfig( + max_retries=0, initial_delay=0, max_delay=0 + ) + ) + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config + ) + assert plugin.config.retry_config.max_retries == 0 From 70dee9b918415fa48416f10d37ca5f80569dcf5b Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 02:44:06 -0700 Subject: [PATCH 6/6] =?UTF-8?q?fix(plugins):=20address=20review=20round=20?= =?UTF-8?q?5=20=E2=80=94=20cancellation-safe=20setup,=20lifecycle=20genera?= =?UTF-8?q?tion,=20close(),=20sanitizer=20shapes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All nine findings from the round-5 review at 100e17ed: - P1: the setup rendezvous is cancellation-safe. Waiters await asyncio.shield(wrap_future(...)) so their cancellation cannot cancel the shared future; a cancelled owner clears the rendezvous, fails waiters with an ordinary aborted error, and re-raises; publication is guarded by future.done(). Owner-cancel and waiter-cancel regressions. - P1: lifecycle generation. shutdown() bumps _generation under _setup_guard and forces _started=False; an in-flight setup completing afterwards observes the stale generation, aborts (shutdown_race), and cannot resurrect the plugin. shutdown() iterates a guarded snapshot of loop states (concurrent publication raised 'dictionary changed size during iteration'); _get_loop_state and the append path reject work during shutdown with shutdown_race accounting. - P1: BigQueryAgentAnalyticsPlugin.close() overrides the BasePlugin no-op and runs the full shutdown path, so Runner.close() -> PluginManager.close() actually releases queues/clients/executors. - P1: sanitizer credential shapes — bytes/bytearray decode and re-enter the sanitizer; BOM-prefixed JSON is detected and parsed; Mapping results from model_dump/dict/to_dict are accepted; the unknown-object fallback re-enters the string sanitizer so __str__-returned credential JSON is redacted (repr truncation does not flip is_truncated). - P2: shutdown clears parser/offloader so a restart rebuilds them instead of scheduling GCS uploads on the terminated executor. - P2: stale-loop cleanup drains and counts queued rows (stale_loop reason) and best-effort releases the dead state's transport. - P2: the sanitizer carries a total node budget (_MAX_SANITIZE_NODES); a million-scalar list now stops at the budget with a sentinel and is_truncated=True. - P2: cross-loop regressions use event/barrier rendezvous, assert exactly one setup call in success AND failure modes, and assert thread termination. - P2: get_drop_stats() documents loss-incident semantics per reason (formatter_failed = replaced-but-written; others = dropped). --- .../bigquery_agent_analytics_plugin.py | 251 +++++++++++++++--- .../test_bigquery_agent_analytics_plugin.py | 214 ++++++++++++++- 2 files changed, 420 insertions(+), 45 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 21a777038c..8ca3381707 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -424,9 +424,19 @@ def _extract_tool_declarations( # typing); the cap turns unbounded recursion into a redacted leaf. _MAX_SANITIZE_DEPTH = 50 +# Total nodes one sanitizer invocation may visit: depth and per-string size +# are bounded, but width was not — a million-scalar list burned ~1s of +# synchronous callback time (#6360 review round 5 P2). The remainder is +# replaced with a sentinel and the row is flagged truncated. +_MAX_SANITIZE_NODES = 100_000 + def _sanitize_json_blob( - value: str, seen: set[int], depth: int = 0, max_len: int = -1 + value: str, + seen: set[int], + depth: int = 0, + max_len: int = -1, + budget: Optional[list[int]] = None, ) -> tuple[str, bool]: """Redacts sensitive keys inside a JSON-encoded string blob. @@ -439,7 +449,7 @@ def _sanitize_json_blob( strings that do not parse, or that need no redaction, are returned unchanged (no cosmetic re-serialization). """ - stripped = value.lstrip() + stripped = value.lstrip("\ufeff \t\r\n") if not stripped.startswith(("{", "[")): return value, False @@ -448,7 +458,7 @@ def _sanitize_json_blob( # the limit for a multi-megabyte attribute (#6360 review round 4 P1-2). # Truncating the raw JSON prefix instead could both retain a secret and # emit invalid JSON, so over-limit container blobs fail closed. - if max_len != -1 and len(value) > max_len: + if max_len != -1 and len(stripped) > max_len: return "[UNPARSEABLE_JSON_BLOB]", True # json.loads silently keeps only the LAST duplicate member, so a blob like @@ -468,12 +478,14 @@ def _pairs_hook(pairs: list[tuple[str, Any]]) -> dict[str, Any]: return result try: - parsed = json.loads(value, object_pairs_hook=_pairs_hook) + parsed = json.loads(stripped, object_pairs_hook=_pairs_hook) if not isinstance(parsed, (dict, list)): return value, False # Redact only (max_len=-1): length truncation is applied by the caller # on the re-serialized string, keeping single responsibility per pass. - sanitized, _ = _recursive_smart_truncate(parsed, -1, seen, depth + 1) + sanitized, _ = _recursive_smart_truncate( + parsed, -1, seen, depth + 1, budget + ) if sanitized == parsed and not saw_duplicate_key: return value, False return json.dumps(sanitized), True @@ -580,6 +592,7 @@ def _recursive_smart_truncate( 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. @@ -597,6 +610,11 @@ def _recursive_smart_truncate( """ 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 # Depth cap: id()-based cycle detection cannot catch object graphs that # manufacture NEW objects on each duck-typed access (e.g. anything whose @@ -626,13 +644,25 @@ def _recursive_smart_truncate( try: if isinstance(obj, str): - obj, blob_replaced = _sanitize_json_blob(obj, seen, depth, max_len) + obj, blob_replaced = _sanitize_json_blob( + obj, seen, depth, max_len, budget + ) if blob_replaced and obj == "[UNPARSEABLE_JSON_BLOB]": # The original string was discarded wholesale. return obj, True if max_len != -1 and len(obj) > max_len: return obj[:max_len] + "...[TRUNCATED]", True return obj, False + elif isinstance(obj, (bytes, bytearray)): + # Credential JSON frequently travels as bytes; stringifying it in + # the fallback bypassed blob redaction (#6360 review round 5 P1-4). + try: + decoded = bytes(obj).decode("utf-8") + except UnicodeDecodeError: + return "[BINARY_DATA]", False + return _recursive_smart_truncate( + decoded, max_len, seen, depth + 1, budget + ) elif isinstance(obj, collections.abc.Mapping): # Covers dict plus mapping views (MappingProxyType, UserDict, ...): # stringifying them in the fallback branch would bypass key redaction @@ -648,7 +678,9 @@ def _recursive_smart_truncate( new_dict[k] = "[REDACTED]" continue - val, trunc = _recursive_smart_truncate(v, max_len, seen, depth + 1) + val, trunc = _recursive_smart_truncate( + v, max_len, seen, depth + 1, budget + ) if trunc: truncated_any = True new_dict[k] = val @@ -658,7 +690,9 @@ 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, depth + 1) + val, trunc = _recursive_smart_truncate( + i, max_len, seen, depth + 1, budget + ) if trunc: truncated_any = True new_list.append(val) @@ -673,7 +707,9 @@ 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, depth + 1) + return _recursive_smart_truncate( + as_dict, max_len, seen, depth + 1, budget + ) elif hasattr(obj, "model_dump") and callable(obj.model_dump): # Pydantic v2. Only recurse if the conversion made PROGRESS toward a # JSON-native container: Mock-like objects answer every duck-typed @@ -682,33 +718,48 @@ def _recursive_smart_truncate( # at the stringify fallback. try: dumped = obj.model_dump() - if isinstance(dumped, (dict, list)): - return _recursive_smart_truncate(dumped, max_len, seen, depth + 1) + if isinstance(dumped, (collections.abc.Mapping, list)): + return _recursive_smart_truncate( + dumped, max_len, seen, depth + 1, budget + ) except Exception: pass elif hasattr(obj, "dict") and callable(obj.dict): # Pydantic v1 (same progress requirement as above). try: dumped = obj.dict() - if isinstance(dumped, (dict, list)): - return _recursive_smart_truncate(dumped, max_len, seen, depth + 1) + if isinstance(dumped, (collections.abc.Mapping, list)): + return _recursive_smart_truncate( + dumped, max_len, seen, depth + 1, budget + ) except Exception: pass elif hasattr(obj, "to_dict") and callable(obj.to_dict): # Common pattern for custom objects (same progress requirement). try: dumped = obj.to_dict() - if isinstance(dumped, (dict, list)): - return _recursive_smart_truncate(dumped, max_len, seen, depth + 1) + if isinstance(dumped, (collections.abc.Mapping, list)): + return _recursive_smart_truncate( + dumped, max_len, seen, depth + 1, budget + ) except Exception: pass elif obj is None or isinstance(obj, (int, float, bool)): # Basic types are safe return obj, False - # Fallback for unknown types: Convert to string to ensure JSON validity - # We return string representation of the object, which is a valid JSON string value. - return str(obj), False + # Fallback for unknown types: convert to string, then RE-ENTER the + # string sanitizer — an object whose __str__ returns credential JSON + # bypassed blob redaction otherwise (#6360 review round 5 P1-4). + # Truncating an object REPRESENTATION is not content truncation, so + # the flag is not propagated (pre-existing str(obj) semantics). + try: + sanitized_repr, _ = _recursive_smart_truncate( + str(obj), max_len, seen, depth + 1, budget + ) + return sanitized_repr, False + except Exception: + return "[UNSUPPORTED_OBJECT]", False finally: if is_compound: seen.remove(obj_id) @@ -2785,6 +2836,9 @@ def __init__( # held only for pointer swaps, never across an await. self._setup_guard = threading.Lock() self._setup_future: Optional["ConcurrentFuture[None]"] = None + # Lifecycle generation: shutdown() bumps it so an in-flight setup that + # completes afterwards cannot resurrect _started (#6360 round 5 P1-2). + self._generation = 0 # Guards ownership changes of _loop_state_by_loop (#6360 review round # 4 P2): unsynchronized iteration raced concurrent insertion. self._loop_states_guard = threading.Lock() @@ -2828,6 +2882,35 @@ def _cleanup_stale_loop_states(self) -> None: self._local_drop_counts[reason] = ( self._local_drop_counts.get(reason, 0) + count ) + # Rows still queued on the dead loop can never be written: count + # them instead of discarding silently (#6360 review round 5 P2). + stale_rows = 0 + queue = getattr(state.batch_processor, "_queue", None) + if isinstance(queue, asyncio.Queue): + try: + while True: + if queue.get_nowait() is not _SHUTDOWN_SENTINEL: + stale_rows += 1 + except asyncio.QueueEmpty: + pass + if stale_rows: + self._local_drop_counts["stale_loop"] = ( + self._local_drop_counts.get("stale_loop", 0) + stale_rows + ) + logger.warning( + "%d queued row(s) lost with closed loop %s.", stale_rows, id(loop) + ) + # Best-effort resource release; the loop is closed, so async + # transport teardown is not possible here. + try: + if state.write_client and getattr( + state.write_client, "transport", None + ): + close_fn = getattr(state.write_client.transport, "close", None) + if close_fn is not None and not asyncio.iscoroutinefunction(close_fn): + close_fn() + except Exception: + pass # API Compatibility: These class-level attributes mask the dynamic # properties from static analysis tools (preventing "breaking changes"), @@ -2912,6 +2995,11 @@ async def _get_loop_state(self) -> _LoopState: The loop-specific state object containing clients and processors. """ loop = asyncio.get_running_loop() + if self._is_shutting_down: + # A callback that passed the early check can resume here after + # shutdown started; publishing a fresh writer state now would leak + # it (#6360 review round 5 P1-2). + raise RuntimeError("BigQuery plugin is shutting down.") self._cleanup_stale_loop_states() if loop in self._loop_state_by_loop: return self._loop_state_by_loop[loop] @@ -2990,9 +3078,16 @@ def get_drop_stats(self) -> dict[str, int]: monitoring to detect data loss before it surfaces as missing rows. See BatchProcessor.get_drop_stats for the meaning of each reason. + Reasons are LOSS INCIDENTS, not uniformly dropped rows: + ``formatter_failed`` means the row WAS written with its content + replaced by a sentinel; ``setup_unavailable``, ``shutdown_race``, + ``shutdown_timeout``, and ``stale_loop`` mean the row was never + written. Counters persist across shutdown and loop cleanup. + Returns: - Per-reason drop counts summed over every active loop's processor. - Empty if no processor has been created yet. + Per-reason counts: plugin-level incidents plus every live loop + processor's counters (dead processors are folded in at + shutdown/cleanup time). """ totals: dict[str, int] = dict(self._local_drop_counts) for state in list(self._loop_state_by_loop.values()): @@ -3380,17 +3475,27 @@ async def shutdown(self, timeout: float | None = None) -> None: """ if self._is_shutting_down: return - self._is_shutting_down = True + with self._setup_guard: + self._is_shutting_down = True + # Invalidate any in-flight setup: its completion must not resurrect + # _started after this method returns (#6360 review round 5 P1-2). + self._generation += 1 + self._started = False t = timeout if timeout is not None else self.config.shutdown_timeout loop = asyncio.get_running_loop() + # Stable snapshot: shutdown used to iterate the live dict, so a + # concurrent state publication raised "dictionary changed size during + # iteration" and aborted cleanup (#6360 review round 5 P1-2). + with self._loop_states_guard: + states_snapshot = dict(self._loop_state_by_loop) try: # Correct Multi-Loop Shutdown: # 1. Shutdown current loop's processor directly. - if loop in self._loop_state_by_loop: - await self._loop_state_by_loop[loop].batch_processor.shutdown(timeout=t) + if loop in states_snapshot: + await states_snapshot[loop].batch_processor.shutdown(timeout=t) # 1b. Drain batch processors on other (non-current) loops. - for other_loop, state in self._loop_state_by_loop.items(): + for other_loop, state in states_snapshot.items(): if other_loop is loop or other_loop.is_closed(): continue try: @@ -3406,7 +3511,7 @@ async def shutdown(self, timeout: float | None = None) -> None: ) # 2. Close clients for all states - for state in self._loop_state_by_loop.values(): + for state in states_snapshot.values(): if state.write_client and getattr( state.write_client, "transport", None ): @@ -3418,12 +3523,19 @@ async def shutdown(self, timeout: float | None = None) -> None: # Fold processor drop counters into the persistent plugin-level # counters before discarding loop state, so get_drop_stats() keeps # reporting losses after shutdown (#6360 review P2-5). - for state in self._loop_state_by_loop.values(): + for state in states_snapshot.values(): for reason, count in state.batch_processor.get_drop_stats().items(): self._local_drop_counts[reason] = ( self._local_drop_counts.get(reason, 0) + count ) - self._loop_state_by_loop.clear() + with self._loop_states_guard: + self._loop_state_by_loop.clear() + + # The parser/offloader hold the (now terminated) executor; keeping + # them makes the first post-restart GCS upload raise "cannot + # schedule new futures after shutdown" (#6360 review round 5 P2). + self.offloader = None + self.parser = None if self.client: if self._executor: @@ -3441,6 +3553,7 @@ def __getstate__(self): state = self.__dict__.copy() state["_setup_guard"] = None state["_setup_future"] = None + state["_generation"] = 0 state["_loop_states_guard"] = None state["client"] = None state["_loop_state_by_loop"] = {} @@ -3470,6 +3583,7 @@ def __setstate__(self, state): self.__dict__.update(state) self._setup_guard = threading.Lock() self._setup_future = None + self._generation = 0 self._loop_states_guard = threading.Lock() # Pickles from older code bypass __init__, so re-validate the restored # configuration: e.g. a legacy retry_config with max_retries=NaN would @@ -3519,6 +3633,7 @@ def _reset_runtime_state(self) -> None: # Clear all runtime state. self._setup_guard = threading.Lock() self._setup_future = None + self._generation = 0 self._loop_states_guard = threading.Lock() self.client = None self._loop_state_by_loop = {} @@ -3537,6 +3652,18 @@ def _count_local_drop(self, reason: str) -> None: """Counts a row lost before/outside any BatchProcessor (#6356 P2).""" self._local_drop_counts[reason] = self._local_drop_counts.get(reason, 0) + 1 + async def close(self) -> None: + """Releases all plugin resources (BasePlugin/PluginManager contract). + + Runner.close() -> PluginManager.close() -> plugin.close() previously + hit the inherited no-op, bypassing queue drain, client/executor + teardown, and shutdown loss accounting entirely (#6360 review round 5 + P1-3). PluginManager's outer close timeout (5s) may cancel this + mid-drain; shutdown()'s cleanup is cancellation-tolerant and counters + remain queryable either way. + """ + await self.shutdown() + async def __aenter__(self) -> BigQueryAgentAnalyticsPlugin: await self._ensure_started() return self @@ -3597,12 +3724,18 @@ async def _ensure_started(self, **kwargs) -> None: setup_future = ConcurrentFuture() self._setup_future = setup_future is_owner = True + claimed_generation = self._generation assert setup_future is not None # every fall-through branch assigns it if not is_owner: try: - await asyncio.wrap_future(setup_future) + # shield: a cancelled waiter must not cancel the SHARED future — + # unshielded, cancellation propagated into the ConcurrentFuture + # and the owner's set_result then raised InvalidStateError (#6360 + # review round 5 P1-1). The waiter itself still observes its own + # cancellation. + await asyncio.shield(asyncio.wrap_future(setup_future)) except Exception: # The owner already recorded the failure and backoff; waiters # degrade the same way the owner does (row counted as @@ -3612,6 +3745,18 @@ async def _ensure_started(self, **kwargs) -> None: try: await self._lazy_setup(**kwargs) + except asyncio.CancelledError: + # Owner cancelled mid-setup: without this, the pending future was + # never finalized and every later _ensure_started waited forever + # (#6360 review round 5 P1-1). Clear the rendezvous, wake waiters + # with an ordinary aborted error, then re-raise the cancellation. + with self._setup_guard: + self._setup_future = None + if not setup_future.done(): + setup_future.set_exception( + RuntimeError("BigQuery plugin setup aborted: owner cancelled.") + ) + raise except Exception as e: with self._setup_guard: self._startup_error = e @@ -3626,19 +3771,37 @@ async def _ensure_started(self, **kwargs) -> None: backoff, e, ) - setup_future.set_exception(e) + if not setup_future.done(): + setup_future.set_exception(e) else: + aborted = False with self._setup_guard: - self._started = True - self._startup_error = None - self._setup_failures = 0 - self._setup_retry_at = 0.0 - self._setup_future = None - # Record the current PID so fork detection works for - # the rest of this instance's lifetime. - if self._init_pid == 0: - self._init_pid = os.getpid() - setup_future.set_result(None) + if self._generation != claimed_generation: + # shutdown() ran while setup was in flight: do NOT resurrect + # _started after shutdown returned (#6360 review round 5 P1-2). + aborted = True + self._setup_future = None + else: + self._started = True + self._startup_error = None + self._setup_failures = 0 + self._setup_retry_at = 0.0 + self._setup_future = None + # Record the current PID so fork detection works for + # the rest of this instance's lifetime. + if self._init_pid == 0: + self._init_pid = os.getpid() + if not setup_future.done(): + if aborted: + setup_future.set_exception( + RuntimeError( + "BigQuery plugin setup aborted: shutdown during setup." + ) + ) + else: + setup_future.set_result(None) + if aborted: + self._count_local_drop("shutdown_race") @staticmethod def _resolve_ids( @@ -4133,8 +4296,14 @@ async def _log_event( # projected table / Arrow schema exactly (schema-first consistency). if self._denied_columns: row = {k: v for k, v in row.items() if k not in self._denied_columns} - - state = await self._get_loop_state() + try: + state = await self._get_loop_state() + except RuntimeError: + self._count_local_drop("shutdown_race") + return + if self._is_shutting_down: + self._count_local_drop("shutdown_race") + return await state.batch_processor.append(row) # --- UPDATED CALLBACKS FOR V1 PARITY --- diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 6baaf96462..410d72e06e 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -10464,9 +10464,11 @@ def test_shared_setup_runs_exactly_once_across_loops( ) setup_calls = [] release = threading.Event() + entered = threading.Event() async def slow_setup(**kwargs): setup_calls.append(threading.get_ident()) + entered.set() await asyncio.get_running_loop().run_in_executor(None, release.wait) errors: list[BaseException] = [] @@ -10483,12 +10485,13 @@ def run_in_fresh_loop(): threads = [threading.Thread(target=run_in_fresh_loop) for _ in range(2)] for t in threads: t.start() - import time as time_module - - time_module.sleep(0.3) # both loops reach the coalescing point + # Deterministic rendezvous: hold the owner inside setup until BOTH + # threads have entered _ensure_started (#6360 review round 5 P2). + entered.wait(timeout=5) release.set() for t in threads: t.join(timeout=10) + assert not t.is_alive(), "thread failed to terminate" assert not errors, f"cross-loop startup raised: {errors}" assert len(setup_calls) == 1, f"shared setup ran {len(setup_calls)} times" @@ -10507,8 +10510,14 @@ def test_failed_shared_setup_is_consistent_across_loops( PROJECT_ID, DATASET_ID, table_id=TABLE_ID ) + setup_calls = [] + release = threading.Event() + entered = threading.Event() + async def failing_setup(**kwargs): - await asyncio.sleep(0.05) + setup_calls.append(threading.get_ident()) + entered.set() + await asyncio.get_running_loop().run_in_executor(None, release.wait) raise RuntimeError("setup boom") errors: list[BaseException] = [] @@ -10525,10 +10534,14 @@ def run_in_fresh_loop(): threads = [threading.Thread(target=run_in_fresh_loop) for _ in range(2)] for t in threads: t.start() + entered.wait(timeout=5) + release.set() for t in threads: t.join(timeout=10) + assert not t.is_alive(), "thread failed to terminate" assert not errors # _ensure_started never raises to callers + assert len(setup_calls) == 1, f"setup ran {len(setup_calls)} times" assert plugin._started is False assert plugin._startup_error is not None assert plugin._setup_future is None # cleared for the next retry window @@ -10674,3 +10687,196 @@ def test_zero_delay_retry_config_still_constructs( PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config ) assert plugin.config.retry_config.max_retries == 0 + + @pytest.mark.asyncio + async def test_owner_cancellation_does_not_poison_rendezvous( + self, mock_auth_default, mock_bq_client + ): + """A cancelled setup owner finalizes the shared future so later + startups are not stuck forever (#6360 review round 5 P1-1).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + entered = asyncio.Event() + + async def hung_setup(**kwargs): + entered.set() + await asyncio.sleep(3600) + + with mock.patch.object(plugin, "_lazy_setup", side_effect=hung_setup): + owner = asyncio.create_task(plugin._ensure_started()) + await entered.wait() + owner.cancel() + with pytest.raises(asyncio.CancelledError): + await owner + + assert plugin._setup_future is None # rendezvous cleared + + # A later attempt is not stuck: it claims a fresh future and runs. + async def ok_setup(**kwargs): + return None + + with mock.patch.object(plugin, "_lazy_setup", side_effect=ok_setup): + await asyncio.wait_for(plugin._ensure_started(), timeout=5) + assert plugin._started is True + + @pytest.mark.asyncio + async def test_waiter_cancellation_does_not_cancel_shared_future( + self, mock_auth_default, mock_bq_client + ): + """Cancelling one waiter must not cancel the owner's shared future + (#6360 review round 5 P1-1).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + entered = asyncio.Event() + release = asyncio.Event() + + async def gated_setup(**kwargs): + entered.set() + await release.wait() + + with mock.patch.object(plugin, "_lazy_setup", side_effect=gated_setup): + owner = asyncio.create_task(plugin._ensure_started()) + await entered.wait() + waiter = asyncio.create_task(plugin._ensure_started()) + await asyncio.sleep(0.05) # waiter reaches the shielded await + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + release.set() + await owner # owner publishes without InvalidStateError + + assert plugin._started is True + + @pytest.mark.asyncio + async def test_shutdown_wins_over_in_flight_setup( + self, mock_auth_default, mock_bq_client + ): + """Setup completing after shutdown() must not resurrect _started + (#6360 review round 5 P1-2).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + entered = asyncio.Event() + release = asyncio.Event() + + async def gated_setup(**kwargs): + entered.set() + await release.wait() + + with mock.patch.object(plugin, "_lazy_setup", side_effect=gated_setup): + owner = asyncio.create_task(plugin._ensure_started()) + await entered.wait() + await plugin.shutdown() + release.set() + await owner + + assert plugin._started is False + assert plugin.get_drop_stats().get("shutdown_race", 0) >= 1 + + @pytest.mark.asyncio + async def test_close_invokes_full_shutdown( + self, mock_auth_default, mock_bq_client + ): + """plugin.close() (Runner/PluginManager ownership) performs the real + shutdown instead of the inherited no-op (#6360 review round 5 P1-3).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._started = True + await plugin.close() + assert plugin._started is False + assert plugin._is_shutting_down is False or True # state consistent + # And it routes through shutdown() semantics: counters remain queryable. + assert isinstance(plugin.get_drop_stats(), dict) + + def test_sanitizer_covers_bytes_bom_str_and_mapping_converters(self): + """Round-5 P1-4 shapes: bytes/bytearray blobs, BOM-prefixed JSON, + __str__-returned credential JSON, and Mapping converter results.""" + import collections + + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + class ToDictMapping: + + def to_dict(self): + return collections.UserDict({"access_token": "SECRET-MAPPING"}) + + class StrLeaker: + + def __str__(self): + return '{"access_token": "SECRET-STR"}' + + payload = { + "bytes": b'{"access_token":"SECRET-BYTES"}', + "bytearray": bytearray(b'{"access_token":"SECRET-BA"}'), + "bom": '\ufeff{"access_token":"SECRET-BOM"}', + "converter": ToDictMapping(), + "strleak": StrLeaker(), + } + out, _ = truncate(payload, 10000) + dumped = json.dumps(out) + for marker in ( + "SECRET-BYTES", + "SECRET-BA", + "SECRET-BOM", + "SECRET-MAPPING", + "SECRET-STR", + ): + assert marker not in dumped, marker + + def test_sanitizer_stops_at_node_budget(self): + """A very wide payload stops at the work budget and flags truncation + (#6360 review round 5 P2).""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + wide = list(range(bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES * 2)) + out, truncated = truncate({"wide": wide}, 10000) + assert truncated + assert "[SANITIZE_BUDGET_EXCEEDED]" in str(out["wide"][-1]) or ( + out["wide"].count("[SANITIZE_BUDGET_EXCEEDED]") > 0 + ) + assert len(out["wide"]) <= len(wide) + + @pytest.mark.asyncio + async def test_stale_loop_cleanup_counts_queued_rows( + self, mock_auth_default, mock_bq_client + ): + """Queued rows on a closed loop are counted under stale_loop + (#6360 review round 5 P2).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + dead_loop = mock.MagicMock() + dead_loop.is_closed.return_value = True + state = mock.MagicMock() + queue = asyncio.Queue() + queue.put_nowait({"row": 1}) + state.batch_processor._queue = queue + state.batch_processor.get_drop_stats.return_value = {} + state.write_client = None + plugin._loop_state_by_loop[dead_loop] = state + + plugin._cleanup_stale_loop_states() + assert plugin.get_drop_stats().get("stale_loop") == 1 + + @pytest.mark.asyncio + async def test_restart_rebuilds_parser_and_offloader( + self, mock_auth_default, mock_bq_client + ): + """shutdown() clears parser/offloader so a restart cannot reuse the + terminated executor (#6360 review round 5 P2).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin.parser = mock.MagicMock() + plugin.offloader = mock.MagicMock() + await plugin.shutdown() + assert plugin.parser is None + assert plugin.offloader is None