diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 65275d3d93..8ca3381707 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -16,6 +16,8 @@ import asyncio import atexit +import collections.abc +from concurrent.futures import Future as ConcurrentFuture from concurrent.futures import ThreadPoolExecutor import contextvars import dataclasses @@ -26,6 +28,7 @@ import functools import json import logging +import math import mimetypes import os import traceback as traceback_module @@ -38,6 +41,7 @@ import random import re +import threading import time from types import MappingProxyType from typing import Any @@ -410,6 +414,170 @@ 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]" + +# 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 + +# 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, + budget: Optional[list[int]] = None, +) -> 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). 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("\ufeff \t\r\n") + 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(stripped) > 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 + # 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(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, budget + ) + if sanitized == parsed and not saw_duplicate_key: + return value, False + return json.dumps(sanitized), True + except (TypeError, ValueError, RecursionError, MemoryError): + # 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 + + +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). + + 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. + """ + _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"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( + "max_content_length must be -1 (unlimited) or >= 1, got" + f" {config.max_content_length}." + ) + 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, -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, -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" + 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 @@ -420,7 +588,11 @@ def _extract_tool_declarations( def _recursive_smart_truncate( - obj: Any, max_len: int, seen: Optional[set[int]] = None + obj: Any, + max_len: int, + seen: Optional[set[int]] = None, + depth: int = 0, + budget: Optional[list[int]] = None, ) -> tuple[Any, bool]: """Recursively truncates string values within a dict or list. @@ -431,12 +603,28 @@ 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). """ 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 + # model_dump()/dict()/to_dict() returns a fresh wrapper — unittest Mocks + # being the canonical case). Without this cap such graphs recurse + # 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]", True obj_id = id(obj) if obj_id in seen: @@ -444,7 +632,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") @@ -456,10 +644,29 @@ def _recursive_smart_truncate( try: if isinstance(obj, str): + 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, dict): + 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 + # (#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. @@ -471,7 +678,9 @@ 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, budget + ) if trunc: truncated_any = True new_dict[k] = val @@ -481,40 +690,76 @@ 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, budget + ) 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) + return _recursive_smart_truncate( + as_dict, max_len, seen, depth + 1, budget + ) 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) + dumped = obj.model_dump() + 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 + # Pydantic v1 (same progress requirement as above). try: - return _recursive_smart_truncate(obj.dict(), max_len, seen) + dumped = obj.dict() + 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 + # Common pattern for custom objects (same progress requirement). try: - return _recursive_smart_truncate(obj.to_dict(), max_len, seen) + dumped = obj.to_dict() + 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) @@ -1195,6 +1440,7 @@ def __init__( "retry_exhausted": 0, "non_retryable": 0, "unexpected_error": 0, + "shutdown_timeout": 0, } async def flush(self) -> None: @@ -1362,8 +1608,21 @@ async def _batch_writer(self) -> None: except asyncio.TimeoutError: continue except asyncio.CancelledError: + # 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), + ) 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 @@ -1539,6 +1798,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) @@ -1563,6 +1839,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) # ============================================================================== @@ -1613,7 +1904,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}" @@ -1644,9 +1942,30 @@ 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, + parse_uid: str = "", + content_ordinal: int = 0, ) -> tuple[str, list[dict[str, Any]], bool]: - """Parses a Content or Part object into summary text and content parts.""" + """Parses a Content or Part object into summary text and content parts. + + ``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. + + ``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 content_parts = [] is_truncated = False summary_text = [] @@ -1673,7 +1992,10 @@ async def _parse_content_object( elif hasattr(part, "inline_data") and part.inline_data: if self.offloader: ext = mimetypes.guess_extension(part.inline_data.mime_type) or ".bin" - path = f"{datetime.now().date()}/{self.trace_id}/{self.span_id}_p{idx}{ext}" + path = ( + f"{datetime.now().date()}/{trace_id}/{span_id}_{parse_uid}" + f"_c{content_ordinal}_p{idx}{ext}" + ) try: uri = await self.offloader.upload_content( part.inline_data.data, part.inline_data.mime_type, path @@ -1712,7 +2034,10 @@ async def _parse_content_object( if self.offloader and (exceeds_inline_byte_limit or exceeds_char_limit): # Text is too big, treat as file - path = f"{datetime.now().date()}/{self.trace_id}/{self.span_id}_p{idx}.txt" + path = ( + f"{datetime.now().date()}/{trace_id}/{span_id}_{parse_uid}" + f"_c{content_ordinal}_p{idx}.txt" + ) try: uri = await self.offloader.upload_content( part.text, "text/plain", path @@ -1760,8 +2085,27 @@ 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 + # 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 json_payload = {} content_parts = [] is_truncated = False @@ -1777,9 +2121,15 @@ def process_text(t: str) -> tuple[str, bool]: if isinstance(content.contents, list) else [content.contents] ) - for c in contents: + for content_idx, c in enumerate(contents): role = getattr(c, "role", "unknown") - summary, parts, trunc = await self._parse_content_object(c) + summary, parts, trunc = await self._parse_content_object( + c, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + content_ordinal=content_idx, + ) if trunc: is_truncated = True content_parts.extend(parts) @@ -1797,14 +2147,25 @@ def process_text(t: str) -> tuple[str, bool]: is_truncated = True json_payload["system_prompt"] = truncated_si else: - summary, parts, trunc = await self._parse_content_object(si) + summary, parts, trunc = await self._parse_content_object( + si, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + content_ordinal=len(contents), + ) if trunc: is_truncated = True content_parts.extend(parts) json_payload["system_prompt"] = summary elif isinstance(content, (types.Content, types.Part)): - summary, parts, trunc = await self._parse_content_object(content) + summary, parts, trunc = await self._parse_content_object( + content, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + ) return {"text_summary": summary}, parts, trunc elif isinstance(content, (dict, list)): @@ -2460,10 +2821,27 @@ 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 + # 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 + # 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() self._credentials = credentials self.client = None self._loop_state_by_loop: dict[asyncio.AbstractEventLoop, _LoopState] = {} @@ -2478,14 +2856,61 @@ 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): 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( "Cleaning up stale loop state for closed loop %s (id=%s).", loop, id(loop), ) - del self._loop_state_by_loop[loop] + # Preserve the dead processor's loss accounting before discarding it + # (#6360 review round 2 P2-6), mirroring shutdown(). + for reason, count in state.batch_processor.get_drop_stats().items(): + 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"), @@ -2570,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] @@ -2619,7 +3049,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)) @@ -2647,11 +3078,18 @@ 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] = {} + 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 @@ -2680,7 +3118,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) @@ -2777,29 +3220,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( @@ -2933,6 +3382,16 @@ def _maybe_upgrade_schema(self, existing_table: bigquery.Table) -> None: e, exc_info=True, ) + 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. @@ -3016,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: @@ -3042,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 ): @@ -3051,7 +3520,22 @@ async def shutdown(self, timeout: float | None = None) -> None: except Exception: pass - self._loop_state_by_loop.clear() + # 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 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 + ) + 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: @@ -3067,7 +3551,10 @@ 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_guard"] = None + state["_setup_future"] = None + state["_generation"] = 0 + state["_loop_states_guard"] = None state["client"] = None state["_loop_state_by_loop"] = {} state["_write_stream_name"] = None @@ -3076,6 +3563,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 @@ -3085,7 +3574,21 @@ 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) + 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_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 + # 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. @@ -3128,7 +3631,10 @@ def _reset_runtime_state(self) -> None: pass # Clear all runtime state. - self._setup_lock = None + 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 = {} self._write_stream_name = None @@ -3137,9 +3643,27 @@ 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 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 @@ -3148,7 +3672,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 @@ -3158,23 +3694,114 @@ 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: - # 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: - if not self._started: - try: - await self._lazy_setup(**kwargs) - self._started = True - self._startup_error = 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() - except Exception as e: - self._startup_error = e - logger.error("Failed to initialize BigQuery Plugin: %s", 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 + claimed_generation = self._generation + + assert setup_future is not None # every fall-through branch assigns it + + if not is_owner: + try: + # 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 + # setup_unavailable by the caller). + pass + return + + 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 + 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, + ) + if not setup_future.done(): + setup_future.set_exception(e) + else: + aborted = False + with self._setup_guard: + 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( @@ -3561,6 +4188,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: @@ -3571,7 +4201,18 @@ async def _log_event( try: raw_content = self.config.content_formatter(raw_content, event_type) except Exception as e: - logger.warning("Content formatter failed: %s", e) + # Fail CLOSED (#6356 P1-1): the formatter is a redaction/privacy + # boundary, so its failure must never fall back to the unformatted + # payload. Log only the exception CLASS — the message or a + # traceback (exc_info) could embed the protected content itself. + logger.warning( + "Content formatter (%s) failed for event %s; writing sentinel" + " instead of original content.", + type(e).__name__, + event_type, + ) + 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 +4231,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 +4252,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) @@ -3642,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 78caca075a..410d72e06e 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 @@ -759,8 +760,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 +1199,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() @@ -2547,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_future is None assert unpickled._executor is None # Start the plugin await plugin._ensure_started() @@ -2558,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_future is None assert unpickled_started._executor is None assert not unpickled_started._loop_state_by_loop finally: @@ -3354,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( @@ -3372,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( @@ -5103,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 = [ @@ -5113,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.""" @@ -6033,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_future = mock.MagicMock() # Keep pure-data fields plugin._schema = ["kept"] plugin.arrow_schema = "kept_arrow" @@ -6048,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_future is None # Pure-data fields are preserved assert plugin._schema == ["kept"] assert plugin.arrow_schema == "kept_arrow" @@ -6386,12 +6411,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() @@ -7629,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 @@ -9698,3 +9729,1154 @@ 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=-1.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. + 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( + 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 + + 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 + ) + + 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 + + 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() + + 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() + 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] = [] + 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() + # 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" + 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 + ) + + setup_calls = [] + release = threading.Event() + entered = threading.Event() + + async def failing_setup(**kwargs): + 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] = [] + + 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() + 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 + + @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 + + @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